.NET: add token and time budgets to LoopAgent

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.
     /// </para>
+    /// <para>
+    /// When the resulting function is invoked during the run of a <see cref="ChatClientAgent"/>, the
+    /// <see cref="AgentResponse.Usage"/> reported by <paramref name="agent"/> 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 <paramref name="agent"/>
+    /// fails or is cancelled, it produces no response and its usage is not counted.
+    /// </para>
     /// </remarks>
     public static AIFunction AsAIFunction(this AIAgent agent, AIFunctionFactoryOptions? options = null, AgentSession? session = null)
     {
@@ -80,7 +86,18 @@ async Task<string> 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;
         }
 

dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs

@@ -230,6 +230,11 @@ protected override async Task<AgentResponse> 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<ChatMessage> inputMessagesForChatClient,
@@ -276,6 +281,13 @@ protected override async Task<AgentResponse> 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<AgentResponseUpdate> 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<ChatMessage> inputMessagesForChatClient,
@@ -361,6 +379,7 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> 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<AgentResponseUpdate> 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<AgentResponseUpdate> 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<AgentResponseUpdate> 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,
+            };
+        }
     }
 
     /// <inheritdoc/>

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<int, CancellationTokenSource> TaskCancellations { get; } = [];
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    /// <remarks>
+    /// 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.
+    /// </remarks>
+    [JsonIgnore]
+    public HashSet<Task<AgentResponse>> UsageChargedRuns { get; } = [];
+
+    /// <summary>
+    /// Gets the usage of runs that were finalized while no parent run was in progress (for example through
+    /// <see cref="BackgroundAgentsProvider.GetIncompleteTasks"/>), to be charged to the next parent run.
+    /// </summary>
+    [JsonIgnore]
+    public List<UsageDetails> UnchargedUsage { get; } = [];
+
     /// <summary>
     /// Gets or sets a value indicating whether this runtime has been released via
     /// <see cref="BackgroundAgentsProvider.ReleaseSessionAsync"/>.

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.
 /// </para>
 /// <para>
+/// The <see cref="AgentResponse.Usage"/> 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.
+/// </para>
+/// <para>
 /// <strong>Security considerations:</strong> 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<AIContext> 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<AIContext>(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
         }
     }
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// 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.
+    /// </para>
+    /// <para>Callers must hold <see cref="BackgroundAgentRuntimeState.SyncRoot"/>.</para>
+    /// </remarks>
+    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<AgentResponse> 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
+            }
+        }
+    }
+
     /// <summary>
     /// Finalizes a task by extracting results from the completed Task and updating the BackgroundTaskInfo.
     /// </summary>
+    /// <param name="taskInfo">The task to finalize.</param>
+    /// <param name="completedTask">The completed run of the task.</param>
+    /// <param name="runtimeState">The runtime state of the session.</param>
+    /// <param name="chargeUsage">
+    /// 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.
+    /// </param>
     /// <remarks>Callers must hold <see cref="BackgroundAgentRuntimeState.SyncRoot"/>.</remarks>
-    private static void FinalizeTask(BackgroundTaskInfo taskInfo, Task<AgentResponse> completedTask, BackgroundAgentRuntimeState runtimeState)
+    private static void FinalizeTask(BackgroundTaskInfo taskInfo, Task<AgentResponse> 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;
         }
     }