shishantbiswas/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-06-02 20:48:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(coderd): implement task to app linking (#20237)

Browse Source 
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
This commit is contained in:
Mathias Fredriksson
2025-10-13 12:57:06 +03:00
committed by GitHub
parent 5dc57da6b4
commit a8f87c2625
4 changed files with 225 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/provisionerdserver/provisionerdserver.go
+78 -3
View File
@@ -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
+26
View File
@@ -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())
}