opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(coderd/provisionerdserver): correct managed agent tracking (#21696)

Browse Source 
Relates to https://github.com/coder/internal/issues/1282

Updates tracking of managed agents to be predicated instead on the
presence of a related `task_id` instead of the presence of a
`coder_ai_task` resource.
This commit is contained in:
Cian Johnston
2026-01-27 12:14:52 +00:00
committed by GitHub
parent c3f41ce08c
commit 7b44976618
3 changed files with 70 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/provisionerdserver/provisionerdserver.go
+15 -16
View File
@@ -2056,13 +2056,11 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
}
var (
hasAITask bool
unknownAppID string
taskAppID uuid.NullUUID
taskAgentID uuid.NullUUID
)
if tasks := jobType.WorkspaceBuild.GetAiTasks(); len(tasks) > 0 {
hasAITask = true
task := tasks[0]
if task == nil {
return xerrors.Errorf("update ai task: task is nil")
@@ -2078,7 +2076,6 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
if !slices.Contains(appIDs, appID) {
unknownAppID = appID
hasAITask = false
} else {
// Only parse for valid app and agent to avoid fk violation.
id, err := uuid.Parse(appID)
@@ -2113,7 +2110,7 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
Level: []database.LogLevel{database.LogLevelWarn, database.LogLevelWarn, database.LogLevelWarn, database.LogLevelWarn},
Stage: []string{"Cleaning Up", "Cleaning Up", "Cleaning Up", "Cleaning Up"},
Output: []string{
fmt.Sprintf("Unknown ai_task_app_id %q. This workspace will be unable to run AI tasks. This may be due to a template configuration issue, please check with the template author.", taskAppID.UUID.String()),
fmt.Sprintf("Unknown ai_task_app_id %q. This workspace will be unable to run AI tasks. This may be due to a template configuration issue, please check with the template author.", unknownAppID),
"Template author: double-check the following:",
" - You have associated the coder_ai_task with a valid coder_app in your template (ref: https://registry.terraform.io/providers/coder/coder/latest/docs/resources/ai_task).",
" - You have associated the coder_agent with at least one other compute resource. Agents with no other associated resources are not inserted into the database.",
@@ -2128,21 +2125,23 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
}
}
if hasAITask && workspaceBuild.Transition == database.WorkspaceTransitionStart {
// Insert usage event for managed agents.
usageInserter := s.UsageInserter.Load()
if usageInserter != nil {
event := usagetypes.DCManagedAgentsV1{
Count: 1,
}
err = (*usageInserter).InsertDiscreteUsageEvent(ctx, db, event)
if err != nil {
return xerrors.Errorf("insert %q event: %w", event.EventType(), err)
var hasAITask bool
if task, err := db.GetTaskByWorkspaceID(ctx, workspace.ID); err == nil {
hasAITask = true
if workspaceBuild.Transition == database.WorkspaceTransitionStart {
// Insert usage event for managed agents.
usageInserter := s.UsageInserter.Load()
if usageInserter != nil {
event := usagetypes.DCManagedAgentsV1{
Count: 1,
}
err = (*usageInserter).InsertDiscreteUsageEvent(ctx, db, event)
if err != nil {
return xerrors.Errorf("insert %q event: %w", event.EventType(), err)
}
}
}
}
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
coderd/provisionerdserver/provisionerdserver_test.go
+37 -76
View File
@@ -2875,7 +2875,7 @@ func TestCompleteJob(t *testing.T) {
sidebarAppID := uuid.New()
for _, tc := range []testcase{
{
name: "has_ai_task is false by default",
name: "has_ai_task is false if task_id is nil",
transition: database.WorkspaceTransitionStart,
input: &proto.CompletedJob_WorkspaceBuild{
// No AiTasks defined.
@@ -2884,6 +2884,37 @@ func TestCompleteJob(t *testing.T) {
expectHasAiTask: false,
expectUsageEvent: false,
},
{
name: "has_ai_task is false even if there are coder_ai_task resources, but no task_id",
transition: database.WorkspaceTransitionStart,
input: &proto.CompletedJob_WorkspaceBuild{
AiTasks: []*sdkproto.AITask{
{
Id: uuid.NewString(),
AppId: sidebarAppID.String(),
},
},
Resources: []*sdkproto.Resource{
{
Agents: []*sdkproto.Agent{
{
Id: uuid.NewString(),
Name: "a",
Apps: []*sdkproto.App{
{
Id: sidebarAppID.String(),
Slug: "test-app",
},
},
},
},
},
},
},
isTask: false,
expectHasAiTask: false,
expectUsageEvent: false,
},
{
name: "has_ai_task is set to true",
transition: database.WorkspaceTransitionStart,
@@ -2961,15 +2992,17 @@ func TestCompleteJob(t *testing.T) {
{
Id: uuid.NewString(),
// Non-existing app ID would previously trigger a FK violation.
// Now it should just be ignored.
// Now it will trigger a warning instead in the provisioner logs.
AppId: sidebarAppID.String(),
},
},
},
isTask: true,
expectTaskStatus: database.TaskStatusInitializing,
expectHasAiTask: false,
expectUsageEvent: false,
// You can still "sort of" use a task in this state, but as we don't have
// the correct app ID you won't be able to communicate with it via Coder.
expectHasAiTask: true,
expectUsageEvent: true,
},
{
name: "has_ai_task is set to true, but transition is not start",
@@ -3004,19 +3037,6 @@ func TestCompleteJob(t *testing.T) {
expectHasAiTask: true,
expectUsageEvent: false,
},
{
name: "current build does not have ai task but previous build did",
seedFunc: seedPreviousWorkspaceStartWithAITask,
transition: database.WorkspaceTransitionStop,
input: &proto.CompletedJob_WorkspaceBuild{
AiTasks: []*sdkproto.AITask{},
Resources: []*sdkproto.Resource{},
},
isTask: true,
expectTaskStatus: database.TaskStatusPaused,
expectHasAiTask: false, // We no longer inherit this from the previous build.
expectUsageEvent: false,
},
} {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
@@ -4407,62 +4427,3 @@ func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ databa
f.collectedEvents = append(f.collectedEvents, event)
return nil
}
func seedPreviousWorkspaceStartWithAITask(ctx context.Context, t testing.TB, db database.Store) error {
t.Helper()
// If the below looks slightly convoluted, that's because it is.
// The workspace doesn't yet have a latest build, so querying all
// workspaces will fail.
tpls, err := db.GetTemplates(ctx)
if err != nil {
return xerrors.Errorf("seedFunc: get template: %w", err)
}
if len(tpls) != 1 {
return xerrors.Errorf("seedFunc: expected exactly one template, got %d", len(tpls))
}
ws, err := db.GetWorkspacesByTemplateID(ctx, tpls[0].ID)
if err != nil {
return xerrors.Errorf("seedFunc: get workspaces: %w", err)
}
if len(ws) != 1 {
return xerrors.Errorf("seedFunc: expected exactly one workspace, got %d", len(ws))
}
w := ws[0]
prevJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
OrganizationID: w.OrganizationID,
InitiatorID: w.OwnerID,
Type: database.ProvisionerJobTypeWorkspaceBuild,
})
tvs, err := db.GetTemplateVersionsByTemplateID(ctx, database.GetTemplateVersionsByTemplateIDParams{
TemplateID: tpls[0].ID,
})
if err != nil {
return xerrors.Errorf("seedFunc: get template version: %w", err)
}
if len(tvs) != 1 {
return xerrors.Errorf("seedFunc: expected exactly one template version, got %d", len(tvs))
}
if tpls[0].ActiveVersionID == uuid.Nil {
return xerrors.Errorf("seedFunc: active version id is nil")
}
res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
JobID: prevJob.ID,
})
agt := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
ResourceID: res.ID,
})
_ = dbgen.WorkspaceApp(t, db, database.WorkspaceApp{
AgentID: agt.ID,
})
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
BuildNumber: 1,
HasAITask: sql.NullBool{Valid: true, Bool: true},
ID: w.ID,
InitiatorID: w.OwnerID,
JobID: prevJob.ID,
TemplateVersionID: tvs[0].ID,
Transition: database.WorkspaceTransitionStart,
WorkspaceID: w.ID,
})
return nil
}