ClaudeDo/src/ClaudeDo.Worker/External/TaskWaitMcpTools.cs

using System.ComponentModel;
using ClaudeDo.Data;
using Microsoft.EntityFrameworkCore;
using ModelContextProtocol.Server;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.External;

public sealed record TaskStatusChangeDto(string TaskId, string Status);
public sealed record WaitForTaskChangeResult(IReadOnlyList<TaskStatusChangeDto> Changed, bool TimedOut);

[McpServerToolType]
public sealed class TaskWaitMcpTools
{
    // Every ClaudeDo-owned launcher (ClaudeProcess for headless runs, InteractiveLaunchSpecService
    // for ConPTY sessions) sets MCP_TOOL_TIMEOUT=930000ms on the claude CLI process; this cap
    // leaves a ~30s margin under that so the tool itself reports TimedOut instead of racing the
    // client's own abort. A caller running claude with a different MCP_TOOL_TIMEOUT (or none --
    // the CLI default is 60s) will see its own client-side timeout fire first; this tool has no
    // way to detect or compensate for that from the server side.
    internal const int MaxTimeoutSeconds = 900;
    private static readonly TimeSpan PollInterval = TimeSpan.FromMilliseconds(500);

    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;

    public TaskWaitMcpTools(IDbContextFactory<ClaudeDoDbContext> dbFactory)
    {
        _dbFactory = dbFactory;
    }

    [McpServerTool, Description(
        "Blocks until at least one of the given tasks leaves Queued/Running, or until timeoutSeconds elapses " +
        "(clamped server-side to 900s). Returns immediately if any task is already outside Queued/Running " +
        "when called (an unknown id is reported as status \"NotFound\" and counts as changed). Use this instead " +
        "of polling get_task in a loop. Pitfall: a planning parent with children goes Running -> " +
        "WaitingForChildren while its children are still working, and by default that counts as \"changed\" -- " +
        "so waiting on a parent returns immediately even though the work isn't done. Set " +
        "treatWaitingForChildrenAsBusy=true to keep waiting through WaitingForChildren; the call then only " +
        "returns once the parent reaches WaitingForReview or a terminal status (default: false, unchanged " +
        "legacy behavior). Requires the calling claude process to run with MCP_TOOL_TIMEOUT >= 930000 (ms) for " +
        "a long wait to actually be held open -- ClaudeDo's own launchers already set this. " +
        "Result: { changed: [{ taskId, status }], timedOut }.")]
    public async Task<WaitForTaskChangeResult> WaitForTaskChange(
        string[] taskIds, int timeoutSeconds = 60, bool treatWaitingForChildrenAsBusy = false,
        CancellationToken cancellationToken = default)
    {
        if (taskIds.Length == 0)
            throw new ArgumentException("taskIds must not be empty.", nameof(taskIds));

        var timeout = TimeSpan.FromSeconds(Math.Clamp(timeoutSeconds, 1, MaxTimeoutSeconds));
        using var timeoutCts = new CancellationTokenSource(timeout);
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);

        try
        {
            while (true)
            {
                var changed = await CheckOnceAsync(taskIds, treatWaitingForChildrenAsBusy, linked.Token);
                if (changed.Count > 0)
                    return new WaitForTaskChangeResult(changed, TimedOut: false);

                await Task.Delay(PollInterval, linked.Token);
            }
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
        {
            return new WaitForTaskChangeResult(Array.Empty<TaskStatusChangeDto>(), TimedOut: true);
        }
    }

    private async Task<IReadOnlyList<TaskStatusChangeDto>> CheckOnceAsync(
        string[] taskIds, bool treatWaitingForChildrenAsBusy, CancellationToken ct)
    {
        await using var ctx = await _dbFactory.CreateDbContextAsync(ct);
        var rows = await ctx.Tasks
            .AsNoTracking()
            .Where(t => taskIds.Contains(t.Id))
            .Select(t => new { t.Id, t.Status })
            .ToListAsync(ct);

        var byId = rows.ToDictionary(r => r.Id, r => r.Status);
        var result = new List<TaskStatusChangeDto>();
        foreach (var id in taskIds)
        {
            if (!byId.TryGetValue(id, out var status))
            {
                result.Add(new TaskStatusChangeDto(id, "NotFound"));
                continue;
            }
            var busy = status == TaskStatus.Queued || status == TaskStatus.Running
                || (treatWaitingForChildrenAsBusy && status == TaskStatus.WaitingForChildren);
            if (!busy)
                result.Add(new TaskStatusChangeDto(id, status.ToString()));
        }
        return result;
    }
}