ClaudeDo/src/ClaudeDo.Worker/External/HandoffRoundTracker.cs

using System.Collections.Concurrent;
using ClaudeDo.Worker.Runner;

namespace ClaudeDo.Worker.External;

// Bounds the list-handler wait<->merge handoff chain server-side, so a model that ignores the
// cooperative "this is the FINAL round" line in handoff.md (InteractiveLaunchSpecService's
// roundInfo, read back via GetState) cannot loop wait->merge->wait->merge forever, each round
// spawning a fresh ConPTY session. Keyed by handler task id -- a "Let Claude handle it" run always creates a fresh
// handler task (InteractiveLaunchSpecService.CreateMergeHelperTaskAsync), so the counter is
// per-run by construction: no reset logic needed, and a later run on the same list starts at 0
// on its own. In-memory only, same tradeoff as TaskRunTokenRegistry -- a worker restart mid-run
// resets the count, which is acceptable since the chain restarts fresh anyway.
public sealed class HandoffRoundTracker
{
    // Two full wait->merge cycles: wait, merge, wait, merge. Only the fallback now -- the live cap
    // comes from AppSettingsEntity.ListHandlerMaxHandoffs, passed per call by HandoffMcpTools, so a
    // run meant to keep going for hours can raise it without the bound disappearing entirely.
    public const int DefaultMaxHandoffs = 4;

    private sealed class TaskState
    {
        public int Count;
        public bool Sealed;
    }

    private readonly ConcurrentDictionary<string, TaskState> _states = new();

    public sealed record Result(bool Rejected, bool Coerced, string EffectivePhase, string? Error);

    /// <summary>Read-only snapshot for a handler task -- does NOT register a round. Used to report
    /// the current round/remaining budget (e.g. in the handoff brief) without mutating state that
    /// only Register should advance. Count is 0 for a task that hasn't handed off yet.</summary>
    public (int Count, bool Sealed) GetState(string taskId)
    {
        if (!_states.TryGetValue(taskId, out var state))
            return (0, false);
        lock (state)
            return (state.Count, state.Sealed);
    }

    // Read-only budget lookup for a handler task that hasn't (yet) called HandoffListHandler --
    // Round/Remaining/Final all report their untouched defaults instead of throwing, since "no
    // handoff yet" is a normal state, not an error.
    public sealed record Status(int Round, int Remaining, bool Final);

    /// <param name="maxHandoffs">The live cap (AppSettingsEntity.ListHandlerMaxHandoffs), passed in
    /// rather than read off a constant so the REPORTED budget can never drift from the cap Register
    /// actually enforces. A brief telling the model "2 rounds left" while the chain really allows 40
    /// is worse than no budget line at all. Clamped like Register's.</param>
    public Status GetStatus(string taskId, int maxHandoffs = DefaultMaxHandoffs)
    {
        var cap = Math.Max(1, maxHandoffs);

        if (!_states.TryGetValue(taskId, out var state))
            return new Status(0, cap, false);

        lock (state)
            return new Status(state.Count, Math.Max(0, cap - state.Count), state.Sealed);
    }

    /// <param name="maxHandoffs">Cap for this chain, normally AppSettingsEntity.ListHandlerMaxHandoffs.
    /// Values below 1 are treated as 1: the cap exists to bound the loop, so "no cap" is not an
    /// option a caller gets to pick by passing 0.</param>
    public Result Register(string taskId, string nextPhase, int maxHandoffs = DefaultMaxHandoffs)
    {
        var cap = Math.Max(1, maxHandoffs);
        var state = _states.GetOrAdd(taskId, _ => new TaskState());
        lock (state)
        {
            if (state.Sealed)
                return new Result(true, false, nextPhase,
                    $"Task {taskId} already completed its merge_final round -- no further handoff is allowed. " +
                    "Start a new handler run (\"Let Claude handle it\") on this list instead.");

            state.Count++;
            var isFinal = nextPhase is MergeHelperPhase.WaitFinal or MergeHelperPhase.MergeFinal;

            var effective = nextPhase;
            if (state.Count > cap && !isFinal)
            {
                effective = nextPhase == MergeHelperPhase.Wait ? MergeHelperPhase.WaitFinal : MergeHelperPhase.MergeFinal;
                isFinal = true;
            }

            if (isFinal && effective == MergeHelperPhase.MergeFinal)
                state.Sealed = true;

            return new Result(false, !string.Equals(effective, nextPhase, StringComparison.Ordinal), effective, null);
        }
    }
}