jobtrackingapp/JobTrackerApi/Services/AiOperationQueue.cs

using JobTrackerApi.Data;
using JobTrackerApi.Models;
using Microsoft.AspNetCore.Identity;
using Microsoft.EntityFrameworkCore;

namespace JobTrackerApi.Services;

public static class AiOperationPriorities
{
    public const int Interactive = 500;
    public const int UserVisible = 400;
    public const int Drafting = 300;
    public const int Scheduled = 200;
    public const int Bulk = 100;
}

public sealed record AiOperationAdmissionResult(UserOperation Operation, bool Created, string StatusUrl);

public sealed class AiOperationAdmissionException(string code, string message, int statusCode, int? retryAfterSeconds = null)
    : Exception(message)
{
    public string Code { get; } = code;
    public int StatusCode { get; } = statusCode;
    public int? RetryAfterSeconds { get; } = retryAfterSeconds;
}

public sealed class AiOperationAdmission(
    UserOperationStore operations,
    JobTrackerContext db,
    ICurrentUserService currentUser,
    UserManager<ApplicationUser> users,
    AiPrivacyPolicy privacy,
    IConfiguration configuration,
    TimeProvider timeProvider,
    AiUsageMeter usage)
{
    // ponytail: process-local gate is sufficient for the current single-backend deployment;
    // replace with a database capacity reservation before running multiple backend replicas.
    private static readonly SemaphoreSlim AdmissionGate = new(1, 1);

    public async Task<AiOperationAdmissionResult> EnqueueAsync(
        string taskType,
        string idempotencyKey,
        string subjectType,
        string subjectId,
        int priority,
        CancellationToken cancellationToken)
    {
        var userId = currentUser.UserId;
        var user = string.IsNullOrWhiteSpace(userId) ? null : await users.FindByIdAsync(userId);
        if (user is null) throw new AiOperationAdmissionException("unauthorized", "Authentication is required.", StatusCodes.Status401Unauthorized);
        var entitlements = AccountPlans.ForRoles(await users.GetRolesAsync(user));
        if (!entitlements.Ai)
            throw new AiOperationAdmissionException(ProEntitlement.RequiredCode, "This AI feature requires Pro.", StatusCodes.Status403Forbidden);
        if (!user.AiEnabled)
            throw new AiOperationAdmissionException(ProEntitlement.DisabledCode, "AI is disabled in your privacy settings.", StatusCodes.Status403Forbidden);

        await AdmissionGate.WaitAsync(cancellationToken);
        try
        {
            var existing = await operations.FindByIdempotencyAsync(taskType, idempotencyKey, cancellationToken);
            if (existing is not null) return Result(existing, false);

            var active = new[] { OperationStatuses.Queued, OperationStatuses.Running, OperationStatuses.WaitingForRetry, OperationStatuses.WaitingForExternalFallback };
            var perUserCapacity = Math.Clamp(configuration.GetValue("AiQueue:PerUserCapacity", 10), 1, 100);
            var globalCapacity = Math.Clamp(configuration.GetValue("AiQueue:GlobalCapacity", 100), perUserCapacity, 10_000);
            var perUserCount = await db.UserOperations.CountAsync(operation => active.Contains(operation.Status), cancellationToken);
            var globalCount = await db.UserOperations.IgnoreQueryFilters().CountAsync(operation => active.Contains(operation.Status), cancellationToken);
            if (perUserCount >= perUserCapacity || globalCount >= globalCapacity)
                throw new AiOperationAdmissionException("ai_queue_full", "AI processing is busy. Try again shortly.", StatusCodes.Status429TooManyRequests, 15);

            var policy = await privacy.EvaluateAsync(user.Id, cancellationToken);
            var usageReservation = AiUsageMeter.ReservationFor(taskType);
            try
            {
                await usage.EnsureCanReserveAsync(user.Id, entitlements, 1, usageReservation.EstimatedTokens, cancellationToken);
            }
            catch (AiUsageLimitException ex)
            {
                throw new AiOperationAdmissionException(ex.Code, ex.Message, StatusCodes.Status429TooManyRequests);
            }
            var deadlineMinutes = Math.Clamp(configuration.GetValue("AiQueue:DeadlineMinutes", 15), 1, 120);
            var created = await operations.CreateAsync(new CreateUserOperation(
                taskType,
                idempotencyKey,
                "pro",
                policy.ExternalProcessingAllowed ? "external_allowed" : "local_only",
                subjectType,
                subjectId,
                priority,
                Math.Clamp(configuration.GetValue("AiQueue:MaxAttempts", 3), 1, 10),
                timeProvider.GetUtcNow().UtcDateTime.AddMinutes(deadlineMinutes),
                usageReservation.InputCharacters,
                usageReservation.EstimatedTokens), cancellationToken);
            return Result(created.Operation, created.Created);
        }
        finally
        {
            AdmissionGate.Release();
        }
    }

    private static AiOperationAdmissionResult Result(UserOperation operation, bool created) =>
        new(operation, created, $"/api/operations/{operation.Id:D}");
}

public sealed record AiOperationExecutionContext(UserOperationLease Lease, string EffectivePrivacyPolicy);
public sealed record AiOperationExecutionResult(
    string? ResultReference,
    string? Provider = null,
    string? Model = null,
    string? RouteReason = null,
    int? UsageInputCharacters = null,
    int? UsageOutputCharacters = null);

public sealed class AiOperationExecutionScope
{
    private readonly AsyncLocal<AiOperationExecutionContext?> _current = new();

    public AiOperationExecutionContext? Current => _current.Value;

    public IDisposable Use(AiOperationExecutionContext context)
    {
        var previous = _current.Value;
        _current.Value = context;
        return new Restore(() => _current.Value = previous);
    }

    private sealed class Restore(Action restore) : IDisposable
    {
        private Action? _restore = restore;
        public void Dispose() => Interlocked.Exchange(ref _restore, null)?.Invoke();
    }
}

public interface IAiOperationHandler
{
    string TaskType { get; }
    Task<AiOperationExecutionResult> ExecuteAsync(AiOperationExecutionContext context, IServiceProvider services, CancellationToken cancellationToken);
}

public sealed class AiOperationFailure(string category, string message, bool retryable) : Exception(message)
{
    public string Category { get; } = category;
    public bool Retryable { get; } = retryable;
}

public sealed class AiOperationWorker(
    IServiceScopeFactory scopes,
    IEnumerable<IAiOperationHandler> registeredHandlers,
    AiPrivacyPolicy privacy,
    AiOperationExecutionScope executionScope,
    IConfiguration configuration)
{
    private readonly IReadOnlyDictionary<string, IAiOperationHandler> _handlers = registeredHandlers
        .ToDictionary(handler => handler.TaskType, StringComparer.Ordinal);

    public async Task<bool> RunOnceAsync(CancellationToken stoppingToken)
    {
        if (_handlers.Count == 0) return false;
        var leaseSeconds = Math.Clamp(configuration.GetValue("AiQueue:LeaseSeconds", 120), 30, 1800);
        await using var claimScope = scopes.CreateAsyncScope();
        var lease = await claimScope.ServiceProvider.GetRequiredService<UserOperationStore>()
            .ClaimNextAsync(TimeSpan.FromSeconds(leaseSeconds), stoppingToken, _handlers.Keys.ToArray());
        if (lease is null) return false;

        await using var ownerScope = scopes.CreateAsyncScope();
        using var owner = ownerScope.ServiceProvider.GetRequiredService<CurrentUserService>().UseBackgroundUser(lease.OwnerUserId);
        var store = ownerScope.ServiceProvider.GetRequiredService<UserOperationStore>();
        var users = ownerScope.ServiceProvider.GetRequiredService<UserManager<ApplicationUser>>();
        var user = await users.FindByIdAsync(lease.OwnerUserId);
        if (user is null || !user.AiEnabled || !AccountPlans.ForRoles(await users.GetRolesAsync(user)).Ai)
        {
            await store.FailAsync(lease.OperationId, lease.LeaseToken, false, "entitlement_changed", "AI access changed before processing began.", TimeSpan.Zero, stoppingToken);
            return true;
        }

        var currentPrivacy = await privacy.EvaluateAsync(user.Id, stoppingToken);
        var effectivePrivacy = lease.PrivacyPolicy == "external_allowed" && currentPrivacy.ExternalProcessingAllowed
            ? "external_allowed"
            : "local_only";
        var timeoutSeconds = Math.Clamp(configuration.GetValue("AiQueue:OperationTimeoutSeconds", 300), 5, 1800);
        using var execution = CancellationTokenSource.CreateLinkedTokenSource(stoppingToken);
        execution.CancelAfter(TimeSpan.FromSeconds(timeoutSeconds));
        using var heartbeatStop = new CancellationTokenSource();
        var heartbeat = MonitorAsync(lease, execution, heartbeatStop.Token);

        try
        {
            var context = new AiOperationExecutionContext(lease, effectivePrivacy);
            using var routing = executionScope.Use(context);
            var result = await _handlers[lease.TaskType].ExecuteAsync(context, ownerScope.ServiceProvider, execution.Token);
            var row = await store.GetAsync(lease.OperationId, stoppingToken);
            if (row?.CancellationRequestedAtUtc is not null)
                await store.AcknowledgeCancellationAsync(lease.OperationId, lease.LeaseToken, stoppingToken);
            else
                await store.CompleteAsync(lease.OperationId, lease.LeaseToken, result.ResultReference,
                    result.Provider, result.Model, result.RouteReason,
                    result.UsageInputCharacters, result.UsageOutputCharacters, stoppingToken);
        }
        catch (AiOperationFailure failure)
        {
            await store.FailAsync(lease.OperationId, lease.LeaseToken, failure.Retryable, failure.Category,
                failure.Message, RetryDelay(lease.AttemptCount), stoppingToken);
        }
        catch (AiGenerationException failure)
        {
            await store.FailAsync(lease.OperationId, lease.LeaseToken, failure.Retryable, failure.Category,
                failure.Message, RetryDelay(lease.AttemptCount), failure.Provider, failure.Model,
                failure.RouteReason, stoppingToken);
        }
        catch (OperationCanceledException) when (!stoppingToken.IsCancellationRequested)
        {
            var row = await store.GetAsync(lease.OperationId, stoppingToken);
            if (row?.CancellationRequestedAtUtc is not null)
                await store.AcknowledgeCancellationAsync(lease.OperationId, lease.LeaseToken, stoppingToken);
            else
                await store.FailAsync(lease.OperationId, lease.LeaseToken, true, "timeout", "AI processing exceeded its deadline.", RetryDelay(lease.AttemptCount), stoppingToken);
        }
        finally
        {
            heartbeatStop.Cancel();
            try { await heartbeat; } catch (OperationCanceledException) { }
        }

        return true;
    }

    private async Task MonitorAsync(UserOperationLease lease, CancellationTokenSource execution, CancellationToken cancellationToken)
    {
        var heartbeatSeconds = Math.Clamp(configuration.GetValue("AiQueue:HeartbeatSeconds", 20), 5, 300);
        using var timer = new PeriodicTimer(TimeSpan.FromSeconds(heartbeatSeconds));
        while (await timer.WaitForNextTickAsync(cancellationToken))
        {
            await using var scope = scopes.CreateAsyncScope();
            using var owner = scope.ServiceProvider.GetRequiredService<CurrentUserService>().UseBackgroundUser(lease.OwnerUserId);
            var store = scope.ServiceProvider.GetRequiredService<UserOperationStore>();
            var row = await store.GetAsync(lease.OperationId, cancellationToken);
            if (row?.CancellationRequestedAtUtc is not null)
            {
                execution.Cancel();
                return;
            }
            await store.HeartbeatAsync(lease.OperationId, lease.LeaseToken,
                TimeSpan.FromSeconds(Math.Clamp(configuration.GetValue("AiQueue:LeaseSeconds", 120), 30, 1800)),
                "processing", null, cancellationToken);
        }
    }

    private static TimeSpan RetryDelay(int attempt) => TimeSpan.FromSeconds(Math.Min(60, (1 << Math.Min(attempt, 5)) + Random.Shared.Next(0, 4)));
}

public sealed class AiOperationHostedService(AiOperationWorker worker, IConfiguration configuration) : BackgroundService
{
    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
    {
        if (!configuration.GetValue("Workers:AiOperationsEnabled", false)) return;
        var concurrency = Math.Clamp(configuration.GetValue("AiQueue:WorkerConcurrency", 1), 1, 4);
        await Task.WhenAll(Enumerable.Range(0, concurrency).Select(_ => RunWorkerAsync(stoppingToken)));
    }

    private async Task RunWorkerAsync(CancellationToken stoppingToken)
    {
        var idleDelayMs = Math.Clamp(configuration.GetValue("AiQueue:IdleDelayMs", 1000), 100, 10_000);
        while (!stoppingToken.IsCancellationRequested)
        {
            if (!await worker.RunOnceAsync(stoppingToken))
                await Task.Delay(idleDelayMs, stoppingToken);
        }
    }
}