ClawdDotNet/tests/ClawdDotNet.Core.Tests/Budget/BudgetReservationTests.cs

using ClawdDotNet.Core.Accounting;
using ClawdDotNet.Core.Budget;
using ClawdDotNet.Core.Config;
using Shouldly;

namespace ClawdDotNet.Core.Tests.Budget;

/// <summary>
/// #716 (the rest of #682): Runs in flight count against the budget with their estimate.
///
/// Before, every run was checked on its own against what was already recorded, so runs
/// starting together all saw the same figure and all passed. #682 capped how many could
/// start at once; the reservation closes the gap.
/// </summary>
public sealed class BudgetReservationTests
{
    private static AgentConfig Agent(decimal dailyCost = 1.00m, long dailyTokens = 0)
        => new() { AgentId = "worker", Budget = new AgentBudget { DailyCostUsd = dailyCost, DailyTokens = dailyTokens } };

    /// <summary>Recent runs averaging 0.30 USD and 3,000 tokens.</summary>
    private static List<RunUsage> History(decimal cost = 0.30m, int tokens = 3_000)
        => [.. Enumerable.Range(0, 5).Select(_ => new RunUsage
        {
            AgentId = "worker", PromptTokens = tokens, CostUsd = cost, CostIsKnown = true, Status = "Completed"
        })];

    [Fact]
    public async Task Runs_in_flight_count_against_the_budget()
    {
        var guard = new BudgetGuard(new FakeUsage(usedCost: 0.60m, History()));

        var (first, a) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
        var (second, b) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
        var (third, c) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);

        first.IsAllowed.ShouldBeTrue("0.60 used, nothing in flight");
        second.IsAllowed.ShouldBeTrue("0.60 used + 0.30 in flight");
        third.IsAllowed.ShouldBeFalse("0.60 used + 0.60 in flight — before #716 this one passed too");
        third.Reason.ShouldNotBeNull().ShouldContain("reserviert");
        c.ShouldBeNull();

        a!.Dispose();
        b!.Dispose();
    }

    [Fact]
    public async Task A_finished_run_frees_its_reservation()
    {
        var guard = new BudgetGuard(new FakeUsage(usedCost: 0.60m, History()));

        var (_, a) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
        var (_, b) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
        a!.Dispose();
        a.Dispose(); // twice is harmless

        var (again, c) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
        again.IsAllowed.ShouldBeTrue();

        b!.Dispose();
        c!.Dispose();
        guard.Reserved("worker").ShouldBe(default(RunEstimate));
        guard.Reserved("").ShouldBe(default(RunEstimate));
    }

    [Fact]
    public async Task A_single_run_is_not_refused_because_of_its_own_estimate()
    {
        // 0.95 of 1.00 used, an average run costs 0.30: a small task may still fit, and
        // refusing it on the average would be guessing.
        var guard = new BudgetGuard(new FakeUsage(usedCost: 0.95m, History()));

        var (status, reservation) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);

        status.IsAllowed.ShouldBeTrue();
        reservation!.Dispose();
    }

    [Fact]
    public async Task Parallel_starts_cannot_all_pass_together()
    {
        var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, History()));

        var results = await Task.WhenAll(Enumerable.Range(0, 10)
            .Select(_ => guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default)));

        // 0.00, 0.30, 0.60, 0.90 in flight pass; from 1.20 on the budget is spoken for.
        results.Count(r => r.Status.IsAllowed).ShouldBe(4);

        foreach (var (_, reservation) in results)
            reservation?.Dispose();
    }

    [Fact]
    public async Task The_token_limit_counts_reservations_too()
    {
        var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, History(cost: 0m, tokens: 4_000)));
        var agent = Agent(dailyCost: 0, dailyTokens: 10_000);

        var (a, ra) = await guard.ReserveAsync(agent, InstanceBudget.Unlimited, default);
        var (b, rb) = await guard.ReserveAsync(agent, InstanceBudget.Unlimited, default);
        var (c, _) = await guard.ReserveAsync(agent, InstanceBudget.Unlimited, default);
        var (d, _) = await guard.ReserveAsync(agent, InstanceBudget.Unlimited, default);

        a.IsAllowed.ShouldBeTrue();
        b.IsAllowed.ShouldBeTrue();
        c.IsAllowed.ShouldBeTrue("8,000 of 10,000 reserved");
        d.IsAllowed.ShouldBeFalse("12,000 reserved");
        d.Limit.ShouldBe(BudgetLimit.AgentTokens);

        ra!.Dispose();
        rb!.Dispose();
    }

    [Fact]
    public async Task The_instance_limit_sees_the_runs_of_all_agents()
    {
        var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, History()));
        var instance = new InstanceBudget { DailyCostUsd = 0.50m };
        var unlimitedAgent = Agent(dailyCost: 0);

        var (a, _) = await guard.ReserveAsync(unlimitedAgent, instance, default);
        var (b, _) = await guard.ReserveAsync(unlimitedAgent, instance, default);
        var (c, _) = await guard.ReserveAsync(unlimitedAgent, instance, default);

        a.IsAllowed.ShouldBeTrue();
        b.IsAllowed.ShouldBeTrue("0.30 in flight");
        c.IsAllowed.ShouldBeFalse("0.60 in flight over an instance limit of 0.50");
        c.Limit.ShouldBe(BudgetLimit.InstanceCost);
    }

    [Fact]
    public async Task Without_any_limit_nothing_is_reserved()
    {
        var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, History()));

        var (status, reservation) = await guard.ReserveAsync(Agent(dailyCost: 0), InstanceBudget.Unlimited, default);

        status.IsAllowed.ShouldBeTrue();
        reservation.ShouldBeNull();
    }

    [Fact]
    public async Task Runs_without_usage_do_not_drag_the_estimate_down()
    {
        var history = History();
        history.Add(new RunUsage { AgentId = "worker", Status = "BudgetExceeded" }); // 0 tokens
        var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, history));

        var estimate = await guard.EstimateRunAsync("worker", default);

        estimate.CostUsd.ShouldBe(0.30m);
        estimate.Tokens.ShouldBe(3_000);
    }

    private sealed class FakeUsage(decimal usedCost, IReadOnlyList<RunUsage> recent) : IUsageRepository
    {
        public Task RecordAsync(RunUsage usage, CancellationToken ct) => Task.CompletedTask;

        public async Task<DailyUsage> GetDailyAsync(DateOnly date, string agentId, CancellationToken ct)
        {
            await Task.Yield(); // like the database: the caller does not get the answer synchronously
            return new DailyUsage(date, 0, 0, usedCost, true, 0);
        }

        public Task<IReadOnlyList<RunUsage>> GetRecentAsync(int limit, CancellationToken ct)
            => Task.FromResult(recent);

        public Task<IReadOnlyList<RunUsage>> GetRecentForAgentAsync(string agentId, int limit, CancellationToken ct)
            => Task.FromResult(recent);

        public Task<int> PurgeBeforeAsync(DateOnly date, CancellationToken ct) => Task.FromResult(0);
    }
}