Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 8a17eff

Browse files
dannykoppingdannykopping
committed
feat: store coder_ai_task resource state
Signed-off-by: Danny Kopping <[email protected]>
1 parent e1ad75b commit 8a17eff

File tree

4 files changed

+125
-3
lines changed

4 files changed

+125
-3
lines changed

‎coderd/provisionerdserver/provisionerdserver.go

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import (
2828
protobuf "google.golang.org/protobuf/proto"
2929

3030
"cdr.dev/slog"
31+
3132
"github.com/coder/coder/v2/coderd/util/slice"
3233

3334
"github.com/coder/coder/v2/codersdk/drpcsdk"
@@ -1654,6 +1655,17 @@ func (s *server) completeTemplateImportJob(ctx context.Context, job database.Pro
16541655
if err != nil {
16551656
return xerrors.Errorf("update template version external auth providers: %w", err)
16561657
}
1658+
err = db.UpdateTemplateVersionAITaskByJobID(ctx, database.UpdateTemplateVersionAITaskByJobIDParams{
1659+
JobID: jobID,
1660+
HasAITask: sql.NullBool{
1661+
Bool: jobType.TemplateImport.HasAiTasks,
1662+
Valid: true,
1663+
},
1664+
UpdatedAt: now,
1665+
})
1666+
if err != nil {
1667+
return xerrors.Errorf("update template version external auth providers: %w", err)
1668+
}
16571669

16581670
// Process terraform values
16591671
plan := jobType.TemplateImport.Plan
@@ -1866,6 +1878,34 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
18661878
}
18671879
}
18681880

1881+
var sidebarAppID uuid.NullUUID
1882+
if len(jobType.WorkspaceBuild.AiTasks) == 1 {
1883+
task := jobType.WorkspaceBuild.AiTasks[0]
1884+
if task.SidebarApp == nil {
1885+
return xerrors.Errorf("update ai task: sidebar app is nil")
1886+
}
1887+
1888+
id, err := uuid.Parse(task.SidebarApp.Id)
1889+
if err != nil {
1890+
return xerrors.Errorf("parse sidebar app id: %w", err)
1891+
}
1892+
1893+
sidebarAppID = uuid.NullUUID{UUID: id, Valid: true}
1894+
}
1895+
1896+
err = db.UpdateWorkspaceBuildAITaskByID(ctx, database.UpdateWorkspaceBuildAITaskByIDParams{
1897+
ID: workspaceBuild.ID,
1898+
HasAITask: sql.NullBool{
1899+
Bool: len(jobType.WorkspaceBuild.AiTasks) > 0,
1900+
Valid: true,
1901+
},
1902+
SidebarAppID: sidebarAppID,
1903+
UpdatedAt: now,
1904+
})
1905+
if err != nil {
1906+
return xerrors.Errorf("update workspace build ai tasks flag: %w", err)
1907+
}
1908+
18691909
// Insert timings inside the transaction now
18701910
// nolint:exhaustruct // The other fields are set further down.
18711911
params := database.InsertProvisionerJobTimingsParams{
@@ -2570,8 +2610,13 @@ func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.
25702610
openIn = database.WorkspaceAppOpenInSlimWindow
25712611
}
25722612

2613+
id, err := uuid.Parse(app.Id)
2614+
if err != nil {
2615+
return xerrors.Errorf("parse app uuid: %w", err)
2616+
}
2617+
25732618
dbApp, err := db.InsertWorkspaceApp(ctx, database.InsertWorkspaceAppParams{
2574-
ID: uuid.New(),
2619+
ID: id,
25752620
CreatedAt: dbtime.Now(),
25762621
AgentID: dbAgent.ID,
25772622
Slug: slug,

‎provisioner/terraform/executor.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -361,6 +361,8 @@ func (e *executor) plan(ctx, killCtx context.Context, env, vars []string, logr l
361361
Plan: planJSON,
362362
ResourceReplacements: resReps,
363363
ModuleFiles: moduleFiles,
364+
HasAiTasks: state.HasAITasks,
365+
AiTasks: state.AITasks,
364366
}
365367

366368
return msg, nil
@@ -577,6 +579,7 @@ func (e *executor) apply(
577579
ExternalAuthProviders: state.ExternalAuthProviders,
578580
State: stateContent,
579581
Timings: e.timings.aggregate(),
582+
AiTasks: state.AITasks,
580583
}, nil
581584
}
582585

‎provisioner/terraform/resources.go

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"strings"
88

99
"github.com/awalterschulze/gographviz"
10+
"github.com/google/uuid"
1011
tfjson "github.com/hashicorp/terraform-json"
1112
"github.com/mitchellh/mapstructure"
1213
"golang.org/x/xerrors"
@@ -93,6 +94,7 @@ type agentDisplayAppsAttributes struct {
9394

9495
// A mapping of attributes on the "coder_app" resource.
9596
type agentAppAttributes struct {
97+
ID string `mapstructure:"id"`
9698
AgentID string `mapstructure:"agent_id"`
9799
// Slug is required in terraform, but to avoid breaking existing users we
98100
// will default to the resource name if it is not specified.
@@ -160,10 +162,29 @@ type State struct {
160162
Parameters []*proto.RichParameter
161163
Presets []*proto.Preset
162164
ExternalAuthProviders []*proto.ExternalAuthProviderResource
165+
AITasks []*proto.AITask
166+
HasAITasks bool
163167
}
164168

165169
var ErrInvalidTerraformAddr = xerrors.New("invalid terraform address")
166170

171+
// hasAITaskResources is used to determine if a template has *any* `coder_ai_task` resources defined. During template
172+
// import, it's possible that none of these have `count=1` since count may be dependent on the value of a `coder_parameter`
173+
// or something else.
174+
// We need to know at template import if these resources exist to inform the frontend of their existence.
175+
func hasAITaskResources(graph *gographviz.Graph) bool {
176+
for _, node := range graph.Nodes.Lookup {
177+
// Check if this node is a coder_ai_task resource
178+
if label, exists := node.Attrs["label"]; exists {
179+
labelValue := strings.Trim(label, `"`)
180+
if strings.Contains(labelValue, "coder_ai_task.") {
181+
return true
182+
}
183+
}
184+
}
185+
return false
186+
}
187+
167188
// ConvertState consumes Terraform state and a GraphViz representation
168189
// produced by `terraform graph` to produce resources consumable by Coder.
169190
// nolint:gocognit // This function makes more sense being large for now, until refactored.
@@ -187,6 +208,7 @@ func ConvertState(ctx context.Context, modules []*tfjson.StateModule, rawGraph s
187208
// Extra array to preserve the order of rich parameters.
188209
tfResourcesRichParameters := make([]*tfjson.StateResource, 0)
189210
tfResourcesPresets := make([]*tfjson.StateResource, 0)
211+
tfResourcesAITasks := make([]*tfjson.StateResource, 0)
190212
var findTerraformResources func(mod *tfjson.StateModule)
191213
findTerraformResources = func(mod *tfjson.StateModule) {
192214
for _, module := range mod.ChildModules {
@@ -199,6 +221,9 @@ func ConvertState(ctx context.Context, modules []*tfjson.StateModule, rawGraph s
199221
if resource.Type == "coder_workspace_preset" {
200222
tfResourcesPresets = append(tfResourcesPresets, resource)
201223
}
224+
if resource.Type == "coder_ai_task" {
225+
tfResourcesAITasks = append(tfResourcesAITasks, resource)
226+
}
202227

203228
label := convertAddressToLabel(resource.Address)
204229
if tfResourcesByLabel[label] == nil {
@@ -522,7 +547,17 @@ func ConvertState(ctx context.Context, modules []*tfjson.StateModule, rawGraph s
522547
continue
523548
}
524549

550+
id := attrs.ID
551+
if id == "" {
552+
// This should never happen since the "id" attribute is set on creation:
553+
// https://github.com/coder/terraform-provider-coder/blob/cfa101df4635e405e66094fa7779f9a89d92f400/provider/app.go#L37
554+
logger.Warn(ctx, "coder_app's id was unexpectedly empty", slog.F("name", attrs.Name))
555+
556+
id = uuid.NewString()
557+
}
558+
525559
agent.Apps = append(agent.Apps, &proto.App{
560+
Id: id,
526561
Slug: attrs.Slug,
527562
DisplayName: attrs.DisplayName,
528563
Command: attrs.Command,
@@ -940,6 +975,27 @@ func ConvertState(ctx context.Context, modules []*tfjson.StateModule, rawGraph s
940975
)
941976
}
942977

978+
// This will only pick up resources which will actually be created.
979+
aiTasks := make([]*proto.AITask, 0, len(tfResourcesAITasks))
980+
for _, resource := range tfResourcesAITasks {
981+
var task provider.AITask
982+
err = mapstructure.Decode(resource.AttributeValues, &task)
983+
if err != nil {
984+
return nil, xerrors.Errorf("decode coder_ai_task attributes: %w", err)
985+
}
986+
987+
if len(task.SidebarApp) < 1 {
988+
return nil, xerrors.Errorf("coder_ai_task has no sidebar_app defined")
989+
}
990+
991+
aiTasks = append(aiTasks, &proto.AITask{
992+
Id: task.ID,
993+
SidebarApp: &proto.AITaskSidebarApp{
994+
Id: task.SidebarApp[0].ID,
995+
},
996+
})
997+
}
998+
943999
// A map is used to ensure we don't have duplicates!
9441000
externalAuthProvidersMap := map[string]*proto.ExternalAuthProviderResource{}
9451001
for _, tfResources := range tfResourcesByLabel {
@@ -975,6 +1031,8 @@ func ConvertState(ctx context.Context, modules []*tfjson.StateModule, rawGraph s
9751031
Parameters: parameters,
9761032
Presets: presets,
9771033
ExternalAuthProviders: externalAuthProviders,
1034+
HasAITasks: hasAITaskResources(graph),
1035+
AITasks: aiTasks,
9781036
}, nil
9791037
}
9801038

‎provisionerd/runner/runner.go

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,13 @@ import (
77
"errors"
88
"fmt"
99
"reflect"
10+
"slices"
1011
"strings"
1112
"sync"
1213
"sync/atomic"
1314
"time"
1415

16+
"github.com/coder/terraform-provider-coder/v2/provider"
1517
"github.com/google/uuid"
1618
"github.com/prometheus/client_golang/prometheus"
1719
"go.opentelemetry.io/otel/attribute"
@@ -584,8 +586,6 @@ func (r *Runner) runTemplateImport(ctx context.Context) (*proto.CompletedJob, *p
584586
externalAuthProviderNames = append(externalAuthProviderNames, it.Id)
585587
}
586588

587-
// fmt.Println("completed job: template import: graph:", startProvision.Graph)
588-
589589
return &proto.CompletedJob{
590590
JobId: r.job.JobId,
591591
Type: &proto.CompletedJob_TemplateImport_{
@@ -603,6 +603,7 @@ func (r *Runner) runTemplateImport(ctx context.Context) (*proto.CompletedJob, *p
603603
ModuleFiles: startProvision.ModuleFiles,
604604
// ModuleFileHash will be populated if the file is uploaded async
605605
ModuleFilesHash: []byte{},
606+
HasAiTasks: startProvision.HasAITasks,
606607
},
607608
},
608609
}, nil
@@ -666,6 +667,7 @@ type templateImportProvision struct {
666667
Presets []*sdkproto.Preset
667668
Plan json.RawMessage
668669
ModuleFiles []byte
670+
HasAITasks bool
669671
}
670672

671673
// Performs a dry-run provision when importing a template.
@@ -799,6 +801,15 @@ func (r *Runner) runTemplateImportProvisionWithRichParameters(
799801
}
800802
}
801803

804+
if c.HasAiTasks {
805+
hasPromptParam := slices.ContainsFunc(c.Parameters, func(param *sdkproto.RichParameter) bool {
806+
return param.Name == provider.TaskPromptParameterName
807+
})
808+
if !hasPromptParam {
809+
return nil, xerrors.Errorf("coder_parameter named '%s' is required when 'coder_ai_task' resource is defined", provider.TaskPromptParameterName)
810+
}
811+
}
812+
802813
return &templateImportProvision{
803814
Resources: c.Resources,
804815
Parameters: c.Parameters,
@@ -807,6 +818,7 @@ func (r *Runner) runTemplateImportProvisionWithRichParameters(
807818
Presets: c.Presets,
808819
Plan: c.Plan,
809820
ModuleFiles: moduleFilesData,
821+
HasAITasks: c.HasAiTasks,
810822
}, nil
811823
default:
812824
return nil, xerrors.Errorf("invalid message type %q received from provisioner",
@@ -1047,6 +1059,9 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
10471059
},
10481060
}
10491061
}
1062+
if len(planComplete.AiTasks) > 1 {
1063+
return nil, r.failedWorkspaceBuildf("only one 'coder_ai_task' resource can be provisioned per template")
1064+
}
10501065

10511066
r.logger.Info(context.Background(), "plan request successful",
10521067
slog.F("resource_count", len(planComplete.Resources)),
@@ -1124,6 +1139,7 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
11241139
Modules: planComplete.Modules,
11251140
// Resource replacements are discovered at plan time, only.
11261141
ResourceReplacements: planComplete.ResourceReplacements,
1142+
AiTasks: applyComplete.AiTasks,
11271143
},
11281144
},
11291145
}, nil

0 commit comments

Comments
 (0)