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 aus der Bestandsaufnahme: 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.
/// </summary>
public sealed class BudgetGuardTests : IDisposable
{
    private readonly string _directory;
    private readonly SqliteUsageRepository _usage;
    private readonly DateTime _heute = 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()
    {
        Microsoft.Data.Sqlite.SqliteConnection.ClearAllPools();
        try { Directory.Delete(_directory, recursive: true); }
        catch { /* Aufräumen ist Nebensache */ }
    }

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

    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 ?? _heute
        }, default);

    // ═══════════════════════════════════════════════════════════
    //  Ohne Grenzen
    // ═══════════════════════════════════════════════════════════

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

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

        status.IsAllowed.ShouldBeTrue();
    }

    // ═══════════════════════════════════════════════════════════
    //  Kostengrenze
    // ═══════════════════════════════════════════════════════════

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

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

        status.IsAllowed.ShouldBeTrue();
    }

    [Fact]
    public async Task Bei_erreichter_Kostengrenze_wird_kein_Lauf_mehr_gestartet()
    {
        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-Grenze — greift auch ohne Preisangaben
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task Die_Tokengrenze_wirkt_auch_wenn_keine_Preise_vorliegen()
    {
        // Der eigentliche Zweck der zweiten Grenze: Bei unbekanntem Modell bleibt die
        // Kostensumme 0 und eine reine Kostengrenze wäre wirkungslos.
        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 Fehlende_Preise_werden_in_der_Begruendung_erwaehnt()
    {
        // Sonst wirkte ein niedriger Verbrauch wie ein noch offener Spielraum.
        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");
    }

    // ═══════════════════════════════════════════════════════════
    //  Instanzgrenze
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task Die_Instanzgrenze_zaehlt_alle_Agenten_zusammen()
    {
        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 Ein_einzelner_Agent_bleibt_unter_der_Instanzgrenze_erlaubt()
    {
        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();
    }

    // ═══════════════════════════════════════════════════════════
    //  Tageswechsel
    // ═══════════════════════════════════════════════════════════

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

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

        status.IsAllowed.ShouldBeTrue();
    }

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

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

        status.IsAllowed.ShouldBeTrue();
    }

    // ═══════════════════════════════════════════════════════════
    //  Restbudget
    // ═══════════════════════════════════════════════════════════

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

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

        remaining.ShouldBe(3.75m);
    }

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

    [Fact]
    public async Task Das_Restbudget_wird_nicht_negativ()
    {
        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()
    {
        Microsoft.Data.Sqlite.SqliteConnection.ClearAllPools();
        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 Betraege_ueberstehen_den_Rundlauf_ohne_Rundungsfehler()
    {
        // Als Fliesskommazahl abgelegt wuerden sich Cent-Betraege ueber tausende
        // Laeufe zu einem sichtbaren Fehler summieren.
        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 Der_Tagesverbrauch_summiert_Tokens_und_Laeufe()
    {
        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 Ohne_Agentenfilter_zaehlt_die_ganze_Instanz()
    {
        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 Ein_Lauf_ohne_Preisangabe_macht_den_Tag_unvollstaendig()
    {
        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 Ein_Tag_ohne_Laeufe_ist_leer()
    {
        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 Alte_Eintraege_lassen_sich_entfernen()
    {
        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 Viele_Laeufe_koennen_gleichzeitig_erfasst_werden()
    {
        // Der alte Weg schrieb TokenUsage.json bei jedem Lauf komplett neu, unter einem
        // globalen Lock — genau das skalierte nicht.
        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);
    }
}