chore: ensure sidebar app ID must be set when has_ai_task is set

Signed-off-by: Danny Kopping <dannykopping@gmail.com>
coder · dannykopping · Jun 24, 2025 · Jun 18, 2025 · Jun 18, 2025 · Jun 18, 2025
commit d7e47f9fd290725834c9dacc2d900ea7818b64f4
diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -1405,7 +1405,7 @@ func (s *MethodTestSuite) TestTemplate() {
 			TemplateID:     uuid.NullUUID{UUID: t.ID, Valid: true},
 		})
 		check.Args(database.UpdateTemplateVersionAITaskByJobIDParams{
-			JobID: job.ID,
+			JobID:     job.ID,
 			HasAITask: sql.NullBool{Bool: true, Valid: true},
 		}).Asserts(t, policy.ActionUpdate)
 	}))

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -12062,6 +12062,13 @@ func (q *FakeQuerier) UpdateWorkspaceAutostart(_ context.Context, arg database.U
 }
 
 func (q *FakeQuerier) UpdateWorkspaceBuildAITaskByID(_ context.Context, arg database.UpdateWorkspaceBuildAITaskByIDParams) error {
+	if arg.HasAITask.Bool && !arg.SidebarAppID.Valid {
+		return xerrors.Errorf("ai_task_sidebar_app_id is required when has_ai_task is true")
+	}
+	if !arg.HasAITask.Valid && arg.SidebarAppID.Valid {
+		return xerrors.Errorf("ai_task_sidebar_app_id is can only be set when has_ai_task is true")
+	}
+
 	err := validateDatabaseType(arg)
 	if err != nil {
 		return err

diff --git a/...1_rename_ai_tasks_sidebar_app_id.down.sql → ...bar_app_id_column_and_constraint.down.sql b/...1_rename_ai_tasks_sidebar_app_id.down.sql → ...bar_app_id_column_and_constraint.down.sql
@@ -1,3 +1,6 @@
+-- Drop the check constraint first
+ALTER TABLE workspace_builds DROP CONSTRAINT workspace_builds_ai_task_sidebar_app_id_required;
+
 -- Revert ai_task_sidebar_app_id back to ai_tasks_sidebar_app_id in workspace_builds table
 ALTER TABLE workspace_builds DROP CONSTRAINT workspace_builds_ai_task_sidebar_app_id_fkey;
 
@@ -46,4 +49,4 @@ FROM
         )
     );
 
-COMMENT ON VIEW workspace_build_with_user IS 'Joins in the username + avatar url of the initiated by user.';
+COMMENT ON VIEW workspace_build_with_user IS 'Joins in the username + avatar url of the initiated by user.';
diff --git a/...341_rename_ai_tasks_sidebar_app_id.up.sql → ...debar_app_id_column_and_constraint.up.sql b/...341_rename_ai_tasks_sidebar_app_id.up.sql → ...debar_app_id_column_and_constraint.up.sql
@@ -5,6 +5,22 @@ ALTER TABLE workspace_builds RENAME COLUMN ai_tasks_sidebar_app_id TO ai_task_si
 
 ALTER TABLE workspace_builds ADD CONSTRAINT workspace_builds_ai_task_sidebar_app_id_fkey FOREIGN KEY (ai_task_sidebar_app_id) REFERENCES workspace_apps(id);
 
+-- if has_ai_task is true, ai_task_sidebar_app_id MUST be set
+-- ai_task_sidebar_app_id can ONLY be set if has_ai_task is true
+--
+--   has_ai_task | ai_task_sidebar_app_id | Result
+--   ------------|------------------------|---------------
+--   NULL        | NULL                   | TRUE (passes)
+--   NULL        | NOT NULL               | FALSE (fails)
+--   FALSE       | NULL                   | TRUE (passes)
+--   FALSE       | NOT NULL               | FALSE (fails)
+--   TRUE        | NULL                   | FALSE (fails)
+--   TRUE        | NOT NULL               | TRUE (passes)
+ALTER TABLE workspace_builds ADD CONSTRAINT workspace_builds_ai_task_sidebar_app_id_required CHECK (
+    (has_ai_task IS NULL OR has_ai_task = false) AND ai_task_sidebar_app_id IS NULL OR
+    (has_ai_task = true AND ai_task_sidebar_app_id IS NOT NULL)
+);
+
 -- Update the workspace_build_with_user view to use the new column name
 DROP VIEW workspace_build_with_user;
 
@@ -46,4 +62,4 @@ FROM
         )
     );
 
-COMMENT ON VIEW workspace_build_with_user IS 'Joins in the username + avatar url of the initiated by user.';
+COMMENT ON VIEW workspace_build_with_user IS 'Joins in the username + avatar url of the initiated by user.';
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -2,7 +2,6 @@ package coderd_test
 
 import (
 	"context"
-	"database/sql"
 	"encoding/json"
 	"fmt"
 	"maps"
@@ -1507,52 +1506,39 @@ func TestWorkspaceAgentAppHealth(t *testing.T) {
 	t.Parallel()
 	client, db := coderdtest.NewWithDatabase(t, nil)
 	user := coderdtest.CreateFirstUser(t, client)
-
-	// When using WithAgent() here and setting some *proto.App instances on the agent, Do() ends up trying to insert duplicate records.
+	apps := []*proto.App{
+		{
+			Slug:    "code-server",
+			Command: "some-command",
+			Url:     "http://localhost:3000",
+			Icon:    "/code.svg",
+		},
+		{
+			Slug:        "code-server-2",
+			DisplayName: "code-server-2",
+			Command:     "some-command",
+			Url:         "http://localhost:3000",
+			Icon:        "/code.svg",
+			Healthcheck: &proto.Healthcheck{
+				Url:       "http://localhost:3000",
+				Interval:  5,
+				Threshold: 6,
+			},
+		},
+	}
 	r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
 		OrganizationID: user.OrganizationID,
 		OwnerID:        user.UserID,
+	}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
+		agents[0].Apps = apps
+		return agents
 	}).Do()
 
-	res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{JobID: r.Build.JobID})
-	agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: res.ID})
-
-	// It's simpler to call db.InsertWorkspaceApp directly than dbgen.WorkspaceApp because it's more terse and direct;
-	// the latter sets a bunch of defaults which make this test hard.
-	_, err := db.InsertWorkspaceApp(dbauthz.AsSystemRestricted(t.Context()), database.InsertWorkspaceAppParams{
-		ID: uuid.New(),
-		Slug:         "code-server",
-		AgentID:      agent.ID,
-		Icon:         "/code.svg",
-		Command:      sql.NullString{String: "some-command", Valid: true},
-		Url:          sql.NullString{String: "http://localhost:3000", Valid: true},
-		SharingLevel: database.AppSharingLevelOwner,
-		Health:       database.WorkspaceAppHealthDisabled,
-		OpenIn:       database.WorkspaceAppOpenInWindow,
-	})
-	require.NoError(t, err)
-	_, err = db.InsertWorkspaceApp(dbauthz.AsSystemRestricted(t.Context()), database.InsertWorkspaceAppParams{
-		ID: uuid.New(),
-		Slug:                 "code-server-2",
-		DisplayName:          "code-server-2",
-		AgentID:              agent.ID,
-		Icon:                 "/code.svg",
-		Command:              sql.NullString{String: "some-command", Valid: true},
-		Url:                  sql.NullString{String: "http://localhost:3000", Valid: true},
-		HealthcheckInterval:  5,
-		HealthcheckUrl:       "http://localhost:3000",
-		HealthcheckThreshold: 6,
-		Health:               database.WorkspaceAppHealthInitializing,
-		SharingLevel:         database.AppSharingLevelOwner,
-		OpenIn:               database.WorkspaceAppOpenInWindow,
-	})
-	require.NoError(t, err)
-
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 
 	agentClient := agentsdk.New(client.URL)
-	agentClient.SetSessionToken(agent.AuthToken.String())
+	agentClient.SetSessionToken(r.AgentToken)
 	conn, err := agentClient.ConnectRPC(ctx)
 	require.NoError(t, err)
 	defer func() {

diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go
@@ -4531,13 +4531,18 @@ func TestWorkspaceFilterHasAITask(t *testing.T) {
 		}
 		job := dbgen.ProvisionerJob(t, db, pubsub, jobConfig)
 
+		res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{JobID: job.ID})
+		agnt := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: res.ID})
+		sidebarApp := dbgen.WorkspaceApp(t, db, database.WorkspaceApp{AgentID: agnt.ID})
+
 		build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
-			WorkspaceID:       ws.ID,
-			TemplateVersionID: version.ID,
-			InitiatorID:       user.UserID,
-			JobID:             job.ID,
-			BuildNumber:       1,
-			HasAITask:         hasAITask,
+			WorkspaceID:        ws.ID,
+			TemplateVersionID:  version.ID,
+			InitiatorID:        user.UserID,
+			JobID:              job.ID,
+			BuildNumber:        1,
+			HasAITask:          hasAITask,
+			AITaskSidebarAppID: uuid.NullUUID{UUID: sidebarApp.ID, Valid: true},
 		})
 
 		if aiTaskPrompt != nil {