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]); + } +} diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index f9b01907f22..3ea45f30b90 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- **agent-framework-core**: Add `AgentBudget` with `max_tokens` and `max_duration` constraints for `AgentLoopMiddleware` autonomous loops; budget stops stamp `loop_exit_reason` (`"token_budget_exceeded"` or `"time_budget_exceeded"`) in `AgentResponse.additional_properties`, and reaching `max_iterations` stamps `"iteration_cap_reached"`. + ## [1.20.0] - 2026-10-02 ### Added diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 3953187578e..efacbf476c3 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -132,6 +132,7 @@ "FileMemoryProvider", ), "._harness._loop": ( + "AgentBudget", "AgentLoopMiddleware", "JudgeVerdict", "background_tasks_running", @@ -432,6 +433,7 @@ "USER_AGENT_KEY", "USER_AGENT_TELEMETRY_DISABLED_ENV_VAR", "Agent", + "AgentBudget", "AgentContext", "AgentEvalConverter", "AgentExecutor", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index e8d1bc6e411..5ba07b7e58a 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -90,6 +90,7 @@ from ._harness._file_access import ( ) from ._harness._file_memory import DEFAULT_FILE_MEMORY_INSTRUCTIONS, DEFAULT_FILE_MEMORY_SOURCE_ID, FileMemoryProvider from ._harness._loop import ( + AgentBudget, AgentLoopMiddleware, JudgeVerdict, background_tasks_running, @@ -387,6 +388,7 @@ __all__ = [ "USER_AGENT_KEY", "USER_AGENT_TELEMETRY_DISABLED_ENV_VAR", "Agent", + "AgentBudget", "AgentContext", "AgentEvalConverter", "AgentExecutor", diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py index b4617a23e01..762f82e55e4 100644 --- a/python/packages/core/agent_framework/_harness/_loop.py +++ b/python/packages/core/agent_framework/_harness/_loop.py @@ -29,7 +29,10 @@ from __future__ import annotations import inspect +import math +import time from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, TypeAlias from pydantic import BaseModel, Field @@ -57,6 +60,7 @@ from .._clients import SupportsChatGetResponse __all__ = [ + "AgentBudget", "AgentLoopMiddleware", "JudgeVerdict", "background_tasks_running", @@ -131,6 +135,58 @@ class JudgeVerdict(BaseModel): DEFAULT_JUDGE_MAX_ITERATIONS = 5 +@experimental(feature_id=ExperimentalFeature.HARNESS) +@dataclass +class AgentBudget: + """Budget constraints for an :class:`AgentLoopMiddleware` autonomous loop. + + Both ``max_tokens`` and ``max_duration`` are optional with no default — they are + workload-specific and must be set explicitly for production autonomous sessions. + When a budget dimension is ``None``, it is unconstrained. + + Examples: + + .. code-block:: python + + from agent_framework import AgentBudget, AgentLoopMiddleware, todos_remaining + + agent = Agent( + client=client, + middleware=[ + AgentLoopMiddleware( + todos_remaining(), + budget=AgentBudget(max_tokens=50_000, max_duration=300.0), + ) + ], + ) + + Note: + Budget stops stamp ``additional_properties["loop_exit_reason"]`` with + ``"token_budget_exceeded"`` or ``"time_budget_exceeded"`` on the returned + :class:`~agent_framework.AgentResponse` (non-streaming only). + """ + + max_tokens: int | None = None + """Maximum total LLM token spend across all loop iterations. + + Counts ``total_token_count`` from each iteration's usage details, falling back to + ``input_token_count + output_token_count`` when ``total_token_count`` is absent. + """ + + max_duration: float | None = None + """Maximum wall-clock time in seconds for the entire loop run. + + Checked at the start of each iteration using :func:`time.monotonic`. + Catches agents blocked on slow tools that token budgets cannot detect. + """ + + def __post_init__(self) -> None: + if self.max_tokens is not None and self.max_tokens < 1: + raise ValueError("AgentBudget.max_tokens must be a positive integer.") + if self.max_duration is not None and (not math.isfinite(self.max_duration) or self.max_duration <= 0): + raise ValueError("AgentBudget.max_duration must be a positive, finite number of seconds.") + + # A callable invoked between iterations. It always receives the loop keyword arguments # (``iteration``, ``last_result``, ``messages``, ``original_messages``, ``session``, ``agent``, # ``progress``, ``feedback``). Callers declare only the keywords they need plus ``**kwargs`` to @@ -285,6 +341,7 @@ def __init__( fresh_context: bool = False, return_final_only: bool = False, additional_instructions: str | None = None, + budget: AgentBudget | None = None, ) -> None: """Initialize the agent loop middleware. @@ -342,6 +399,9 @@ def __init__( messages, so it is preserved across ``fresh_context`` resets and (with a session) persists server-side across iterations. Used by :meth:`with_judge` to tell the agent about the criteria its response must satisfy, but available to any loop. + budget: Optional :class:`AgentBudget` bounding total token spend and/or wall-clock time. + Checked after each iteration, before ``should_continue`` is evaluated. When exceeded + the loop stops and (non-streaming only) stamps ``loop_exit_reason``. Raises: ValueError: If ``max_iterations`` is not ``None`` and is less than 1. @@ -357,6 +417,7 @@ def __init__( self.fresh_context = fresh_context self.return_final_only = return_final_only self.additional_instructions = additional_instructions + self.budget: AgentBudget | None = budget @classmethod def with_judge( @@ -370,6 +431,7 @@ def with_judge( max_iterations: int | None = DEFAULT_JUDGE_MAX_ITERATIONS, next_message: NextMessageCallable | None = None, fresh_context: bool = False, + budget: AgentBudget | None = None, ) -> Self: """Create a loop that continues until a judge chat client decides the request was answered. @@ -424,6 +486,7 @@ def with_judge( conversation; an attached session is snapshotted before the loop and restored to that baseline between iterations. See :meth:`__init__` for the full semantics. Defaults to ``False``. + budget: Optional :class:`AgentBudget`. See :meth:`__init__`. """ judge_instructions = (instructions or DEFAULT_JUDGE_INSTRUCTIONS).replace( CRITERIA_PLACEHOLDER, _render_criteria_block(criteria) @@ -440,6 +503,7 @@ def with_judge( next_message=next_message or judge_next_message, fresh_context=fresh_context, additional_instructions=_criteria_agent_instruction(criteria) if criteria else None, + budget=budget, ) async def process( @@ -538,6 +602,7 @@ async def _process_non_streaming( snapshot: dict[str, Any] | None, ) -> None: iteration = 0 + cap_fired = False work_iterations = 0 progress: list[str] = [] # Aggregated transcript across iterations: each iteration's response messages plus the @@ -545,6 +610,9 @@ async def _process_non_streaming( aggregated: list[Message] = [] aggregated_usage: UsageDetails | None = None final_result: AgentResponse | None = None + tokens_used: int = 0 + budget_start: float = time.monotonic() + budget_exit: str | None = None stamped_options = dict(context.options) if context.options is not None else {} stamped_options[_LOOP_ITERATION_TOKEN_KEY] = object() context.options = stamped_options @@ -565,6 +633,13 @@ async def _process_non_streaming( if result.usage_details is not None: aggregated_usage = add_usage_details(aggregated_usage, result.usage_details) + # Accumulate token spend for budget enforcement. + if self.budget is not None and self.budget.max_tokens is not None and result.usage_details is not None: + ud = result.usage_details + tokens_used += ud.get("total_token_count") or ( + (ud.get("input_token_count") or 0) + (ud.get("output_token_count") or 0) + ) + # Escape hatch: if this iteration is asking for tool approval, stop and return the # response so the caller can approve, instead of continuing or injecting next_message. if self._has_pending_approval_request(result): @@ -581,9 +656,21 @@ async def _process_non_streaming( ) work_iterations += 1 + # Budget pre-check — runs before should_continue to avoid invoking an expensive + # predicate/judge after a budget has already been exhausted. + if self.budget is not None: + elapsed = time.monotonic() - budget_start + if self.budget.max_duration is not None and elapsed >= self.budget.max_duration: + budget_exit = "time_budget_exceeded" + break + if self.budget.max_tokens is not None and tokens_used >= self.budget.max_tokens: + budget_exit = "token_budget_exceeded" + break # Decide whether to stop and capture any feedback from should_continue first, so the # feedback is available to both the progress and next-message callables this iteration. stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations) + if stop and self.max_iterations is not None and work_iterations >= self.max_iterations: + cap_fired = True loop_kwargs = self._build_loop_kwargs( context=context, iteration=iteration, @@ -617,6 +704,11 @@ async def _process_non_streaming( finally: context.options.pop(_LOOP_ITERATION_TOKEN_KEY, None) + if cap_fired: + final_result.additional_properties.setdefault("loop_exit_reason", "iteration_cap_reached") + if budget_exit is not None: + final_result.additional_properties.setdefault("loop_exit_reason", budget_exit) + if not self.return_final_only: context.result = self._aggregate_response(final_result, aggregated, aggregated_usage) await self._fire_turn_scoped_after_providers( @@ -638,8 +730,11 @@ def _process_streaming( async def _generator() -> Any: iteration = 0 + cap_fired = False work_iterations = 0 progress: list[str] = [] + tokens_used: int = 0 + budget_start: float = time.monotonic() stamped_options = dict(context.options) if context.options is not None else {} stamped_options[_LOOP_ITERATION_TOKEN_KEY] = object() context.options = stamped_options @@ -683,9 +778,25 @@ async def _generator() -> Any: ) work_iterations += 1 + # Budget pre-check (before should_continue). Streaming cannot stamp + # ``loop_exit_reason`` on the response (the caller assembles it from updates), so + # the loop simply stops, consistent with the iteration-cap limitation. + if self.budget is not None: + if self.budget.max_tokens is not None and final.usage_details: + ud = final.usage_details + tokens_used += ud.get("total_token_count") or ( + (ud.get("input_token_count") or 0) + (ud.get("output_token_count") or 0) + ) + elapsed = time.monotonic() - budget_start + if self.budget.max_duration is not None and elapsed >= self.budget.max_duration: + return + if self.budget.max_tokens is not None and tokens_used >= self.budget.max_tokens: + return # Decide whether to stop and capture any feedback from should_continue first, so the # feedback is available to both the progress and next-message callables this iteration. stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations) + if stop and self.max_iterations is not None and work_iterations >= self.max_iterations: + cap_fired = True loop_kwargs = self._build_loop_kwargs( context=context, iteration=iteration, @@ -721,6 +832,8 @@ async def _generator() -> Any: yield self._message_to_update(message) finally: context.options.pop(_LOOP_ITERATION_TOKEN_KEY, None) + if cap_fired and holder["final"] is not None: + holder["final"].additional_properties.setdefault("loop_exit_reason", "iteration_cap_reached") await self._fire_turn_scoped_after_providers(context, holder["final"], original_messages) def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse: diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py index 3b87544d804..be2b392b271 100644 --- a/python/packages/core/tests/core/test_harness_loop.py +++ b/python/packages/core/tests/core/test_harness_loop.py @@ -2,6 +2,7 @@ from __future__ import annotations +import time from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence from typing import Any @@ -10,6 +11,7 @@ from agent_framework import ( Agent, + AgentBudget, AgentContext, AgentMiddleware, AgentModeProvider, @@ -29,6 +31,7 @@ ResponseStream, TodoItem, TodoProvider, + UsageDetails, background_tasks_running, background_tasks_running_message, set_agent_mode, @@ -61,8 +64,10 @@ def __init__( texts: list[str] | None = None, honor_response_format: bool = False, service_mode: bool = False, + usage_per_call: int | None = None, ) -> None: super().__init__() + self._usage_per_call = usage_per_call self.call_count: int = 0 self.received_messages: list[list[str]] = [] self.received_response_formats: list[Any] = [] @@ -108,6 +113,9 @@ async def _get() -> ChatResponse: messages=Message(role="assistant", contents=[self._next_text(messages)]), response_format=response_format if self._honor_response_format else None, conversation_id=conversation_id, + usage_details=( + UsageDetails(total_token_count=self._usage_per_call) if self._usage_per_call is not None else None + ), ) return _get() @@ -1697,3 +1705,90 @@ async def test_concurrent_same_agent_run_during_stream_pause_is_not_suppressed() # once for the loop boundary, once for the concurrent run's own boundary assert turn_scoped.after_calls == 2 + + +# region budget + + +async def test_budget_stops_on_token_limit() -> None: + client = RecordingChatClient(usage_per_call=50) + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=10, budget=AgentBudget(max_tokens=100))], + ) + + response = await agent.run("task") + + assert client.call_count == 2 + assert response.additional_properties.get("loop_exit_reason") == "token_budget_exceeded" + + +async def test_budget_stops_on_duration_limit() -> None: + class _SlowClient(RecordingChatClient): + def _inner_get_response(self, **kwargs: Any) -> Any: + time.sleep(0.05) # exceed the 1ms budget (Windows clock tick is ~15ms) + return super()._inner_get_response(**kwargs) + + client = _SlowClient() + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=10, budget=AgentBudget(max_duration=0.001))], + ) + + response = await agent.run("task") + + assert client.call_count == 1 + assert response.additional_properties.get("loop_exit_reason") == "time_budget_exceeded" + + +async def test_budget_not_triggered_when_none() -> None: + client = RecordingChatClient(usage_per_call=50) + agent = Agent(client=client, middleware=[AgentLoopMiddleware(always_continue, max_iterations=3)]) + + response = await agent.run("task") + + assert client.call_count == 3 + assert response.additional_properties.get("loop_exit_reason") == "iteration_cap_reached" + + +async def test_budget_token_check_before_should_continue() -> None: + client = RecordingChatClient(usage_per_call=50) + calls: list[int] = [] + + def should_continue(*, iteration: int, **kwargs: Any) -> bool: + calls.append(iteration) + return True + + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(should_continue, max_iterations=10, budget=AgentBudget(max_tokens=100))], + ) + + await agent.run("task") + + assert client.call_count == 2 + # should_continue ran only for iteration 1; the budget stop pre-empted it on iteration 2. + assert calls == [1] + + +async def test_budget_exit_reason_in_additional_properties() -> None: + client = RecordingChatClient(usage_per_call=60) + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, budget=AgentBudget(max_tokens=50))], + ) + + response = await agent.run("task") + + assert client.call_count == 1 + assert response.additional_properties["loop_exit_reason"] == "token_budget_exceeded" + + +def test_agent_budget_validates_max_tokens() -> None: + with pytest.raises(ValueError, match="max_tokens"): + AgentBudget(max_tokens=0) + + +def test_agent_budget_validates_max_duration() -> None: + with pytest.raises(ValueError, match="max_duration"): + AgentBudget(max_duration=-1.0)