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.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 service = new QueueService(dbFactory, runner, _cfg, NullLogger<QueueService>.Instance, waker, picker,
            overrideSlot, built.State, built.RunCancels, new FakeUsageGate(), new UsageState(), broadcaster);
        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();

        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.Failed) 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_does_not_mark_the_task_Failed()
    {
        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();

        // Wait for the slot to be claimed and then released again (RunInSlotAsync's
        // ContinueWith removes it once the catch block — ours or a stray one — finishes).
        var deadline = DateTime.UtcNow.AddSeconds(10);
        while (service.GetActive().Any(a => a.taskId == taskId) && DateTime.UtcNow < deadline)
            await Task.Delay(25);
        await Task.Delay(100); // let the fire-and-forget continuation fully settle

        TaskEntity? reloaded;
        using (var verify = _db.CreateContext())
            reloaded = await new TaskRepository(verify).GetByIdAsync(taskId);

        // The picker's atomic claim already flipped it to Running; the cancelled slot run must
        // leave it there rather than flipping it to Failed.
        Assert.Equal(TaskStatus.Running, reloaded!.Status);
        Assert.DoesNotContain(hub.Proxy.Calls,
            c => c.Method == "TaskUpdated" && (string)c.Args[0]! == taskId);
    }
}