From 2d23301a4b5246bee09a37555ecbad720c6f1908 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:37:51 +0000 Subject: [PATCH 1/4] Ensure usage is merged for all looping components --- dotnet/eng/MSBuild/Shared.props | 3 + .../MessageMerger.cs | 39 +- .../Microsoft.Agents.AI.Workflows.csproj | 1 + .../ChatClient/MessageInjectingChatClient.cs | 10 +- .../Harness/Loop/LoopAgent.cs | 37 +- .../Harness/ToolApproval/ToolApprovalAgent.cs | 10 +- .../Microsoft.Agents.AI.csproj | 1 + .../Usage/UsageAggregationExtensions.cs | 194 +++++++ ...equiredFunctionBypassingChatClientTests.cs | 61 +++ .../MessageInjectingChatClientTests.cs | 235 ++++++++ ...allChatHistoryPersistingChatClientTests.cs | 73 +++ .../Harness/Loop/LoopAgentTests.cs | 301 +++++++++++ .../ToolApproval/ToolApprovalAgentTests.cs | 130 +++++ .../Shared/UsageAggregationExtensionsTests.cs | 501 ++++++++++++++++++ .../MessageMergerTests.cs | 60 +++ 15 files changed, 1594 insertions(+), 62 deletions(-) create mode 100644 dotnet/src/Shared/Usage/UsageAggregationExtensions.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs diff --git a/dotnet/eng/MSBuild/Shared.props b/dotnet/eng/MSBuild/Shared.props index 6cf07927292..8d32d1113af 100644 --- a/dotnet/eng/MSBuild/Shared.props +++ b/dotnet/eng/MSBuild/Shared.props @@ -26,6 +26,9 @@ + + + diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs index f656b1aa068..bf83e2f7d1a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs @@ -132,7 +132,7 @@ public AgentResponse ComputeMerged(string primaryResponseId, string? primaryAgen _ = finishReasons.Add(response.FinishReason.Value); } - usage = MergeUsage(usage, response.Usage); + usage = UsageAggregationExtensions.MergeUsage(usage, response.Usage); additionalProperties = MergeProperties(additionalProperties, response.AdditionalProperties); } @@ -219,7 +219,7 @@ static AgentResponse MergeResponses(AgentResponse? current, AgentResponse incomi Messages = current.Messages.Concat(incoming.Messages).ToList(), ResponseId = current.ResponseId, RawRepresentation = rawRepresentation, - Usage = MergeUsage(current.Usage, incoming.Usage), + Usage = UsageAggregationExtensions.MergeUsage(current.Usage, incoming.Usage), }; } @@ -269,40 +269,5 @@ static IEnumerable GetMessagesWithCreatedAt(AgentResponse response) return merged; } - - static UsageDetails? MergeUsage(UsageDetails? current, UsageDetails? incoming) - { - if (current is null) - { - return incoming; - } - - AdditionalPropertiesDictionary? additionalCounts = current.AdditionalCounts; - if (incoming is null) - { - return current; - } - - if (additionalCounts is null) - { - additionalCounts = incoming.AdditionalCounts; - } - else if (incoming.AdditionalCounts is not null) - { - foreach (string key in incoming.AdditionalCounts.Keys) - { - additionalCounts[key] = incoming.AdditionalCounts[key] + - (additionalCounts.TryGetValue(key, out long? existingCount) ? existingCount.Value : 0); - } - } - - return new UsageDetails - { - InputTokenCount = current.InputTokenCount + incoming.InputTokenCount, - OutputTokenCount = current.OutputTokenCount + incoming.OutputTokenCount, - TotalTokenCount = current.TotalTokenCount + incoming.TotalTokenCount, - AdditionalCounts = additionalCounts, - }; - } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj b/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj index 9bd247be911..374bdbacb52 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj @@ -8,6 +8,7 @@ true true + true true true true diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs index 5251b301447..b46f4930bb2 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs @@ -77,22 +77,28 @@ public override async Task GetResponseAsync( // are pending but new messages have been injected into the queue, we call the service again // so the model can process them. The loop exits when the response contains actionable // function calls (handed off to the parent FunctionInvokingChatClient) or the queue is empty. + // Usage is accumulated across every iteration so the returned response reports the token cost + // of all service calls made, not just the last one. + UsageDetails? aggregatedUsage = null; + while (true) { var response = await base.GetResponseAsync(newMessages, options, cancellationToken).ConfigureAwait(false); + UsageAggregationExtensions.AccumulateUsage(ref aggregatedUsage, response.Usage); + // If the response contains actionable function calls, the parent FunctionInvokingChatClient // loop will iterate — return immediately so it can process them. if (HasActionableFunctionCalls(response.Messages)) { - return response; + return response.WithAggregatedUsage(aggregatedUsage); } // No actionable function calls. If there are pending injected messages, loop again // to send them to the service. Otherwise, we're done. if (await this.IsQueueEmptyAsync(session, cancellationToken).ConfigureAwait(false)) { - return response; + return response.WithAggregatedUsage(aggregatedUsage); } // Propagate any ConversationId returned by the service so subsequent iterations diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs index 009a214d713..a4087f7c6d5 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs @@ -163,6 +163,10 @@ protected override async Task RunCoreAsync( // followed by that iteration's response messages. Unused when only the final response is returned. List transcript = []; + // Accumulates usage across every inner invocation so the returned response reports the token cost of the + // whole run rather than only its final iteration. Aggregated even when only the last response is returned. + UsageDetails? aggregatedUsage = null; + // The loop-synthesized on-behalf-of messages that drive the current iteration (none for the first iteration). IReadOnlyList currentSurfaced = []; @@ -174,6 +178,8 @@ protected override async Task RunCoreAsync( AgentResponse response = await this.InnerAgent.RunAsync(currentMessages, activeSession, options, cancellationToken).ConfigureAwait(false); iteration++; + UsageAggregationExtensions.AccumulateUsage(ref aggregatedUsage, response.Usage); + // Record this iteration's on-behalf-of input (before the response it elicited) and the response itself. transcript.AddRange(currentSurfaced); transcript.AddRange(response.Messages); @@ -189,21 +195,21 @@ protected override async Task RunCoreAsync( // Stop and surface the response when the agent is waiting for a tool approval. if (HasPendingApprovalRequests(response)) { - return this.BuildResult(response, transcript); + return this.BuildResult(response, transcript, aggregatedUsage); } // Enforce the global safety cap regardless of what the evaluators want. if (iteration >= this._maxIterations) { this.LogMaxIterationsReached(iteration); - return this.BuildResult(response, transcript); + return this.BuildResult(response, transcript, aggregatedUsage); } // Ask the evaluators whether to continue; stop when none of them request a re-invocation. LoopNextStep step = await this.EvaluateAndBuildNextAsync(context, feedbackLog, initialSessionSnapshot, cancellationToken).ConfigureAwait(false); if (!step.ShouldContinue) { - return this.BuildResult(response, transcript); + return this.BuildResult(response, transcript, aggregatedUsage); } currentMessages = step.Messages; @@ -447,26 +453,13 @@ private static AgentResponseUpdate CreateOnBehalfOfUpdate(ChatMessage message, s /// /// Produces the non-streaming run result: either the final iteration's response (when configured) or an - /// aggregated response carrying the full transcript with the final response's metadata. + /// aggregated response carrying the full transcript with the final response's metadata. In both cases the + /// usage reported is , covering every iteration of the run. /// - private AgentResponse BuildResult(AgentResponse lastResponse, List transcript) - { - if (this._nonStreamingReturnsLastResponseOnly) - { - return lastResponse; - } - - return new AgentResponse(transcript) - { - AgentId = lastResponse.AgentId, - ResponseId = lastResponse.ResponseId, - CreatedAt = lastResponse.CreatedAt, - FinishReason = lastResponse.FinishReason, - Usage = lastResponse.Usage, - AdditionalProperties = lastResponse.AdditionalProperties, - ContinuationToken = lastResponse.ContinuationToken, - }; - } + private AgentResponse BuildResult(AgentResponse lastResponse, List transcript, UsageDetails? aggregatedUsage) + => this._nonStreamingReturnsLastResponseOnly + ? lastResponse.WithAggregatedUsage(aggregatedUsage) + : lastResponse.WithAggregatedUsage(aggregatedUsage, transcript); private static bool HasPendingApprovalRequests(AgentResponse response) { diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs b/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs index ce840d03f50..365ed97f392 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs @@ -140,6 +140,11 @@ protected override async Task RunCoreAsync( // invocation, so a per-request cap (FunctionInvokingChatClient.MaximumIterationsPerRequest) // restarts every time and cannot bound it; without a cap here a model that keeps // requesting an auto-approved tool bills indefinitely. + // + // Usage is accumulated across every re-invocation so the caller sees the token cost + // of the whole run, not just its final inner call. + UsageDetails? aggregatedUsage = null; + for (int iteration = 0; ; iteration++) { // Inject any collected approval responses as a user message ahead of the caller's messages. @@ -156,13 +161,16 @@ protected override async Task RunCoreAsync( var response = await this.InnerAgent.RunAsync(processedMessages, session, options, cancellationToken).ConfigureAwait(false); + UsageAggregationExtensions.AccumulateUsage(ref aggregatedUsage, response.Usage); + // Classify approval requests: auto-approve matching, queue excess, keep first unapproved. bool allAutoApproved = await this.ProcessAndQueueOutboundApprovalRequestsAsync(response.Messages, state, session, options, requestMessages).ConfigureAwait(false); if (!allAutoApproved) { // Response has real content or an unapproved approval request — return to caller. - return response; + // Return a copy carrying the aggregated usage rather than mutating the inner agent's response. + return response.WithAggregatedUsage(aggregatedUsage); } // All approval requests were auto-approved. Loop to re-invoke with them injected. diff --git a/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj b/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj index c95207ef641..61a3225bbd2 100644 --- a/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj +++ b/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj @@ -9,6 +9,7 @@ true true true + true true true true diff --git a/dotnet/src/Shared/Usage/UsageAggregationExtensions.cs b/dotnet/src/Shared/Usage/UsageAggregationExtensions.cs new file mode 100644 index 00000000000..2f6128e6cde --- /dev/null +++ b/dotnet/src/Shared/Usage/UsageAggregationExtensions.cs @@ -0,0 +1,194 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Helpers for aggregating across multiple service or agent invocations that +/// make up a single logical run. +/// +/// +/// Several components re-invoke an inner agent or chat client in a loop within a single run (for example +/// when auto-approving tool calls, injecting messages, or re-running an agent until an evaluator is +/// satisfied). Each inner invocation reports its own usage, and the aggregate must be surfaced to the +/// caller so that the reported token counts reflect the entire run rather than just its final step. +/// +internal static class UsageAggregationExtensions +{ + /// + /// Combines two instances into a new instance containing their summed values. + /// + /// The running aggregate, or if nothing has been accumulated yet. + /// The usage reported by the latest invocation, or if none was reported. + /// + /// A new containing the summed token counts and additional counts, or + /// when both and are + /// . + /// + /// + /// Neither argument is mutated, and neither argument is ever returned by reference, since both may be + /// owned and observed by callers. Every strongly-typed counter exposed by is + /// summed, matching the set covered by , so that no provider-reported + /// counter is lost when a merged instance replaces the original. Token counts are summed in a null-aware + /// manner: combining a count with a non-null count yields the non-null count, and + /// combining two counts yields . Entries in + /// are summed per key so that provider-specific counters + /// (such as cached, reasoning, or cost counters) aggregate correctly. + /// + public static UsageDetails? MergeUsage(UsageDetails? current, UsageDetails? incoming) + { + if (current is null && incoming is null) + { + return null; + } + + var merged = new UsageDetails + { + InputTokenCount = AddCounts(current?.InputTokenCount, incoming?.InputTokenCount), + OutputTokenCount = AddCounts(current?.OutputTokenCount, incoming?.OutputTokenCount), + TotalTokenCount = AddCounts(current?.TotalTokenCount, incoming?.TotalTokenCount), + CachedInputTokenCount = AddCounts(current?.CachedInputTokenCount, incoming?.CachedInputTokenCount), + ReasoningTokenCount = AddCounts(current?.ReasoningTokenCount, incoming?.ReasoningTokenCount), + InputAudioTokenCount = AddCounts(current?.InputAudioTokenCount, incoming?.InputAudioTokenCount), + InputTextTokenCount = AddCounts(current?.InputTextTokenCount, incoming?.InputTextTokenCount), + OutputAudioTokenCount = AddCounts(current?.OutputAudioTokenCount, incoming?.OutputAudioTokenCount), + OutputTextTokenCount = AddCounts(current?.OutputTextTokenCount, incoming?.OutputTextTokenCount), + }; + + AdditionalPropertiesDictionary? additionalCounts = MergeAdditionalCounts(current?.AdditionalCounts, incoming?.AdditionalCounts); + if (additionalCounts is not null) + { + merged.AdditionalCounts = additionalCounts; + } + + return merged; + } + + /// + /// Adds the usage into the running aggregate referenced by + /// , replacing it with a new combined instance. + /// + /// The running aggregate to update. May be . + /// The usage reported by the latest invocation, or if none was reported. + public static void AccumulateUsage(ref UsageDetails? current, UsageDetails? incoming) + => current = MergeUsage(current, incoming); + + /// + /// Returns a reporting in place of the usage + /// carried by , which typically covers only the final service call of a run. + /// + /// The response produced by the final call of the run. + /// The usage accumulated across every call that made up the run. + /// + /// The messages the returned response should carry, or to keep those of + /// . Used when a run returns a transcript spanning multiple calls. + /// + /// + /// A copy is returned rather than the usage being assigned onto , because the + /// inner client may still own and observe that instance. The original is returned unchanged only when it + /// already carries exactly the usage to report and no message substitution is requested. + /// + public static ChatResponse WithAggregatedUsage(this ChatResponse response, UsageDetails? aggregatedUsage, IList? messages = null) + { + if (messages is null && ReferenceEquals(response.Usage, aggregatedUsage)) + { + return response; + } + + return new ChatResponse(messages ?? response.Messages) + { + ResponseId = response.ResponseId, + ConversationId = response.ConversationId, + ModelId = response.ModelId, + CreatedAt = response.CreatedAt, + FinishReason = response.FinishReason, + Usage = aggregatedUsage, + ContinuationToken = response.ContinuationToken, + RawRepresentation = response.RawRepresentation, + AdditionalProperties = response.AdditionalProperties, + }; + } + + /// + /// Returns an reporting in place of the usage + /// carried by , which typically covers only the final invocation of a run. + /// + /// The response produced by the final invocation of the run. + /// The usage accumulated across every invocation that made up the run. + /// + /// The messages the returned response should carry, or to keep those of + /// . Used when a run returns a transcript spanning multiple invocations. + /// + /// + /// A copy is returned rather than the usage being assigned onto , because the + /// inner agent may still own and observe that instance. The original is returned unchanged only when it + /// already carries exactly the usage to report and no message substitution is requested. + /// + public static AgentResponse WithAggregatedUsage(this AgentResponse response, UsageDetails? aggregatedUsage, IList? messages = null) + { + if (messages is null && ReferenceEquals(response.Usage, aggregatedUsage)) + { + return response; + } + + return new AgentResponse(messages ?? response.Messages) + { + AgentId = response.AgentId, + ResponseId = response.ResponseId, + CreatedAt = response.CreatedAt, + FinishReason = response.FinishReason, + Usage = aggregatedUsage, + ContinuationToken = response.ContinuationToken, + RawRepresentation = response.RawRepresentation, + AdditionalProperties = response.AdditionalProperties, + }; + } + + /// + /// Adds two nullable counts, treating as "not reported" rather than as zero so + /// that an aggregate only reports a count when at least one contributor reported one. + /// + private static long? AddCounts(long? current, long? incoming) + => current is null ? incoming : incoming is null ? current : current + incoming; + + /// + /// Produces a new dictionary containing the per-key sums of the supplied additional counts, or + /// when neither side has any entries. + /// + private static AdditionalPropertiesDictionary? MergeAdditionalCounts( + AdditionalPropertiesDictionary? current, + AdditionalPropertiesDictionary? incoming) + { + bool hasCurrent = current is { Count: > 0 }; + bool hasIncoming = incoming is { Count: > 0 }; + + if (!hasCurrent && !hasIncoming) + { + return null; + } + + var merged = new AdditionalPropertiesDictionary(); + + if (hasCurrent) + { + foreach (var entry in current!) + { + merged[entry.Key] = entry.Value; + } + } + + if (hasIncoming) + { + foreach (var entry in incoming!) + { + merged[entry.Key] = merged.TryGetValue(entry.Key, out long existing) + ? existing + entry.Value + : entry.Value; + } + } + + return merged; + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalNotRequiredFunctionBypassingChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalNotRequiredFunctionBypassingChatClientTests.cs index fb2b30629eb..1a49d48f2b2 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalNotRequiredFunctionBypassingChatClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalNotRequiredFunctionBypassingChatClientTests.cs @@ -579,6 +579,67 @@ public void WithDefaultAgentMiddleware_DisableApprovalNotRequiredFunctionBypassi #endregion + #region Usage Pass-Through Tests + + /// + /// Verifies that usage reported by the inner client is surfaced unchanged, since this decorator + /// makes exactly one inner call and must not drop or alter usage. + /// + [Fact] + public async Task GetResponseAsync_PassesUsageThroughUnchangedAsync() + { + // Arrange + var usage = new UsageDetails { InputTokenCount = 11, OutputTokenCount = 7, TotalTokenCount = 18 }; + var innerClient = CreateMockChatClient((_, _, _) => + Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Hello")]) { Usage = usage })); + + var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient); + var session = new ChatClientAgentSession(); + + // Act + var response = await RunWithAgentContextAsync(decorator, session); + + // Assert + Assert.NotNull(response.Usage); + Assert.Equal(11, response.Usage!.InputTokenCount); + Assert.Equal(7, response.Usage.OutputTokenCount); + Assert.Equal(18, response.Usage.TotalTokenCount); + } + + /// + /// Verifies that a streaming update carrying both auto-approved approval content and usage content + /// still surfaces its usage after the approval content is stripped. + /// + [Fact] + public async Task GetStreamingResponseAsync_UpdateWithAutoApprovedAndUsage_StillSurfacesUsageAsync() + { + // Arrange + var noApprovalTool = AIFunctionFactory.Create(() => "result", "plainTool"); + var approval = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "plainTool")); + var usage = new UsageDetails { InputTokenCount = 9, OutputTokenCount = 4, TotalTokenCount = 13 }; + + var innerClient = CreateMockStreamingChatClient((_, _, _) => + ToAsyncEnumerableAsync(new ChatResponseUpdate(ChatRole.Assistant, [approval, new UsageContent(usage)]))); + + var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient); + var session = new ChatClientAgentSession(); + var options = new ChatOptions { Tools = [noApprovalTool] }; + + // Act + List updates = []; + await RunStreamingWithAgentContextAsync(decorator, session, updates, options); + + // Assert — the approval request was bypassed but the usage survived. + Assert.DoesNotContain(updates.SelectMany(static u => u.Contents), static c => c is ToolApprovalRequestContent); + var response = updates.ToChatResponse(); + Assert.NotNull(response.Usage); + Assert.Equal(9, response.Usage!.InputTokenCount); + Assert.Equal(4, response.Usage.OutputTokenCount); + Assert.Equal(13, response.Usage.TotalTokenCount); + } + + #endregion + #region Helpers private static async Task RunWithAgentContextAsync( diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs index 4bd69b089d2..9e8159d6f20 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs @@ -2,6 +2,7 @@ using System.Collections.Generic; using System.Linq; +using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -288,6 +289,69 @@ public async Task RunAsync_DoesNotLoopInternally_WhenActionableFCCPresentAsync() Assert.Equal(2, serviceCallCount); } + /// + /// Verifies that usage is aggregated when an actionable exits the + /// injected-message loop on a later iteration. + /// + [Fact] + public async Task RunAsync_ActionableFCCOnLaterIteration_AggregatesUsageAcrossInjectedLoopExitAsync() + { + // Arrange + int serviceCallCount = 0; + Mock mockService = new(); + MessageInjectingChatClient? injectorRef = null; + ChatClientAgentSession? sessionRef = null; + + mockService.Setup( + s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(async (IEnumerable msgs, ChatOptions? _, CancellationToken ct) => + { + serviceCallCount++; + if (serviceCallCount == 1) + { + await injectorRef!.EnqueueMessagesAsync(sessionRef!, [new ChatMessage(ChatRole.User, "injected")], ct); + return new ChatResponse([new(ChatRole.Assistant, "queued")]) { Usage = CreateUsageForCall(serviceCallCount) }; + } + + if (serviceCallCount == 2) + { + return new ChatResponse([new(ChatRole.Assistant, + [new FunctionCallContent("call1", "myTool", new Dictionary())])]) + { + Usage = CreateUsageForCall(serviceCallCount) + }; + } + + return new ChatResponse([new(ChatRole.Assistant, "final")]) { Usage = CreateUsageForCall(serviceCallCount) }; + }); + + var tool = AIFunctionFactory.Create(() => "tool result", "myTool", "A test tool"); + ChatClientAgent agent = new(mockService.Object, options: new() + { + ChatOptions = new() { Tools = [tool] }, + RequirePerServiceCallChatHistoryPersistence = true, + EnableMessageInjection = true, + }, services: new ServiceCollection().BuildServiceProvider()); + + injectorRef = agent.ChatClient.GetService()!; + + // Act + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + sessionRef = session; + var response = await agent.RunAsync([new(ChatRole.User, "original")], session); + + // Assert + Assert.Equal(3, serviceCallCount); + Assert.Equal("final", response.Text); + Assert.NotNull(response.Usage); + Assert.Equal(42, response.Usage!.InputTokenCount); + Assert.Equal(15, response.Usage.OutputTokenCount); + Assert.Equal(57, response.Usage.TotalTokenCount); + } + /// /// Verifies that the internal loop fires when the response contains only InformationalOnly /// FunctionCallContent (which are not actionable) and there are pending injected messages. @@ -537,4 +601,175 @@ public async Task EnqueueMessages_ConcurrentEnqueues_DoesNotLoseMessagesAsync() IReadOnlyList pending = await injector!.GetPendingMessagesAsync(session); Assert.Equal(ThreadCount * MessagesPerThread, pending.Count); } + + /// + /// Verifies that usage from every internal loop iteration is summed into the returned response, + /// rather than only the final service call's usage being reported. + /// + [Fact] + public async Task RunAsync_LoopsInternally_AggregatesUsageAcrossIterationsAsync() + { + // Arrange + int serviceCallCount = 0; + Mock mockService = new(); + MessageInjectingChatClient? injectorRef = null; + ChatClientAgentSession? sessionRef = null; + + mockService.Setup( + s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(async (IEnumerable msgs, ChatOptions? _, CancellationToken ct) => + { + serviceCallCount++; + if (serviceCallCount < 3) + { + // Enqueue a message so the injection loop runs again. + await injectorRef!.EnqueueMessagesAsync(sessionRef!, [new ChatMessage(ChatRole.User, $"injected {serviceCallCount}")], ct); + } + + return new ChatResponse([new(ChatRole.Assistant, $"response {serviceCallCount}")]) + { + Usage = CreateUsageForCall(serviceCallCount), + }; + }); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + EnableMessageInjection = true, + }); + + injectorRef = agent.ChatClient.GetService()!; + + // Act + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + sessionRef = session; + var response = await agent.RunAsync([new(ChatRole.User, "original")], session); + + // Assert + Assert.Equal(3, serviceCallCount); + Assert.NotNull(response.Usage); + Assert.Equal(42, response.Usage!.InputTokenCount); + Assert.Equal(15, response.Usage.OutputTokenCount); + Assert.Equal(57, response.Usage.TotalTokenCount); + } + + /// + /// Verifies that a single service call surfaces its usage unchanged and that the underlying + /// client's usage instance is not mutated. + /// + [Fact] + public async Task RunAsync_SingleServiceCall_SurfacesUsageWithoutMutatingServiceUsageAsync() + { + // Arrange + UsageDetails serviceUsage = new() { InputTokenCount = 12, OutputTokenCount = 3, TotalTokenCount = 15 }; + Mock mockService = new(); + mockService.Setup( + s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "done")]) { Usage = serviceUsage }); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + EnableMessageInjection = true, + }); + + // Act + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + var response = await agent.RunAsync([new(ChatRole.User, "original")], session); + + // Assert + Assert.NotNull(response.Usage); + Assert.NotSame(serviceUsage, response.Usage); + Assert.Equal(12, response.Usage!.InputTokenCount); + Assert.Equal(3, response.Usage.OutputTokenCount); + Assert.Equal(15, response.Usage.TotalTokenCount); + Assert.Equal(12, serviceUsage.InputTokenCount); + Assert.Equal(3, serviceUsage.OutputTokenCount); + Assert.Equal(15, serviceUsage.TotalTokenCount); + } + + /// + /// Verifies that the streaming path surfaces usage from every internal loop iteration so an + /// aggregated response reports the usage of all service calls. + /// + [Fact] + public async Task RunStreamingAsync_LoopsInternally_SurfacesUsageFromEveryIterationAsync() + { + // Arrange + int serviceCallCount = 0; + Mock mockService = new(); + MessageInjectingChatClient? injectorRef = null; + ChatClientAgentSession? sessionRef = null; + + mockService.Setup( + s => s.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns((IEnumerable msgs, ChatOptions? _, CancellationToken ct) => + { + serviceCallCount++; + int call = serviceCallCount; + return StreamWithUsageAsync(call, injectorRef!, sessionRef!, ct); + }); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + EnableMessageInjection = true, + }); + + injectorRef = agent.ChatClient.GetService()!; + + // Act + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + sessionRef = session; + List updates = []; + await foreach (var update in agent.RunStreamingAsync([new(ChatRole.User, "original")], session)) + { + updates.Add(update); + } + + // Assert + Assert.Equal(3, serviceCallCount); + var response = updates.ToAgentResponse(); + Assert.NotNull(response.Usage); + Assert.Equal(42, response.Usage!.InputTokenCount); + Assert.Equal(15, response.Usage.OutputTokenCount); + Assert.Equal(57, response.Usage.TotalTokenCount); + } + + /// + /// Creates distinct usage values for a service call so aggregation tests cannot pass by multiplying + /// the final usage by the call count. + /// + private static UsageDetails CreateUsageForCall(int call) + => new() + { + InputTokenCount = call is 1 ? 2 : call is 2 ? 11 : 29, + OutputTokenCount = call is 1 ? 3 : call is 2 ? 5 : 7, + TotalTokenCount = call is 1 ? 5 : call is 2 ? 16 : 36, + }; + + /// + /// Streams a text update followed by a usage update, enqueuing an injected message on the first two + /// calls so the injection loop runs three times in total. + /// + private static async IAsyncEnumerable StreamWithUsageAsync( + int call, + MessageInjectingChatClient injector, + ChatClientAgentSession session, + [EnumeratorCancellation] CancellationToken cancellationToken) + { + if (call < 3) + { + await injector.EnqueueMessagesAsync(session, [new ChatMessage(ChatRole.User, $"injected {call}")], cancellationToken); + } + + yield return new ChatResponseUpdate(ChatRole.Assistant, $"response {call}"); + yield return new ChatResponseUpdate(ChatRole.Assistant, [new UsageContent(CreateUsageForCall(call))]); + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/PerServiceCallChatHistoryPersistingChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/PerServiceCallChatHistoryPersistingChatClientTests.cs index 8fcf99e17c9..a24b8453e1a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/PerServiceCallChatHistoryPersistingChatClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/PerServiceCallChatHistoryPersistingChatClientTests.cs @@ -718,6 +718,79 @@ private static async IAsyncEnumerable CreateAsyncEnumerableA await Task.CompletedTask; } + /// + /// Verifies that usage reported by the inner client is surfaced unchanged. This decorator makes + /// exactly one inner call per invocation and must never drop usage. + /// + [Fact] + public async Task RunAsync_PassesUsageThroughUnchangedAsync() + { + // Arrange + Mock mockService = new(); + mockService.Setup( + s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]) + { + Usage = new UsageDetails { InputTokenCount = 21, OutputTokenCount = 9, TotalTokenCount = 30 }, + }); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + RequirePerServiceCallChatHistoryPersistence = true, + }); + + // Act + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + var response = await agent.RunAsync([new(ChatRole.User, "test")], session); + + // Assert + Assert.NotNull(response.Usage); + Assert.Equal(21, response.Usage!.InputTokenCount); + Assert.Equal(9, response.Usage.OutputTokenCount); + Assert.Equal(30, response.Usage.TotalTokenCount); + } + + /// + /// Verifies that usage streamed by the inner client survives the decorator's streaming pipeline. + /// + [Fact] + public async Task RunStreamingAsync_PassesUsageThroughUnchangedAsync() + { + // Arrange + Mock mockService = new(); + mockService.Setup( + s => s.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(CreateAsyncEnumerableAsync( + new ChatResponseUpdate(ChatRole.Assistant, "response"), + new ChatResponseUpdate(ChatRole.Assistant, [new UsageContent(new UsageDetails { InputTokenCount = 21, OutputTokenCount = 9, TotalTokenCount = 30 })]))); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + RequirePerServiceCallChatHistoryPersistence = true, + }); + + // Act + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + List updates = []; + await foreach (var update in agent.RunStreamingAsync([new(ChatRole.User, "test")], session)) + { + updates.Add(update); + } + + // Assert + var response = updates.ToAgentResponse(); + Assert.NotNull(response.Usage); + Assert.Equal(21, response.Usage!.InputTokenCount); + Assert.Equal(9, response.Usage.OutputTokenCount); + Assert.Equal(30, response.Usage.TotalTokenCount); + } + /// /// Verifies that when per-service-call persistence is active and no real conversation ID exists, /// sets the diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentTests.cs index 428298f1d62..866445cd1eb 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentTests.cs @@ -164,6 +164,30 @@ public async Task RunAsync_PredicateLoopsUntilFalse_AggregatesAllIterationsAsync Assert.Equal(["iteration 1", "iteration 2", "iteration 3"], response.Messages.Select(static m => m.Text)); } + /// + /// Verify that the aggregated transcript is returned even when no iteration reports usage. The shared + /// WithAggregatedUsage helper has a fast path that returns the original response when its usage is + /// already reference-equal to the aggregate (which is the case when both are ), so + /// this pins that the fast path can never suppress transcript aggregation. + /// + [Fact] + public async Task RunAsync_AggregatedTranscript_ReturnsFullTranscriptWhenNoUsageReportedAsync() + { + // Arrange + var capture = new InnerAgentCapture(call => + new AgentResponse([new ChatMessage(ChatRole.Assistant, $"iteration {call}")]) { Usage = null }); + var evaluator = While(ctx => ctx.LastResponse.Text != "iteration 3"); + var agent = new LoopAgent(capture.Agent, evaluator); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, capture.CallCount); + Assert.Equal(["iteration 1", "iteration 2", "iteration 3"], response.Messages.Select(static m => m.Text)); + Assert.Null(response.Usage); + } + /// /// Verify that returns only the final /// iteration's response instead of the aggregated transcript. @@ -1006,6 +1030,283 @@ public async Task RunAsync_ExcludeOnBehalfOfMessages_OmitsThemFromResponseAsync( #endregion + #region Usage aggregation + + /// + /// Verify that usage from every iteration is summed into the aggregated response rather than only the + /// final iteration's usage being reported. + /// + [Fact] + public async Task RunAsync_MultipleIterations_AggregatesUsageAcrossIterationsAsync() + { + // Arrange + var capture = new InnerAgentCapture(call => + new AgentResponse([new ChatMessage(ChatRole.Assistant, $"iteration {call}")]) + { + Usage = new UsageDetails { InputTokenCount = call * 10, OutputTokenCount = call, TotalTokenCount = (call * 10) + call }, + }); + var evaluator = While(ctx => ctx.Iteration < 3); + var agent = new LoopAgent(capture.Agent, evaluator); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, capture.CallCount); + Assert.NotNull(response.Usage); + Assert.Equal(60, response.Usage!.InputTokenCount); + Assert.Equal(6, response.Usage.OutputTokenCount); + Assert.Equal(66, response.Usage.TotalTokenCount); + } + + /// + /// Verify that usage covers the whole run even when only the final iteration's response is returned. + /// + [Fact] + public async Task RunAsync_LastResponseOnly_StillAggregatesUsageAcrossIterationsAsync() + { + // Arrange + var capture = new InnerAgentCapture(call => + new AgentResponse([new ChatMessage(ChatRole.Assistant, $"iteration {call}")]) + { + Usage = CreateUsageForCall(call), + }); + var evaluator = While(ctx => ctx.Iteration < 3); + var options = new LoopAgentOptions { NonStreamingReturnsLastResponseOnly = true }; + var agent = new LoopAgent(capture.Agent, evaluator, options); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, capture.CallCount); + Assert.Single(response.Messages); + Assert.Equal("iteration 3", response.Text); + Assert.NotNull(response.Usage); + Assert.Equal(42, response.Usage!.InputTokenCount); + Assert.Equal(15, response.Usage.OutputTokenCount); + Assert.Equal(57, response.Usage.TotalTokenCount); + } + + /// + /// Verify that usage is aggregated across all iterations before the global safety cap stops the loop. + /// + [Fact] + public async Task RunAsync_AlwaysContinue_StopsAtGlobalCapAndAggregatesUsageAsync() + { + // Arrange + var capture = new InnerAgentCapture(call => + new AgentResponse([new ChatMessage(ChatRole.Assistant, $"iteration {call}")]) + { + Usage = CreateUsageForCall(call), + }); + var evaluator = While(static _ => true); + var options = new LoopAgentOptions { MaxIterations = 3 }; + var agent = new LoopAgent(capture.Agent, evaluator, options); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, capture.CallCount); + Assert.NotNull(response.Usage); + Assert.Equal(42, response.Usage!.InputTokenCount); + Assert.Equal(15, response.Usage.OutputTokenCount); + Assert.Equal(57, response.Usage.TotalTokenCount); + } + + /// + /// Verify that usage is aggregated when a later iteration stops the loop with a pending approval request. + /// + [Fact] + public async Task RunAsync_PendingApprovalRequestOnLaterIteration_AggregatesUsageAsync() + { + // Arrange + var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); + var capture = new InnerAgentCapture(call => + new AgentResponse( + [call < 3 + ? new ChatMessage(ChatRole.Assistant, $"iteration {call}") + : new ChatMessage(ChatRole.Assistant, [approvalRequest])]) + { + Usage = CreateUsageForCall(call), + }); + var evaluator = While(static _ => true); + var agent = new LoopAgent(capture.Agent, evaluator); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, capture.CallCount); + Assert.Contains(response.Messages.SelectMany(static m => m.Contents), static c => c is ToolApprovalRequestContent); + Assert.NotNull(response.Usage); + Assert.Equal(42, response.Usage!.InputTokenCount); + Assert.Equal(15, response.Usage.OutputTokenCount); + Assert.Equal(57, response.Usage.TotalTokenCount); + } + + /// + /// Verify that iterations reporting no usage are skipped rather than zeroing out the aggregate. + /// + [Fact] + public async Task RunAsync_SomeIterationsWithoutUsage_AggregatesReportedUsageOnlyAsync() + { + // Arrange + var capture = new InnerAgentCapture(call => + new AgentResponse([new ChatMessage(ChatRole.Assistant, $"iteration {call}")]) + { + Usage = call == 2 ? null : new UsageDetails { InputTokenCount = 7, TotalTokenCount = 7 }, + }); + var evaluator = While(ctx => ctx.Iteration < 3); + var agent = new LoopAgent(capture.Agent, evaluator); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, capture.CallCount); + Assert.NotNull(response.Usage); + Assert.Equal(14, response.Usage!.InputTokenCount); + Assert.Equal(14, response.Usage.TotalTokenCount); + Assert.Null(response.Usage.OutputTokenCount); + } + + /// + /// Verify that aggregating usage never mutates the usage instances owned by the inner agent's responses. + /// + [Fact] + public async Task RunAsync_AggregatingUsage_DoesNotMutateInnerResponseUsageAsync() + { + // Arrange + var innerUsages = new List(); + var capture = new InnerAgentCapture(call => + { + var usage = new UsageDetails { InputTokenCount = 10, OutputTokenCount = 5, TotalTokenCount = 15 }; + innerUsages.Add(usage); + return new AgentResponse([new ChatMessage(ChatRole.Assistant, $"iteration {call}")]) { Usage = usage }; + }); + var evaluator = While(ctx => ctx.Iteration < 3); + var agent = new LoopAgent(capture.Agent, evaluator); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.All(innerUsages, usage => + { + Assert.Equal(10, usage.InputTokenCount); + Assert.Equal(5, usage.OutputTokenCount); + Assert.Equal(15, usage.TotalTokenCount); + }); + Assert.DoesNotContain(innerUsages, usage => ReferenceEquals(usage, response.Usage)); + } + + /// + /// Verify that usage lives on and is not also duplicated into the + /// aggregated transcript messages, which would cause downstream consumers to double count it. + /// + [Fact] + public async Task RunAsync_AggregatedTranscript_DoesNotDuplicateUsageIntoMessagesAsync() + { + // Arrange + var capture = new InnerAgentCapture(call => + new AgentResponse([new ChatMessage(ChatRole.Assistant, $"iteration {call}")]) + { + Usage = new UsageDetails { InputTokenCount = 10, TotalTokenCount = 10 }, + }); + var evaluator = While(ctx => ctx.Iteration < 2); + var agent = new LoopAgent(capture.Agent, evaluator); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.Equal(20, response.Usage!.InputTokenCount); + Assert.DoesNotContain(response.Messages.SelectMany(static m => m.Contents), static c => c is UsageContent); + } + + /// + /// Verify that the streaming path surfaces every iteration's usage so an aggregated response built from + /// the updates reports the usage of the whole run. + /// + [Fact] + public async Task RunStreamingAsync_MultipleIterations_SurfacesUsageFromEveryIterationAsync() + { + // Arrange + var capture = new InnerStreamingCapture(call => + [ + new AgentResponseUpdate(ChatRole.Assistant, $"chunk {call}"), + new AgentResponseUpdate(ChatRole.Assistant, [new UsageContent(CreateUsageForCall(call))]), + ]); + var evaluator = While(ctx => ctx.Iteration < 3); + var agent = new LoopAgent(capture.Agent, evaluator); + + // Act + var updates = new List(); + await foreach (var update in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession())) + { + updates.Add(update); + } + + // Assert + Assert.Equal(3, capture.CallCount); + var response = updates.ToAgentResponse(); + Assert.NotNull(response.Usage); + Assert.Equal(42, response.Usage!.InputTokenCount); + Assert.Equal(15, response.Usage.OutputTokenCount); + Assert.Equal(57, response.Usage.TotalTokenCount); + } + + /// + /// Creates distinct usage values for an inner call so aggregation tests cannot pass by multiplying + /// the final usage by the call count. + /// + private static UsageDetails CreateUsageForCall(int call) + => new() + { + InputTokenCount = call is 1 ? 2 : call is 2 ? 11 : 29, + OutputTokenCount = call is 1 ? 3 : call is 2 ? 5 : 7, + TotalTokenCount = call is 1 ? 5 : call is 2 ? 16 : 36, + }; + + /// + /// Verify that provider-reported counters beyond the three headline token counts (such as cached and + /// reasoning tokens, both of which real providers populate) survive aggregation end to end, rather than + /// being dropped when the aggregated response replaces the inner response's usage instance. + /// + [Fact] + public async Task RunAsync_MultipleIterations_AggregatesAllProviderReportedCountersAsync() + { + // Arrange + var capture = new InnerAgentCapture(call => + new AgentResponse([new ChatMessage(ChatRole.Assistant, $"iteration {call}")]) + { + Usage = new UsageDetails + { + InputTokenCount = 100, + CachedInputTokenCount = 3, + ReasoningTokenCount = 11, + AdditionalCounts = new() { ["cost"] = 2 }, + }, + }); + var evaluator = While(ctx => ctx.Iteration < 3); + var agent = new LoopAgent(capture.Agent, evaluator); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, capture.CallCount); + Assert.NotNull(response.Usage); + Assert.Equal(300, response.Usage!.InputTokenCount); + Assert.Equal(9, response.Usage.CachedInputTokenCount); + Assert.Equal(33, response.Usage.ReasoningTokenCount); + Assert.Equal(6, response.Usage.AdditionalCounts!["cost"]); + } + + #endregion + #region RunStreamingAsync /// diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs index 5b14963cdc5..1bc9d515750 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs @@ -1782,6 +1782,136 @@ [new ChatMessage(ChatRole.User, [alwaysApprove])], #endregion + #region Usage aggregation + + /// + /// Verify that usage from every auto-approval re-invocation of the inner agent is summed into the + /// response returned to the caller. + /// + [Fact] + public async Task RunAsync_AutoApprovalLoop_AggregatesUsageAcrossInvocationsAsync() + { + // Arrange + var session = new ChatClientAgentSession(); + var callCount = 0; + var innerAgent = new Mock(); + innerAgent + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(() => + { + callCount++; + var usage = new UsageDetails { InputTokenCount = callCount is 1 ? 2 : callCount is 2 ? 11 : 29, OutputTokenCount = callCount is 1 ? 3 : callCount is 2 ? 5 : 7, TotalTokenCount = callCount is 1 ? 5 : callCount is 2 ? 16 : 36 }; + if (callCount < 3) + { + var request = new ToolApprovalRequestContent($"req{callCount}", new FunctionCallContent($"call{callCount}", "DangerousTool")); + return new AgentResponse([new ChatMessage(ChatRole.Assistant, [request])]) { Usage = usage }; + } + + return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]) { Usage = usage }; + }); + + var options = new ToolApprovalAgentOptions + { + AutoApprovalRules = [ToolApprovalAgent.AllToolsAutoApprovalRule] + }; + var agent = new ToolApprovalAgent(innerAgent.Object, options); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session); + + // Assert + Assert.Equal(3, callCount); + Assert.Equal("Done", response.Text); + Assert.NotNull(response.Usage); + Assert.Equal(42, response.Usage!.InputTokenCount); + Assert.Equal(15, response.Usage.OutputTokenCount); + Assert.Equal(57, response.Usage.TotalTokenCount); + } + + /// + /// Verify that a single inner invocation still surfaces its usage unchanged and does not mutate the + /// inner agent's usage instance. + /// + [Fact] + public async Task RunAsync_SingleInvocation_SurfacesUsageWithoutMutatingInnerResponseAsync() + { + // Arrange + var session = new ChatClientAgentSession(); + var innerUsage = new UsageDetails { InputTokenCount = 12, OutputTokenCount = 3, TotalTokenCount = 15 }; + var innerResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]) { Usage = innerUsage }; + var agent = new ToolApprovalAgent(CreateMockAgent(innerResponse).Object); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session); + + // Assert + Assert.NotNull(response.Usage); + Assert.NotSame(innerUsage, response.Usage); + Assert.Equal(12, response.Usage!.InputTokenCount); + Assert.Equal(3, response.Usage.OutputTokenCount); + Assert.Equal(15, response.Usage.TotalTokenCount); + Assert.Equal(12, innerUsage.InputTokenCount); + Assert.Equal(3, innerUsage.OutputTokenCount); + Assert.Equal(15, innerUsage.TotalTokenCount); + } + + /// + /// Verify that the streaming path surfaces every auto-approval iteration's usage so an aggregated + /// response built from the updates reports the usage of the whole run. + /// + [Fact] + public async Task RunStreamingAsync_AutoApprovalLoop_SurfacesUsageFromEveryInvocationAsync() + { + // Arrange + var session = new ChatClientAgentSession(); + var callCount = 0; + var innerAgent = new Mock(); + innerAgent + .Protected() + .Setup>("RunCoreStreamingAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns, AgentSession?, AgentRunOptions?, CancellationToken>((_, _, _, ct) => + { + callCount++; + var usageUpdate = new AgentResponseUpdate(ChatRole.Assistant, [new UsageContent(new UsageDetails { InputTokenCount = callCount is 1 ? 2 : callCount is 2 ? 11 : 29, OutputTokenCount = callCount is 1 ? 3 : callCount is 2 ? 5 : 7, TotalTokenCount = callCount is 1 ? 5 : callCount is 2 ? 16 : 36 })]); + AgentResponseUpdate[] updates = callCount < 3 + ? [new AgentResponseUpdate(ChatRole.Assistant, [new ToolApprovalRequestContent($"req{callCount}", new FunctionCallContent($"call{callCount}", "DangerousTool"))]), usageUpdate] + : [new AgentResponseUpdate(ChatRole.Assistant, "Done"), usageUpdate]; + return ToAsyncEnumerableAsync(updates, ct); + }); + + var options = new ToolApprovalAgentOptions + { + AutoApprovalRules = [ToolApprovalAgent.AllToolsAutoApprovalRule] + }; + var agent = new ToolApprovalAgent(innerAgent.Object, options); + + // Act + var updates = new List(); + await foreach (var update in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hi")], session)) + { + updates.Add(update); + } + + // Assert + Assert.Equal(3, callCount); + var response = updates.ToAgentResponse(); + Assert.NotNull(response.Usage); + Assert.Equal(42, response.Usage!.InputTokenCount); + Assert.Equal(15, response.Usage.OutputTokenCount); + Assert.Equal(57, response.Usage.TotalTokenCount); + } + + #endregion + #region Helpers private static Mock CreateMockAgent(AgentResponse response) diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs new file mode 100644 index 00000000000..5c0c126ce0c --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs @@ -0,0 +1,501 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Reflection; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.UnitTests; + +/// +/// Unit tests for . +/// +public class UsageAggregationExtensionsTests +{ + /// + /// Verify that merging two null usage values returns null. + /// + [Fact] + public void MergeUsage_BothInputsNull_ReturnsNull() + { + // Arrange, Act + UsageDetails? result = UsageAggregationExtensions.MergeUsage(null, null); + + // Assert + Assert.Null(result); + } + + /// + /// Verify that merging one null usage value returns a new copy of the non-null usage value. + /// + [Theory] + [InlineData(true)] + [InlineData(false)] + public void MergeUsage_OneInputNull_ReturnsNewCopy(bool currentIsNull) + { + // Arrange + UsageDetails usage = CreateUsage(2, 3, 5, new() { ["cached"] = 7 }); + UsageSnapshot before = UsageSnapshot.Capture(usage); + + // Act + UsageDetails? result = currentIsNull + ? UsageAggregationExtensions.MergeUsage(null, usage) + : UsageAggregationExtensions.MergeUsage(usage, null); + + // Assert + Assert.NotNull(result); + Assert.NotSame(usage, result); + Assert.NotSame(usage.AdditionalCounts, result!.AdditionalCounts); + Assert.Equal(2, result.InputTokenCount); + Assert.Equal(3, result.OutputTokenCount); + Assert.Equal(5, result.TotalTokenCount); + Assert.Equal(7, result.AdditionalCounts!["cached"]); + Assert.Equal(before, UsageSnapshot.Capture(usage)); + } + + /// + /// Verify that token counts are summed while preserving null as not reported. + /// + [Fact] + public void MergeUsage_TokenCounts_SumsNullAware() + { + // Arrange + UsageDetails current = CreateUsage(input: 2, output: null, total: null); + UsageDetails incoming = CreateUsage(input: 11, output: 5, total: null); + + // Act + UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.NotSame(current, result); + Assert.NotSame(incoming, result); + Assert.Equal(13, result!.InputTokenCount); + Assert.Equal(5, result.OutputTokenCount); + Assert.Null(result.TotalTokenCount); + } + + /// + /// Verify that additional counts are summed per key and preserve disjoint keys. + /// + [Fact] + public void MergeUsage_AdditionalCounts_SumsOverlappingAndUnionsDisjointKeys() + { + // Arrange + UsageDetails current = CreateUsage(input: null, output: null, total: null, new() { ["cached"] = 2, ["reasoning"] = 3 }); + UsageDetails incoming = CreateUsage(input: null, output: null, total: null, new() { ["cached"] = 11, ["audio"] = 29 }); + + // Act + UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.NotSame(current.AdditionalCounts, result!.AdditionalCounts); + Assert.NotSame(incoming.AdditionalCounts, result.AdditionalCounts); + Assert.Equal(13, result.AdditionalCounts!["cached"]); + Assert.Equal(3, result.AdditionalCounts["reasoning"]); + Assert.Equal(29, result.AdditionalCounts["audio"]); + } + + /// + /// Verify that additional counts are handled when one or both sides do not report any keys. + /// + [Fact] + public void MergeUsage_AdditionalCounts_HandlesNullDictionaries() + { + // Arrange + UsageDetails withCounts = CreateUsage(input: null, output: null, total: null, new() { ["cached"] = 7 }); + UsageDetails withoutCounts = CreateUsage(input: 1, output: 2, total: 3); + + // Act + UsageDetails? oneSide = UsageAggregationExtensions.MergeUsage(withoutCounts, withCounts); + UsageDetails? bothSides = UsageAggregationExtensions.MergeUsage(withoutCounts, CreateUsage(input: null, output: null, total: null)); + + // Assert + Assert.NotNull(oneSide); + Assert.Equal(7, oneSide!.AdditionalCounts!["cached"]); + Assert.NotNull(bothSides); + Assert.Null(bothSides!.AdditionalCounts); + } + + /// + /// Verify that merging does not mutate either input usage or additional-count dictionary. + /// + [Fact] + public void MergeUsage_DoesNotMutateInputs() + { + // Arrange + UsageDetails current = CreateUsage(2, 3, 5, new() { ["cached"] = 7, ["reasoning"] = 11 }); + UsageDetails incoming = CreateUsage(13, null, 17, new() { ["cached"] = 19, ["audio"] = 23 }); + UsageSnapshot currentBefore = UsageSnapshot.Capture(current); + UsageSnapshot incomingBefore = UsageSnapshot.Capture(incoming); + + // Act + UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.Equal(currentBefore, UsageSnapshot.Capture(current)); + Assert.Equal(incomingBefore, UsageSnapshot.Capture(incoming)); + } + + /// + /// Verify that accumulating usage replaces the running aggregate with new combined instances. + /// + [Fact] + public void AccumulateUsage_SeveralAccumulations_UpdatesAggregate() + { + // Arrange + UsageDetails? aggregate = null; + UsageDetails first = CreateUsage(2, 3, 5, new() { ["cached"] = 7 }); + UsageDetails second = CreateUsage(11, null, 16, new() { ["cached"] = 13, ["audio"] = 17 }); + UsageDetails third = CreateUsage(null, 29, null); + + // Act + UsageAggregationExtensions.AccumulateUsage(ref aggregate, first); + UsageDetails firstAggregate = aggregate!; + UsageAggregationExtensions.AccumulateUsage(ref aggregate, null); + UsageDetails secondAggregate = aggregate!; + UsageAggregationExtensions.AccumulateUsage(ref aggregate, second); + UsageAggregationExtensions.AccumulateUsage(ref aggregate, third); + + // Assert + Assert.NotSame(first, firstAggregate); + Assert.NotSame(firstAggregate, secondAggregate); + Assert.Equal(13, aggregate!.InputTokenCount); + Assert.Equal(32, aggregate.OutputTokenCount); + Assert.Equal(21, aggregate.TotalTokenCount); + Assert.Equal(20, aggregate.AdditionalCounts!["cached"]); + Assert.Equal(17, aggregate.AdditionalCounts["audio"]); + } + + /// + /// Verify that every strongly-typed counter exposed by is summed, not just the + /// three headline token counts. Providers such as the GitHub Copilot agent report + /// , and reasoning tokens are common for OpenAI-family + /// models, so dropping any of these would silently lose provider-reported data. + /// + [Fact] + public void MergeUsage_SumsAllStronglyTypedCounters() + { + // Arrange + UsageDetails current = new() + { + InputTokenCount = 1, + OutputTokenCount = 2, + TotalTokenCount = 3, + CachedInputTokenCount = 4, + ReasoningTokenCount = 5, + InputAudioTokenCount = 6, + InputTextTokenCount = 7, + OutputAudioTokenCount = 8, + OutputTextTokenCount = 9, + }; + UsageDetails incoming = new() + { + InputTokenCount = 10, + OutputTokenCount = 20, + TotalTokenCount = 30, + CachedInputTokenCount = 40, + ReasoningTokenCount = 50, + InputAudioTokenCount = 60, + InputTextTokenCount = 70, + OutputAudioTokenCount = 80, + OutputTextTokenCount = 90, + }; + + // Act + UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.Equal(11, result!.InputTokenCount); + Assert.Equal(22, result.OutputTokenCount); + Assert.Equal(33, result.TotalTokenCount); + Assert.Equal(44, result.CachedInputTokenCount); + Assert.Equal(55, result.ReasoningTokenCount); + Assert.Equal(66, result.InputAudioTokenCount); + Assert.Equal(77, result.InputTextTokenCount); + Assert.Equal(88, result.OutputAudioTokenCount); + Assert.Equal(99, result.OutputTextTokenCount); + } + + /// + /// Verify that the extra strongly-typed counters survive a merge where only one side reports them, which + /// is the common case when a single iteration of a loop reports cached or reasoning tokens. + /// + [Fact] + public void MergeUsage_OneSideOnlyReportsExtraCounters_PreservesThem() + { + // Arrange + UsageDetails current = new() { InputTokenCount = 5 }; + UsageDetails incoming = new() { InputTokenCount = 6, CachedInputTokenCount = 3, ReasoningTokenCount = 4 }; + + // Act + UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.Equal(11, result!.InputTokenCount); + Assert.Equal(3, result.CachedInputTokenCount); + Assert.Equal(4, result.ReasoningTokenCount); + } + + /// + /// intentionally mirrors the semantics of + /// (which is what FunctionInvokingChatClient uses to aggregate usage + /// across its own function-calling turns) while avoiding that method's in-place mutation. This asserts the + /// two agree for every combination of reported and unreported counters. + /// + [Theory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + public void MergeUsage_MatchesUsageDetailsAddSemantics(bool currentReported, bool incomingReported) + { + // Arrange + UsageDetails? current = currentReported ? CreateFullyPopulatedUsage(1) : null; + UsageDetails? incoming = incomingReported ? CreateFullyPopulatedUsage(100) : null; + + UsageDetails expected = new(); + if (current is not null) + { + expected.Add(current); + } + + if (incoming is not null) + { + expected.Add(incoming); + } + + // Act + UsageDetails? actual = UsageAggregationExtensions.MergeUsage(current, incoming); + + // Assert + Assert.NotNull(actual); + foreach (var property in GetTokenCountProperties()) + { + Assert.Equal((long?)property.GetValue(expected), (long?)property.GetValue(actual)); + } + + Assert.Equal( + expected.AdditionalCounts?.OrderBy(static e => e.Key).Select(static e => $"{e.Key}:{e.Value}") ?? [], + actual.AdditionalCounts?.OrderBy(static e => e.Key).Select(static e => $"{e.Key}:{e.Value}") ?? []); + } + + /// + /// Guards against a future counter being added upstream without being summed here. + /// Because every fix site replaces the inner response's usage with a merged instance, an unmerged counter + /// would be silently dropped even on single-iteration runs. + /// + [Fact] + public void MergeUsage_SumsEveryTokenCountPropertyExposedByUsageDetails() + { + // Arrange + UsageDetails current = CreateFullyPopulatedUsage(1); + UsageDetails incoming = CreateFullyPopulatedUsage(100); + + // Act + UsageDetails? merged = UsageAggregationExtensions.MergeUsage(current, incoming); + + // Assert + var properties = GetTokenCountProperties().ToList(); + Assert.NotEmpty(properties); + Assert.NotNull(merged); + foreach (var property in properties) + { + long? currentValue = (long?)property.GetValue(current); + long? incomingValue = (long?)property.GetValue(incoming); + Assert.Equal(currentValue + incomingValue, (long?)property.GetValue(merged)); + } + } + + private static IEnumerable GetTokenCountProperties() + => typeof(UsageDetails) + .GetProperties(BindingFlags.Public | BindingFlags.Instance) + .Where(static p => p.PropertyType == typeof(long?) && p.CanRead && p.CanWrite); + + /// + /// Assigns a distinct value to every settable counter so that a counter which is not + /// summed by produces a detectable mismatch. + /// + private static UsageDetails CreateFullyPopulatedUsage(long seed) + { + UsageDetails usage = new(); + long offset = 0; + foreach (var property in GetTokenCountProperties()) + { + property.SetValue(usage, seed + offset++); + } + + usage.AdditionalCounts = new() { ["provider_specific"] = seed }; + return usage; + } + + /// + /// Every settable property other than must survive the copy, otherwise + /// replacing the inner client's response with an aggregated one would silently discard response metadata. + /// + [Fact] + public void WithAggregatedUsage_ChatResponse_CopiesEverySettablePropertyExceptUsage() + { + // Arrange + ChatMessage message = new(ChatRole.Assistant, "hello"); + ChatResponse original = new([message]) + { + ResponseId = "resp-1", + ConversationId = "conv-1", + ModelId = "model-1", + CreatedAt = DateTimeOffset.UnixEpoch, + FinishReason = ChatFinishReason.Stop, + Usage = CreateUsage(1, 1, 2), + ContinuationToken = new TestContinuationToken(), + RawRepresentation = new object(), + AdditionalProperties = new() { ["key"] = "value" }, + }; + + UsageDetails aggregated = CreateUsage(10, 20, 30); + + // Act + ChatResponse copy = original.WithAggregatedUsage(aggregated); + + // Assert + Assert.NotSame(original, copy); + Assert.Same(aggregated, copy.Usage); + AssertAllSettablePropertiesCopied(original, copy); + Assert.Equal([message], copy.Messages); + } + + /// + /// Every settable property other than must survive the copy, otherwise + /// replacing the inner agent's response with an aggregated one would silently discard response metadata. + /// + [Fact] + public void WithAggregatedUsage_AgentResponse_CopiesEverySettablePropertyExceptUsage() + { + // Arrange + ChatMessage message = new(ChatRole.Assistant, "hello"); + AgentResponse original = new([message]) + { + AgentId = "agent-1", + ResponseId = "resp-1", + CreatedAt = DateTimeOffset.UnixEpoch, + FinishReason = ChatFinishReason.Stop, + Usage = CreateUsage(1, 1, 2), + ContinuationToken = new TestContinuationToken(), + RawRepresentation = new object(), + AdditionalProperties = new() { ["key"] = "value" }, + }; + + UsageDetails aggregated = CreateUsage(10, 20, 30); + + // Act + AgentResponse copy = original.WithAggregatedUsage(aggregated); + + // Assert + Assert.NotSame(original, copy); + Assert.Same(aggregated, copy.Usage); + AssertAllSettablePropertiesCopied(original, copy); + Assert.Equal([message], copy.Messages); + } + + /// + /// When a run returns a transcript spanning multiple invocations, the supplied messages replace those of + /// the final response while the remaining metadata is still carried over. + /// + [Fact] + public void WithAggregatedUsage_AgentResponse_SubstitutesSuppliedMessagesAndRetainsMetadata() + { + // Arrange + AgentResponse original = new([new ChatMessage(ChatRole.Assistant, "last")]) + { + AgentId = "agent-1", + RawRepresentation = new object(), + }; + + List transcript = + [ + new(ChatRole.Assistant, "first"), + new(ChatRole.Assistant, "last"), + ]; + + // Act + AgentResponse copy = original.WithAggregatedUsage(null, transcript); + + // Assert + Assert.Equal(transcript, copy.Messages); + Assert.Equal("agent-1", copy.AgentId); + Assert.Same(original.RawRepresentation, copy.RawRepresentation); + Assert.Null(copy.Usage); + } + + /// + /// The copy must never alias the inner response, since replacing usage on a shared instance is exactly the + /// mutation hazard these helpers exist to avoid. A substitution request therefore always copies. + /// + [Fact] + public void WithAggregatedUsage_ReturnsOriginalOnlyWhenUsageAlreadyMatchesAndNoMessagesSupplied() + { + // Arrange + UsageDetails usage = CreateUsage(1, 2, 3); + ChatResponse chatResponse = new([new ChatMessage(ChatRole.Assistant, "hi")]) { Usage = usage }; + AgentResponse agentResponse = new([new ChatMessage(ChatRole.Assistant, "hi")]) { Usage = usage }; + + // Act & Assert + Assert.Same(chatResponse, chatResponse.WithAggregatedUsage(usage)); + Assert.Same(agentResponse, agentResponse.WithAggregatedUsage(usage)); + Assert.NotSame(chatResponse, chatResponse.WithAggregatedUsage(CreateUsage(1, 2, 3))); + Assert.NotSame(agentResponse, agentResponse.WithAggregatedUsage(usage, [new ChatMessage(ChatRole.Assistant, "other")])); + } + + private static void AssertAllSettablePropertiesCopied(T original, T copy) + { + var properties = typeof(T) + .GetProperties(BindingFlags.Public | BindingFlags.Instance) + .Where(static p => p.CanWrite && p.Name != nameof(AgentResponse.Usage) && p.Name != nameof(AgentResponse.Messages)) + .ToList(); + + Assert.NotEmpty(properties); + foreach (var property in properties) + { + object? expected = property.GetValue(original); + Assert.NotNull(expected); + Assert.Equal(expected, property.GetValue(copy)); + } + } + + private sealed class TestContinuationToken : ResponseContinuationToken + { + public override ReadOnlyMemory ToBytes() => new([1, 2, 3]); + } + + private static UsageDetails CreateUsage(long? input, long? output, long? total, AdditionalPropertiesDictionary? additionalCounts = null) + { + UsageDetails usage = new() + { + InputTokenCount = input, + OutputTokenCount = output, + TotalTokenCount = total, + }; + + if (additionalCounts is not null) + { + usage.AdditionalCounts = additionalCounts; + } + + return usage; + } + + private sealed record UsageSnapshot(long? Input, long? Output, long? Total, string AdditionalCounts) + { + public static UsageSnapshot Capture(UsageDetails usage) + => new( + usage.InputTokenCount, + usage.OutputTokenCount, + usage.TotalTokenCount, + string.Join( + "|", + usage.AdditionalCounts?.OrderBy(static entry => entry.Key).Select(static entry => $"{entry.Key}:{entry.Value}") ?? [])); + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs index 30c8a6941cd..93fb5680ed4 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/MessageMergerTests.cs @@ -275,6 +275,66 @@ public void Test_MessageMerger_DoesNotFoldIdentifierlessReasoningIntoDifferentRo response.Messages[1].Role.Should().Be(ChatRole.Tool); } + /// + /// Verify that usage from merged response buckets is aggregated with distinct token values and additional counts. + /// + [Fact] + public void Test_MessageMerger_AggregatesUsageAndAdditionalCounts() + { + // Arrange + const string ResponseId1 = "response-1"; + const string ResponseId2 = "response-2"; + MessageMerger merger = new(); + + merger.AddUpdate(new AgentResponseUpdate(ChatRole.Assistant, "first") + { + ResponseId = ResponseId1, + MessageId = "message-1", + }); + merger.AddUpdate(new AgentResponseUpdate(ChatRole.Assistant, + [new UsageContent(new UsageDetails + { + InputTokenCount = 2, + OutputTokenCount = 3, + TotalTokenCount = 5, + AdditionalCounts = new() { ["cached"] = 7, ["reasoning"] = 11 }, + })]) + { + ResponseId = ResponseId1, + MessageId = "message-1", + }); + merger.AddUpdate(new AgentResponseUpdate(ChatRole.Assistant, "second") + { + ResponseId = ResponseId2, + MessageId = "message-2", + }); + merger.AddUpdate(new AgentResponseUpdate(ChatRole.Assistant, + [new UsageContent(new UsageDetails + { + InputTokenCount = 29, + OutputTokenCount = 7, + TotalTokenCount = 36, + AdditionalCounts = new() { ["cached"] = 13, ["audio"] = 17 }, + })]) + { + ResponseId = ResponseId2, + MessageId = "message-2", + }); + + // Act + AgentResponse response = merger.ComputeMerged(ResponseId1); + + // Assert + response.Usage.Should().NotBeNull(); + response.Usage!.InputTokenCount.Should().Be(31); + response.Usage.OutputTokenCount.Should().Be(10); + response.Usage.TotalTokenCount.Should().Be(41); + response.Usage.AdditionalCounts.Should().NotBeNull(); + response.Usage.AdditionalCounts!["cached"].Should().Be(20); + response.Usage.AdditionalCounts["reasoning"].Should().Be(11); + response.Usage.AdditionalCounts["audio"].Should().Be(17); + } + private static void AddTextMessage(MessageMerger merger, string responseId, string text, DateTimeOffset? createdAt = null) { merger.AddUpdate(new AgentResponseUpdate From a20aed55286cf934c715c7651c1c9b18c2e33d39 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:00:17 +0000 Subject: [PATCH 2/4] Add max tool approval loop fixes --- .../Harness/ToolApproval/ToolApprovalAgent.cs | 8 +- .../ToolApproval/ToolApprovalAgentTests.cs | 116 ++++++++++++++++++ 2 files changed, 123 insertions(+), 1 deletion(-) diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs b/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs index 365ed97f392..fda35d081a0 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs @@ -156,7 +156,13 @@ protected override async Task RunCoreAsync( // request it surfaces goes to the caller to decide rather than continuing the chain. // Returning here without this call would hand back a response whose approval requests // were already stripped — the empty response the loop exists to avoid. - return await this.InnerAgent.RunAsync(processedMessages, session, options, cancellationToken).ConfigureAwait(false); + var cappedResponse = await this.InnerAgent.RunAsync(processedMessages, session, options, cancellationToken).ConfigureAwait(false); + + // This turn is still part of the same run, so its usage joins the aggregate rather + // than replacing it; otherwise hitting the cap would discard every prior turn's cost. + UsageAggregationExtensions.AccumulateUsage(ref aggregatedUsage, cappedResponse.Usage); + + return cappedResponse.WithAggregatedUsage(aggregatedUsage); } var response = await this.InnerAgent.RunAsync(processedMessages, session, options, cancellationToken).ConfigureAwait(false); diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs index 1bc9d515750..6e0b164d3ae 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs @@ -2614,6 +2614,122 @@ public async Task RunStreamingAsync_AutoApprovedToolRequestedForever_StopsAtIter Assert.NotEmpty(updates.SelectMany(u => u.Contents).OfType()); } + /// + /// Verify that hitting still reports the + /// usage of the entire run. The capped path takes an extra final turn outside the accumulating loop, so a + /// naive early return there would discard every prior turn's cost — the exact under-reporting this + /// aggregation exists to prevent, and the case most likely to involve a large token spend. + /// + [Fact] + public async Task RunAsync_StopsAtIterationCap_AggregatesUsageAcrossEveryTurnAsync() + { + // Arrange — always asks for an auto-approved tool, reporting distinct usage per turn. + var session = new ChatClientAgentSession(); + var callCount = 0; + var innerAgent = new Mock(); + innerAgent + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(() => + { + callCount++; + return new AgentResponse([new ChatMessage(ChatRole.Assistant, + [new ToolApprovalRequestContent($"req{callCount}", new FunctionCallContent($"call{callCount}", "load_skill"))])]) + { + Usage = new UsageDetails + { + InputTokenCount = callCount, + OutputTokenCount = callCount * 10, + TotalTokenCount = callCount * 11, + }, + }; + }); + + var options = new ToolApprovalAgentOptions + { + AutoApprovalRules = [ToolApprovalAgent.AllToolsAutoApprovalRule], + MaxAutoApprovalIterations = 3, + }; + var agent = new ToolApprovalAgent(innerAgent.Object, options); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session); + + // Assert — 3 auto-approving passes plus the final capped turn, all four counted. + Assert.Equal(4, callCount); + Assert.NotNull(response.Usage); + Assert.Equal(1 + 2 + 3 + 4, response.Usage!.InputTokenCount); + Assert.Equal(10 + 20 + 30 + 40, response.Usage.OutputTokenCount); + Assert.Equal(11 + 22 + 33 + 44, response.Usage.TotalTokenCount); + } + + /// + /// Verify the streaming path reports the whole run's usage when the cap is hit. Streaming aggregates by + /// passing every through to the caller, including those from the final capped + /// turn, so no update may be swallowed by the cap branch. + /// + [Fact] + public async Task RunStreamingAsync_StopsAtIterationCap_AggregatesUsageAcrossEveryTurnAsync() + { + // Arrange + var session = new ChatClientAgentSession(); + var callCount = 0; + var innerAgent = new Mock(); + innerAgent + .Protected() + .Setup>("RunCoreStreamingAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns(() => + { + callCount++; + int call = callCount; + return ToAsyncEnumerableAsync([ + new AgentResponseUpdate(ChatRole.Assistant, + new List { new ToolApprovalRequestContent($"req{call}", new FunctionCallContent($"call{call}", "load_skill")) }), + new AgentResponseUpdate(ChatRole.Assistant, + new List + { + new UsageContent(new UsageDetails + { + InputTokenCount = call, + OutputTokenCount = call * 10, + TotalTokenCount = call * 11, + }), + }), + ]); + }); + + var options = new ToolApprovalAgentOptions + { + AutoApprovalRules = [ToolApprovalAgent.AllToolsAutoApprovalRule], + MaxAutoApprovalIterations = 3, + }; + var agent = new ToolApprovalAgent(innerAgent.Object, options); + + // Act + var updates = new List(); + await foreach (var update in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hi")], session)) + { + updates.Add(update); + } + + var response = updates.ToAgentResponse(); + + // Assert + Assert.Equal(4, callCount); + Assert.NotNull(response.Usage); + Assert.Equal(1 + 2 + 3 + 4, response.Usage!.InputTokenCount); + Assert.Equal(10 + 20 + 30 + 40, response.Usage.OutputTokenCount); + Assert.Equal(11 + 22 + 33 + 44, response.Usage.TotalTokenCount); + } + /// /// Verify that the cap defaults to /// when the caller does not configure one. From 0f98c350af78a620ed2d304ec293a80f798b9db0 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:12:30 +0000 Subject: [PATCH 3/4] Fix net472 build break in usage aggregation tests DateTimeOffset.UnixEpoch is not available on .NET Framework 4.7.2, so the WithAggregatedUsage copy tests failed to compile for that target framework. Use an explicit DateTimeOffset instead; the specific instant is irrelevant, the value only needs to be non-default so the copy assertion is meaningful. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Shared/UsageAggregationExtensionsTests.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs index 5c0c126ce0c..fffc1fe282b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs @@ -347,7 +347,7 @@ public void WithAggregatedUsage_ChatResponse_CopiesEverySettablePropertyExceptUs ResponseId = "resp-1", ConversationId = "conv-1", ModelId = "model-1", - CreatedAt = DateTimeOffset.UnixEpoch, + CreatedAt = new DateTimeOffset(2024, 1, 2, 3, 4, 5, TimeSpan.Zero), FinishReason = ChatFinishReason.Stop, Usage = CreateUsage(1, 1, 2), ContinuationToken = new TestContinuationToken(), @@ -380,7 +380,7 @@ public void WithAggregatedUsage_AgentResponse_CopiesEverySettablePropertyExceptU { AgentId = "agent-1", ResponseId = "resp-1", - CreatedAt = DateTimeOffset.UnixEpoch, + CreatedAt = new DateTimeOffset(2024, 1, 2, 3, 4, 5, TimeSpan.Zero), FinishReason = ChatFinishReason.Stop, Usage = CreateUsage(1, 1, 2), ContinuationToken = new TestContinuationToken(), From 59343cd63891d68967e067b009c12ab7863660e6 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 7 Aug 2026 09:06:52 +0000 Subject: [PATCH 4/4] Address PR comments --- .../MessageMerger.cs | 4 +- .../ChatClient/MessageInjectingChatClient.cs | 6 +- .../Harness/Loop/LoopAgent.cs | 12 +- .../Harness/ToolApproval/ToolApprovalAgent.cs | 12 +- .../Usage/UsageAggregationExtensions.cs | 186 ++----- dotnet/src/Shared/Usage/UsageAggregator.cs | 123 +++++ .../MessageInjectingChatClientTests.cs | 55 ++ .../Harness/Loop/LoopAgentTests.cs | 38 +- .../ToolApproval/ToolApprovalAgentTests.cs | 48 ++ .../Shared/DerivedResponseTestTypes.cs | 32 ++ .../Shared/UsageAggregationExtensionsTests.cs | 491 ++++-------------- .../Shared/UsageAggregatorTests.cs | 363 +++++++++++++ 12 files changed, 807 insertions(+), 563 deletions(-) create mode 100644 dotnet/src/Shared/Usage/UsageAggregator.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/DerivedResponseTestTypes.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregatorTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs index bf83e2f7d1a..4dee01a93b2 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/MessageMerger.cs @@ -132,7 +132,7 @@ public AgentResponse ComputeMerged(string primaryResponseId, string? primaryAgen _ = finishReasons.Add(response.FinishReason.Value); } - usage = UsageAggregationExtensions.MergeUsage(usage, response.Usage); + usage = UsageAggregator.Combine(usage, response.Usage); additionalProperties = MergeProperties(additionalProperties, response.AdditionalProperties); } @@ -219,7 +219,7 @@ static AgentResponse MergeResponses(AgentResponse? current, AgentResponse incomi Messages = current.Messages.Concat(incoming.Messages).ToList(), ResponseId = current.ResponseId, RawRepresentation = rawRepresentation, - Usage = UsageAggregationExtensions.MergeUsage(current.Usage, incoming.Usage), + Usage = UsageAggregator.Combine(current.Usage, incoming.Usage), }; } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs index b46f4930bb2..4bbd7620ae6 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/MessageInjectingChatClient.cs @@ -85,20 +85,20 @@ public override async Task GetResponseAsync( { var response = await base.GetResponseAsync(newMessages, options, cancellationToken).ConfigureAwait(false); - UsageAggregationExtensions.AccumulateUsage(ref aggregatedUsage, response.Usage); + UsageAggregator.Accumulate(ref aggregatedUsage, response.Usage); // If the response contains actionable function calls, the parent FunctionInvokingChatClient // loop will iterate — return immediately so it can process them. if (HasActionableFunctionCalls(response.Messages)) { - return response.WithAggregatedUsage(aggregatedUsage); + return response.ApplyAggregatedUsage(aggregatedUsage); } // No actionable function calls. If there are pending injected messages, loop again // to send them to the service. Otherwise, we're done. if (await this.IsQueueEmptyAsync(session, cancellationToken).ConfigureAwait(false)) { - return response.WithAggregatedUsage(aggregatedUsage); + return response.ApplyAggregatedUsage(aggregatedUsage); } // Propagate any ConversationId returned by the service so subsequent iterations diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs index a4087f7c6d5..398dc061565 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs @@ -178,7 +178,7 @@ protected override async Task RunCoreAsync( AgentResponse response = await this.InnerAgent.RunAsync(currentMessages, activeSession, options, cancellationToken).ConfigureAwait(false); iteration++; - UsageAggregationExtensions.AccumulateUsage(ref aggregatedUsage, response.Usage); + UsageAggregator.Accumulate(ref aggregatedUsage, response.Usage); // Record this iteration's on-behalf-of input (before the response it elicited) and the response itself. transcript.AddRange(currentSurfaced); @@ -452,14 +452,14 @@ private static AgentResponseUpdate CreateOnBehalfOfUpdate(ChatMessage message, s } /// - /// Produces the non-streaming run result: either the final iteration's response (when configured) or an - /// aggregated response carrying the full transcript with the final response's metadata. In both cases the - /// usage reported is , covering every iteration of the run. + /// Produces the non-streaming run result from the final iteration's response, which carries either its own + /// messages (when configured) or the full transcript of the run. In both cases the usage reported is + /// , covering every iteration of the run. /// private AgentResponse BuildResult(AgentResponse lastResponse, List transcript, UsageDetails? aggregatedUsage) => this._nonStreamingReturnsLastResponseOnly - ? lastResponse.WithAggregatedUsage(aggregatedUsage) - : lastResponse.WithAggregatedUsage(aggregatedUsage, transcript); + ? lastResponse.ApplyAggregatedUsage(aggregatedUsage) + : lastResponse.ApplyAggregatedUsage(aggregatedUsage, transcript); private static bool HasPendingApprovalRequests(AgentResponse response) { diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs b/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs index fda35d081a0..79a2a160d0d 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/ToolApproval/ToolApprovalAgent.cs @@ -160,23 +160,23 @@ protected override async Task RunCoreAsync( // This turn is still part of the same run, so its usage joins the aggregate rather // than replacing it; otherwise hitting the cap would discard every prior turn's cost. - UsageAggregationExtensions.AccumulateUsage(ref aggregatedUsage, cappedResponse.Usage); + UsageAggregator.Accumulate(ref aggregatedUsage, cappedResponse.Usage); - return cappedResponse.WithAggregatedUsage(aggregatedUsage); + return cappedResponse.ApplyAggregatedUsage(aggregatedUsage); } var response = await this.InnerAgent.RunAsync(processedMessages, session, options, cancellationToken).ConfigureAwait(false); - UsageAggregationExtensions.AccumulateUsage(ref aggregatedUsage, response.Usage); + UsageAggregator.Accumulate(ref aggregatedUsage, response.Usage); // Classify approval requests: auto-approve matching, queue excess, keep first unapproved. bool allAutoApproved = await this.ProcessAndQueueOutboundApprovalRequestsAsync(response.Messages, state, session, options, requestMessages).ConfigureAwait(false); if (!allAutoApproved) { - // Response has real content or an unapproved approval request — return to caller. - // Return a copy carrying the aggregated usage rather than mutating the inner agent's response. - return response.WithAggregatedUsage(aggregatedUsage); + // Response has real content or an unapproved approval request — return to caller, + // reporting the usage accumulated across every turn of the run. + return response.ApplyAggregatedUsage(aggregatedUsage); } // All approval requests were auto-approved. Loop to re-invoke with them injected. diff --git a/dotnet/src/Shared/Usage/UsageAggregationExtensions.cs b/dotnet/src/Shared/Usage/UsageAggregationExtensions.cs index 2f6128e6cde..066070cb1c9 100644 --- a/dotnet/src/Shared/Usage/UsageAggregationExtensions.cs +++ b/dotnet/src/Shared/Usage/UsageAggregationExtensions.cs @@ -6,189 +6,65 @@ namespace Microsoft.Agents.AI; /// -/// Helpers for aggregating across multiple service or agent invocations that -/// make up a single logical run. +/// Extension methods for reporting usage aggregated by on the response that +/// concludes a run. /// -/// -/// Several components re-invoke an inner agent or chat client in a loop within a single run (for example -/// when auto-approving tool calls, injecting messages, or re-running an agent until an evaluator is -/// satisfied). Each inner invocation reports its own usage, and the aggregate must be surfaced to the -/// caller so that the reported token counts reflect the entire run rather than just its final step. -/// internal static class UsageAggregationExtensions { /// - /// Combines two instances into a new instance containing their summed values. - /// - /// The running aggregate, or if nothing has been accumulated yet. - /// The usage reported by the latest invocation, or if none was reported. - /// - /// A new containing the summed token counts and additional counts, or - /// when both and are - /// . - /// - /// - /// Neither argument is mutated, and neither argument is ever returned by reference, since both may be - /// owned and observed by callers. Every strongly-typed counter exposed by is - /// summed, matching the set covered by , so that no provider-reported - /// counter is lost when a merged instance replaces the original. Token counts are summed in a null-aware - /// manner: combining a count with a non-null count yields the non-null count, and - /// combining two counts yields . Entries in - /// are summed per key so that provider-specific counters - /// (such as cached, reasoning, or cost counters) aggregate correctly. - /// - public static UsageDetails? MergeUsage(UsageDetails? current, UsageDetails? incoming) - { - if (current is null && incoming is null) - { - return null; - } - - var merged = new UsageDetails - { - InputTokenCount = AddCounts(current?.InputTokenCount, incoming?.InputTokenCount), - OutputTokenCount = AddCounts(current?.OutputTokenCount, incoming?.OutputTokenCount), - TotalTokenCount = AddCounts(current?.TotalTokenCount, incoming?.TotalTokenCount), - CachedInputTokenCount = AddCounts(current?.CachedInputTokenCount, incoming?.CachedInputTokenCount), - ReasoningTokenCount = AddCounts(current?.ReasoningTokenCount, incoming?.ReasoningTokenCount), - InputAudioTokenCount = AddCounts(current?.InputAudioTokenCount, incoming?.InputAudioTokenCount), - InputTextTokenCount = AddCounts(current?.InputTextTokenCount, incoming?.InputTextTokenCount), - OutputAudioTokenCount = AddCounts(current?.OutputAudioTokenCount, incoming?.OutputAudioTokenCount), - OutputTextTokenCount = AddCounts(current?.OutputTextTokenCount, incoming?.OutputTextTokenCount), - }; - - AdditionalPropertiesDictionary? additionalCounts = MergeAdditionalCounts(current?.AdditionalCounts, incoming?.AdditionalCounts); - if (additionalCounts is not null) - { - merged.AdditionalCounts = additionalCounts; - } - - return merged; - } - - /// - /// Adds the usage into the running aggregate referenced by - /// , replacing it with a new combined instance. - /// - /// The running aggregate to update. May be . - /// The usage reported by the latest invocation, or if none was reported. - public static void AccumulateUsage(ref UsageDetails? current, UsageDetails? incoming) - => current = MergeUsage(current, incoming); - - /// - /// Returns a reporting in place of the usage - /// carried by , which typically covers only the final service call of a run. + /// Reports on in place of the usage it + /// carries, which typically covers only the final service call of a run. /// /// The response produced by the final call of the run. /// The usage accumulated across every call that made up the run. /// - /// The messages the returned response should carry, or to keep those of - /// . Used when a run returns a transcript spanning multiple calls. + /// The messages the response should carry, or to keep those it already has. Used + /// when a run returns a transcript spanning multiple calls. /// + /// The same instance, updated in place. /// - /// A copy is returned rather than the usage being assigned onto , because the - /// inner client may still own and observe that instance. The original is returned unchanged only when it - /// already carries exactly the usage to report and no message substitution is requested. + /// The supplied response is updated rather than copied, so that a derived response type returned by an + /// inner client (along with any state it carries) survives the aggregation. This matches how + /// reports the usage it accumulates across function-calling + /// iterations. Only the response is mutated: is a freshly combined + /// instance, so no owned by an inner client is modified. /// - public static ChatResponse WithAggregatedUsage(this ChatResponse response, UsageDetails? aggregatedUsage, IList? messages = null) + public static ChatResponse ApplyAggregatedUsage(this ChatResponse response, UsageDetails? aggregatedUsage, IList? messages = null) { - if (messages is null && ReferenceEquals(response.Usage, aggregatedUsage)) + if (messages is not null) { - return response; + response.Messages = messages; } - return new ChatResponse(messages ?? response.Messages) - { - ResponseId = response.ResponseId, - ConversationId = response.ConversationId, - ModelId = response.ModelId, - CreatedAt = response.CreatedAt, - FinishReason = response.FinishReason, - Usage = aggregatedUsage, - ContinuationToken = response.ContinuationToken, - RawRepresentation = response.RawRepresentation, - AdditionalProperties = response.AdditionalProperties, - }; + response.Usage = aggregatedUsage; + return response; } /// - /// Returns an reporting in place of the usage - /// carried by , which typically covers only the final invocation of a run. + /// Reports on in place of the usage it + /// carries, which typically covers only the final invocation of a run. /// /// The response produced by the final invocation of the run. /// The usage accumulated across every invocation that made up the run. /// - /// The messages the returned response should carry, or to keep those of - /// . Used when a run returns a transcript spanning multiple invocations. + /// The messages the response should carry, or to keep those it already has. Used + /// when a run returns a transcript spanning multiple invocations. /// + /// The same instance, updated in place. /// - /// A copy is returned rather than the usage being assigned onto , because the - /// inner agent may still own and observe that instance. The original is returned unchanged only when it - /// already carries exactly the usage to report and no message substitution is requested. + /// The supplied response is updated rather than copied, so that a derived response type returned by an + /// inner agent (such as , along with any state it carries) survives the + /// aggregation. Only the response is mutated: is a freshly combined + /// instance, so no owned by an inner agent is modified. /// - public static AgentResponse WithAggregatedUsage(this AgentResponse response, UsageDetails? aggregatedUsage, IList? messages = null) + public static AgentResponse ApplyAggregatedUsage(this AgentResponse response, UsageDetails? aggregatedUsage, IList? messages = null) { - if (messages is null && ReferenceEquals(response.Usage, aggregatedUsage)) - { - return response; - } - - return new AgentResponse(messages ?? response.Messages) - { - AgentId = response.AgentId, - ResponseId = response.ResponseId, - CreatedAt = response.CreatedAt, - FinishReason = response.FinishReason, - Usage = aggregatedUsage, - ContinuationToken = response.ContinuationToken, - RawRepresentation = response.RawRepresentation, - AdditionalProperties = response.AdditionalProperties, - }; - } - - /// - /// Adds two nullable counts, treating as "not reported" rather than as zero so - /// that an aggregate only reports a count when at least one contributor reported one. - /// - private static long? AddCounts(long? current, long? incoming) - => current is null ? incoming : incoming is null ? current : current + incoming; - - /// - /// Produces a new dictionary containing the per-key sums of the supplied additional counts, or - /// when neither side has any entries. - /// - private static AdditionalPropertiesDictionary? MergeAdditionalCounts( - AdditionalPropertiesDictionary? current, - AdditionalPropertiesDictionary? incoming) - { - bool hasCurrent = current is { Count: > 0 }; - bool hasIncoming = incoming is { Count: > 0 }; - - if (!hasCurrent && !hasIncoming) - { - return null; - } - - var merged = new AdditionalPropertiesDictionary(); - - if (hasCurrent) - { - foreach (var entry in current!) - { - merged[entry.Key] = entry.Value; - } - } - - if (hasIncoming) + if (messages is not null) { - foreach (var entry in incoming!) - { - merged[entry.Key] = merged.TryGetValue(entry.Key, out long existing) - ? existing + entry.Value - : entry.Value; - } + response.Messages = messages; } - return merged; + response.Usage = aggregatedUsage; + return response; } } diff --git a/dotnet/src/Shared/Usage/UsageAggregator.cs b/dotnet/src/Shared/Usage/UsageAggregator.cs new file mode 100644 index 00000000000..d34b406db84 --- /dev/null +++ b/dotnet/src/Shared/Usage/UsageAggregator.cs @@ -0,0 +1,123 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Combines reported by the individual service or agent invocations that make up +/// a single logical run. +/// +/// +/// Several components re-invoke an inner agent or chat client in a loop within a single run (for example +/// when auto-approving tool calls, injecting messages, or re-running an agent until an evaluator is +/// satisfied). Each inner invocation reports its own usage, and the aggregate must be surfaced to the +/// caller so that the reported token counts reflect the entire run rather than just its final step. +/// +internal static class UsageAggregator +{ + /// + /// Combines two instances into a new instance containing their summed values. + /// + /// The running aggregate, or if nothing has been accumulated yet. + /// The usage reported by the latest invocation, or if none was reported. + /// + /// A new containing the summed token counts and additional counts, or + /// when both and are + /// . + /// + /// + /// Neither argument is mutated, and neither argument is ever returned by reference, since both may be + /// owned and observed by callers and combines in place. Every + /// strongly-typed counter exposed by is summed, matching the set covered by + /// , so that no provider-reported counter is lost when a combined instance + /// replaces the original. Token counts are summed in a null-aware manner: combining a + /// count with a non-null count yields the non-null count, and combining two + /// counts yields . Entries in + /// are summed per key so that provider-specific counters + /// (such as cached, reasoning, or cost counters) aggregate correctly. + /// + public static UsageDetails? Combine(UsageDetails? current, UsageDetails? incoming) + { + if (current is null && incoming is null) + { + return null; + } + + var combined = new UsageDetails + { + InputTokenCount = AddCounts(current?.InputTokenCount, incoming?.InputTokenCount), + OutputTokenCount = AddCounts(current?.OutputTokenCount, incoming?.OutputTokenCount), + TotalTokenCount = AddCounts(current?.TotalTokenCount, incoming?.TotalTokenCount), + CachedInputTokenCount = AddCounts(current?.CachedInputTokenCount, incoming?.CachedInputTokenCount), + ReasoningTokenCount = AddCounts(current?.ReasoningTokenCount, incoming?.ReasoningTokenCount), + InputAudioTokenCount = AddCounts(current?.InputAudioTokenCount, incoming?.InputAudioTokenCount), + InputTextTokenCount = AddCounts(current?.InputTextTokenCount, incoming?.InputTextTokenCount), + OutputAudioTokenCount = AddCounts(current?.OutputAudioTokenCount, incoming?.OutputAudioTokenCount), + OutputTextTokenCount = AddCounts(current?.OutputTextTokenCount, incoming?.OutputTextTokenCount), + }; + + AdditionalPropertiesDictionary? additionalCounts = CombineAdditionalCounts(current?.AdditionalCounts, incoming?.AdditionalCounts); + if (additionalCounts is not null) + { + combined.AdditionalCounts = additionalCounts; + } + + return combined; + } + + /// + /// Adds the usage into the running aggregate referenced by + /// , replacing it with a new combined instance. + /// + /// The running aggregate to update. May be . + /// The usage reported by the latest invocation, or if none was reported. + public static void Accumulate(ref UsageDetails? current, UsageDetails? incoming) + => current = Combine(current, incoming); + + /// + /// Adds two nullable counts, treating as "not reported" rather than as zero so + /// that an aggregate only reports a count when at least one contributor reported one. + /// + private static long? AddCounts(long? current, long? incoming) + => current is null ? incoming : incoming is null ? current : current + incoming; + + /// + /// Produces a new dictionary containing the per-key sums of the supplied additional counts, or + /// when neither side has any entries. + /// + private static AdditionalPropertiesDictionary? CombineAdditionalCounts( + AdditionalPropertiesDictionary? current, + AdditionalPropertiesDictionary? incoming) + { + bool hasCurrent = current is { Count: > 0 }; + bool hasIncoming = incoming is { Count: > 0 }; + + if (!hasCurrent && !hasIncoming) + { + return null; + } + + var combined = new AdditionalPropertiesDictionary(); + + if (hasCurrent) + { + foreach (var entry in current!) + { + combined[entry.Key] = entry.Value; + } + } + + if (hasIncoming) + { + foreach (var entry in incoming!) + { + combined[entry.Key] = combined.TryGetValue(entry.Key, out long existing) + ? existing + entry.Value + : entry.Value; + } + } + + return combined; + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs index 9e8159d6f20..ce6f20681b8 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/MessageInjectingChatClientTests.cs @@ -655,6 +655,61 @@ public async Task RunAsync_LoopsInternally_AggregatesUsageAcrossIterationsAsync( Assert.Equal(57, response.Usage.TotalTokenCount); } + /// + /// Verifies that a derived returned by the underlying client survives the + /// injection loop. Aggregated usage is reported by updating the client's own response rather than by + /// building a replacement, so a subclass keeps its runtime type and the state it carries. + /// + [Fact] + public async Task RunAsync_InnerClientReturnsDerivedResponse_PreservesRuntimeTypeWhileAggregatingUsageAsync() + { + // Arrange + int serviceCallCount = 0; + Mock mockService = new(); + MessageInjectingChatClient? injectorRef = null; + ChatClientAgentSession? sessionRef = null; + + mockService.Setup( + s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(async (IEnumerable msgs, ChatOptions? _, CancellationToken ct) => + { + serviceCallCount++; + if (serviceCallCount < 3) + { + // Enqueue a message so the injection loop runs again. + await injectorRef!.EnqueueMessagesAsync(sessionRef!, [new ChatMessage(ChatRole.User, $"injected {serviceCallCount}")], ct); + } + + return new TestDerivedChatResponse([new(ChatRole.Assistant, $"response {serviceCallCount}")]) + { + DerivedState = $"call {serviceCallCount}", + Usage = new UsageDetails { InputTokenCount = serviceCallCount, OutputTokenCount = serviceCallCount * 10 }, + }; + }); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + EnableMessageInjection = true, + }); + + injectorRef = agent.ChatClient.GetService()!; + + // Act + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + sessionRef = session; + var response = await agent.RunAsync([new(ChatRole.User, "original")], session); + + // Assert — the underlying response type reaches the caller intact, carrying the whole run's usage. + Assert.Equal(3, serviceCallCount); + var derived = Assert.IsType(response.RawRepresentation); + Assert.Equal("call 3", derived.DerivedState); + Assert.Equal(1 + 2 + 3, derived.Usage!.InputTokenCount); + Assert.Equal(10 + 20 + 30, derived.Usage.OutputTokenCount); + } + /// /// Verifies that a single service call surfaces its usage unchanged and that the underlying /// client's usage instance is not mutated. diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentTests.cs index 866445cd1eb..b6161da73a0 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentTests.cs @@ -165,10 +165,8 @@ public async Task RunAsync_PredicateLoopsUntilFalse_AggregatesAllIterationsAsync } /// - /// Verify that the aggregated transcript is returned even when no iteration reports usage. The shared - /// WithAggregatedUsage helper has a fast path that returns the original response when its usage is - /// already reference-equal to the aggregate (which is the case when both are ), so - /// this pins that the fast path can never suppress transcript aggregation. + /// Verify that the aggregated transcript is returned even when no iteration reports usage, so that + /// transcript aggregation is never coupled to whether any usage happened to be reported. /// [Fact] public async Task RunAsync_AggregatedTranscript_ReturnsFullTranscriptWhenNoUsageReportedAsync() @@ -211,6 +209,38 @@ public async Task RunAsync_LastResponseOnly_ReturnsFinalResponseAsync() Assert.Single(response.Messages); } + /// + /// Verify that a derived returned by the inner agent survives the loop in both + /// result modes. The transcript and aggregated usage are applied to the final iteration's response rather + /// than to a replacement, so a subclass keeps its runtime type and the state it carries. + /// + [Theory] + [InlineData(true)] + [InlineData(false)] + public async Task RunAsync_InnerAgentReturnsDerivedResponse_PreservesRuntimeTypeWhileAggregatingUsageAsync(bool lastResponseOnly) + { + // Arrange + var capture = new InnerAgentCapture(call => + new TestDerivedAgentResponse([new ChatMessage(ChatRole.Assistant, $"iteration {call}")]) + { + DerivedState = $"iteration {call}", + Usage = new UsageDetails { InputTokenCount = call, OutputTokenCount = call * 10 }, + }); + var evaluator = While(ctx => ctx.LastResponse.Text != "iteration 3"); + var options = new LoopAgentOptions { NonStreamingReturnsLastResponseOnly = lastResponseOnly }; + var agent = new LoopAgent(capture.Agent, evaluator, options); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "go")], new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, capture.CallCount); + var derived = Assert.IsType(response); + Assert.Equal("iteration 3", derived.DerivedState); + Assert.Equal(1 + 2 + 3, derived.Usage!.InputTokenCount); + Assert.Equal(10 + 20 + 30, derived.Usage.OutputTokenCount); + } + /// /// Verify that the caller's initial messages are sent once and a re-invocation without feedback sends none. /// diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs index 6e0b164d3ae..b24f7338dff 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/ToolApproval/ToolApprovalAgentTests.cs @@ -2667,6 +2667,54 @@ public async Task RunAsync_StopsAtIterationCap_AggregatesUsageAcrossEveryTurnAsy Assert.Equal(11 + 22 + 33 + 44, response.Usage.TotalTokenCount); } + /// + /// Verify that a derived returned by the inner agent survives the + /// auto-approval loop. Usage is reported by updating the inner agent's response rather than by building a + /// replacement, so a subclass such as AgentResponse<T> keeps its runtime type and the state + /// it carries instead of being silently downgraded to a base response. + /// + [Fact] + public async Task RunAsync_InnerAgentReturnsDerivedResponse_PreservesRuntimeTypeWhileAggregatingUsageAsync() + { + // Arrange — turn 1 asks for an auto-approved tool, turn 2 answers; both return a derived response. + var session = new ChatClientAgentSession(); + var callCount = 0; + var innerAgent = new Mock(); + innerAgent + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(() => + { + callCount++; + IList messages = callCount == 1 + ? [new ChatMessage(ChatRole.Assistant, [new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "load_skill"))])] + : [new ChatMessage(ChatRole.Assistant, "done")]; + + return new TestDerivedAgentResponse(messages) + { + DerivedState = $"turn{callCount}", + Usage = new UsageDetails { InputTokenCount = callCount, OutputTokenCount = callCount * 10 }, + }; + }); + + var options = new ToolApprovalAgentOptions { AutoApprovalRules = [ToolApprovalAgent.AllToolsAutoApprovalRule] }; + var agent = new ToolApprovalAgent(innerAgent.Object, options); + + // Act + var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session); + + // Assert + Assert.Equal(2, callCount); + var derived = Assert.IsType(response); + Assert.Equal("turn2", derived.DerivedState); + Assert.Equal(1 + 2, derived.Usage!.InputTokenCount); + Assert.Equal(10 + 20, derived.Usage.OutputTokenCount); + } + /// /// Verify the streaming path reports the whole run's usage when the cap is hit. Streaming aggregates by /// passing every through to the caller, including those from the final capped diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/DerivedResponseTestTypes.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/DerivedResponseTestTypes.cs new file mode 100644 index 00000000000..f38eb505ede --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/DerivedResponseTestTypes.cs @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.UnitTests; + +/// +/// An subclass carrying extra state, standing in for the derived response types an +/// inner agent may return (such as AgentResponse<T>, which carries a deserialized result). +/// +/// +/// Used to pin that components which adjust a response on the way out — for example to report usage aggregated +/// across a loop — do so without downgrading it to a base . +/// +internal sealed class TestDerivedAgentResponse(IList messages) : AgentResponse(messages) +{ + public string? DerivedState { get; set; } +} + +/// +/// A subclass carrying extra state, standing in for the derived response types a +/// custom may return. +/// +/// +/// Used to pin that components which adjust a response on the way out — for example to report usage aggregated +/// across a loop — do so without downgrading it to a base . +/// +internal sealed class TestDerivedChatResponse(IList messages) : ChatResponse(messages) +{ + public string? DerivedState { get; set; } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs index fffc1fe282b..aa3af6e520c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregationExtensionsTests.cs @@ -14,331 +14,12 @@ namespace Microsoft.Agents.AI.UnitTests; public class UsageAggregationExtensionsTests { /// - /// Verify that merging two null usage values returns null. + /// Aggregated usage is reported by updating the response in place rather than by copying it, so that a + /// derived response type returned by an inner client survives. Only is + /// touched; every other property is left exactly as the inner client set it. /// [Fact] - public void MergeUsage_BothInputsNull_ReturnsNull() - { - // Arrange, Act - UsageDetails? result = UsageAggregationExtensions.MergeUsage(null, null); - - // Assert - Assert.Null(result); - } - - /// - /// Verify that merging one null usage value returns a new copy of the non-null usage value. - /// - [Theory] - [InlineData(true)] - [InlineData(false)] - public void MergeUsage_OneInputNull_ReturnsNewCopy(bool currentIsNull) - { - // Arrange - UsageDetails usage = CreateUsage(2, 3, 5, new() { ["cached"] = 7 }); - UsageSnapshot before = UsageSnapshot.Capture(usage); - - // Act - UsageDetails? result = currentIsNull - ? UsageAggregationExtensions.MergeUsage(null, usage) - : UsageAggregationExtensions.MergeUsage(usage, null); - - // Assert - Assert.NotNull(result); - Assert.NotSame(usage, result); - Assert.NotSame(usage.AdditionalCounts, result!.AdditionalCounts); - Assert.Equal(2, result.InputTokenCount); - Assert.Equal(3, result.OutputTokenCount); - Assert.Equal(5, result.TotalTokenCount); - Assert.Equal(7, result.AdditionalCounts!["cached"]); - Assert.Equal(before, UsageSnapshot.Capture(usage)); - } - - /// - /// Verify that token counts are summed while preserving null as not reported. - /// - [Fact] - public void MergeUsage_TokenCounts_SumsNullAware() - { - // Arrange - UsageDetails current = CreateUsage(input: 2, output: null, total: null); - UsageDetails incoming = CreateUsage(input: 11, output: 5, total: null); - - // Act - UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); - - // Assert - Assert.NotNull(result); - Assert.NotSame(current, result); - Assert.NotSame(incoming, result); - Assert.Equal(13, result!.InputTokenCount); - Assert.Equal(5, result.OutputTokenCount); - Assert.Null(result.TotalTokenCount); - } - - /// - /// Verify that additional counts are summed per key and preserve disjoint keys. - /// - [Fact] - public void MergeUsage_AdditionalCounts_SumsOverlappingAndUnionsDisjointKeys() - { - // Arrange - UsageDetails current = CreateUsage(input: null, output: null, total: null, new() { ["cached"] = 2, ["reasoning"] = 3 }); - UsageDetails incoming = CreateUsage(input: null, output: null, total: null, new() { ["cached"] = 11, ["audio"] = 29 }); - - // Act - UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); - - // Assert - Assert.NotNull(result); - Assert.NotSame(current.AdditionalCounts, result!.AdditionalCounts); - Assert.NotSame(incoming.AdditionalCounts, result.AdditionalCounts); - Assert.Equal(13, result.AdditionalCounts!["cached"]); - Assert.Equal(3, result.AdditionalCounts["reasoning"]); - Assert.Equal(29, result.AdditionalCounts["audio"]); - } - - /// - /// Verify that additional counts are handled when one or both sides do not report any keys. - /// - [Fact] - public void MergeUsage_AdditionalCounts_HandlesNullDictionaries() - { - // Arrange - UsageDetails withCounts = CreateUsage(input: null, output: null, total: null, new() { ["cached"] = 7 }); - UsageDetails withoutCounts = CreateUsage(input: 1, output: 2, total: 3); - - // Act - UsageDetails? oneSide = UsageAggregationExtensions.MergeUsage(withoutCounts, withCounts); - UsageDetails? bothSides = UsageAggregationExtensions.MergeUsage(withoutCounts, CreateUsage(input: null, output: null, total: null)); - - // Assert - Assert.NotNull(oneSide); - Assert.Equal(7, oneSide!.AdditionalCounts!["cached"]); - Assert.NotNull(bothSides); - Assert.Null(bothSides!.AdditionalCounts); - } - - /// - /// Verify that merging does not mutate either input usage or additional-count dictionary. - /// - [Fact] - public void MergeUsage_DoesNotMutateInputs() - { - // Arrange - UsageDetails current = CreateUsage(2, 3, 5, new() { ["cached"] = 7, ["reasoning"] = 11 }); - UsageDetails incoming = CreateUsage(13, null, 17, new() { ["cached"] = 19, ["audio"] = 23 }); - UsageSnapshot currentBefore = UsageSnapshot.Capture(current); - UsageSnapshot incomingBefore = UsageSnapshot.Capture(incoming); - - // Act - UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); - - // Assert - Assert.NotNull(result); - Assert.Equal(currentBefore, UsageSnapshot.Capture(current)); - Assert.Equal(incomingBefore, UsageSnapshot.Capture(incoming)); - } - - /// - /// Verify that accumulating usage replaces the running aggregate with new combined instances. - /// - [Fact] - public void AccumulateUsage_SeveralAccumulations_UpdatesAggregate() - { - // Arrange - UsageDetails? aggregate = null; - UsageDetails first = CreateUsage(2, 3, 5, new() { ["cached"] = 7 }); - UsageDetails second = CreateUsage(11, null, 16, new() { ["cached"] = 13, ["audio"] = 17 }); - UsageDetails third = CreateUsage(null, 29, null); - - // Act - UsageAggregationExtensions.AccumulateUsage(ref aggregate, first); - UsageDetails firstAggregate = aggregate!; - UsageAggregationExtensions.AccumulateUsage(ref aggregate, null); - UsageDetails secondAggregate = aggregate!; - UsageAggregationExtensions.AccumulateUsage(ref aggregate, second); - UsageAggregationExtensions.AccumulateUsage(ref aggregate, third); - - // Assert - Assert.NotSame(first, firstAggregate); - Assert.NotSame(firstAggregate, secondAggregate); - Assert.Equal(13, aggregate!.InputTokenCount); - Assert.Equal(32, aggregate.OutputTokenCount); - Assert.Equal(21, aggregate.TotalTokenCount); - Assert.Equal(20, aggregate.AdditionalCounts!["cached"]); - Assert.Equal(17, aggregate.AdditionalCounts["audio"]); - } - - /// - /// Verify that every strongly-typed counter exposed by is summed, not just the - /// three headline token counts. Providers such as the GitHub Copilot agent report - /// , and reasoning tokens are common for OpenAI-family - /// models, so dropping any of these would silently lose provider-reported data. - /// - [Fact] - public void MergeUsage_SumsAllStronglyTypedCounters() - { - // Arrange - UsageDetails current = new() - { - InputTokenCount = 1, - OutputTokenCount = 2, - TotalTokenCount = 3, - CachedInputTokenCount = 4, - ReasoningTokenCount = 5, - InputAudioTokenCount = 6, - InputTextTokenCount = 7, - OutputAudioTokenCount = 8, - OutputTextTokenCount = 9, - }; - UsageDetails incoming = new() - { - InputTokenCount = 10, - OutputTokenCount = 20, - TotalTokenCount = 30, - CachedInputTokenCount = 40, - ReasoningTokenCount = 50, - InputAudioTokenCount = 60, - InputTextTokenCount = 70, - OutputAudioTokenCount = 80, - OutputTextTokenCount = 90, - }; - - // Act - UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); - - // Assert - Assert.NotNull(result); - Assert.Equal(11, result!.InputTokenCount); - Assert.Equal(22, result.OutputTokenCount); - Assert.Equal(33, result.TotalTokenCount); - Assert.Equal(44, result.CachedInputTokenCount); - Assert.Equal(55, result.ReasoningTokenCount); - Assert.Equal(66, result.InputAudioTokenCount); - Assert.Equal(77, result.InputTextTokenCount); - Assert.Equal(88, result.OutputAudioTokenCount); - Assert.Equal(99, result.OutputTextTokenCount); - } - - /// - /// Verify that the extra strongly-typed counters survive a merge where only one side reports them, which - /// is the common case when a single iteration of a loop reports cached or reasoning tokens. - /// - [Fact] - public void MergeUsage_OneSideOnlyReportsExtraCounters_PreservesThem() - { - // Arrange - UsageDetails current = new() { InputTokenCount = 5 }; - UsageDetails incoming = new() { InputTokenCount = 6, CachedInputTokenCount = 3, ReasoningTokenCount = 4 }; - - // Act - UsageDetails? result = UsageAggregationExtensions.MergeUsage(current, incoming); - - // Assert - Assert.NotNull(result); - Assert.Equal(11, result!.InputTokenCount); - Assert.Equal(3, result.CachedInputTokenCount); - Assert.Equal(4, result.ReasoningTokenCount); - } - - /// - /// intentionally mirrors the semantics of - /// (which is what FunctionInvokingChatClient uses to aggregate usage - /// across its own function-calling turns) while avoiding that method's in-place mutation. This asserts the - /// two agree for every combination of reported and unreported counters. - /// - [Theory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - public void MergeUsage_MatchesUsageDetailsAddSemantics(bool currentReported, bool incomingReported) - { - // Arrange - UsageDetails? current = currentReported ? CreateFullyPopulatedUsage(1) : null; - UsageDetails? incoming = incomingReported ? CreateFullyPopulatedUsage(100) : null; - - UsageDetails expected = new(); - if (current is not null) - { - expected.Add(current); - } - - if (incoming is not null) - { - expected.Add(incoming); - } - - // Act - UsageDetails? actual = UsageAggregationExtensions.MergeUsage(current, incoming); - - // Assert - Assert.NotNull(actual); - foreach (var property in GetTokenCountProperties()) - { - Assert.Equal((long?)property.GetValue(expected), (long?)property.GetValue(actual)); - } - - Assert.Equal( - expected.AdditionalCounts?.OrderBy(static e => e.Key).Select(static e => $"{e.Key}:{e.Value}") ?? [], - actual.AdditionalCounts?.OrderBy(static e => e.Key).Select(static e => $"{e.Key}:{e.Value}") ?? []); - } - - /// - /// Guards against a future counter being added upstream without being summed here. - /// Because every fix site replaces the inner response's usage with a merged instance, an unmerged counter - /// would be silently dropped even on single-iteration runs. - /// - [Fact] - public void MergeUsage_SumsEveryTokenCountPropertyExposedByUsageDetails() - { - // Arrange - UsageDetails current = CreateFullyPopulatedUsage(1); - UsageDetails incoming = CreateFullyPopulatedUsage(100); - - // Act - UsageDetails? merged = UsageAggregationExtensions.MergeUsage(current, incoming); - - // Assert - var properties = GetTokenCountProperties().ToList(); - Assert.NotEmpty(properties); - Assert.NotNull(merged); - foreach (var property in properties) - { - long? currentValue = (long?)property.GetValue(current); - long? incomingValue = (long?)property.GetValue(incoming); - Assert.Equal(currentValue + incomingValue, (long?)property.GetValue(merged)); - } - } - - private static IEnumerable GetTokenCountProperties() - => typeof(UsageDetails) - .GetProperties(BindingFlags.Public | BindingFlags.Instance) - .Where(static p => p.PropertyType == typeof(long?) && p.CanRead && p.CanWrite); - - /// - /// Assigns a distinct value to every settable counter so that a counter which is not - /// summed by produces a detectable mismatch. - /// - private static UsageDetails CreateFullyPopulatedUsage(long seed) - { - UsageDetails usage = new(); - long offset = 0; - foreach (var property in GetTokenCountProperties()) - { - property.SetValue(usage, seed + offset++); - } - - usage.AdditionalCounts = new() { ["provider_specific"] = seed }; - return usage; - } - - /// - /// Every settable property other than must survive the copy, otherwise - /// replacing the inner client's response with an aggregated one would silently discard response metadata. - /// - [Fact] - public void WithAggregatedUsage_ChatResponse_CopiesEverySettablePropertyExceptUsage() + public void ApplyAggregatedUsage_ChatResponse_UpdatesInPlaceAndLeavesEveryOtherPropertyUntouched() { // Arrange ChatMessage message = new(ChatRole.Assistant, "hello"); @@ -355,24 +36,25 @@ public void WithAggregatedUsage_ChatResponse_CopiesEverySettablePropertyExceptUs AdditionalProperties = new() { ["key"] = "value" }, }; + PropertySnapshot before = PropertySnapshot.Capture(original); UsageDetails aggregated = CreateUsage(10, 20, 30); // Act - ChatResponse copy = original.WithAggregatedUsage(aggregated); + ChatResponse result = original.ApplyAggregatedUsage(aggregated); // Assert - Assert.NotSame(original, copy); - Assert.Same(aggregated, copy.Usage); - AssertAllSettablePropertiesCopied(original, copy); - Assert.Equal([message], copy.Messages); + Assert.Same(original, result); + Assert.Same(aggregated, result.Usage); + Assert.Equal([message], result.Messages); + before.AssertUnchangedExceptUsage(result); } /// - /// Every settable property other than must survive the copy, otherwise - /// replacing the inner agent's response with an aggregated one would silently discard response metadata. + /// The overload behaves identically: the inner agent's response instance is + /// updated in place so that its runtime type and any state it carries survive. /// [Fact] - public void WithAggregatedUsage_AgentResponse_CopiesEverySettablePropertyExceptUsage() + public void ApplyAggregatedUsage_AgentResponse_UpdatesInPlaceAndLeavesEveryOtherPropertyUntouched() { // Arrange ChatMessage message = new(ChatRole.Assistant, "hello"); @@ -388,24 +70,25 @@ public void WithAggregatedUsage_AgentResponse_CopiesEverySettablePropertyExceptU AdditionalProperties = new() { ["key"] = "value" }, }; + PropertySnapshot before = PropertySnapshot.Capture(original); UsageDetails aggregated = CreateUsage(10, 20, 30); // Act - AgentResponse copy = original.WithAggregatedUsage(aggregated); + AgentResponse result = original.ApplyAggregatedUsage(aggregated); // Assert - Assert.NotSame(original, copy); - Assert.Same(aggregated, copy.Usage); - AssertAllSettablePropertiesCopied(original, copy); - Assert.Equal([message], copy.Messages); + Assert.Same(original, result); + Assert.Same(aggregated, result.Usage); + Assert.Equal([message], result.Messages); + before.AssertUnchangedExceptUsage(result); } /// /// When a run returns a transcript spanning multiple invocations, the supplied messages replace those of - /// the final response while the remaining metadata is still carried over. + /// the final response while every other property is left as the inner agent set it. /// [Fact] - public void WithAggregatedUsage_AgentResponse_SubstitutesSuppliedMessagesAndRetainsMetadata() + public void ApplyAggregatedUsage_AgentResponse_SubstitutesSuppliedMessagesAndRetainsMetadata() { // Arrange AgentResponse original = new([new ChatMessage(ChatRole.Assistant, "last")]) @@ -413,6 +96,7 @@ public void WithAggregatedUsage_AgentResponse_SubstitutesSuppliedMessagesAndReta AgentId = "agent-1", RawRepresentation = new object(), }; + object rawRepresentation = original.RawRepresentation; List transcript = [ @@ -421,48 +105,75 @@ public void WithAggregatedUsage_AgentResponse_SubstitutesSuppliedMessagesAndReta ]; // Act - AgentResponse copy = original.WithAggregatedUsage(null, transcript); + AgentResponse result = original.ApplyAggregatedUsage(null, transcript); // Assert - Assert.Equal(transcript, copy.Messages); - Assert.Equal("agent-1", copy.AgentId); - Assert.Same(original.RawRepresentation, copy.RawRepresentation); - Assert.Null(copy.Usage); + Assert.Same(original, result); + Assert.Same(transcript, result.Messages); + Assert.Equal("agent-1", result.AgentId); + Assert.Same(rawRepresentation, result.RawRepresentation); + Assert.Null(result.Usage); } /// - /// The copy must never alias the inner response, since replacing usage on a shared instance is exactly the - /// mutation hazard these helpers exist to avoid. A substitution request therefore always copies. + /// Omitting the messages argument must leave the response's existing messages alone, since most callers + /// only need to correct the reported usage. /// [Fact] - public void WithAggregatedUsage_ReturnsOriginalOnlyWhenUsageAlreadyMatchesAndNoMessagesSupplied() + public void ApplyAggregatedUsage_NoMessagesSupplied_LeavesMessagesUntouched() { // Arrange - UsageDetails usage = CreateUsage(1, 2, 3); - ChatResponse chatResponse = new([new ChatMessage(ChatRole.Assistant, "hi")]) { Usage = usage }; - AgentResponse agentResponse = new([new ChatMessage(ChatRole.Assistant, "hi")]) { Usage = usage }; + List messages = [new(ChatRole.Assistant, "hi")]; + ChatResponse chatResponse = new(messages); + AgentResponse agentResponse = new(messages); + + // Act + ChatResponse chatResult = chatResponse.ApplyAggregatedUsage(CreateUsage(1, 2, 3)); + AgentResponse agentResult = agentResponse.ApplyAggregatedUsage(CreateUsage(1, 2, 3)); - // Act & Assert - Assert.Same(chatResponse, chatResponse.WithAggregatedUsage(usage)); - Assert.Same(agentResponse, agentResponse.WithAggregatedUsage(usage)); - Assert.NotSame(chatResponse, chatResponse.WithAggregatedUsage(CreateUsage(1, 2, 3))); - Assert.NotSame(agentResponse, agentResponse.WithAggregatedUsage(usage, [new ChatMessage(ChatRole.Assistant, "other")])); + // Assert + Assert.Same(messages, chatResult.Messages); + Assert.Same(messages, agentResult.Messages); } - private static void AssertAllSettablePropertiesCopied(T original, T copy) + /// + /// A derived returned by an inner chat client must survive usage aggregation + /// with its additional state intact. Building a replacement base response would silently downgrade it. + /// + [Fact] + public void ApplyAggregatedUsage_DerivedChatResponse_PreservesRuntimeTypeAndDerivedState() { - var properties = typeof(T) - .GetProperties(BindingFlags.Public | BindingFlags.Instance) - .Where(static p => p.CanWrite && p.Name != nameof(AgentResponse.Usage) && p.Name != nameof(AgentResponse.Messages)) - .ToList(); + // Arrange + TestDerivedChatResponse original = new([new ChatMessage(ChatRole.Assistant, "hi")]) { DerivedState = "custom" }; - Assert.NotEmpty(properties); - foreach (var property in properties) - { - object? expected = property.GetValue(original); - Assert.NotNull(expected); - Assert.Equal(expected, property.GetValue(copy)); - } + // Act + ChatResponse result = original.ApplyAggregatedUsage(CreateUsage(1, 2, 3), [new ChatMessage(ChatRole.Assistant, "transcript")]); + + // Assert + TestDerivedChatResponse derived = Assert.IsType(result); + Assert.Same(original, derived); + Assert.Equal("custom", derived.DerivedState); + Assert.Equal(1, derived.Usage!.InputTokenCount); + } + + /// + /// A derived such as AgentResponse<T> must likewise survive usage + /// aggregation, since replacing it with a base response would discard its deserialized result. + /// + [Fact] + public void ApplyAggregatedUsage_DerivedAgentResponse_PreservesRuntimeTypeAndDerivedState() + { + // Arrange + TestDerivedAgentResponse original = new([new ChatMessage(ChatRole.Assistant, "hi")]) { DerivedState = "custom" }; + + // Act + AgentResponse result = original.ApplyAggregatedUsage(CreateUsage(1, 2, 3), [new ChatMessage(ChatRole.Assistant, "transcript")]); + + // Assert + TestDerivedAgentResponse derived = Assert.IsType(result); + Assert.Same(original, derived); + Assert.Equal("custom", derived.DerivedState); + Assert.Equal(1, derived.Usage!.InputTokenCount); } private sealed class TestContinuationToken : ResponseContinuationToken @@ -470,32 +181,38 @@ private sealed class TestContinuationToken : ResponseContinuationToken public override ReadOnlyMemory ToBytes() => new([1, 2, 3]); } - private static UsageDetails CreateUsage(long? input, long? output, long? total, AdditionalPropertiesDictionary? additionalCounts = null) + /// + /// Captures every settable property except and + /// , so that a helper which starts writing to any of them is caught. + /// + private sealed class PropertySnapshot(Dictionary values) { - UsageDetails usage = new() - { - InputTokenCount = input, - OutputTokenCount = output, - TotalTokenCount = total, - }; + public static PropertySnapshot Capture(T response) + where T : notnull + => new(GetProperties().ToDictionary(static p => p, p => p.GetValue(response))); - if (additionalCounts is not null) + public void AssertUnchangedExceptUsage(T response) + where T : notnull { - usage.AdditionalCounts = additionalCounts; + Assert.NotEmpty(values); + foreach (var entry in values) + { + Assert.NotNull(entry.Value); + Assert.Equal(entry.Value, entry.Key.GetValue(response)); + } } - return usage; + private static IEnumerable GetProperties() + => typeof(T) + .GetProperties(BindingFlags.Public | BindingFlags.Instance) + .Where(static p => p.CanWrite && p.Name is not (nameof(AgentResponse.Usage) or nameof(AgentResponse.Messages))); } - private sealed record UsageSnapshot(long? Input, long? Output, long? Total, string AdditionalCounts) - { - public static UsageSnapshot Capture(UsageDetails usage) - => new( - usage.InputTokenCount, - usage.OutputTokenCount, - usage.TotalTokenCount, - string.Join( - "|", - usage.AdditionalCounts?.OrderBy(static entry => entry.Key).Select(static entry => $"{entry.Key}:{entry.Value}") ?? [])); - } + private static UsageDetails CreateUsage(long? input, long? output, long? total) + => new() + { + InputTokenCount = input, + OutputTokenCount = output, + TotalTokenCount = total, + }; } diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregatorTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregatorTests.cs new file mode 100644 index 00000000000..1307f242716 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Shared/UsageAggregatorTests.cs @@ -0,0 +1,363 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using System.Linq; +using System.Reflection; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.UnitTests; + +/// +/// Unit tests for . +/// +public class UsageAggregatorTests +{ + /// + /// Verify that combining two null usage values returns null. + /// + [Fact] + public void Combine_BothInputsNull_ReturnsNull() + { + // Arrange, Act + UsageDetails? result = UsageAggregator.Combine(null, null); + + // Assert + Assert.Null(result); + } + + /// + /// Verify that combining one null usage value returns a new copy of the non-null usage value. + /// + [Theory] + [InlineData(true)] + [InlineData(false)] + public void Combine_OneInputNull_ReturnsNewCopy(bool currentIsNull) + { + // Arrange + UsageDetails usage = CreateUsage(2, 3, 5, new() { ["cached"] = 7 }); + UsageSnapshot before = UsageSnapshot.Capture(usage); + + // Act + UsageDetails? result = currentIsNull + ? UsageAggregator.Combine(null, usage) + : UsageAggregator.Combine(usage, null); + + // Assert + Assert.NotNull(result); + Assert.NotSame(usage, result); + Assert.NotSame(usage.AdditionalCounts, result!.AdditionalCounts); + Assert.Equal(2, result.InputTokenCount); + Assert.Equal(3, result.OutputTokenCount); + Assert.Equal(5, result.TotalTokenCount); + Assert.Equal(7, result.AdditionalCounts!["cached"]); + Assert.Equal(before, UsageSnapshot.Capture(usage)); + } + + /// + /// Verify that token counts are summed while preserving null as not reported. + /// + [Fact] + public void Combine_TokenCounts_SumsNullAware() + { + // Arrange + UsageDetails current = CreateUsage(input: 2, output: null, total: null); + UsageDetails incoming = CreateUsage(input: 11, output: 5, total: null); + + // Act + UsageDetails? result = UsageAggregator.Combine(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.NotSame(current, result); + Assert.NotSame(incoming, result); + Assert.Equal(13, result!.InputTokenCount); + Assert.Equal(5, result.OutputTokenCount); + Assert.Null(result.TotalTokenCount); + } + + /// + /// Verify that additional counts are summed per key and preserve disjoint keys. + /// + [Fact] + public void Combine_AdditionalCounts_SumsOverlappingAndUnionsDisjointKeys() + { + // Arrange + UsageDetails current = CreateUsage(input: null, output: null, total: null, new() { ["cached"] = 2, ["reasoning"] = 3 }); + UsageDetails incoming = CreateUsage(input: null, output: null, total: null, new() { ["cached"] = 11, ["audio"] = 29 }); + + // Act + UsageDetails? result = UsageAggregator.Combine(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.NotSame(current.AdditionalCounts, result!.AdditionalCounts); + Assert.NotSame(incoming.AdditionalCounts, result.AdditionalCounts); + Assert.Equal(13, result.AdditionalCounts!["cached"]); + Assert.Equal(3, result.AdditionalCounts["reasoning"]); + Assert.Equal(29, result.AdditionalCounts["audio"]); + } + + /// + /// Verify that additional counts are handled when one or both sides do not report any keys. + /// + [Fact] + public void Combine_AdditionalCounts_HandlesNullDictionaries() + { + // Arrange + UsageDetails withCounts = CreateUsage(input: null, output: null, total: null, new() { ["cached"] = 7 }); + UsageDetails withoutCounts = CreateUsage(input: 1, output: 2, total: 3); + + // Act + UsageDetails? oneSide = UsageAggregator.Combine(withoutCounts, withCounts); + UsageDetails? bothSides = UsageAggregator.Combine(withoutCounts, CreateUsage(input: null, output: null, total: null)); + + // Assert + Assert.NotNull(oneSide); + Assert.Equal(7, oneSide!.AdditionalCounts!["cached"]); + Assert.NotNull(bothSides); + Assert.Null(bothSides!.AdditionalCounts); + } + + /// + /// Verify that combining does not mutate either input usage or additional-count dictionary. + /// + [Fact] + public void Combine_DoesNotMutateInputs() + { + // Arrange + UsageDetails current = CreateUsage(2, 3, 5, new() { ["cached"] = 7, ["reasoning"] = 11 }); + UsageDetails incoming = CreateUsage(13, null, 17, new() { ["cached"] = 19, ["audio"] = 23 }); + UsageSnapshot currentBefore = UsageSnapshot.Capture(current); + UsageSnapshot incomingBefore = UsageSnapshot.Capture(incoming); + + // Act + UsageDetails? result = UsageAggregator.Combine(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.Equal(currentBefore, UsageSnapshot.Capture(current)); + Assert.Equal(incomingBefore, UsageSnapshot.Capture(incoming)); + } + + /// + /// Verify that accumulating usage replaces the running aggregate with new combined instances. + /// + [Fact] + public void Accumulate_SeveralAccumulations_UpdatesAggregate() + { + // Arrange + UsageDetails? aggregate = null; + UsageDetails first = CreateUsage(2, 3, 5, new() { ["cached"] = 7 }); + UsageDetails second = CreateUsage(11, null, 16, new() { ["cached"] = 13, ["audio"] = 17 }); + UsageDetails third = CreateUsage(null, 29, null); + + // Act + UsageAggregator.Accumulate(ref aggregate, first); + UsageDetails firstAggregate = aggregate!; + UsageAggregator.Accumulate(ref aggregate, null); + UsageDetails secondAggregate = aggregate!; + UsageAggregator.Accumulate(ref aggregate, second); + UsageAggregator.Accumulate(ref aggregate, third); + + // Assert + Assert.NotSame(first, firstAggregate); + Assert.NotSame(firstAggregate, secondAggregate); + Assert.Equal(13, aggregate!.InputTokenCount); + Assert.Equal(32, aggregate.OutputTokenCount); + Assert.Equal(21, aggregate.TotalTokenCount); + Assert.Equal(20, aggregate.AdditionalCounts!["cached"]); + Assert.Equal(17, aggregate.AdditionalCounts["audio"]); + } + + /// + /// Verify that every strongly-typed counter exposed by is summed, not just the + /// three headline token counts. Providers such as the GitHub Copilot agent report + /// , and reasoning tokens are common for OpenAI-family + /// models, so dropping any of these would silently lose provider-reported data. + /// + [Fact] + public void Combine_SumsAllStronglyTypedCounters() + { + // Arrange + UsageDetails current = new() + { + InputTokenCount = 1, + OutputTokenCount = 2, + TotalTokenCount = 3, + CachedInputTokenCount = 4, + ReasoningTokenCount = 5, + InputAudioTokenCount = 6, + InputTextTokenCount = 7, + OutputAudioTokenCount = 8, + OutputTextTokenCount = 9, + }; + UsageDetails incoming = new() + { + InputTokenCount = 10, + OutputTokenCount = 20, + TotalTokenCount = 30, + CachedInputTokenCount = 40, + ReasoningTokenCount = 50, + InputAudioTokenCount = 60, + InputTextTokenCount = 70, + OutputAudioTokenCount = 80, + OutputTextTokenCount = 90, + }; + + // Act + UsageDetails? result = UsageAggregator.Combine(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.Equal(11, result!.InputTokenCount); + Assert.Equal(22, result.OutputTokenCount); + Assert.Equal(33, result.TotalTokenCount); + Assert.Equal(44, result.CachedInputTokenCount); + Assert.Equal(55, result.ReasoningTokenCount); + Assert.Equal(66, result.InputAudioTokenCount); + Assert.Equal(77, result.InputTextTokenCount); + Assert.Equal(88, result.OutputAudioTokenCount); + Assert.Equal(99, result.OutputTextTokenCount); + } + + /// + /// Verify that the extra strongly-typed counters survive a merge where only one side reports them, which + /// is the common case when a single iteration of a loop reports cached or reasoning tokens. + /// + [Fact] + public void Combine_OneSideOnlyReportsExtraCounters_PreservesThem() + { + // Arrange + UsageDetails current = new() { InputTokenCount = 5 }; + UsageDetails incoming = new() { InputTokenCount = 6, CachedInputTokenCount = 3, ReasoningTokenCount = 4 }; + + // Act + UsageDetails? result = UsageAggregator.Combine(current, incoming); + + // Assert + Assert.NotNull(result); + Assert.Equal(11, result!.InputTokenCount); + Assert.Equal(3, result.CachedInputTokenCount); + Assert.Equal(4, result.ReasoningTokenCount); + } + + /// + /// intentionally mirrors the semantics of + /// (which is what FunctionInvokingChatClient uses to aggregate usage + /// across its own function-calling turns) while avoiding that method's in-place mutation. This asserts the + /// two agree for every combination of reported and unreported counters. + /// + [Theory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + public void Combine_MatchesUsageDetailsAddSemantics(bool currentReported, bool incomingReported) + { + // Arrange + UsageDetails? current = currentReported ? CreateFullyPopulatedUsage(1) : null; + UsageDetails? incoming = incomingReported ? CreateFullyPopulatedUsage(100) : null; + + UsageDetails expected = new(); + if (current is not null) + { + expected.Add(current); + } + + if (incoming is not null) + { + expected.Add(incoming); + } + + // Act + UsageDetails? actual = UsageAggregator.Combine(current, incoming); + + // Assert + Assert.NotNull(actual); + foreach (var property in GetTokenCountProperties()) + { + Assert.Equal((long?)property.GetValue(expected), (long?)property.GetValue(actual)); + } + + Assert.Equal( + expected.AdditionalCounts?.OrderBy(static e => e.Key).Select(static e => $"{e.Key}:{e.Value}") ?? [], + actual.AdditionalCounts?.OrderBy(static e => e.Key).Select(static e => $"{e.Key}:{e.Value}") ?? []); + } + + /// + /// Guards against a future counter being added upstream without being summed here. + /// Because every fix site reports a freshly combined instance on the response, an uncombined counter + /// would be silently dropped even on single-iteration runs. + /// + [Fact] + public void Combine_SumsEveryTokenCountPropertyExposedByUsageDetails() + { + // Arrange + UsageDetails current = CreateFullyPopulatedUsage(1); + UsageDetails incoming = CreateFullyPopulatedUsage(100); + + // Act + UsageDetails? merged = UsageAggregator.Combine(current, incoming); + + // Assert + var properties = GetTokenCountProperties().ToList(); + Assert.NotEmpty(properties); + Assert.NotNull(merged); + foreach (var property in properties) + { + long? currentValue = (long?)property.GetValue(current); + long? incomingValue = (long?)property.GetValue(incoming); + Assert.Equal(currentValue + incomingValue, (long?)property.GetValue(merged)); + } + } + + private static IEnumerable GetTokenCountProperties() + => typeof(UsageDetails) + .GetProperties(BindingFlags.Public | BindingFlags.Instance) + .Where(static p => p.PropertyType == typeof(long?) && p.CanRead && p.CanWrite); + + /// + /// Assigns a distinct value to every settable counter so that a counter which is not + /// summed by produces a detectable mismatch. + /// + private static UsageDetails CreateFullyPopulatedUsage(long seed) + { + UsageDetails usage = new(); + long offset = 0; + foreach (var property in GetTokenCountProperties()) + { + property.SetValue(usage, seed + offset++); + } + + usage.AdditionalCounts = new() { ["provider_specific"] = seed }; + return usage; + } + + private sealed record UsageSnapshot(long? Input, long? Output, long? Total, string AdditionalCounts) + { + public static UsageSnapshot Capture(UsageDetails usage) + => new( + usage.InputTokenCount, + usage.OutputTokenCount, + usage.TotalTokenCount, + string.Join( + "|", + usage.AdditionalCounts?.OrderBy(static entry => entry.Key).Select(static entry => $"{entry.Key}:{entry.Value}") ?? [])); + } + + private static UsageDetails CreateUsage(long? input, long? output, long? total, AdditionalPropertiesDictionary? additionalCounts = null) + { + UsageDetails usage = new() + { + InputTokenCount = input, + OutputTokenCount = output, + TotalTokenCount = total, + }; + + if (additionalCounts is not null) + { + usage.AdditionalCounts = additionalCounts; + } + + return usage; + } +}