ClaudeDo/src/ClaudeDo.Worker/Runner/WorktreeManager.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Git;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Config;
using ClaudeDo.Worker.Git;
using ClaudeDo.Worker.Hub;
using Microsoft.EntityFrameworkCore;

namespace ClaudeDo.Worker.Runner;

public sealed record WorktreeContext(string WorktreePath, string BranchName, string BaseCommit);

public sealed class WorktreeManager
{
    // Phase token for the OperationProgress broadcast while `git worktree add` (+ the
    // self-heal retry) is in flight -- stable identifier, localized by the UI via the
    // pre-existing "ops.worker.creatingWorktree" key. opKey is the task id.
    public const string PhaseCreatingWorktree = "creating_worktree";

    private readonly GitService _git;
    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;
    private readonly WorkerConfig _cfg;
    private readonly ILogger<WorktreeManager> _logger;
    private readonly HubBroadcaster? _broadcaster;

    public WorktreeManager(
        GitService git, IDbContextFactory<ClaudeDoDbContext> dbFactory, WorkerConfig cfg,
        ILogger<WorktreeManager> logger, HubBroadcaster? broadcaster = null)
    {
        _git = git;
        _dbFactory = dbFactory;
        _cfg = cfg;
        _logger = logger;
        _broadcaster = broadcaster;
    }

    private Task BroadcastCreating(string taskId) =>
        _broadcaster?.OperationProgress(taskId, PhaseCreatingWorktree, 0, 0) ?? Task.CompletedTask;

    public async Task<WorktreeContext> CreateAsync(TaskEntity task, ListEntity list, CancellationToken ct)
    {
        var workingDir = list.WorkingDir
            ?? throw new InvalidOperationException("list.WorkingDir is null");

        if (!await _git.IsGitRepoAsync(workingDir, ct))
            throw new InvalidOperationException($"working_dir is not a git repository: {workingDir}");

        // A fresh `git init` repo has no HEAD commit to base a worktree on — bootstrap one.
        if (await _git.EnsureHeadCommitAsync(workingDir, ct))
            _logger.LogInformation("Repository at {Dir} had no commits; created an empty bootstrap commit", workingDir);

        var baseCommit = await ResolveBaseCommitAsync(task, workingDir, ct);
        // Use the full task id (dashes stripped) in the branch name so
        // two GUIDs sharing an 8-char prefix cannot collide on the same branch.
        var idForBranch = task.Id.Replace("-", "");
        var branchName = $"claudedo/{idForBranch}";
        var slug = CommitMessageBuilder.ToSlug(list.Name);

        string strategy;
        string? centralRoot;
        using (var settingsCtx = _dbFactory.CreateDbContext())
        {
            var settings = await new AppSettingsRepository(settingsCtx).GetAsync(ct);
            strategy = settings.WorktreeStrategy;
            centralRoot = !string.IsNullOrWhiteSpace(settings.CentralWorktreeRoot)
                ? settings.CentralWorktreeRoot
                : _cfg.CentralWorktreeRoot;
        }

        var worktreePath = strategy.Equals("central", StringComparison.OrdinalIgnoreCase)
            ? Path.Combine(centralRoot ?? _cfg.CentralWorktreeRoot, slug, task.Id)
            : Path.Combine(WorktreeRootResolver.ResolveSiblingRoot(workingDir), ".claudedo-worktrees", slug, task.Id);

        worktreePath = Path.GetFullPath(worktreePath);
        WorktreeRootResolver.EnsureOutsideWorkingDir(worktreePath, workingDir);

        // Ensure parent directory exists.
        Directory.CreateDirectory(Path.GetDirectoryName(worktreePath)!);

        // Create the worktree. If a stale branch from a previous run remains
        // (e.g. after force-remove), delete it and retry once.
        await BroadcastCreating(task.Id);
        try
        {
            await _git.WorktreeAddAsync(workingDir, branchName, worktreePath, baseCommit, ct);
        }
        catch (InvalidOperationException ex) when (ex.Message.Contains("already exists", StringComparison.OrdinalIgnoreCase))
        {
            _logger.LogWarning("Branch {Branch} already exists; cleaning phantom worktrees and retrying", branchName);
            // Self-heal is the slow section (list -> remove -> prune -> branch delete -> retry);
            // report it as its own section rather than staying silent through the retry.
            await BroadcastCreating(task.Id);

            // Find and forcefully remove any existing worktree registered against this branch.
            List<string> stalePaths;
            try { stalePaths = await _git.ListWorktreePathsForBranchAsync(workingDir, branchName, ct); }
            catch (Exception listEx)
            {
                _logger.LogWarning(listEx, "git worktree list failed during self-heal");
                stalePaths = new();
            }
            foreach (var stalePath in stalePaths)
            {
                try { await _git.WorktreeRemoveAsync(workingDir, stalePath, force: true, ct); }
                catch (Exception wrEx) { _logger.LogWarning(wrEx, "Failed to remove stale worktree at {Path}", stalePath); }
            }

            try { await _git.WorktreePruneAsync(workingDir, ct); }
            catch (Exception pruneEx) { _logger.LogWarning(pruneEx, "git worktree prune failed during self-heal"); }
            try { await _git.BranchDeleteAsync(workingDir, branchName, force: true, ct); }
            catch (Exception delEx) { _logger.LogWarning(delEx, "git branch -D failed during self-heal"); }

            await _git.WorktreeAddAsync(workingDir, branchName, worktreePath, baseCommit, ct);
        }

        // Insert worktrees row AFTER git succeeds — if git throws, no row is created.
        // If the insert itself fails, the worktree is already on disk with nothing
        // tracking it: remove it (best-effort, non-cancellable) before rethrowing.
        try
        {
            using var context = _dbFactory.CreateDbContext();
            var wtRepo = new WorktreeRepository(context);
            // Drop any stale row from a prior run (force-remove may have left the DB side behind).
            await wtRepo.DeleteAsync(task.Id, ct);
            await wtRepo.AddAsync(new WorktreeEntity
            {
                TaskId = task.Id,
                Path = worktreePath,
                BranchName = branchName,
                BaseCommit = baseCommit,
                HeadCommit = null,
                DiffStat = null,
                State = WorktreeState.Active,
                CreatedAt = DateTime.UtcNow,
            }, ct);
        }
        catch (Exception dbEx)
        {
            _logger.LogError(dbEx,
                "Failed to record worktree row for task {TaskId}; removing orphaned worktree at {Path}",
                task.Id, worktreePath);
            try { await _git.WorktreeRemoveAsync(workingDir, worktreePath, force: true, CancellationToken.None); }
            catch (Exception rmEx) { _logger.LogWarning(rmEx, "Failed to remove orphaned worktree at {Path}", worktreePath); }
            try { await _git.BranchDeleteAsync(workingDir, branchName, force: true, CancellationToken.None); }
            catch (Exception delEx) { _logger.LogWarning(delEx, "Failed to delete orphaned branch {Branch}", branchName); }
            throw;
        }

        _logger.LogInformation("Created worktree for task {TaskId} at {Path} (branch {Branch}, base {Base})",
            task.Id, worktreePath, branchName, baseCommit);

        return new WorktreeContext(worktreePath, branchName, baseCommit);
    }

    /// <returns>true if a commit was made; false if no changes.</returns>
    public async Task<bool> CommitIfChangedAsync(WorktreeContext ctx, TaskEntity task, ListEntity list, CancellationToken ct)
    {
        if (!await _git.HasChangesAsync(ctx.WorktreePath, ct))
        {
            _logger.LogInformation("No changes in worktree for task {TaskId}, skipping commit", task.Id);
            return false;
        }

        await _git.AddAllAsync(ctx.WorktreePath, ct);

        var message = CommitMessageBuilder.Build(
            task.CommitType, list.Name, task.Title, task.Description, task.Id);

        await _git.CommitAsync(ctx.WorktreePath, message, ct);

        var head = await _git.RevParseHeadAsync(ctx.WorktreePath, ct);
        var diffStat = await _git.DiffStatAsync(ctx.WorktreePath, ctx.BaseCommit, head, ct: ct);

        using var context = _dbFactory.CreateDbContext();
        var wtRepo = new WorktreeRepository(context);
        await wtRepo.UpdateHeadAsync(task.Id, head, diffStat, ct);

        _logger.LogInformation("Committed changes for task {TaskId}: {Head}", task.Id, head);
        return true;
    }

    public async Task DiscardAsync(WorktreeEntity wt, string workingDir, CancellationToken ct)
    {
        try
        {
            await _git.WorktreeRemoveAsync(workingDir, wt.Path, force: true, ct);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "git worktree remove failed for {Path}", wt.Path);
            throw;
        }

        try
        {
            await _git.BranchDeleteAsync(workingDir, wt.BranchName, force: true, ct);
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "git branch -D {Branch} failed after worktree removal; continuing", wt.BranchName);
        }

        using var context = _dbFactory.CreateDbContext();
        var wtRepo = new WorktreeRepository(context);
        await wtRepo.SetStateAsync(wt.TaskId, WorktreeState.Discarded, ct);

        _logger.LogInformation("Discarded worktree for task {TaskId} (branch {Branch})", wt.TaskId, wt.BranchName);
    }

    // Improvement children (parent is a non-planning task with its own worktree) branch
    // from the parent's recorded HEAD so they build on the parent's not-yet-merged work.
    // Planning children and standalone tasks base off the list's current HEAD.
    private async Task<string> ResolveBaseCommitAsync(TaskEntity task, string workingDir, CancellationToken ct)
    {
        if (task.ParentTaskId is not null)
        {
            using var ctx = _dbFactory.CreateDbContext();
            var parent = await ctx.Tasks.AsNoTracking()
                .FirstOrDefaultAsync(t => t.Id == task.ParentTaskId, ct);
            if (parent is not null && parent.PlanningPhase == PlanningPhase.None)
            {
                var parentWt = await new WorktreeRepository(ctx).GetByTaskIdAsync(task.ParentTaskId, ct);
                var parentHead = parentWt?.HeadCommit ?? parentWt?.BaseCommit;
                if (parentHead is not null)
                    return parentHead;
            }
        }
        return await _git.RevParseHeadAsync(workingDir, ct);
    }
}