ai-workflow-skill/internal/cli/orch/worktree.go

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