chore(coderd/provisionerdserver): avoid fk error on invalid ai_task_sidebar_app_id (#19253)

This is a workaround for https://github.com/coder/coder/issues/18776

We avoid the foreign key issue by checking the previously inserted
workspace applications before calling UpdateWorkspaceAITask. Now,
affected workspaces will show as "not running an AI task" on the single
task view, which is technically correct.

We also insert a provisioner job log at WARN level to ensure that the
user sees some information that they have run into this issue, as well
as logging on the server side.

Longer term, we plan to modify how the workspace tasks view is
presented. This is a stopgap measure until we solidify that plan.

NOTE: this does **not** address the fact that stopping a workspace with
`has_ai_task: true` will result in the completed stop build no longer
having `has_ai_task: true`, resulting in tasks "disappearing" on stop.

coderd/provisionerdserver/provisionerdserver.go

@@ -1925,12 +1925,16 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
 			return xerrors.Errorf("update workspace build deadline: %w", err)
 		}
 
+		appIDs := make([]string, 0)
 		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 {
 				dur := time.Duration(protoAgent.GetConnectionTimeoutSeconds()) * time.Second
 				agentTimeouts[dur] = true
+				for _, app := range protoAgent.GetApps() {
+					appIDs = append(appIDs, app.GetId())
+				}
 			}
 
 			err = InsertWorkspaceResource(ctx, db, job.ID, workspaceBuild.Transition, protoResource, telemetrySnapshot)
@@ -1945,14 +1949,21 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
 		}
 
 		var sidebarAppID uuid.NullUUID
-		hasAITask := len(jobType.WorkspaceBuild.AiTasks) == 1
-		if hasAITask {
-			task := jobType.WorkspaceBuild.AiTasks[0]
-			if task.SidebarApp == nil {
-				return xerrors.Errorf("update ai task: sidebar app is nil")
+		var hasAITask bool
+		var warnUnknownSidebarAppID bool
+		if tasks := jobType.WorkspaceBuild.GetAiTasks(); len(tasks) > 0 {
+			hasAITask = true
+			task := tasks[0]
+			if task == nil || task.GetSidebarApp() == nil || len(task.GetSidebarApp().GetId()) == 0 {
+				return xerrors.Errorf("update ai task: sidebar app is nil or empty")
 			}
 
-			id, err := uuid.Parse(task.SidebarApp.Id)
+			sidebarTaskID := task.GetSidebarApp().GetId()
+			if !slices.Contains(appIDs, sidebarTaskID) {
+				warnUnknownSidebarAppID = true
+			}
+
+			id, err := uuid.Parse(task.GetSidebarApp().GetId())
 			if err != nil {
 				return xerrors.Errorf("parse sidebar app id: %w", err)
 			}
@@ -1960,9 +1971,48 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
 			sidebarAppID = uuid.NullUUID{UUID: id, Valid: true}
 		}
 
+		if warnUnknownSidebarAppID {
+			// Ref: https://github.com/coder/coder/issues/18776
+			// This can happen for a number of reasons:
+			// 1. Misconfigured template
+			// 2. Count=0 on the agent due to stop transition, meaning the associated coder_app was not inserted.
+			// Failing the build at this point is not ideal, so log a warning instead.
+			s.Logger.Warn(ctx, "unknown ai_task_sidebar_app_id",
+				slog.F("ai_task_sidebar_app_id", sidebarAppID.UUID.String()),
+				slog.F("job_id", job.ID.String()),
+				slog.F("workspace_id", workspace.ID),
+				slog.F("workspace_build_id", workspaceBuild.ID),
+				slog.F("transition", string(workspaceBuild.Transition)),
+			)
+			// In order to surface this to the user, we will also insert a warning into the build logs.
+			if _, err := db.InsertProvisionerJobLogs(ctx, database.InsertProvisionerJobLogsParams{
+				JobID:     jobID,
+				CreatedAt: []time.Time{now, now, now, now},
+				Source:    []database.LogSource{database.LogSourceProvisionerDaemon, database.LogSourceProvisionerDaemon, database.LogSourceProvisionerDaemon, database.LogSourceProvisionerDaemon},
+				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_sidebar_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.", sidebarAppID.UUID.String()),
+					"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.",
+				},
+			}); err != nil {
+				s.Logger.Error(ctx, "insert provisioner job log for ai task sidebar app id warning",
+					slog.F("job_id", jobID),
+					slog.F("workspace_id", workspace.ID),
+					slog.F("workspace_build_id", workspaceBuild.ID),
+					slog.F("transition", string(workspaceBuild.Transition)),
+				)
+			}
+			// Important: reset hasAITask and sidebarAppID so that we don't run into a fk constraint violation.
+			hasAITask = false
+			sidebarAppID = uuid.NullUUID{}
+		}
+
 		// 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.
-		err = db.UpdateWorkspaceBuildAITaskByID(ctx, database.UpdateWorkspaceBuildAITaskByIDParams{
+		if err := db.UpdateWorkspaceBuildAITaskByID(ctx, database.UpdateWorkspaceBuildAITaskByIDParams{
 			ID: workspaceBuild.ID,
 			HasAITask: sql.NullBool{
 				Bool:  hasAITask,
@@ -1970,8 +2020,7 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
 			},
 			SidebarAppID: sidebarAppID,
 			UpdatedAt:    now,
-		})
-		if err != nil {
+		}); err != nil {
 			return xerrors.Errorf("update workspace build ai tasks flag: %w", err)
 		}
 

coderd/provisionerdserver/provisionerdserver_test.go

@@ -2794,6 +2794,22 @@ func TestCompleteJob(t *testing.T) {
 					},
 					expected: true,
 				},
+				// Checks regression for https://github.com/coder/coder/issues/18776
+				{
+					name: "non-existing app",
+					input: &proto.CompletedJob_WorkspaceBuild{
+						AiTasks: []*sdkproto.AITask{
+							{
+								Id: uuid.NewString(),
+								SidebarApp: &sdkproto.AITaskSidebarApp{
+									// Non-existing app ID would previously trigger a FK violation.
+									Id: uuid.NewString(),
+								},
+							},
+						},
+					},
+					expected: false,
+				},
 			} {
 				t.Run(tc.name, func(t *testing.T) {
 					t.Parallel()