From a8f87c2625643b42e37f4fb12d48de5811cb427d Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 13 Oct 2025 12:57:06 +0300 Subject: [PATCH] feat(coderd): implement task to app linking (#20237) This change adds workspace build/agent/app linking to tasks and wires it into `wsbuilder` and `provisionerdserver`. Closes coder/internal#948 Closes coder/coder#20212 Closes coder/coder#19773 --- .../provisionerdserver/provisionerdserver.go | 81 +++++++++++++- .../provisionerdserver_test.go | 26 +++++ coderd/wsbuilder/wsbuilder.go | 16 +++ coderd/wsbuilder/wsbuilder_test.go | 105 ++++++++++++++++++ 4 files changed, 225 insertions(+), 3 deletions(-) diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go index f8d361f537..8e07eb4dc5 100644 --- a/coderd/provisionerdserver/provisionerdserver.go +++ b/coderd/provisionerdserver/provisionerdserver.go @@ -1964,18 +1964,41 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro } appIDs := make([]string, 0) + agentIDByAppID := make(map[string]uuid.UUID) agentTimeouts := make(map[time.Duration]bool) // A set of agent timeouts. // This could be a bulk insert to improve performance. for _, protoResource := range jobType.WorkspaceBuild.Resources { - for _, protoAgent := range protoResource.Agents { + for _, protoAgent := range protoResource.GetAgents() { + if protoAgent == nil { + continue + } + // By default InsertWorkspaceResource ignores the protoAgent.Id + // and generates a new one, but we will insert these using the + // InsertWorkspaceResourceWithAgentIDsFromProto option so that + // we can properly map agent IDs to app IDs. This is needed for + // task linking. + agentID := uuid.New() + protoAgent.Id = agentID.String() + dur := time.Duration(protoAgent.GetConnectionTimeoutSeconds()) * time.Second agentTimeouts[dur] = true for _, app := range protoAgent.GetApps() { appIDs = append(appIDs, app.GetId()) + agentIDByAppID[app.GetId()] = agentID } } - err = InsertWorkspaceResource(ctx, db, job.ID, workspaceBuild.Transition, protoResource, telemetrySnapshot) + err = InsertWorkspaceResource( + ctx, + db, + job.ID, + workspaceBuild.Transition, + protoResource, + telemetrySnapshot, + // Ensure that the agent IDs we set previously + // are written to the database. + InsertWorkspaceResourceWithAgentIDsFromProto(), + ) if err != nil { return xerrors.Errorf("insert provisioner job: %w", err) } @@ -1987,6 +2010,7 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro } var taskAppID uuid.NullUUID + var taskAgentID uuid.NullUUID var hasAITask bool var warnUnknownTaskAppID bool if tasks := jobType.WorkspaceBuild.GetAiTasks(); len(tasks) > 0 { @@ -2014,6 +2038,9 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro } taskAppID = uuid.NullUUID{UUID: id, Valid: true} + + agentID, ok := agentIDByAppID[appID] + taskAgentID = uuid.NullUUID{UUID: agentID, Valid: ok} } // This is a hacky workaround for the issue with tasks 'disappearing' on stop: @@ -2108,6 +2135,27 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro } } + if task, err := db.GetTaskByWorkspaceID(ctx, workspace.ID); err == nil { + // Irrespective of whether the agent or sidebar app is present, + // perform the upsert to ensure a link between the task and + // workspace build. Linking the task to the build is typically + // already established by wsbuilder. + _, err = db.UpsertTaskWorkspaceApp( + ctx, + database.UpsertTaskWorkspaceAppParams{ + TaskID: task.ID, + WorkspaceBuildNumber: workspaceBuild.BuildNumber, + WorkspaceAgentID: taskAgentID, + WorkspaceAppID: taskAppID, + }, + ) + if err != nil { + return xerrors.Errorf("upsert task workspace app: %w", err) + } + } else if !errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf("get task by workspace id: %w", err) + } + // Regardless of whether there is an AI task or not, update the field to indicate one way or the other since it // always defaults to nil. ONLY if has_ai_task=true MUST ai_task_sidebar_app_id be set. if err := db.UpdateWorkspaceBuildFlagsByID(ctx, database.UpdateWorkspaceBuildFlagsByIDParams{ @@ -2578,7 +2626,28 @@ func InsertWorkspacePresetAndParameters(ctx context.Context, db database.Store, return nil } -func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.UUID, transition database.WorkspaceTransition, protoResource *sdkproto.Resource, snapshot *telemetry.Snapshot) error { +type insertWorkspaceResourceOptions struct { + useAgentIDsFromProto bool +} + +// InsertWorkspaceResourceOption represents a functional option for +// InsertWorkspaceResource. +type InsertWorkspaceResourceOption func(*insertWorkspaceResourceOptions) + +// InsertWorkspaceResourceWithAgentIDsFromProto allows inserting agents into the +// database using the agent IDs defined in the proto resource. +func InsertWorkspaceResourceWithAgentIDsFromProto() InsertWorkspaceResourceOption { + return func(opts *insertWorkspaceResourceOptions) { + opts.useAgentIDsFromProto = true + } +} + +func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.UUID, transition database.WorkspaceTransition, protoResource *sdkproto.Resource, snapshot *telemetry.Snapshot, opt ...InsertWorkspaceResourceOption) error { + opts := &insertWorkspaceResourceOptions{} + for _, o := range opt { + o(opts) + } + resource, err := db.InsertWorkspaceResource(ctx, database.InsertWorkspaceResourceParams{ ID: uuid.New(), CreatedAt: dbtime.Now(), @@ -2675,6 +2744,12 @@ func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid. } agentID := uuid.New() + if opts.useAgentIDsFromProto { + agentID, err = uuid.Parse(prAgent.Id) + if err != nil { + return xerrors.Errorf("invalid agent ID format; must be uuid: %w", err) + } + } dbAgent, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{ ID: agentID, ParentID: uuid.NullUUID{}, diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go index 6409ba9b1d..0cfb83faa4 100644 --- a/coderd/provisionerdserver/provisionerdserver_test.go +++ b/coderd/provisionerdserver/provisionerdserver_test.go @@ -2850,6 +2850,8 @@ func TestCompleteJob(t *testing.T) { seedFunc func(context.Context, testing.TB, database.Store) error // If you need to insert other resources transition database.WorkspaceTransition input *proto.CompletedJob_WorkspaceBuild + isTask bool + expectTaskStatus database.TaskStatus expectHasAiTask bool expectUsageEvent bool } @@ -2862,6 +2864,7 @@ func TestCompleteJob(t *testing.T) { input: &proto.CompletedJob_WorkspaceBuild{ // No AiTasks defined. }, + isTask: false, expectHasAiTask: false, expectUsageEvent: false, }, @@ -2894,6 +2897,8 @@ func TestCompleteJob(t *testing.T) { }, }, }, + isTask: true, + expectTaskStatus: database.TaskStatusInitializing, expectHasAiTask: true, expectUsageEvent: true, }, @@ -2912,6 +2917,8 @@ func TestCompleteJob(t *testing.T) { }, }, }, + isTask: true, + expectTaskStatus: database.TaskStatusInitializing, expectHasAiTask: false, expectUsageEvent: false, }, @@ -2944,6 +2951,8 @@ func TestCompleteJob(t *testing.T) { }, }, }, + isTask: true, + expectTaskStatus: database.TaskStatusPaused, expectHasAiTask: true, expectUsageEvent: false, }, @@ -2955,6 +2964,8 @@ func TestCompleteJob(t *testing.T) { AiTasks: []*sdkproto.AITask{}, Resources: []*sdkproto.Resource{}, }, + isTask: true, + expectTaskStatus: database.TaskStatusPaused, expectHasAiTask: true, expectUsageEvent: false, }, @@ -2992,6 +3003,15 @@ func TestCompleteJob(t *testing.T) { OwnerID: user.ID, OrganizationID: pd.OrganizationID, }) + var taskTable database.TaskTable + if tc.isTask { + taskTable = dbgen.Task(t, db, database.TaskTable{ + OwnerID: user.ID, + OrganizationID: pd.OrganizationID, + WorkspaceID: uuid.NullUUID{UUID: workspaceTable.ID, Valid: true}, + TemplateVersionID: version.ID, + }) + } ctx := testutil.Context(t, testutil.WaitShort) if tc.seedFunc != nil { @@ -3060,6 +3080,12 @@ func TestCompleteJob(t *testing.T) { require.True(t, build.HasAITask.Valid) // We ALWAYS expect a value to be set, therefore not nil, i.e. valid = true. require.Equal(t, tc.expectHasAiTask, build.HasAITask.Bool) + if tc.isTask { + task, err := db.GetTaskByID(ctx, taskTable.ID) + require.NoError(t, err) + require.Equal(t, tc.expectTaskStatus, task.Status) + } + if tc.expectHasAiTask && build.Transition != database.WorkspaceTransitionStop { require.Equal(t, sidebarAppID, build.AITaskSidebarAppID.UUID.String()) } diff --git a/coderd/wsbuilder/wsbuilder.go b/coderd/wsbuilder/wsbuilder.go index 223b8bec08..6aef8c2c2a 100644 --- a/coderd/wsbuilder/wsbuilder.go +++ b/coderd/wsbuilder/wsbuilder.go @@ -6,6 +6,7 @@ import ( "context" "database/sql" "encoding/json" + "errors" "fmt" "net/http" "time" @@ -488,6 +489,21 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object return BuildError{code, "insert workspace build", err} } + // If this is a task workspace, link it to the latest workspace build. + if task, err := store.GetTaskByWorkspaceID(b.ctx, b.workspace.ID); err == nil { + _, err = store.UpsertTaskWorkspaceApp(b.ctx, database.UpsertTaskWorkspaceAppParams{ + TaskID: task.ID, + WorkspaceBuildNumber: buildNum, + WorkspaceAgentID: uuid.NullUUID{}, // Updated by the provisioner upon job completion. + WorkspaceAppID: uuid.NullUUID{}, // Updated by the provisioner upon job completion. + }) + if err != nil { + return BuildError{http.StatusInternalServerError, "upsert task workspace app", err} + } + } else if !errors.Is(err, sql.ErrNoRows) { + return BuildError{http.StatusInternalServerError, "get task by workspace id", err} + } + err = store.InsertWorkspaceBuildParameters(b.ctx, database.InsertWorkspaceBuildParametersParams{ WorkspaceBuildID: workspaceBuildID, Name: names, diff --git a/coderd/wsbuilder/wsbuilder_test.go b/coderd/wsbuilder/wsbuilder_test.go index b862e6459c..3a8921dd6d 100644 --- a/coderd/wsbuilder/wsbuilder_test.go +++ b/coderd/wsbuilder/wsbuilder_test.go @@ -47,6 +47,7 @@ var ( lastBuildJobID = uuid.MustParse("12341234-0000-0000-000c-000000000000") otherUserID = uuid.MustParse("12341234-0000-0000-000d-000000000000") presetID = uuid.MustParse("12341234-0000-0000-000e-000000000000") + taskID = uuid.MustParse("12341234-0000-0000-000f-000000000000") ) func TestBuilder_NoOptions(t *testing.T) { @@ -94,6 +95,7 @@ func TestBuilder_NoOptions(t *testing.T) { asrt.Equal(buildID, bld.ID) }), withBuild, + withNoTask, expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) { asrt.Equal(buildID, params.WorkspaceBuildID) asrt.Empty(params.Name) @@ -140,6 +142,7 @@ func TestBuilder_Initiator(t *testing.T) { expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) { }), withBuild, + withNoTask, ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -188,6 +191,7 @@ func TestBuilder_Baggage(t *testing.T) { expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) { }), withBuild, + withNoTask, ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -229,6 +233,7 @@ func TestBuilder_Reason(t *testing.T) { expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) { }), withBuild, + withNoTask, ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -275,6 +280,7 @@ func TestBuilder_ActiveVersion(t *testing.T) { expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) { }), withBuild, + withNoTask, ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -391,6 +397,7 @@ func TestWorkspaceBuildWithTags(t *testing.T) { expectBuildParameters(func(_ database.InsertWorkspaceBuildParametersParams) { }), withBuild, + withNoTask, expectFindMatchingPresetID(uuid.Nil, sql.ErrNoRows), ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -476,6 +483,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) { } }), withBuild, + withNoTask, expectFindMatchingPresetID(uuid.Nil, sql.ErrNoRows), ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -526,6 +534,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) { } }), withBuild, + withNoTask, expectFindMatchingPresetID(uuid.Nil, sql.ErrNoRows), ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -669,6 +678,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) { } }), withBuild, + withNoTask, expectFindMatchingPresetID(uuid.Nil, sql.ErrNoRows), ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -735,6 +745,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) { } }), withBuild, + withNoTask, ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -798,6 +809,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) { } }), withBuild, + withNoTask, ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -860,6 +872,7 @@ func TestWorkspaceBuildWithPreset(t *testing.T) { asrt.Equal(presetID, bld.TemplateVersionPresetID.UUID) }), withBuild, + withNoTask, expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) { asrt.Equal(buildID, params.WorkspaceBuildID) asrt.Empty(params.Name) @@ -929,6 +942,7 @@ func TestWorkspaceBuildDeleteOrphan(t *testing.T) { asrt.Equal(buildID, bld.ID) }), withBuild, + withNoTask, expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) { asrt.Equal(buildID, params.WorkspaceBuildID) asrt.Empty(params.Name) @@ -992,6 +1006,7 @@ func TestWorkspaceBuildDeleteOrphan(t *testing.T) { asrt.Equal(buildID, bld.ID) }), withBuild, + withNoTask, expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) { asrt.Equal(buildID, params.WorkspaceBuildID) asrt.Empty(params.Name) @@ -1057,6 +1072,7 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) { expectFindMatchingPresetID(uuid.Nil, sql.ErrNoRows), expectBuild(func(bld database.InsertWorkspaceBuildParams) {}), withBuild, + withNoTask, expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) {}), ) fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) @@ -1133,6 +1149,59 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) { } } +func TestWorkspaceBuildWithTask(t *testing.T) { + t.Parallel() + req := require.New(t) + asrt := assert.New(t) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + testTask := database.Task{ + ID: taskID, + OrganizationID: orgID, + OwnerID: userID, + Name: "test-task", + WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, + TemplateVersionID: activeVersionID, + CreatedAt: dbtime.Now(), + } + + mDB := expectDB(t, + // Inputs + withTemplate, + withInactiveVersion(nil), + withLastBuildFound, + withTemplateVersionVariables(inactiveVersionID, nil), + withRichParameters(nil), + withParameterSchemas(inactiveJobID, nil), + withWorkspaceTags(inactiveVersionID, nil), + withProvisionerDaemons([]database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow{}), + + // Outputs + expectProvisionerJob(func(job database.InsertProvisionerJobParams) {}), + withInTx, + expectFindMatchingPresetID(uuid.Nil, sql.ErrNoRows), + expectBuild(func(bld database.InsertWorkspaceBuildParams) {}), + withBuild, + withTask(testTask), + expectUpsertTaskWorkspaceApp(func(params database.UpsertTaskWorkspaceAppParams) { + asrt.Equal(taskID, params.TaskID) + asrt.Equal(int32(2), params.WorkspaceBuildNumber) + asrt.False(params.WorkspaceAgentID.Valid, "workspace_agent_id should be NULL initially") + asrt.False(params.WorkspaceAppID.Valid, "workspace_app_id should be NULL initially") + }), + expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) {}), + ) + fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) + + ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID} + uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}) + // nolint: dogsled + _, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{}) + req.NoError(err) +} + func TestWsbuildError(t *testing.T) { t.Parallel() @@ -1514,3 +1583,39 @@ type fakeUsageChecker struct { func (f *fakeUsageChecker) CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error) { return f.checkBuildUsageFunc(ctx, store, templateVersion) } + +func withNoTask(mTx *dbmock.MockStore) { + mTx.EXPECT().GetTaskByWorkspaceID(gomock.Any(), gomock.Any()).Times(1). + DoAndReturn(func(ctx context.Context, id uuid.UUID) (database.Task, error) { + return database.Task{}, sql.ErrNoRows + }) +} + +func withTask(task database.Task) func(mTx *dbmock.MockStore) { + return func(mTx *dbmock.MockStore) { + mTx.EXPECT().GetTaskByWorkspaceID(gomock.Any(), gomock.Any()).Times(1). + DoAndReturn(func(ctx context.Context, id uuid.UUID) (database.Task, error) { + return task, nil + }) + } +} + +func expectUpsertTaskWorkspaceApp( + assertions func(database.UpsertTaskWorkspaceAppParams), +) func(mTx *dbmock.MockStore) { + return func(mTx *dbmock.MockStore) { + mTx.EXPECT().UpsertTaskWorkspaceApp(gomock.Any(), gomock.Any()). + Times(1). + DoAndReturn( + func(ctx context.Context, params database.UpsertTaskWorkspaceAppParams) (database.TaskWorkspaceApp, error) { + assertions(params) + return database.TaskWorkspaceApp{ + TaskID: params.TaskID, + WorkspaceBuildNumber: params.WorkspaceBuildNumber, + WorkspaceAgentID: params.WorkspaceAgentID, + WorkspaceAppID: params.WorkspaceAppID, + }, nil + }, + ) + } +}