diff --git a/internal/jobs/runners/worktree_prune.go b/internal/jobs/runners/worktree_prune.go new file mode 100644 index 0000000..79295aa --- /dev/null +++ b/internal/jobs/runners/worktree_prune.go @@ -0,0 +1,122 @@ +package runners + +import ( + "context" + "log/slog" + "time" + + "github.com/google/uuid" + + "github.com/yan1h/agent-coding-workflow/internal/infra/notify" +) + +// WorktreeRow projects the worktree fields needed for prune. OwnerID and +// MainPath come from a JOIN performed in the production adapter (workspace +// + project), not from the worktrees table directly. +type WorktreeRow struct { + ID uuid.UUID + OwnerID uuid.UUID + Branch string + MainPath string + Path string +} + +// WorktreePruneRepo is the minimal repository surface used by this runner. +type WorktreePruneRepo interface { + ListWorktreesNeedingPruneWarning(ctx context.Context, lastUsedBefore time.Time) ([]WorktreeRow, error) + ListWorktreesNeedingPrune(ctx context.Context, lastUsedBefore time.Time) ([]WorktreeRow, error) + MarkPruneWarning(ctx context.Context, id uuid.UUID) error + TryStartPrune(ctx context.Context, id uuid.UUID) (WorktreeRow, error) + FinishPruneSuccess(ctx context.Context, id uuid.UUID) error + FinishPruneRollback(ctx context.Context, id uuid.UUID) error +} + +// WorktreePruner abstracts the git CLI surface needed. +type WorktreePruner interface { + WorktreeRemove(ctx context.Context, mainPath, worktreePath string) error +} + +// WorktreePruneConfig matches jobs.WorktreePruneConfig minus enabled/interval +// (managed by scheduler). +type WorktreePruneConfig struct { + IdleThreshold time.Duration + WarningLead time.Duration +} + +// WorktreePrune is the B runner. +type WorktreePrune struct { + repo WorktreePruneRepo + gitr WorktreePruner + disp Dispatcher + cfg WorktreePruneConfig + log *slog.Logger +} + +// NewWorktreePrune constructs the runner. +func NewWorktreePrune(repo WorktreePruneRepo, gitr WorktreePruner, disp Dispatcher, cfg WorktreePruneConfig, log *slog.Logger) *WorktreePrune { + if log == nil { + log = slog.Default() + } + return &WorktreePrune{repo: repo, gitr: gitr, disp: disp, cfg: cfg, log: log} +} + +// Run executes warning phase then execution phase. +func (r *WorktreePrune) Run(ctx context.Context) { + r.warningPhase(ctx) + r.executionPhase(ctx) +} + +func (r *WorktreePrune) warningPhase(ctx context.Context) { + cutoff := time.Now().Add(-(r.cfg.IdleThreshold - r.cfg.WarningLead)) + rows, err := r.repo.ListWorktreesNeedingPruneWarning(ctx, cutoff) + if err != nil { + r.log.Error("worktree_prune.list_warning", "err", err.Error()) + return + } + for _, w := range rows { + if err := r.repo.MarkPruneWarning(ctx, w.ID); err != nil { + r.log.Error("worktree_prune.mark_warning", "id", w.ID, "err", err.Error()) + continue + } + if derr := r.disp.Dispatch(ctx, notify.Message{ + UserID: w.OwnerID, + Topic: "worktree.prune_pending", + Severity: notify.SeverityWarning, + Title: "Worktree will be pruned soon", + Body: "Branch " + w.Branch + " has been idle and will be removed in 3 days unless used.", + Metadata: map[string]any{ + "worktree_id": w.ID.String(), + "branch": w.Branch, + }, + CreatedAt: time.Now().UTC(), + }); derr != nil { + r.log.Error("worktree_prune.dispatch", "id", w.ID, "err", derr.Error()) + } + } +} + +func (r *WorktreePrune) executionPhase(ctx context.Context) { + cutoff := time.Now().Add(-r.cfg.IdleThreshold) + rows, err := r.repo.ListWorktreesNeedingPrune(ctx, cutoff) + if err != nil { + r.log.Error("worktree_prune.list_prune", "err", err.Error()) + return + } + for _, w := range rows { + started, err := r.repo.TryStartPrune(ctx, w.ID) + if err != nil { + // CAS lost (idle → other status), skip + continue + } + if rmErr := r.gitr.WorktreeRemove(ctx, started.MainPath, started.Path); rmErr != nil { + r.log.Error("worktree_prune.remove_failed", "id", w.ID, "err", rmErr.Error()) + if rbErr := r.repo.FinishPruneRollback(ctx, w.ID); rbErr != nil { + r.log.Error("worktree_prune.rollback", "id", w.ID, "err", rbErr.Error()) + } + continue + } + if err := r.repo.FinishPruneSuccess(ctx, w.ID); err != nil { + r.log.Error("worktree_prune.finish_success", "id", w.ID, "err", err.Error()) + } + } +} diff --git a/internal/jobs/runners/worktree_prune_test.go b/internal/jobs/runners/worktree_prune_test.go new file mode 100644 index 0000000..b6d5240 --- /dev/null +++ b/internal/jobs/runners/worktree_prune_test.go @@ -0,0 +1,180 @@ +package runners_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/yan1h/agent-coding-workflow/internal/infra/notify" + "github.com/yan1h/agent-coding-workflow/internal/jobs/runners" +) + +// fakeWorktreeRepo is a fresh fake distinct from fakeWorkspaceFetchRepo. +type fakeWorktreeRepo struct { + warnList []runners.WorktreeRow + pruneList []runners.WorktreeRow + warnListErr error + pruneListErr error + + markedWarning []uuid.UUID + markWarningErr error + + tryStartResult runners.WorktreeRow + tryStartErr error + tryStartCalled map[uuid.UUID]bool + + finalized map[uuid.UUID]string // "ok" | "rollback" + finishSuccessErr error + finishRollbackErr error +} + +func (f *fakeWorktreeRepo) ListWorktreesNeedingPruneWarning(_ context.Context, _ time.Time) ([]runners.WorktreeRow, error) { + return f.warnList, f.warnListErr +} + +func (f *fakeWorktreeRepo) ListWorktreesNeedingPrune(_ context.Context, _ time.Time) ([]runners.WorktreeRow, error) { + return f.pruneList, f.pruneListErr +} + +func (f *fakeWorktreeRepo) MarkPruneWarning(_ context.Context, id uuid.UUID) error { + f.markedWarning = append(f.markedWarning, id) + return f.markWarningErr +} + +func (f *fakeWorktreeRepo) TryStartPrune(_ context.Context, id uuid.UUID) (runners.WorktreeRow, error) { + if f.tryStartCalled == nil { + f.tryStartCalled = map[uuid.UUID]bool{} + } + f.tryStartCalled[id] = true + if f.tryStartErr != nil { + return runners.WorktreeRow{}, f.tryStartErr + } + return f.tryStartResult, nil +} + +func (f *fakeWorktreeRepo) FinishPruneSuccess(_ context.Context, id uuid.UUID) error { + if f.finalized == nil { + f.finalized = map[uuid.UUID]string{} + } + f.finalized[id] = "ok" + return f.finishSuccessErr +} + +func (f *fakeWorktreeRepo) FinishPruneRollback(_ context.Context, id uuid.UUID) error { + if f.finalized == nil { + f.finalized = map[uuid.UUID]string{} + } + f.finalized[id] = "rollback" + return f.finishRollbackErr +} + +// fakePruner records WorktreeRemove calls. +type fakePruner struct { + err error + calls int +} + +func (f *fakePruner) WorktreeRemove(_ context.Context, _, _ string) error { + f.calls++ + return f.err +} + +func TestWorktreePrune_WarningPhase(t *testing.T) { + t.Parallel() + id := uuid.New() + ownerID := uuid.New() + row := runners.WorktreeRow{ID: id, OwnerID: ownerID, Branch: "feat-x", MainPath: "/main", Path: "/main/.worktrees/feat-x"} + + repo := &fakeWorktreeRepo{warnList: []runners.WorktreeRow{row}} + disp := &recordingDispatcher{} + gitr := &fakePruner{} + cfg := runners.WorktreePruneConfig{IdleThreshold: 7 * 24 * time.Hour, WarningLead: 3 * 24 * time.Hour} + + r := runners.NewWorktreePrune(repo, gitr, disp, cfg, discardLogger()) + r.Run(context.Background()) + + require.Len(t, repo.markedWarning, 1) + assert.Equal(t, id, repo.markedWarning[0]) + + require.Len(t, disp.msgs, 1) + msg := disp.msgs[0] + assert.Equal(t, ownerID, msg.UserID) + assert.Equal(t, "worktree.prune_pending", msg.Topic) + assert.Equal(t, notify.SeverityWarning, msg.Severity) + + // execution phase: pruneList is empty, so no prune calls. + assert.Zero(t, gitr.calls) +} + +func TestWorktreePrune_ExecutionPhaseSuccess(t *testing.T) { + t.Parallel() + id := uuid.New() + ownerID := uuid.New() + row := runners.WorktreeRow{ID: id, OwnerID: ownerID, Branch: "feat-y", MainPath: "/main", Path: "/main/.worktrees/feat-y"} + + repo := &fakeWorktreeRepo{ + pruneList: []runners.WorktreeRow{row}, + tryStartResult: row, + } + gitr := &fakePruner{} + cfg := runners.WorktreePruneConfig{IdleThreshold: 7 * 24 * time.Hour, WarningLead: 3 * 24 * time.Hour} + + r := runners.NewWorktreePrune(repo, gitr, &recordingDispatcher{}, cfg, discardLogger()) + r.Run(context.Background()) + + assert.True(t, repo.tryStartCalled[id]) + assert.Equal(t, 1, gitr.calls) + require.NotNil(t, repo.finalized) + assert.Equal(t, "ok", repo.finalized[id]) +} + +func TestWorktreePrune_ExecutionFailureRollsBack(t *testing.T) { + t.Parallel() + id := uuid.New() + ownerID := uuid.New() + row := runners.WorktreeRow{ID: id, OwnerID: ownerID, Branch: "feat-z", MainPath: "/main", Path: "/main/.worktrees/feat-z"} + + repo := &fakeWorktreeRepo{ + pruneList: []runners.WorktreeRow{row}, + tryStartResult: row, + } + gitr := &fakePruner{err: errors.New("git worktree remove failed")} + cfg := runners.WorktreePruneConfig{IdleThreshold: 7 * 24 * time.Hour, WarningLead: 3 * 24 * time.Hour} + + r := runners.NewWorktreePrune(repo, gitr, &recordingDispatcher{}, cfg, discardLogger()) + r.Run(context.Background()) + + assert.True(t, repo.tryStartCalled[id]) + assert.Equal(t, 1, gitr.calls) + require.NotNil(t, repo.finalized) + assert.Equal(t, "rollback", repo.finalized[id]) +} + +func TestWorktreePrune_CASLostSkipsSilently(t *testing.T) { + t.Parallel() + id := uuid.New() + ownerID := uuid.New() + row := runners.WorktreeRow{ID: id, OwnerID: ownerID, Branch: "feat-cas", MainPath: "/main", Path: "/main/.worktrees/feat-cas"} + + repo := &fakeWorktreeRepo{ + pruneList: []runners.WorktreeRow{row}, + tryStartErr: errors.New("worktree no longer idle"), + } + gitr := &fakePruner{} + cfg := runners.WorktreePruneConfig{IdleThreshold: 7 * 24 * time.Hour, WarningLead: 3 * 24 * time.Hour} + + r := runners.NewWorktreePrune(repo, gitr, &recordingDispatcher{}, cfg, discardLogger()) + r.Run(context.Background()) + + // TryStartPrune was called but failed. + assert.True(t, repo.tryStartCalled[id]) + // No git call. + assert.Zero(t, gitr.calls) + // No finalize for that ID. + assert.Empty(t, repo.finalized) +} diff --git a/internal/workspace/domain.go b/internal/workspace/domain.go index 44d7187..b7685f7 100644 --- a/internal/workspace/domain.go +++ b/internal/workspace/domain.go @@ -83,9 +83,10 @@ type Worktree struct { Path string Status WorktreeStatus ActiveHolder string // "user:" / "session:" / 空 - AcquiredAt *time.Time - LastUsedAt time.Time - CreatedAt time.Time + AcquiredAt *time.Time + LastUsedAt time.Time + PruneWarningAt *time.Time + CreatedAt time.Time } // Credential 是 git_credentials 表的领域投影。Secret 仅在创建/更新时有值, diff --git a/internal/workspace/queries/worktrees.sql b/internal/workspace/queries/worktrees.sql index 255e070..fbd38e7 100644 --- a/internal/workspace/queries/worktrees.sql +++ b/internal/workspace/queries/worktrees.sql @@ -49,3 +49,35 @@ RETURNING *; -- name: DeleteWorktree :exec DELETE FROM workspace_worktrees WHERE id = $1; + +-- name: ListWorktreesNeedingPruneWarning :many +SELECT * FROM workspace_worktrees + WHERE status = 'idle' + AND prune_warning_at IS NULL + AND last_used_at < $1 + ORDER BY last_used_at ASC; + +-- name: ListWorktreesNeedingPrune :many +SELECT * FROM workspace_worktrees + WHERE status = 'idle' + AND last_used_at < $1 + ORDER BY last_used_at ASC; + +-- name: MarkPruneWarning :exec +UPDATE workspace_worktrees + SET prune_warning_at = now() + WHERE id = $1 AND prune_warning_at IS NULL; + +-- name: TryStartPrune :one +UPDATE workspace_worktrees + SET status = 'pruning' + WHERE id = $1 AND status = 'idle' +RETURNING *; + +-- name: FinishPruneSuccess :exec +DELETE FROM workspace_worktrees WHERE id = $1; + +-- name: FinishPruneRollback :exec +UPDATE workspace_worktrees + SET status = 'idle' + WHERE id = $1 AND status = 'pruning'; diff --git a/internal/workspace/repository.go b/internal/workspace/repository.go index 97113aa..a023687 100644 --- a/internal/workspace/repository.go +++ b/internal/workspace/repository.go @@ -50,6 +50,12 @@ type Repository interface { GetWorktreeByID(ctx context.Context, id uuid.UUID) (*Worktree, error) ListWorktreesByWorkspace(ctx context.Context, wsID uuid.UUID) ([]*Worktree, error) DeleteWorktree(ctx context.Context, id uuid.UUID) error + ListWorktreesNeedingPruneWarning(ctx context.Context, lastUsedBefore time.Time) ([]*Worktree, error) + ListWorktreesNeedingPrune(ctx context.Context, lastUsedBefore time.Time) ([]*Worktree, error) + MarkPruneWarning(ctx context.Context, id uuid.UUID) error + TryStartPrune(ctx context.Context, id uuid.UUID) (*Worktree, error) + FinishPruneSuccess(ctx context.Context, id uuid.UUID) error + FinishPruneRollback(ctx context.Context, id uuid.UUID) error // Credential UpsertCredential(ctx context.Context, c *Credential) (*Credential, error) @@ -240,15 +246,16 @@ func (r *pgRepo) ResetStuckSyncStatuses(ctx context.Context) error { func rowToWorktree(r workspacesqlc.WorkspaceWorktree) *Worktree { return &Worktree{ - ID: fromPgUUID(r.ID), - WorkspaceID: fromPgUUID(r.WorkspaceID), - Branch: r.Branch, - Path: r.Path, - Status: WorktreeStatus(r.Status), - ActiveHolder: fromPtrString(r.ActiveHolder), - AcquiredAt: fromPgTimePtr(r.AcquiredAt), - LastUsedAt: r.LastUsedAt.Time, - CreatedAt: r.CreatedAt.Time, + ID: fromPgUUID(r.ID), + WorkspaceID: fromPgUUID(r.WorkspaceID), + Branch: r.Branch, + Path: r.Path, + Status: WorktreeStatus(r.Status), + ActiveHolder: fromPtrString(r.ActiveHolder), + AcquiredAt: fromPgTimePtr(r.AcquiredAt), + LastUsedAt: r.LastUsedAt.Time, + PruneWarningAt: fromPgTimePtr(r.PruneWarningAt), + CreatedAt: r.CreatedAt.Time, } } @@ -283,6 +290,64 @@ func (r *pgRepo) DeleteWorktree(ctx context.Context, id uuid.UUID) error { return nil } +func (r *pgRepo) ListWorktreesNeedingPruneWarning(ctx context.Context, lastUsedBefore time.Time) ([]*Worktree, error) { + ts := pgtype.Timestamptz{Time: lastUsedBefore, Valid: true} + rows, err := r.q.ListWorktreesNeedingPruneWarning(ctx, ts) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list worktrees needing prune warning") + } + out := make([]*Worktree, 0, len(rows)) + for _, row := range rows { + out = append(out, rowToWorktree(row)) + } + return out, nil +} + +func (r *pgRepo) ListWorktreesNeedingPrune(ctx context.Context, lastUsedBefore time.Time) ([]*Worktree, error) { + ts := pgtype.Timestamptz{Time: lastUsedBefore, Valid: true} + rows, err := r.q.ListWorktreesNeedingPrune(ctx, ts) + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "list worktrees needing prune") + } + out := make([]*Worktree, 0, len(rows)) + for _, row := range rows { + out = append(out, rowToWorktree(row)) + } + return out, nil +} + +func (r *pgRepo) MarkPruneWarning(ctx context.Context, id uuid.UUID) error { + if err := r.q.MarkPruneWarning(ctx, toPgUUID(id)); err != nil { + return errs.Wrap(err, errs.CodeInternal, "mark prune warning") + } + return nil +} + +func (r *pgRepo) TryStartPrune(ctx context.Context, id uuid.UUID) (*Worktree, error) { + row, err := r.q.TryStartPrune(ctx, toPgUUID(id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeWorktreeAlreadyActive, "worktree no longer idle") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "try start prune") + } + return rowToWorktree(row), nil +} + +func (r *pgRepo) FinishPruneSuccess(ctx context.Context, id uuid.UUID) error { + if err := r.q.FinishPruneSuccess(ctx, toPgUUID(id)); err != nil { + return errs.Wrap(err, errs.CodeInternal, "finish prune success") + } + return nil +} + +func (r *pgRepo) FinishPruneRollback(ctx context.Context, id uuid.UUID) error { + if err := r.q.FinishPruneRollback(ctx, toPgUUID(id)); err != nil { + return errs.Wrap(err, errs.CodeInternal, "finish prune rollback") + } + return nil +} + // ===== Credential ===== func rowToCredential(r workspacesqlc.GitCredential) *Credential { diff --git a/internal/workspace/repository_test.go b/internal/workspace/repository_test.go index a35481c..7d6d27c 100644 --- a/internal/workspace/repository_test.go +++ b/internal/workspace/repository_test.go @@ -13,6 +13,7 @@ import ( tcpg "github.com/testcontainers/testcontainers-go/modules/postgres" "github.com/yan1h/agent-coding-workflow/internal/infra/db" + "github.com/yan1h/agent-coding-workflow/internal/infra/errs" ) // setupRepo 起一个临时 postgres 容器、跑迁移、注入一条 user/project 行做外键, @@ -186,3 +187,83 @@ func TestRepo_SetWorktreeActive_ClearsPruneWarning(t *testing.T) { require.NoError(t, pool.QueryRow(ctx, `SELECT prune_warning_at FROM workspace_worktrees WHERE id = $1`, wtID).Scan(&pwAt)) require.Nil(t, pwAt, "SetWorktreeIdle should clear prune_warning_at") } + +// TestRepo_TryStartPrune_CASLossReturnsAlreadyActive 验证:对一个已变为 active 的 +// worktree 调用 TryStartPrune 会返回 CodeWorktreeAlreadyActive。 +func TestRepo_TryStartPrune_CASLossReturnsAlreadyActive(t *testing.T) { + t.Parallel() + repo, pool, pid := setupRepo(t) + ctx := context.Background() + wsID := uuid.New() + _, err := repo.CreateWorkspace(ctx, &Workspace{ + ID: wsID, ProjectID: pid, Slug: "ws-cas", Name: "CAS", + GitRemoteURL: "https://x", DefaultBranch: "main", MainPath: "/x", SyncStatus: SyncStatusIdle, + }) + require.NoError(t, err) + + wtID := uuid.New() + require.NoError(t, repo.InTx(ctx, func(tx Tx) error { + _, e := tx.InsertWorktree(ctx, &Worktree{ + ID: wtID, WorkspaceID: wsID, Branch: "cas-branch", + Path: "/x/.worktrees/cas", Status: WorktreeStatusIdle, + }) + return e + })) + + // 绕过 service 直接把 status 改成 active,模拟 CAS 竞争丢失。 + _, err = pool.Exec(ctx, `UPDATE workspace_worktrees SET status = 'active' WHERE id = $1`, wtID) + require.NoError(t, err) + + _, err = repo.TryStartPrune(ctx, wtID) + require.Error(t, err) + appErr, ok := errs.As(err) + require.True(t, ok, "expected AppError") + require.Equal(t, errs.CodeWorktreeAlreadyActive, appErr.Code) +} + +// TestRepo_ListWorktreesNeedingPruneWarning_FiltersExisting 验证:已有 +// prune_warning_at 的 worktree 不会被返回,仅返回尚未标记的。 +func TestRepo_ListWorktreesNeedingPruneWarning_FiltersExisting(t *testing.T) { + t.Parallel() + repo, pool, pid := setupRepo(t) + ctx := context.Background() + wsID := uuid.New() + _, err := repo.CreateWorkspace(ctx, &Workspace{ + ID: wsID, ProjectID: pid, Slug: "ws-pwf", Name: "PWF", + GitRemoteURL: "https://x", DefaultBranch: "main", MainPath: "/x", SyncStatus: SyncStatusIdle, + }) + require.NoError(t, err) + + cutoff := time.Now().Add(-24 * time.Hour) + + // worktree A:没有 prune_warning_at,last_used_at 在 cutoff 之前 → 应返回。 + wtA := uuid.New() + require.NoError(t, repo.InTx(ctx, func(tx Tx) error { + _, e := tx.InsertWorktree(ctx, &Worktree{ + ID: wtA, WorkspaceID: wsID, Branch: "prune-a", + Path: "/x/.worktrees/a", Status: WorktreeStatusIdle, + }) + return e + })) + _, err = pool.Exec(ctx, `UPDATE workspace_worktrees SET last_used_at = $1 WHERE id = $2`, cutoff.Add(-time.Hour), wtA) + require.NoError(t, err) + + // worktree B:已经有 prune_warning_at,last_used_at 同样在 cutoff 前 → 不应返回。 + wtB := uuid.New() + require.NoError(t, repo.InTx(ctx, func(tx Tx) error { + _, e := tx.InsertWorktree(ctx, &Worktree{ + ID: wtB, WorkspaceID: wsID, Branch: "prune-b", + Path: "/x/.worktrees/b", Status: WorktreeStatusIdle, + }) + return e + })) + _, err = pool.Exec(ctx, + `UPDATE workspace_worktrees SET last_used_at = $1, prune_warning_at = now() WHERE id = $2`, + cutoff.Add(-time.Hour), wtB) + require.NoError(t, err) + + rows, err := repo.ListWorktreesNeedingPruneWarning(ctx, cutoff) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, wtA, rows[0].ID) +} diff --git a/internal/workspace/sqlc/worktrees.sql.go b/internal/workspace/sqlc/worktrees.sql.go index c7f4eba..c99a50d 100644 --- a/internal/workspace/sqlc/worktrees.sql.go +++ b/internal/workspace/sqlc/worktrees.sql.go @@ -63,6 +63,26 @@ func (q *Queries) DeleteWorktree(ctx context.Context, id pgtype.UUID) error { return err } +const finishPruneRollback = `-- name: FinishPruneRollback :exec +UPDATE workspace_worktrees + SET status = 'idle' + WHERE id = $1 AND status = 'pruning' +` + +func (q *Queries) FinishPruneRollback(ctx context.Context, id pgtype.UUID) error { + _, err := q.db.Exec(ctx, finishPruneRollback, id) + return err +} + +const finishPruneSuccess = `-- name: FinishPruneSuccess :exec +DELETE FROM workspace_worktrees WHERE id = $1 +` + +func (q *Queries) FinishPruneSuccess(ctx context.Context, id pgtype.UUID) error { + _, err := q.db.Exec(ctx, finishPruneSuccess, id) + return err +} + const getWorktreeByBranchForUpdate = `-- name: GetWorktreeByBranchForUpdate :one SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at, prune_warning_at FROM workspace_worktrees WHERE workspace_id = $1 AND branch = $2 @@ -174,6 +194,94 @@ func (q *Queries) ListWorktreesByWorkspace(ctx context.Context, workspaceID pgty return items, nil } +const listWorktreesNeedingPrune = `-- name: ListWorktreesNeedingPrune :many +SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at, prune_warning_at FROM workspace_worktrees + WHERE status = 'idle' + AND last_used_at < $1 + ORDER BY last_used_at ASC +` + +func (q *Queries) ListWorktreesNeedingPrune(ctx context.Context, lastUsedAt pgtype.Timestamptz) ([]WorkspaceWorktree, error) { + rows, err := q.db.Query(ctx, listWorktreesNeedingPrune, lastUsedAt) + if err != nil { + return nil, err + } + defer rows.Close() + var items []WorkspaceWorktree + for rows.Next() { + var i WorkspaceWorktree + if err := rows.Scan( + &i.ID, + &i.WorkspaceID, + &i.Branch, + &i.Path, + &i.Status, + &i.ActiveHolder, + &i.AcquiredAt, + &i.LastUsedAt, + &i.CreatedAt, + &i.PruneWarningAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listWorktreesNeedingPruneWarning = `-- name: ListWorktreesNeedingPruneWarning :many +SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at, prune_warning_at FROM workspace_worktrees + WHERE status = 'idle' + AND prune_warning_at IS NULL + AND last_used_at < $1 + ORDER BY last_used_at ASC +` + +func (q *Queries) ListWorktreesNeedingPruneWarning(ctx context.Context, lastUsedAt pgtype.Timestamptz) ([]WorkspaceWorktree, error) { + rows, err := q.db.Query(ctx, listWorktreesNeedingPruneWarning, lastUsedAt) + if err != nil { + return nil, err + } + defer rows.Close() + var items []WorkspaceWorktree + for rows.Next() { + var i WorkspaceWorktree + if err := rows.Scan( + &i.ID, + &i.WorkspaceID, + &i.Branch, + &i.Path, + &i.Status, + &i.ActiveHolder, + &i.AcquiredAt, + &i.LastUsedAt, + &i.CreatedAt, + &i.PruneWarningAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const markPruneWarning = `-- name: MarkPruneWarning :exec +UPDATE workspace_worktrees + SET prune_warning_at = now() + WHERE id = $1 AND prune_warning_at IS NULL +` + +func (q *Queries) MarkPruneWarning(ctx context.Context, id pgtype.UUID) error { + _, err := q.db.Exec(ctx, markPruneWarning, id) + return err +} + const setWorktreeActive = `-- name: SetWorktreeActive :one UPDATE workspace_worktrees SET status = 'active', @@ -261,3 +369,28 @@ func (q *Queries) SetWorktreePruning(ctx context.Context, id pgtype.UUID) (Works ) return i, err } + +const tryStartPrune = `-- name: TryStartPrune :one +UPDATE workspace_worktrees + SET status = 'pruning' + WHERE id = $1 AND status = 'idle' +RETURNING id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at, prune_warning_at +` + +func (q *Queries) TryStartPrune(ctx context.Context, id pgtype.UUID) (WorkspaceWorktree, error) { + row := q.db.QueryRow(ctx, tryStartPrune, id) + var i WorkspaceWorktree + err := row.Scan( + &i.ID, + &i.WorkspaceID, + &i.Branch, + &i.Path, + &i.Status, + &i.ActiveHolder, + &i.AcquiredAt, + &i.LastUsedAt, + &i.CreatedAt, + &i.PruneWarningAt, + ) + return i, err +} diff --git a/internal/workspace/workspace_service_test.go b/internal/workspace/workspace_service_test.go index 0609c05..ae5502a 100644 --- a/internal/workspace/workspace_service_test.go +++ b/internal/workspace/workspace_service_test.go @@ -86,6 +86,18 @@ func (f *fakeRepo) ListWorktreesByWorkspace(_ context.Context, _ uuid.UUID) ([]* return nil, nil } func (f *fakeRepo) DeleteWorktree(_ context.Context, _ uuid.UUID) error { return nil } +func (f *fakeRepo) ListWorktreesNeedingPruneWarning(_ context.Context, _ time.Time) ([]*Worktree, error) { + return nil, nil +} +func (f *fakeRepo) ListWorktreesNeedingPrune(_ context.Context, _ time.Time) ([]*Worktree, error) { + return nil, nil +} +func (f *fakeRepo) MarkPruneWarning(_ context.Context, _ uuid.UUID) error { return nil } +func (f *fakeRepo) TryStartPrune(_ context.Context, _ uuid.UUID) (*Worktree, error) { + return nil, nil +} +func (f *fakeRepo) FinishPruneSuccess(_ context.Context, _ uuid.UUID) error { return nil } +func (f *fakeRepo) FinishPruneRollback(_ context.Context, _ uuid.UUID) error { return nil } func (f *fakeRepo) UpsertCredential(_ context.Context, c *Credential) (*Credential, error) { cp := *c f.creds[c.WorkspaceID] = &cp