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

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

namespace ClawdDotNet.Core.Tests.Budget;

/// <summary>
/// K5 from the inventory: there was no upper limit. An agent in a loop —
/// for example through mutual send_message calls — could burn credit
/// unsupervised; the credits display was purely informational.
/// </summary>
public sealed class BudgetGuardTests : IDisposable
{
    private readonly string _directory;
    private readonly SqliteUsageRepository _usage;
    private readonly DateTime _today = new(2026, 7, 28, 14, 0, 0, DateTimeKind.Local);

    public BudgetGuardTests()
    {
        _directory = Path.Combine(Path.GetTempPath(), "clawd-tests", Guid.NewGuid().ToString("N"));
        _usage = new SqliteUsageRepository(new SqliteStorage(Path.Combine(_directory, "state.db")));
    }

    public void Dispose()
    {
        ClawdDotNet.Core.Storage.SqliteStorage.ReleasePooledConnections(_directory);
        try { Directory.Delete(_directory, recursive: true); }
        catch { /* Cleaning up is secondary */ }
    }

    private BudgetGuard Guard() => new(_usage, () => _today);

    private static AgentConfig Agent(decimal costLimit = 0, long tokenLimit = 0, string id = "agent-a") => new()
    {
        AgentId = id,
        Budget = new AgentBudget { DailyCostUsd = costLimit, DailyTokens = tokenLimit }
    };

    private Task Record(decimal cost, int prompt = 1000, int completion = 100,
        string agentId = "agent-a", bool costKnown = true, DateTime? when = null)
        => _usage.RecordAsync(new RunUsage
        {
            AgentId = agentId,
            Model = "test/model",
            PromptTokens = prompt,
            CompletionTokens = completion,
            CostUsd = cost,
            CostIsKnown = costKnown,
            OccurredAt = when ?? _today
        }, default);

    // ═══════════════════════════════════════════════════════════
    //  Without limits
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task Without_a_set_limit_running_is_always_allowed()
    {
        await Record(cost: 1_000m);

        var status = await Guard().CheckAsync(Agent(), InstanceBudget.Unlimited, default);

        status.IsAllowed.ShouldBeTrue();
    }

    // ═══════════════════════════════════════════════════════════
    //  Cost limit
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task Below_the_cost_limit_the_agent_keeps_running()
    {
        await Record(cost: 1.50m);

        var status = await Guard().CheckAsync(Agent(costLimit: 5m), InstanceBudget.Unlimited, default);

        status.IsAllowed.ShouldBeTrue();
    }

    [Fact]
    public async Task Once_the_cost_limit_is_reached_no_run_is_started()
    {
        await Record(cost: 3m);
        await Record(cost: 2m);

        var status = await Guard().CheckAsync(Agent(costLimit: 5m), InstanceBudget.Unlimited, default);

        status.IsAllowed.ShouldBeFalse();
        status.Limit.ShouldBe(BudgetLimit.AgentCost);
        status.Reason.ShouldNotBeNull().ShouldContain("5");
    }

    // ═══════════════════════════════════════════════════════════
    //  Token limit — applies even without price information
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task The_token_limit_applies_even_without_prices()
    {
        // The actual purpose of the second limit: with an unknown model the
        // cost sum stays 0 and a pure cost limit would be ineffective.
        await Record(cost: 0m, prompt: 60_000, completion: 5_000, costKnown: false);

        var agent = Agent(costLimit: 10m, tokenLimit: 50_000);
        var status = await Guard().CheckAsync(agent, InstanceBudget.Unlimited, default);

        status.IsAllowed.ShouldBeFalse();
        status.Limit.ShouldBe(BudgetLimit.AgentTokens);
    }

    [Fact]
    public async Task Missing_prices_are_mentioned_in_the_reasoning()
    {
        // Otherwise a low consumption would look like room still left.
        await Record(cost: 0m, prompt: 10_000, completion: 1_000, costKnown: false);
        await Record(cost: 5m);

        var status = await Guard().CheckAsync(Agent(costLimit: 5m), InstanceBudget.Unlimited, default);

        status.IsAllowed.ShouldBeFalse();
        status.Reason.ShouldNotBeNull().ShouldContain("keine Preise");
    }

    // ═══════════════════════════════════════════════════════════
    //  Instance limit
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task The_instance_limit_counts_all_agents_together()
    {
        await Record(cost: 3m, agentId: "agent-a");
        await Record(cost: 3m, agentId: "agent-b");
        await Record(cost: 3m, agentId: "agent-c");

        var instance = new InstanceBudget { DailyCostUsd = 8m };
        var status = await Guard().CheckAsync(Agent(), instance, default);

        status.IsAllowed.ShouldBeFalse();
        status.Limit.ShouldBe(BudgetLimit.InstanceCost);
    }

    [Fact]
    public async Task A_single_agent_stays_allowed_under_the_instance_limit()
    {
        await Record(cost: 2m, agentId: "agent-a");

        var instance = new InstanceBudget { DailyCostUsd = 10m };
        var status = await Guard().CheckAsync(Agent(costLimit: 5m), instance, default);

        status.IsAllowed.ShouldBeTrue();
    }

    // ═══════════════════════════════════════════════════════════
    //  Day change
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task The_previous_days_usage_does_not_count()
    {
        await Record(cost: 100m, when: _today.AddDays(-1));

        var status = await Guard().CheckAsync(Agent(costLimit: 5m), InstanceBudget.Unlimited, default);

        status.IsAllowed.ShouldBeTrue();
    }

    [Fact]
    public async Task On_the_next_day_the_budget_is_available_again()
    {
        await Record(cost: 10m);

        var tomorrow = new BudgetGuard(_usage, () => _today.AddDays(1));
        var status = await tomorrow.CheckAsync(Agent(costLimit: 5m), InstanceBudget.Unlimited, default);

        status.IsAllowed.ShouldBeTrue();
    }

    // ═══════════════════════════════════════════════════════════
    //  Remaining budget
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task The_remaining_budget_is_calculated_correctly()
    {
        await Record(cost: 1.25m);

        var remaining = await Guard().GetRemainingCostAsync(Agent(costLimit: 5m), default);

        remaining.ShouldBe(3.75m);
    }

    [Fact]
    public async Task Without_cost_limit_there_is_no_remaining_budget()
    {
        (await Guard().GetRemainingCostAsync(Agent(), default)).ShouldBeNull();
    }

    [Fact]
    public async Task The_remaining_budget_does_not_go_negative()
    {
        await Record(cost: 20m);

        (await Guard().GetRemainingCostAsync(Agent(costLimit: 5m), default)).ShouldBe(0m);
    }
}

public sealed class UsageRepositoryTests : IDisposable
{
    private readonly string _directory;
    private readonly SqliteUsageRepository _usage;
    private static readonly DateTime Tag = new(2026, 7, 28, 10, 0, 0, DateTimeKind.Local);

    public UsageRepositoryTests()
    {
        _directory = Path.Combine(Path.GetTempPath(), "clawd-tests", Guid.NewGuid().ToString("N"));
        _usage = new SqliteUsageRepository(new SqliteStorage(Path.Combine(_directory, "state.db")));
    }

    public void Dispose()
    {
        ClawdDotNet.Core.Storage.SqliteStorage.ReleasePooledConnections(_directory);
        try { Directory.Delete(_directory, recursive: true); }
        catch { }
    }

    private Task Record(decimal cost, string agentId = "a", DateTime? when = null, bool known = true)
        => _usage.RecordAsync(new RunUsage
        {
            AgentId = agentId,
            Model = "test/model",
            PromptTokens = 1000,
            CompletionTokens = 200,
            CostUsd = cost,
            CostIsKnown = known,
            Status = "Completed",
            OccurredAt = when ?? Tag
        }, default);

    [Fact]
    public async Task Amounts_survive_the_round_trip_without_rounding_errors()
    {
        // Stored as floating point, cent amounts would add up over thousands
        // of runs to a visible error.
        for (var i = 0; i < 1000; i++)
            await Record(0.0001m);

        var daily = await _usage.GetDailyAsync(DateOnly.FromDateTime(Tag), "a", default);

        daily.CostUsd.ShouldBe(0.1m);
    }

    [Fact]
    public async Task The_daily_usage_sums_tokens_and_runs()
    {
        await Record(1m);
        await Record(2m);

        var daily = await _usage.GetDailyAsync(DateOnly.FromDateTime(Tag), "a", default);

        daily.RunCount.ShouldBe(2);
        daily.PromptTokens.ShouldBe(2000);
        daily.CompletionTokens.ShouldBe(400);
        daily.TotalTokens.ShouldBe(2400);
    }

    [Fact]
    public async Task Without_agent_filter_the_whole_instance_counts()
    {
        await Record(1m, agentId: "a");
        await Record(2m, agentId: "b");

        var daily = await _usage.GetDailyAsync(DateOnly.FromDateTime(Tag), agentId: "", default);

        daily.CostUsd.ShouldBe(3m);
        daily.RunCount.ShouldBe(2);
    }

    [Fact]
    public async Task A_run_without_price_makes_the_day_incomplete()
    {
        await Record(1m);
        await Record(0m, known: false);

        var daily = await _usage.GetDailyAsync(DateOnly.FromDateTime(Tag), "a", default);

        daily.CostIsComplete.ShouldBeFalse();
    }

    [Fact]
    public async Task A_day_without_runs_is_empty()
    {
        var daily = await _usage.GetDailyAsync(DateOnly.FromDateTime(Tag), "a", default);

        daily.RunCount.ShouldBe(0);
        daily.CostUsd.ShouldBe(0m);
        daily.CostIsComplete.ShouldBeTrue();
    }

    [Fact]
    public async Task Old_entries_can_be_removed()
    {
        await Record(1m, when: Tag.AddDays(-40));
        await Record(2m, when: Tag);

        var removed = await _usage.PurgeBeforeAsync(DateOnly.FromDateTime(Tag.AddDays(-30)), default);

        removed.ShouldBe(1);
        (await _usage.GetRecentAsync(10, default)).Count.ShouldBe(1);
    }

    [Fact]
    public async Task Many_runs_can_be_recorded_concurrently()
    {
        // The old way rewrote TokenUsage.json completely on every run, under a
        // global lock — exactly that did not scale.
        await Task.WhenAll(Enumerable.Range(0, 100).Select(i => Record(0.01m, agentId: $"agent-{i % 8}")));

        var daily = await _usage.GetDailyAsync(DateOnly.FromDateTime(Tag), agentId: "", default);

        daily.RunCount.ShouldBe(100);
        daily.CostUsd.ShouldBe(1.00m);
    }
}