ClawdDotNet/src/ClawdDotNet.Core/Budget/BudgetGuard.cs

using ClawdDotNet.Core.Config;
using ClawdDotNet.Core.Accounting;

namespace ClawdDotNet.Core.Budget;

/// <summary>Welche Grenze überschritten wurde.</summary>
public enum BudgetLimit
{
    None,
    AgentCost,
    AgentTokens,
    InstanceCost,
    InstanceTokens
}

/// <summary>#716: What a run is expected to use — the agent's average over recent runs.</summary>
public readonly record struct RunEstimate(decimal CostUsd, long Tokens)
{
    public static RunEstimate operator +(RunEstimate a, RunEstimate b) => new(a.CostUsd + b.CostUsd, a.Tokens + b.Tokens);
    public static RunEstimate operator -(RunEstimate a, RunEstimate b) => new(a.CostUsd - b.CostUsd, a.Tokens - b.Tokens);
}

public sealed record BudgetStatus(bool IsAllowed, BudgetLimit Limit, string? Reason)
{
    public static BudgetStatus Allowed { get; } = new(true, BudgetLimit.None, null);
}

/// <summary>
/// Wacht über das Tagesbudget.
///
/// Hintergrund (K5): Es gab keine Obergrenze. Ein Agent in einer Schleife — etwa durch
/// gegenseitige send_message-Aufrufe — konnte unbeaufsichtigt Guthaben verbrennen; die
/// Credits-Anzeige war rein informativ.
///
/// Zwei Arten von Grenzen, weil sich Kosten nicht immer beziffern lassen: Liefert der
/// Anbieter für ein Modell keine Preise, greift die Kostengrenze nicht — die
/// Token-Grenze dagegen immer. Wer sich absichern will, setzt beide.
/// </summary>
public sealed class BudgetGuard
{
    private readonly IUsageRepository _usage;
    private readonly Func<DateTime> _now;

    /// <summary>How many recent runs of an agent make up its per-run estimate (#716).</summary>
    public const int EstimateSample = 20;

    /// <summary>
    /// #716: Check and reserve happen under this lock. Without it, two runs starting at
    /// the same moment would both read the same "used" figure and both pass.
    /// </summary>
    private readonly SemaphoreSlim _reserveGate = new(1, 1);

    private readonly Lock _reservedLock = new();
    private readonly Dictionary<string, RunEstimate> _reservedByAgent = new(StringComparer.OrdinalIgnoreCase);
    private RunEstimate _reservedInstance;

    /// <summary>#717: What runs in flight have reserved of each special budget.</summary>
    private readonly Dictionary<string, RunEstimate> _reservedByPot = new(StringComparer.OrdinalIgnoreCase);

    public BudgetGuard(IUsageRepository usage, Func<DateTime>? now = null)
    {
        _usage = usage;
        _now = now ?? (() => DateTime.Now);
    }

    /// <summary>
    /// Prüft, ob der Agent heute noch laufen darf.
    /// </summary>
    public Task<BudgetStatus> CheckAsync(
        AgentConfig agent, InstanceBudget instanceBudget, CancellationToken ct)
        => CheckCoreAsync(agent, instanceBudget, default, default, ct);

    /// <summary>
    /// #716: Checks like <see cref="CheckAsync"/> and, if the run may start, reserves its
    /// estimated usage until the returned handle is disposed — i.e. until the run's usage
    /// is recorded.
    ///
    /// <para><b>What changes.</b> Before, each run was checked on its own against what was
    /// already recorded. Runs in flight were invisible, so several starting together all
    /// passed and exceeded the budget together (#682 capped how many; this closes it).
    /// Now the runs in flight count with their estimate.</para>
    ///
    /// <para><b>What does not change.</b> A single run is still refused only when the
    /// budget is already used up — its own estimate is not added. Otherwise a small task
    /// would be refused because the agent's average run is large.</para>
    /// </summary>
    public async Task<(BudgetStatus Status, IDisposable? Reservation)> ReserveAsync(
        AgentConfig agent, InstanceBudget instanceBudget, CancellationToken ct)
    {
        if (!agent.Budget.HasAnyLimit && !instanceBudget.HasAnyLimit)
            return (BudgetStatus.Allowed, null);

        var estimate = await EstimateRunAsync(agent.AgentId, ct);

        await _reserveGate.WaitAsync(ct);
        try
        {
            RunEstimate agentReserved, instanceReserved;
            lock (_reservedLock)
            {
                agentReserved = _reservedByAgent.GetValueOrDefault(agent.AgentId);
                instanceReserved = _reservedInstance;
            }

            var status = await CheckCoreAsync(agent, instanceBudget, agentReserved, instanceReserved, ct);
            if (!status.IsAllowed)
                return (status, null);

            lock (_reservedLock)
            {
                _reservedByAgent[agent.AgentId] = _reservedByAgent.GetValueOrDefault(agent.AgentId) + estimate;
                _reservedInstance += estimate;
            }

            return (status, new Reservation(() => Release(agent.AgentId, estimate)));
        }
        finally
        {
            _reserveGate.Release();
        }
    }

    /// <summary>
    /// #717: Reserves a run on a special budget instead of the daily limits. Refused once what
    /// the pot's runs used, plus what runs in flight reserved, has reached the amount — the
    /// same rule as for the daily limit: a run is not refused for its own estimate.
    /// </summary>
    public async Task<IDisposable?> ReservePotAsync(SpecialBudget pot, decimal spent, string agentId, CancellationToken ct)
    {
        var estimate = await EstimateRunAsync(agentId, ct);

        await _reserveGate.WaitAsync(ct);
        try
        {
            lock (_reservedLock)
            {
                var reserved = _reservedByPot.GetValueOrDefault(pot.TaskId);
                if (spent + reserved.CostUsd >= pot.AmountUsd)
                    return null;
                _reservedByPot[pot.TaskId] = reserved + estimate;
            }

            return new Reservation(() => ReleasePot(pot.TaskId, estimate));
        }
        finally
        {
            _reserveGate.Release();
        }
    }

    private void ReleasePot(string potId, RunEstimate estimate)
    {
        lock (_reservedLock)
        {
            var left = _reservedByPot.GetValueOrDefault(potId) - estimate;
            if (left.CostUsd <= 0 && left.Tokens <= 0)
                _reservedByPot.Remove(potId);
            else
                _reservedByPot[potId] = left;
        }
    }

    /// <summary>Today's usage of an agent, or with an empty id of the whole instance.</summary>
    public Task<DailyUsage> GetUsedTodayAsync(string agentId, CancellationToken ct)
        => _usage.GetDailyForBudgetAsync(DateOnly.FromDateTime(_now()), agentId, ct);

    /// <summary>What is currently reserved by runs in flight — for the agent, or with an
    /// empty id for the whole instance.</summary>
    public RunEstimate Reserved(string agentId)
    {
        lock (_reservedLock)
            return agentId.Length == 0 ? _reservedInstance : _reservedByAgent.GetValueOrDefault(agentId);
    }

    /// <summary>
    /// The average of the agent's recent runs. Runs without usage (refusals, aborted
    /// before the first call) do not count — they would drag the estimate towards zero.
    /// </summary>
    public async Task<RunEstimate> EstimateRunAsync(string agentId, CancellationToken ct)
    {
        var recent = await _usage.GetRecentForAgentAsync(agentId, EstimateSample, ct);
        var runs = recent.Where(r => r.TotalTokens > 0).ToList();
        if (runs.Count == 0)
            return default;

        return new RunEstimate(
            runs.Average(r => r.CostUsd),
            (long)Math.Ceiling(runs.Average(r => (double)r.TotalTokens)));
    }

    private void Release(string agentId, RunEstimate estimate)
    {
        lock (_reservedLock)
        {
            var left = _reservedByAgent.GetValueOrDefault(agentId) - estimate;
            if (left.CostUsd <= 0 && left.Tokens <= 0)
                _reservedByAgent.Remove(agentId);
            else
                _reservedByAgent[agentId] = left;

            _reservedInstance -= estimate;
        }
    }

    private sealed class Reservation(Action release) : IDisposable
    {
        private int _disposed;

        public void Dispose()
        {
            if (Interlocked.Exchange(ref _disposed, 1) == 0)
                release();
        }
    }

    private async Task<BudgetStatus> CheckCoreAsync(
        AgentConfig agent, InstanceBudget instanceBudget,
        RunEstimate agentReserved, RunEstimate instanceReserved, CancellationToken ct)
    {
        var today = DateOnly.FromDateTime(_now());

        // Agentengrenze
        var agentBudget = agent.Budget;
        if (agentBudget.HasAnyLimit)
        {
            var used = await _usage.GetDailyForBudgetAsync(today, agent.AgentId, ct);

            if (Exceeds(agentBudget.DailyCostUsd, used.CostUsd + agentReserved.CostUsd))
            {
                return new BudgetStatus(false, BudgetLimit.AgentCost,
                    $"Tagesbudget des Agenten erschöpft: {used.CostUsd:F4} von maximal " +
                    $"{agentBudget.DailyCostUsd:F2} USD verbraucht." +
                    ReservedHint(agentReserved.CostUsd) +
                    UnpricedHint(used));
            }

            if (Exceeds(agentBudget.DailyTokens, used.TotalTokens + agentReserved.Tokens))
            {
                return new BudgetStatus(false, BudgetLimit.AgentTokens,
                    $"Token-Tageslimit des Agenten erschöpft: {used.TotalTokens:N0} von maximal " +
                    $"{agentBudget.DailyTokens:N0}." +
                    ReservedTokensHint(agentReserved.Tokens));
            }
        }

        // Instanzgrenze — gilt für alle Agenten zusammen
        if (instanceBudget.HasAnyLimit)
        {
            var used = await _usage.GetDailyForBudgetAsync(today, agentId: "", ct);

            if (Exceeds(instanceBudget.DailyCostUsd, used.CostUsd + instanceReserved.CostUsd))
            {
                return new BudgetStatus(false, BudgetLimit.InstanceCost,
                    $"Tagesbudget der Instanz erschöpft: {used.CostUsd:F4} von maximal " +
                    $"{instanceBudget.DailyCostUsd:F2} USD verbraucht." +
                    ReservedHint(instanceReserved.CostUsd) +
                    UnpricedHint(used));
            }

            if (Exceeds(instanceBudget.DailyTokens, used.TotalTokens + instanceReserved.Tokens))
            {
                return new BudgetStatus(false, BudgetLimit.InstanceTokens,
                    $"Token-Tageslimit der Instanz erschöpft: {used.TotalTokens:N0} von maximal " +
                    $"{instanceBudget.DailyTokens:N0}." +
                    ReservedTokensHint(instanceReserved.Tokens));
            }
        }

        return BudgetStatus.Allowed;
    }

    /// <summary>Verbleibendes Kostenbudget des Agenten für heute; null bei fehlender Grenze.</summary>
    public async Task<decimal?> GetRemainingCostAsync(AgentConfig agent, CancellationToken ct)
    {
        if (agent.Budget.DailyCostUsd <= 0)
            return null;

        var used = await _usage.GetDailyForBudgetAsync(DateOnly.FromDateTime(_now()), agent.AgentId, ct);
        return Math.Max(0m, agent.Budget.DailyCostUsd - used.CostUsd);
    }

    /// <summary>0 oder kleiner bedeutet: keine Grenze gesetzt.</summary>
    private static bool Exceeds(decimal limit, decimal used) => limit > 0 && used >= limit;

    private static bool Exceeds(long limit, long used) => limit > 0 && used >= limit;

    /// <summary>#716: Without this, a refusal at "3.20 of 5.00 USD used" would look wrong.</summary>
    private static string ReservedHint(decimal reserved)
        => reserved > 0 ? $" Weitere {reserved:F4} USD sind für laufende Läufe reserviert." : "";

    private static string ReservedTokensHint(long reserved)
        => reserved > 0 ? $" Weitere {reserved:N0} Tokens sind für laufende Läufe reserviert." : "";

    /// <summary>
    /// Ohne Preisangaben ist die verbrauchte Summe unvollständig — sonst wirkte ein
    /// niedriger Wert wie ein noch offener Spielraum.
    /// </summary>
    private static string UnpricedHint(DailyUsage used)
        => used.CostIsComplete
            ? ""
            : " Achtung: Für mindestens ein Modell lagen keine Preise vor, der tatsächliche " +
              "Verbrauch liegt höher.";
}