ClaudeDo/tests/ClaudeDo.Worker.Tests/Services/QueueServiceSlotFailureTests.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Git;
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.Data.Sqlite;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging.Abstractions;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Tests.Services;

// The queue picker's raw-SQL claim commits status='running' before the runner starts. If
// anything then throws before the runner's own terminal-status write, the task used to stay
// Running forever with the UI never notified (RunInSlotAsync's catch only logged the error).
// It must now mark the task Failed for a real exception (which broadcasts TaskUpdated), but
// must NOT do so for a cancellation — the cancel path already wrote the terminal status.
//
// These drive the real QueueService end to end (StartAsync + the waker), not just the
// FailAsync contract, so they actually exercise the fixed catch block.
public sealed class QueueServiceSlotFailureTests : IDisposable
{
    private readonly DbFixture _db = new();
    private readonly string _tempDir;
    private readonly WorkerConfig _cfg;

    public QueueServiceSlotFailureTests()
    {
        _tempDir = Path.Combine(Path.GetTempPath(), $"claudedo_slotfail_{Guid.NewGuid():N}");
        Directory.CreateDirectory(_tempDir);
        _cfg = new WorkerConfig
        {
            SandboxRoot = Path.Combine(_tempDir, "sandbox"),
            LogRoot = Path.Combine(_tempDir, "logs"),
            QueueBackstopIntervalMs = 50, // fast for tests
        };
    }

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

    // Mirrors QueueServiceTests.CreateService but takes the picker as a parameter so each test
    // can engineer the exact failure path it needs to exercise.
    // Build() wires its own CapturingHubContext internally and hands it back as .Hub — the
    // broadcaster inside TaskStateService (and therefore FailAsync's TaskUpdated) uses that
    // exact instance, so everything else here must share it too rather than constructing a
    // second CapturingHubContext that would silently miss FailAsync's broadcast.
    private (QueueService service, CapturingHubContext hub, QueueWaker waker) CreateService(IQueuePicker picker)
    {
        var dbFactory = _db.CreateFactory();
        var built = TaskStateServiceBuilder.Build(dbFactory);
        var broadcaster = new HubBroadcaster(built.Hub);
        var wtManager = new WorktreeManager(new GitService(), dbFactory, _cfg, NullLogger<WorktreeManager>.Instance);
        var argsBuilder = new ClaudeArgsBuilder();
        var runner = new TaskRunner(new FakeClaudeProcess(), dbFactory, broadcaster, wtManager, argsBuilder, _cfg,
            NullLogger<TaskRunner>.Instance, built.State, new TaskRunTokenRegistry(), new AttachmentStore(),
            new FakeSessionSkillSeeder(), new FakeTranscriptUsageReader());
        var waker = new QueueWaker();
        var overrideSlot = new OverrideSlotService(dbFactory, runner, NullLogger<OverrideSlotService>.Instance, built.RunCancels);
        var usageState = new UsageState();
        var service = new QueueService(dbFactory, runner, _cfg, NullLogger<QueueService>.Instance, waker, picker,
            overrideSlot, built.State, built.RunCancels, new FakeUsageGate(), usageState, broadcaster,
            new UsageLimitAutoContinueCoordinator(dbFactory, usageState, overrideSlot, broadcaster, NullLogger<UsageLimitAutoContinueCoordinator>.Instance));
        return (service, built.Hub, waker);
    }

    private async Task<string> SeedListAsync()
    {
        var listId = Guid.NewGuid().ToString();
        using var ctx = _db.CreateContext();
        ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", CreatedAt = DateTime.UtcNow });
        await ctx.SaveChangesAsync();
        return listId;
    }

    // Directly rewrites the task's list_id via a raw connection with FK enforcement off,
    // reproducing "the list vanished between the queue claim and the run" without going
    // through EF's foreign-key-checked connections (which would reject the write).
    private void OrphanTaskListId(string taskId)
    {
        using var conn = new SqliteConnection($"Data Source={_db.DbPath}");
        conn.Open();
        using (var pragmaCmd = conn.CreateCommand())
        {
            pragmaCmd.CommandText = "PRAGMA foreign_keys=OFF;";
            pragmaCmd.ExecuteNonQuery();
        }
        using var cmd = conn.CreateCommand();
        cmd.CommandText = "UPDATE tasks SET list_id = 'orphaned-missing-list' WHERE id = $id;";
        cmd.Parameters.AddWithValue("$id", taskId);
        cmd.ExecuteNonQuery();
    }

    [Fact]
    public async Task A_throwing_slot_run_marks_the_task_Failed_and_broadcasts_TaskUpdated()
    {
        var listId = await SeedListAsync();
        var taskId = Guid.NewGuid().ToString();

        using (var ctx = _db.CreateContext())
        {
            ctx.Tasks.Add(new TaskEntity
            {
                Id = taskId, ListId = listId, Title = "T", Status = TaskStatus.Queued,
                ReviewFeedback = "please fix", CreatedAt = DateTime.UtcNow,
            });
            await ctx.SaveChangesAsync();

            // A prior run with a session id routes RunInSlotAsync into TaskRunner.ContinueAsync
            // instead of RunAsync.
            await new TaskRunRepository(ctx).AddAsync(new TaskRunEntity
            {
                Id = Guid.NewGuid().ToString(), TaskId = taskId, RunNumber = 1, IsRetry = false,
                Prompt = "original", SessionId = "sess-1", StartedAt = DateTime.UtcNow.AddMinutes(-5),
            });
        }

        // ContinueAsync's setup block reads the list *before* its own try/catch starts
        // (TaskRunner.cs, ContinueAsync ~line 232-234) and throws InvalidOperationException
        // ("List not found.") straight past TaskRunner's own protection. That's the exact gap
        // QueueService.RunInSlotAsync's own catch now has to cover.
        OrphanTaskListId(taskId);

        var (service, hub, waker) = CreateService(new QueuePicker(_db.CreateFactory()));

        using var cts = new CancellationTokenSource();
        await service.StartAsync(cts.Token);
        waker.Wake();

        // FailAsync (TaskStateService.cs:236-249) commits the DB status flip via
        // ExecuteUpdateAsync *before* it calls the broadcaster's TaskUpdated — so a poll that
        // breaks the instant it observes Status==Failed can race ahead of the broadcast still
        // landing in hub.Proxy.Calls. Wait for both signals together so the assertions below
        // never sample a genuinely-not-yet-broadcast window as a failure.
        TaskEntity? reloaded = null;
        var deadline = DateTime.UtcNow.AddSeconds(10);
        while (DateTime.UtcNow < deadline)
        {
            using var verify = _db.CreateContext();
            reloaded = await new TaskRepository(verify).GetByIdAsync(taskId);
            var broadcastSeen = hub.Proxy.Calls.Any(
                c => c.Method == "TaskUpdated" && (string)c.Args[0]! == taskId);
            if (reloaded!.Status == TaskStatus.Failed && broadcastSeen) break;
            await Task.Delay(25);
        }
        cts.Cancel();

        Assert.Equal(TaskStatus.Failed, reloaded!.Status);
        Assert.Contains(hub.Proxy.Calls,
            c => c.Method == "TaskUpdated" && (string)c.Args[0]! == taskId);
    }

    // A fake IQueuePicker that performs the real atomic claim (so the DB row transitions
    // Queued->Running exactly like production) and then, synchronously before returning,
    // cancels the token QueueService's per-slot CTS is linked from. By the time
    // QueueService.ExecuteAsync creates that linked CTS and dispatches RunInSlotAsync, the
    // token is already cancelled — deterministic, no timing race required.
    private sealed class ClaimThenCancelPicker : IQueuePicker
    {
        private readonly IQueuePicker _inner;
        private readonly CancellationTokenSource _cancelAfterClaim;

        public ClaimThenCancelPicker(IQueuePicker inner, CancellationTokenSource cancelAfterClaim)
        {
            _inner = inner;
            _cancelAfterClaim = cancelAfterClaim;
        }

        public async Task<TaskEntity?> ClaimNextAsync(DateTime now, CancellationToken ct)
        {
            var claimed = await _inner.ClaimNextAsync(now, ct);
            if (claimed is not null) _cancelAfterClaim.Cancel();
            return claimed;
        }
    }

    [Fact]
    public async Task A_cancelled_slot_run_is_marked_Cancelled_not_left_Running()
    {
        var listId = await SeedListAsync();
        var taskId = Guid.NewGuid().ToString();

        using (var ctx = _db.CreateContext())
        {
            ctx.Tasks.Add(new TaskEntity
            {
                Id = taskId, ListId = listId, Title = "T", Status = TaskStatus.Queued,
                CreatedAt = DateTime.UtcNow,
            });
            await ctx.SaveChangesAsync();
        }

        var outerCts = new CancellationTokenSource();
        var realPicker = new QueuePicker(_db.CreateFactory());
        var picker = new ClaimThenCancelPicker(realPicker, outerCts);
        var (service, hub, waker) = CreateService(picker);

        await service.StartAsync(outerCts.Token);
        waker.Wake();

        var reloaded = await PollUntilLeftQueuedAsync(taskId);

        // This reproduces a plain Stop (QueueService.CancelTask cancels the slot's CTS
        // directly, without going through TaskStateService) landing during the pre-dispatch DB
        // reads — before TaskRunner's own claim/try-catch ever starts. Nothing else writes a
        // terminal status for that window, so the OCE catch itself must close it.
        Assert.Equal(TaskStatus.Cancelled, reloaded!.Status);
        Assert.Contains(hub.Proxy.Calls,
            c => c.Method == "TaskUpdated" && (string)c.Args[0]! == taskId);
    }

    // A_cancelled_slot_run_is_marked_Cancelled_not_left_Running above reproduces a plain Stop.
    // TaskStateService.CancelAsync (the hub's CancelReview action) is a different origin: it
    // writes Cancelled *before* it cancels the run's CTS via RunCancellationRegistry, so by the
    // time RunInSlotAsync's OCE catch runs, the task has already left Running. The catch must
    // recognize that and not blindly stomp whatever terminal status is already there.
    [Fact]
    public async Task A_cancelled_slot_run_does_not_stomp_a_status_already_written_terminal()
    {
        var listId = await SeedListAsync();
        var taskId = Guid.NewGuid().ToString();

        using (var ctx = _db.CreateContext())
        {
            ctx.Tasks.Add(new TaskEntity
            {
                Id = taskId, ListId = listId, Title = "T", Status = TaskStatus.Queued,
                CreatedAt = DateTime.UtcNow,
            });
            await ctx.SaveChangesAsync();
        }

        var outerCts = new CancellationTokenSource();
        var realPicker = new QueuePicker(_db.CreateFactory());
        var picker = new ClaimMarkFailedThenCancelPicker(realPicker, outerCts, _db);
        var (service, hub, waker) = CreateService(picker);

        await service.StartAsync(outerCts.Token);
        waker.Wake();

        var reloaded = await PollUntilLeftQueuedAsync(taskId);

        // Some other terminal write (simulating CancelAsync/FailAsync having already run) must
        // survive — the OCE catch must not overwrite it with Cancelled just because it observed
        // a cancellation.
        Assert.Equal(TaskStatus.Failed, reloaded!.Status);
    }

    // Polls until the task leaves its initial Queued status (the claim + dispatch + OCE-catch
    // chain settling), or gives up at the deadline; then waits a further grace period so the
    // fire-and-forget continuation (queue-slot cleanup, any terminal-status write) fully lands
    // before the caller reads the final state.
    private async Task<TaskEntity?> PollUntilLeftQueuedAsync(string taskId)
    {
        TaskEntity? reloaded = null;
        var deadline = DateTime.UtcNow.AddSeconds(10);
        while (DateTime.UtcNow < deadline)
        {
            using var verify = _db.CreateContext();
            reloaded = await new TaskRepository(verify).GetByIdAsync(taskId);
            if (reloaded!.Status != TaskStatus.Queued) break;
            await Task.Delay(25);
        }
        await Task.Delay(150); // let the fire-and-forget continuation fully settle

        using var final = _db.CreateContext();
        return await new TaskRepository(final).GetByIdAsync(taskId);
    }

    // Simulates a terminal status already having been written (by TaskStateService, from some
    // other origin) between the picker's claim and RunInSlotAsync's OCE catch running.
    private sealed class ClaimMarkFailedThenCancelPicker : IQueuePicker
    {
        private readonly IQueuePicker _inner;
        private readonly CancellationTokenSource _cancelAfterClaim;
        private readonly DbFixture _db;

        public ClaimMarkFailedThenCancelPicker(IQueuePicker inner, CancellationTokenSource cancelAfterClaim, DbFixture db)
        {
            _inner = inner;
            _cancelAfterClaim = cancelAfterClaim;
            _db = db;
        }

        public async Task<TaskEntity?> ClaimNextAsync(DateTime now, CancellationToken ct)
        {
            var claimed = await _inner.ClaimNextAsync(now, ct);
            if (claimed is not null)
            {
                using (var ctx = _db.CreateContext())
                {
                    await ctx.Tasks.Where(t => t.Id == claimed.Id)
                        .ExecuteUpdateAsync(s => s
                            .SetProperty(t => t.Status, TaskStatus.Failed)
                            .SetProperty(t => t.FinishedAt, DateTime.UtcNow));
                }
                _cancelAfterClaim.Cancel();
            }
            return claimed;
        }
    }
}