Added the exact token count where the provider reports it (#989)

Co-authored-by: Thorsten Sommer <SommerEngineering@users.noreply.github.com>

app/Tests/Chat/ChatThreadReportedHistoryTests.cs

@@ -0,0 +1,184 @@
+using AIStudio.Chat;
+using AIStudio.Provider;
+
+namespace AIStudio.Tests.Chat;
+
+/// <summary>
+/// Checks when what a provider reported for a conversation still describes it.
+/// </summary>
+/// <remarks>
+/// A reported number is shown as exact, so a wrong one does more harm than no number at all: it
+/// tells a person that their chat fits while it does not. Every case below is one where the thread
+/// moved on after the report, and the estimate has to take over -- or one where it came back to
+/// exactly the state the report was taken in, and the report counts again.
+/// </remarks>
+[TestFixture]
+public sealed class ChatThreadReportedHistoryTests
+{
+    private static readonly DateTimeOffset START = new(2026, 9, 23, 10, 0, 0, TimeSpan.Zero);
+
+    private static readonly Model MODEL = new("model-a", null);
+
+    private static readonly Model OTHER_MODEL = new("model-b", null);
+
+    [Test]
+    public void TheReportOfTheLastAnswerCounts()
+    {
+        var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2));
+
+        var history = thread.ReportedHistoryFor(MODEL);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(history.IsKnown, Is.True);
+            Assert.That(history.PromptTokens, Is.EqualTo(1200));
+
+            //
+            // The answer comes along as text rather than as the provider's number for it, which
+            // would include the reasoning the next request never carries:
+            //
+            Assert.That(history.LastAnswer, Is.EqualTo("Answer 2"));
+        });
+    }
+
+    [Test]
+    public void AnAnswerWithoutTextAddsNothingToThePrompt()
+    {
+        //
+        // An answer which consists of reasoning only. It may be kept to be read, but without any
+        // text it is never sent, so what the provider counted is all the next request carries of
+        // the conversation so far.
+        //
+        var reasoningOnly = new ContentText { Text = string.Empty };
+        reasoningOnly.RecordReportedUsage(TokenUsage.Of(1200), MODEL.Id, 2);
+        var thread = Thread(Question(1), Block(ChatRole.AI, reasoningOnly, 2));
+
+        var history = thread.ReportedHistoryFor(MODEL);
+        Assert.Multiple(() =>
+        {
+            Assert.That(history.IsKnown, Is.True);
+            Assert.That(history.PromptTokens, Is.EqualTo(1200));
+            Assert.That(history.LastAnswer, Is.Empty);
+        });
+    }
+
+    [Test]
+    public void AMessageWhoseRequestFailedIsNotLeftOut()
+    {
+        //
+        // A request which was turned down leaves the question standing and removes the empty
+        // answer. The report before it knows nothing about that question -- which may well be the
+        // very message that made the chat too large.
+        //
+        var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), Question(3));
+
+        Assert.That(thread.ReportedHistoryFor(MODEL).IsKnown, Is.False);
+    }
+
+    [Test]
+    public void AnAnswerBeingWrittenDoesNotBorrowTheReportBeforeIt()
+    {
+        var streaming = Block(ChatRole.AI, new ContentText { Text = "Half an ans", IsStreaming = true }, 4);
+        var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), Question(3), streaming);
+
+        Assert.That(thread.ReportedHistoryFor(MODEL).IsKnown, Is.False);
+    }
+
+    [Test]
+    public void AnAnswerWithoutAReportDoesNotBorrowTheReportBeforeIt()
+    {
+        //
+        // What an answer looks like whose provider reports nothing, or whose API is not read for a
+        // report yet: finished, but without a report of its own.
+        //
+        var withoutReport = Block(ChatRole.AI, new ContentText { Text = "Second answer" }, 4);
+        var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), Question(3), withoutReport);
+
+        Assert.That(thread.ReportedHistoryFor(MODEL).IsKnown, Is.False);
+    }
+
+    [Test]
+    public void DeletingAnEarlierMessageOutdatesTheReport()
+    {
+        //
+        // Deleting a large message to make room is exactly when somebody watches the number, and
+        // the deleted message is still inside what the provider reported.
+        //
+        var firstQuestion = Question(1);
+        var thread = Thread(firstQuestion, Answer(2, promptTokens: 1200, blockCount: 2), Question(3), Answer(4, promptTokens: 2600, blockCount: 4));
+
+        thread.Remove(firstQuestion.Content!);
+
+        Assert.That(thread.ReportedHistoryFor(MODEL).IsKnown, Is.False);
+    }
+
+    [Test]
+    public void AnotherModelOutdatesTheReport()
+    {
+        var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2));
+
+        Assert.That(thread.ReportedHistoryFor(OTHER_MODEL).IsKnown, Is.False);
+    }
+
+    [Test]
+    public void EditingTheLastMessageBringsTheReportBeforeItBack()
+    {
+        var lastQuestion = Question(3);
+        var lastAnswer = Answer(4, promptTokens: 2600, blockCount: 4);
+        var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2), lastQuestion, lastAnswer);
+
+        //
+        // What editing the last message does to the thread: the question goes back into the
+        // composer, and its answer goes with it.
+        //
+        thread.Remove(lastQuestion.Content!);
+        thread.Remove(lastAnswer.Content!);
+
+        var history = thread.ReportedHistoryFor(MODEL);
+        Assert.Multiple(() =>
+        {
+            Assert.That(history.IsKnown, Is.True, "The first answer is the last block again, with exactly the blocks it was reported for.");
+            Assert.That(history.PromptTokens, Is.EqualTo(1200));
+            Assert.That(history.LastAnswer, Is.EqualTo("Answer 2"));
+        });
+    }
+
+    [Test]
+    public void RollingBackToAnAnswerBringsItsReportBack()
+    {
+        var firstAnswer = Answer(2, promptTokens: 1200, blockCount: 2);
+        var thread = Thread(Question(1), firstAnswer, Question(3), Answer(4, promptTokens: 2600, blockCount: 4));
+
+        thread.RollBackTo(firstAnswer.Content!);
+
+        var history = thread.ReportedHistoryFor(MODEL);
+        Assert.Multiple(() =>
+        {
+            Assert.That(history.IsKnown, Is.True);
+            Assert.That(history.PromptTokens, Is.EqualTo(1200));
+            Assert.That(history.LastAnswer, Is.EqualTo("Answer 2"));
+        });
+    }
+
+    private static ChatThread Thread(params ContentBlock[] blocks) => new()
+    {
+        Blocks = [..blocks],
+    };
+
+    private static ContentBlock Question(int minute) => Block(ChatRole.USER, new ContentText { Text = $"Question {minute}" }, minute);
+
+    private static ContentBlock Answer(int minute, int promptTokens, int blockCount)
+    {
+        var answer = new ContentText { Text = $"Answer {minute}" };
+        answer.RecordReportedUsage(TokenUsage.Of(promptTokens), MODEL.Id, blockCount);
+        return Block(ChatRole.AI, answer, minute);
+    }
+
+    private static ContentBlock Block(ChatRole role, ContentText content, int minute) => new()
+    {
+        Time = START.AddMinutes(minute),
+        ContentType = ContentType.TEXT,
+        Content = content,
+        Role = role,
+    };
+}

app/Tests/Chat/ConversationPartsTests.cs

@@ -56,7 +56,8 @@ public sealed class ConversationPartsTests
         Assert.Multiple(() =>
         {
             Assert.That(parts.Texts, Is.EqualTo(new[] { "You are helpful.", "What is the capital of France?", "Paris." }));
-            Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "And of Italy?" }));
+            Assert.That(parts.DraftText, Is.EqualTo("And of Italy?"));
+            Assert.That(parts.GrowingTexts, Is.Empty, "The draft is a part of its own, and it stands nowhere else.");
         });
     }
 
@@ -77,7 +78,8 @@ public sealed class ConversationPartsTests
         Assert.Multiple(() =>
         {
             Assert.That(parts.Texts, Is.EqualTo(new[] { "A question." }));
-            Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "The answer so far", "a draft" }));
+            Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "The answer so far" }));
+            Assert.That(parts.DraftText, Is.EqualTo("a draft"));
         });
     }
 
@@ -135,7 +137,7 @@ public sealed class ConversationPartsTests
         Assert.Multiple(() =>
         {
             Assert.That(parts.Texts, Is.Empty);
-            Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "Hello" }));
+            Assert.That(parts.DraftText, Is.EqualTo("Hello"));
         });
     }
 
@@ -149,6 +151,7 @@ public sealed class ConversationPartsTests
         {
             Assert.That(parts.Texts, Is.Empty);
             Assert.That(parts.GrowingTexts, Is.Empty);
+            Assert.That(parts.DraftText, Is.Empty);
         });
     }
 
@@ -188,7 +191,11 @@ public sealed class ConversationPartsTests
 
         var parts = ConversationParts.Of(new() { Blocks = [block] }, string.Empty, "And this one.", [FileAttachment.FromPath(draft)], imagesAreSent: true, toolDefinitions: null);
 
-        Assert.That(parts.Documents.Select(document => document.FileName), Is.EqualTo(new[] { "older.txt", "draft.txt" }));
+        Assert.Multiple(() =>
+        {
+            Assert.That(parts.Documents.Select(document => document.FileName), Is.EqualTo(new[] { "older.txt" }));
+            Assert.That(parts.DraftDocuments.Select(document => document.FileName), Is.EqualTo(new[] { "draft.txt" }), "Both count, each in the part it belongs to.");
+        });
     }
 
     [Test]
@@ -201,7 +208,7 @@ public sealed class ConversationPartsTests
 
         var parts = ConversationParts.Of(null, string.Empty, "Here", [attachment], imagesAreSent: true, toolDefinitions: null);
 
-        Assert.That(parts.Documents, Is.Empty);
+        Assert.That(parts.DraftDocuments, Is.Empty);
     }
 
     [Test]
@@ -214,8 +221,29 @@ public sealed class ConversationPartsTests
 
         Assert.Multiple(() =>
         {
-            Assert.That(parts.Documents.Select(entry => entry.FileName), Is.EqualTo(new[] { "notes.txt" }));
+            Assert.That(parts.DraftDocuments.Select(entry => entry.FileName), Is.EqualTo(new[] { "notes.txt" }));
+            Assert.That(parts.DraftImages, Is.EqualTo(1));
+        });
+    }
+
+    [Test]
+    public void ImagesOfTheConversationAndOfTheDraftAreKeptApart()
+    {
+        //
+        // A provider which reported the conversation so far counted the pictures in it as well,
+        // but never the ones which are still waiting in the composer.
+        //
+        var sent = this.WriteFile("sent.png", "not really a png");
+        var waiting = this.WriteFile("waiting.png", "not really a png either");
+        var block = Block("Look at this.");
+        ((ContentText)block.Content!).FileAttachments.Add(FileAttachment.FromPath(sent));
+
+        var parts = ConversationParts.Of(new() { Blocks = [block] }, string.Empty, "And at this.", [FileAttachment.FromPath(waiting)], imagesAreSent: true, toolDefinitions: null);
+
+        Assert.Multiple(() =>
+        {
             Assert.That(parts.Images, Is.EqualTo(1));
+            Assert.That(parts.DraftImages, Is.EqualTo(1));
         });
     }
 
@@ -230,7 +258,7 @@ public sealed class ConversationPartsTests
 
         var parts = ConversationParts.Of(null, string.Empty, "Look", [FileAttachment.FromPath(image)], imagesAreSent: false, toolDefinitions: null);
 
-        Assert.That(parts.Images, Is.Zero);
+        Assert.That(parts.DraftImages, Is.Zero);
     }
 
     [Test]
@@ -250,7 +278,8 @@ public sealed class ConversationPartsTests
         Assert.Multiple(() =>
         {
             Assert.That(parts.Texts, Is.Empty);
-            Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "What the web search found.", "What the page said." }));
+            Assert.That(parts.ToolConversation, Is.EqualTo(new[] { "What the web search found.", "What the page said." }));
+            Assert.That(parts.GrowingTexts, Is.Empty, "The tool conversation is a part of its own, so that its share can be named.");
         });
     }
 
@@ -267,7 +296,7 @@ public sealed class ConversationPartsTests
 
         var parts = ConversationParts.Of(new() { Blocks = [running] }, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);
 
-        Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "The same page.", "The same page." }));
+        Assert.That(parts.ToolConversation, Is.EqualTo(new[] { "The same page.", "The same page." }));
     }
 
     [Test]
@@ -288,7 +317,7 @@ public sealed class ConversationPartsTests
         Assert.Multiple(() =>
         {
             Assert.That(parts.Texts, Is.EqualTo(new[] { "Here is what I found." }));
-            Assert.That(parts.GrowingTexts, Is.Empty);
+            Assert.That(parts.ToolConversation, Is.Empty);
         });
     }
 

app/Tests/Chat/ConversationTokensTests.cs

@@ -25,7 +25,7 @@ public sealed class ConversationTokensTests
         var counted = new ConversationTokens
         {
             IsKnown = true,
-            UncountedImages = images,
+            Images = images,
             ImageLimits = new ImageLimits(null, allowed),
         };
 
@@ -42,7 +42,7 @@ public sealed class ConversationTokensTests
         var counted = new ConversationTokens
         {
             IsKnown = true,
-            UncountedImages = 500,
+            Images = 500,
             ImageLimits = ImageLimits.UNKNOWN,
         };
 
@@ -60,7 +60,7 @@ public sealed class ConversationTokensTests
         var counted = new ConversationTokens
         {
             IsKnown = true,
-            UncountedImages = 20,
+            Images = 20,
             ImageLimits = new ImageLimits(8, 100),
         };
 
@@ -77,7 +77,7 @@ public sealed class ConversationTokensTests
         var counted = new ConversationTokens
         {
             IsKnown = true,
-            UncountedImages = 0,
+            Images = 0,
             ImageLimits = new ImageLimits(null, 0),
         };
 
@@ -93,4 +93,64 @@ public sealed class ConversationTokensTests
         //
         Assert.That(ConversationTokens.UNAVAILABLE.TooManyImages, Is.False);
     }
+
+    [Test]
+    public void PicturesAProviderCountedAreNotCalledUncounted()
+    {
+        //
+        // What the provider reported for the conversation so far includes its pictures, however
+        // it charges them. Only those still waiting in the composer are left for nobody to count.
+        //
+        var reported = new ConversationTokens
+        {
+            IsKnown = true,
+            HistoryIsReported = true,
+            Images = 3,
+            DraftImages = 1,
+        };
+
+        var estimated = reported with { HistoryIsReported = false };
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(reported.UncountedImages, Is.EqualTo(1), "The provider counted the two which were sent.");
+            Assert.That(estimated.UncountedImages, Is.EqualTo(3), "Without a report, nobody counted any of them.");
+        });
+    }
+
+    [Test]
+    public void TooManyPicturesStaysTooManyWhenTheProviderCountedThem()
+    {
+        //
+        // The limit is on how many pictures travel, not on what they cost. A provider which has
+        // counted them still refuses the request which carries one too many.
+        //
+        var counted = new ConversationTokens
+        {
+            IsKnown = true,
+            HistoryIsReported = true,
+            Images = 10,
+            ImageLimits = new ImageLimits(8, null),
+        };
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(counted.UncountedImages, Is.Zero);
+            Assert.That(counted.TooManyImages, Is.True);
+        });
+    }
+
+    [Test]
+    public void TheWholeNumberIsTheConversationAndTheDraft()
+    {
+        var counted = new ConversationTokens
+        {
+            IsKnown = true,
+            HistoryTokens = 12_400,
+            ToolTokens = 9_000,
+            DraftTokens = 340,
+        };
+
+        Assert.That(counted.Tokens, Is.EqualTo(12_740), "The tools' share is part of the conversation, not added on top of it.");
+    }
 }

app/Tests/Chat/TokenAmountTests.cs

@@ -51,4 +51,4 @@ public sealed class TokenAmountTests
         //
         Assert.That(TokenAmount.Format(tokens, GERMAN), Is.EqualTo(wanted));
     }
-}
+}

app/Tests/Provider/ChatCompletionUsageTests.cs

@@ -0,0 +1,129 @@
+using System.Text.Json;
+
+using AIStudio.Provider;
+using AIStudio.Provider.OpenAI;
+
+namespace AIStudio.Tests.Provider;
+
+/// <summary>
+/// Checks that what a provider says a request cost is read off the stream, and asked for.
+/// </summary>
+/// <remarks>
+/// Both halves matter and neither is visible from the other: an OpenAI-compatible provider says
+/// nothing about the cost of a streamed request unless the request asks for it, and the line it
+/// then sends carries no content, so the reading side has to look for it apart from the text.
+/// Get either half wrong and the app silently keeps estimating, which looks exactly like a
+/// provider which reports nothing.
+/// </remarks>
+[TestFixture]
+public sealed class ChatCompletionUsageTests
+{
+    /// <summary>
+    /// The last line of a streamed answer at a provider which was asked for the usage.
+    /// </summary>
+    private const string USAGE_LINE =
+        """
+        {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[],"usage":{"prompt_tokens":1200,"completion_tokens":345,"total_tokens":1545}}
+        """;
+
+    /// <summary>
+    /// An ordinary line carrying a piece of the answer.
+    /// </summary>
+    private const string CONTENT_LINE =
+        """
+        {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{"content":"Hi"}}]}
+        """;
+
+    [Test]
+    public void TheFinalLineStatesWhatTheRequestCost()
+    {
+        var line = JsonSerializer.Deserialize<ChatCompletionDeltaStreamLine>(USAGE_LINE, ProviderJsonOptions.OPTIONS);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(line!.GetUsage().IsKnown, Is.True);
+            Assert.That(line.GetUsage().PromptTokens, Is.EqualTo(1200));
+
+            //
+            // The line which carries the usage carries no answer, which is why it has to be read
+            // before the content check drops it:
+            //
+            Assert.That(line.ContainsContent(), Is.False);
+        });
+    }
+
+    [Test]
+    public void ALineOfTheAnswerStatesNoCost()
+    {
+        var line = JsonSerializer.Deserialize<ChatCompletionDeltaStreamLine>(CONTENT_LINE, ProviderJsonOptions.OPTIONS);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(line!.GetUsage().IsKnown, Is.False);
+            Assert.That(line.ContainsContent(), Is.True);
+        });
+    }
+
+    [Test]
+    public void AStreamedRequestAsksForTheUsage()
+    {
+        var request = new ChatCompletionAPIRequest("gpt-5", [], true);
+        var json = JsonSerializer.Serialize(request, ProviderJsonOptions.OPTIONS);
+
+        Assert.That(json, Does.Contain("""
+                                       "stream_options":{"include_usage":true}
+                                       """));
+    }
+
+    [Test]
+    public void ARequestWhichIsNotStreamedDoesNot()
+    {
+        var request = new ChatCompletionAPIRequest("gpt-5", [], false);
+        var json = JsonSerializer.Serialize(request, ProviderJsonOptions.OPTIONS);
+
+        Assert.That(json, Does.Not.Contain("stream_options"));
+    }
+
+    /// <summary>
+    /// A provider which sends the block but fills in nothing usable states nothing.
+    /// </summary>
+    /// <remarks>
+    /// Several OpenAI-compatible servers send an empty or zeroed usage block on every line while
+    /// streaming and the real numbers only at the end. Reading a zero as a fact would replace an
+    /// estimate with a statement that the conversation costs nothing.
+    /// </remarks>
+    [Test]
+    public void AnEmptyUsageBlockStatesNothing()
+    {
+        var line = JsonSerializer.Deserialize<ChatCompletionDeltaStreamLine>(
+            """
+            {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[],"usage":{"prompt_tokens":0,"completion_tokens":0}}
+            """, ProviderJsonOptions.OPTIONS);
+
+        Assert.That(line!.GetUsage().IsKnown, Is.False);
+    }
+
+    /// <summary>
+    /// The line a real LM Studio server sends, taken off the wire.
+    /// </summary>
+    /// <remarks>
+    /// It carries far more than we read, and it shows why the prompt is all we take: 102 of the 116
+    /// completion tokens are the model's reasoning, which the next request never carries. Counting
+    /// the completion would have put the history at 133 tokens, when what travels on is the prompt
+    /// and an answer of a handful of tokens.
+    /// </remarks>
+    [Test]
+    public void ARealServerLineIsRead()
+    {
+        var line = JsonSerializer.Deserialize<ChatCompletionDeltaStreamLine>(
+            """
+            {"id":"chatcmpl-xb8mn282eff3tiu46xz8t3","object":"chat.completion.chunk","created":1789917744,"model":"google/gemma-4-12b-qat","system_fingerprint":"google/gemma-4-12b-qat","choices":[],"usage":{"prompt_tokens":17,"completion_tokens":116,"total_tokens":133,"completion_tokens_details":{"reasoning_tokens":102}}}
+            """, ProviderJsonOptions.OPTIONS);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(line!.GetUsage().IsKnown, Is.True);
+            Assert.That(line.GetUsage().PromptTokens, Is.EqualTo(17));
+        });
+    }
+}

app/Tests/Provider/ToolCalling/ChatCompletionToolCallAccumulatorTests.cs

@@ -183,7 +183,38 @@ public sealed class ChatCompletionToolCallAccumulatorTests
         
         Assert.That(part.Sources.Select(x => x.URL), Is.EqualTo(new[] { "https://example.org/" }), "Whatever the provider announced on that line reaches the user with it.");
     }
-    
+
+    [Test]
+    public void TheLineWithoutChoicesStatesWhatTheRequestCost()
+    {
+        var accumulator = new ChatCompletionToolCallAccumulator();
+        var part = accumulator.Process(Event("""{"choices":[],"usage":{"prompt_tokens":1200,"completion_tokens":345,"total_tokens":1545}}"""));
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(part.Usage.IsKnown, Is.True, "The last line of the stream has no choices, and it must not be dropped for that.");
+            Assert.That(part.Usage.PromptTokens, Is.EqualTo(1200));
+            Assert.That(part.HasContent, Is.False, "It has nothing to show, though.");
+        });
+    }
+
+    [Test]
+    public void AUsageNextToTheLastTextIsReadAsWell()
+    {
+        //
+        // Some providers put the usage on the line which carries the last piece of the answer
+        // rather than on a line of its own.
+        //
+        var accumulator = new ChatCompletionToolCallAccumulator();
+        var part = accumulator.Process(Event("""{"choices":[{"index":0,"delta":{"content":"Bye"}}],"usage":{"prompt_tokens":1200,"completion_tokens":2}}"""));
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(part.TextDelta, Is.EqualTo("Bye"));
+            Assert.That(part.Usage.PromptTokens, Is.EqualTo(1200));
+        });
+    }
+
     private static ChatCompletionResponseMessage? Read(params string[] data)
     {
         var accumulator = new ChatCompletionToolCallAccumulator();

app/Tests/Provider/ToolCalling/ChatCompletionToolCallingAdapterTests.cs

@@ -0,0 +1,169 @@
+using System.Runtime.CompilerServices;
+using System.Text.Json;
+
+using AIStudio.Provider;
+using AIStudio.Provider.OpenAI;
+
+using Microsoft.Extensions.Logging.Abstractions;
+
+namespace AIStudio.Tests.Provider.ToolCalling;
+
+/// <summary>
+/// Checks what a round of a tool calling conversation asks for, and what it passes on.
+/// </summary>
+/// <remarks>
+/// Every round of a tool conversation is a request of its own, and every one of them reports what
+/// it cost. Only the first one describes what the next question will be sent after: every later
+/// round carries the tool calls and their results on top, none of which is sent again once the
+/// answer stands. Passing on the last report instead would put the chat at the size of everything
+/// the tools returned, which is the one number a person watching their context window must not
+/// see as exact.
+///
+/// What a round asks for is one tool call at a time, wherever the provider lets it ask: a provider
+/// which rejects the question fails the whole request, so it is not asked at all.
+/// </remarks>
+[TestFixture]
+public sealed class ChatCompletionToolCallingAdapterTests
+{
+    private const string FIRST_ROUND_USAGE = """{"choices":[],"usage":{"prompt_tokens":1200,"completion_tokens":20}}""";
+
+    private const string SECOND_ROUND_USAGE = """{"choices":[],"usage":{"prompt_tokens":9800,"completion_tokens":150}}""";
+
+    [Test]
+    public async Task OnlyTheFirstRoundPassesOnWhatItsRequestCost()
+    {
+        var adapter = Adapter(
+        [
+            """{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"web_search","arguments":"{\"query\":\"weather\"}"}}]}}]}""",
+            FIRST_ROUND_USAGE,
+            "[DONE]",
+        ],
+        [
+            """{"choices":[{"index":0,"delta":{"content":"It is sunny."}}]}""",
+            SECOND_ROUND_USAGE,
+            "[DONE]",
+        ]);
+
+        var firstRound = await Usages(adapter);
+
+        //
+        // What the loop does between two rounds: the model's turn and the tool's result become part
+        // of the next request.
+        //
+        adapter.RecordAssistantTurn();
+        adapter.RecordToolResult("call_1", "Sunny, 24 degrees.");
+
+        var secondRound = await Usages(adapter);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(firstRound, Is.EqualTo(new[] { 1200 }), "The first round's prompt is the conversation up to the question.");
+            Assert.That(secondRound, Is.Empty, "The second round's prompt holds the tool result as well, which the next question is not sent with.");
+        });
+    }
+
+    [Test]
+    public async Task ARoundWithoutToolCallsPassesItOnAsWell()
+    {
+        //
+        // Offering tools does not mean the model uses them. Then the first round is the only one,
+        // and its report is as good as the one of a request which offered none.
+        //
+        var adapter = Adapter(
+        [
+            """{"choices":[{"index":0,"delta":{"content":"Hello."}}]}""",
+            FIRST_ROUND_USAGE,
+            "[DONE]",
+        ]);
+
+        Assert.That(await Usages(adapter), Is.EqualTo(new[] { 1200 }));
+    }
+
+    [Test]
+    public async Task ARoundWhichOffersToolsAsksForOneCallAtATime()
+    {
+        Assert.That(await SentRequest(mayAskForSequentialToolCalls: true, includeTools: true), Does.Contain("\"parallel_tool_calls\":false"));
+    }
+
+    [Test]
+    public async Task AProviderWhichRejectsTheQuestionIsNotAskedIt()
+    {
+        //
+        // Hugging Face answers the question with a bad request. Its models may then ask for several
+        // calls at once, which the loop works through one by one anyway.
+        //
+        Assert.That(await SentRequest(mayAskForSequentialToolCalls: false, includeTools: true), Does.Not.Contain("parallel_tool_calls"));
+    }
+
+    [Test]
+    public async Task ARoundWithoutToolsDoesNotAskAboutToolCalls()
+    {
+        Assert.That(await SentRequest(mayAskForSequentialToolCalls: true, includeTools: false), Does.Not.Contain("parallel_tool_calls"));
+    }
+
+    /// <summary>
+    /// Runs one round and returns the request it sent, as it goes over the wire.
+    /// </summary>
+    private static async Task<string> SentRequest(bool mayAskForSequentialToolCalls, bool includeTools)
+    {
+        ChatCompletionAPIRequest? sent = null;
+        var adapter = Adapter(mayAskForSequentialToolCalls, request => sent = request, ["[DONE]"]);
+        await foreach (var _ in adapter.ExecuteRoundAsync(null, includeTools))
+        {
+        }
+
+        return JsonSerializer.Serialize(sent, ProviderJsonOptions.OPTIONS);
+    }
+
+    /// <summary>
+    /// Runs the next round and returns the prompt of every usage it passed on.
+    /// </summary>
+    private static async Task<List<int>> Usages(ChatCompletionToolCallingAdapter<ChatCompletionAPIRequest> adapter)
+    {
+        var usages = new List<int>();
+        await foreach (var streamEvent in adapter.ExecuteRoundAsync(null, true))
+            if (streamEvent.Delta is { Usage.IsKnown: true } delta)
+                usages.Add(delta.Usage.PromptTokens);
+
+        return usages;
+    }
+
+    /// <summary>
+    /// Builds an adapter whose requests are answered by the given rounds, one after another.
+    /// </summary>
+    private static ChatCompletionToolCallingAdapter<ChatCompletionAPIRequest> Adapter(params string[][] rounds) => Adapter(true, _ => { }, rounds);
+
+    /// <summary>
+    /// Builds an adapter whose requests are answered by the given rounds, and which hands every
+    /// request it sends to the given observer.
+    /// </summary>
+    private static ChatCompletionToolCallingAdapter<ChatCompletionAPIRequest> Adapter(bool mayAskForSequentialToolCalls, Action<ChatCompletionAPIRequest> sent, params string[][] rounds)
+    {
+        var nextRound = 0;
+        return new(
+            (_, _, tools) => Task.FromResult(new ChatCompletionAPIRequest("model-a", [], true) { Tools = tools }),
+            new TextMessage("You are a helpful assistant.", "system"),
+            new Dictionary<string, object>(),
+            [],
+            mayAskForSequentialToolCalls,
+            [],
+            (request, token) =>
+            {
+                sent(request);
+                return Lines(rounds[nextRound++], token);
+            },
+            _ => [],
+            NullLogger.Instance);
+    }
+
+    private static async IAsyncEnumerable<ServerSentEvent> Lines(string[] data, [EnumeratorCancellation] CancellationToken token = default)
+    {
+        foreach (var line in data)
+        {
+            token.ThrowIfCancellationRequested();
+            yield return new ServerSentEvent($"data: {line}", line);
+        }
+
+        await Task.CompletedTask;
+    }
+}