diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs index a472ae53552..094d44bde6e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs @@ -348,7 +348,6 @@ public override async IAsyncEnumerable 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 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()) diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ConsentAwareMcpClientAIFunction.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ConsentAwareMcpClientAIFunction.cs index 5f3ec0ed9b1..eeb1518d2c8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ConsentAwareMcpClientAIFunction.cs +++ b/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 is cancelled. This causes /// to stop the tool loop (it guards /// exceptions with when (!ct.IsCancellationRequested)) and surfaces an -/// to the handler. The handler then emits the -/// mcp_approval_request output item and marks the response as incomplete. +/// to the handler. The handler then emits the +/// oauth_consent_request output item and marks the response as incomplete. /// /// 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(); } diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/McpConsentContext.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/McpConsentContext.cs index 96ea08383b1..e0ce82436ea 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/McpConsentContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/McpConsentContext.cs @@ -39,7 +39,7 @@ internal static class McpConsentContext { /// /// Holds the shared 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. /// internal static readonly AsyncLocal Current = new(); } diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OAuthConsentPolicyHandlerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OAuthConsentPolicyHandlerTests.cs index d346db9ae36..1ce0d2bf4db 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OAuthConsentPolicyHandlerTests.cs +++ b/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 events, string rejectedLink) { Assert.DoesNotContain(events, e => e is ResponseOutputItemAddedEvent { Item: OAuthConsentRequestOutputItem }); @@ -211,7 +237,8 @@ private static async Task> 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> 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 RunCoreStreamingA } } + private sealed class ConsentContextTrackingAgent : AgentBase + { + public RequestConsentState? RunState { get; private set; } + + public RequestConsentState? DisposeState { get; private set; } + + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable 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; + } + } + } + /// /// Mirrors on a -32006 response: records the /// consent on the request state, cancels the tool loop, and surfaces the cancellation. diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OAuthConsentToolInvocationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OAuthConsentToolInvocationTests.cs new file mode 100644 index 00000000000..e03a68e0299 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/OAuthConsentToolInvocationTests.cs @@ -0,0 +1,302 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.AI.AgentServer.Responses; +using Azure.AI.AgentServer.Responses.Models; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging.Abstractions; +using ModelContextProtocol.Client; +using Moq; + +namespace Microsoft.Agents.AI.Foundry.Hosting.UnitTests; + +public sealed class OAuthConsentToolInvocationTests +{ + private const string ConsentLink = "https://auth.example.com/authorize?state=test"; + + [Theory] + [InlineData("graph-retrieval___SearchInternalKnowledge")] + [InlineData("call_tool")] + public async Task CreateAsync_ToolRequiresConsentAfterStreamingOutput_SurfacesConsentAsync(string toolName) + { + // Arrange + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(15)); + using var peer = new ConsentPeer(ConsentLink); + using var httpClient = new HttpClient(peer); + await using var transport = CreateTransport(httpClient); + await using var client = await McpClient.CreateAsync(transport, cancellationToken: timeout.Token); + + // Act + var (events, modelCalls) = await RunAsync(client, toolName, timeout.Token); + + // Assert + Assert.Equal(1, peer.ToolCalls); + var added = Assert.Single(events.OfType(), e => e.Item is OAuthConsentRequestOutputItem); + var consent = Assert.IsType(added.Item); + Assert.Equal(ConsentLink, consent.ConsentLink); + Assert.Equal(toolName, consent.ServerLabel); + var done = Assert.Single(events.OfType(), e => e.Item is OAuthConsentRequestOutputItem); + Assert.Equal(consent.Id, Assert.IsType(done.Item).Id); + Assert.IsType(events[^1]); + Assert.Equal(1, modelCalls); + } + + [Fact] + public async Task CreateAsync_ConsentWithoutStructuredDetail_UsesMessageUrlAsync() + { + // Arrange + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(15)); + using var peer = new ConsentPeer(ConsentLink, includeDetail: false); + using var httpClient = new HttpClient(peer); + await using var transport = CreateTransport(httpClient); + await using var client = await McpClient.CreateAsync(transport, cancellationToken: timeout.Token); + + // Act + var (events, _) = await RunAsync(client, "call_tool", timeout.Token); + + // Assert + var added = Assert.Single(events.OfType(), e => e.Item is OAuthConsentRequestOutputItem); + Assert.Equal(ConsentLink, Assert.IsType(added.Item).ConsentLink); + } + + [Theory] + [InlineData(null)] + [InlineData(42)] + [InlineData("http://auth.example.com/authorize")] + public async Task CreateAsync_InvalidStructuredConsentDetail_FailsWithoutUsingMessageUrlAsync(object? detail) + { + // Arrange + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(15)); + using var peer = new ConsentPeer(ConsentLink) { ErrorDetail = detail }; + using var httpClient = new HttpClient(peer); + await using var transport = CreateTransport(httpClient); + await using var client = await McpClient.CreateAsync(transport, cancellationToken: timeout.Token); + + // Act + var (events, modelCalls) = await RunAsync(client, "call_tool", timeout.Token); + + // Assert + Assert.DoesNotContain(events, e => e is ResponseOutputItemAddedEvent { Item: OAuthConsentRequestOutputItem }); + var failed = Assert.IsType(events[^1]); + Assert.Contains("consent link policy", failed.Response.Error.Message); + Assert.DoesNotContain(ConsentLink, failed.Response.Error.Message); + Assert.Equal(1, modelCalls); + } + + [Fact] + public async Task CreateAsync_ConcurrentConsentRequests_KeepTheirOwnLinksAndCancellationAsync() + { + // Arrange + const string OtherConsentLink = "https://auth.example.com/authorize?state=other"; + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(15)); + var bothCallsEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + int callsEntered = 0; + using var firstPeer = new ConsentPeer(ConsentLink) { BeforeToolCall = WaitForBothCallsAsync }; + using var secondPeer = new ConsentPeer(OtherConsentLink) { BeforeToolCall = WaitForBothCallsAsync }; + using var firstHttpClient = new HttpClient(firstPeer); + using var secondHttpClient = new HttpClient(secondPeer); + await using var firstTransport = CreateTransport(firstHttpClient); + await using var secondTransport = CreateTransport(secondHttpClient); + await using var firstClient = await McpClient.CreateAsync(firstTransport, cancellationToken: timeout.Token); + await using var secondClient = await McpClient.CreateAsync(secondTransport, cancellationToken: timeout.Token); + using var callerCancellation = new CancellationTokenSource(); + var callerState = new RequestConsentState { CancellationSource = callerCancellation }; + var previousState = McpConsentContext.Current.Value; + McpConsentContext.Current.Value = callerState; + try + { + // Act + var results = await Task.WhenAll( + RunAsync(firstClient, "call_tool", timeout.Token), + RunAsync(secondClient, "call_tool", timeout.Token)); + + // Assert + Assert.Equal(ConsentLink, GetConsentLink(results[0].Events)); + Assert.Equal(OtherConsentLink, GetConsentLink(results[1].Events)); + Assert.All(results, result => + { + Assert.IsType(result.Events[^1]); + Assert.Equal(1, result.ModelCalls); + }); + Assert.Same(callerState, McpConsentContext.Current.Value); + Assert.Null(callerState.Pending); + Assert.False(callerCancellation.IsCancellationRequested); + } + finally + { + McpConsentContext.Current.Value = previousState; + } + + async Task WaitForBothCallsAsync(CancellationToken cancellationToken) + { + if (Interlocked.Increment(ref callsEntered) == 2) + { + bothCallsEntered.TrySetResult(); + } + + await bothCallsEntered.Task.WaitAsync(cancellationToken); + } + } + + private static string GetConsentLink(List events) => + Assert.IsType( + Assert.Single(events.OfType(), e => e.Item is OAuthConsentRequestOutputItem).Item).ConsentLink; + + private static HttpClientTransport CreateTransport(HttpClient httpClient) => new( + new HttpClientTransportOptions + { + Endpoint = new Uri("https://toolbox.example.com/mcp"), + TransportMode = HttpTransportMode.StreamableHttp, + }, + httpClient); + + private static async Task<(List Events, int ModelCalls)> RunAsync( + McpClient client, + string toolName, + CancellationToken cancellationToken) + { + var tool = new ConsentAwareMcpClientAIFunction( + new McpClientTool(client, new ModelContextProtocol.Protocol.Tool + { + Name = toolName, + Description = "Invokes the OAuth-protected tool.", + InputSchema = JsonSerializer.SerializeToElement(new { type = "object" }), + }, serializerOptions: null), + "test-toolbox"); + var modelCalls = 0; + var chatClient = new Mock(MockBehavior.Loose); + chatClient.Setup(c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(() => Interlocked.Increment(ref modelCalls) == 1 + ? YieldUpdateAsync(new ChatResponseUpdate( + ChatRole.Assistant, + [new FunctionCallContent("call_test", toolName)]) + { + MessageId = "tool_message", + FinishReason = ChatFinishReason.ToolCalls, + }) + : YieldUpdateAsync(new ChatResponseUpdate(ChatRole.Assistant, "The tool failed.") + { + MessageId = "final_message", + FinishReason = ChatFinishReason.Stop, + })); + + var agent = new ChatClientAgent(chatClient.Object, new ChatClientAgentOptions + { + ChatOptions = new ChatOptions { Tools = [tool] }, + }); + var services = new ServiceCollection(); + services.AddSingleton(new InMemoryAgentSessionStore()); + services.AddSingleton(agent); + services.AddSingleton(new FakeHostedSessionIsolationKeyProvider()); + using var provider = services.BuildServiceProvider(); + var handler = new AgentFrameworkResponseHandler(provider, NullLogger.Instance); + var context = new Mock(MockBehavior.Loose, "resp_" + new string('0', 46)) { CallBase = true }; + context.Setup(x => x.GetHistoryAsync(It.IsAny())).ReturnsAsync(Array.Empty()); + context.Setup(x => x.GetInputItemsAsync(It.IsAny(), It.IsAny())).ReturnsAsync(Array.Empty()); + var events = new List(); + await foreach (var evt in handler.CreateAsync(new CreateResponse { Model = "test" }, context.Object, cancellationToken)) + { + events.Add(evt); + } + + return (events, modelCalls); + } + + private static async IAsyncEnumerable YieldUpdateAsync(ChatResponseUpdate update) + { + await Task.Yield(); + yield return update; + } + + private sealed class ConsentPeer(string consentLink, bool includeDetail = true) : HttpMessageHandler + { + public int ToolCalls { get; private set; } + + public object? ErrorDetail { get; set; } = consentLink; + + public Func? BeforeToolCall { get; init; } + + protected override async Task SendAsync( + HttpRequestMessage request, + CancellationToken cancellationToken) + { + using var body = JsonDocument.Parse( + await request.Content!.ReadAsStringAsync(cancellationToken)); + var root = body.RootElement; + if (!root.TryGetProperty("id", out var id)) + { + return new HttpResponseMessage(HttpStatusCode.Accepted); + } + + string json; + switch (root.GetProperty("method").GetString()) + { + case "server/discover": + json = JsonSerializer.Serialize(new + { + jsonrpc = "2.0", + id, + error = new { code = -32601, message = "Method not found" }, + }); + break; + case "initialize": + json = JsonSerializer.Serialize(new + { + jsonrpc = "2.0", + id, + result = new + { + protocolVersion = "2025-03-26", + capabilities = new { tools = new { } }, + serverInfo = new { name = "consent-test", version = "1.0" }, + }, + }); + break; + case "tools/call": + this.ToolCalls++; + if (this.BeforeToolCall is { } beforeToolCall) + { + await beforeToolCall(cancellationToken); + } + + var data = new Dictionary { ["grpcStatus"] = "FailedPrecondition" }; + if (includeDetail) + { + data["detail"] = this.ErrorDetail; + } + + json = JsonSerializer.Serialize(new + { + jsonrpc = "2.0", + id, + error = new + { + code = -32006, + message = consentLink, + data, + }, + }); + break; + default: + throw new InvalidOperationException("Unexpected MCP method: " + root.GetProperty("method")); + } + + return new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent(json, Encoding.UTF8, "application/json"), + }; + } + } +}