From c23ce9e12a637c437d0f0b8616baefc8209be601 Mon Sep 17 00:00:00 2001 From: Thorsten Sommer Date: Wed, 23 Sep 2026 18:36:25 +0200 Subject: [PATCH] Fixed the exact token count outliving the conversation it was reported for --- app/MindWork AI Studio/Chat/ChatThread.cs | 42 +++++ app/MindWork AI Studio/Chat/ContentText.cs | 27 ++- .../Chat/ReportedTokenUsage.cs | 14 ++ .../Components/ChatComponent.razor.cs | 42 +---- .../Tools/AIJobs/AIJobService.cs | 11 +- .../Chat/ChatThreadReportedUsageTests.cs | 155 ++++++++++++++++++ 6 files changed, 239 insertions(+), 52 deletions(-) create mode 100644 app/Tests/Chat/ChatThreadReportedUsageTests.cs diff --git a/app/MindWork AI Studio/Chat/ChatThread.cs b/app/MindWork AI Studio/Chat/ChatThread.cs index 466fbe44..7dbbe536 100644 --- a/app/MindWork AI Studio/Chat/ChatThread.cs +++ b/app/MindWork AI Studio/Chat/ChatThread.cs @@ -392,6 +392,48 @@ public sealed record ChatThread return true; } + /// + /// Finds what a provider reported for this conversation, as long as the report still describes it. + /// + /// + /// A report describes one request: the conversation up to the answer which carries it. It is + /// worth something only while the thread still is that conversation, so only the last block is + /// asked, and no earlier answer ever stands in for it. Whatever came after an older report -- a + /// message whose request was turned down, an answer which is still being written, an answer + /// without a report of its own -- is missing from that report's number, and the estimate is + /// closer to the truth than a figure which leaves it out. Answers written while tools were + /// offered carry no report yet, so for them the estimate always takes over.

+ /// + /// A report also stops counting when the thread holds a different number of blocks than the + /// request did, which means an earlier message was deleted, and when the next request goes to + /// another model, which counts the same conversation with another tokenizer. Editing the last + /// message or rolling the chat back makes an earlier answer the last block again, with exactly + /// the blocks it was reported for, so its report counts once more.

+ /// + /// What goes unnoticed is a change beside the messages: another system prompt, profile, or + /// selection of tools. That shows only with the next answer. Noticing it would take a + /// fingerprint of the system prompt as it was sent, after the data sources added to it, which + /// is a lot of machinery for a number which corrects itself one answer later. + /// + /// The model the next request would go to. + /// What the provider reported, or unknown when no report describes this conversation. + public TokenUsage ReportedUsageFor(Model model) + { + if (this.Blocks.Count is 0) + return TokenUsage.UNKNOWN; + + if (this.Blocks[^1].Content is not ContentText { IsStreaming: false, ReportedUsage: { } reported }) + return TokenUsage.UNKNOWN; + + if (reported.BlockCount != this.Blocks.Count) + return TokenUsage.UNKNOWN; + + if (!string.Equals(reported.ModelId, model.Id, StringComparison.Ordinal)) + return TokenUsage.UNKNOWN; + + return reported.ToTokenUsage(); + } + private static void DeleteManagedAttachments(ContentBlock block) { if (block.Content is not ContentText textContent) diff --git a/app/MindWork AI Studio/Chat/ContentText.cs b/app/MindWork AI Studio/Chat/ContentText.cs index 48e54c14..f74f3e9b 100644 --- a/app/MindWork AI Studio/Chat/ContentText.cs +++ b/app/MindWork AI Studio/Chat/ContentText.cs @@ -58,10 +58,12 @@ public sealed class ContentText : IContent /// /// Kept on the answer rather than beside the chat, so that it is stored, loaded, and exported /// with the message it belongs to -- and so that it goes away when the message does. An edited - /// or regenerated answer removes its block, which removes these numbers with it, and the chat - /// falls back to the estimate instead of carrying a figure for a conversation which no longer - /// exists. Null for every answer written before this was recorded, and at every provider which - /// reports nothing. + /// or regenerated answer removes its block, which removes these numbers with it. Null for every + /// answer written before this was recorded, and at every provider which reports nothing. + /// + /// Having a report is not the same as the report still being true. Whether it still describes + /// the conversation is decided by ChatThread.ReportedUsageFor, which looks at the thread around + /// the answer, not just at the answer. /// public ReportedTokenUsage? ReportedUsage { get; set; } @@ -121,7 +123,8 @@ public sealed class ContentText : IContent /// /// What the current chunk of the stream says, which is mostly nothing. /// The model the request went to. - public void RecordReportedUsage(TokenUsage usage, string modelId) + /// How many blocks the conversation had when the request went out, this answer included. + public void RecordReportedUsage(TokenUsage usage, string modelId, int blockCount) { if (!usage.IsKnown) return; @@ -131,6 +134,7 @@ public sealed class ContentText : IContent PromptTokens = usage.PromptTokens, CompletionTokens = usage.CompletionTokens, ModelId = modelId, + BlockCount = blockCount, }; } @@ -194,7 +198,10 @@ public sealed class ContentText : IContent // Get the settings manager: var settings = Program.SERVICE_PROVIDER.GetService()!; - + + // What the request carries, for telling later whether a report still describes this chat: + var blocksSent = chatThread.Blocks.Count; + // Start another thread by using a task to uncouple // the UI thread from the AI processing: try @@ -224,7 +231,7 @@ public sealed class ContentText : IContent this.Sources.MergeSources(contentStreamChunk.Sources); // Keep what the provider says the request cost: - this.RecordReportedUsage(contentStreamChunk.Usage, chatModel.Id); + this.RecordReportedUsage(contentStreamChunk.Usage, chatModel.Id, blocksSent); // Notify the UI that the content has changed, // depending on the energy saving mode: @@ -349,6 +356,12 @@ public sealed class ContentText : IContent } /// + /// + /// The reported usage stays behind on purpose. A clone continues somewhere else -- as the + /// example conversation of a chat template, or as an assistant's conversation carried over into + /// a chat -- with another system prompt around it, and what the provider stated was for the + /// request this answer came out of. + /// public IContent DeepClone() => new ContentText { Text = this.Text, diff --git a/app/MindWork AI Studio/Chat/ReportedTokenUsage.cs b/app/MindWork AI Studio/Chat/ReportedTokenUsage.cs index 90f7c5dd..1654ded6 100644 --- a/app/MindWork AI Studio/Chat/ReportedTokenUsage.cs +++ b/app/MindWork AI Studio/Chat/ReportedTokenUsage.cs @@ -37,6 +37,20 @@ public sealed record ReportedTokenUsage /// public string ModelId { get; init; } = string.Empty; + /// + /// How many blocks the conversation had when the request went out, this answer included. + /// + /// + /// What tells an outdated report apart without anybody having to remember anything about it. + /// Blocks are only ever added at the end, so while this answer is the last block, a thread with + /// the same count is the thread the provider saw. A lower count means an earlier message was + /// deleted, and that message is still inside the reported number. + /// + /// Taken when the request is sent rather than when the report arrives: whatever is deleted + /// while the answer streams in was still part of what the provider counted. + /// + public int BlockCount { get; init; } + /// /// States the stored numbers as a usage again. /// diff --git a/app/MindWork AI Studio/Components/ChatComponent.razor.cs b/app/MindWork AI Studio/Components/ChatComponent.razor.cs index 05ec1864..750e4e45 100644 --- a/app/MindWork AI Studio/Components/ChatComponent.razor.cs +++ b/app/MindWork AI Studio/Components/ChatComponent.razor.cs @@ -1614,7 +1614,7 @@ public partial class ChatComponent : MSGComponentBase var toolDefinitions = this.GetRunnableToolDefinitions(); provider = this.Provider; parts = ConversationParts.Of(thread, this.BuildSystemPromptFor(thread, toolDefinitions), this.UserInput, this.ComposerState.FileAttachments, provider.SupportsImageInput(), toolDefinitions); - reported = LastReportedTokensOf(thread, provider.Model); + reported = thread.ReportedUsageFor(provider.Model); }); var counted = await this.ConversationTokenCounter.CountAsync(provider, parts, reported, token); @@ -1631,46 +1631,6 @@ public partial class ChatComponent : MSGComponentBase }); } - /// - /// Finds what a provider last said a request of this conversation cost. - /// - /// - /// The last answer which carries such a number decides, and nothing else has to be remembered - /// for it: the number lives on the answer, so editing, regenerating, or deleting a message - /// takes it along and an earlier answer -- or none at all -- becomes the one which counts. - /// - /// An answer still being written is passed over. Its number arrives with the last line of the - /// stream, and until then it states what the request before it cost, which is a conversation - /// shorter than the one on the screen. - /// - /// The conversation to look through. - /// The model the next request would go to. - /// What the provider reported, or unknown when none of them did. - private static TokenUsage LastReportedTokensOf(ChatThread thread, Model model) - { - for (var index = thread.Blocks.Count - 1; index >= 0; index--) - { - if (thread.Blocks[index].Content is not ContentText { IsStreaming: false, ReportedUsage: { } reported }) - continue; - - var usage = reported.ToTokenUsage(); - if (!usage.IsKnown) - continue; - - // - // A number charged for another model says nothing about this one: another model counts - // the same conversation with another tokenizer. Reading on would only find older - // answers of that same other model, so the search ends here and the estimate takes - // over until this model has answered once. - // - return string.Equals(reported.ModelId, model.Id, StringComparison.Ordinal) - ? usage - : TokenUsage.UNKNOWN; - } - - return TokenUsage.UNKNOWN; - } - /// /// Works out the system prompt a thread would send. /// diff --git a/app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs b/app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs index 58b3644d..3f85d62c 100644 --- a/app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs +++ b/app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs @@ -300,11 +300,14 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes var lastStreamingEvent = DateTimeOffset.MinValue; if (!TrySetWaitingForRemote(state, token)) return; - + + // What the request carries, for telling later whether a report still describes this chat: + var blocksSent = chatThread.Blocks.Count; + await this.NotifyChangedAsync(state); await foreach (var contentStreamChunk in provider.StreamChatCompletion(request.ProviderSettings.Model, chatThread, settingsManager, token)) { - if (!TryApplyStreamChunk(state, contentStreamChunk, token)) + if (!TryApplyStreamChunk(state, contentStreamChunk, blocksSent, token)) break; var now = DateTimeOffset.Now; @@ -443,7 +446,7 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes } } - private static bool TryApplyStreamChunk(AIJobState state, ContentStreamChunk contentStreamChunk, CancellationToken token) + private static bool TryApplyStreamChunk(AIJobState state, ContentStreamChunk contentStreamChunk, int blocksSent, CancellationToken token) { lock (state.SyncRoot) { @@ -455,7 +458,7 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes aiText.IsStreaming = true; aiText.Text += contentStreamChunk; aiText.Sources.MergeSources(contentStreamChunk.Sources); - aiText.RecordReportedUsage(contentStreamChunk.Usage, state.ChatGenerationRequest.ProviderSettings.Model.Id); + aiText.RecordReportedUsage(contentStreamChunk.Usage, state.ChatGenerationRequest.ProviderSettings.Model.Id, blocksSent); if (state.Snapshot.Status is not AIJobStatus.RUNNING) { diff --git a/app/Tests/Chat/ChatThreadReportedUsageTests.cs b/app/Tests/Chat/ChatThreadReportedUsageTests.cs new file mode 100644 index 00000000..baa0663d --- /dev/null +++ b/app/Tests/Chat/ChatThreadReportedUsageTests.cs @@ -0,0 +1,155 @@ +using AIStudio.Chat; +using AIStudio.Provider; + +namespace AIStudio.Tests.Chat; + +/// +/// Checks when what a provider reported for a conversation still describes it. +/// +/// +/// 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. +/// +[TestFixture] +public sealed class ChatThreadReportedUsageTests +{ + 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 usage = thread.ReportedUsageFor(MODEL); + + Assert.Multiple(() => + { + Assert.That(usage.IsKnown, Is.True); + Assert.That(usage.PromptTokens, Is.EqualTo(1200)); + }); + } + + [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.ReportedUsageFor(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.ReportedUsageFor(MODEL).IsKnown, Is.False); + } + + [Test] + public void AnAnswerWithoutAReportDoesNotBorrowTheReportBeforeIt() + { + // + // What an answer written while tools were offered looks like today: 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.ReportedUsageFor(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.ReportedUsageFor(MODEL).IsKnown, Is.False); + } + + [Test] + public void AnotherModelOutdatesTheReport() + { + var thread = Thread(Question(1), Answer(2, promptTokens: 1200, blockCount: 2)); + + Assert.That(thread.ReportedUsageFor(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 usage = thread.ReportedUsageFor(MODEL); + Assert.Multiple(() => + { + Assert.That(usage.IsKnown, Is.True, "The first answer is the last block again, with exactly the blocks it was reported for."); + Assert.That(usage.PromptTokens, Is.EqualTo(1200)); + }); + } + + [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 usage = thread.ReportedUsageFor(MODEL); + Assert.Multiple(() => + { + Assert.That(usage.IsKnown, Is.True); + Assert.That(usage.PromptTokens, Is.EqualTo(1200)); + }); + } + + 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, 100), 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, + }; +} \ No newline at end of file