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"
)

// IndexBuildHook is a best-effort callback fired after a worktree is created (or
// a workspace is synced) to enqueue a code-index build. Injected from app.go via
// SetIndexBuildHook to avoid a workspace→codeindex import cycle. May be nil.
type IndexBuildHook func(ctx context.Context, wsID uuid.UUID, worktreeID *uuid.UUID, commitSHA string) error

type worktreeService struct {
	repo      Repository
	rec       audit.Recorder
	gitr      git.Runner
	pa        ProjectAccess
	dataDir   string
	indexHook IndexBuildHook
}

// 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}
}

// SetIndexBuildHook injects the code-index enqueue callback (app wiring). Safe to
// call once at construction; nil disables index builds on worktree create.
func (s *worktreeService) SetIndexBuildHook(h IndexBuildHook) { s.indexHook = h }

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, "origin/"+base); 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})
	// Best-effort: enqueue a code-index build for the new worktree at its HEAD.
	// Failures are logged via audit but never block worktree creation.
	if s.indexHook != nil {
		wtID := saved.ID
		if herr := s.indexHook(ctx, wsID, &wtID, ""); herr != nil {
			s.audit(ctx, &c.UserID, "worktree.index_enqueue_failed", "worktree", saved.ID.String(), map[string]any{"err": herr.Error()})
		}
	}
	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 {
	ws, err := s.authorizeWorktreeWrite(ctx, c, wtID)
	if err != nil {
		return err
	}
	// check-and-transition 必须原子：在 tx 内对行加 FOR UPDATE，确认 idle 后才转 pruning。
	// 否则并发 Acquire（FOR UPDATE）可能刚把它置为 active，而这里用非锁读看到旧的 idle，
	// 进而 git rm 掉 agent 正在使用的目录。仅 idle→pruning 在锁下完成。
	var wtPath string
	err = s.repo.InTx(ctx, func(tx Tx) error {
		wt, e := tx.GetWorktreeByIDForUpdate(ctx, wtID)
		if e != nil {
			return e
		}
		if wt.Status != WorktreeStatusIdle {
			return errs.New(errs.CodeWorktreeAlreadyActive, "release before delete")
		}
		pruned, e := tx.SetWorktreePruning(ctx, wtID)
		if e != nil {
			return e
		}
		wtPath = pruned.Path
		return nil
	})
	if err != nil {
		return err
	}
	// 仅在守卫后的转移成功后才动磁盘。中途失败 row 留 pruning。
	if rmErr := s.gitr.WorktreeRemove(ctx, ws.MainPath, wtPath); 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) {
	if _, err := s.authorizeWorktreeWrite(ctx, c, wtID); err != nil {
		return nil, err
	}
	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) {
	if _, err := s.authorizeWorktreeWrite(ctx, c, wtID); err != nil {
		return nil, err
	}
	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
}

// ReleaseByHolder 释放给定 holder 字符串持有的所有活跃 worktree。
// 实现仅委托 repo；audit 由 caller 在更高层根据自己上下文写入。
func (s *worktreeService) ReleaseByHolder(ctx context.Context, holder string) ([]*Worktree, error) {
	return s.repo.ReleaseWorktreesByHolder(ctx, holder)
}

// authorizeWorktreeWrite 把 worktree → workspace → project 解析出来并要求 write 权限，
// 与 Delete/Create 的鉴权方式一致。返回 worktree 所属的 workspace 供调用方使用
// （如取 MainPath）。Acquire/Release 此前缺失该校验，导致任意已登录用户可凭 UUID
// 跨租户锁定/释放任意 worktree。
func (s *worktreeService) authorizeWorktreeWrite(ctx context.Context, c Caller, wtID uuid.UUID) (*Workspace, error) {
	wt, err := s.repo.GetWorktreeByID(ctx, wtID)
	if err != nil {
		return nil, err
	}
	ws, err := s.repo.GetWorkspaceByID(ctx, wt.WorkspaceID)
	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")
	}
	return ws, 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