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

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

namespace ClaudeDo.Worker.External;

// BlockedReason is set only when Status is "Blocked" -- a Queued task the picker will not
// claim yet, either because of a planning-chain predecessor or an unmet depends-on link.
public sealed record TaskStatusChangeDto(string TaskId, string Status, string? BlockedReason = null);
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 wall-clock 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.
    //
    // MCP_TOOL_TIMEOUT is a separate mechanism from Claude Code's idle-silence abort
    // (CLAUDE_CODE_MCP_TOOL_IDLE_TIMEOUT, default 300s for HTTP-transport MCP servers like this
    // one) -- no ClaudeDo launcher raises that one. Without periodic activity, a wait anywhere
    // near this cap gets killed at ~300s regardless of launcher. WaitForTaskChange survives that
    // by reporting MCP progress well under 300s apart (see ProgressReportInterval below), which
    // resets Claude Code's idle timer.
    internal const int MaxTimeoutSeconds = 900;
    private static readonly TimeSpan PollInterval = TimeSpan.FromMilliseconds(500);

    // Not readonly: tests shrink this to observe a progress report without waiting 30s.
    internal static TimeSpan ProgressReportInterval = TimeSpan.FromSeconds(30);

    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 -- use this instead of " +
        "polling get_task in a loop. Returns immediately if a task is already outside Queued/Running " +
        "(an unknown id reports status \"NotFound\" and counts as changed). A Queued task the picker " +
        "will not claim yet (a planning-chain predecessor, or a depends_on link whose target isn't " +
        "Done) also reports immediately as status \"Blocked\" with blockedReason set, instead of " +
        "silently waiting out the full timeout. Pitfall: a planning parent goes Running -> " +
        "WaitingForChildren while its children are still working, so by default waiting on a parent " +
        "returns early; see treatWaitingForChildrenAsBusy. Sends MCP progress " +
        "pings every 30s while waiting so a long wait survives the calling client's own idle-silence " +
        "abort (Claude Code defaults to killing an MCP call after ~300s of silence) -- this is not " +
        "guaranteed by every possible MCP client.")]
    public async Task<WaitForTaskChangeResult> WaitForTaskChange(
        string[] taskIds,
        [Description(
            "How long to wait, in seconds, before giving up. Clamped server-side to 900s (15 min) " +
            "regardless of what's passed.")]
        int timeoutSeconds = 60,
        [Description(
            "When true, WaitingForChildren still counts as busy, so waiting on a planning parent " +
            "continues until it reaches WaitingForReview or a terminal status instead of returning " +
            "as soon as it leaves Running.")]
        bool treatWaitingForChildrenAsBusy = false,
        IProgress<ProgressNotificationValue>? progress = null,
        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
        {
            var lastProgressAt = DateTime.UtcNow;
            while (true)
            {
                var changed = await CheckOnceAsync(taskIds, treatWaitingForChildrenAsBusy, linked.Token);
                if (changed.Count > 0)
                    return new WaitForTaskChangeResult(changed, TimedOut: false);

                var now = DateTime.UtcNow;
                if (now - lastProgressAt >= ProgressReportInterval)
                {
                    lastProgressAt = now;
                    progress?.Report(new ProgressNotificationValue { Progress = 0, Message = "Still waiting for a task status change." });
                }

                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, t.BlockedByTaskId, t.DependsOnTaskId })
            .ToListAsync(ct);

        var byId = rows.ToDictionary(r => r.Id, r => r);

        // A Queued task with an unmet depends_on link never becomes "changed" by itself -- the
        // picker will not touch it. Batch-resolve those dependencies' statuses once instead of a
        // query per candidate.
        var dependencyIds = rows
            .Where(r => r.Status == TaskStatus.Queued && r.BlockedByTaskId is null && r.DependsOnTaskId is not null)
            .Select(r => r.DependsOnTaskId!)
            .Distinct()
            .ToList();
        var dependencyStatuses = dependencyIds.Count == 0
            ? new Dictionary<string, TaskStatus>()
            : await ctx.Tasks.AsNoTracking()
                .Where(t => dependencyIds.Contains(t.Id))
                .Select(t => new { t.Id, t.Status })
                .ToDictionaryAsync(t => t.Id, t => t.Status, ct);

        var result = new List<TaskStatusChangeDto>();
        foreach (var id in taskIds)
        {
            if (!byId.TryGetValue(id, out var row))
            {
                result.Add(new TaskStatusChangeDto(id, "NotFound"));
                continue;
            }

            if (row.Status == TaskStatus.Queued)
            {
                if (row.BlockedByTaskId is not null)
                {
                    result.Add(new TaskStatusChangeDto(id, "Blocked",
                        $"Blocked by planning-chain predecessor {row.BlockedByTaskId}."));
                    continue;
                }
                if (row.DependsOnTaskId is not null)
                {
                    var known = dependencyStatuses.TryGetValue(row.DependsOnTaskId, out var depStatus);
                    if (!known || depStatus != TaskStatus.Done)
                    {
                        result.Add(new TaskStatusChangeDto(id, "Blocked",
                            $"Blocked: depends on task {row.DependsOnTaskId} (status: " +
                            (known ? depStatus.ToString() : "not found") + ")."));
                        continue;
                    }
                }
            }

            var busy = row.Status == TaskStatus.Queued || row.Status == TaskStatus.Running
                || (treatWaitingForChildrenAsBusy && row.Status == TaskStatus.WaitingForChildren);
            if (!busy)
                result.Add(new TaskStatusChangeDto(id, row.Status.ToString()));
        }
        return result;
    }
}