ClaudeDo/src/ClaudeDo.Worker/Queue/QueueService.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Config;
using ClaudeDo.Worker.Hub;
using ClaudeDo.Worker.Runner;
using ClaudeDo.Worker.State;
using ClaudeDo.Worker.Usage;
using ClaudeDo.Worker.Usage.Interfaces;
using Microsoft.EntityFrameworkCore;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Queue;

public sealed class QueueService : BackgroundService
{
    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;
    private readonly TaskRunner _runner;
    private readonly WorkerConfig _cfg;
    private readonly ILogger<QueueService> _logger;
    private readonly QueueWaker _waker;
    private readonly IQueuePicker _picker;
    private readonly OverrideSlotService _override;
    private readonly ITaskStateService _state;
    private readonly RunCancellationRegistry _runCancels;
    private readonly IUsageGate _usageGate;
    private readonly UsageState _usageState;
    private readonly HubBroadcaster _broadcaster;

    private readonly object _lock = new();
    private readonly Dictionary<string, QueueSlotState> _queueSlots = new();
    private bool _usageGateBlocked;

    public QueueService(
        IDbContextFactory<ClaudeDoDbContext> dbFactory,
        TaskRunner runner,
        WorkerConfig cfg,
        ILogger<QueueService> logger,
        QueueWaker waker,
        IQueuePicker picker,
        OverrideSlotService overrideSlot,
        ITaskStateService state,
        RunCancellationRegistry runCancels,
        IUsageGate usageGate,
        UsageState usageState,
        HubBroadcaster broadcaster)
    {
        _dbFactory = dbFactory;
        _runner = runner;
        _cfg = cfg;
        _logger = logger;
        _waker = waker;
        _picker = picker;
        _override = overrideSlot;
        _state = state;
        _runCancels = runCancels;
        _usageGate = usageGate;
        _usageState = usageState;
        _broadcaster = broadcaster;
    }

    public IReadOnlyList<(string slot, string taskId, DateTime startedAt)> GetActive()
    {
        var list = new List<(string, string, DateTime)>();
        lock (_lock)
        {
            foreach (var slot in _queueSlots.Values)
                list.Add(("queue", slot.TaskId, slot.StartedAt));
        }
        var o = _override.CurrentSlot;
        if (o is not null) list.Add(("override", o.TaskId, o.StartedAt));
        return list;
    }

    public Task RunNow(string taskId)
    {
        EnsureNotInQueueSlot(taskId);
        return _override.RunNow(taskId);
    }

    public Task<string> ContinueTask(string taskId, string followUpPrompt)
    {
        EnsureNotInQueueSlot(taskId);
        return _override.ContinueTask(taskId, followUpPrompt);
    }

    private void EnsureNotInQueueSlot(string taskId)
    {
        lock (_lock)
        {
            if (_queueSlots.ContainsKey(taskId))
                throw new InvalidOperationException("task is already running in queue slot");
        }
    }

    public bool CancelTask(string taskId)
    {
        if (_override.TryCancel(taskId)) return true;

        lock (_lock)
        {
            if (_queueSlots.TryGetValue(taskId, out var slot))
            {
                slot.Cts.Cancel();
                return true;
            }
        }
        return false;
    }

    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
    {
        _logger.LogInformation("QueueService started");

        using var timer = new PeriodicTimer(TimeSpan.FromMilliseconds(_cfg.QueueBackstopIntervalMs));

        while (!stoppingToken.IsCancellationRequested)
        {
            try
            {
                // Wait for wake signal or backstop timer.
                var wakeTask = _waker.WaitAsync(stoppingToken);
                var timerTask = timer.WaitForNextTickAsync(stoppingToken).AsTask();

                await Task.WhenAny(wakeTask, timerTask);

                var maxParallel = await GetMaxParallelAsync(stoppingToken);

                var gateDecision = await _usageGate.EvaluateAsync(stoppingToken);
                await ReportUsageGateTransitionAsync(gateDecision);

                // Only queue refill is gated. Runs already in flight (RunNow, ContinueTask,
                // interactive sessions, planning, daily prep) keep going regardless.
                if (!gateDecision.IsBlocked)
                {
                    // Fill as many free slots as the limit allows.
                    while (!stoppingToken.IsCancellationRequested)
                    {
                        lock (_lock)
                        {
                            if (_queueSlots.Count >= maxParallel) break;
                        }

                        var task = await _picker.ClaimNextAsync(DateTime.UtcNow, stoppingToken);
                        if (task is null) break;

                        lock (_lock)
                        {
                            var cts = CancellationTokenSource.CreateLinkedTokenSource(stoppingToken);
                            _queueSlots[task.Id] = new QueueSlotState { TaskId = task.Id, StartedAt = DateTime.UtcNow, Cts = cts };
                            _runCancels.Register(task.Id, cts);

                            _ = RunInSlotAsync(task.Id, cts.Token).ContinueWith(t =>
                            {
                                if (t.IsFaulted)
                                    _logger.LogError(t.Exception, "RunInSlotAsync failed for task {TaskId} in queue slot", task.Id);
                                lock (_lock) { _queueSlots.Remove(task.Id); }
                                _runCancels.Unregister(task.Id, cts);
                                cts.Dispose();
                                _waker.Wake(); // Check for next task immediately.
                            }, TaskScheduler.Default);
                        }
                    }
                }
            }
            catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
            {
                break;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "QueueService loop error");
            }
        }

        _logger.LogInformation("QueueService stopping");
    }

    private async Task ReportUsageGateTransitionAsync(UsageGateDecision decision)
    {
        if (decision.IsBlocked == _usageGateBlocked) return;
        _usageGateBlocked = decision.IsBlocked;

        if (decision.IsBlocked)
        {
            _logger.LogInformation("QueueService: usage gate blocking queue refill ({Reason})", decision.Reason);
            await _broadcaster.WorkerLog($"Queue pausiert: {decision.Reason}", WorkerLogLevel.Warn, DateTime.UtcNow);
        }
        else
        {
            var snapshot = _usageState.Snapshot;
            var message = snapshot?.FiveHour is not null && snapshot.SevenDay is not null
                ? $"Queue fortgesetzt: 5h {snapshot.FiveHour.Utilization:0}%, 7d {snapshot.SevenDay.Utilization:0}%"
                : "Queue fortgesetzt";

            _logger.LogInformation("QueueService: usage gate cleared, queue refill resumed");
            await _broadcaster.WorkerLog(message, WorkerLogLevel.Info, DateTime.UtcNow);
        }
    }

    private async Task<int> GetMaxParallelAsync(CancellationToken ct)
    {
        try
        {
            using var context = _dbFactory.CreateDbContext();
            var settings = await new AppSettingsRepository(context).GetAsync(ct);
            return Math.Max(1, settings.MaxParallelExecutions);
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "Failed to read max parallel executions; defaulting to 1");
            return 1;
        }
    }

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

            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.");
            }

            // A task re-queued from review carries reviewer feedback. Resume the prior
            // Claude session with that feedback as the next turn when a session exists;
            // otherwise fall back to a fresh run with the feedback folded into the prompt.
            if (!string.IsNullOrWhiteSpace(task.ReviewFeedback))
            {
                var feedback = task.ReviewFeedback!;
                string? sessionId;
                using (var context = _dbFactory.CreateDbContext())
                    sessionId = (await new TaskRunRepository(context).GetLatestByTaskIdAsync(taskId, ct))?.SessionId;

                if (sessionId is not null)
                {
                    await _runner.ContinueAsync(taskId, feedback, "queue", ct, alreadyClaimed: true);
                }
                else
                {
                    task.Description = string.IsNullOrWhiteSpace(task.Description)
                        ? $"Reviewer feedback: {feedback}"
                        : $"{task.Description}\n\nReviewer feedback: {feedback}";
                    await _runner.RunAsync(task, "queue", ct, alreadyClaimed: true);
                }

                // Clear the consumed feedback only once the run reached a successful
                // terminal state, so a failed or cancelled run keeps it for a manual retry.
                TaskStatus statusAfter;
                using (var context = _dbFactory.CreateDbContext())
                    statusAfter = await context.Tasks.Where(t => t.Id == taskId)
                        .Select(t => t.Status).FirstAsync(CancellationToken.None);
                if (statusAfter is TaskStatus.WaitingForReview or TaskStatus.Done)
                    await _state.ClearReviewFeedbackAsync(taskId, CancellationToken.None);
                return;
            }

            await _runner.RunAsync(task, "queue", ct, alreadyClaimed: true);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Slot runner error for task {TaskId}", taskId);
        }
    }
}