ClaudeDo/src/ClaudeDo.Worker/External/WorktreeMcpTools.cs

using System.ComponentModel;
using System.Diagnostics;
using ClaudeDo.Data;
using ClaudeDo.Data.Git;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Git;
using ClaudeDo.Worker.Lifecycle;
using ClaudeDo.Worker.Worktrees;
using Microsoft.EntityFrameworkCore;
using ModelContextProtocol;
using ModelContextProtocol.Server;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.External;

public sealed record WorktreeInfoDto(
    string Path, string Branch, string HeadCommit, string BaseCommit,
    int Ahead, int Behind, bool IsDirty, string? MergeCommit = null);

public sealed record TaskDiffDto(
    string Content, IReadOnlyList<string> Files, bool Truncated, int TotalBytes);

public sealed record WorktreeListItemDto(
    string? TaskId, int? Number, string Path, string Branch,
    string HeadCommit, bool IsDirty, bool MergedIntoMain);

public sealed record CleanupWorktreeResult(
    bool Removed, string WorktreePath, bool BranchDeleted, int? Number = null);

[McpServerToolType]
public sealed class WorktreeMcpTools
{
    // Mirrors TaskMergeService.ProgressReportInterval / TaskWaitMcpTools.ProgressReportInterval:
    // a separate field (not shared) so shrinking one for a test can't race another's tests. Used
    // by the single-element long-runners here (GetTaskDiff, CleanupTaskWorktree) that wrap a bare
    // GitService/WorktreeMaintenanceService call with no natural i/n to report instead.
    internal static TimeSpan ProgressReportInterval = TimeSpan.FromSeconds(30);

    private readonly TaskRepository _tasks;
    private readonly GitService _git;
    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;
    private readonly WorktreeMaintenanceService _maintenance;

    public WorktreeMcpTools(
        TaskRepository tasks,
        GitService git,
        IDbContextFactory<ClaudeDoDbContext> dbFactory,
        WorktreeMaintenanceService maintenance)
    {
        _tasks = tasks;
        _git = git;
        _dbFactory = dbFactory;
        _maintenance = maintenance;
    }

    [McpServerTool, Description(
        "Get a task's git worktree state — path, branch, base/head commit, ahead/behind counts, isDirty, and the " +
        "mergeCommit its branch produced once merged. behind is 0 when the 'main' ref is unreachable, so do not " +
        "read 0 as \"up to date\" without checking. A null mergeCommit means revert_merge cannot act on this task. " +
        "Throws if the task or its worktree does not exist.")]
    public async Task<WorktreeInfoDto> GetTaskWorktree(string taskId, CancellationToken cancellationToken)
    {
        taskId = await TaskIdResolver.ResolveAsync(_tasks, taskId, cancellationToken);
        var (_, _, wt) = await LoadWorktreeContextAsync(taskId, cancellationToken);

        var headCommit = !string.IsNullOrWhiteSpace(wt.HeadCommit)
            ? wt.HeadCommit
            : await TryRunGitAsync(wt.Path, ["rev-parse", "HEAD"], cancellationToken) ?? wt.BaseCommit;

        var isDirty = Directory.Exists(wt.Path) && await _git.HasChangesAsync(wt.Path, cancellationToken);
        var ahead  = await GitRevListCountAsync(wt.Path, $"{wt.BaseCommit}..HEAD", cancellationToken);
        var behind = await GitRevListCountAsync(wt.Path, "HEAD..main", cancellationToken);

        return new WorktreeInfoDto(wt.Path, wt.BranchName, headCommit!, wt.BaseCommit, ahead, behind, isDirty, wt.MergeCommit);
    }

    [McpServerTool, Description(
        "Read what a task actually changed — the diff of its worktree against its base commit (for a worktree-less " +
        "list-handler host task, the fixed HandlerBaseCommit..HandlerHeadCommit range over the list's working dir " +
        "instead). files lists the changed paths in either mode; truncated=true means the diff was capped and " +
        "totalBytes holds its real size. Throws if the task has no worktree/review range, or the relevant " +
        "directory is missing from disk.")]
    public async Task<TaskDiffDto> GetTaskDiff(
        string taskId,
        [Description("false (default): the full unified diff, capped at 200 KB. true: a --stat summary with " +
                     "per-file insertion/deletion counts — start here when the diff may be large.")]
        bool stat = false,
        [Description("Restrict the diff to these paths (relative to the repo root), e.g. after a --stat pass or a " +
                     "conflict report already narrowed down which files matter. Omit/empty for every changed file " +
                     "— the default and the only prior behavior. Works in both stat and full-diff mode.")]
        IReadOnlyList<string>? paths = null,
        CancellationToken cancellationToken = default,
        IProgress<ProgressNotificationValue>? progress = null)
    {
        taskId = await TaskIdResolver.ResolveAsync(_tasks, taskId, cancellationToken);
        var (repoPath, baseCommit, headCommit) = await LoadDiffRangeAsync(taskId, cancellationToken);

        const int maxBytes = 200 * 1024;

        if (stat)
        {
            var diffStat = await ProgressReporter.RunAsync(
                _git.DiffStatAsync(repoPath, baseCommit, headCommit ?? "HEAD", paths, cancellationToken),
                ProgressReportInterval, progress, "computing diff stat");
            return new TaskDiffDto(diffStat, ExternalMcpService.ParseDiffStatFileNames(diffStat), false, diffStat.Length);
        }

        var diff = headCommit is null
            ? await ProgressReporter.RunAsync(
                _git.GetBranchDiffAsync(repoPath, baseCommit, paths, cancellationToken), ProgressReportInterval, progress, "computing diff")
            : await ProgressReporter.RunAsync(
                _git.GetCommitRangeDiffAsync(repoPath, baseCommit, headCommit, paths, cancellationToken), ProgressReportInterval, progress, "computing diff");
        var files = ParseDiffFileNames(diff);

        if (diff.Length <= maxBytes)
            return new TaskDiffDto(diff, files, false, diff.Length);

        return new TaskDiffDto(diff[..maxBytes], files, true, diff.Length);
    }

    // Resolves where a task's diff lives: a live worktree (repo path = worktree path, diffed
    // against HEAD) or, for a worktree-less list-handler host task, the fixed
    // HandlerBaseCommit..HandlerHeadCommit range over the list's working dir (headCommit
    // non-null signals "fixed range" to the caller instead of "diff against live HEAD").
    private async Task<(string RepoPath, string BaseCommit, string? HeadCommit)> LoadDiffRangeAsync(
        string taskId, CancellationToken ct)
    {
        using var ctx = _dbFactory.CreateDbContext();
        var task = await new TaskRepository(ctx).GetByIdAsync(taskId, ct)
            ?? throw new InvalidOperationException($"Task {taskId} not found.");
        var wt = await new WorktreeRepository(ctx).GetByTaskIdAsync(taskId, ct);

        if (wt is not null)
        {
            if (!Directory.Exists(wt.Path))
                throw new InvalidOperationException($"Worktree directory does not exist on disk: {wt.Path}");
            return (wt.Path, wt.BaseCommit, null);
        }

        if (task.HandlerBaseCommit is { Length: > 0 } handlerBase && task.HandlerHeadCommit is { Length: > 0 } handlerHead)
        {
            var list = await new ListRepository(ctx).GetByIdAsync(task.ListId, ct)
                ?? throw new InvalidOperationException("List not found.");
            if (string.IsNullOrEmpty(list.WorkingDir) || !Directory.Exists(list.WorkingDir))
                throw new InvalidOperationException("The list's working directory no longer exists.");
            return (list.WorkingDir, handlerBase, handlerHead);
        }

        throw new InvalidOperationException($"Task {taskId} has no worktree.");
    }

    [McpServerTool, Description(
        "Survey every worktree ClaudeDo tracks — use it to find leftovers to clean up. Only worktrees recorded in " +
        "the ClaudeDo database appear here, and headCommit is empty when the path is missing from disk." +
        McpToolDocs.ProgressHint)]
    public async Task<IReadOnlyList<WorktreeListItemDto>> ListWorktrees(
        CancellationToken cancellationToken = default,
        IProgress<ProgressNotificationValue>? progress = null)
    {
        var rows = await _maintenance.GetOverviewAsync(null, cancellationToken);
        // One git status + one rev-parse per row, run concurrently -- with many tracked worktrees
        // (this tool's whole point is surfacing leftovers nobody cleaned up) that can still take a
        // while, so this reports on the same elapsed-time cadence as a single long git call rather
        // than per-row (rows finish out of order under Task.WhenAll, so there's no natural i/n).
        var results = await ProgressReporter.RunAsync(
            Task.WhenAll(rows.Select(async row =>
            {
                var isDirty = row.PathExistsOnDisk && await TryGetIsDirtyAsync(row.Path, cancellationToken);
                var headCommit = row.PathExistsOnDisk
                    ? (await TryRunGitAsync(row.Path, ["rev-parse", "HEAD"], cancellationToken) ?? "")
                    : "";
                return new WorktreeListItemDto(
                    row.TaskId, row.TaskNumber, row.Path, row.BranchName, headCommit,
                    isDirty, row.State == WorktreeState.Merged);
            })),
            ProgressReportInterval, progress, "surveying worktrees");
        return results;
    }

    [McpServerTool, Description(
        "Remove a task's worktree directory and delete its git branch. branchDeleted reports whether the branch " +
        "went too." + McpToolDocs.NotWhileRunning)]
    public async Task<CleanupWorktreeResult> CleanupTaskWorktree(
        string taskId,
        [Description("false (default): refuse a worktree with uncommitted changes. true: remove it anyway, losing " +
                     "those changes.")]
        bool force = false,
        CancellationToken cancellationToken = default,
        IProgress<ProgressNotificationValue>? progress = null)
    {
        taskId = await TaskIdResolver.ResolveAsync(_tasks, taskId, cancellationToken);

        using var ctx = _dbFactory.CreateDbContext();
        var task = await new TaskRepository(ctx).GetByIdAsync(taskId, cancellationToken)
            ?? throw new InvalidOperationException($"Task {taskId} not found.");
        var wt = await new WorktreeRepository(ctx).GetByTaskIdAsync(taskId, cancellationToken)
            ?? throw new InvalidOperationException($"Task {taskId} has no worktree.");

        if (task.Status == TaskStatus.Running)
            throw new InvalidOperationException("Cannot remove worktree of a running task.");

        if (!force && Directory.Exists(wt.Path))
        {
            var isDirty = await ProgressReporter.RunAsync(
                _git.HasChangesAsync(wt.Path, cancellationToken), ProgressReportInterval, progress, "checking worktree for changes");
            if (isDirty)
                throw new InvalidOperationException(
                    "Worktree has uncommitted changes. Use force=true to remove anyway (changes will be lost).");
        }

        var path = wt.Path;
        var result = await ProgressReporter.RunAsync(
            _maintenance.ForceRemoveAsync(taskId, cancellationToken), ProgressReportInterval, progress, "removing worktree");
        return new CleanupWorktreeResult(result.Removed, path, result.BranchDeleted, task.Number);
    }

    private async Task<(TaskEntity Task, ListEntity List, WorktreeEntity Wt)> LoadWorktreeContextAsync(
        string taskId, CancellationToken ct)
    {
        using var ctx = _dbFactory.CreateDbContext();
        var task = await new TaskRepository(ctx).GetByIdAsync(taskId, ct)
            ?? throw new InvalidOperationException($"Task {taskId} not found.");
        var list = await new ListRepository(ctx).GetByIdAsync(task.ListId, ct)
            ?? throw new InvalidOperationException("List not found.");
        var wt = await new WorktreeRepository(ctx).GetByTaskIdAsync(taskId, ct)
            ?? throw new InvalidOperationException($"Task {taskId} has no worktree.");
        return (task, list, wt);
    }

    private async Task<bool> TryGetIsDirtyAsync(string path, CancellationToken ct)
    {
        try { return await _git.HasChangesAsync(path, ct); }
        catch { return false; }
    }

    // Minimal git runner for operations not covered by GitService (rev-list --count, rev-parse from worktree).
    private static async Task<string?> TryRunGitAsync(string dir, string[] args, CancellationToken ct)
    {
        try
        {
            var psi = new ProcessStartInfo("git")
            {
                UseShellExecute = false,
                RedirectStandardOutput = true,
                CreateNoWindow = true,
            };
            psi.ArgumentList.Add("-C");
            psi.ArgumentList.Add(dir);
            foreach (var a in args) psi.ArgumentList.Add(a);
            using var proc = Process.Start(psi)!;
            await using var _ = ct.Register(() => { try { proc.Kill(entireProcessTree: true); } catch { } });
            var stdout = await proc.StandardOutput.ReadToEndAsync();
            await proc.WaitForExitAsync(CancellationToken.None);
            ct.ThrowIfCancellationRequested();
            return proc.ExitCode == 0 ? stdout.Trim() : null;
        }
        catch (OperationCanceledException) { throw; }
        catch { return null; }
    }

    internal static async Task<int> GitRevListCountAsync(string dir, string range, CancellationToken ct)
    {
        var result = await TryRunGitAsync(dir, ["rev-list", "--count", range], ct);
        return int.TryParse(result, out var n) ? n : 0;
    }

    private static IReadOnlyList<string> ParseDiffFileNames(string diff)
    {
        var files = new List<string>();
        foreach (var line in diff.Split('\n'))
        {
            var s = line.TrimEnd('\r');
            if (s.StartsWith("+++ b/", StringComparison.Ordinal))
                files.Add(s[6..]);
        }
        return files;
    }
}