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

using ClawdDotNet.Core.Accounting;
using ClawdDotNet.Core.Config;
using ClawdDotNet.Core.Engine;
using ClawdDotNet.Core.Tests.Infrastructure;
using Shouldly;

namespace ClawdDotNet.Core.Tests.Budget;

/// <summary>
/// #608: A run rejected because of an exhausted budget left only a log warning -
/// no receipt, no indication on the task. In test operation all
/// weather data was therefore missing from 18:20 on, without it being visible outside the log.
/// </summary>
public sealed class BudgetBlockVisibleTests
{
    /// <summary>Reports for every agent a daily consumption above every usual limit.</summary>
    private sealed class ExhaustedBudget : IUsageRepository
    {
        public Task RecordAsync(RunUsage usage, CancellationToken ct) => Task.CompletedTask;

        public Task<DailyUsage> GetDailyAsync(DateOnly date, string agentId, CancellationToken ct)
            => Task.FromResult(new DailyUsage(date, 400_000, 20_000, 0.408m, true, 28));

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

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

    private static (EngineFixture Fixture, InMemoryAuditRepository Audit, AgentConfig Agent) Setup()
    {
        var audit = new InMemoryAuditRepository();
        var fixture = new EngineFixture(audit, usage: new ExhaustedBudget());
        var agent = fixture.AddAgent("wetterdienst");
        agent.Budget.DailyCostUsd = 0.40m;
        return (fixture, audit, agent);
    }

    [Fact]
    public async Task Rejected_fresh_task_run_leaves_a_receipt_on_the_task()
    {
        var (fixture, audit, agent) = Setup();

        var result = await fixture.Engine.RunAsync(
            agent, "Hole die Warnungen", "inst", CancellationToken.None, source: "task", taskId: "t-unwetter");

        result.Status.ShouldBe(AgentRunStatus.BudgetExceeded);
        var receipt = audit.Receipts.ShouldHaveSingleItem();
        receipt.TaskId.ShouldBe("t-unwetter");
        receipt.AgentId.ShouldBe("wetterdienst");
        receipt.Status.ShouldBe(nameof(AgentRunStatus.BudgetExceeded));
        receipt.ResultRef.ShouldContain("Budget erschöpft");
        receipt.ResultRef.ShouldContain("0,40", Case.Insensitive);
        receipt.PromptTokens.ShouldBe(0);
        receipt.CostUsd.ShouldBe(0m);
    }

    [Fact]
    public async Task Rejected_run_in_chat_history_also_leaves_a_receipt()
    {
        var (fixture, audit, agent) = Setup();

        var result = await fixture.Engine.ChatAsync(
            agent, "Digest", "inst", CancellationToken.None, source: "task", taskId: "t-digest");

        result.Status.ShouldBe(AgentRunStatus.BudgetExceeded);
        audit.Receipts.ShouldHaveSingleItem().TaskId.ShouldBe("t-digest");
    }

    [Fact]
    public async Task Every_further_rejection_of_the_day_stands_as_its_own_receipt()
    {
        var (fixture, audit, agent) = Setup();

        for (var i = 0; i < 3; i++)
            await fixture.Engine.RunAsync(agent, "Los", "inst", CancellationToken.None, source: "task", taskId: "t-" + i);

        audit.Receipts.Select(r => r.TaskId).ShouldBe(["t-0", "t-1", "t-2"]);
    }
}