From a7fb68c34d2ff9b552f62064c48451a2e6891263 Mon Sep 17 00:00:00 2001 From: manjunathshiva Date: Sat, 12 Sep 2026 22:59:22 +0530 Subject: [PATCH] .NET: Preserve cached and reasoning token counts in Foundry Hosting --- .../OutputConverter.cs | 9 +- .../OutputConverterTests.cs | 101 ++++++++++++++++++ 2 files changed, 106 insertions(+), 4 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/OutputConverter.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/OutputConverter.cs index 2a4682b5296..8cbb0aa2ce5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/OutputConverter.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/OutputConverter.cs @@ -557,10 +557,11 @@ private static ResponseUsage ConvertUsage(UsageDetails details, ResponseUsage? e var outputTokens = details.OutputTokenCount ?? 0; var totalTokens = details.TotalTokenCount ?? 0; - var cachedTokens = details.AdditionalCounts?.TryGetValue("InputTokenDetails.CachedTokenCount", out var cached) ?? false - ? cached : 0; - var reasoningTokens = details.AdditionalCounts?.TryGetValue("OutputTokenDetails.ReasoningTokenCount", out var reasoning) ?? false - ? reasoning : 0; + // Prefer the dedicated counters, retaining the legacy dictionary keys as a fallback. + var cachedTokens = details.CachedInputTokenCount ?? + (details.AdditionalCounts?.TryGetValue("InputTokenDetails.CachedTokenCount", out var cached) is true ? cached : 0); + var reasoningTokens = details.ReasoningTokenCount ?? + (details.AdditionalCounts?.TryGetValue("OutputTokenDetails.ReasoningTokenCount", out var reasoning) is true ? reasoning : 0); if (existing is not null) { diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OutputConverterTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OutputConverterTests.cs index db036cfa73f..0aca1b238eb 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OutputConverterTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OutputConverterTests.cs @@ -184,6 +184,107 @@ public async Task ConvertUpdatesToEventsAsync_UsageContent_IncludesUsageInComple Assert.NotNull(completedEvent); } + [Theory] + [InlineData(34304L, 128L, null, null, 34304L, 128L)] + [InlineData(34304L, 128L, 999L, 999L, 34304L, 128L)] + [InlineData(0L, 0L, 999L, 999L, 0L, 0L)] + [InlineData(null, null, null, null, 0L, 0L)] + [InlineData(null, null, 34304L, 128L, 34304L, 128L)] + [InlineData(34304L, null, 999L, 128L, 34304L, 128L)] + [InlineData(null, 128L, 34304L, 999L, 34304L, 128L)] + public async Task ConvertUpdatesToEventsAsync_UsageCounters_PreservesDetailsAsync( + long? cachedTokens, long? reasoningTokens, long? legacyCachedTokens, long? legacyReasoningTokens, + long expectedCachedTokens, long expectedReasoningTokens) + { + // Arrange + var (stream, _) = CreateTestStream(); + var details = new UsageDetails + { + InputTokenCount = 34847, + OutputTokenCount = 1000, + TotalTokenCount = 35847, + CachedInputTokenCount = cachedTokens, + ReasoningTokenCount = reasoningTokens, + }; + if (legacyCachedTokens is { } cached) + { + (details.AdditionalCounts ??= [])["InputTokenDetails.CachedTokenCount"] = cached; + } + + if (legacyReasoningTokens is { } reasoning) + { + (details.AdditionalCounts ??= [])["OutputTokenDetails.ReasoningTokenCount"] = reasoning; + } + + var updates = new[] { new AgentResponseUpdate { Contents = [new UsageContent(details)] } }; + + // Act + var events = new List(); + await foreach (var evt in OutputConverter.ConvertUpdatesToEventsAsync(ToAsync(updates), stream)) + { + events.Add(evt); + } + + // Assert + var completed = Assert.Single(events.OfType()); + var usage = Assert.IsType(completed.Response.Usage); + Assert.Equal(34847, usage.InputTokens); + Assert.Equal(1000, usage.OutputTokens); + Assert.Equal(35847, usage.TotalTokens); + Assert.Equal(expectedCachedTokens, usage.InputTokensDetails.CachedTokens); + Assert.Equal(expectedReasoningTokens, usage.OutputTokensDetails.ReasoningTokens); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ConvertUpdatesToEventsAsync_UsageCounters_AccumulatesInTerminalEventAsync(bool fail) + { + // Arrange + var (stream, _) = CreateTestStream(); + var updates = new List + { + new() + { + Contents = [new UsageContent(new UsageDetails + { + InputTokenCount = 100, OutputTokenCount = 10, TotalTokenCount = 110, + CachedInputTokenCount = 64, ReasoningTokenCount = 2, + })], + }, + new() + { + Contents = [new UsageContent(new UsageDetails + { + InputTokenCount = 200, OutputTokenCount = 20, TotalTokenCount = 220, + CachedInputTokenCount = 128, ReasoningTokenCount = 4, + })], + }, + }; + if (fail) + { + updates.Add(new AgentResponseUpdate { Contents = [new ErrorContent("Test failure")] }); + } + + // Act + var events = new List(); + await foreach (var evt in OutputConverter.ConvertUpdatesToEventsAsync(ToAsync(updates), stream)) + { + events.Add(evt); + } + + // Assert + var terminal = Assert.Single(events); + var usage = Assert.IsType(fail + ? Assert.IsType(terminal).Response.Usage + : Assert.IsType(terminal).Response.Usage); + Assert.Equal(300, usage.InputTokens); + Assert.Equal(30, usage.OutputTokens); + Assert.Equal(330, usage.TotalTokens); + Assert.Equal(192, usage.InputTokensDetails.CachedTokens); + Assert.Equal(6, usage.OutputTokensDetails.ReasoningTokens); + } + [Fact] public async Task ConvertUpdatesToEventsAsync_ReasoningContent_EmitsReasoningEventsAsync() {