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.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) BuildCoordinator() { 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(); 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); } 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, 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); } [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); } }