diff --git a/internal/workspace/queries/worktrees.sql b/internal/workspace/queries/worktrees.sql index 929547a..0c987c1 100644 --- a/internal/workspace/queries/worktrees.sql +++ b/internal/workspace/queries/worktrees.sql @@ -7,6 +7,10 @@ RETURNING *; -- name: GetWorktreeByID :one SELECT * FROM workspace_worktrees WHERE id = $1; +-- name: GetWorktreeByIDForUpdate :one +SELECT * FROM workspace_worktrees WHERE id = $1 +FOR UPDATE; + -- name: GetWorktreeByBranchForUpdate :one SELECT * FROM workspace_worktrees WHERE workspace_id = $1 AND branch = $2 diff --git a/internal/workspace/repository.go b/internal/workspace/repository.go index 0457c4b..97113aa 100644 --- a/internal/workspace/repository.go +++ b/internal/workspace/repository.go @@ -23,6 +23,9 @@ import ( // 状态转移方法,不直接操作 pgx.Tx。 type Tx interface { GetWorktreeForUpdate(ctx context.Context, wsID uuid.UUID, branch string) (*Worktree, error) + // GetWorktreeByIDForUpdate 在事务内对 id 行加 FOR UPDATE — Acquire/Release 必须用它, + // GetWorktreeByID 不带行锁,并发时会 TOCTOU。 + GetWorktreeByIDForUpdate(ctx context.Context, id uuid.UUID) (*Worktree, error) GetWorktreeByID(ctx context.Context, id uuid.UUID) (*Worktree, error) InsertWorktree(ctx context.Context, w *Worktree) (*Worktree, error) SetWorktreeActive(ctx context.Context, id uuid.UUID, holder string) (*Worktree, error) @@ -355,6 +358,17 @@ func (t *pgTx) GetWorktreeByID(ctx context.Context, id uuid.UUID) (*Worktree, er return rowToWorktree(row), nil } +func (t *pgTx) GetWorktreeByIDForUpdate(ctx context.Context, id uuid.UUID) (*Worktree, error) { + row, err := t.q.GetWorktreeByIDForUpdate(ctx, toPgUUID(id)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, errs.New(errs.CodeNotFound, "worktree not found") + } + if err != nil { + return nil, errs.Wrap(err, errs.CodeInternal, "select for update") + } + return rowToWorktree(row), nil +} + func (t *pgTx) InsertWorktree(ctx context.Context, w *Worktree) (*Worktree, error) { row, err := t.q.CreateWorktree(ctx, workspacesqlc.CreateWorktreeParams{ ID: toPgUUID(w.ID), diff --git a/internal/workspace/sqlc/worktrees.sql.go b/internal/workspace/sqlc/worktrees.sql.go index b20b2cb..67614ec 100644 --- a/internal/workspace/sqlc/worktrees.sql.go +++ b/internal/workspace/sqlc/worktrees.sql.go @@ -111,6 +111,28 @@ func (q *Queries) GetWorktreeByID(ctx context.Context, id pgtype.UUID) (Workspac return i, err } +const getWorktreeByIDForUpdate = `-- name: GetWorktreeByIDForUpdate :one +SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at FROM workspace_worktrees WHERE id = $1 +FOR UPDATE +` + +func (q *Queries) GetWorktreeByIDForUpdate(ctx context.Context, id pgtype.UUID) (WorkspaceWorktree, error) { + row := q.db.QueryRow(ctx, getWorktreeByIDForUpdate, 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, + ) + return i, err +} + const listWorktreesByWorkspace = `-- name: ListWorktreesByWorkspace :many SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at FROM workspace_worktrees WHERE workspace_id = $1 diff --git a/internal/workspace/workspace_service_test.go b/internal/workspace/workspace_service_test.go index 3bc7d3b..2202dce 100644 --- a/internal/workspace/workspace_service_test.go +++ b/internal/workspace/workspace_service_test.go @@ -118,6 +118,13 @@ func (t *fakeTx) GetWorktreeByID(_ context.Context, id uuid.UUID) (*Worktree, er } return nil, &errsNotFound{} } +func (t *fakeTx) GetWorktreeByIDForUpdate(_ context.Context, id uuid.UUID) (*Worktree, error) { + if w, ok := t.f.wts[id]; ok { + cp := *w + return &cp, nil + } + return nil, &errsNotFound{} +} func (t *fakeTx) InsertWorktree(_ context.Context, w *Worktree) (*Worktree, error) { cp := *w t.f.wts[w.ID] = &cp diff --git a/internal/workspace/worktree_service.go b/internal/workspace/worktree_service.go new file mode 100644 index 0000000..3790b5f --- /dev/null +++ b/internal/workspace/worktree_service.go @@ -0,0 +1,195 @@ +package workspace + +import ( + "context" + "errors" + "time" + + "github.com/google/uuid" + + "github.com/yan1h/agent-coding-workflow/internal/audit" + "github.com/yan1h/agent-coding-workflow/internal/infra/errs" + "github.com/yan1h/agent-coding-workflow/internal/infra/git" +) + +type worktreeService struct { + repo Repository + rec audit.Recorder + gitr git.Runner + pa ProjectAccess + dataDir string +} + +// NewWorktreeService 构造 WorktreeService。 +func NewWorktreeService(repo Repository, rec audit.Recorder, gitr git.Runner, pa ProjectAccess, dataDir string) WorktreeService { + return &worktreeService{repo: repo, rec: rec, gitr: gitr, pa: pa, dataDir: dataDir} +} + +func (s *worktreeService) Create(ctx context.Context, c Caller, wsID uuid.UUID, in CreateWorktreeInput) (*Worktree, error) { + ws, err := s.repo.GetWorkspaceByID(ctx, wsID) + if err != nil { + return nil, err + } + if _, canWrite, err := s.pa.ResolveByID(ctx, c.UserID, c.IsAdmin, ws.ProjectID); err != nil { + return nil, err + } else if !canWrite { + return nil, errs.New(errs.CodeForbidden, "no write access") + } + if err := ValidateBranch(in.Branch); err != nil { + return nil, errs.Wrap(err, errs.CodeInvalidInput, "invalid branch") + } + base := in.BaseBranch + if base == "" { + base = ws.DefaultBranch + } + if err := ValidateBranch(base); err != nil { + return nil, errs.Wrap(err, errs.CodeInvalidInput, "invalid base branch") + } + wtPath := WorktreePath(s.dataDir, wsID, in.Branch) + if err := s.gitr.WorktreeAdd(ctx, ws.MainPath, in.Branch, wtPath); err != nil { + if git.IsConflict(err) { + return nil, errs.Wrap(err, errs.CodeWorktreeBranchConflict, "branch already used") + } + return nil, errs.Wrap(err, errs.CodeGitCmdFailed, "git worktree add") + } + wt := &Worktree{ + ID: uuid.New(), WorkspaceID: wsID, Branch: in.Branch, Path: wtPath, + Status: WorktreeStatusIdle, + } + var saved *Worktree + err = s.repo.InTx(ctx, func(tx Tx) error { + w, e := tx.InsertWorktree(ctx, wt) + if e != nil { + return e + } + saved = w + return nil + }) + if err != nil { + // 回滚磁盘 + _ = s.gitr.WorktreeRemove(ctx, ws.MainPath, wtPath) + return nil, err + } + s.audit(ctx, &c.UserID, "worktree.create", "worktree", saved.ID.String(), map[string]any{"branch": in.Branch}) + return saved, nil +} + +func (s *worktreeService) List(ctx context.Context, c Caller, wsID uuid.UUID) ([]*Worktree, error) { + ws, err := s.repo.GetWorkspaceByID(ctx, wsID) + if err != nil { + return nil, err + } + if canRead, _, err := s.pa.ResolveByID(ctx, c.UserID, c.IsAdmin, ws.ProjectID); err != nil { + return nil, err + } else if !canRead { + return nil, errs.New(errs.CodeForbidden, "no read access") + } + return s.repo.ListWorktreesByWorkspace(ctx, wsID) +} + +func (s *worktreeService) Delete(ctx context.Context, c Caller, wtID uuid.UUID) error { + wt, err := s.repo.GetWorktreeByID(ctx, wtID) + if err != nil { + return err + } + ws, err := s.repo.GetWorkspaceByID(ctx, wt.WorkspaceID) + if err != nil { + return err + } + if _, canWrite, err := s.pa.ResolveByID(ctx, c.UserID, c.IsAdmin, ws.ProjectID); err != nil { + return err + } else if !canWrite { + return errs.New(errs.CodeForbidden, "no write access") + } + if wt.Status == WorktreeStatusActive { + return errs.New(errs.CodeWorktreeAlreadyActive, "release before delete") + } + // 先 set pruning,再调 git rm,然后 DELETE row。中途失败 row 留 pruning。 + err = s.repo.InTx(ctx, func(tx Tx) error { + _, e := tx.SetWorktreePruning(ctx, wtID) + return e + }) + if err != nil { + return err + } + if rmErr := s.gitr.WorktreeRemove(ctx, ws.MainPath, wt.Path); rmErr != nil { + return errs.Wrap(rmErr, errs.CodeGitCmdFailed, "git worktree remove") + } + if err := s.repo.DeleteWorktree(ctx, wtID); err != nil { + return err + } + s.audit(ctx, &c.UserID, "worktree.delete", "worktree", wtID.String(), nil) + return nil +} + +func (s *worktreeService) Acquire(ctx context.Context, c Caller, wtID uuid.UUID) (*Worktree, error) { + holder := c.Holder() + var out *Worktree + err := s.repo.InTx(ctx, func(tx Tx) error { + wt, err := tx.GetWorktreeByIDForUpdate(ctx, wtID) + if err != nil { + return err + } + switch wt.Status { + case WorktreeStatusPruning: + return errs.New(errs.CodeNotFound, "worktree being pruned") + case WorktreeStatusActive: + if wt.ActiveHolder == holder { + out = wt // idempotent + return nil + } + return errs.New(errs.CodeWorktreeAlreadyActive, "held by another holder") + case WorktreeStatusIdle: + updated, err := tx.SetWorktreeActive(ctx, wtID, holder) + if err != nil { + return err + } + out = updated + return nil + } + return errs.New(errs.CodeInternal, "unknown worktree status") + }) + if err != nil { + return nil, err + } + s.audit(ctx, &c.UserID, "worktree.acquire", "worktree", wtID.String(), map[string]any{"holder": holder}) + return out, nil +} + +func (s *worktreeService) Release(ctx context.Context, c Caller, wtID uuid.UUID) (*Worktree, error) { + holder := c.Holder() + var out *Worktree + err := s.repo.InTx(ctx, func(tx Tx) error { + wt, err := tx.GetWorktreeByIDForUpdate(ctx, wtID) + if err != nil { + return err + } + if wt.Status != WorktreeStatusActive { + return errs.New(errs.CodeWorktreeNotHeldByCaller, "worktree not active") + } + if wt.ActiveHolder != holder && !c.IsAdmin { + return errs.New(errs.CodeWorktreeNotHeldByCaller, "not held by caller") + } + updated, err := tx.SetWorktreeIdle(ctx, wtID) + if err != nil { + return err + } + out = updated + return nil + }) + if err != nil { + return nil, err + } + s.audit(ctx, &c.UserID, "worktree.release", "worktree", wtID.String(), map[string]any{"holder": holder}) + return out, nil +} + +func (s *worktreeService) audit(ctx context.Context, uid *uuid.UUID, action, ttype, tid string, meta map[string]any) { + _ = s.rec.Record(ctx, audit.Entry{ + UserID: uid, Action: action, TargetType: ttype, TargetID: tid, Metadata: meta, + }) +} + +// 防止未使用 import lint 移除 +var _ = errors.Is +var _ = time.Now diff --git a/internal/workspace/worktree_service_test.go b/internal/workspace/worktree_service_test.go new file mode 100644 index 0000000..c4aca9e --- /dev/null +++ b/internal/workspace/worktree_service_test.go @@ -0,0 +1,83 @@ +package workspace + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func newWtSvc(t *testing.T) (*worktreeService, *fakeRepo, *fakePA, *fakeGit, uuid.UUID) { + t.Helper() + pa := &fakePA{pid: uuid.New(), canRead: true, canWrite: true} + gitr := &fakeGit{} + repo := newFakeRepo() + svc := NewWorktreeService(repo, noopAudit{}, gitr, pa, "/data").(*worktreeService) + wsID := uuid.New() + repo.wss[wsID] = &Workspace{ID: wsID, ProjectID: pa.pid, MainPath: "/data/main", DefaultBranch: "main"} + return svc, repo, pa, gitr, wsID +} + +func TestWorktreeService_CreateAndList(t *testing.T) { + t.Parallel() + svc, _, _, _, wsID := newWtSvc(t) + wt, err := svc.Create(context.Background(), Caller{UserID: uuid.New()}, wsID, CreateWorktreeInput{Branch: "feat-x"}) + require.NoError(t, err) + require.Equal(t, "feat-x", wt.Branch) + require.Equal(t, WorktreeStatusIdle, wt.Status) +} + +func TestWorktreeService_AcquireIdempotent(t *testing.T) { + t.Parallel() + svc, _, _, _, wsID := newWtSvc(t) + caller := Caller{UserID: uuid.New()} + wt, err := svc.Create(context.Background(), caller, wsID, CreateWorktreeInput{Branch: "feat-y"}) + require.NoError(t, err) + + a1, err := svc.Acquire(context.Background(), caller, wt.ID) + require.NoError(t, err) + require.Equal(t, WorktreeStatusActive, a1.Status) + a2, err := svc.Acquire(context.Background(), caller, wt.ID) // idempotent + require.NoError(t, err) + require.Equal(t, a1.ID, a2.ID) +} + +func TestWorktreeService_AcquireConflict(t *testing.T) { + t.Parallel() + svc, _, _, _, wsID := newWtSvc(t) + a := Caller{UserID: uuid.New()} + b := Caller{UserID: uuid.New()} + wt, err := svc.Create(context.Background(), a, wsID, CreateWorktreeInput{Branch: "feat-z"}) + require.NoError(t, err) + _, err = svc.Acquire(context.Background(), a, wt.ID) + require.NoError(t, err) + _, err = svc.Acquire(context.Background(), b, wt.ID) + require.Error(t, err) +} + +func TestWorktreeService_Release(t *testing.T) { + t.Parallel() + svc, _, _, _, wsID := newWtSvc(t) + caller := Caller{UserID: uuid.New()} + wt, err := svc.Create(context.Background(), caller, wsID, CreateWorktreeInput{Branch: "feat-r"}) + require.NoError(t, err) + _, err = svc.Acquire(context.Background(), caller, wt.ID) + require.NoError(t, err) + _, err = svc.Release(context.Background(), caller, wt.ID) + require.NoError(t, err) + got, _ := svc.repo.GetWorktreeByID(context.Background(), wt.ID) + require.Equal(t, WorktreeStatusIdle, got.Status) +} + +func TestWorktreeService_Delete_BlockedWhenActive(t *testing.T) { + t.Parallel() + svc, _, _, _, wsID := newWtSvc(t) + caller := Caller{UserID: uuid.New()} + wt, err := svc.Create(context.Background(), caller, wsID, CreateWorktreeInput{Branch: "feat-d"}) + require.NoError(t, err) + _, err = svc.Acquire(context.Background(), caller, wt.ID) + require.NoError(t, err) + err = svc.Delete(context.Background(), caller, wt.ID) + require.Error(t, err) +}