using ClaudeDo.Data; using ClaudeDo.Data.Models; using ClaudeDo.Data.Repositories; using ClaudeDo.Worker.External; using ClaudeDo.Worker.Hub; using ClaudeDo.Worker.Usage; using Microsoft.EntityFrameworkCore; namespace ClaudeDo.Worker.Runner; /// /// Hands a PowerUse list-handler round over cleanly across the 5h usage limit, instead of letting it /// die mid-merge: /// /// 1. Warn — once per round, as soon as the 5h window crosses /// AppSettings.PowerUseHandoffPct, a message is typed straight into the live ConPTY session /// (HandlerUsageWarning → Mission Control) telling it to finish what it is doing and call /// handoff_list_handler itself. A session that hands off on its own leaves no half-merged /// tree behind; an auto-restart after the fact cannot clean one up. /// 2. Withhold — the resulting handoff is not broadcast immediately (the next session would /// open straight into the same exhausted window) but parked in /// until the window resets. Same slot the auto-restart of a session that *did* die uses /// (), so the two can never reopen one round twice. /// 3. Fire — this service's tick is also what releases a parked handoff. Deliberately a poll /// over the persisted state rather than a timer: a worker restart mid-wait re-picks the parked /// handoff up on the next tick, and an overdue one goes out immediately, for free. /// public sealed class HandlerUsageGuard : BackgroundService { internal static readonly TimeSpan TickInterval = TimeSpan.FromSeconds(30); // Small cushion past ResetsAt: the reported reset time is the server's, and reopening a session // one second early buys nothing but another dead tile. internal static readonly TimeSpan ResetGrace = TimeSpan.FromSeconds(60); private readonly IDbContextFactory _dbFactory; private readonly UsageState _usageState; private readonly HandlerRoundState _rounds; private readonly HandoffRoundTracker _handoffRounds; private readonly HubBroadcaster _broadcaster; private readonly ILogger _logger; public HandlerUsageGuard( IDbContextFactory dbFactory, UsageState usageState, HandlerRoundState rounds, HandoffRoundTracker handoffRounds, HubBroadcaster broadcaster, ILogger logger) { _dbFactory = dbFactory; _usageState = usageState; _rounds = rounds; _handoffRounds = handoffRounds; _broadcaster = broadcaster; _logger = logger; } /// Test seam — a real 5h wait is not a test. internal Func NowUtc { get; set; } = () => DateTimeOffset.UtcNow; /// PowerUse verdict for one handler task: both PowerUse switches on, a 5h bucket /// available, plus the configured threshold and the task's display number. internal sealed record Policy(int HandoffPct, UsageBucket FiveHour, int Number); protected override async Task ExecuteAsync(CancellationToken stoppingToken) { while (!stoppingToken.IsCancellationRequested) { try { await TickAsync(stoppingToken); await Task.Delay(TickInterval, stoppingToken); } catch (OperationCanceledException) { return; } catch (Exception ex) { _logger.LogWarning(ex, "HandlerUsageGuard: tick failed"); await Task.Delay(TickInterval, stoppingToken); } } } internal async Task TickAsync(CancellationToken ct) { foreach (var round in _rounds.All()) { await FireIfDueAsync(round); await WarnIfNearLimitAsync(round, ct); } } /// Called by handoff_list_handler. Returns the time the handoff was parked until /// when it must be withheld, or null when it should go out now (PowerUse off for this list, well /// below the threshold, or no reset time to wait for). public async Task TryDeferHandoffAsync( string handlerTaskId, IReadOnlyList survivingTaskIds, string phase, CancellationToken ct) { var policy = await EvaluateAsync(handlerTaskId, ct); if (policy is null) return null; // The warning flag matters on its own: utilization can dip back under the threshold between // the warning and the session actually finishing its merge, and a handoff produced *because* // of the warning must still wait for the reset. var warned = _rounds.Get(handlerTaskId)?.UsageWarningSent == true; if (policy.FiveHour.Utilization < policy.HandoffPct && !warned) return null; if (policy.FiveHour.ResetsAt is not { } resetsAt) { _logger.LogWarning( "HandlerUsageGuard: 5h window at {Pct}% has no reset time; handing off {TaskId} immediately", policy.FiveHour.Utilization, handlerTaskId); return null; } var dueAt = resetsAt + ResetGrace; _rounds.SetDeferredHandoff(handlerTaskId, dueAt, phase, survivingTaskIds); await _broadcaster.WorkerLog( $"List handler #{policy.Number}: handoff to phase {phase} withheld until " + $"{dueAt.ToLocalTime():HH:mm} (5h window at {policy.FiveHour.Utilization:0}%)", WorkerLogLevel.Warn, DateTime.UtcNow); return dueAt; } private async Task FireIfDueAsync(HandlerRound round) { if (round.DeferredHandoffAt is not { } dueAt || dueAt > NowUtc()) return; if (round.DeferredPhase is not { Length: > 0 } phase || round.DeferredTaskIds is not { Count: > 0 } taskIds) { _rounds.ClearDeferredHandoff(round.HandlerTaskId, spendAutoRestart: false); return; } // Cleared BEFORE the broadcast so a throw on the way out can't leave a slot that re-fires on // every tick; the restart budget is spent here too, which is the point of the whole // rework -- it is now only spent once a handoff really went out. _rounds.ClearDeferredHandoff(round.HandlerTaskId, spendAutoRestart: round.DeferredIsAutoRestart); await _broadcaster.HandoffRequested(round.HandlerTaskId, taskIds, phase); await _broadcaster.WorkerLog( $"Reopening list handler at phase {phase} after the usage-limit reset", WorkerLogLevel.Info, DateTime.UtcNow); } private async Task WarnIfNearLimitAsync(HandlerRound round, CancellationToken ct) { if (round.UsageWarningSent) return; // A sealed chain has no further handoff to ask for -- telling it to call // handoff_list_handler would only earn the session an exception. if (_handoffRounds.GetState(round.HandlerTaskId).Sealed) return; var policy = await EvaluateAsync(round.HandlerTaskId, ct); if (policy is null || policy.FiveHour.Utilization < policy.HandoffPct) return; if (!_rounds.TryMarkUsageWarned(round.HandlerTaskId)) return; await _broadcaster.HandlerUsageWarning( round.HandlerTaskId, BuildWarning(round.HandlerTaskId, round.Phase, policy.FiveHour.Utilization)); await _broadcaster.WorkerLog( $"List handler #{policy.Number}: warned the session at 5h {policy.FiveHour.Utilization:0}% " + $"(handoff threshold {policy.HandoffPct}%)", WorkerLogLevel.Warn, DateTime.UtcNow); } // Plain ASCII on purpose: this is typed character-by-character into a PTY, not rendered by us. internal static string BuildWarning(string handlerTaskId, string phase, double utilization) => $"USAGE LIMIT IMMINENT ({utilization:0}% of the 5h window). Finish or abort the merge you are " + "in - never leave conflict markers in the tree - then call " + $"handoff_list_handler(taskId: \"{handlerTaskId}\", survivingTaskIds: , " + $"nextPhase: \"{NextPhaseFor(phase)}\") right now. Do not start anything new. " + "The next session opens automatically after the limit resets."; // The session should resume where it is, so the warning names its own phase back. The opening // triage round is the exception: "triage" is not a handoff phase, and what follows triage is // always the wait phase. private static string NextPhaseFor(string phase) => MergeHelperPhase.All.Contains(phase, StringComparer.Ordinal) ? phase : MergeHelperPhase.Wait; /// Both PowerUse switches, exactly as ComposeHandlerPromptAsync and /// gate them — the global toggle alone must not make an /// unrelated list's handler session get typed into. Null = leave this round alone. private async Task EvaluateAsync(string handlerTaskId, CancellationToken ct) { if (_usageState.Snapshot?.FiveHour is not { } fiveHour) return null; try { await using var ctx = await _dbFactory.CreateDbContextAsync(ct); var task = await new TaskRepository(ctx).GetByIdAsync(handlerTaskId, ct); if (task is null) return null; var settings = await new AppSettingsRepository(ctx).GetAsync(ct); if (!settings.PowerUseEnabled) return null; if ((await new ListRepository(ctx).GetConfigAsync(task.ListId, ct))?.PowerUseAllowed != true) return null; return new Policy(Math.Clamp(settings.PowerUseHandoffPct, 50, 99), fiveHour, task.Number); } catch (Exception ex) { _logger.LogWarning(ex, "HandlerUsageGuard: could not evaluate handler task {TaskId}", handlerTaskId); return null; } } }