coder · david-fraley · Jan 29, 2026 · Dec 9, 2025 · Jan 27, 2026 · Jan 28, 2026
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -83,6 +83,7 @@ import (
 	"github.com/coder/coder/v2/coderd/schedule"
 	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/coderd/updatecheck"
+	"github.com/coder/coder/v2/coderd/usage"
 	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/coderd/webpush"
 	"github.com/coder/coder/v2/coderd/workspaceapps"
@@ -186,6 +187,7 @@ type Options struct {
 	TelemetryReporter                  telemetry.Reporter
 
 	ProvisionerdServerMetrics *provisionerdserver.Metrics
+	UsageInserter             usage.Inserter
 }
 
 // New constructs a codersdk client connected to an in-memory API instance.
@@ -266,6 +268,11 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 		}
 	}
 
+	var usageInserter *atomic.Pointer[usage.Inserter]
+	if options.UsageInserter != nil {
+		usageInserter = &atomic.Pointer[usage.Inserter]{}
+		usageInserter.Store(&options.UsageInserter)
+	}
 	if options.Database == nil {
 		options.Database, options.Pubsub = dbtestutil.NewDB(t)
 	}
@@ -559,6 +566,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			Database:                       options.Database,
 			Pubsub:                         options.Pubsub,
 			ExternalAuthConfigs:            options.ExternalAuthConfigs,
+			UsageInserter:                  usageInserter,
 
 			Auditor:                            options.Auditor,
 			ConnectionLogger:                   options.ConnectionLogger,

diff --git a/coderd/coderdtest/usage.go b/coderd/coderdtest/usage.go
@@ -0,0 +1,44 @@
+package coderdtest
+
+import (
+	"context"
+	"sync"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/usage"
+	"github.com/coder/coder/v2/coderd/usage/usagetypes"
+)
+
+var _ usage.Inserter = (*UsageInserter)(nil)
+
+type UsageInserter struct {
+	sync.Mutex
+	events []usagetypes.DiscreteEvent
+}
+
+func NewUsageInserter() *UsageInserter {
+	return &UsageInserter{
+		events: []usagetypes.DiscreteEvent{},
+	}
+}
+
+func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
+	u.Lock()
+	defer u.Unlock()
+	u.events = append(u.events, event)
+	return nil
+}
+
+func (u *UsageInserter) GetEvents() []usagetypes.DiscreteEvent {
+	u.Lock()
+	defer u.Unlock()
+	eventsCopy := make([]usagetypes.DiscreteEvent, len(u.events))
+	copy(eventsCopy, u.events)
+	return eventsCopy
+}
+
+func (u *UsageInserter) Reset() {
+	u.Lock()
+	defer u.Unlock()
+	u.events = []usagetypes.DiscreteEvent{}
+}
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -2026,13 +2026,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")
@@ -2048,7 +2046,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)
@@ -2083,7 +2080,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.",
@@ -2098,21 +2095,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
@@ -2878,7 +2878,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.
@@ -2887,6 +2887,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,
@@ -2964,15 +2995,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",
@@ -3007,19 +3040,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()
@@ -4410,62 +4430,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/coderd/wsbuilder/wsbuilder.go b/coderd/wsbuilder/wsbuilder.go
@@ -6,7 +6,6 @@ import (
 	"context"
 	"database/sql"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"net/http"
 	"time"
@@ -87,13 +86,15 @@ type Builder struct {
 	templateVersionPresetParameterValues *[]database.TemplateVersionPresetParameter
 	parameterRender                      dynamicparameters.Renderer
 	workspaceTags                        *map[string]string
+	task                                 *database.Task
+	hasTask                              *bool // A workspace without a task will have a nil `task` and false `hasTask`.
 
 	prebuiltWorkspaceBuildStage  sdkproto.PrebuiltWorkspaceBuildStage
 	verifyNoLegacyParametersOnce bool
 }
 
 type UsageChecker interface {
-	CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (UsageCheckResponse, error)
+	CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, task *database.Task, transition database.WorkspaceTransition) (UsageCheckResponse, error)
 }
 
 type UsageCheckResponse struct {
@@ -105,7 +106,7 @@ type NoopUsageChecker struct{}
 
 var _ UsageChecker = NoopUsageChecker{}
 
-func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion) (UsageCheckResponse, error) {
+func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ *database.Task, _ database.WorkspaceTransition) (UsageCheckResponse, error) {
 	return UsageCheckResponse{
 		Permitted: true,
 	}, nil
@@ -489,8 +490,12 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
 			return BuildError{code, "insert workspace build", err}
 		}
 
+		task, err := b.getWorkspaceTask()
+		if err != nil {
+			return BuildError{http.StatusInternalServerError, "get task by workspace id", 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 {
+		if task != nil {
 			_, err = store.UpsertTaskWorkspaceApp(b.ctx, database.UpsertTaskWorkspaceAppParams{
 				TaskID:               task.ID,
 				WorkspaceBuildNumber: buildNum,
@@ -500,8 +505,6 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
 			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{
@@ -634,6 +637,27 @@ func (b *Builder) getTemplateVersionID() (uuid.UUID, error) {
 	return bld.TemplateVersionID, nil
 }
 
+// getWorkspaceTask returns the task associated with the workspace, if any.
+// If no task exists, it returns (nil, nil).
+func (b *Builder) getWorkspaceTask() (*database.Task, error) {
+	if b.hasTask != nil {
+		return b.task, nil
+	}
+	t, err := b.store.GetTaskByWorkspaceID(b.ctx, b.workspace.ID)
+	if err != nil {
+		if xerrors.Is(err, sql.ErrNoRows) {
+			b.hasTask = ptr.Ref(false)
+			//nolint:nilnil // No task exists.
+			return nil, nil
+		}
+		return nil, xerrors.Errorf("get task: %w", err)
+	}
+
+	b.task = &t
+	b.hasTask = ptr.Ref(true)
+	return b.task, nil
+}
+
 func (b *Builder) getTemplateTerraformValues() (*database.TemplateVersionTerraformValue, error) {
 	if b.terraformValues != nil {
 		return b.terraformValues, nil
@@ -1307,7 +1331,12 @@ func (b *Builder) checkUsage() error {
 		return BuildError{http.StatusInternalServerError, "Failed to fetch template version", err}
 	}
 
-	resp, err := b.usageChecker.CheckBuildUsage(b.ctx, b.store, templateVersion)
+	task, err := b.getWorkspaceTask()
+	if err != nil {
+		return BuildError{http.StatusInternalServerError, "Failed to fetch workspace task", err}
+	}
+
+	resp, err := b.usageChecker.CheckBuildUsage(b.ctx, b.store, templateVersion, task, b.trans)
 	if err != nil {
 		return BuildError{http.StatusInternalServerError, "Failed to check build usage", err}
 	}