You've already forked agentic-coding-workflow
feat(workspace): WorktreeService with acquire/release state machine
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user