ClaudeDo/tests/ClaudeDo.Worker.Tests/Services/QueueServiceSlotGuardTests.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.Extensions.Logging.Abstractions;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Tests.Services;

public sealed class QueueServiceSlotGuardTests : IDisposable
{
    private readonly DbFixture _db = new();
    private readonly ClaudeDoDbContext _ctx;
    private readonly TaskRepository _taskRepo;
    private readonly ListRepository _listRepo;
    private readonly WorkerConfig _cfg;
    private readonly string _tempDir;

    public QueueServiceSlotGuardTests()
    {
        _ctx = _db.CreateContext();
        _taskRepo = new TaskRepository(_ctx);
        _listRepo = new ListRepository(_ctx);
        _tempDir = Path.Combine(Path.GetTempPath(), $"claudedo_slotguard_{Guid.NewGuid():N}");
        Directory.CreateDirectory(_tempDir);
        _cfg = new WorkerConfig
        {
            SandboxRoot = Path.Combine(_tempDir, "sandbox"),
            LogRoot = Path.Combine(_tempDir, "logs"),
            QueueBackstopIntervalMs = 50,
        };
    }

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

    private QueueWaker _waker = null!;

    private (QueueService service, FakeClaudeProcess fakeProcess) CreateService(
        Func<string, string, IReadOnlyList<string>, Func<string, Task>, CancellationToken, Task<RunResult>>? handler = null)
    {
        var fake = new FakeClaudeProcess(handler);
        var broadcaster = new HubBroadcaster(new CapturingHubContext());
        var dbFactory = _db.CreateFactory();
        var wtManager = new WorktreeManager(new ClaudeDo.Data.Git.GitService(), dbFactory, _cfg, NullLogger<WorktreeManager>.Instance);
        var argsBuilder = new ClaudeArgsBuilder();
        var built = TaskStateServiceBuilder.Build(dbFactory);
        var state = built.State;
        var runner = new TaskRunner(fake, dbFactory, broadcaster, wtManager, argsBuilder, _cfg,
            NullLogger<TaskRunner>.Instance, state, new TaskRunTokenRegistry(), new AttachmentStore(), new FakeSessionSkillSeeder(), new FakeTranscriptUsageReader());
        _waker = new QueueWaker();
        var picker = new QueuePicker(dbFactory);
        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, state, built.RunCancels,
            new FakeUsageGate(), usageState, broadcaster,
            new UsageLimitAutoContinueCoordinator(dbFactory, usageState, overrideSlot, broadcaster, NullLogger<UsageLimitAutoContinueCoordinator>.Instance));
        return (service, fake);
    }

    private async Task<string> SeedListAsync()
    {
        var listId = Guid.NewGuid().ToString();
        await _listRepo.AddAsync(new ListEntity { Id = listId, Name = "Test", CreatedAt = DateTime.UtcNow });
        return listId;
    }

    private async Task<TaskEntity> SeedQueuedTaskAsync(string listId)
    {
        var task = new TaskEntity
        {
            Id = Guid.NewGuid().ToString(),
            ListId = listId,
            Title = "Guard test task",
            Description = "Test",
            Status = TaskStatus.Queued,
            CreatedAt = DateTime.UtcNow,
        };
        await _taskRepo.AddAsync(task);
        return task;
    }

    [Fact]
    public async Task RunNow_Throws_When_Task_Already_Running_In_Queue_Slot()
    {
        var listId = await SeedListAsync();
        var task = await SeedQueuedTaskAsync(listId);

        // Gate keeps the queue slot occupied indefinitely.
        var tcs = new TaskCompletionSource<RunResult>();
        var queuePickedUp = new TaskCompletionSource();

        var (service, _) = CreateService(async (_, _, _, _, ct) =>
        {
            queuePickedUp.TrySetResult();
            return await tcs.Task;
        });

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

        // Wait until the queue slot has actually picked up the task.
        await queuePickedUp.Task.WaitAsync(TimeSpan.FromSeconds(5));

        // Now the same taskId is in the queue slot — RunNow must reject it.
        var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => service.RunNow(task.Id));
        Assert.Contains("already running", ex.Message);

        tcs.SetResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok" });
        cts.Cancel();
    }

    [Fact]
    public async Task ContinueTask_Throws_When_Task_Already_Running_In_Queue_Slot()
    {
        var listId = await SeedListAsync();
        var task = await SeedQueuedTaskAsync(listId);

        var tcs = new TaskCompletionSource<RunResult>();
        var queuePickedUp = new TaskCompletionSource();

        var (service, _) = CreateService(async (_, _, _, _, ct) =>
        {
            queuePickedUp.TrySetResult();
            return await tcs.Task;
        });

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

        await queuePickedUp.Task.WaitAsync(TimeSpan.FromSeconds(5));

        var ex = await Assert.ThrowsAsync<InvalidOperationException>(() =>
            service.ContinueTask(task.Id, "follow-up"));
        Assert.Contains("already running", ex.Message);

        tcs.SetResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok" });
        cts.Cancel();
    }
}