.NET: Preserve toolbox OAuth consent context

dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs

@@ -348,7 +348,6 @@ public override async IAsyncEnumerable<ResponseStreamEvent> CreateAsync(
             cancellationToken,
             context.Shutdown);
         var consentState = new RequestConsentState { CancellationSource = consentCts };
-        McpConsentContext.Current.Value = consentState;
 
         // 7. Run the agent and convert output
         // NOTE: C# forbids 'yield return' inside a try block that has a catch clause,
@@ -397,9 +396,20 @@ public override async IAsyncEnumerable<ResponseStreamEvent> CreateAsync(
                 {
                     if (!pendingEvents.TryDequeue(out evt))
                     {
-                        if (!await enumerator.MoveNextAsync().ConfigureAwait(false))
+                        // AsyncLocal assignments do not survive iterator yields. Bind the same
+                        // consent state on every advance without replacing the caller's context.
+                        var previousConsentState = McpConsentContext.Current.Value;
+                        try
                         {
-                            break;
+                            McpConsentContext.Current.Value = consentState;
+                            if (!await enumerator.MoveNextAsync().ConfigureAwait(false))
+                            {
+                                break;
+                            }
+                        }
+                        finally
+                        {
+                            McpConsentContext.Current.Value = previousConsentState;
                         }
 
                         switch (enumerator.Current)
@@ -592,7 +602,16 @@ await PersistWorkflowCheckpointAsync(
         }
         finally
         {
-            await enumerator.DisposeAsync().ConfigureAwait(false);
+            var previousConsentState = McpConsentContext.Current.Value;
+            try
+            {
+                McpConsentContext.Current.Value = consentState;
+                await enumerator.DisposeAsync().ConfigureAwait(false);
+            }
+            finally
+            {
+                McpConsentContext.Current.Value = previousConsentState;
+            }
 
             // Only after the the agent ran when can check precisely if the session had been used to store messages in the backend for validation.
             if (CheckNotAllowedStoreUsage())

dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ConsentAwareMcpClientAIFunction.cs

@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
@@ -21,8 +22,8 @@ namespace Microsoft.Agents.AI.Foundry.Hosting;
 /// and the per-request <see cref="RequestConsentState.CancellationSource"/> is cancelled. This causes
 /// <see cref="FunctionInvokingChatClient"/> to stop the tool loop (it guards
 /// exceptions with <c>when (!ct.IsCancellationRequested)</c>) and surfaces an
-/// <see cref="System.OperationCanceledException"/> to the handler. The handler then emits the
-/// <c>mcp_approval_request</c> output item and marks the response as <c>incomplete</c>.
+/// <see cref="OperationCanceledException"/> to the handler. The handler then emits the
+/// <c>oauth_consent_request</c> output item and marks the response as <c>incomplete</c>.
 /// </para>
 /// </remarks>
 internal sealed class ConsentAwareMcpClientAIFunction : AIFunction
@@ -59,7 +60,19 @@ internal ConsentAwareMcpClientAIFunction(McpClientTool inner, string toolboxName
             var state = McpConsentContext.Current.Value;
             if (state is not null)
             {
-                state.Pending = new McpConsentInfo(this._toolboxName, this._inner.Name, ex.Message);
+                const string RemoteErrorPrefix = "Request failed (remote): ";
+                var consentUrl = ex.Data["detail"] switch
+                {
+                    string detail => detail,
+                    JsonElement { ValueKind: JsonValueKind.String } detail => detail.GetString()!,
+                    _ when !ex.Data.Contains("detail") => ex.Message.StartsWith(RemoteErrorPrefix, StringComparison.Ordinal)
+                        ? ex.Message[RemoteErrorPrefix.Length..]
+                        : ex.Message,
+                    // The detail field exists but is not a string. Use an empty URL so the
+                    // handler rejects this consent response instead of falling back to the message.
+                    _ => string.Empty,
+                };
+                state.Pending = new McpConsentInfo(this._toolboxName, this._inner.Name, consentUrl);
                 state.CancellationSource?.Cancel();
             }
 

dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/McpConsentContext.cs

@@ -39,7 +39,7 @@ internal static class McpConsentContext
 {
     /// <summary>
     /// Holds the shared <see cref="RequestConsentState"/> for the current request.
-    /// Set once by the handler; read and mutated by the tool wrapper.
+    /// Bound while advancing or disposing the agent stream; read and mutated by the tool wrapper.
     /// </summary>
     internal static readonly AsyncLocal<RequestConsentState?> Current = new();
 }

dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OAuthConsentPolicyHandlerTests.cs

@@ -148,6 +148,32 @@ public async Task CreateAsync_PerCallConsentMatchingAllowlist_SavesSessionAsync(
         Assert.Equal(1, sessionStore.SaveAttempts);
     }
 
+    [Fact]
+    public async Task CreateAsync_ConsumerStopsReading_DisposesAgentWithRequestConsentContextAsync()
+    {
+        // Arrange
+        await using var service = await CreateStartedServiceAsync(allowedOrigins: null, toolboxNames: [], opener: null);
+        var agent = new ConsentContextTrackingAgent();
+        var callerState = new RequestConsentState();
+        var previousState = McpConsentContext.Current.Value;
+        McpConsentContext.Current.Value = callerState;
+        try
+        {
+            // Act
+            await RunAsync(agent, service, stopAfterFirstOutput: true);
+
+            // Assert
+            Assert.NotNull(agent.RunState);
+            Assert.NotSame(callerState, agent.RunState);
+            Assert.Same(agent.RunState, agent.DisposeState);
+            Assert.Same(callerState, McpConsentContext.Current.Value);
+        }
+        finally
+        {
+            McpConsentContext.Current.Value = previousState;
+        }
+    }
+
     private static void AssertFailedWithoutLink(List<ResponseStreamEvent> events, string rejectedLink)
     {
         Assert.DoesNotContain(events, e => e is ResponseOutputItemAddedEvent { Item: OAuthConsentRequestOutputItem });
@@ -211,7 +237,8 @@ private static async Task<List<ResponseStreamEvent>> RunAsync(
         AIAgent agent,
         FoundryToolboxService toolboxService,
         CreateResponse? request = null,
-        AgentSessionStore? sessionStore = null)
+        AgentSessionStore? sessionStore = null,
+        bool stopAfterFirstOutput = false)
     {
         var services = new ServiceCollection();
         services.AddSingleton(sessionStore ?? new InMemoryAgentSessionStore());
@@ -230,6 +257,10 @@ private static async Task<List<ResponseStreamEvent>> RunAsync(
         await foreach (var evt in handler.CreateAsync(request ?? CreateRequest(), context.Object, CancellationToken.None))
         {
             events.Add(evt);
+            if (stopAfterFirstOutput && evt is ResponseOutputItemAddedEvent)
+            {
+                break;
+            }
         }
 
         return events;
@@ -307,6 +338,31 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingA
         }
     }
 
+    private sealed class ConsentContextTrackingAgent : AgentBase
+    {
+        public RequestConsentState? RunState { get; private set; }
+
+        public RequestConsentState? DisposeState { get; private set; }
+
+        protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
+            IEnumerable<ChatMessage> messages,
+            AgentSession? session,
+            AgentRunOptions? options,
+            [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            this.RunState = McpConsentContext.Current.Value;
+            try
+            {
+                await Task.Yield();
+                yield return new AgentResponseUpdate { MessageId = "resp_msg_1", Contents = [new MeaiTextContent("hi")] };
+            }
+            finally
+            {
+                this.DisposeState = McpConsentContext.Current.Value;
+            }
+        }
+    }
+
     /// <summary>
     /// Mirrors <see cref="ConsentAwareMcpClientAIFunction"/> on a <c>-32006</c> response: records the
     /// consent on the request state, cancels the tool loop, and surfaces the cancellation.