ClaudeDo/src/ClaudeDo.Worker/External/BatchMcpTools.cs

using System.ComponentModel;
using System.Text.Json;
using ClaudeDo.Worker.Git;
using ClaudeDo.Worker.Lifecycle;
using ModelContextProtocol;
using ModelContextProtocol.Server;

namespace ClaudeDo.Worker.External;

public sealed record BatchAddTaskInput(
    string Title,
    [property: Description("Task description/instructions for the agent.")] string? Description = null,
    [property: Description("Model override: haiku|sonnet|opus. Blank inherits the list/global default.")] string? Model = null);
public sealed record BatchSetMyDayInput(
    string TaskId,
    [property: Description("true to add the task to My Day, false to remove it.")] bool IsMyDay,
    [property: Description("Position within My Day; omit to append at the end.")] int? SortOrder = null);

// task is populated when found and includeDescription=false (the default, lean reference);
// taskFull is populated when found and includeDescription=true (full task incl.
// Description/Result, possibly narrowed by `fields` and/or truncated). Both are null when
// found=false.
public sealed record BatchGetTaskResult(string Id, bool Found, TaskRefDto? Task, BatchTaskDetailDto? TaskFull, string? Error);

// Batch-only counterpart to TaskDto (deliberately not TaskDto itself — get_task's own shape stays
// untouched). Every field defaults to null; BatchGetTasks only populates the ones the caller asked
// for via `fields` (or all of them when `fields` is omitted). Description/Result are cut to
// descriptionMaxChars with the *Truncated/*FullLength pair telling the caller it happened and how
// much text there really is. RoadblockText is the tail of Result after TaskRunner's roadblock
// marker (see ComposeReviewResult in Runner/TaskRunner.cs) — the few bullet lines a roadblocked
// run reported, without pulling in the rest of Result/Description to get them.
public sealed record BatchTaskDetailDto(
    string Id,
    int Number,
    string? ListId = null,
    string? Title = null,
    string? Description = null,
    string? Status = null,
    string? Result = null,
    string? CreatedBy = null,
    DateTime? CreatedAt = null,
    DateTime? StartedAt = null,
    DateTime? FinishedAt = null,
    bool? IsMyDay = null,
    int? SortOrder = null,
    int? RoadblockCount = null,
    string? RoadblockText = null,
    string? FailureReason = null,
    bool DescriptionTruncated = false,
    int? DescriptionFullLength = null,
    bool ResultTruncated = false,
    int? ResultFullLength = null);
public sealed record BatchAddTaskResult(
    int Index, string Title, bool Ok, TaskRefDto? Task,
    IReadOnlyList<PossibleDuplicateDto>? PossibleDuplicates, string? Error);
// BaseDirty mirrors TaskRefDto.BaseDirty: populated only for BatchUpdateTaskStatus items that
// just transitioned to Queued against a list whose working dir has uncommitted changes.
public sealed record BatchTaskResult(string TaskId, bool Ok, string? Error, DirtyBaseWarning? BaseDirty = null, int? Number = null);
public sealed record BatchCancelResult(string TaskId, bool Ok, bool Cancelled, string? Error, int? Number = null);
public sealed record BatchCleanupResult(string TaskId, bool Ok, bool Removed, bool BranchDeleted, string? Error, int? Number = null);

/// <summary>
/// Batch variants of the single-entity tools on <see cref="ExternalMcpService"/>.
/// Every tool is best-effort: each item is attempted independently and reported in a
/// per-item result array (a failing item never aborts the rest). Operations run
/// sequentially — the underlying repositories share one scoped DbContext, which is not
/// thread-safe. Merge/review stay single-task (conflicts need the interactive resolver).
/// </summary>
[McpServerToolType]
public sealed class BatchMcpTools
{
    private const int MaxBatchSize = 100;

    // Roughly 6k tokens of headroom under a typical 25k-token tool-result limit — the incident
    // this guards against was a single 51k-char line from 8 tasks' full Description/Result.
    private const int MaxResponseChars = 25_000;

    private const string RoadblockMarker = "Roadblocks reported during the run:";

    private static readonly string[] KnownDetailFields =
    {
        "listId", "title", "description", "status", "result", "createdBy", "createdAt",
        "startedAt", "finishedAt", "isMyDay", "sortOrder", "roadblockCount", "roadblockText",
        "failureReason",
    };

    private readonly ExternalMcpService _svc;
    private readonly WorktreeMcpTools _worktree;

    public BatchMcpTools(ExternalMcpService svc, WorktreeMcpTools worktree)
    {
        _svc = svc;
        _worktree = worktree;
    }

    [McpServerTool, Description(
        "Fetch a snapshot of many tasks in one call — use for an overview or polling a fan-out instead of " +
        "calling get_task per id. A missing id comes back as found=false, not an error; error is only set " +
        "for an unexpected failure. A Failed task's task/taskFull carries failureReason (see get_task). " +
        "When includeDescription=true, Description/Result are cut to " +
        "descriptionMaxChars (default 1500 — was unlimited); check *Truncated/*FullLength before assuming " +
        "you got the whole text. Use `fields` to fetch only what you need (e.g. just roadblockText) instead " +
        "of raising the cap. The whole response is still capped — if it's too big even so, lower " +
        "descriptionMaxChars, narrow `fields`, or split taskIds into a smaller batch." +
        McpToolDocs.MaxBatch + McpToolDocs.BatchProgressHint)]
    public async Task<IReadOnlyList<BatchGetTaskResult>> BatchGetTasks(
        string[] taskIds,
        [Description("If true, return the full task (incl. Description/Result) in `taskFull`; if false " +
            "(default), return a lean reference in `task`.")] bool includeDescription = false,
        [Description("Max chars kept for Description and for Result (each independently) when " +
            "includeDescription=true; longer text is cut and flagged via *Truncated/*FullLength. Ignored " +
            "when includeDescription=false.")] int descriptionMaxChars = 1500,
        [Description("Restrict `taskFull` to these field names (e.g. ['title','status','roadblockText']) " +
            "instead of returning all of them; unknown names are rejected. Ignored when " +
            "includeDescription=false.")] string[]? fields = null,
        CancellationToken cancellationToken = default,
        IProgress<ProgressNotificationValue>? progress = null)
    {
        EnsureWithinCap(taskIds, nameof(taskIds));
        if (descriptionMaxChars < 0)
            throw new InvalidOperationException($"{nameof(descriptionMaxChars)} must be >= 0.");
        ValidateFields(fields);

        var results = new List<BatchGetTaskResult>(taskIds.Length);
        for (var i = 0; i < taskIds.Length; i++)
        {
            var id = taskIds[i];
            try
            {
                if (includeDescription)
                {
                    var task = await _svc.GetTask(id, cancellationToken);
                    results.Add(new BatchGetTaskResult(id, true, null, BuildDetail(task, fields, descriptionMaxChars), null));
                }
                else
                {
                    var taskRef = await _svc.GetTaskRefAsync(id, cancellationToken);
                    results.Add(new BatchGetTaskResult(id, true, taskRef, null, null));
                }
            }
            catch (OperationCanceledException) { throw; }
            catch (InvalidOperationException)
            {
                results.Add(new BatchGetTaskResult(id, false, null, null, null));
            }
            catch (Exception ex)
            {
                results.Add(new BatchGetTaskResult(id, false, null, null, ex.Message));
            }
            ProgressReporter.ReportItem(progress, "Fetching tasks", i + 1, taskIds.Length);
        }

        var responseLength = JsonSerializer.Serialize(results).Length;
        if (responseLength > MaxResponseChars)
            throw new InvalidOperationException(
                $"Response too large: {responseLength} chars (max {MaxResponseChars}). Lower " +
                $"{nameof(descriptionMaxChars)}, narrow `fields`, or split {nameof(taskIds)} into a smaller batch.");

        return results;
    }

    private static void ValidateFields(string[]? fields)
    {
        if (fields is null) return;
        var unknown = fields.Where(f => !KnownDetailFields.Contains(f, StringComparer.OrdinalIgnoreCase)).ToArray();
        if (unknown.Length > 0)
            throw new InvalidOperationException(
                $"Unknown field(s) in `fields`: {string.Join(", ", unknown)}. Known fields: " +
                string.Join(", ", KnownDetailFields));
    }

    private static BatchTaskDetailDto BuildDetail(TaskDto t, string[]? fields, int descriptionMaxChars)
    {
        bool Want(string name) => fields is null || fields.Contains(name, StringComparer.OrdinalIgnoreCase);

        var (description, descTruncated, descFullLength) =
            Want("description") ? Truncate(t.Description, descriptionMaxChars) : (null, false, null);
        var (result, resultTruncated, resultFullLength) =
            Want("result") ? Truncate(t.Result, descriptionMaxChars) : (null, false, null);

        return new BatchTaskDetailDto(
            Id: t.Id,
            Number: t.Number,
            ListId: Want("listId") ? t.ListId : null,
            Title: Want("title") ? t.Title : null,
            Description: description,
            Status: Want("status") ? t.Status : null,
            Result: result,
            CreatedBy: Want("createdBy") ? t.CreatedBy : null,
            CreatedAt: Want("createdAt") ? t.CreatedAt : null,
            StartedAt: Want("startedAt") ? t.StartedAt : null,
            FinishedAt: Want("finishedAt") ? t.FinishedAt : null,
            IsMyDay: Want("isMyDay") ? t.IsMyDay : null,
            SortOrder: Want("sortOrder") ? t.SortOrder : null,
            RoadblockCount: Want("roadblockCount") ? t.RoadblockCount : null,
            RoadblockText: Want("roadblockText") ? ExtractRoadblockText(t.Result) : null,
            FailureReason: Want("failureReason") ? t.FailureReason : null,
            DescriptionTruncated: descTruncated,
            DescriptionFullLength: descFullLength,
            ResultTruncated: resultTruncated,
            ResultFullLength: resultFullLength);
    }

    private static (string? Text, bool Truncated, int? FullLength) Truncate(string? text, int maxChars)
    {
        if (text is null || text.Length <= maxChars) return (text, false, null);
        return (text[..maxChars], true, text.Length);
    }

    private static string? ExtractRoadblockText(string? result)
    {
        if (string.IsNullOrEmpty(result)) return null;
        var idx = result.IndexOf(RoadblockMarker, StringComparison.Ordinal);
        return idx < 0 ? null : result[(idx + RoadblockMarker.Length)..].Trim();
    }

    [McpServerTool, Description(
        "Create many tasks in one list at once — use instead of repeated add_task calls when seeding a list. " +
        "Every item is still created even if it looks like a duplicate; possibleDuplicates is a non-blocking " +
        "heads-up (up to 3 similar open tasks in the list) worth mentioning to the caller, not an error." +
        McpToolDocs.LeanTaskRef + McpToolDocs.MaxBatch + McpToolDocs.BatchProgressHint)]
    public async Task<IReadOnlyList<BatchAddTaskResult>> BatchAddTasks(
        string listId,
        BatchAddTaskInput[] tasks,
        string? createdBy = null,
        [Description("If true, enqueue every created task immediately instead of leaving it Idle.")] bool queueImmediately = false,
        CancellationToken cancellationToken = default,
        IProgress<ProgressNotificationValue>? progress = null)
    {
        EnsureWithinCap(tasks, nameof(tasks));

        var results = new List<BatchAddTaskResult>(tasks.Length);
        for (var i = 0; i < tasks.Length; i++)
        {
            var item = tasks[i];
            try
            {
                var created = await _svc.AddTask(
                    listId, item.Title, item.Description, createdBy,
                    queueImmediately, item.Model, cancellationToken: cancellationToken);
                results.Add(new BatchAddTaskResult(i, item.Title, true, created.Task, created.PossibleDuplicates, null));
            }
            catch (OperationCanceledException) { throw; }
            catch (Exception ex)
            {
                results.Add(new BatchAddTaskResult(i, item.Title, false, null, null, ex.Message));
            }
            ProgressReporter.ReportItem(progress, "Adding tasks", i + 1, tasks.Length);
        }
        return results;
    }

    [McpServerTool, Description(
        "Set the status of many tasks at once — use for bulk queue/cancel/done actions instead of calling " +
        "update_task_status per task. 'Done' is refused per-item for a task with an active worktree. " +
        "baseDirty on a 'Queued' item means that task's list has uncommitted changes in its working dir right " +
        "now — a new worktree forks from the last commit and won't include them; non-blocking, but worth " +
        "checking before assuming a fresh worktree starts from what's on disk." +
        McpToolDocs.MaxBatch + McpToolDocs.BatchProgressHint)]
    public async Task<IReadOnlyList<BatchTaskResult>> BatchUpdateTaskStatus(
        string[] taskIds,
        [Description("One of 'Idle', 'Queued', 'Cancelled', or 'Done'.")] string status,
        CancellationToken cancellationToken,
        IProgress<ProgressNotificationValue>? progress = null)
    {
        EnsureWithinCap(taskIds, nameof(taskIds));

        var results = new List<BatchTaskResult>(taskIds.Length);
        for (var i = 0; i < taskIds.Length; i++)
        {
            var id = taskIds[i];
            try
            {
                var task = await _svc.UpdateTaskStatus(id, status, cancellationToken);
                results.Add(new BatchTaskResult(id, true, null, task.BaseDirty, task.Number));
            }
            catch (OperationCanceledException) { throw; }
            catch (Exception ex)
            {
                results.Add(new BatchTaskResult(id, false, ex.Message));
            }
            ProgressReporter.ReportItem(progress, "Updating task status", i + 1, taskIds.Length);
        }
        return results;
    }

    [McpServerTool, Description(
        "Cancel many running tasks at once — use to bulk-stop tasks instead of calling cancel_task per id. " +
        "ok=true with cancelled=false just means the task wasn't running." +
        McpToolDocs.MaxBatch + McpToolDocs.BatchProgressHint)]
    public async Task<IReadOnlyList<BatchCancelResult>> BatchCancelTasks(
        string[] taskIds, CancellationToken cancellationToken, IProgress<ProgressNotificationValue>? progress = null)
    {
        EnsureWithinCap(taskIds, nameof(taskIds));

        var results = new List<BatchCancelResult>(taskIds.Length);
        for (var i = 0; i < taskIds.Length; i++)
        {
            var id = taskIds[i];
            try
            {
                var r = await _svc.CancelTask(id, cancellationToken);
                results.Add(new BatchCancelResult(id, true, r.Cancelled, null, r.Number));
            }
            catch (OperationCanceledException) { throw; }
            catch (Exception ex)
            {
                results.Add(new BatchCancelResult(id, false, false, ex.Message));
            }
            ProgressReporter.ReportItem(progress, "Cancelling tasks", i + 1, taskIds.Length);
        }
        return results;
    }

    [McpServerTool, Description(
        "Delete many tasks at once — use for bulk cleanup instead of calling delete_task per id." +
        McpToolDocs.NotWhileRunning + McpToolDocs.MaxBatch + McpToolDocs.BatchProgressHint)]
    public async Task<IReadOnlyList<BatchTaskResult>> BatchDeleteTasks(
        string[] taskIds, CancellationToken cancellationToken, IProgress<ProgressNotificationValue>? progress = null)
    {
        EnsureWithinCap(taskIds, nameof(taskIds));
        return await RunPerTaskAsync(taskIds,
            async (id, ct) => (await _svc.DeleteTask(id, ct)).Number, cancellationToken, progress, "Deleting tasks");
    }

    [McpServerTool, Description(
        "Set or clear MyDay (daily prep) for many tasks at once — use instead of calling set_my_day per task. " +
        "Still cap-guarded: items that would exceed DailyPrepMaxTasks open MyDay tasks fail individually " +
        "(ok=false) without blocking the rest." + McpToolDocs.MaxBatch + McpToolDocs.BatchProgressHint)]
    public async Task<IReadOnlyList<BatchTaskResult>> BatchSetMyDay(
        BatchSetMyDayInput[] items, CancellationToken cancellationToken, IProgress<ProgressNotificationValue>? progress = null)
    {
        EnsureWithinCap(items, nameof(items));

        var results = new List<BatchTaskResult>(items.Length);
        for (var i = 0; i < items.Length; i++)
        {
            var item = items[i];
            try
            {
                var task = await _svc.SetMyDay(item.TaskId, item.IsMyDay, item.SortOrder, cancellationToken);
                results.Add(new BatchTaskResult(item.TaskId, true, null, Number: task.Number));
            }
            catch (OperationCanceledException) { throw; }
            catch (Exception ex)
            {
                results.Add(new BatchTaskResult(item.TaskId, false, ex.Message));
            }
            ProgressReporter.ReportItem(progress, "Setting My Day", i + 1, items.Length);
        }
        return results;
    }

    [McpServerTool, Description(
        "Remove the worktrees (directory + git branch) of many tasks at once — use for bulk cleanup instead " +
        "of calling cleanup_task_worktree per id." +
        McpToolDocs.NotWhileRunning + McpToolDocs.MaxBatch + McpToolDocs.BatchProgressHint)]
    public async Task<IReadOnlyList<BatchCleanupResult>> BatchCleanupTaskWorktrees(
        string[] taskIds,
        [Description("If true, also remove a dirty worktree, losing uncommitted changes; a Running task " +
            "is still refused either way.")] bool force = false,
        CancellationToken cancellationToken = default,
        IProgress<ProgressNotificationValue>? progress = null)
    {
        EnsureWithinCap(taskIds, nameof(taskIds));

        var results = new List<BatchCleanupResult>(taskIds.Length);
        for (var i = 0; i < taskIds.Length; i++)
        {
            var id = taskIds[i];
            try
            {
                var r = await _worktree.CleanupTaskWorktree(id, force, cancellationToken);
                results.Add(new BatchCleanupResult(id, true, r.Removed, r.BranchDeleted, null, r.Number));
            }
            catch (OperationCanceledException) { throw; }
            catch (Exception ex)
            {
                results.Add(new BatchCleanupResult(id, false, false, false, ex.Message));
            }
            ProgressReporter.ReportItem(progress, "Cleaning up worktrees", i + 1, taskIds.Length);
        }
        return results;
    }

    private static async Task<IReadOnlyList<BatchTaskResult>> RunPerTaskAsync(
        string[] taskIds, Func<string, CancellationToken, Task<int?>> op, CancellationToken cancellationToken,
        IProgress<ProgressNotificationValue>? progress, string progressAction)
    {
        var results = new List<BatchTaskResult>(taskIds.Length);
        for (var i = 0; i < taskIds.Length; i++)
        {
            var id = taskIds[i];
            try
            {
                var number = await op(id, cancellationToken);
                results.Add(new BatchTaskResult(id, true, null, Number: number));
            }
            catch (OperationCanceledException) { throw; }
            catch (Exception ex)
            {
                results.Add(new BatchTaskResult(id, false, ex.Message));
            }
            ProgressReporter.ReportItem(progress, progressAction, i + 1, taskIds.Length);
        }
        return results;
    }

    private static void EnsureWithinCap<T>(IReadOnlyCollection<T>? items, string name)
    {
        if (items is null || items.Count == 0)
            throw new InvalidOperationException($"{name} is required and must contain at least one item.");
        if (items.Count > MaxBatchSize)
            throw new InvalidOperationException(
                $"Batch too large: {items.Count} items (max {MaxBatchSize}). Split into smaller batches.");
    }
}