fix(coderd): clean up task snapshots on task deletion (#21949)

Task snapshots were orphaned when tasks were soft-deleted. The
`task_snapshots` table has an `ON DELETE CASCADE` foreign key, but
that only fires on hard deletes.

Modified DeleteTask to use a CTE that atomically soft-deletes the
task and removes its snapshot in a single transaction. The query now
returns just the task UUID instead of the full row.

Closes coder/internal#1283

coderd/aitasks_test.go

@@ -398,6 +398,144 @@ func TestTasks(t *testing.T) {
 			require.NoError(t, err, "should be possible to delete a task with no workspace")
 		})
 
+		t.Run("SnapshotCleanupOnDeletion", func(t *testing.T) {
+			t.Parallel()
+
+			client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+			user := coderdtest.CreateFirstUser(t, client)
+			template := createAITemplate(t, client, user)
+
+			ctx := testutil.Context(t, testutil.WaitLong)
+
+			userObj, err := client.User(ctx, user.UserID.String())
+			require.NoError(t, err)
+			userSubject := coderdtest.AuthzUserSubject(userObj)
+
+			task, err := client.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
+				TemplateVersionID: template.ActiveVersionID,
+				Input:             "delete me with snapshot",
+			})
+			require.NoError(t, err)
+			ws, err := client.Workspace(ctx, task.WorkspaceID.UUID)
+			require.NoError(t, err)
+			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
+
+			// Create a snapshot for the task.
+			snapshotJSON := `{"format":"agentapi","data":{"messages":[{"role":"user","content":"test"}]}}`
+			err = db.UpsertTaskSnapshot(dbauthz.As(ctx, userSubject), database.UpsertTaskSnapshotParams{
+				TaskID:               task.ID,
+				LogSnapshot:          json.RawMessage(snapshotJSON),
+				LogSnapshotCreatedAt: dbtime.Now(),
+			})
+			require.NoError(t, err)
+
+			// Verify snapshot exists.
+			_, err = db.GetTaskSnapshot(dbauthz.As(ctx, userSubject), task.ID)
+			require.NoError(t, err)
+
+			// Delete the task.
+			err = client.DeleteTask(ctx, "me", task.ID)
+			require.NoError(t, err, "delete task request should be accepted")
+
+			// Verify snapshot no longer exists.
+			_, err = db.GetTaskSnapshot(dbauthz.As(ctx, userSubject), task.ID)
+			require.ErrorIs(t, err, sql.ErrNoRows, "snapshot should be deleted with task")
+		})
+
+		t.Run("DeletionWithoutSnapshot", func(t *testing.T) {
+			t.Parallel()
+
+			client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+			user := coderdtest.CreateFirstUser(t, client)
+			template := createAITemplate(t, client, user)
+
+			ctx := testutil.Context(t, testutil.WaitLong)
+
+			userObj, err := client.User(ctx, user.UserID.String())
+			require.NoError(t, err)
+			userSubject := coderdtest.AuthzUserSubject(userObj)
+
+			task, err := client.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
+				TemplateVersionID: template.ActiveVersionID,
+				Input:             "delete me without snapshot",
+			})
+			require.NoError(t, err)
+			ws, err := client.Workspace(ctx, task.WorkspaceID.UUID)
+			require.NoError(t, err)
+			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
+
+			// Verify no snapshot exists.
+			_, err = db.GetTaskSnapshot(dbauthz.As(ctx, userSubject), task.ID)
+			require.ErrorIs(t, err, sql.ErrNoRows, "snapshot should not exist initially")
+
+			// Delete the task (should succeed even without snapshot).
+			err = client.DeleteTask(ctx, "me", task.ID)
+			require.NoError(t, err, "delete task should succeed even without snapshot")
+		})
+
+		t.Run("PreservesOtherTaskSnapshots", func(t *testing.T) {
+			t.Parallel()
+
+			client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+			user := coderdtest.CreateFirstUser(t, client)
+			template := createAITemplate(t, client, user)
+
+			ctx := testutil.Context(t, testutil.WaitLong)
+
+			userObj, err := client.User(ctx, user.UserID.String())
+			require.NoError(t, err)
+			userSubject := coderdtest.AuthzUserSubject(userObj)
+
+			// Create task A.
+			taskA, err := client.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
+				TemplateVersionID: template.ActiveVersionID,
+				Input:             "task A",
+			})
+			require.NoError(t, err)
+			wsA, err := client.Workspace(ctx, taskA.WorkspaceID.UUID)
+			require.NoError(t, err)
+			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, wsA.LatestBuild.ID)
+
+			// Create task B.
+			taskB, err := client.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
+				TemplateVersionID: template.ActiveVersionID,
+				Input:             "task B",
+			})
+			require.NoError(t, err)
+			wsB, err := client.Workspace(ctx, taskB.WorkspaceID.UUID)
+			require.NoError(t, err)
+			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, wsB.LatestBuild.ID)
+
+			// Create snapshots for both tasks.
+			snapshotJSONA := `{"format":"agentapi","data":{"messages":[{"role":"user","content":"task A"}]}}`
+			err = db.UpsertTaskSnapshot(dbauthz.As(ctx, userSubject), database.UpsertTaskSnapshotParams{
+				TaskID:               taskA.ID,
+				LogSnapshot:          json.RawMessage(snapshotJSONA),
+				LogSnapshotCreatedAt: dbtime.Now(),
+			})
+			require.NoError(t, err)
+
+			snapshotJSONB := `{"format":"agentapi","data":{"messages":[{"role":"user","content":"task B"}]}}`
+			err = db.UpsertTaskSnapshot(dbauthz.As(ctx, userSubject), database.UpsertTaskSnapshotParams{
+				TaskID:               taskB.ID,
+				LogSnapshot:          json.RawMessage(snapshotJSONB),
+				LogSnapshotCreatedAt: dbtime.Now(),
+			})
+			require.NoError(t, err)
+
+			// Delete task A.
+			err = client.DeleteTask(ctx, "me", taskA.ID)
+			require.NoError(t, err, "delete task A should succeed")
+
+			// Verify task A's snapshot is removed.
+			_, err = db.GetTaskSnapshot(dbauthz.As(ctx, userSubject), taskA.ID)
+			require.ErrorIs(t, err, sql.ErrNoRows, "task A snapshot should be deleted")
+
+			// Verify task B's snapshot still exists.
+			_, err = db.GetTaskSnapshot(dbauthz.As(ctx, userSubject), taskB.ID)
+			require.NoError(t, err, "task B snapshot should still exist")
+		})
+
 		t.Run("DeletingTaskWorkspaceDeletesTask", func(t *testing.T) {
 			t.Parallel()
 

coderd/database/dbauthz/dbauthz.go

@@ -1932,14 +1932,14 @@ func (q *querier) DeleteTailnetTunnel(ctx context.Context, arg database.DeleteTa
 	return q.db.DeleteTailnetTunnel(ctx, arg)
 }
 
-func (q *querier) DeleteTask(ctx context.Context, arg database.DeleteTaskParams) (database.TaskTable, error) {
+func (q *querier) DeleteTask(ctx context.Context, arg database.DeleteTaskParams) (uuid.UUID, error) {
 	task, err := q.db.GetTaskByID(ctx, arg.ID)
 	if err != nil {
-		return database.TaskTable{}, err
+		return uuid.UUID{}, err
 	}
 
 	if err := q.authorizeContext(ctx, policy.ActionDelete, task.RBACObject()); err != nil {
-		return database.TaskTable{}, err
+		return uuid.UUID{}, err
 	}
 
 	return q.db.DeleteTask(ctx, arg)

coderd/database/dbauthz/dbauthz_test.go

@@ -2517,8 +2517,8 @@ func (s *MethodTestSuite) TestTasks() {
 			DeletedAt: dbtime.Now(),
 		}
 		dbm.EXPECT().GetTaskByID(gomock.Any(), task.ID).Return(task, nil).AnyTimes()
-		dbm.EXPECT().DeleteTask(gomock.Any(), arg).Return(database.TaskTable{}, nil).AnyTimes()
-		check.Args(arg).Asserts(task, policy.ActionDelete).Returns(database.TaskTable{})
+		dbm.EXPECT().DeleteTask(gomock.Any(), arg).Return(task.ID, nil).AnyTimes()
+		check.Args(arg).Asserts(task, policy.ActionDelete).Returns(task.ID)
 	}))
 	s.Run("InsertTask", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
 		tpl := testutil.Fake(s.T(), faker, database.Template{})

coderd/database/dbmetrics/querymetrics.go

@@ -575,7 +575,7 @@ func (m queryMetricsStore) DeleteTailnetTunnel(ctx context.Context, arg database
 	return r0, r1
 }
 
-func (m queryMetricsStore) DeleteTask(ctx context.Context, arg database.DeleteTaskParams) (database.TaskTable, error) {
+func (m queryMetricsStore) DeleteTask(ctx context.Context, arg database.DeleteTaskParams) (uuid.UUID, error) {
 	start := time.Now()
 	r0, r1 := m.s.DeleteTask(ctx, arg)
 	m.queryLatencies.WithLabelValues("DeleteTask").Observe(time.Since(start).Seconds())

coderd/database/dbmock/dbmock.go

@@ -941,10 +941,10 @@ func (mr *MockStoreMockRecorder) DeleteTailnetTunnel(ctx, arg any) *gomock.Call
 }
 
 // DeleteTask mocks base method.
-func (m *MockStore) DeleteTask(ctx context.Context, arg database.DeleteTaskParams) (database.TaskTable, error) {
+func (m *MockStore) DeleteTask(ctx context.Context, arg database.DeleteTaskParams) (uuid.UUID, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "DeleteTask", ctx, arg)
-	ret0, _ := ret[0].(database.TaskTable)
+	ret0, _ := ret[0].(uuid.UUID)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }

coderd/database/querier.go

@@ -132,7 +132,7 @@ type sqlcQuerier interface {
 	DeleteRuntimeConfig(ctx context.Context, key string) error
 	DeleteTailnetPeer(ctx context.Context, arg DeleteTailnetPeerParams) (DeleteTailnetPeerRow, error)
 	DeleteTailnetTunnel(ctx context.Context, arg DeleteTailnetTunnelParams) (DeleteTailnetTunnelRow, error)
-	DeleteTask(ctx context.Context, arg DeleteTaskParams) (TaskTable, error)
+	DeleteTask(ctx context.Context, arg DeleteTaskParams) (uuid.UUID, error)
 	DeleteUserSecret(ctx context.Context, id uuid.UUID) error
 	DeleteWebpushSubscriptionByUserIDAndEndpoint(ctx context.Context, arg DeleteWebpushSubscriptionByUserIDAndEndpointParams) error
 	DeleteWebpushSubscriptions(ctx context.Context, ids []uuid.UUID) error

coderd/database/queries.sql.go

@@ -13151,13 +13151,19 @@ func (q *sqlQuerier) UpsertTailnetTunnel(ctx context.Context, arg UpsertTailnetT
 }
 
 const deleteTask = `-- name: DeleteTask :one
-UPDATE tasks
-SET
-	deleted_at = $1::timestamptz
-WHERE
-	id = $2::uuid
-	AND deleted_at IS NULL
-RETURNING id, organization_id, owner_id, name, workspace_id, template_version_id, template_parameters, prompt, created_at, deleted_at, display_name
+WITH deleted_task AS (
+	UPDATE tasks
+	SET
+		deleted_at = $1::timestamptz
+	WHERE
+		id = $2::uuid
+		AND deleted_at IS NULL
+	RETURNING id
+), deleted_snapshot AS (
+	DELETE FROM task_snapshots
+	WHERE task_id = $2::uuid
+)
+SELECT id FROM deleted_task
 `
 
 type DeleteTaskParams struct {
@@ -13165,23 +13171,11 @@ type DeleteTaskParams struct {
 	ID        uuid.UUID `db:"id" json:"id"`
 }
 
-func (q *sqlQuerier) DeleteTask(ctx context.Context, arg DeleteTaskParams) (TaskTable, error) {
+func (q *sqlQuerier) DeleteTask(ctx context.Context, arg DeleteTaskParams) (uuid.UUID, error) {
 	row := q.db.QueryRowContext(ctx, deleteTask, arg.DeletedAt, arg.ID)
-	var i TaskTable
-	err := row.Scan(
-		&i.ID,
-		&i.OrganizationID,
-		&i.OwnerID,
-		&i.Name,
-		&i.WorkspaceID,
-		&i.TemplateVersionID,
-		&i.TemplateParameters,
-		&i.Prompt,
-		&i.CreatedAt,
-		&i.DeletedAt,
-		&i.DisplayName,
-	)
-	return i, err
+	var id uuid.UUID
+	err := row.Scan(&id)
+	return id, err
 }
 
 const getTaskByID = `-- name: GetTaskByID :one

coderd/database/queries/tasks.sql

@@ -57,13 +57,19 @@ AND CASE WHEN @status::text != '' THEN tws.status = @status::task_status ELSE TR
 ORDER BY tws.created_at DESC;
 
 -- name: DeleteTask :one
-UPDATE tasks
-SET
-	deleted_at = @deleted_at::timestamptz
-WHERE
-	id = @id::uuid
-	AND deleted_at IS NULL
-RETURNING *;
+WITH deleted_task AS (
+	UPDATE tasks
+	SET
+		deleted_at = @deleted_at::timestamptz
+	WHERE
+		id = @id::uuid
+		AND deleted_at IS NULL
+	RETURNING id
+), deleted_snapshot AS (
+	DELETE FROM task_snapshots
+	WHERE task_id = @id::uuid
+)
+SELECT id FROM deleted_task;
 
 
 -- name: UpdateTaskPrompt :one