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

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{},

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())
 					}

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,

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
+				},
+			)
+	}
+}