Add orch strict worktree dispatch

2026-03-19 13:52:51 +08:00
parent 07f4a6fdae
commit 1b0cd723d7
9 changed files with 776 additions and 48 deletions
@@ -0,0 +1,303 @@
+package orch
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+
+	"ai-workflow-skill/internal/protocol"
+	"ai-workflow-skill/internal/store"
+
+	"github.com/spf13/cobra"
+)
+
+func newDispatchWorkspacePreparer(cmd *cobra.Command, opts dispatchOptions) store.DispatchWorkspacePreparer {
+	if !dispatchUsesWorktree(opts) {
+		return nil
+	}
+
+	ctx := cmd.Context()
+
+	return func(task store.Task, attemptNo int) (store.DispatchWorkspace, func(), error) {
+		return provisionDispatchWorkspace(ctx, opts, task, attemptNo)
+	}
+}
+
+func dispatchUsesWorktree(opts dispatchOptions) bool {
+	return strings.TrimSpace(opts.repoPath) != "" ||
+		strings.TrimSpace(opts.workspaceRoot) != "" ||
+		opts.strictWorktree
+}
+
+func provisionDispatchWorkspace(ctx context.Context, opts dispatchOptions, task store.Task, attemptNo int) (store.DispatchWorkspace, func(), error) {
+	repoRoot, err := resolveRepoRoot(ctx, opts.repoPath)
+	if err != nil {
+		return store.DispatchWorkspace{}, nil, err
+	}
+
+	workspaceRoot := resolveWorkspaceRoot(repoRoot, opts.workspaceRoot)
+	if err := ensureWorkspaceRootIgnored(repoRoot, workspaceRoot); err != nil {
+		return store.DispatchWorkspace{}, nil, err
+	}
+
+	baseRef, baseCommit, err := resolveDispatchBase(ctx, repoRoot, workspaceRoot, opts.baseRef, opts.strictWorktree)
+	if err != nil {
+		return store.DispatchWorkspace{}, nil, err
+	}
+
+	branchName := buildAttemptBranchName(task.RunID, task.TaskID, attemptNo)
+	worktreePath := buildAttemptWorktreePath(workspaceRoot, task.RunID, task.TaskID, attemptNo)
+
+	if err := os.MkdirAll(filepath.Dir(worktreePath), 0o755); err != nil {
+		return store.DispatchWorkspace{}, nil, fmt.Errorf("create worktree parent dir: %w", err)
+	}
+	if _, err := os.Stat(worktreePath); err == nil {
+		return store.DispatchWorkspace{}, nil, fmt.Errorf("%w: worktree path already exists: %s", store.ErrInvalidState, worktreePath)
+	} else if err != nil && !os.IsNotExist(err) {
+		return store.DispatchWorkspace{}, nil, fmt.Errorf("stat worktree path: %w", err)
+	}
+
+	if _, _, err := runGit(ctx, repoRoot, "worktree", "add", "-b", branchName, worktreePath, baseCommit); err != nil {
+		return store.DispatchWorkspace{}, nil, err
+	}
+
+	cleanup := func() {
+		_, _, _ = runGit(context.Background(), repoRoot, "worktree", "remove", "--force", worktreePath)
+		_, _, _ = runGit(context.Background(), repoRoot, "branch", "-D", branchName)
+		_ = os.RemoveAll(worktreePath)
+	}
+
+	return store.DispatchWorkspace{
+		BaseRef:         baseRef,
+		BaseCommit:      baseCommit,
+		BranchName:      branchName,
+		WorktreePath:    worktreePath,
+		WorkspaceStatus: "created",
+	}, cleanup, nil
+}
+
+func resolveRepoRoot(ctx context.Context, repoPath string) (string, error) {
+	startPath := strings.TrimSpace(repoPath)
+	if startPath == "" {
+		var err error
+		startPath, err = os.Getwd()
+		if err != nil {
+			return "", fmt.Errorf("get current working directory: %w", err)
+		}
+	}
+
+	absPath, err := filepath.Abs(startPath)
+	if err != nil {
+		return "", fmt.Errorf("resolve repo path: %w", err)
+	}
+
+	stdout, _, err := runGit(ctx, absPath, "rev-parse", "--show-toplevel")
+	if err != nil {
+		return "", protocol.InvalidInput("repo-path must point to a Git worktree", err)
+	}
+	return strings.TrimSpace(stdout), nil
+}
+
+func resolveDispatchBase(ctx context.Context, repoRoot, workspaceRoot, requestedBaseRef string, strict bool) (string, string, error) {
+	baseRef := strings.TrimSpace(requestedBaseRef)
+	if baseRef != "" {
+		baseCommit, err := resolveCommit(ctx, repoRoot, baseRef)
+		if err != nil {
+			return "", "", protocol.InvalidInput("base-ref must resolve to a commit", err)
+		}
+		return baseRef, baseCommit, nil
+	}
+
+	if strict {
+		dirty, err := repoHasUncommittedChanges(ctx, repoRoot, workspaceRoot)
+		if err != nil {
+			return "", "", err
+		}
+		if dirty {
+			return "", "", fmt.Errorf("%w: repository has uncommitted changes; specify --base-ref or clean the repo", store.ErrInvalidState)
+		}
+	}
+
+	baseCommit, err := resolveCommit(ctx, repoRoot, "HEAD")
+	if err != nil {
+		return "", "", protocol.InvalidInput("failed to resolve HEAD commit", err)
+	}
+	return "HEAD", baseCommit, nil
+}
+
+func resolveCommit(ctx context.Context, repoRoot, ref string) (string, error) {
+	stdout, _, err := runGit(ctx, repoRoot, "rev-parse", "--verify", ref+"^{commit}")
+	if err != nil {
+		return "", err
+	}
+	return strings.TrimSpace(stdout), nil
+}
+
+func repoHasUncommittedChanges(ctx context.Context, repoRoot, workspaceRoot string) (bool, error) {
+	stdout, _, err := runGit(ctx, repoRoot, "status", "--porcelain")
+	if err != nil {
+		return false, fmt.Errorf("check repository status: %w", err)
+	}
+
+	for _, line := range strings.Split(strings.TrimSpace(stdout), "\n") {
+		line = strings.TrimSpace(line)
+		if line == "" {
+			continue
+		}
+		if len(line) >= 3 {
+			path := strings.TrimSpace(line[3:])
+			if shouldIgnoreStatusPath(repoRoot, workspaceRoot, path) {
+				continue
+			}
+		}
+		return true, nil
+	}
+
+	return false, nil
+}
+
+func resolveWorkspaceRoot(repoRoot, configuredRoot string) string {
+	root := strings.TrimSpace(configuredRoot)
+	if root == "" {
+		return filepath.Join(repoRoot, ".orch", "worktrees")
+	}
+	if filepath.IsAbs(root) {
+		return root
+	}
+	return filepath.Join(repoRoot, root)
+}
+
+func ensureWorkspaceRootIgnored(repoRoot, workspaceRoot string) error {
+	relative, err := filepath.Rel(repoRoot, workspaceRoot)
+	if err != nil {
+		return fmt.Errorf("resolve workspace root exclude path: %w", err)
+	}
+	if relative == "." || strings.HasPrefix(relative, "..") {
+		return nil
+	}
+
+	pattern := filepath.ToSlash(relative)
+	if !strings.HasSuffix(pattern, "/") {
+		pattern += "/"
+	}
+
+	excludePath := filepath.Join(repoRoot, ".git", "info", "exclude")
+	if err := os.MkdirAll(filepath.Dir(excludePath), 0o755); err != nil {
+		return fmt.Errorf("create git info dir: %w", err)
+	}
+
+	content, err := os.ReadFile(excludePath)
+	if err != nil && !os.IsNotExist(err) {
+		return fmt.Errorf("read git exclude file: %w", err)
+	}
+	if strings.Contains(string(content), pattern) {
+		return nil
+	}
+
+	appendContent := pattern + "\n"
+	if len(content) > 0 && !strings.HasSuffix(string(content), "\n") {
+		appendContent = "\n" + appendContent
+	}
+	file, err := os.OpenFile(excludePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
+	if err != nil {
+		return fmt.Errorf("open git exclude file: %w", err)
+	}
+	defer file.Close()
+
+	if _, err := file.WriteString(appendContent); err != nil {
+		return fmt.Errorf("append git exclude file: %w", err)
+	}
+	return nil
+}
+
+func shouldIgnoreStatusPath(repoRoot, workspaceRoot, statusPath string) bool {
+	relative, err := filepath.Rel(repoRoot, workspaceRoot)
+	if err != nil || relative == "." || strings.HasPrefix(relative, "..") {
+		return false
+	}
+
+	relative = filepath.ToSlash(relative)
+	statusPath = filepath.ToSlash(strings.Trim(statusPath, `"`))
+	return statusPath == relative || strings.HasPrefix(statusPath, relative+"/")
+}
+
+func buildAttemptBranchName(runID, taskID string, attemptNo int) string {
+	return fmt.Sprintf(
+		"orch/%s/%s/attempt-%d",
+		sanitizeGitSegment(runID),
+		sanitizeGitSegment(taskID),
+		attemptNo,
+	)
+}
+
+func buildAttemptWorktreePath(workspaceRoot, runID, taskID string, attemptNo int) string {
+	return filepath.Join(
+		workspaceRoot,
+		sanitizePathSegment(runID),
+		sanitizePathSegment(taskID),
+		fmt.Sprintf("attempt-%d", attemptNo),
+	)
+}
+
+func sanitizeGitSegment(value string) string {
+	return sanitizeSegment(value)
+}
+
+func sanitizePathSegment(value string) string {
+	return sanitizeSegment(value)
+}
+
+func sanitizeSegment(value string) string {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return "item"
+	}
+
+	var b strings.Builder
+	lastDash := false
+	for _, r := range value {
+		if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') {
+			b.WriteRune(r)
+			lastDash = false
+			continue
+		}
+		if r == '-' || r == '_' || r == '.' {
+			if !lastDash {
+				b.WriteByte('-')
+				lastDash = true
+			}
+			continue
+		}
+		if !lastDash {
+			b.WriteByte('-')
+			lastDash = true
+		}
+	}
+
+	result := strings.Trim(b.String(), "-.")
+	if result == "" {
+		return "item"
+	}
+	if strings.HasSuffix(result, ".lock") {
+		result = strings.TrimSuffix(result, ".lock") + "-lock"
+	}
+	return result
+}
+
+func runGit(ctx context.Context, repoRoot string, args ...string) (string, string, error) {
+	cmdArgs := append([]string{"-C", repoRoot}, args...)
+	cmd := exec.CommandContext(ctx, "git", cmdArgs...)
+	output, err := cmd.CombinedOutput()
+	if err == nil {
+		return string(output), "", nil
+	}
+
+	message := strings.TrimSpace(string(output))
+	if message == "" {
+		message = err.Error()
+	}
+	return "", message, fmt.Errorf("git %s: %s", strings.Join(args, " "), message)
+}