using ClaudeDo.Data.Models; using ClaudeDo.Data.Repositories; using ClaudeDo.Worker.External; using ClaudeDo.Worker.Hub; using ClaudeDo.Worker.Runner; using ClaudeDo.Worker.Tests.Infrastructure; using ClaudeDo.Worker.Usage; using Microsoft.EntityFrameworkCore; using Microsoft.Extensions.Logging.Abstractions; namespace ClaudeDo.Worker.Tests.Runner; /// Handing a PowerUse list-handler round over cleanly across the 5h usage limit: warn the live /// session once, then withhold the handoff it produces until the window resets. The UI half (the /// warning actually reaching the ConPTY session) lives in Ui.Tests' MissionControlViewModel tests. public sealed class HandlerUsageGuardTests : IDisposable { private readonly DbFixture _db = new(); private readonly string _tempDir; private readonly CapturingHubContext _hub = new(); private readonly HandlerRoundState _rounds; private readonly HandoffRoundTracker _handoffRounds = new(); private static readonly DateTimeOffset ResetsAt = new(2026, 9, 3, 2, 0, 0, TimeSpan.Zero); private static readonly DateTimeOffset DueAt = ResetsAt + HandlerUsageGuard.ResetGrace; private static readonly DateTimeOffset Now = ResetsAt.AddHours(-3); public HandlerUsageGuardTests() { _tempDir = Path.Combine(Path.GetTempPath(), $"cd_usageguard_{Guid.NewGuid():N}"); Directory.CreateDirectory(_tempDir); _rounds = new HandlerRoundState(RoundsPath); } private string RoundsPath => Path.Combine(_tempDir, "handler-rounds.json"); public void Dispose() { _db.Dispose(); try { Directory.Delete(_tempDir, true); } catch { } } private static UsageState Usage(double percent, DateTimeOffset? resetsAt) => Snapshot(new UsageBucket(percent, resetsAt)); private static UsageState Snapshot(UsageBucket? fiveHour) { var state = new UsageState(); state.ReportSuccess(new UsageSnapshot(fiveHour, null, Array.Empty(), DateTime.UtcNow)); return state; } private HandlerUsageGuard BuildGuard(UsageState usage, DateTimeOffset? now = null, HandlerRoundState? rounds = null) => new(_db.CreateFactory(), usage, rounds ?? _rounds, _handoffRounds, new HubBroadcaster(_hub), NullLogger.Instance) { NowUtc = () => now ?? Now }; private int _nextNumber = 1; private async Task SeedHandlerTaskAsync(bool powerUseEnabled = true, bool powerUseAllowed = true) { var listId = Guid.NewGuid().ToString(); var taskId = Guid.NewGuid().ToString(); using (var ctx = _db.CreateContext()) { ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", WorkingDir = _tempDir, CreatedAt = DateTime.UtcNow }); ctx.ListConfigs.Add(new ListConfigEntity { ListId = listId, PowerUseAllowed = powerUseAllowed }); ctx.Tasks.Add(new TaskEntity { Id = taskId, ListId = listId, Title = "Handler", Number = _nextNumber++, IsManual = true, CreatedAt = DateTime.UtcNow, }); await ctx.SaveChangesAsync(); // Written straight to the row so the fixture doesn't depend on the repository's // clamping/normalizing copy block. await new AppSettingsRepository(ctx).GetAsync(); var settings = await ctx.AppSettings.FirstAsync(s => s.Id == AppSettingsEntity.SingletonId); settings.PowerUseEnabled = powerUseEnabled; await ctx.SaveChangesAsync(); } return taskId; } private IReadOnlyList Calls(string method) => _hub.Proxy.Calls.Where(c => c.Method == method).ToList(); private IReadOnlyList Warnings => Calls("HandlerUsageWarning"); private IReadOnlyList HandoffCalls => Calls("HandoffRequested"); // ── acceptance 1: exactly one warning per round ──────────────────────────────────────────── [Fact] public async Task AtTheThreshold_WarnsTheSessionExactlyOncePerRound() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.Merge, new[] { "a" }, DateTime.UtcNow)); var guard = BuildGuard(Usage(95, ResetsAt)); await guard.TickAsync(default); await guard.TickAsync(default); await BuildGuard(Usage(99, ResetsAt)).TickAsync(default); var warning = Assert.Single(Warnings); Assert.Equal(taskId, warning.Args[0]); var message = (string)warning.Args[1]!; Assert.Contains("USAGE LIMIT IMMINENT", message); // Names the session's own phase back, so it resumes where it is. Assert.Contains($"nextPhase: \"{MergeHelperPhase.Merge}\"", message); Assert.Contains(taskId, message); } [Fact] public async Task BelowTheThreshold_NoWarning() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.Merge, new[] { "a" }, DateTime.UtcNow)); await BuildGuard(Usage(94, ResetsAt)).TickAsync(default); Assert.Empty(Warnings); } // The triage round has no "triage" handoff phase -- what follows it is always the wait phase. [Fact] public async Task TriageRound_IsToldToHandOffToTheWaitPhase() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound( taskId, InteractiveLaunchSpecService.TriagePhase, new[] { "a" }, DateTime.UtcNow)); await BuildGuard(Usage(96, ResetsAt)).TickAsync(default); Assert.Contains($"nextPhase: \"{MergeHelperPhase.Wait}\"", (string)Assert.Single(Warnings).Args[1]!); } // A sealed chain can't hand off at all -- telling it to would only earn the session an exception. [Fact] public async Task SealedChain_IsNotWarned() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.MergeFinal, new[] { "a" }, DateTime.UtcNow)); _handoffRounds.Register(taskId, MergeHelperPhase.MergeFinal); await BuildGuard(Usage(99, ResetsAt)).TickAsync(default); Assert.Empty(Warnings); } // ── acceptance 5: PowerUse gates both halves ─────────────────────────────────────────────── [Theory] [InlineData(false, true)] [InlineData(true, false)] public async Task PowerUseOff_NeitherWarnsNorWithholds(bool powerUseEnabled, bool powerUseAllowed) { var taskId = await SeedHandlerTaskAsync(powerUseEnabled, powerUseAllowed); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.Merge, new[] { "a" }, DateTime.UtcNow)); var guard = BuildGuard(Usage(99, ResetsAt)); await guard.TickAsync(default); Assert.Empty(Warnings); Assert.Null(await guard.TryDeferHandoffAsync(taskId, new[] { "a" }, MergeHelperPhase.Merge, default)); Assert.Null(_rounds.Get(taskId)!.DeferredHandoffAt); } // ── acceptance 2 + 3: withholding the handoff ────────────────────────────────────────────── [Fact] public async Task BelowTheThreshold_HandoffIsNotWithheld() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.Merge, new[] { "a" }, DateTime.UtcNow)); var deferred = await BuildGuard(Usage(94, ResetsAt)) .TryDeferHandoffAsync(taskId, new[] { "a" }, MergeHelperPhase.Merge, default); Assert.Null(deferred); Assert.Null(_rounds.Get(taskId)!.DeferredHandoffAt); } [Fact] public async Task AtTheThreshold_HandoffIsWithheldUntilTheResetPlusGrace() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.Wait, new[] { "a" }, DateTime.UtcNow)); var guard = BuildGuard(Usage(95, ResetsAt)); var deferred = await guard.TryDeferHandoffAsync(taskId, new[] { "a", "b" }, MergeHelperPhase.Merge, default); Assert.Equal(DueAt, deferred); Assert.Empty(HandoffCalls); // Still parked before it is due, sent afterwards -- with the handed-off phase and task set, // not the round's own. await BuildGuard(Usage(95, ResetsAt), DueAt.AddMinutes(-1)).TickAsync(default); Assert.Empty(HandoffCalls); await BuildGuard(Usage(95, ResetsAt), DueAt).TickAsync(default); var call = Assert.Single(HandoffCalls); Assert.Equal(taskId, call.Args[0]); Assert.Equal(new[] { "a", "b" }, (IReadOnlyList)call.Args[1]!); Assert.Equal(MergeHelperPhase.Merge, call.Args[2]); // A clean handoff must not cost the next session its auto-restart budget. Assert.False(_rounds.Get(taskId)!.AutoRestartUsed); } // Utilization can dip back below the threshold while the session finishes its merge; the handoff // that the warning ASKED for still has to wait for the reset. [Fact] public async Task AfterAWarning_HandoffIsWithheldEvenIfUtilizationDropped() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.Merge, new[] { "a" }, DateTime.UtcNow)); await BuildGuard(Usage(97, ResetsAt)).TickAsync(default); Assert.Single(Warnings); var deferred = await BuildGuard(Usage(80, ResetsAt)) .TryDeferHandoffAsync(taskId, new[] { "a" }, MergeHelperPhase.Merge, default); Assert.Equal(DueAt, deferred); } // No reset time = nothing to wait for. Sending now beats parking a handoff that never fires. [Fact] public async Task NoResetTime_HandsOffImmediately() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.Merge, new[] { "a" }, DateTime.UtcNow)); var deferred = await BuildGuard(Usage(99, null)) .TryDeferHandoffAsync(taskId, new[] { "a" }, MergeHelperPhase.Merge, default); Assert.Null(deferred); Assert.Null(_rounds.Get(taskId)!.DeferredHandoffAt); } // ── acceptance 4: a parked handoff survives a worker restart ─────────────────────────────── [Fact] public async Task WorkerRestart_ReschedulesAParkedHandoff_AndSendsAnOverdueOneAtOnce() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.Merge, new[] { "a" }, DateTime.UtcNow)); await BuildGuard(Usage(99, ResetsAt)) .TryDeferHandoffAsync(taskId, new[] { "a" }, MergeHelperPhase.Merge, default); // Worker restart: nothing but the file carries over. var afterRestart = new HandlerRoundState(RoundsPath); Assert.Equal(DueAt, afterRestart.Get(taskId)!.DeferredHandoffAt); // Not due yet -> still parked. await BuildGuard(Usage(99, ResetsAt), Now, afterRestart).TickAsync(default); Assert.Empty(HandoffCalls); // Long overdue (the worker was down past the reset) -> goes out on the first tick. await BuildGuard(Usage(99, ResetsAt), DueAt.AddHours(2), afterRestart).TickAsync(default); Assert.Single(HandoffCalls); Assert.Null(afterRestart.Get(taskId)!.DeferredHandoffAt); } // One trigger, one state: a parked handoff blocks the auto-restart of the same round, so the // round can never be reopened twice. [Fact] public async Task AParkedHandoff_BlocksTheAutoRestartOfTheSameRound() { var taskId = await SeedHandlerTaskAsync(); _rounds.Save(new HandlerRound(taskId, MergeHelperPhase.Merge, new[] { "a" }, DateTime.UtcNow)); var usage = Usage(99, ResetsAt); await BuildGuard(usage).TryDeferHandoffAsync(taskId, new[] { "a" }, MergeHelperPhase.Merge, default); var coordinator = new HandlerRestartCoordinator( _db.CreateFactory(), usage, _rounds, new HubBroadcaster(_hub), NullLogger.Instance); Assert.False(await coordinator.OnSessionExitedAsync(taskId, 1)); await BuildGuard(usage, DueAt).TickAsync(default); Assert.Single(HandoffCalls); } }