ClawdDotNet/src/ClawdDotNet.Core/Budget/BudgetHeadroom.cs

using ClawdDotNet.Core.Config;
using ClawdDotNet.Core.Tasks;

namespace ClawdDotNet.Core.Budget;

/// <summary>
/// #716: Decides whether a late run — carried-over work or a catch-up — may start now,
/// without taking budget the rest of the day still needs.
///
/// <para><b>The rule: the new day has right of way.</b> Headroom is the limit minus what
/// was used, minus what runs in flight have reserved, minus what the day's remaining
/// scheduled runs will need, minus a safety margin. A late run starts only if its
/// estimate fits into that. Otherwise one overdrawn day would push its shortfall into the
/// next, and the next, instead of resolving it.</para>
///
/// <para>This replaces the flat "at most 25 % of the new day" from the concept: a fixed
/// share is either too much on a busy day or needlessly little on a quiet one. The same
/// calculation is what fill work (#717) will run on.</para>
///
/// <para><b>Estimates, honestly.</b> A run is estimated from the agent's recent runs
/// (<see cref="BudgetGuard.EstimateRunAsync"/>). Without history the estimate is zero —
/// then only the margin keeps a reserve. That errs towards letting late work through;
/// the budget check at the start of the run still refuses once the limit is reached.</para>
/// </summary>
public sealed class BudgetHeadroom : ILateRunGate, IBudgetOutlook
{
    /// <summary>Share of each limit kept free as a reserve for surprises.</summary>
    public const decimal SafetyShare = 0.10m;

    private readonly BudgetGuard _guard;
    private readonly Func<IReadOnlyList<AgentConfig>> _agents;
    private readonly Func<InstanceBudget> _instanceBudget;
    private readonly TimeProvider _clock;

    public BudgetHeadroom(
        BudgetGuard guard,
        Func<IReadOnlyList<AgentConfig>> agents,
        Func<InstanceBudget> instanceBudget,
        TimeProvider? clock = null)
    {
        _guard = guard;
        _agents = agents;
        _instanceBudget = instanceBudget;
        _clock = clock ?? TimeProvider.System;
    }

    public async Task<bool> HasRoomAsync(
        TaskItem lateTask, IReadOnlyList<TaskItem> runnable, DateTime nowUtc, CancellationToken ct)
    {
        var agents = _agents();
        var agent = FindAgent(agents, lateTask);
        var instanceBudget = _instanceBudget();

        // Without a known agent there is nothing to calculate against; the budget check at
        // the start of the run still applies.
        if (agent is null || (!agent.Budget.HasAnyLimit && !instanceBudget.HasAnyLimit))
            return true;

        var endOfDay = StartOfNextLocalDay(nowUtc);
        var estimates = new Dictionary<string, RunEstimate>(StringComparer.OrdinalIgnoreCase);

        async Task<RunEstimate> EstimateFor(string agentId)
        {
            if (!estimates.TryGetValue(agentId, out var estimate))
                estimates[agentId] = estimate = await _guard.EstimateRunAsync(agentId, ct);
            return estimate;
        }

        var ownEstimate = await EstimateFor(agent.AgentId);

        // Regular work of the rest of the day: everything scheduled except late work
        // itself — that is what has to wait for headroom, not what headroom is kept for.
        var regular = runnable.Where(t => t.Id != lateTask.Id && !IsLate(t)).ToList();

        if (agent.Budget.HasAnyLimit)
        {
            var demand = default(RunEstimate);
            foreach (var task in regular.Where(t => IsOf(t, agent.AgentId)))
                demand += Times(ownEstimate, TaskSchedule.OccurrencesBetween(task, nowUtc, endOfDay));

            var used = await _guard.GetUsedTodayAsync(agent.AgentId, ct);
            var reserved = _guard.Reserved(agent.AgentId);

            if (!Fits(agent.Budget.DailyCostUsd, used.CostUsd, reserved.CostUsd, demand.CostUsd, ownEstimate.CostUsd)
                || !Fits(agent.Budget.DailyTokens, used.TotalTokens, reserved.Tokens, demand.Tokens, ownEstimate.Tokens))
                return false;
        }

        if (instanceBudget.HasAnyLimit)
        {
            var demand = default(RunEstimate);
            foreach (var task in regular)
            {
                var agentId = TaskAssignee.AgentId(task.Assignee);
                if (agentId.Length == 0)
                    continue;
                demand += Times(await EstimateFor(agentId), TaskSchedule.OccurrencesBetween(task, nowUtc, endOfDay));
            }

            var used = await _guard.GetUsedTodayAsync("", ct);
            var reserved = _guard.Reserved("");

            if (!Fits(instanceBudget.DailyCostUsd, used.CostUsd, reserved.CostUsd, demand.CostUsd, ownEstimate.CostUsd)
                || !Fits(instanceBudget.DailyTokens, used.TotalTokens, reserved.Tokens, demand.Tokens, ownEstimate.Tokens))
                return false;
        }

        return true;
    }

    /// <summary>
    /// #716: Walks through the rest of today's schedule in time order and finds, per agent
    /// with a limit and for the instance, the first run the budget would no longer cover.
    /// Late work is left out: it only ever runs in headroom, so it cannot be what makes the
    /// day's own runs fail.
    /// </summary>
    public async Task<IReadOnlyList<BudgetShortfall>> ShortfallsAsync(
        IReadOnlyList<TaskItem> runnable, DateTime nowUtc, CancellationToken ct)
    {
        var agents = _agents();
        var instanceBudget = _instanceBudget();
        var endOfDay = StartOfNextLocalDay(nowUtc);
        var result = new List<BudgetShortfall>();

        // Every planned run until midnight, with whose it is.
        var planned = new List<(DateTime At, string AgentId)>();
        foreach (var task in runnable.Where(t => !IsLate(t)))
        {
            var agent = FindAgent(agents, task);
            if (agent is null)
                continue;
            foreach (var at in TaskSchedule.OccurrenceTimesBetween(task, nowUtc, endOfDay))
                planned.Add((at, agent.AgentId));
        }

        planned.Sort((a, b) => a.At.CompareTo(b.At));

        var estimates = new Dictionary<string, RunEstimate>(StringComparer.OrdinalIgnoreCase);
        foreach (var agentId in planned.Select(p => p.AgentId).Distinct(StringComparer.OrdinalIgnoreCase))
            estimates[agentId] = await _guard.EstimateRunAsync(agentId, ct);

        foreach (var agent in agents.Where(a => a.Budget.HasAnyLimit))
        {
            var runs = planned.Where(p => string.Equals(p.AgentId, agent.AgentId, StringComparison.OrdinalIgnoreCase)).ToList();
            var used = await _guard.GetUsedTodayAsync(agent.AgentId, ct);
            var reserved = _guard.Reserved(agent.AgentId);
            if (FirstRefused(runs, estimates, agent.Budget.DailyCostUsd, agent.Budget.DailyTokens, used, reserved) is { } index)
                result.Add(new BudgetShortfall(agent.AgentId, runs[index].At, runs.Count - index, runs.Count));
        }

        if (instanceBudget.HasAnyLimit)
        {
            var used = await _guard.GetUsedTodayAsync("", ct);
            var reserved = _guard.Reserved("");
            if (FirstRefused(planned, estimates, instanceBudget.DailyCostUsd, instanceBudget.DailyTokens, used, reserved) is { } index)
                result.Add(new BudgetShortfall("", planned[index].At, planned.Count - index, planned.Count));
        }

        return result;
    }

    /// <summary>
    /// The index of the first run that would find the budget used up, or <c>null</c>.
    /// Mirrors the check at the start of a run: a run is refused once usage has reached
    /// the limit, not merely because its own estimate would overshoot.
    /// </summary>
    private static int? FirstRefused(
        List<(DateTime At, string AgentId)> runs, Dictionary<string, RunEstimate> estimates,
        decimal costLimit, long tokenLimit, Accounting.DailyUsage used, RunEstimate reserved)
    {
        var cost = used.CostUsd + reserved.CostUsd;
        long tokens = used.TotalTokens + reserved.Tokens;

        for (var i = 0; i < runs.Count; i++)
        {
            if ((costLimit > 0 && cost >= costLimit) || (tokenLimit > 0 && tokens >= tokenLimit))
                return i;

            var estimate = estimates.GetValueOrDefault(runs[i].AgentId);
            cost += estimate.CostUsd;
            tokens += estimate.Tokens;
        }

        return null;
    }

    /// <summary>Carried-over work and pending catch-ups are the late runs.</summary>
    public static bool IsLate(TaskItem task) => task.IsCarriedOver || task.CatchUpPending;

    /// <summary>0 or less means: no limit of this kind.</summary>
    private static bool Fits(decimal limit, decimal used, decimal reserved, decimal demand, decimal estimate)
        => limit <= 0 || limit - used - reserved - demand - limit * SafetyShare >= estimate;

    private static bool Fits(long limit, long used, long reserved, long demand, long estimate)
        => limit <= 0 || limit - used - reserved - demand - (long)(limit * SafetyShare) >= estimate;

    private static RunEstimate Times(RunEstimate estimate, int runs)
        => new(estimate.CostUsd * runs, estimate.Tokens * runs);

    private static bool IsOf(TaskItem task, string agentId)
        => string.Equals(TaskAssignee.AgentId(task.Assignee), agentId, StringComparison.OrdinalIgnoreCase);

    private static AgentConfig? FindAgent(IReadOnlyList<AgentConfig> agents, TaskItem task)
    {
        var agentId = TaskAssignee.AgentId(task.Assignee);
        return agentId.Length == 0
            ? (agents.Count == 1 ? agents[0] : null)
            : agents.FirstOrDefault(a => string.Equals(a.AgentId, agentId, StringComparison.OrdinalIgnoreCase));
    }

    private DateTime StartOfNextLocalDay(DateTime nowUtc)
    {
        var zone = _clock.LocalTimeZone;
        var local = TimeZoneInfo.ConvertTimeFromUtc(DateTime.SpecifyKind(nowUtc, DateTimeKind.Utc), zone);
        var midnight = DateTime.SpecifyKind(local.Date.AddDays(1), DateTimeKind.Unspecified);
        while (zone.IsInvalidTime(midnight))
            midnight = midnight.AddMinutes(1);
        return TimeZoneInfo.ConvertTimeToUtc(midnight, zone);
    }
}