ClaudeDo/src/ClaudeDo.Worker/Queue/OverrideSlotService.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Runner;
using Microsoft.EntityFrameworkCore;

namespace ClaudeDo.Worker.Queue;

public sealed class OverrideSlotService
{
    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;
    private readonly TaskRunner _runner;
    private readonly ILogger<OverrideSlotService> _logger;
    private readonly RunCancellationRegistry _runCancels;

    private readonly object _lock = new();
    private volatile QueueSlotState? _slot;

    public OverrideSlotService(
        IDbContextFactory<ClaudeDoDbContext> dbFactory,
        TaskRunner runner,
        ILogger<OverrideSlotService> logger,
        RunCancellationRegistry runCancels)
    {
        _dbFactory = dbFactory;
        _runner = runner;
        _logger = logger;
        _runCancels = runCancels;
    }

    public QueueSlotState? CurrentSlot => _slot;

    public async Task RunNow(string taskId)
    {
        using (var context = _dbFactory.CreateDbContext())
        {
            var task = await new TaskRepository(context).GetByIdAsync(taskId);
            if (task is null)
                throw new KeyNotFoundException($"Task '{taskId}' not found.");

            // Fast-fail precheck only — not the race guard. A task sitting Queued still
            // passes here and can race the queue picker's atomic claim; TaskRunner.RunAsync
            // resolves that by claiming Running before creating any resources.
            if (task.Status == Data.Models.TaskStatus.Running)
                throw new InvalidOperationException("task is already running");
        }

        StartInSlot(taskId, ct => RunInSlotAsync(taskId, ct), "RunInSlotAsync failed for task {TaskId}");
    }

    public async Task<string> ContinueTask(string taskId, string followUpPrompt)
    {
        using (var context = _dbFactory.CreateDbContext())
        {
            var task = await new TaskRepository(context).GetByIdAsync(taskId)
                ?? throw new KeyNotFoundException($"Task '{taskId}' not found.");

            if (task.Status == Data.Models.TaskStatus.Running)
                throw new InvalidOperationException("task is already running");
        }

        StartInSlot(taskId, ct => RunContinueInSlotAsync(taskId, followUpPrompt, ct),
            "RunContinueInSlotAsync failed for task {TaskId}");

        return taskId;
    }

    // Claims the single override slot under lock, runs <work> in the background,
    // and releases the slot when it completes. Throws if the slot is already busy.
    private void StartInSlot(string taskId, Func<CancellationToken, Task> work, string faultMessage)
    {
        lock (_lock)
        {
            if (_slot is not null)
                throw new InvalidOperationException("override slot busy");

            var cts = new CancellationTokenSource();
            if (!_runCancels.Register(taskId, cts))
            {
                cts.Dispose();
                throw new InvalidOperationException(
                    $"Task '{taskId}' lost the double-dispatch race to the queue picker; it is already running there.");
            }

            _slot = new QueueSlotState { TaskId = taskId, StartedAt = DateTime.UtcNow, Cts = cts };

            _ = work(cts.Token).ContinueWith(t =>
            {
                if (t.IsFaulted)
                    _logger.LogError(t.Exception, faultMessage, taskId);
                lock (_lock) { _slot = null; }
                _runCancels.Unregister(taskId, cts);
                cts.Dispose();
            }, TaskScheduler.Default);
        }
    }

    public bool TryCancel(string taskId)
    {
        lock (_lock)
        {
            if (_slot is not null && _slot.TaskId == taskId)
            {
                _slot.Cts.Cancel();
                return true;
            }
        }
        return false;
    }

    private async Task RunInSlotAsync(string taskId, CancellationToken ct)
    {
        try
        {
            _logger.LogInformation("Starting task {TaskId} in override slot", taskId);

            Data.Models.TaskEntity task;
            using (var context = _dbFactory.CreateDbContext())
            {
                var taskRepo = new TaskRepository(context);
                task = await taskRepo.GetByIdAsync(taskId, ct)
                    ?? throw new KeyNotFoundException($"Task '{taskId}' not found.");
            }

            await _runner.RunAsync(task, "override", ct);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Override slot runner error for task {TaskId}", taskId);
        }
    }

    private async Task RunContinueInSlotAsync(string taskId, string followUpPrompt, CancellationToken ct)
    {
        try
        {
            _logger.LogInformation("Continuing task {TaskId} in override slot", taskId);
            await _runner.ContinueAsync(taskId, followUpPrompt, "override", ct);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Continue runner error for task {TaskId}", taskId);
        }
    }
}