agentic-coding-workflow/internal/workspace/worktree_service.go

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