Oct 9, 2025 · Oct 9, 2025 · Oct 9, 2025
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
 }

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

 var taskAppID uuid.NullUUID
 var taskAgentID uuid.NullUUID
 var hasAITask bool
 var warnUnknownTaskAppID bool
 if tasks := jobType.WorkspaceBuild.GetAiTasks(); len(tasks) > 0 {
 }

 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:
 }
 }

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

 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
 seedFuncfunc(context.Context, testing.TB, database.Store)error// If you need to insert other resources
 transition       database.WorkspaceTransition
 input*proto.CompletedJob_WorkspaceBuild
 isTaskbool
 expectTaskStatus database.TaskStatus
 expectHasAiTaskbool
 expectUsageEventbool
 }
 input:&proto.CompletedJob_WorkspaceBuild{
 // No AiTasks defined.
 },
 isTask:false,
 expectHasAiTask:false,
 expectUsageEvent:false,
 },
 },
 },
 },
 isTask:true,
 expectTaskStatus:database.TaskStatusInitializing,
 expectHasAiTask:true,
 expectUsageEvent:true,
 },
 },
 },
 },
 isTask:true,
 expectTaskStatus:database.TaskStatusInitializing,
 expectHasAiTask:false,
 expectUsageEvent:false,
 },
 },
 },
 },
 isTask:true,
 expectTaskStatus:database.TaskStatusPaused,
 expectHasAiTask:true,
 expectUsageEvent:false,
 },
 AiTasks:   []*sdkproto.AITask{},
 Resources: []*sdkproto.Resource{},
 },
 isTask:true,
 expectTaskStatus:database.TaskStatusPaused,
 expectHasAiTask:true,
 expectUsageEvent:false,
 },
 OwnerID:user.ID,
 OrganizationID:pd.OrganizationID,
 })
 vartaskTable database.TaskTable
 iftc.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)
 iftc.seedFunc!=nil {
 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)

 iftc.isTask {
 task,err:=db.GetTaskByID(ctx,taskTable.ID)
 require.NoError(t,err)
 require.Equal(t,tc.expectTaskStatus,task.Status)
 }

 iftc.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
 "context"
 "database/sql"
 "encoding/json"
 "errors"
 "fmt"
 "net/http"
 "time"
 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,
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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)
		}
Expand All		@@ -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 {
Expand DownExpand Up		@@ -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:
Expand DownExpand Up		@@ -2108,6 +2135,27 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
		}
		}

		if task, err := db.GetTaskByWorkspaceID(ctx, workspace.ID); err == nil {
Copy link MemberAuthor mafredriOct 9, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. Review: Previous sidebar app behavior is retained for the time being, we're just improving the task data model in parallel. johnstcn reacted with thumbs up emoji
		// 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{
Expand DownExpand Up		@@ -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(),
Expand DownExpand Up		@@ -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{},
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -2850,6 +2850,8 @@ func TestCompleteJob(t *testing.T) {
		seedFuncfunc(context.Context, testing.TB, database.Store)error// If you need to insert other resources
		transition database.WorkspaceTransition
		input*proto.CompletedJob_WorkspaceBuild
		isTaskbool
		expectTaskStatus database.TaskStatus
		expectHasAiTaskbool
		expectUsageEventbool
		}
Expand All		@@ -2862,6 +2864,7 @@ func TestCompleteJob(t *testing.T) {
		input:&proto.CompletedJob_WorkspaceBuild{
		// No AiTasks defined.
		},
		isTask:false,
		expectHasAiTask:false,
		expectUsageEvent:false,
		},
Expand DownExpand Up		@@ -2894,6 +2897,8 @@ func TestCompleteJob(t *testing.T) {
		},
		},
		},
		isTask:true,
		expectTaskStatus:database.TaskStatusInitializing,
		expectHasAiTask:true,
		expectUsageEvent:true,
		},
Expand All		@@ -2912,6 +2917,8 @@ func TestCompleteJob(t *testing.T) {
		},
		},
		},
		isTask:true,
		expectTaskStatus:database.TaskStatusInitializing,
		expectHasAiTask:false,
		expectUsageEvent:false,
		},
Expand DownExpand Up		@@ -2944,6 +2951,8 @@ func TestCompleteJob(t *testing.T) {
		},
		},
		},
		isTask:true,
		expectTaskStatus:database.TaskStatusPaused,
		expectHasAiTask:true,
		expectUsageEvent:false,
		},
Expand All		@@ -2955,6 +2964,8 @@ func TestCompleteJob(t *testing.T) {
		AiTasks: []*sdkproto.AITask{},
		Resources: []*sdkproto.Resource{},
		},
		isTask:true,
		expectTaskStatus:database.TaskStatusPaused,
		expectHasAiTask:true,
		expectUsageEvent:false,
		},
Expand DownExpand Up		@@ -2992,6 +3003,15 @@ func TestCompleteJob(t *testing.T) {
		OwnerID:user.ID,
		OrganizationID:pd.OrganizationID,
		})
		vartaskTable database.TaskTable
		iftc.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)
		iftc.seedFunc!=nil {
Expand DownExpand Up		@@ -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)

		iftc.isTask {
		task,err:=db.GetTaskByID(ctx,taskTable.ID)
		require.NoError(t,err)
		require.Equal(t,tc.expectTaskStatus,task.Status)
		}

		iftc.expectHasAiTask&&build.Transition!=database.WorkspaceTransitionStop {
		require.Equal(t,sidebarAppID,build.AITaskSidebarAppID.UUID.String())
		}
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -6,6 +6,7 @@ import (
		"context"
		"database/sql"
		"encoding/json"
		"errors"
		"fmt"
		"net/http"
		"time"
Expand DownExpand Up		@@ -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,
Expand Down