ClaudeDo/src/ClaudeDo.Worker/Worktrees/WorktreeMaintenanceService.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Git;
using ClaudeDo.Data.Models;
using Microsoft.EntityFrameworkCore;

namespace ClaudeDo.Worker.Worktrees;

public sealed class WorktreeMaintenanceService
{
    public sealed record CleanupResult(int Removed, IReadOnlyList<string> RemovedTaskIds);
    public sealed record ResetResult(int Removed, int TasksAffected, bool Blocked, int RunningTasks, IReadOnlyList<string> RemovedTaskIds);
    public sealed record ForceRemoveResult(bool Removed, string? Reason, bool BranchDeleted);

    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;
    private readonly GitService _git;
    private readonly ILogger<WorktreeMaintenanceService> _logger;

    public WorktreeMaintenanceService(
        IDbContextFactory<ClaudeDoDbContext> dbFactory,
        GitService git,
        ILogger<WorktreeMaintenanceService> logger)
    {
        _dbFactory = dbFactory;
        _git = git;
        _logger = logger;
    }

    public async Task<CleanupResult> CleanupFinishedAsync(string? listId = null, CancellationToken ct = default)
    {
        using var context = _dbFactory.CreateDbContext();
        var query = from w in context.Worktrees
                    join t in context.Tasks on w.TaskId equals t.Id
                    join l in context.Lists on t.ListId equals l.Id
                    where w.State == WorktreeState.Merged || w.State == WorktreeState.Discarded
                    select new { Row = new WorktreeRow(w.TaskId, w.Path, w.BranchName, l.WorkingDir), ListId = t.ListId };

        if (!string.IsNullOrEmpty(listId))
            query = query.Where(x => x.ListId == listId);

        var rows = await query.AsNoTracking().Select(x => x.Row).ToListAsync(ct);

        int removed = 0;
        var removedTaskIds = new List<string>();
        foreach (var row in rows)
        {
            var (rowRemoved, _) = await TryRemoveAsync(row, force: false, ct);
            if (rowRemoved)
            {
                removed++;
                removedTaskIds.Add(row.TaskId);
            }
        }
        return new CleanupResult(removed, removedTaskIds);
    }

    public async Task<ResetResult> ResetAllAsync(CancellationToken ct = default)
    {
        using var context = _dbFactory.CreateDbContext();
        var running = await context.Tasks.AsNoTracking()
            .CountAsync(t => t.Status == ClaudeDo.Data.Models.TaskStatus.Running, ct);
        if (running > 0)
            return new ResetResult(0, 0, Blocked: true, RunningTasks: running, Array.Empty<string>());

        var rows = await (from w in context.Worktrees
                          join t in context.Tasks on w.TaskId equals t.Id
                          join l in context.Lists on t.ListId equals l.Id
                          select new WorktreeRow(w.TaskId, w.Path, w.BranchName, l.WorkingDir))
                         .AsNoTracking()
                         .ToListAsync(ct);

        int removed = 0;
        var removedTaskIds = new List<string>();
        foreach (var row in rows)
        {
            var (rowRemoved, _) = await TryRemoveAsync(row, force: true, ct);
            if (rowRemoved)
            {
                removed++;
                removedTaskIds.Add(row.TaskId);
            }
        }
        return new ResetResult(removed, rows.Count, Blocked: false, RunningTasks: 0, removedTaskIds);
    }

    public async Task<IReadOnlyList<WorktreeOverviewRow>> GetOverviewAsync(
        string? listId, CancellationToken ct = default)
    {
        using var context = _dbFactory.CreateDbContext();
        var query = from w in context.Worktrees
                    join t in context.Tasks on w.TaskId equals t.Id
                    join l in context.Lists on t.ListId equals l.Id
                    select new
                    {
                        w.TaskId, t.Number, t.Title, t.Status, ListId = l.Id, ListName = l.Name,
                        w.Path, w.BranchName, w.BaseCommit, w.State, w.DiffStat, w.CreatedAt,
                    };

        if (!string.IsNullOrEmpty(listId))
            query = query.Where(x => x.ListId == listId);

        var rows = await query.AsNoTracking().ToListAsync(ct);

        return rows.Select(x => new WorktreeOverviewRow(
            x.TaskId, x.Number, x.Title, x.Status, x.ListId, x.ListName,
            x.Path, x.BranchName, x.BaseCommit ?? "", x.State, x.DiffStat, x.CreatedAt,
            PathExistsOnDisk: !string.IsNullOrWhiteSpace(x.Path) && Directory.Exists(x.Path))).ToList();
    }

    public async Task<ForceRemoveResult> ForceRemoveAsync(string taskId, CancellationToken ct = default)
    {
        using var context = _dbFactory.CreateDbContext();

        var row = await (from w in context.Worktrees
                         join t in context.Tasks on w.TaskId equals t.Id
                         join l in context.Lists on t.ListId equals l.Id
                         where w.TaskId == taskId
                         select new { Row = new WorktreeRow(w.TaskId, w.Path, w.BranchName, l.WorkingDir),
                                      Status = t.Status })
                        .AsNoTracking()
                        .FirstOrDefaultAsync(ct);

        if (row is null)
            return new ForceRemoveResult(false, "worktree not found", false);

        if (row.Status == ClaudeDo.Data.Models.TaskStatus.Running)
            return new ForceRemoveResult(false, "task is currently running", false);

        var (ok, branchDeleted) = await TryRemoveAsync(row.Row, force: true, ct);
        return new ForceRemoveResult(ok, ok ? null : "remove failed", branchDeleted);
    }

    private async Task<(bool Removed, bool BranchDeleted)> TryRemoveAsync(WorktreeRow row, bool force, CancellationToken ct)
    {
        var repoDirExists = !string.IsNullOrWhiteSpace(row.WorkingDir) && Directory.Exists(row.WorkingDir);
        bool dirRemoved;

        if (repoDirExists)
        {
            dirRemoved = true;
            try
            {
                await _git.WorktreeRemoveAsync(row.WorkingDir!, row.Path, force, ct);
            }
            catch (Exception ex)
            {
                // "is not a working tree" means git already forgot this worktree (pruned, or
                // removed by hand). Deleting the directory is then the whole job, not a fallback
                // from a real failure -- warning about it floods the footer log strip when a
                // batch of stale rows is cleaned up.
                if (IsAlreadyUnregistered(ex))
                    _logger.LogDebug(ex, "git no longer tracks {Path}; deleting the directory", row.Path);
                else
                    _logger.LogWarning(ex,
                        "git worktree remove failed for {Path}; falling back to directory delete", row.Path);
                try { if (Directory.Exists(row.Path)) Directory.Delete(row.Path, recursive: true); }
                catch (Exception delEx)
                {
                    _logger.LogError(delEx, "Directory.Delete fallback also failed for {Path}", row.Path);
                    dirRemoved = false;
                }
            }
            if (Directory.Exists(row.Path)) dirRemoved = false;
        }
        else
        {
            dirRemoved = true;
            try { if (Directory.Exists(row.Path)) Directory.Delete(row.Path, recursive: true); }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Directory.Delete failed for {Path}", row.Path);
                dirRemoved = false;
            }
        }

        // Drop the DB row only when the on-disk worktree is gone; otherwise we'd silently
        // strand a directory while reporting success.
        if (!dirRemoved) return (false, false);

        // Branch cleanup: otherwise rerunning the task hits "branch already exists".
        // Prune first so git no longer thinks the branch is checked out by a phantom worktree.
        bool branchDeleted = false;
        if (repoDirExists)
        {
            try { await _git.WorktreePruneAsync(row.WorkingDir!, ct); }
            catch (Exception ex) { _logger.LogWarning(ex, "git worktree prune failed for {Repo}", row.WorkingDir); }

            if (!string.IsNullOrWhiteSpace(row.BranchName))
            {
                try
                {
                    await _git.BranchDeleteAsync(row.WorkingDir!, row.BranchName, force: true, ct);
                    branchDeleted = true;
                }
                catch (Exception ex)
                {
                    // Already deleted -- the normal case once the branch has been merged.
                    if (IsBranchAlreadyGone(ex))
                        _logger.LogDebug(ex, "branch {Branch} was already deleted", row.BranchName);
                    else
                        _logger.LogWarning(ex, "Failed to delete branch {Branch} for worktree {Path}",
                            row.BranchName, row.Path);
                }
            }
        }

        using var context = _dbFactory.CreateDbContext();
        await context.Worktrees.Where(w => w.TaskId == row.TaskId).ExecuteDeleteAsync(ct);
        return (true, branchDeleted);
    }

    // Both are "already done" outcomes, not failures — see the call sites.
    private static bool IsAlreadyUnregistered(Exception ex) =>
        ex.Message.Contains("is not a working tree", StringComparison.OrdinalIgnoreCase);

    private static bool IsBranchAlreadyGone(Exception ex) =>
        ex.Message.Contains("' not found", StringComparison.Ordinal);

    private sealed record WorktreeRow(string TaskId, string Path, string BranchName, string? WorkingDir);
}