From 3d8a4b8d7c8195266df46853bf71bbc3ae4f7096 Mon Sep 17 00:00:00 2001 From: Mark Licata Date: Wed, 7 Oct 2026 08:33:38 -0700 Subject: [PATCH 1/4] .NET: add token and time budgets to LoopAgent Co-Authored-By: Claude Opus 5.5 --- .../Harness/Loop/LoopAgent.cs | 51 ++++++++- .../Harness/Loop/LoopAgentOptions.cs | 41 +++++++ .../Harness/Loop/LoopExitReason.cs | 34 ++++++ .../PublicAPI/net10.0/PublicAPI.Unshipped.txt | 9 ++ .../PublicAPI/net472/PublicAPI.Unshipped.txt | 9 ++ .../PublicAPI/net8.0/PublicAPI.Unshipped.txt | 9 ++ .../PublicAPI/net9.0/PublicAPI.Unshipped.txt | 9 ++ .../netstandard2.0/PublicAPI.Unshipped.txt | 9 ++ .../Harness/Loop/LoopAgentBudgetTests.cs | 100 ++++++++++++++++++ 9 files changed, 270 insertions(+), 1 deletion(-) create mode 100644 dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopExitReason.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs index 398dc061565..809ef8259d7 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System.Collections.Generic; +using System.Diagnostics; using System.Diagnostics.CodeAnalysis; using System.Linq; using System.Runtime.CompilerServices; @@ -69,6 +70,8 @@ public sealed class LoopAgent : DelegatingAIAgent private readonly IReadOnlyList _evaluators; private readonly int _maxIterations; + private readonly long? _maxTokens; + private readonly System.TimeSpan? _maxDuration; private readonly bool _freshContextPerIteration; private readonly string? _onBehalfOfAuthorName; private readonly bool _excludeOnBehalfOfMessages; @@ -121,6 +124,8 @@ public LoopAgent(AIAgent innerAgent, IEnumerable evaluators, Loop this._evaluators = evaluatorArray; this._maxIterations = Throw.IfLessThan(options?.MaxIterations ?? DefaultMaxIterations, 1); + this._maxTokens = options?.MaxTokens; + this._maxDuration = options?.MaxDuration; this._freshContextPerIteration = options?.FreshContextPerIteration ?? false; this._onBehalfOfAuthorName = options?.OnBehalfOfAuthorName; this._excludeOnBehalfOfMessages = options?.ExcludeOnBehalfOfMessages ?? false; @@ -159,6 +164,10 @@ protected override async Task RunCoreAsync( IEnumerable currentMessages = initialMessages; int iteration = 0; + // Budget tracking (tokens and wall-clock time) for the whole run. + long tokensUsed = 0; + Stopwatch? budgetStopwatch = this._maxDuration.HasValue ? Stopwatch.StartNew() : null; + // Aggregates the full transcript across iterations: each iteration's surfaced on-behalf-of input messages // followed by that iteration's response messages. Unused when only the final response is returned. List transcript = []; @@ -180,6 +189,12 @@ protected override async Task RunCoreAsync( UsageAggregator.Accumulate(ref aggregatedUsage, response.Usage); + if (this._maxTokens.HasValue && response.Usage is { } usage) + { + tokensUsed += usage.TotalTokenCount + ?? ((usage.InputTokenCount ?? 0) + (usage.OutputTokenCount ?? 0)); + } + // Record this iteration's on-behalf-of input (before the response it elicited) and the response itself. transcript.AddRange(currentSurfaced); transcript.AddRange(response.Messages); @@ -198,11 +213,28 @@ protected override async Task RunCoreAsync( return this.BuildResult(response, transcript, aggregatedUsage); } + // Budget pre-checks — run before the cap and evaluators so an expensive judge is not called once a budget is exhausted. + if (budgetStopwatch is not null && budgetStopwatch.Elapsed >= this._maxDuration) + { + AgentResponse budgetResult = this.BuildResult(response, transcript, aggregatedUsage); + (budgetResult.AdditionalProperties ??= new())[LoopExitReason.AdditionalPropertiesKey] = LoopExitReason.TimeBudgetExceeded; + return budgetResult; + } + + if (this._maxTokens.HasValue && tokensUsed >= this._maxTokens.Value) + { + AgentResponse budgetResult = this.BuildResult(response, transcript, aggregatedUsage); + (budgetResult.AdditionalProperties ??= new())[LoopExitReason.AdditionalPropertiesKey] = LoopExitReason.TokenBudgetExceeded; + return budgetResult; + } + // Enforce the global safety cap regardless of what the evaluators want. if (iteration >= this._maxIterations) { this.LogMaxIterationsReached(iteration); - return this.BuildResult(response, transcript, aggregatedUsage); + AgentResponse capResult = this.BuildResult(response, transcript, aggregatedUsage); + (capResult.AdditionalProperties ??= new())[LoopExitReason.AdditionalPropertiesKey] = LoopExitReason.IterationCapReached; + return capResult; } // Ask the evaluators whether to continue; stop when none of them request a re-invocation. @@ -247,6 +279,10 @@ protected override async IAsyncEnumerable RunCoreStreamingA IEnumerable currentMessages = initialMessages; int iteration = 0; + // Budget tracking (tokens and wall-clock time) for the whole run. + long tokensUsed = 0; + Stopwatch? budgetStopwatch = this._maxDuration.HasValue ? Stopwatch.StartNew() : null; + // The loop-synthesized on-behalf-of messages that drive the current iteration (none for the first iteration). IReadOnlyList currentSurfaced = []; @@ -307,6 +343,19 @@ protected override async IAsyncEnumerable RunCoreStreamingA yield break; } + // Budget pre-checks (no exit-reason stamping in streaming; consistent with the iteration cap). + if (this._maxTokens.HasValue && response.Usage is { } streamUsage) + { + tokensUsed += streamUsage.TotalTokenCount + ?? ((streamUsage.InputTokenCount ?? 0) + (streamUsage.OutputTokenCount ?? 0)); + } + + if ((budgetStopwatch is not null && budgetStopwatch.Elapsed >= this._maxDuration) + || (this._maxTokens.HasValue && tokensUsed >= this._maxTokens.Value)) + { + yield break; + } + // Enforce the global safety cap regardless of what the evaluators want. if (iteration >= this._maxIterations) { diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs index ec009b4594c..a00e1d35f6b 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs @@ -114,4 +114,45 @@ public sealed class LoopAgentOptions /// it fires multiple times, the most recent invocation carries the session the loop is currently using. /// public Func? SessionCreatedCallback { get; set; } + + /// Gets or sets the maximum total LLM token spend across all iterations, or for unconstrained. + /// + /// Counts TotalTokenCount from each iteration's AgentResponse.Usage, falling back to + /// InputTokenCount + OutputTokenCount when TotalTokenCount is absent. + /// When exceeded, the loop stamps AgentResponse.AdditionalProperties["loop_exit_reason"] = "token_budget_exceeded" + /// (non-streaming only) and returns the last response. + /// + public long? MaxTokens + { + get; + set + { + if (value.HasValue && value.Value < 1) + { + throw new ArgumentOutOfRangeException(nameof(value), "MaxTokens must be a positive integer."); + } + + field = value; + } + } + + /// Gets or sets the maximum wall-clock duration for the entire loop run, or for unconstrained. + /// + /// Checked after each iteration. Catches agents blocked on slow tools that token budgets cannot detect. + /// When exceeded, stamps AgentResponse.AdditionalProperties["loop_exit_reason"] = "time_budget_exceeded" + /// (non-streaming only). + /// + public TimeSpan? MaxDuration + { + get; + set + { + if (value.HasValue && value.Value <= TimeSpan.Zero) + { + throw new ArgumentOutOfRangeException(nameof(value), "MaxDuration must be a positive duration."); + } + + field = value; + } + } } diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopExitReason.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopExitReason.cs new file mode 100644 index 00000000000..12c0720f43d --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopExitReason.cs @@ -0,0 +1,34 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Diagnostics.CodeAnalysis; +using Microsoft.Shared.DiagnosticIds; + +namespace Microsoft.Agents.AI; + +/// +/// Provides typed exit-reason strings for runs. +/// Values are plain strings to remain JSON-transparent. The reason is surfaced under +/// in AgentResponse.AdditionalProperties. +/// +/// +/// +/// Exit reasons are stamped by on non-streaming runs only. +/// In streaming mode the response is assembled from yielded updates by the caller, so there is no +/// single AgentResponse object available inside the generator to stamp. +/// +/// +[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] +public static class LoopExitReason +{ + /// The key used in AgentResponse.AdditionalProperties to surface the loop exit reason. + public const string AdditionalPropertiesKey = "loop_exit_reason"; + + /// The cap was reached before the evaluators stopped the loop. + public const string IterationCapReached = "iteration_cap_reached"; + + /// The token budget () was exhausted before the loop completed. + public const string TokenBudgetExceeded = "token_budget_exceeded"; + + /// The wall-clock time budget () was exhausted before the loop completed. + public const string TimeBudgetExceeded = "time_budget_exceeded"; +} diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt index fada975c025..7676ecacf21 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -24,3 +24,12 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.LoopExitReason +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.AdditionalPropertiesKey = "loop_exit_reason" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TokenBudgetExceeded = "token_budget_exceeded" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TimeBudgetExceeded = "time_budget_exceeded" -> string! +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.get -> long? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.set -> void +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.get -> System.TimeSpan? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.set -> void +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.IterationCapReached = "iteration_cap_reached" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt index fada975c025..7676ecacf21 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -24,3 +24,12 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.LoopExitReason +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.AdditionalPropertiesKey = "loop_exit_reason" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TokenBudgetExceeded = "token_budget_exceeded" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TimeBudgetExceeded = "time_budget_exceeded" -> string! +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.get -> long? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.set -> void +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.get -> System.TimeSpan? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.set -> void +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.IterationCapReached = "iteration_cap_reached" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt index fada975c025..7676ecacf21 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -24,3 +24,12 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.LoopExitReason +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.AdditionalPropertiesKey = "loop_exit_reason" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TokenBudgetExceeded = "token_budget_exceeded" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TimeBudgetExceeded = "time_budget_exceeded" -> string! +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.get -> long? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.set -> void +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.get -> System.TimeSpan? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.set -> void +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.IterationCapReached = "iteration_cap_reached" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt index fada975c025..7676ecacf21 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -24,3 +24,12 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.LoopExitReason +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.AdditionalPropertiesKey = "loop_exit_reason" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TokenBudgetExceeded = "token_budget_exceeded" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TimeBudgetExceeded = "time_budget_exceeded" -> string! +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.get -> long? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.set -> void +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.get -> System.TimeSpan? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.set -> void +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.IterationCapReached = "iteration_cap_reached" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index fada975c025..7676ecacf21 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -24,3 +24,12 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.LoopExitReason +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.AdditionalPropertiesKey = "loop_exit_reason" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TokenBudgetExceeded = "token_budget_exceeded" -> string! +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.TimeBudgetExceeded = "time_budget_exceeded" -> string! +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.get -> long? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxTokens.set -> void +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.get -> System.TimeSpan? +[MAAI001]Microsoft.Agents.AI.LoopAgentOptions.MaxDuration.set -> void +[MAAI001]const Microsoft.Agents.AI.LoopExitReason.IterationCapReached = "iteration_cap_reached" -> string! diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs new file mode 100644 index 00000000000..655b4377986 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs @@ -0,0 +1,100 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.UnitTests; + +/// +/// Unit tests for the token and wall-clock budget support of . +/// +public class LoopAgentBudgetTests +{ + private static AgentResponse ResponseWithTokens(long totalTokens) => + new(new ChatMessage(ChatRole.Assistant, "working")) { Usage = new UsageDetails { TotalTokenCount = totalTokens } }; + + /// + /// Verify that the loop stops once cumulative token spend reaches MaxTokens and stamps the exit reason. + /// + [Fact] + public async Task RunAsync_TokenBudgetExceeded_StopsAndStampsReasonAsync() + { + // Arrange + var inner = new InnerAgentCapture(_ => ResponseWithTokens(30)); + var agent = new LoopAgent(inner.Agent, LoopTestHelpers.While(_ => true), new LoopAgentOptions { MaxTokens = 50, MaxIterations = 10 }); + + // Act + AgentResponse response = await agent.RunAsync("task", new ChatClientAgentSession()); + + // Assert: 30 < 50 after call 1, 60 >= 50 after call 2. + Assert.Equal(2, inner.CallCount); + Assert.NotNull(response.AdditionalProperties); + Assert.Equal(LoopExitReason.TokenBudgetExceeded, response.AdditionalProperties![LoopExitReason.AdditionalPropertiesKey]); + Assert.Equal("token_budget_exceeded", response.AdditionalProperties[LoopExitReason.AdditionalPropertiesKey]); + } + + /// + /// Verify that the loop stops after the first iteration when the duration budget is already exhausted. + /// + [Fact] + public async Task RunAsync_TimeBudgetExceeded_StopsAndStampsReasonAsync() + { + // Arrange + var inner = new InnerAgentCapture(_ => ResponseWithTokens(1)); + var agent = new LoopAgent(inner.Agent, LoopTestHelpers.While(_ => true), new LoopAgentOptions { MaxDuration = TimeSpan.FromTicks(1), MaxIterations = 10 }); + + // Act + AgentResponse response = await agent.RunAsync("task", new ChatClientAgentSession()); + + // Assert + Assert.Equal(1, inner.CallCount); + Assert.NotNull(response.AdditionalProperties); + Assert.Equal(LoopExitReason.TimeBudgetExceeded, response.AdditionalProperties![LoopExitReason.AdditionalPropertiesKey]); + } + + /// + /// Verify that without a budget the loop runs to the iteration cap. + /// + [Fact] + public async Task RunAsync_NoBudget_RunsToCapAsync() + { + // Arrange + var inner = new InnerAgentCapture(_ => ResponseWithTokens(1_000_000)); + var agent = new LoopAgent(inner.Agent, LoopTestHelpers.While(_ => true), new LoopAgentOptions { MaxIterations = 3 }); + + // Act + AgentResponse response = await agent.RunAsync("task", new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, inner.CallCount); + Assert.Equal(LoopExitReason.IterationCapReached, response.AdditionalProperties![LoopExitReason.AdditionalPropertiesKey]); + } + + /// + /// Verify that the budget stop fires before evaluators are consulted. + /// + [Fact] + public async Task RunAsync_BudgetCheckBeforeEvaluatorsAsync() + { + // Arrange + int evaluatorCalls = 0; + var inner = new InnerAgentCapture(_ => ResponseWithTokens(60)); + var agent = new LoopAgent( + inner.Agent, + LoopTestHelpers.While(_ => + { + evaluatorCalls++; + return true; + }), + new LoopAgentOptions { MaxTokens = 50, MaxIterations = 10 }); + + // Act + AgentResponse response = await agent.RunAsync("task", new ChatClientAgentSession()); + + // Assert + Assert.Equal(1, inner.CallCount); + Assert.Equal(0, evaluatorCalls); + Assert.Equal(LoopExitReason.TokenBudgetExceeded, response.AdditionalProperties![LoopExitReason.AdditionalPropertiesKey]); + } +} From f41c72e5948695f575d8df8db1e091d6749c5c2f Mon Sep 17 00:00:00 2001 From: Mark Licata Date: Wed, 7 Oct 2026 09:04:17 -0700 Subject: [PATCH 2/4] Address review feedback: document MaxDuration semantics, add streaming budget tests Co-Authored-By: Claude Opus 5.5 --- .../Harness/Loop/LoopAgent.cs | 14 ++ .../Harness/Loop/LoopAgentOptions.cs | 15 +- .../Harness/Loop/LoopAgentBudgetTests.cs | 154 ++++++++++++++++++ 3 files changed, 180 insertions(+), 3 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs index 809ef8259d7..5eabbfec798 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs @@ -244,6 +244,14 @@ protected override async Task RunCoreAsync( return this.BuildResult(response, transcript, aggregatedUsage); } + // Re-check the duration budget so a slow evaluator cannot exhaust it and still start another iteration. + if (budgetStopwatch is not null && budgetStopwatch.Elapsed >= this._maxDuration) + { + AgentResponse budgetResult = this.BuildResult(response, transcript, aggregatedUsage); + (budgetResult.AdditionalProperties ??= new())[LoopExitReason.AdditionalPropertiesKey] = LoopExitReason.TimeBudgetExceeded; + return budgetResult; + } + currentMessages = step.Messages; currentSurfaced = step.SurfacedMessages; } @@ -370,6 +378,12 @@ protected override async IAsyncEnumerable RunCoreStreamingA yield break; } + // Re-check the duration budget so a slow evaluator cannot exhaust it and still start another iteration. + if (budgetStopwatch is not null && budgetStopwatch.Elapsed >= this._maxDuration) + { + yield break; + } + currentMessages = step.Messages; currentSurfaced = step.SurfacedMessages; } diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs index a00e1d35f6b..5a5b205b584 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs @@ -138,9 +138,18 @@ public long? MaxTokens /// Gets or sets the maximum wall-clock duration for the entire loop run, or for unconstrained. /// - /// Checked after each iteration. Catches agents blocked on slow tools that token budgets cannot detect. - /// When exceeded, stamps AgentResponse.AdditionalProperties["loop_exit_reason"] = "time_budget_exceeded" - /// (non-streaming only). + /// + /// The duration is checked between iterations: after each inner agent call (or stream) completes, and again after + /// the evaluators run. Once it is exceeded, the loop does not start another iteration and returns the last response, + /// stamping AgentResponse.AdditionalProperties["loop_exit_reason"] = "time_budget_exceeded" (non-streaming only). + /// This bounds the number of slow iterations a run can start, which token budgets cannot detect. + /// + /// + /// This is not a hard deadline: it does not interrupt an in-flight inner agent call, stream, or evaluator, so a run can + /// exceed the duration by as long as that work takes. Callers that need a hard deadline should pass a + /// that cancels after a timeout (for example via + /// ). + /// /// public TimeSpan? MaxDuration { diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs index 655b4377986..d1420a76b4c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Collections.Generic; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -14,6 +15,23 @@ public class LoopAgentBudgetTests private static AgentResponse ResponseWithTokens(long totalTokens) => new(new ChatMessage(ChatRole.Assistant, "working")) { Usage = new UsageDetails { TotalTokenCount = totalTokens } }; + private static AgentResponseUpdate[] UpdatesWithUsage(int call, UsageDetails usage) => + [ + new AgentResponseUpdate(ChatRole.Assistant, $"chunk {call}"), + new AgentResponseUpdate(ChatRole.Assistant, [new UsageContent(usage)]), + ]; + + private static async Task> CollectAsync(IAsyncEnumerable stream) + { + List updates = []; + await foreach (AgentResponseUpdate update in stream) + { + updates.Add(update); + } + + return updates; + } + /// /// Verify that the loop stops once cumulative token spend reaches MaxTokens and stamps the exit reason. /// @@ -97,4 +115,140 @@ public async Task RunAsync_BudgetCheckBeforeEvaluatorsAsync() Assert.Equal(0, evaluatorCalls); Assert.Equal(LoopExitReason.TokenBudgetExceeded, response.AdditionalProperties![LoopExitReason.AdditionalPropertiesKey]); } + + /// + /// Verify that a slow evaluator which exhausts the duration budget does not start another iteration. + /// + [Fact] + public async Task RunAsync_TimeBudgetExceededDuringEvaluation_DoesNotStartAnotherIterationAsync() + { + // Arrange + int evaluatorCalls = 0; + var inner = new InnerAgentCapture(_ => ResponseWithTokens(1)); + var evaluator = new DelegateLoopEvaluator(async (_, ct) => + { + evaluatorCalls++; + await Task.Delay(TimeSpan.FromMilliseconds(1200), ct); + return LoopEvaluation.Continue(); + }); + var agent = new LoopAgent(inner.Agent, evaluator, new LoopAgentOptions { MaxDuration = TimeSpan.FromSeconds(1), MaxIterations = 10 }); + + // Act + AgentResponse response = await agent.RunAsync("task", new ChatClientAgentSession()); + + // Assert: the evaluator ran once and asked to continue, but the budget expired while it ran. + Assert.Equal(1, evaluatorCalls); + Assert.Equal(1, inner.CallCount); + Assert.Equal(LoopExitReason.TimeBudgetExceeded, response.AdditionalProperties![LoopExitReason.AdditionalPropertiesKey]); + } + + /// + /// Verify that the streaming loop stops once cumulative token spend across iterations reaches MaxTokens, without + /// consulting the evaluators or starting another iteration after the budget is exhausted. + /// + [Fact] + public async Task RunStreamingAsync_TokenBudgetReachedCumulatively_StopsWithoutEvaluatingAsync() + { + // Arrange + int evaluatorCalls = 0; + var inner = new InnerStreamingCapture(call => UpdatesWithUsage(call, new UsageDetails { TotalTokenCount = 30 })); + var agent = new LoopAgent( + inner.Agent, + LoopTestHelpers.While(_ => + { + evaluatorCalls++; + return true; + }), + new LoopAgentOptions { MaxTokens = 50, MaxIterations = 10 }); + + // Act + List updates = await CollectAsync(agent.RunStreamingAsync("task", new ChatClientAgentSession())); + + // Assert: 30 < 50 after call 1 (evaluator consulted), 60 >= 50 after call 2 (evaluator skipped, loop stops). + Assert.Equal(2, inner.CallCount); + Assert.Equal(1, evaluatorCalls); + Assert.Contains(updates, u => u.Text == "chunk 1"); + Assert.Contains(updates, u => u.Text == "chunk 2"); + } + + /// + /// Verify that the streaming loop counts InputTokenCount + OutputTokenCount when TotalTokenCount is absent. + /// + [Fact] + public async Task RunStreamingAsync_TotalTokenCountAbsent_FallsBackToInputPlusOutputAsync() + { + // Arrange + int evaluatorCalls = 0; + var inner = new InnerStreamingCapture(call => UpdatesWithUsage(call, new UsageDetails { InputTokenCount = 20, OutputTokenCount = 15 })); + var agent = new LoopAgent( + inner.Agent, + LoopTestHelpers.While(_ => + { + evaluatorCalls++; + return true; + }), + new LoopAgentOptions { MaxTokens = 70, MaxIterations = 10 }); + + // Act + List updates = await CollectAsync(agent.RunStreamingAsync("task", new ChatClientAgentSession())); + + // Assert: 35 < 70 after call 1, 70 >= 70 after call 2. Without the fallback the loop would run to the cap. + Assert.Equal(2, inner.CallCount); + Assert.Equal(1, evaluatorCalls); + Assert.DoesNotContain(updates, u => u.Text == "chunk 3"); + } + + /// + /// Verify that the streaming loop stops after the first iteration, without consulting evaluators, when that + /// iteration alone exhausts the token budget. + /// + [Fact] + public async Task RunStreamingAsync_TokenBudgetExhaustedByFirstIteration_SkipsEvaluatorsAsync() + { + // Arrange + int evaluatorCalls = 0; + var inner = new InnerStreamingCapture(call => UpdatesWithUsage(call, new UsageDetails { TotalTokenCount = 60 })); + var agent = new LoopAgent( + inner.Agent, + LoopTestHelpers.While(_ => + { + evaluatorCalls++; + return true; + }), + new LoopAgentOptions { MaxTokens = 50, MaxIterations = 10 }); + + // Act + List updates = await CollectAsync(agent.RunStreamingAsync("task", new ChatClientAgentSession())); + + // Assert + Assert.Equal(1, inner.CallCount); + Assert.Equal(0, evaluatorCalls); + Assert.Equal(2, updates.Count); + } + + /// + /// Verify that in streaming mode a slow evaluator which exhausts the duration budget does not start another iteration. + /// + [Fact] + public async Task RunStreamingAsync_TimeBudgetExceededDuringEvaluation_DoesNotStartAnotherIterationAsync() + { + // Arrange + int evaluatorCalls = 0; + var inner = new InnerStreamingCapture(call => UpdatesWithUsage(call, new UsageDetails { TotalTokenCount = 1 })); + var evaluator = new DelegateLoopEvaluator(async (_, ct) => + { + evaluatorCalls++; + await Task.Delay(TimeSpan.FromMilliseconds(1200), ct); + return LoopEvaluation.Continue(); + }); + var agent = new LoopAgent(inner.Agent, evaluator, new LoopAgentOptions { MaxDuration = TimeSpan.FromSeconds(1), MaxIterations = 10 }); + + // Act + List updates = await CollectAsync(agent.RunStreamingAsync("task", new ChatClientAgentSession())); + + // Assert + Assert.Equal(1, evaluatorCalls); + Assert.Equal(1, inner.CallCount); + Assert.DoesNotContain(updates, u => u.Text == "chunk 2"); + } } From ce755bf5cbd74a77306324cf94180f12615b864b Mon Sep 17 00:00:00 2001 From: Mark Licata Date: Tue, 13 Oct 2026 08:42:33 -0700 Subject: [PATCH 3/4] Address review feedback: Throw helper, warn when token usage is missing - Use Throw.IfLessThanOrEqual for MaxTokens validation. - Log a warning (once per run) when MaxTokens is set but the inner agent's response reports no token usage, in both RunAsync and RunStreamingAsync. - Document what MaxTokens counts: inner agent input+output tokens (TotalTokenCount, falling back to Input+Output); evaluator and judge usage is not counted. - Add tests for the warning, option validation, and judge usage exclusion. Co-Authored-By: Claude Opus 5.5 --- .../Harness/Loop/LoopAgent.cs | 44 +++++- .../Harness/Loop/LoopAgentOptions.cs | 17 +- .../Harness/Loop/LoopAgentBudgetTests.cs | 149 ++++++++++++++++++ 3 files changed, 200 insertions(+), 10 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs index 5eabbfec798..08395205826 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgent.cs @@ -166,6 +166,7 @@ protected override async Task RunCoreAsync( // Budget tracking (tokens and wall-clock time) for the whole run. long tokensUsed = 0; + bool tokenUsageMissingLogged = false; Stopwatch? budgetStopwatch = this._maxDuration.HasValue ? Stopwatch.StartNew() : null; // Aggregates the full transcript across iterations: each iteration's surfaced on-behalf-of input messages @@ -189,10 +190,9 @@ protected override async Task RunCoreAsync( UsageAggregator.Accumulate(ref aggregatedUsage, response.Usage); - if (this._maxTokens.HasValue && response.Usage is { } usage) + if (this._maxTokens.HasValue) { - tokensUsed += usage.TotalTokenCount - ?? ((usage.InputTokenCount ?? 0) + (usage.OutputTokenCount ?? 0)); + tokensUsed += this.CountTokens(response.Usage, iteration, ref tokenUsageMissingLogged); } // Record this iteration's on-behalf-of input (before the response it elicited) and the response itself. @@ -289,6 +289,7 @@ protected override async IAsyncEnumerable RunCoreStreamingA // Budget tracking (tokens and wall-clock time) for the whole run. long tokensUsed = 0; + bool tokenUsageMissingLogged = false; Stopwatch? budgetStopwatch = this._maxDuration.HasValue ? Stopwatch.StartNew() : null; // The loop-synthesized on-behalf-of messages that drive the current iteration (none for the first iteration). @@ -352,10 +353,9 @@ protected override async IAsyncEnumerable RunCoreStreamingA } // Budget pre-checks (no exit-reason stamping in streaming; consistent with the iteration cap). - if (this._maxTokens.HasValue && response.Usage is { } streamUsage) + if (this._maxTokens.HasValue) { - tokensUsed += streamUsage.TotalTokenCount - ?? ((streamUsage.InputTokenCount ?? 0) + (streamUsage.OutputTokenCount ?? 0)); + tokensUsed += this.CountTokens(response.Usage, iteration, ref tokenUsageMissingLogged); } if ((budgetStopwatch is not null && budgetStopwatch.Elapsed >= this._maxDuration) @@ -540,6 +540,38 @@ private static bool HasPendingApprovalRequests(AgentResponse response) return false; } + /// + /// Returns the number of tokens reported by : , falling + /// back to plus . When the response + /// carries no token counts, returns zero and logs a warning (once per run) because the token budget cannot be enforced. + /// + private long CountTokens(UsageDetails? usage, int iteration, ref bool tokenUsageMissingLogged) + { + if (usage is not null && (usage.TotalTokenCount.HasValue || usage.InputTokenCount.HasValue || usage.OutputTokenCount.HasValue)) + { + return usage.TotalTokenCount ?? ((usage.InputTokenCount ?? 0) + (usage.OutputTokenCount ?? 0)); + } + + if (!tokenUsageMissingLogged) + { + tokenUsageMissingLogged = true; + this.LogTokenUsageMissing(iteration); + } + + return 0; + } + + private void LogTokenUsageMissing(int iteration) + { + if (this._logger.IsEnabled(LogLevel.Warning)) + { + this._logger.LogWarning( + "LoopAgent has MaxTokens set to {MaxTokens}, but the inner agent's response for iteration {Iteration} reported no token usage, so the token budget cannot be enforced for that response.", + this._maxTokens, + iteration); + } + } + private void LogMaxIterationsReached(int iteration) { if (this._logger.IsEnabled(LogLevel.Information)) diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs index 5a5b205b584..7461abe73ce 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs @@ -5,6 +5,7 @@ using System.Threading; using System.Threading.Tasks; using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; @@ -117,19 +118,27 @@ public sealed class LoopAgentOptions /// Gets or sets the maximum total LLM token spend across all iterations, or for unconstrained. /// - /// Counts TotalTokenCount from each iteration's AgentResponse.Usage, falling back to - /// InputTokenCount + OutputTokenCount when TotalTokenCount is absent. + /// + /// Counts the input and output tokens used by the inner agent, including any sub-agents whose usage the inner agent + /// reports in its response. For each iteration, the loop reads TotalTokenCount from the inner agent's + /// AgentResponse.Usage, falling back to InputTokenCount + OutputTokenCount when TotalTokenCount + /// is absent. Tokens spent by instances, including judge model calls made by + /// , are not counted. What is counted cannot be configured. + /// + /// + /// If the inner agent reports no token usage, the budget cannot be enforced, and the loop logs a warning (once per run). /// When exceeded, the loop stamps AgentResponse.AdditionalProperties["loop_exit_reason"] = "token_budget_exceeded" /// (non-streaming only) and returns the last response. + /// /// public long? MaxTokens { get; set { - if (value.HasValue && value.Value < 1) + if (value.HasValue) { - throw new ArgumentOutOfRangeException(nameof(value), "MaxTokens must be a positive integer."); + _ = Throw.IfLessThanOrEqual(value.Value, 0, nameof(value)); } field = value; diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs index d1420a76b4c..5e8a50df13a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/Loop/LoopAgentBudgetTests.cs @@ -2,8 +2,11 @@ using System; using System.Collections.Generic; +using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging; +using Moq; namespace Microsoft.Agents.AI.UnitTests; @@ -21,6 +24,25 @@ private static AgentResponseUpdate[] UpdatesWithUsage(int call, UsageDetails usa new AgentResponseUpdate(ChatRole.Assistant, [new UsageContent(usage)]), ]; + private static (Mock Logger, ILoggerFactory Factory) CreateLogger() + { + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + return (logger, loggerFactory.Object); + } + + private static void VerifyWarningLogged(Mock logger, Times times) => + logger.Verify( + l => l.Log( + LogLevel.Warning, + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny>()), + times); + private static async Task> CollectAsync(IAsyncEnumerable stream) { List updates = []; @@ -251,4 +273,131 @@ public async Task RunStreamingAsync_TimeBudgetExceededDuringEvaluation_DoesNotSt Assert.Equal(1, inner.CallCount); Assert.DoesNotContain(updates, u => u.Text == "chunk 2"); } + + /// + /// Verify that non-positive MaxTokens values are rejected. + /// + [Theory] + [InlineData(0L)] + [InlineData(-1L)] + public void MaxTokens_NonPositive_Throws(long value) + { + var options = new LoopAgentOptions(); + ArgumentOutOfRangeException ex = Assert.Throws(() => options.MaxTokens = value); + Assert.Equal("value", ex.ParamName); + } + + /// + /// Verify that non-positive MaxDuration values are rejected. + /// + [Theory] + [InlineData(0L)] + [InlineData(-1L)] + public void MaxDuration_NonPositive_Throws(long ticks) + { + var options = new LoopAgentOptions(); + ArgumentOutOfRangeException ex = Assert.Throws(() => options.MaxDuration = TimeSpan.FromTicks(ticks)); + Assert.Equal("value", ex.ParamName); + } + + /// + /// Verify that a warning is logged once per run when MaxTokens is set but the inner agent reports no token usage. + /// + [Fact] + public async Task RunAsync_MaxTokensSetWithoutUsage_LogsWarningOnceAsync() + { + // Arrange + var (logger, loggerFactory) = CreateLogger(); + var inner = new InnerAgentCapture(_ => new AgentResponse(new ChatMessage(ChatRole.Assistant, "working"))); + var agent = new LoopAgent(inner.Agent, LoopTestHelpers.While(_ => true), new LoopAgentOptions { MaxTokens = 50, MaxIterations = 3 }, loggerFactory); + + // Act + await agent.RunAsync("task", new ChatClientAgentSession()); + + // Assert + Assert.Equal(3, inner.CallCount); + VerifyWarningLogged(logger, Times.Once()); + } + + /// + /// Verify that no warning is logged when MaxTokens is set and the inner agent reports token usage. + /// + [Fact] + public async Task RunAsync_MaxTokensSetWithUsage_DoesNotLogWarningAsync() + { + // Arrange + var (logger, loggerFactory) = CreateLogger(); + var inner = new InnerAgentCapture(_ => ResponseWithTokens(10)); + var agent = new LoopAgent(inner.Agent, LoopTestHelpers.While(_ => true), new LoopAgentOptions { MaxTokens = 50, MaxIterations = 3 }, loggerFactory); + + // Act + await agent.RunAsync("task", new ChatClientAgentSession()); + + // Assert + VerifyWarningLogged(logger, Times.Never()); + } + + /// + /// Verify that the streaming loop logs a warning once per run when MaxTokens is set but no token usage is reported. + /// + [Fact] + public async Task RunStreamingAsync_MaxTokensSetWithoutUsage_LogsWarningOnceAsync() + { + // Arrange + var (logger, loggerFactory) = CreateLogger(); + var inner = new InnerStreamingCapture(call => [new AgentResponseUpdate(ChatRole.Assistant, $"chunk {call}")]); + var agent = new LoopAgent(inner.Agent, LoopTestHelpers.While(_ => true), new LoopAgentOptions { MaxTokens = 50, MaxIterations = 3 }, loggerFactory); + + // Act + await CollectAsync(agent.RunStreamingAsync("task", new ChatClientAgentSession())); + + // Assert + Assert.Equal(3, inner.CallCount); + VerifyWarningLogged(logger, Times.Once()); + } + + /// + /// Verify that the streaming loop does not log a warning when MaxTokens is set and token usage is reported. + /// + [Fact] + public async Task RunStreamingAsync_MaxTokensSetWithUsage_DoesNotLogWarningAsync() + { + // Arrange + var (logger, loggerFactory) = CreateLogger(); + var inner = new InnerStreamingCapture(call => UpdatesWithUsage(call, new UsageDetails { TotalTokenCount = 10 })); + var agent = new LoopAgent(inner.Agent, LoopTestHelpers.While(_ => true), new LoopAgentOptions { MaxTokens = 50, MaxIterations = 3 }, loggerFactory); + + // Act + await CollectAsync(agent.RunStreamingAsync("task", new ChatClientAgentSession())); + + // Assert + VerifyWarningLogged(logger, Times.Never()); + } + + /// + /// Verify that tokens spent by the judge evaluator are not charged against MaxTokens; only the inner agent's usage counts. + /// + [Fact] + public async Task RunAsync_JudgeEvaluatorUsage_NotCountedTowardsTokenBudgetAsync() + { + // Arrange: the judge always asks for more and reports a large usage that alone would exhaust the budget. + var judgeMock = new Mock(); + judgeMock.Setup(c => c.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(() => new ChatResponse(new ChatMessage(ChatRole.Assistant, "{\"answered\":false,\"gapAnalysis\":\"more\"}")) + { + Usage = new UsageDetails { TotalTokenCount = 1000 }, + }); + var inner = new InnerAgentCapture(_ => ResponseWithTokens(10)); + var agent = new LoopAgent(inner.Agent, new AIJudgeLoopEvaluator(judgeMock.Object), new LoopAgentOptions { MaxTokens = 50, MaxIterations = 10 }); + + // Act + AgentResponse response = await agent.RunAsync("task", new ChatClientAgentSession()); + + // Assert: 10 tokens per iteration reaches 50 after five iterations; judge usage would have stopped it after one. + Assert.Equal(5, inner.CallCount); + Assert.Equal(LoopExitReason.TokenBudgetExceeded, response.AdditionalProperties![LoopExitReason.AdditionalPropertiesKey]); + } } From 13fad81c519a878a09677d2773292afe40d7e4c3 Mon Sep 17 00:00:00 2001 From: Mark Licata Date: Tue, 13 Oct 2026 09:39:37 -0700 Subject: [PATCH 4/4] .NET: roll up sub-agent token usage into the calling agent's response ChatClientAgent now adds the usage of the sub-agents it invokes to its own AgentResponse.Usage, in both RunAsync and RunStreamingAsync, so that LoopAgentOptions.MaxTokens also counts sub-agent tokens. - An internal AsyncLocal SubAgentUsageScope (modelled on ServedModelScope) holds a lock-based usage accumulator per ChatClientAgent run. It is set up before context providers run and re-established after each streaming yield. - AsAIFunction adds the sub-agent's reported usage to the caller's scope and clears the scope for the sub-agent's own run, so A -> B -> C counts each layer once. Concurrent tool invocations are summed safely. - FunctionInvokingChatClient overwrites response usage with its own model call totals, so the merge happens after it returns. Streaming runs end with a usage-only update carrying the sub-agent total. - BackgroundAgentsProvider charges each successful background run once, to the parent run that first observes its completion (context provision or the wait/result tools). Runs that finish after the last parent run, or are only finalized at release, are not counted. - OpenTelemetryAgent keeps reporting each agent's own usage on its invoke_agent span; the rolled-up sub-agent usage is excluded. - Failed or cancelled sub-agents and agents invoked directly from custom tools are not counted. No public API changes. Co-Authored-By: Claude Opus 5.5 --- .../Microsoft.Agents.AI/AgentExtensions.cs | 17 + .../ChatClient/ChatClientAgent.cs | 40 +- .../BackgroundAgentRuntimeState.cs | 19 + .../BackgroundAgentsProvider.cs | 89 +++- .../Harness/Loop/LoopAgentOptions.cs | 18 +- .../Microsoft.Agents.AI/OpenTelemetryAgent.cs | 87 +++- .../Microsoft.Agents.AI/SubAgentUsageScope.cs | 141 +++++ .../ChatClientAgent_SubAgentUsageTests.cs | 481 ++++++++++++++++++ 8 files changed, 881 insertions(+), 11 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI/SubAgentUsageScope.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_SubAgentUsageTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI/AgentExtensions.cs b/dotnet/src/Microsoft.Agents.AI/AgentExtensions.cs index 260756e3055..b8fc0295668 100644 --- a/dotnet/src/Microsoft.Agents.AI/AgentExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI/AgentExtensions.cs @@ -65,6 +65,12 @@ public static AIAgentBuilder AsBuilder(this AIAgent innerAgent) /// in multiple conversations or in requests where the parallel function calls may result in concurrent usage of the session, /// as that could lead to undefined and unpredictable behavior. /// + /// + /// When the resulting function is invoked during the run of a , the + /// reported by is added to the usage reported by that run, + /// so that the calling agent's usage includes the tokens consumed by the agents it invokes. If + /// fails or is cancelled, it produces no response and its usage is not counted. + /// /// public static AIFunction AsAIFunction(this AIAgent agent, AIFunctionFactoryOptions? options = null, AgentSession? session = null) { @@ -80,7 +86,18 @@ async Task InvokeAgentAsync( ? new AgentRunOptions { AdditionalProperties = dict } : null; + // Capture the calling agent's run, if any, so the sub-agent's usage can be rolled up into it. The scope is + // then cleared for the sub-agent's own run, so that only the usage the sub-agent reports in its response is + // counted, exactly once, at this call boundary. A sub-agent that is a ChatClientAgent establishes its own + // scope, so the usage it reports already includes that of its own sub-agents. + SubAgentUsageScope.UsageAccumulator? parentUsage = SubAgentUsageScope.Current; + SubAgentUsageScope.Current = null; + var response = await agent.RunAsync(query, session: session, options: agentRunOptions, cancellationToken: cancellationToken).ConfigureAwait(false); + + // Failed or cancelled sub-agent runs produce no response, so their usage is not counted. + parentUsage?.Add(response.Usage); + return response.Text; } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index 5988737629e..8ea5737f69f 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -230,6 +230,11 @@ protected override async Task RunCoreAsync( FeatureUsage.MarkUsed((int)FeatureIndex.CoreAgent); #pragma warning restore MAAI001 + // Collect the usage reported by sub-agents invoked during this run (including by context providers), + // so that it can be rolled up into this run's usage. + var subAgentUsage = new SubAgentUsageScope.UsageAccumulator(); + SubAgentUsageScope.Current = subAgentUsage; + (ChatClientAgentSession safeSession, ChatOptions? chatOptions, List inputMessagesForChatClient, @@ -276,6 +281,13 @@ protected override async Task RunCoreAsync( // is unreliable (the caller may stop consuming the stream before the decorator can persist). await this.NotifyProvidersOfNewMessagesAtEndOfRunAsync(safeSession, inputMessagesForChatClient, chatResponse.Messages, chatOptions, cancellationToken, forceNotify: forceEndOfRunPersistence).ConfigureAwait(false); + // Roll up the usage of any sub-agents invoked during this run. The function-invoking client only reports the + // usage of the model calls it made itself, so this has to happen after it returns. + if (subAgentUsage.GetTotal() is { } subAgentTotal) + { + chatResponse.Usage = SubAgentUsageScope.Merge(chatResponse.Usage, subAgentTotal); + } + return new AgentResponse(chatResponse) { AgentId = this.Id, @@ -320,6 +332,12 @@ protected override async IAsyncEnumerable RunCoreStreamingA FeatureUsage.MarkUsed((int)FeatureIndex.CoreAgent); #pragma warning restore MAAI001 + // Collect the usage reported by sub-agents invoked during this run (including by context providers), + // so that it can be rolled up into this run's usage. The scope is re-established after each yield, since + // the code resumes in the caller's execution context. + var subAgentUsage = new SubAgentUsageScope.UsageAccumulator(); + SubAgentUsageScope.Current = subAgentUsage; + (ChatClientAgentSession safeSession, ChatOptions? chatOptions, List inputMessagesForChatClient, @@ -361,6 +379,7 @@ protected override async IAsyncEnumerable RunCoreStreamingA // left suspended at `yield return`, never running their finally blocks, and any in-flight // FunctionResultContent / FunctionCallContent state would not be persisted before the next // turn, leaving the next request to the model with dangling tool calls. + ChatClientAgentContinuationToken? lastContinuationToken = null; try { bool hasUpdates; @@ -384,10 +403,11 @@ protected override async IAsyncEnumerable RunCoreStreamingA responseUpdates.Add(update); + lastContinuationToken = WrapContinuationToken(update.ContinuationToken, GetInputMessages(inputMessages, continuationToken), responseUpdates); yield return new(update) { AgentId = this.Id, - ContinuationToken = WrapContinuationToken(update.ContinuationToken, GetInputMessages(inputMessages, continuationToken), responseUpdates) + ContinuationToken = lastContinuationToken }; } @@ -397,6 +417,9 @@ protected override async IAsyncEnumerable RunCoreStreamingA // The base class RunStreamingAsync restores the original context (potentially with // null session) after each yield, so we must re-establish it for the decorator. EnsureRunContextHasSession(safeSession); + + // Likewise, re-establish the sub-agent usage scope for any tools invoked while resuming. + SubAgentUsageScope.Current = subAgentUsage; hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); } catch (Exception ex) @@ -423,6 +446,21 @@ protected override async IAsyncEnumerable RunCoreStreamingA { await responseUpdatesEnumerator.DisposeAsync().ConfigureAwait(false); } + + // Roll up the usage of any sub-agents invoked during this run as a final update carrying only usage, so that + // aggregating the stream (e.g. with ToAgentResponse) includes it. The function-invoking client only reports + // the usage of the model calls it made itself. + if (subAgentUsage.GetTotal() is { } subAgentTotal) + { + // No role, author name or message id is set, so that aggregation attaches it to the last message rather + // than starting a new one. The continuation token of the last update is preserved so that the final + // update still reflects whether the response can be resumed. + yield return new AgentResponseUpdate(role: null, contents: [SubAgentUsageScope.CreateSubAgentUsageContent(subAgentTotal)]) + { + AgentId = this.Id, + ContinuationToken = lastContinuationToken, + }; + } } /// diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentRuntimeState.cs b/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentRuntimeState.cs index 8a5dae51710..a4d1c001f8b 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentRuntimeState.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentRuntimeState.cs @@ -4,6 +4,7 @@ using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; +using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI; @@ -54,6 +55,24 @@ internal sealed class BackgroundAgentRuntimeState [JsonIgnore] public Dictionary TaskCancellations { get; } = []; + /// + /// Gets the in-flight runs that finished successfully and whose usage has already been charged to a parent run, + /// but that have not been finalized yet. + /// + /// + /// Each run's usage is charged once, to whichever parent run first observes that it finished. A run is removed + /// from this set when it is finalized, so that it is never charged again. + /// + [JsonIgnore] + public HashSet> UsageChargedRuns { get; } = []; + + /// + /// Gets the usage of runs that were finalized while no parent run was in progress (for example through + /// ), to be charged to the next parent run. + /// + [JsonIgnore] + public List UnchargedUsage { get; } = []; + /// /// Gets or sets a value indicating whether this runtime has been released via /// . diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProvider.cs b/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProvider.cs index 41ab04014db..493c318ea50 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProvider.cs @@ -40,6 +40,13 @@ namespace Microsoft.Agents.AI; /// abandoned work does not continue to invoke models and tools in the background. /// /// +/// The of each successful background run is added once to the usage of the parent +/// agent run that first observes that it finished: either when the parent run starts, or when one of the tools +/// above finalizes the task. Usage of runs that finish after the last parent run, or that are only finalized while +/// the session is released, is not added to any parent run, and nor is the usage of failed or cancelled runs, which +/// produce no response. +/// +/// /// Security considerations: The agents passed to the constructor are delegated /// arbitrary work by the parent agent — the parent sends them text input (which may include content /// derived from the parent's own untrusted context) and receives back whatever text they produce. A @@ -148,6 +155,11 @@ protected override ValueTask ProvideAIContextAsync(InvokingContext co BackgroundAgentState state = this._sessionState.GetOrInitializeState(context.Session); BackgroundAgentRuntimeState runtimeState = this._runtimeSessionState.GetOrInitializeState(context.Session); + lock (runtimeState.SyncRoot) + { + ChargeUsageOfFinishedRuns(runtimeState); + } + return new ValueTask(new AIContext { Instructions = this._instructions, @@ -327,7 +339,8 @@ public async Task ReleaseSessionAsync( continue; } - FinalizeTask(tracked, kvp.Value, runtimeState); + // Usage of runs that are only finalized during release is not charged to any parent run. + FinalizeTask(tracked, kvp.Value, runtimeState, chargeUsage: false); if (kvp.Value.IsCanceled && pendingTaskIds.Contains(kvp.Key)) { @@ -343,6 +356,8 @@ public async Task ReleaseSessionAsync( runtimeState.InFlightTasks.Clear(); runtimeState.BackgroundTaskSessions.Clear(); + runtimeState.UsageChargedRuns.Clear(); + runtimeState.UnchargedUsage.Clear(); // Anything still running was abandoned (for example after the timeout elapsed). foreach (BackgroundTaskInfo task in state.Tasks) @@ -517,14 +532,74 @@ private void TryRefreshTaskState(BackgroundAgentState state, BackgroundAgentRunt } } + /// + /// Charges the usage of background runs that have finished successfully, but have not been charged yet, to the + /// parent agent run in progress, if any. + /// + /// + /// + /// Each background run's usage is charged once, to whichever parent run first observes that it finished, either + /// here (when the parent run starts) or when the run is finalized by one of the tools. Runs that finish after the + /// last parent run, or that are only finalized while the session is released, are not charged to any parent run. + /// + /// Callers must hold . + /// + private static void ChargeUsageOfFinishedRuns(BackgroundAgentRuntimeState runtimeState) + { + if (SubAgentUsageScope.Current is not { } parentUsage) + { + return; + } + + foreach (UsageDetails usage in runtimeState.UnchargedUsage) + { + parentUsage.Add(usage); + } + + runtimeState.UnchargedUsage.Clear(); + + foreach (Task run in runtimeState.InFlightTasks.Values) + { + if (run.Status == TaskStatus.RanToCompletion && runtimeState.UsageChargedRuns.Add(run)) + { +#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits — task is already completed + parentUsage.Add(run.Result.Usage); +#pragma warning restore VSTHRD002 + } + } + } + /// /// Finalizes a task by extracting results from the completed Task and updating the BackgroundTaskInfo. /// + /// The task to finalize. + /// The completed run of the task. + /// The runtime state of the session. + /// + /// Whether the usage of a successful run that has not been charged yet should be charged to the parent agent run + /// in progress, or, when there is none, to the next parent run that starts. + /// /// Callers must hold . - private static void FinalizeTask(BackgroundTaskInfo taskInfo, Task completedTask, BackgroundAgentRuntimeState runtimeState) + private static void FinalizeTask(BackgroundTaskInfo taskInfo, Task completedTask, BackgroundAgentRuntimeState runtimeState, bool chargeUsage = true) { + bool usageAlreadyCharged = runtimeState.UsageChargedRuns.Remove(completedTask); + if (completedTask.Status == TaskStatus.RanToCompletion) { +#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits — task is already completed + if (chargeUsage && !usageAlreadyCharged && completedTask.Result.Usage is { } usage) +#pragma warning restore VSTHRD002 + { + if (SubAgentUsageScope.Current is { } parentUsage) + { + parentUsage.Add(usage); + } + else + { + runtimeState.UnchargedUsage.Add(usage); + } + } + taskInfo.Status = BackgroundTaskStatus.Completed; #pragma warning disable VSTHRD002 // Avoid problematic synchronous waits — task is already completed taskInfo.ResultText = completedTask.Result.Text; @@ -598,7 +673,15 @@ private static bool StartTrackedRun(BackgroundAgentRuntimeState runtimeState, in // this isolation, the background agent's RunAsync would overwrite the outer (calling) // agent's CurrentRunContext, corrupting all subsequent tool invocations in the // same FICC batch. - runtimeState.InFlightTasks[taskId] = Task.Run(() => agent.RunAsync(input, subSession, cancellationToken: cts.Token), cts.Token); + // The background run's usage is charged to a parent run when it is observed to have finished (see + // ChargeUsageOfFinishedRuns), so the calling run's sub-agent usage scope is cleared for it. + runtimeState.InFlightTasks[taskId] = Task.Run( + () => + { + SubAgentUsageScope.Current = null; + return agent.RunAsync(input, subSession, cancellationToken: cts.Token); + }, + cts.Token); return true; } } diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs index 7461abe73ce..0c182384b1c 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/Loop/LoopAgentOptions.cs @@ -119,11 +119,19 @@ public sealed class LoopAgentOptions /// Gets or sets the maximum total LLM token spend across all iterations, or for unconstrained. /// /// - /// Counts the input and output tokens used by the inner agent, including any sub-agents whose usage the inner agent - /// reports in its response. For each iteration, the loop reads TotalTokenCount from the inner agent's - /// AgentResponse.Usage, falling back to InputTokenCount + OutputTokenCount when TotalTokenCount - /// is absent. Tokens spent by instances, including judge model calls made by - /// , are not counted. What is counted cannot be configured. + /// Counts the input and output tokens used by the inner agent. For each iteration, the loop reads + /// TotalTokenCount from the inner agent's AgentResponse.Usage, falling back to + /// InputTokenCount + OutputTokenCount when TotalTokenCount is absent. What is counted cannot be configured. + /// + /// + /// A includes in its AgentResponse.Usage the usage of the sub-agents it invokes + /// through tools and through , so + /// those tokens count toward this budget. A background task's usage is counted once, in the run that first observes + /// that the task finished; tasks that finish after the last run, or are only finalized when the session is released, + /// are not counted. Also not counted: sub-agents that fail or are cancelled (they produce no response), agents that + /// custom tools invoke directly through RunAsync, sub-agent usage that another kind of inner agent does not + /// report in its response, and tokens spent by instances, including judge model calls made + /// by . /// /// /// If the inner agent reports no token usage, the budget cannot be enforced, and the loop logs a warning (once per run). diff --git a/dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs b/dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs index bb94f4674a3..072ceda0c3a 100644 --- a/dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs @@ -311,7 +311,7 @@ public async Task GetResponseAsync( var response = await parentAgent.InnerAgent.RunAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false); // Wrap the response in a ChatResponse so we can pass it back through OpenTelemetryChatClient. - return response.AsChatResponse(); + return ExcludeSubAgentUsage(response.AsChatResponse(), response); } public async IAsyncEnumerable GetStreamingResponseAsync( @@ -327,10 +327,93 @@ public async IAsyncEnumerable GetStreamingResponseAsync( await foreach (var update in parentAgent.InnerAgent.RunStreamingAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false)) { // Wrap the response updates in ChatResponseUpdates so we can pass them back through OpenTelemetryChatClient. - yield return update.AsChatResponseUpdate(); + yield return ExcludeSubAgentUsage(update.AsChatResponseUpdate(), update); } } + /// + /// Returns , or, when the agent's usage includes usage rolled up from sub-agents, a + /// copy of it that reports the agent's own usage only, since the sub-agents report their usage on their own spans. + /// + /// + /// The copy keeps as its raw representation, so that hands + /// the inner agent's response, including the rolled-up usage, back to the caller unchanged. + /// + private static ChatResponse ExcludeSubAgentUsage(ChatResponse chatResponse, AgentResponse response) + { + UsageDetails? ownUsage = SubAgentUsageScope.GetOwnUsage(chatResponse.Usage); + if (ReferenceEquals(ownUsage, chatResponse.Usage)) + { + return chatResponse; + } + + return new(chatResponse.Messages) + { + AdditionalProperties = chatResponse.AdditionalProperties, + ContinuationToken = chatResponse.ContinuationToken, + ConversationId = chatResponse.ConversationId, + CreatedAt = chatResponse.CreatedAt, + FinishReason = chatResponse.FinishReason, + ModelId = chatResponse.ModelId, + RawRepresentation = response, + ResponseId = chatResponse.ResponseId, + Usage = ownUsage, + }; + } + + /// + /// Returns , or, when it carries usage rolled up from sub-agents, a copy of it without + /// that usage, so that the span reports the agent's own usage only. + /// + /// + /// The copy keeps as its raw representation, so that + /// hands the inner agent's update, including the rolled-up usage, back to the caller unchanged. + /// + private static ChatResponseUpdate ExcludeSubAgentUsage(ChatResponseUpdate chatUpdate, AgentResponseUpdate update) + { + List? filtered = null; + for (int i = 0; i < chatUpdate.Contents.Count; i++) + { + AIContent content = chatUpdate.Contents[i]; + if (SubAgentUsageScope.IsSubAgentUsageContent(content)) + { + if (filtered is null) + { + filtered = new(chatUpdate.Contents.Count - 1); + for (int j = 0; j < i; j++) + { + filtered.Add(chatUpdate.Contents[j]); + } + } + } + else + { + filtered?.Add(content); + } + } + + if (filtered is null) + { + return chatUpdate; + } + + return new() + { + AdditionalProperties = chatUpdate.AdditionalProperties, + AuthorName = chatUpdate.AuthorName, + Contents = filtered, + ContinuationToken = chatUpdate.ContinuationToken, + ConversationId = chatUpdate.ConversationId, + CreatedAt = chatUpdate.CreatedAt, + FinishReason = chatUpdate.FinishReason, + MessageId = chatUpdate.MessageId, + ModelId = chatUpdate.ModelId, + RawRepresentation = update, + ResponseId = chatUpdate.ResponseId, + Role = chatUpdate.Role, + }; + } + public object? GetService(Type serviceType, object? serviceKey = null) => // Delegate any inquiries made by the OpenTelemetryChatClient back to the parent agent. parentAgent.GetService(serviceType, serviceKey); diff --git a/dotnet/src/Microsoft.Agents.AI/SubAgentUsageScope.cs b/dotnet/src/Microsoft.Agents.AI/SubAgentUsageScope.cs new file mode 100644 index 00000000000..c5ced87b181 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/SubAgentUsageScope.cs @@ -0,0 +1,141 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Runtime.CompilerServices; +using System.Threading; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// AsyncLocal carrier that rolls up the token usage of sub-agents into the usage reported by the agent run that +/// invoked them. +/// +/// +/// +/// pushes a fresh onto this scope at the start of every +/// run, before its instances are invoked. Because +/// mutations inside a child async method do not propagate back to the caller (copy-on-write semantics), the +/// accumulator is a reference type: code running within the run, such as the function produced by +/// , captures and adds the sub-agent's reported +/// into it. When the run finishes, the agent merges the accumulated total into its +/// own . +/// +/// +/// Every sub-agent run establishes its own accumulator, so the usage it reports already includes its own sub-agents. +/// Only the call boundary reports usage to the parent, so in a chain A → B → C every agent's tokens are counted +/// exactly once. +/// +/// +/// The portion of a merged usage that came from sub-agents is tracked here, keyed by object identity, so that +/// can keep reporting only the agent's own usage on its invoke_agent span. +/// +/// +internal static class SubAgentUsageScope +{ + private static readonly AsyncLocal s_current = new(); + + /// Maps a merged instance to the agent's own usage before the merge. + private static readonly ConditionalWeakTable> s_ownUsageByMergedUsage = new(); + + /// Tracks the instances that carry only sub-agent usage on a streamed update. + private static readonly ConditionalWeakTable s_subAgentUsageContents = new(); + + /// + /// Gets or sets the accumulator for the current agent run, or when no run that rolls up + /// sub-agent usage is in progress. + /// + public static UsageAccumulator? Current + { + get => s_current.Value; + set => s_current.Value = value; + } + + /// + /// Creates a new that is the sum of and + /// , remembering so that it can be recovered with + /// . + /// + public static UsageDetails Merge(UsageDetails? ownUsage, UsageDetails subAgentUsage) + { + UsageDetails merged = new(); + if (ownUsage is not null) + { + merged.Add(ownUsage); + } + + merged.Add(subAgentUsage); + + s_ownUsageByMergedUsage.Add(merged, new StrongBox(ownUsage)); + return merged; + } + + /// + /// Returns the agent's own usage for a produced by , or + /// itself when it was not produced by a merge. + /// + public static UsageDetails? GetOwnUsage(UsageDetails? usage) => + usage is not null && s_ownUsageByMergedUsage.TryGetValue(usage, out StrongBox? own) + ? own.Value + : usage; + + /// + /// Creates a carrying sub-agent usage on a streamed update, marked so that it can be + /// recognized with . + /// + public static UsageContent CreateSubAgentUsageContent(UsageDetails subAgentUsage) + { + UsageContent content = new(subAgentUsage); + s_subAgentUsageContents.Add(content, content); + return content; + } + + /// + /// Gets a value indicating whether was created by . + /// + public static bool IsSubAgentUsageContent(AIContent content) => + s_subAgentUsageContents.TryGetValue(content, out _); + + /// + /// A thread-safe running total of the usage reported by sub-agents during a single agent run. + /// + /// + /// Sub-agents may report concurrently, for example when the function-invoking client runs tool calls in parallel. + /// + internal sealed class UsageAccumulator + { + private readonly object _lock = new(); + private UsageDetails? _total; + + /// Adds to the running total. A value is ignored. + public void Add(UsageDetails? usage) + { + if (usage is null) + { + return; + } + + lock (this._lock) + { + (this._total ??= new()).Add(usage); + } + } + + /// + /// Gets a snapshot of the running total, or if no sub-agent has reported usage. + /// + public UsageDetails? GetTotal() + { + lock (this._lock) + { + if (this._total is null) + { + return null; + } + + UsageDetails snapshot = new(); + snapshot.Add(this._total); + return snapshot; + } + } + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_SubAgentUsageTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_SubAgentUsageTests.cs new file mode 100644 index 00000000000..aa9f87c2814 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_SubAgentUsageTests.cs @@ -0,0 +1,481 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Linq; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Moq; +using OpenTelemetry.Trace; + +namespace Microsoft.Agents.AI.UnitTests; + +/// +/// Unit tests verifying that a rolls up the usage of the sub-agents it invokes into its +/// own . +/// +public class ChatClientAgent_SubAgentUsageTests +{ + // Every scripted model call made by a tool-calling agent reports this usage, and such an agent makes two model + // calls per run (one requesting the tool calls, one producing the final answer). + private const long OwnInputPerCall = 10; + private const long OwnOutputPerCall = 5; + private const long OwnInputPerRun = 2 * OwnInputPerCall; + private const long OwnOutputPerRun = 2 * OwnOutputPerCall; + + /// + /// Verify that the usage of a sub-agent invoked through AsAIFunction is added to the parent's usage. + /// + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Run_SubAgentViaAsAIFunction_ParentUsageIncludesSubAgentUsageAsync(bool streaming) + { + // Arrange + ChatClientAgent child = CreateTextAgent("Child", Usage(100, 50)); + ChatClientAgent parent = CreateToolCallingAgent("Parent", child.AsAIFunction()); + + // Act + AgentResponse response = await RunAsync(parent, streaming); + + // Assert + AssertUsage(response.Usage, OwnInputPerRun + 100, OwnOutputPerRun + 50); + } + + /// + /// Verify that a streaming run ends with an update carrying only the rolled-up sub-agent usage. + /// + [Fact] + public async Task RunStreamingAsync_SubAgentViaAsAIFunction_LastUpdateCarriesSubAgentUsageAsync() + { + // Arrange + ChatClientAgent child = CreateTextAgent("Child", Usage(100, 50)); + ChatClientAgent parent = CreateToolCallingAgent("Parent", child.AsAIFunction()); + + // Act + List updates = []; + await foreach (AgentResponseUpdate update in parent.RunStreamingAsync("question")) + { + updates.Add(update); + } + + // Assert + AgentResponseUpdate last = updates[updates.Count - 1]; + Assert.Null(last.Role); + Assert.Equal(parent.Id, last.AgentId); + UsageContent usageContent = Assert.IsType(Assert.Single(last.Contents)); + Assert.True(SubAgentUsageScope.IsSubAgentUsageContent(usageContent)); + AssertUsage(usageContent.Details, 100, 50); + } + + /// + /// Verify that in a chain A → B → C every agent's usage is counted exactly once. + /// + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Run_NestedSubAgents_CountsEachLayerOnceAsync(bool streaming) + { + // Arrange + ChatClientAgent c = CreateTextAgent("C", Usage(1000, 500)); + ChatClientAgent b = CreateToolCallingAgent("B", c.AsAIFunction()); + ChatClientAgent a = CreateToolCallingAgent("A", b.AsAIFunction()); + + // Act + AgentResponse response = await RunAsync(a, streaming); + + // Assert: A's own usage + B's own usage + C's usage, each once. + AssertUsage(response.Usage, (2 * OwnInputPerRun) + 1000, (2 * OwnOutputPerRun) + 500); + } + + /// + /// Verify that the usage of sub-agents invoked concurrently is summed correctly. + /// + [Fact] + public async Task RunAsync_ConcurrentSubAgentCalls_SumsAllSubAgentUsageAsync() + { + // Arrange: three tool calls in one response, invoked concurrently. Each child waits until all three have started, + // so the usage is reported concurrently. + const int CallCount = 3; + int started = 0; + var allStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var child = new TestAIAgent + { + NameFunc = () => "Child", + RunAsyncFunc = async (_, _, _, cancellationToken) => + { + if (Interlocked.Increment(ref started) == CallCount) + { + allStarted.TrySetResult(true); + } + + await Task.WhenAny(allStarted.Task, Task.Delay(TimeSpan.FromSeconds(10), cancellationToken)); + return new AgentResponse(new ChatMessage(ChatRole.Assistant, "child answer")) { Usage = Usage(100, 50) }; + }, + }; + + IChatClient chatClient = new ChatClientBuilder(CreateToolCallingClient("Child", CallCount)) + .UseFunctionInvocation(configure: f => f.AllowConcurrentInvocation = true) + .Build(); + var parent = new ChatClientAgent(chatClient, new ChatClientAgentOptions + { + Name = "Parent", + ChatOptions = new() { Tools = [child.AsAIFunction()] }, + }); + + // Act + AgentResponse response = await parent.RunAsync("question"); + + // Assert + Assert.Equal(CallCount, started); + AssertUsage(response.Usage, OwnInputPerRun + (CallCount * 100), OwnOutputPerRun + (CallCount * 50)); + } + + /// + /// Verify that a sub-agent reporting no usage leaves the parent's usage unchanged. + /// + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Run_SubAgentWithNullUsage_ParentUsageUnchangedAsync(bool streaming) + { + // Arrange + ChatClientAgent child = CreateTextAgent("Child", usage: null); + ChatClientAgent parent = CreateToolCallingAgent("Parent", child.AsAIFunction()); + + // Act + List updates = []; + AgentResponse response; + if (streaming) + { + await foreach (AgentResponseUpdate update in parent.RunStreamingAsync("question")) + { + updates.Add(update); + } + + response = updates.ToAgentResponse(); + } + else + { + response = await parent.RunAsync("question"); + } + + // Assert + AssertUsage(response.Usage, OwnInputPerRun, OwnOutputPerRun); + Assert.DoesNotContain(updates, u => u.Contents.Any(SubAgentUsageScope.IsSubAgentUsageContent)); + } + + /// + /// Verify that a function created by AsAIFunction does nothing with usage when no agent run is in progress. + /// + [Fact] + public async Task AsAIFunction_NoScope_DoesNotRecordUsageAsync() + { + // Arrange + SubAgentUsageScope.Current = null; + AIFunction function = CreateTextAgent("Child", Usage(100, 50)).AsAIFunction(); + + // Act + object? result = await function.InvokeAsync(new AIFunctionArguments { ["query"] = "question" }); + + // Assert + Assert.NotNull(result); + Assert.Null(SubAgentUsageScope.Current); + } + + /// + /// Verify that a function created by AsAIFunction adds the sub-agent's usage to the scope in progress, and that the + /// sub-agent itself runs without the caller's scope, so its usage is only counted at the call boundary. + /// + [Fact] + public async Task AsAIFunction_WithScope_AddsSubAgentUsageAndClearsScopeForSubAgentAsync() + { + // Arrange + SubAgentUsageScope.UsageAccumulator? scopeSeenByChild = new(); + var child = new TestAIAgent + { + NameFunc = () => "Child", + RunAsyncFunc = (_, _, _, _) => + { + scopeSeenByChild = SubAgentUsageScope.Current; + return Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "child answer")) { Usage = Usage(100, 50) }); + }, + }; + var accumulator = new SubAgentUsageScope.UsageAccumulator(); + SubAgentUsageScope.Current = accumulator; + + // Act + await child.AsAIFunction().InvokeAsync(new AIFunctionArguments { ["query"] = "question" }); + + // Assert + Assert.Null(scopeSeenByChild); + Assert.Same(accumulator, SubAgentUsageScope.Current); + AssertUsage(accumulator.GetTotal(), 100, 50); + } + + /// + /// Verify that the invoke_agent span reports the agent's own usage only, while the response includes the sub-agent usage. + /// + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Run_WithOpenTelemetry_SpanExcludesSubAgentUsageAsync(bool streaming) + { + // Arrange + string sourceName = Guid.NewGuid().ToString(); + var activities = new List(); + using TracerProvider tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder() + .AddSource(sourceName) + .AddInMemoryExporter(activities) + .Build(); + + ChatClientAgent child = CreateTextAgent("Child", Usage(100, 50)); + AIAgent parent = CreateToolCallingAgent("Parent", child.AsAIFunction()) + .AsBuilder() + .UseOpenTelemetry(sourceName) + .Build(); + + // Act + AgentResponse response = await RunAsync(parent, streaming); + + // Assert + AssertUsage(response.Usage, OwnInputPerRun + 100, OwnOutputPerRun + 50); + + Activity invokeAgent; + lock (activities) + { + invokeAgent = Assert.Single(activities, a => (a.GetTagItem("gen_ai.operation.name") as string) == "invoke_agent"); + } + + Assert.Equal(OwnInputPerRun, Convert.ToInt64(invokeAgent.GetTagItem("gen_ai.usage.input_tokens"))); + Assert.Equal(OwnOutputPerRun, Convert.ToInt64(invokeAgent.GetTagItem("gen_ai.usage.output_tokens"))); + } + + /// + /// Verify that the usage of a background agent is charged to the parent run that observes its completion through the + /// wait tool, and not again in a later run. + /// + [Fact] + public async Task RunAsync_BackgroundAgentObservedByWaitTool_ChargedOnceAsync() + { + // Arrange + ChatClientAgent worker = CreateTextAgent("Worker", Usage(100, 50)); + var provider = new BackgroundAgentsProvider([worker]); + + int callIndex = 0; + var parentClient = new ScriptedChatClient(_ => + { + ChatMessage message = Interlocked.Increment(ref callIndex) switch + { + 1 => new(ChatRole.Assistant, [new FunctionCallContent("call1", "background_agents_start_task", new Dictionary + { + ["agentName"] = "Worker", + ["input"] = "do the work", + ["description"] = "work", + })]), + 2 => new(ChatRole.Assistant, [new FunctionCallContent("call2", "background_agents_wait_for_first_completion", new Dictionary + { + ["taskIds"] = new List { 1 }, + })]), + _ => new(ChatRole.Assistant, "done"), + }; + + return new ChatResponse(message) { Usage = Usage(OwnInputPerCall, OwnOutputPerCall) }; + }); + + var parent = new ChatClientAgent(parentClient, new ChatClientAgentOptions + { + Name = "Parent", + AIContextProviders = [provider], + }); + AgentSession session = await parent.CreateSessionAsync(); + + // Act + AgentResponse first = await parent.RunAsync("start the work", session); + AgentResponse second = await parent.RunAsync("anything else?", session); + + // Assert + Assert.Equal(4, callIndex); + AssertUsage(first.Usage, (3 * OwnInputPerCall) + 100, (3 * OwnOutputPerCall) + 50); + AssertUsage(second.Usage, OwnInputPerCall, OwnOutputPerCall); + } + + /// + /// Verify that the usage of a background agent that finished between parent runs is charged once, to the next parent + /// run, and not again when the task is later finalized by a tool. + /// + [Fact] + public async Task BackgroundAgentsProvider_RunFinishedBetweenParentRuns_ChargedOnceToNextRunAsync() + { + // Arrange + ChatClientAgent worker = CreateTextAgent("Worker", Usage(100, 50)); + var provider = new BackgroundAgentsProvider([worker]); + var session = new ChatClientAgentSession(); + + var firstRun = new SubAgentUsageScope.UsageAccumulator(); + SubAgentUsageScope.Current = firstRun; + AIContext firstContext = await InvokeProviderAsync(provider, session); + await GetTool(firstContext, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Worker", + ["input"] = "do the work", + ["description"] = "work", + }); + + // Wait for the background run to finish without finalizing it. + string runtimeStateKey = provider.StateKeys[1]; + Assert.True(session.StateBag.TryGetValue(runtimeStateKey, out BackgroundAgentRuntimeState? runtimeState, AgentJsonUtilities.DefaultOptions)); + Task backgroundRun; + lock (runtimeState!.SyncRoot) + { + backgroundRun = runtimeState.InFlightTasks[1]; + } + + await backgroundRun; + + // Act: the next parent run observes the finished run when its context is provided. + var secondRun = new SubAgentUsageScope.UsageAccumulator(); + SubAgentUsageScope.Current = secondRun; + await InvokeProviderAsync(provider, session); + + // A third parent run observes it again and then finalizes it through a tool. + var thirdRun = new SubAgentUsageScope.UsageAccumulator(); + SubAgentUsageScope.Current = thirdRun; + AIContext thirdContext = await InvokeProviderAsync(provider, session); + object? result = await GetTool(thirdContext, "background_agents_get_task_results").InvokeAsync(new AIFunctionArguments { ["taskId"] = 1 }); + + // Assert + Assert.Null(firstRun.GetTotal()); + AssertUsage(secondRun.GetTotal(), 100, 50); + Assert.Null(thirdRun.GetTotal()); + Assert.Contains("child answer", result?.ToString()); + } + + /// + /// Verify that a sub-agent's usage counts toward the LoopAgent token budget. + /// + [Fact] + public async Task LoopAgent_SubAgentUsage_TriggersTokenBudgetAsync() + { + // Arrange: each iteration costs 30 own tokens plus 150 sub-agent tokens. Without the sub-agent usage, the budget + // would not be reached before the iteration cap. + int childCalls = 0; + var child = new TestAIAgent + { + NameFunc = () => "Child", + RunAsyncFunc = (_, _, _, _) => + { + Interlocked.Increment(ref childCalls); + return Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "child answer")) { Usage = Usage(100, 50) }); + }, + }; + ChatClientAgent parent = CreateToolCallingAgent("Parent", child.AsAIFunction()); + var loop = new LoopAgent(parent, LoopTestHelpers.While(_ => true), new LoopAgentOptions { MaxTokens = 200, MaxIterations = 10 }); + + // Act + AgentResponse response = await loop.RunAsync("task", new ChatClientAgentSession()); + + // Assert: 180 < 200 after iteration 1, 360 >= 200 after iteration 2. + Assert.Equal(2, childCalls); + Assert.NotNull(response.AdditionalProperties); + Assert.Equal(LoopExitReason.TokenBudgetExceeded, response.AdditionalProperties![LoopExitReason.AdditionalPropertiesKey]); + } + + private static UsageDetails Usage(long input, long output) => + new() { InputTokenCount = input, OutputTokenCount = output, TotalTokenCount = input + output }; + + private static void AssertUsage(UsageDetails? usage, long expectedInput, long expectedOutput) + { + Assert.NotNull(usage); + Assert.Equal(expectedInput, usage!.InputTokenCount); + Assert.Equal(expectedOutput, usage.OutputTokenCount); + Assert.Equal(expectedInput + expectedOutput, usage.TotalTokenCount); + } + + private static async Task RunAsync(AIAgent agent, bool streaming) + { + if (!streaming) + { + return await agent.RunAsync("question"); + } + + List updates = []; + await foreach (AgentResponseUpdate update in agent.RunStreamingAsync("question")) + { + updates.Add(update); + } + + return updates.ToAgentResponse(); + } + + /// Creates an agent that answers with text, reporting . + private static ChatClientAgent CreateTextAgent(string name, UsageDetails? usage) => + new( + new ScriptedChatClient(_ => new ChatResponse(new ChatMessage(ChatRole.Assistant, "child answer")) + { + Usage = usage is null ? null : Usage(usage.InputTokenCount!.Value, usage.OutputTokenCount!.Value), + }), + new ChatClientAgentOptions { Name = name }); + + /// Creates an agent that calls once and then answers with text. + private static ChatClientAgent CreateToolCallingAgent(string name, AIFunction tool) => + new(CreateToolCallingClient(tool.Name, callCount: 1), new ChatClientAgentOptions + { + Name = name, + ChatOptions = new() { Tools = [tool] }, + }); + + /// + /// Creates a chat client that requests calls to when the last + /// message is not a tool result, and answers with text otherwise. Every call reports the same own usage. + /// + private static ScriptedChatClient CreateToolCallingClient(string toolName, int callCount) => + new(messages => + { + ChatMessage message = messages[messages.Count - 1].Contents.OfType().Any() + ? new(ChatRole.Assistant, "parent answer") + : new(ChatRole.Assistant, Enumerable.Range(0, callCount) + .Select(i => (AIContent)new FunctionCallContent($"call{i}", toolName, new Dictionary { ["query"] = "sub question" })) + .ToList()); + + return new ChatResponse(message) { Usage = Usage(OwnInputPerCall, OwnOutputPerCall) }; + }); + + private static async Task InvokeProviderAsync(BackgroundAgentsProvider provider, AgentSession session) + { +#pragma warning disable MAAI001 + var context = new AIContextProvider.InvokingContext(new Mock().Object, session, new AIContext()); +#pragma warning restore MAAI001 + return await provider.InvokingAsync(context); + } + + private static AIFunction GetTool(AIContext context, string name) => + (AIFunction)context.Tools!.First(t => t is AIFunction f && f.Name == name); + + /// A chat client whose responses are produced by a callback, streamed as the equivalent updates. + private sealed class ScriptedChatClient(Func, ChatResponse> handler) : IChatClient + { + public Task GetResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) => + Task.FromResult(handler(messages.ToList())); + + public async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + await Task.Yield(); + foreach (ChatResponseUpdate update in handler(messages.ToList()).ToChatResponseUpdates()) + { + yield return update; + } + } + + public object? GetService(Type serviceType, object? serviceKey = null) => + serviceKey is null && serviceType.IsInstanceOfType(this) ? this : null; + + public void Dispose() + { + } + } +}