ClaudeDo/src/ClaudeDo.Worker/Runner/TaskRunMcpService.cs

using System.ComponentModel;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Hub;
using ModelContextProtocol.Server;

namespace ClaudeDo.Worker.Runner;

public sealed record SuggestedImprovementDto(string ChildTaskId);

[McpServerToolType]
public sealed class TaskRunMcpService
{
    private readonly TaskRepository _tasks;
    private readonly TaskRunMcpContextAccessor _ctx;
    private readonly HubBroadcaster _broadcaster;
    private readonly PendingQuestionRegistry _pending;

    // How long a running task blocks waiting for the user to answer an AskUser question
    // before it gives up and proceeds autonomously. NOTE: the spawned claude process must
    // run with MCP_TOOL_TIMEOUT raised above this (ClaudeProcess sets it) — Claude Code
    // otherwise caps HTTP MCP tool calls at 60 s and would kill the call early.
    internal static readonly TimeSpan QuestionWindow = TimeSpan.FromMinutes(3);

    internal const string TimeoutFallback =
        "No response received within 3 minutes — proceed using your best judgment, "
        + "note the assumption you made, and continue.";

    public TaskRunMcpService(
        TaskRepository tasks,
        TaskRunMcpContextAccessor ctx,
        HubBroadcaster broadcaster,
        PendingQuestionRegistry pending)
    {
        _tasks = tasks;
        _ctx = ctx;
        _broadcaster = broadcaster;
        _pending = pending;
    }

    [McpServerTool, Description(
        "File an out-of-scope improvement as a child task of the current task. The child runs " +
        "automatically after this task finishes and is surfaced for review alongside it. Use ONLY " +
        "for work that is genuinely outside this task's scope (a refactor, follow-up, or tech debt) " +
        "— never for work that belongs to the current task. Set model to the cheapest model that can " +
        "do the follow-up well — 'haiku' for trivial/mechanical work, 'sonnet' for normal coding, " +
        "'opus' only for complex work. Leave model null to inherit the list/global default.")]
    public async Task<SuggestedImprovementDto> SuggestImprovement(
        string title,
        string description,
        string? model,
        CancellationToken cancellationToken)
    {
        var callerId = _ctx.Current.CallerTaskId;
        var caller = await _tasks.GetByIdAsync(callerId, cancellationToken)
            ?? throw new InvalidOperationException("Calling task not found.");
        if (caller.ParentTaskId is not null)
            throw new InvalidOperationException(
                "A child task cannot suggest further improvements (improvements are one layer deep).");

        var child = await _tasks.CreateChildAsync(
            callerId, title, description, commitType: null, createdBy: callerId, model: model, ct: cancellationToken);
        await _broadcaster.TaskUpdated(child.Id);
        await _broadcaster.TaskUpdated(callerId);
        return new SuggestedImprovementDto(child.Id);
    }

    [McpServerTool, Description(
        "Ask the user a question and wait up to 3 minutes for their answer. Use this ONLY "
        + "when you genuinely need a human decision to proceed correctly and a wrong guess "
        + "would be costly or hard to undo (an irreversible action, contradictory "
        + "requirements, or a real fork where both options have meaningful consequences). "
        + "Do NOT use it for routine choices you can reasonably make yourself — for those, "
        + "pick the most sensible option and continue. The returned string is the user's "
        + "answer; if no one responds in time it tells you to proceed on your own judgment.")]
    public Task<string> AskUser(string question, CancellationToken cancellationToken)
    {
        var callerId = _ctx.Current.CallerTaskId;
        return AwaitAnswerAsync(
            _pending, callerId, question ?? string.Empty,
            (questionId, q) => _broadcaster.TaskQuestionAsked(callerId, questionId, q),
            questionId => _broadcaster.TaskQuestionResolved(callerId, questionId),
            QuestionWindow, cancellationToken);
    }

    // Registers the question, signals it, and blocks until the user answers, the window
    // elapses (returns the fallback), or the run is cancelled (rethrows). Pure except for
    // the two broadcast callbacks, so it is unit-testable without a real hub or DbContext.
    internal static async Task<string> AwaitAnswerAsync(
        PendingQuestionRegistry pending,
        string taskId,
        string question,
        Func<string, string, Task> onAsked,
        Func<string, Task> onResolved,
        TimeSpan timeout,
        CancellationToken ct)
    {
        var (questionId, answer) = pending.Register(taskId, question);
        await onAsked(questionId, question);
        try
        {
            using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
            timeoutCts.CancelAfter(timeout);
            try
            {
                return await answer.WaitAsync(timeoutCts.Token);
            }
            catch (OperationCanceledException) when (!ct.IsCancellationRequested)
            {
                // The window elapsed (not a run cancellation) — let the agent carry on.
                return TimeoutFallback;
            }
        }
        finally
        {
            pending.Remove(taskId, questionId);
            await onResolved(questionId);
        }
    }
}