ClaudeDo/tests/ClaudeDo.Worker.Tests/Runner/HandlerUsageGuardTests.cs

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<UsageLimitRow>(), 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<HandlerUsageGuard>.Instance)
        { NowUtc = () => now ?? Now };

    private int _nextNumber = 1;

    private async Task<string> 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<CapturedHubCall> Calls(string method) =>
        _hub.Proxy.Calls.Where(c => c.Method == method).ToList();

    private IReadOnlyList<CapturedHubCall> Warnings => Calls("HandlerUsageWarning");
    private IReadOnlyList<CapturedHubCall> 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<string>)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<HandlerRestartCoordinator>.Instance);
        Assert.False(await coordinator.OnSessionExitedAsync(taskId, 1));

        await BuildGuard(usage, DueAt).TickAsync(default);
        Assert.Single(HandoffCalls);
    }
}