ClaudeDo/tests/ClaudeDo.Worker.Tests/Queue/UsageLimitAutoContinueCoordinatorTests.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Config;
using ClaudeDo.Worker.Hub;
using ClaudeDo.Worker.Queue;
using ClaudeDo.Worker.Runner;
using ClaudeDo.Worker.Tests.Infrastructure;
using ClaudeDo.Worker.Usage;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging.Abstractions;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Tests.Queue;

/// Covers the "Continue on session limit reset" toggle's coordinator in isolation from
/// QueueService's own timer plumbing (see QueueService.ScheduleResetWake for that part).
public sealed class UsageLimitAutoContinueCoordinatorTests : IDisposable
{
    private readonly DbFixture _db = new();
    private readonly string _tempDir;
    private readonly WorkerConfig _cfg;

    public UsageLimitAutoContinueCoordinatorTests()
    {
        _tempDir = Path.Combine(Path.GetTempPath(), $"cd_usagelimit_{Guid.NewGuid():N}");
        Directory.CreateDirectory(_tempDir);
        _cfg = new WorkerConfig { SandboxRoot = _tempDir, LogRoot = _tempDir };
    }

    public void Dispose() { _db.Dispose(); try { Directory.Delete(_tempDir, true); } catch { } }

    private (UsageLimitAutoContinueCoordinator Coordinator, UsageState UsageState, FakeClaudeProcess Process, OverrideSlotService Slot) BuildCoordinator(
        Func<string, string, IReadOnlyList<string>, Func<string, Task>, CancellationToken, Task<RunResult>>? handler = null)
    {
        var dbFactory = _db.CreateFactory();
        var state = TaskStateServiceBuilder.Build(dbFactory).State;
        var wt = new WorktreeManager(new ClaudeDo.Data.Git.GitService(), dbFactory, _cfg, NullLogger<WorktreeManager>.Instance);
        var fake = new FakeClaudeProcess(handler);
        var runner = new TaskRunner(fake, dbFactory, new HubBroadcaster(new CapturingHubContext()), wt,
            new ClaudeArgsBuilder(), _cfg, NullLogger<TaskRunner>.Instance, state, new TaskRunTokenRegistry(),
            new AttachmentStore(), new FakeSessionSkillSeeder(), new FakeTranscriptUsageReader());
        var runCancels = new RunCancellationRegistry(NullLogger<RunCancellationRegistry>.Instance);
        var overrideSlot = new OverrideSlotService(dbFactory, runner, NullLogger<OverrideSlotService>.Instance, runCancels);
        var usageState = new UsageState();
        var coordinator = new UsageLimitAutoContinueCoordinator(
            dbFactory, usageState, overrideSlot,
            new HubBroadcaster(new CapturingHubContext()), NullLogger<UsageLimitAutoContinueCoordinator>.Instance);
        return (coordinator, usageState, fake, overrideSlot);
    }

    // Raw inserts bypass TaskNumberAllocator, so hand out distinct Numbers — tasks.number is unique.
    private int _nextNumber = 1;

    private async Task<string> SeedUsageLimitTaskAsync(bool withSessionId = 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", CreatedAt = DateTime.UtcNow });
            ctx.Tasks.Add(new TaskEntity
            {
                Id = taskId, ListId = listId, Title = "T", Status = TaskStatus.Failed, Number = _nextNumber++,
                FailureReason = "usage_limit", CreatedAt = DateTime.UtcNow, FinishedAt = DateTime.UtcNow,
            });
            if (withSessionId)
            {
                ctx.TaskRuns.Add(new TaskRunEntity
                {
                    Id = Guid.NewGuid().ToString(), TaskId = taskId, RunNumber = 1, IsRetry = false, Prompt = "p",
                    LogPath = Path.Combine(_tempDir, "log.ndjson"), StartedAt = DateTime.UtcNow,
                    FinishedAt = DateTime.UtcNow, SessionId = "sess-1", ExitCode = 1,
                });
            }
            await ctx.SaveChangesAsync();
        }
        return taskId;
    }

    private async Task EnableToggleAsync()
    {
        using var ctx = _db.CreateContext();
        var repo = new AppSettingsRepository(ctx);
        var settings = await repo.GetAsync();
        settings.AutoContinueOnUsageLimit = true;
        await repo.UpdateAsync(settings);
    }

    private async Task<TaskEntity?> PollUntilCalledAsync(string taskId, FakeClaudeProcess fake)
    {
        var deadline = DateTime.UtcNow.AddSeconds(10);
        while (DateTime.UtcNow < deadline && fake.CallCount == 0)
            await Task.Delay(25);
        await Task.Delay(150); // let the fire-and-forget continuation fully settle
        using var ctx = _db.CreateContext();
        return await new TaskRepository(ctx).GetByIdAsync(taskId);
    }

    [Fact]
    public async Task Toggle_Off_Never_Schedules_Or_Continues()
    {
        var (coordinator, usageState, fake, _) = BuildCoordinator();
        var taskId = await SeedUsageLimitTaskAsync();
        usageState.ReportSuccess(new UsageSnapshot(
            new UsageBucket(85, DateTimeOffset.UtcNow.AddSeconds(-1)), null, Array.Empty<UsageLimitRow>(), DateTime.UtcNow));

        var wakeAt = await coordinator.GetScheduledWakeAtAsync(default);
        await coordinator.RunAsync(default);
        await Task.Delay(150);

        Assert.Null(wakeAt);
        Assert.Equal(0, fake.CallCount);
        using var verify = _db.CreateContext();
        var task = await new TaskRepository(verify).GetByIdAsync(taskId);
        Assert.Null(task!.UsageLimitAutoContinuedAt);
    }

    [Fact]
    public async Task Toggle_On_No_Snapshot_Yet_Schedules_Nothing_And_Does_Not_Crash()
    {
        var (coordinator, _, fake, _) = BuildCoordinator();
        await SeedUsageLimitTaskAsync();
        await EnableToggleAsync();

        var wakeAt = await coordinator.GetScheduledWakeAtAsync(default);
        await coordinator.RunAsync(default); // must not throw despite no usage snapshot

        Assert.Null(wakeAt);
        Assert.Equal(0, fake.CallCount);
    }

    [Fact]
    public async Task Toggle_On_Reset_In_Future_Schedules_Wake_But_Does_Not_Continue_Yet()
    {
        var (coordinator, usageState, fake, _) = BuildCoordinator();
        await SeedUsageLimitTaskAsync();
        await EnableToggleAsync();
        var resetsAt = DateTimeOffset.UtcNow.AddMinutes(5);
        usageState.ReportSuccess(new UsageSnapshot(
            new UsageBucket(85, resetsAt), null, Array.Empty<UsageLimitRow>(), DateTime.UtcNow));

        var wakeAt = await coordinator.GetScheduledWakeAtAsync(default);
        await coordinator.RunAsync(default);
        await Task.Delay(150);

        Assert.Equal(resetsAt, wakeAt);
        Assert.Equal(0, fake.CallCount);
    }

    [Fact]
    public async Task Toggle_On_Reset_Already_Passed_Continues_Exactly_Once()
    {
        var (coordinator, usageState, fake, _) = BuildCoordinator();
        var taskId = await SeedUsageLimitTaskAsync();
        await EnableToggleAsync();
        usageState.ReportSuccess(new UsageSnapshot(
            new UsageBucket(85, DateTimeOffset.UtcNow.AddSeconds(-1)), null, Array.Empty<UsageLimitRow>(), DateTime.UtcNow));

        await coordinator.RunAsync(default);
        var task = await PollUntilCalledAsync(taskId, fake);

        Assert.Equal(1, fake.CallCount);
        Assert.NotNull(task!.UsageLimitAutoContinuedAt);

        // A second tick (e.g. the next backstop) must not fire it again.
        await coordinator.RunAsync(default);
        await Task.Delay(150);
        Assert.Equal(1, fake.CallCount);
    }

    /// A usage limit stops every running task at once, so "more candidates than override slots"
    /// is the normal case, not an edge case. The claim has to be taken before the run starts
    /// (it is the dedupe guard), so a candidate that never got a slot must not keep it — the
    /// candidate query filters on that marker and would exclude the task forever.
    [Fact]
    public async Task Busy_Override_Slot_Leaves_The_Other_Candidate_For_A_Later_Tick()
    {
        var gate = new TaskCompletionSource();
        var (coordinator, usageState, fake, slot) = BuildCoordinator(async (_, _, _, _, _) =>
        {
            await gate.Task;
            return new RunResult { ExitCode = 0, ResultMarkdown = "ok" };
        });
        await SeedUsageLimitTaskAsync();
        await SeedUsageLimitTaskAsync();
        await EnableToggleAsync();
        usageState.ReportSuccess(new UsageSnapshot(
            new UsageBucket(85, DateTimeOffset.UtcNow.AddSeconds(-1)), null, Array.Empty<UsageLimitRow>(), DateTime.UtcNow));

        // Tick 1 — the first candidate takes the single slot and holds it on the gate.
        await coordinator.RunAsync(default);
        Assert.Equal(1, await CountClaimedAsync());

        // Tick 2 while that slot is still busy — the second candidate must stay unclaimed.
        await coordinator.RunAsync(default);
        Assert.Equal(1, await CountClaimedAsync());

        // Let the first run finish, freeing the slot.
        await PollAsync(() => fake.CallCount >= 1);
        gate.SetResult();
        await PollAsync(() => slot.CurrentSlot is null);

        // Tick 3 — the leftover candidate finally gets its turn.
        await coordinator.RunAsync(default);
        await PollAsync(() => fake.CallCount >= 2);
        Assert.Equal(2, fake.CallCount);
        Assert.Equal(2, await CountClaimedAsync());
    }

    private async Task<int> CountClaimedAsync()
    {
        using var ctx = _db.CreateContext();
        return await ctx.Tasks.CountAsync(t => t.UsageLimitAutoContinuedAt != null);
    }

    private static async Task PollAsync(Func<bool> condition)
    {
        var deadline = DateTime.UtcNow.AddSeconds(10);
        while (DateTime.UtcNow < deadline && !condition())
            await Task.Delay(25);
        Assert.True(condition(), "condition did not become true within 10s");
    }

    [Fact]
    public async Task Missing_Reset_Time_Never_Continues_And_Never_Throws()
    {
        var (coordinator, usageState, fake, _) = BuildCoordinator();
        await SeedUsageLimitTaskAsync();
        await EnableToggleAsync();
        // A snapshot exists, but this bucket has no reset time — fail open.
        usageState.ReportSuccess(new UsageSnapshot(
            new UsageBucket(85, null), null, Array.Empty<UsageLimitRow>(), DateTime.UtcNow));

        var wakeAt = await coordinator.GetScheduledWakeAtAsync(default);
        await coordinator.RunAsync(default);
        await Task.Delay(150);

        Assert.Null(wakeAt);
        Assert.Equal(0, fake.CallCount);
    }

    [Fact]
    public async Task Non_UsageLimit_Failure_Is_Never_Touched()
    {
        var (coordinator, usageState, fake, _) = BuildCoordinator();
        var listId = Guid.NewGuid().ToString();
        var taskId = Guid.NewGuid().ToString();
        using (var ctx = _db.CreateContext())
        {
            ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", CreatedAt = DateTime.UtcNow });
            ctx.Tasks.Add(new TaskEntity { Id = taskId, ListId = listId, Title = "T", Status = TaskStatus.Failed,
                FailureReason = "error", CreatedAt = DateTime.UtcNow });
            await ctx.SaveChangesAsync();
        }
        await EnableToggleAsync();
        usageState.ReportSuccess(new UsageSnapshot(
            new UsageBucket(85, DateTimeOffset.UtcNow.AddSeconds(-1)), null, Array.Empty<UsageLimitRow>(), DateTime.UtcNow));

        await coordinator.RunAsync(default);
        await Task.Delay(150);

        Assert.Equal(0, fake.CallCount);
    }
}