ClawdDotNet/tests/ClawdDotNet.Core.Tests/Engine/TokenAccountingTests.cs

using ClawdDotNet.Core.Api.Models;
using ClawdDotNet.Core.Config;
using ClawdDotNet.Core.Engine;
using ClawdDotNet.Core.Tests.Infrastructure;
using Shouldly;

namespace ClawdDotNet.Core.Tests.Engine;

public sealed class TokenAccountingTests
{
    // ═══════════════════════════════════════════════════════════
    //  B4 — record prompt and completion separately
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task Prompt_and_completion_are_summed_separately()
    {
        // So far the UI estimated 50/50. In reality the ratio is more like 95:5 —
        // and since output tokens cost a multiple, the cost display was
        // off by several times.
        var fixture = new EngineFixture().WithTool(FakeTool.Returning("ok"));
        var agent = fixture.AddAgent("agent-tokens", "TestTool");

        fixture.Client
            .RespondsWithText("Fertig", new Usage
            {
                PromptTokens = 9_500,
                CompletionTokens = 500,
                TotalTokens = 10_000
            });

        var result = await fixture.Engine.ChatAsync(agent, "Frage", "test-instance", default);

        result.PromptTokens.ShouldBe(9_500);
        result.CompletionTokens.ShouldBe(500);
        result.TokensUsed.ShouldBe(10_000);
    }

    [Fact]
    public async Task Across_several_steps_it_sums_correctly()
    {
        var fixture = new EngineFixture().WithTool(FakeTool.Returning("ok"));
        var agent = fixture.AddAgent("agent-sum", "TestTool");

        fixture.Client
            .RespondsWithToolCall("TestTool")        // 100 / 20 / 120 according to the fake
            .RespondsWithToolCall("TestTool")
            .RespondsWithText("Fertig");

        var result = await fixture.Engine.ChatAsync(agent, "Frage", "test-instance", default);

        result.StepCount.ShouldBe(3);
        result.PromptTokens.ShouldBe(300);
        result.CompletionTokens.ShouldBe(60);
        result.TokensUsed.ShouldBe(360);
    }

    // ═══════════════════════════════════════════════════════════
    //  T9 — make the cache effect measurable
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task Cached_tokens_are_passed_through()
    {
        // Without this number it could not be proven whether prompt caching works —
        // a silently ineffective cache would otherwise go unnoticed.
        var fixture = new EngineFixture();
        var agent = fixture.AddAgent("agent-cached");

        fixture.Client.RespondsWithText("Fertig", new Usage
        {
            PromptTokens = 20_000,
            CompletionTokens = 300,
            TotalTokens = 20_300,
            PromptTokensDetails = new PromptTokensDetails { CachedTokens = 18_000 }
        });

        var result = await fixture.Engine.ChatAsync(agent, "Frage", "test-instance", default);

        result.CachedTokens.ShouldBe(18_000);
    }

    [Fact]
    public void Missing_cache_values_give_null_instead_of_an_error()
    {
        var usage = new Usage { PromptTokens = 100, CompletionTokens = 10, TotalTokens = 110 };

        usage.CachedTokens.ShouldBe(0);
    }

    // ═══════════════════════════════════════════════════════════
    //  T3 — compaction runs with the cheap model
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task The_summary_uses_the_cheap_model()
    {
        var client = new FakeChatClient().AlwaysRespondsWithText("- Zusammenfassung.");
        var compactor = new ContextCompactor(client, TestLogging.Factory);

        var guard = new LoopGuardConfig
        {
            MaxContextTokens = 100,
            CompactionThreshold = 0.5,
            SummaryModel = "google/gemini-2.5-flash"
        };

        var messages = Conversation.Start().Repeat(10).Build();

        await compactor.CompactIfNeededAsync(messages, guard, "anthropic/claude-opus-4", default);

        client.ReceivedRequests.ShouldNotBeEmpty();
        client.ReceivedRequests[0].Model.ShouldBe("google/gemini-2.5-flash",
            "Zusammenfassen ist anspruchslos und darf nicht das teure Agentenmodell belegen");
    }

    [Fact]
    public async Task Without_configured_SummaryModel_the_agent_model_is_used()
    {
        var client = new FakeChatClient().AlwaysRespondsWithText("- Zusammenfassung.");
        var compactor = new ContextCompactor(client, TestLogging.Factory);

        var guard = new LoopGuardConfig
        {
            MaxContextTokens = 100,
            CompactionThreshold = 0.5,
            SummaryModel = ""
        };

        var messages = Conversation.Start().Repeat(10).Build();

        await compactor.CompactIfNeededAsync(messages, guard, "anthropic/claude-opus-4", default);

        client.ReceivedRequests[0].Model.ShouldBe("anthropic/claude-opus-4");
    }

    // ═══════════════════════════════════════════════════════════
    //  Only the part that goes away is summarized
    // ═══════════════════════════════════════════════════════════

    [Fact]
    public async Task The_kept_tail_does_not_go_into_the_summary_call()
    {
        // The tail is kept verbatim — summarizing it additionally
        // would be context paid for twice.
        var client = new FakeChatClient().AlwaysRespondsWithText("- Zusammenfassung.");
        var compactor = new ContextCompactor(client, TestLogging.Factory);

        var guard = new LoopGuardConfig { MaxContextTokens = 100, CompactionThreshold = 0.5 };

        var messages = Conversation.Start()
            .Repeat(8)
            .User("EINZIGARTIGE-LETZTE-NACHRICHT")
            .Build();

        await compactor.CompactIfNeededAsync(messages, guard, "test/model", default);

        var summaryPrompt = client.ReceivedRequests[0].Messages.Last().Content!;
        summaryPrompt.ShouldNotContain("EINZIGARTIGE-LETZTE-NACHRICHT");
    }
}