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, Func, CancellationToken, Task>? handler = null) { var dbFactory = _db.CreateFactory(); var state = TaskStateServiceBuilder.Build(dbFactory).State; var wt = new WorktreeManager(new ClaudeDo.Data.Git.GitService(), dbFactory, _cfg, NullLogger.Instance); var fake = new FakeClaudeProcess(handler); var runner = new TaskRunner(fake, dbFactory, new HubBroadcaster(new CapturingHubContext()), wt, new ClaudeArgsBuilder(), _cfg, NullLogger.Instance, state, new TaskRunTokenRegistry(), new AttachmentStore(), new FakeSessionSkillSeeder(), new FakeTranscriptUsageReader()); var runCancels = new RunCancellationRegistry(NullLogger.Instance); var overrideSlot = new OverrideSlotService(dbFactory, runner, NullLogger.Instance, runCancels); var usageState = new UsageState(); var coordinator = new UsageLimitAutoContinueCoordinator( dbFactory, usageState, overrideSlot, new HubBroadcaster(new CapturingHubContext()), NullLogger.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 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 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(), 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(), 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(), 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(), 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 CountClaimedAsync() { using var ctx = _db.CreateContext(); return await ctx.Tasks.CountAsync(t => t.UsageLimitAutoContinuedAt != null); } private static async Task PollAsync(Func 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(), 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(), DateTime.UtcNow)); await coordinator.RunAsync(default); await Task.Delay(150); Assert.Equal(0, fake.CallCount); } }