From 7b44976618851c199a6e33f0d54925fb5fbcbefc Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Tue, 27 Jan 2026 12:14:52 +0000 Subject: [PATCH] fix(coderd/provisionerdserver): correct managed agent tracking (#21696) 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. --- .../provisionerdserver/provisionerdserver.go | 31 +++-- .../provisionerdserver_test.go | 113 ++++++------------ enterprise/coderd/coderd_test.go | 26 ++-- 3 files changed, 70 insertions(+), 100 deletions(-) diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go index 1ce46670a9..79018f5df7 100644 --- a/coderd/provisionerdserver/provisionerdserver.go +++ b/coderd/provisionerdserver/provisionerdserver.go @@ -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 diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go index 77a2023537..524fa7e030 100644 --- a/coderd/provisionerdserver/provisionerdserver_test.go +++ b/coderd/provisionerdserver/provisionerdserver_test.go @@ -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 -} diff --git a/enterprise/coderd/coderd_test.go b/enterprise/coderd/coderd_test.go index 3588199941..19f81bd51d 100644 --- a/enterprise/coderd/coderd_test.go +++ b/enterprise/coderd/coderd_test.go @@ -623,7 +623,7 @@ func TestManagedAgentLimit(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) - cli, _ := coderdenttest.New(t, &coderdenttest.Options{ + cli, owner := coderdenttest.New(t, &coderdenttest.Options{ Options: &coderdtest.Options{ IncludeProvisionerDaemon: true, }, @@ -708,15 +708,25 @@ func TestManagedAgentLimit(t *testing.T) { noAiTemplate := coderdtest.CreateTemplate(t, cli, uuid.Nil, noAiVersion.ID) // Create one AI workspace, which should succeed. - workspace := coderdtest.CreateWorkspace(t, cli, aiTemplate.ID) + task, err := cli.CreateTask(ctx, owner.UserID.String(), codersdk.CreateTaskRequest{ + Name: namesgenerator.UniqueNameWith("-"), + TemplateVersionID: aiTemplate.ActiveVersionID, + TemplateVersionPresetID: uuid.Nil, + Input: "hi", + DisplayName: namesgenerator.UniqueName(), + }) + require.NoError(t, err, "creating task for AI workspace must succeed") + workspace, err := cli.Workspace(ctx, task.WorkspaceID.UUID) + require.NoError(t, err, "fetching AI workspace must succeed") coderdtest.AwaitWorkspaceBuildJobCompleted(t, cli, workspace.LatestBuild.ID) - // Create a second AI workspace, which should fail. This needs to be done - // manually because coderdtest.CreateWorkspace expects it to succeed. - _, err = cli.CreateUserWorkspace(ctx, codersdk.Me, codersdk.CreateWorkspaceRequest{ //nolint:gocritic // owners must still be subject to the limit - TemplateID: aiTemplate.ID, - Name: coderdtest.RandomUsername(t), - AutomaticUpdates: codersdk.AutomaticUpdatesNever, + // Create a second AI workspace, which should fail. + _, err = cli.CreateTask(ctx, owner.UserID.String(), codersdk.CreateTaskRequest{ + Name: namesgenerator.UniqueNameWith("-"), + TemplateVersionID: aiTemplate.ActiveVersionID, + TemplateVersionPresetID: uuid.Nil, + Input: "hi", + DisplayName: namesgenerator.UniqueName(), }) require.ErrorContains(t, err, "You have breached the managed agent limit in your license")