diff --git a/coderd/aitasks_test.go b/coderd/aitasks_test.go index 7037d92b2b..3fd278dc76 100644 --- a/coderd/aitasks_test.go +++ b/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() diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 84560dc9b7..eb3248c3d6 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/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) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index c6f1ea59dd..31128c0c7e 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/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{}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index f2289e2b0c..ad3afd180b 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/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()) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 5aa7c74dc9..d4de074dcf 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/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 } diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 6e4d070851..f39abd822a 100644 --- a/coderd/database/querier.go +++ b/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 diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 68c7473370..9812f6df49 100644 --- a/coderd/database/queries.sql.go +++ b/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 diff --git a/coderd/database/queries/tasks.sql b/coderd/database/queries/tasks.sql index 8deda80a2b..012acbd976 100644 --- a/coderd/database/queries/tasks.sql +++ b/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