From b31fdddfee245206700c89157ee4353a853d0d7d Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Mon, 19 Oct 2026 08:43:53 +0000 Subject: [PATCH 1/6] .NET: Honor hosted agent lifetimes and add request-time agent selection for A2A hosting Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- dotnet/eng/MSBuild/Shared.props | 3 + .../A2AEndpointRouteBuilderExtensions.cs | 121 ++++++ ...ft.Agents.AI.Hosting.A2A.AspNetCore.csproj | 4 +- .../SelectedAgentA2ARequestHandler.cs | 160 +++++++ .../A2AAgentHandler.cs | 28 +- .../A2AServerServiceCollectionExtensions.cs | 113 +++-- .../Microsoft.Agents.AI.Hosting.A2A.csproj | 2 + .../ScopedA2AAgentHandler.cs | 84 ++++ .../HostedAgents/HostedAgentResolution.cs | 52 +++ .../A2ARequestScopedAgentTests.cs | 396 ++++++++++++++++++ ...AServerServiceCollectionExtensionsTests.cs | 4 +- 11 files changed, 935 insertions(+), 32 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ScopedA2AAgentHandler.cs create mode 100644 dotnet/src/Shared/HostedAgents/HostedAgentResolution.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs diff --git a/dotnet/eng/MSBuild/Shared.props b/dotnet/eng/MSBuild/Shared.props index 10db0a47ef4..3c8f9670284 100644 --- a/dotnet/eng/MSBuild/Shared.props +++ b/dotnet/eng/MSBuild/Shared.props @@ -38,6 +38,9 @@ + + + diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs index 6dad579786e..975a7d6e1dc 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs @@ -6,6 +6,8 @@ using A2A.AspNetCore; using Microsoft.Agents.AI; using Microsoft.Agents.AI.Hosting; +using Microsoft.Agents.AI.Hosting.A2A; +using Microsoft.AspNetCore.Http; using Microsoft.AspNetCore.Routing; using Microsoft.Extensions.DependencyInjection; using Microsoft.Shared.DiagnosticIds; @@ -93,6 +95,56 @@ public static IEndpointConventionBuilder MapA2AHttpJson(this IEndpointRouteBuild return endpoint; } + /// + /// Maps A2A HTTP+JSON endpoints that select the agent to serve for each request. + /// + /// The to add the A2A endpoints to. + /// The route path prefix for A2A endpoints. + /// + /// A delegate invoked once per request that returns the agent to serve, or when + /// no agent matches the request. The delegate can use route values, claims, headers, and + /// to select or resolve the agent. + /// + /// An optional callback to configure the applied to every selected agent. + /// An for further endpoint configuration. + /// + /// + /// When returns , the endpoint responds with + /// 404 Not Found before any session or task state is loaded. No AddA2AServer registration is required. + /// + /// + /// The selected agent's is its stable logical identity. It keys the agent's + /// and registrations, partitions persisted sessions, and + /// identifies the A2A task state shared by the requests that select it through this endpoint, so conversations + /// and tasks continue across requests even when the selector returns a new agent instance each time. The selected + /// agent must therefore have a non-empty name. The session store is resolved from + /// , so the run cannot outlive the request: when the configured + /// returns a task, non-streaming requests with ReturnImmediately enabled are + /// rejected with . + /// + /// + /// Selection is not an authorization boundary. See + /// for endpoint authorization and caller isolation requirements. + /// + /// + public static IEndpointConventionBuilder MapA2AHttpJson( + this IEndpointRouteBuilder endpoints, + string path, + Func agentSelector, + Action? configureOptions = null) + { + ArgumentNullException.ThrowIfNull(endpoints); + ArgumentException.ThrowIfNullOrWhiteSpace(path); + ArgumentNullException.ThrowIfNull(agentSelector); + + var requestHandler = CreateSelectedAgentRequestHandler(endpoints, configureOptions); + + IEndpointConventionBuilder endpoint = endpoints.MapHttpA2A(requestHandler, path); + AddAgentSelection(endpoint, requestHandler, agentSelector); + MarkFeatureUsed(); + return endpoint; + } + /// /// Maps A2A JSON-RPC endpoints for the specified agent to the given path. /// An for the agent must be registered first by calling @@ -163,6 +215,75 @@ public static IEndpointConventionBuilder MapA2AJsonRpc(this IEndpointRouteBuilde return endpoint; } + /// + /// Maps A2A JSON-RPC endpoints that select the agent to serve for each request. + /// + /// The to add the A2A endpoints to. + /// The route path prefix for A2A endpoints. + /// + /// A delegate invoked once per request that returns the agent to serve, or when + /// no agent matches the request. The delegate can use route values, claims, headers, and + /// to select or resolve the agent. + /// + /// An optional callback to configure the applied to every selected agent. + /// An for further endpoint configuration. + /// + /// See + /// for selection, identity, and lifetime behavior, and + /// for endpoint authorization and caller isolation requirements, which also apply to this JSON-RPC binding. + /// + public static IEndpointConventionBuilder MapA2AJsonRpc( + this IEndpointRouteBuilder endpoints, + string path, + Func agentSelector, + Action? configureOptions = null) + { + ArgumentNullException.ThrowIfNull(endpoints); + ArgumentException.ThrowIfNullOrWhiteSpace(path); + ArgumentNullException.ThrowIfNull(agentSelector); + + var requestHandler = CreateSelectedAgentRequestHandler(endpoints, configureOptions); + + IEndpointConventionBuilder endpoint = endpoints.MapA2A(requestHandler, path); + AddAgentSelection(endpoint, requestHandler, agentSelector); + MarkFeatureUsed(); + return endpoint; + } + + private static SelectedAgentA2ARequestHandler CreateSelectedAgentRequestHandler( + IEndpointRouteBuilder endpoints, + Action? configureOptions) + { + A2AServerRegistrationOptions? options = null; + if (configureOptions is not null) + { + options = new A2AServerRegistrationOptions(); + configureOptions(options); + } + + return new SelectedAgentA2ARequestHandler(endpoints.ServiceProvider, options); + } + + private static void AddAgentSelection( + IEndpointConventionBuilder endpoint, + SelectedAgentA2ARequestHandler requestHandler, + Func agentSelector) + { + // The A2A endpoints call the request handler without request context, so select the agent for the + // asynchronous flow of the request before the endpoint runs. + endpoint.AddEndpointFilter(async (invocationContext, next) => + { + HttpContext httpContext = invocationContext.HttpContext; + if (agentSelector(httpContext) is not { } agent) + { + return Results.NotFound(); + } + + requestHandler.Select(agent, httpContext.RequestServices); + return await next(invocationContext).ConfigureAwait(false); + }); + } + private static void MarkFeatureUsed() { #pragma warning disable MAAI001 diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/Microsoft.Agents.AI.Hosting.A2A.AspNetCore.csproj b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/Microsoft.Agents.AI.Hosting.A2A.AspNetCore.csproj index b91ea40baa7..0333f13f23e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/Microsoft.Agents.AI.Hosting.A2A.AspNetCore.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/Microsoft.Agents.AI.Hosting.A2A.AspNetCore.csproj @@ -12,8 +12,8 @@ - true - true + true diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs new file mode 100644 index 00000000000..d29d81ec256 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs @@ -0,0 +1,160 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using A2A; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Shared.DiagnosticIds; + +namespace Microsoft.Agents.AI.Hosting.A2A; + +/// +/// An that dispatches each A2A operation to the agent selected for the current request. +/// +/// +/// The A2A endpoints invoke this handler without request context, so the endpoint filter records the selected agent +/// for the asynchronous flow of the request before the endpoint runs. Each selected agent name owns an +/// whose task state is shared by all requests that select that name. +/// +[Experimental(DiagnosticIds.Experiments.AIResponseContinuations)] +internal sealed class SelectedAgentA2ARequestHandler : IA2ARequestHandler +{ + private readonly AsyncLocal _selection = new(); + private readonly ConcurrentDictionary _servers = new(StringComparer.Ordinal); + private readonly IServiceProvider _applicationServices; + private readonly A2AServerRegistrationOptions? _options; + + /// + /// Initializes a new instance of the class. + /// + /// The application service provider that owns the shared task state. + /// The optional registration options applied to every selected agent. + public SelectedAgentA2ARequestHandler(IServiceProvider applicationServices, A2AServerRegistrationOptions? options) + { + this._applicationServices = applicationServices; + this._options = options; + } + + /// + /// Makes the target of the A2A operations dispatched in the current asynchronous flow. + /// + /// The agent selected for the request. + /// The request services that own the agent's session services. + public void Select(AIAgent agent, IServiceProvider requestServices) + { + string agentName = HostedAgentResolution.GetSelectedAgentName(agent); + SelectedAgentServer server = this._servers.GetOrAdd( + agentName, + static (name, state) => new SelectedAgentServer(state._applicationServices, name, state._options), + this); + + this._selection.Value = new Selection(agent, agentName, requestServices, server, this._options?.AgentRunMode ?? AgentRunMode.ReturnMessage); + } + + /// + public Task SendMessageAsync(SendMessageRequest request, CancellationToken cancellationToken = default) + { + Selection selection = this.GetSelection(); + selection.BindAgent(request.Message); + return selection.A2AServer.SendMessageAsync(request, cancellationToken); + } + + /// + public IAsyncEnumerable SendStreamingMessageAsync(SendMessageRequest request, CancellationToken cancellationToken = default) + { + Selection selection = this.GetSelection(); + selection.BindAgent(request.Message); + return selection.A2AServer.SendStreamingMessageAsync(request, cancellationToken); + } + + /// + public Task GetTaskAsync(GetTaskRequest request, CancellationToken cancellationToken = default) + => this.GetSelection().A2AServer.GetTaskAsync(request, cancellationToken); + + /// + public Task ListTasksAsync(ListTasksRequest request, CancellationToken cancellationToken = default) + => this.GetSelection().A2AServer.ListTasksAsync(request, cancellationToken); + + /// + public Task CancelTaskAsync(CancelTaskRequest request, CancellationToken cancellationToken = default) + => this.GetSelection().A2AServer.CancelTaskAsync(request, cancellationToken); + + /// + public IAsyncEnumerable SubscribeToTaskAsync(SubscribeToTaskRequest request, CancellationToken cancellationToken = default) + => this.GetSelection().A2AServer.SubscribeToTaskAsync(request, cancellationToken); + + /// + public Task CreateTaskPushNotificationConfigAsync(TaskPushNotificationConfig config, CancellationToken cancellationToken = default) + => this.GetSelection().A2AServer.CreateTaskPushNotificationConfigAsync(config, cancellationToken); + + /// + public Task GetTaskPushNotificationConfigAsync(GetTaskPushNotificationConfigRequest request, CancellationToken cancellationToken = default) + => this.GetSelection().A2AServer.GetTaskPushNotificationConfigAsync(request, cancellationToken); + + /// + public Task ListTaskPushNotificationConfigsAsync(ListTaskPushNotificationConfigsRequest request, CancellationToken cancellationToken = default) + => this.GetSelection().A2AServer.ListTaskPushNotificationConfigsAsync(request, cancellationToken); + + /// + public Task DeleteTaskPushNotificationConfigAsync(DeleteTaskPushNotificationConfigRequest request, CancellationToken cancellationToken = default) + => this.GetSelection().A2AServer.DeleteTaskPushNotificationConfigAsync(request, cancellationToken); + + /// + public Task GetExtendedAgentCardAsync(GetExtendedAgentCardRequest request, CancellationToken cancellationToken = default) + => this.GetSelection().A2AServer.GetExtendedAgentCardAsync(request, cancellationToken); + + private Selection GetSelection() => + this._selection.Value ?? throw new InvalidOperationException("No agent was selected for the current A2A request."); + + private sealed class Selection( + AIAgent agent, + string agentName, + IServiceProvider requestServices, + SelectedAgentServer server, + AgentRunMode runMode) + { + public A2AServer A2AServer => server.A2AServer; + + public void BindAgent(Message message) + { + AIHostAgent hostAgent = A2AServerServiceCollectionExtensions.CreateHostAgent(requestServices, agent, agentName, sessionStorageIdentity: agentName); + server.Bind(message, new A2AAgentHandler(hostAgent, runMode, allowReturnImmediately: false)); + } + } + + /// + /// Owns the task state of one selected agent name and runs each message with the agent selected by its request. + /// + private sealed class SelectedAgentServer : IAgentHandler + { + // A2AServer passes the incoming message to the handler, including when it runs the handler in the background. + private readonly ConditionalWeakTable _handlers = new(); + + public SelectedAgentServer(IServiceProvider applicationServices, string agentName, A2AServerRegistrationOptions? options) + { + this.A2AServer = A2AServerServiceCollectionExtensions.CreateA2AServer(applicationServices, agentName, this, options); + } + + public A2AServer A2AServer { get; } + + public void Bind(Message message, IAgentHandler handler) => this._handlers.AddOrUpdate(message, handler); + + public Task ExecuteAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) + { + if (context.Message is null || !this._handlers.TryGetValue(context.Message, out IAgentHandler? handler)) + { + throw new InvalidOperationException("The A2A message is not associated with the agent selected for its request."); + } + + return handler.ExecuteAsync(context, eventQueue, cancellationToken); + } + + public Task CancelAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) + => A2AAgentHandler.CancelTaskAsync(context, eventQueue, cancellationToken); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs index 8baf274020d..b3ec933be67 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs @@ -28,21 +28,28 @@ internal sealed class A2AAgentHandler : IAgentHandler private readonly AIHostAgent _hostAgent; private readonly AgentRunMode _runMode; + private readonly bool _allowReturnImmediately; /// /// Initializes a new instance of the class. /// /// The hosted agent that provides the execution logic. /// Controls which A2A artifact the agent response is returned as. + /// + /// to reject non-streaming requests with ReturnImmediately = true when + /// returns a task, because the run would continue after the request completes. + /// public A2AAgentHandler( AIHostAgent hostAgent, - AgentRunMode runMode) + AgentRunMode runMode, + bool allowReturnImmediately = true) { ArgumentNullException.ThrowIfNull(hostAgent); ArgumentNullException.ThrowIfNull(runMode); this._hostAgent = hostAgent; this._runMode = runMode; + this._allowReturnImmediately = allowReturnImmediately; } /// @@ -67,7 +74,17 @@ public Task ExecuteAsync(RequestContext context, AgentEventQueue eventQueue, Can } /// - public async Task CancelAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) + public Task CancelAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) + => CancelTaskAsync(context, eventQueue, cancellationToken); + + /// + /// Transitions the task identified by to the canceled state. + /// + /// The request context of the task to cancel. + /// The queue the cancellation event is written to. + /// A to cancel the operation. + /// A task that represents the asynchronous operation. + internal static async Task CancelTaskAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) { var taskUpdater = new TaskUpdater(eventQueue, context.TaskId, context.ContextId); await taskUpdater.CancelAsync(cancellationToken: cancellationToken).ConfigureAwait(false); @@ -140,6 +157,13 @@ private async Task HandleNewMessageAsync(RequestContext context, AgentEventQueue var decisionContext = new A2ARunDecisionContext(context); var returnTask = await this._runMode.ShouldReturnTaskAsync(decisionContext, cancellationToken).ConfigureAwait(false); + if (returnTask && !aggregateTaskUpdates && !context.StreamingResponse && !this._allowReturnImmediately) + { + throw new A2AException( + "ReturnImmediately is not supported when the agent is selected for each request, because the selected agent and its services are bound to the HTTP request. Send the message with ReturnImmediately disabled.", + A2AErrorCode.UnsupportedOperation); + } + var options = CreateRunOptions(context); var updates = this._hostAgent.RunStreamingAsync(chatMessages, session, options, cancellationToken); diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs index b3fc808b213..c466e394de4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs @@ -30,6 +30,11 @@ public static class A2AServerServiceCollectionExtensions /// The for chaining. /// /// + /// Singleton agent registrations are resolved once when the server is created. Scoped and transient registrations + /// are resolved for each A2A operation as described in + /// . + /// + /// /// Trust model. The A2A contextId and taskId arrive /// from the wire and are treated as chain-resume identifiers — not as /// authorization tokens. accepts an explicit user partition, @@ -55,7 +60,7 @@ public static IHostedAgentBuilder AddA2AServer(this IHostedAgentBuilder agentBui { ArgumentNullException.ThrowIfNull(agentBuilder); - agentBuilder.ServiceCollection.AddA2AServer(agentBuilder.Name, configureOptions); + AddA2AServer(agentBuilder.ServiceCollection, agentBuilder.Name, agentBuilder.Lifetime, configureOptions); return agentBuilder; } @@ -121,13 +126,25 @@ public static IHostApplicationBuilder AddA2AServer(this IHostApplicationBuilder /// An optional callback to configure . /// The for chaining. /// + /// + /// The server and its task store are singletons shared by all requests. The keyed , its keyed + /// , and an optional keyed are resolved from a new + /// service scope for each A2A operation, so the configured DI lifetimes of these registrations are honored and + /// operations that continue in the background keep their services alive until they complete. Persisted sessions + /// are partitioned by , which remains stable across scoped or transient agent instances. + /// + /// /// See the trust-model remarks on /// for guidance on multi-user hosts (the wire contextId and taskId /// are chain-resume identifiers, not authorization tokens; multi-user hosts must /// supply a trusted user partition via UseClaimsBasedAgentIsolation(...) or /// a custom ). + /// /// public static IServiceCollection AddA2AServer(this IServiceCollection services, string agentName, Action? configureOptions = null) + => AddA2AServer(services, agentName, agentLifetime: null, configureOptions); + + private static IServiceCollection AddA2AServer(IServiceCollection services, string agentName, ServiceLifetime? agentLifetime, Action? configureOptions) { ArgumentNullException.ThrowIfNull(services); ArgumentException.ThrowIfNullOrWhiteSpace(agentName); @@ -141,8 +158,19 @@ public static IServiceCollection AddA2AServer(this IServiceCollection services, services.AddKeyedSingleton(agentName, (sp, _) => { - var agent = sp.GetRequiredKeyedService(agentName); - return CreateA2AServer(sp, agent, options); + // Custom handlers replace the agent, so the agent registration is only required without one. + bool hasCustomHandler = sp.GetService()?.IsKeyedService(typeof(IAgentHandler), agentName) ?? false; + AIAgent? agent = null; + if (!hasCustomHandler) + { + HostedAgentResolution.EnsureAgentRegistered(sp, agentName); + if (agentLifetime == ServiceLifetime.Singleton) + { + agent = sp.GetRequiredKeyedService(agentName); + } + } + + return CreateA2AServer(sp, agentName, agent, sessionStorageIdentity: agentName, options); }); return services; @@ -178,39 +206,45 @@ public static IServiceCollection AddA2AServer(this IServiceCollection services, configureOptions(options); } - services.AddKeyedSingleton(agent.Name, (sp, _) => CreateA2AServer(sp, agent, options)); + services.AddKeyedSingleton(agent.Name, (sp, _) => CreateA2AServer(sp, agent.Name, agent, sessionStorageIdentity: null, options)); return services; } - private static A2AServer CreateA2AServer(IServiceProvider serviceProvider, AIAgent agent, A2AServerRegistrationOptions? options) + private static A2AServer CreateA2AServer( + IServiceProvider serviceProvider, + string agentName, + AIAgent? agent, + string? sessionStorageIdentity, + A2AServerRegistrationOptions? options) { - var isolationKeyProvider = serviceProvider.GetService(); - - var agentHandler = serviceProvider.GetKeyedService(agent.Name); - if (agentHandler is null) - { - var agentSessionStore = serviceProvider.GetKeyedService(agent.Name); - var runMode = options?.AgentRunMode ?? AgentRunMode.ReturnMessage; - - // Ensure that we have an IsolationKeyScopedAgentSessionStore registered. - if (agentSessionStore?.GetService() is null) - { - agentSessionStore ??= new NoopAgentSessionStore(); - agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = isolationKeyProvider != null }); - } + // The server owns task state shared across requests, while the agent and its session services are + // resolved for each operation so that scoped and transient registrations keep their lifetimes. + var agentHandler = new ScopedA2AAgentHandler( + serviceProvider.GetRequiredService(), + agentName, + agent, + sessionStorageIdentity, + options?.AgentRunMode ?? AgentRunMode.ReturnMessage); - var hostAgent = new AIHostAgent( - innerAgent: agent, - sessionStore: agentSessionStore); - - agentHandler = new A2AAgentHandler(hostAgent, runMode); - } + return CreateA2AServer(serviceProvider, agentName, agentHandler, options); + } + /// + /// Creates an for the specified agent name with its keyed or default task store. + /// + /// The application service provider. + /// The name that keys the agent's hosted services. + /// The handler that executes A2A operations. + /// The optional registration options. + /// The created server. + internal static A2AServer CreateA2AServer(IServiceProvider serviceProvider, string agentName, IAgentHandler agentHandler, A2AServerRegistrationOptions? options) + { + var isolationKeyProvider = serviceProvider.GetService(); var loggerFactory = serviceProvider.GetService() ?? NullLoggerFactory.Instance; - ITaskStore taskStore = serviceProvider.GetKeyedService(agent.Name) ?? new InMemoryTaskStore(); + ITaskStore taskStore = serviceProvider.GetKeyedService(agentName) ?? new InMemoryTaskStore(); - // Wrap the task store with isolation key scoping, same as the session store above. + // Wrap the task store with isolation key scoping, same as the session store. if (taskStore is not IsolationKeyScopedTaskStore) { taskStore = new IsolationKeyScopedTaskStore(taskStore, isolationKeyProvider, strict: isolationKeyProvider != null); @@ -223,4 +257,29 @@ private static A2AServer CreateA2AServer(IServiceProvider serviceProvider, AIAge loggerFactory.CreateLogger(), options?.ServerOptions); } + + /// + /// Creates the hosting wrapper that persists the sessions of an agent served over A2A. + /// + /// The service provider that owns the agent's session services. + /// The agent to host. + /// The name that keys the agent's session store. + /// + /// The stable logical identity that partitions persisted sessions, or to use the agent instance identity. + /// + /// The hosting wrapper for . + internal static AIHostAgent CreateHostAgent(IServiceProvider services, AIAgent agent, string agentName, string? sessionStorageIdentity) + { + var isolationKeyProvider = services.GetService(); + var agentSessionStore = services.GetKeyedService(agentName); + + // Ensure that we have an IsolationKeyScopedAgentSessionStore registered. + if (agentSessionStore?.GetService() is null) + { + agentSessionStore ??= new NoopAgentSessionStore(); + agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = isolationKeyProvider != null }); + } + + return new AIHostAgent(agent, agentSessionStore, sessionStorageIdentity); + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/Microsoft.Agents.AI.Hosting.A2A.csproj b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/Microsoft.Agents.AI.Hosting.A2A.csproj index 2db2d6b6f6a..627d123bbd5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/Microsoft.Agents.AI.Hosting.A2A.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/Microsoft.Agents.AI.Hosting.A2A.csproj @@ -13,6 +13,7 @@ true true true + true @@ -28,6 +29,7 @@ + diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ScopedA2AAgentHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ScopedA2AAgentHandler.cs new file mode 100644 index 00000000000..49352b2f8bf --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ScopedA2AAgentHandler.cs @@ -0,0 +1,84 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Diagnostics.CodeAnalysis; +using System.Threading; +using System.Threading.Tasks; +using A2A; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Shared.DiagnosticIds; + +namespace Microsoft.Agents.AI.Hosting.A2A; + +/// +/// An that resolves the agent and its session services from a new service scope +/// for each A2A operation, so the configured DI lifetimes of the hosted registration are honored. +/// +/// +/// A2A operations can continue in the background after the originating HTTP request completes, so each +/// operation owns its scope instead of borrowing the request scope. +/// +[Experimental(DiagnosticIds.Experiments.AIResponseContinuations)] +internal sealed class ScopedA2AAgentHandler : IAgentHandler +{ + private readonly IServiceScopeFactory _scopeFactory; + private readonly string _agentName; + private readonly AIAgent? _agent; + private readonly string? _sessionStorageIdentity; + private readonly AgentRunMode _runMode; + + /// + /// Initializes a new instance of the class. + /// + /// The factory used to create a scope for each operation. + /// The name that keys the agent registration and its hosted services. + /// A fixed agent instance, or to resolve the keyed registration in each scope. + /// + /// The stable logical identity that partitions persisted sessions, or to use the agent instance identity. + /// + /// Controls which A2A artifact the agent response is returned as. + public ScopedA2AAgentHandler(IServiceScopeFactory scopeFactory, string agentName, AIAgent? agent, string? sessionStorageIdentity, AgentRunMode runMode) + { + this._scopeFactory = scopeFactory; + this._agentName = agentName; + this._agent = agent; + this._sessionStorageIdentity = sessionStorageIdentity; + this._runMode = runMode; + } + + /// + public async Task ExecuteAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) + { + AsyncServiceScope scope = this._scopeFactory.CreateAsyncScope(); + await using (scope.ConfigureAwait(false)) + { + await this.CreateHandler(scope.ServiceProvider).ExecuteAsync(context, eventQueue, cancellationToken).ConfigureAwait(false); + } + } + + /// + public async Task CancelAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) + { + AsyncServiceScope scope = this._scopeFactory.CreateAsyncScope(); + await using (scope.ConfigureAwait(false)) + { + // Canceling a task does not involve the agent, so only a custom handler needs to be resolved. + Task cancellation = scope.ServiceProvider.GetKeyedService(this._agentName) is { } agentHandler + ? agentHandler.CancelAsync(context, eventQueue, cancellationToken) + : A2AAgentHandler.CancelTaskAsync(context, eventQueue, cancellationToken); + await cancellation.ConfigureAwait(false); + } + } + + private IAgentHandler CreateHandler(IServiceProvider services) + { + if (services.GetKeyedService(this._agentName) is { } agentHandler) + { + return agentHandler; + } + + AIAgent agent = this._agent ?? services.GetRequiredKeyedService(this._agentName); + AIHostAgent hostAgent = A2AServerServiceCollectionExtensions.CreateHostAgent(services, agent, this._agentName, this._sessionStorageIdentity); + return new A2AAgentHandler(hostAgent, this._runMode); + } +} diff --git a/dotnet/src/Shared/HostedAgents/HostedAgentResolution.cs b/dotnet/src/Shared/HostedAgents/HostedAgentResolution.cs new file mode 100644 index 00000000000..42d9638d61a --- /dev/null +++ b/dotnet/src/Shared/HostedAgents/HostedAgentResolution.cs @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using Microsoft.Extensions.DependencyInjection; + +namespace Microsoft.Agents.AI.Hosting; + +/// +/// Provides validation shared by protocol hosting integrations that resolve agents for each request. +/// +internal static class HostedAgentResolution +{ + /// + /// Verifies that a keyed registration exists without resolving it. + /// + /// The application service provider. + /// The name of the keyed agent registration. + /// No agent is registered with . + /// + /// Resolving the agent from the application service provider would capture scoped or transient registrations, + /// so the check is skipped when the provider cannot report keyed registrations. + /// + public static void EnsureAgentRegistered(IServiceProvider services, string agentName) + { + if (services.GetService() is { } isKeyedService && + !isKeyedService.IsKeyedService(typeof(AIAgent), agentName)) + { + throw new InvalidOperationException( + $"No {nameof(AIAgent)} is registered with the name '{agentName}'. " + + $"Register it with AddAIAgent(\"{agentName}\", ...) before mapping the endpoint."); + } + } + + /// + /// Gets the stable logical identity of an agent returned by a request-time agent selector. + /// + /// The selected agent. + /// The agent name, which identifies the agent's hosted services and persisted sessions. + /// The agent does not have a name. + public static string GetSelectedAgentName(AIAgent agent) + { + string? agentName = agent.Name; + if (string.IsNullOrWhiteSpace(agentName)) + { + throw new InvalidOperationException( + $"The agent selector returned an {nameof(AIAgent)} without a {nameof(AIAgent.Name)}. " + + "Selected agents must have a stable name, which identifies their hosted services and persisted sessions across requests."); + } + + return agentName!; + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs new file mode 100644 index 00000000000..05afb712c60 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs @@ -0,0 +1,396 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Globalization; +using System.Linq; +using System.Net; +using System.Net.Http; +using System.Runtime.CompilerServices; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using A2A; +using Microsoft.AspNetCore.Builder; +using Microsoft.AspNetCore.Http; +using Microsoft.AspNetCore.Routing; +using Microsoft.AspNetCore.TestHost; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Hosting; +using Moq; + +namespace Microsoft.Agents.AI.Hosting.A2A.UnitTests; + +/// +/// Verifies that A2A hosting honors agent registration lifetimes and selects agents for each request. +/// +public sealed class A2ARequestScopedAgentTests +{ + [Fact] + public async Task AddA2AServer_WithScopedAgent_ResolvesAgentForEachOperationAsync() + { + // Arrange + int createdAgents = 0; + ServiceCollection services = new(); + services.AddLogging(); + services.AddScoped(); + services.AddAIAgent( + "assistant", + (sp, name) => + { + _ = sp.GetRequiredService(); + Interlocked.Increment(ref createdAgents); + return new TurnCountingAgent(name); + }, + ServiceLifetime.Scoped) + .WithInMemorySessionStore(withIsolation: false); + services.AddA2AServer("assistant"); + await using ServiceProvider provider = services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }); + A2AServer server = provider.GetRequiredKeyedService("assistant"); + + // Act + SendMessageResponse firstResponse = await server.SendMessageAsync(CreateRequest("ctx-1")); + SendMessageResponse secondResponse = await server.SendMessageAsync(CreateRequest("ctx-1")); + + // Assert + Assert.Equal(2, createdAgents); + Assert.Equal("assistant turn 1", GetText(firstResponse)); + Assert.Equal("assistant turn 2", GetText(secondResponse)); + } + + [Fact] + public async Task AddA2AServer_WithSingletonHostedAgentBuilder_ResolvesAgentOnceAsync() + { + // Arrange + int createdAgents = 0; + ServiceCollection services = new(); + services.AddLogging(); + services.AddAIAgent( + "assistant", + (_, name) => + { + Interlocked.Increment(ref createdAgents); + return new TurnCountingAgent(name); + }) + .WithInMemorySessionStore(withIsolation: false) + .AddA2AServer(); + await using ServiceProvider provider = services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }); + A2AServer server = provider.GetRequiredKeyedService("assistant"); + int createdAtServerCreation = createdAgents; + + // Act + await server.SendMessageAsync(CreateRequest("ctx-1")); + SendMessageResponse secondResponse = await server.SendMessageAsync(CreateRequest("ctx-1")); + + // Assert + Assert.Equal(1, createdAtServerCreation); + Assert.Equal(1, createdAgents); + Assert.Equal("assistant turn 2", GetText(secondResponse)); + } + + [Fact] + public async Task AddA2AServer_WithUnregisteredAgent_ThrowsWhenServerIsResolvedAsync() + { + // Arrange + ServiceCollection services = new(); + services.AddLogging(); + services.AddA2AServer("missing"); + await using ServiceProvider provider = services.BuildServiceProvider(); + + // Act & Assert + InvalidOperationException exception = Assert.Throws(() => provider.GetRequiredKeyedService("missing")); + Assert.Contains("missing", exception.Message, StringComparison.Ordinal); + } + + [Fact] + public async Task AddA2AServer_WithCustomAgentHandler_DoesNotRequireAgentRegistrationAsync() + { + // Arrange + ServiceCollection services = new(); + services.AddLogging(); + services.AddKeyedSingleton("custom", Mock.Of()); + services.AddA2AServer("custom"); + await using ServiceProvider provider = services.BuildServiceProvider(); + + // Act + A2AServer? server = provider.GetKeyedService("custom"); + + // Assert + Assert.NotNull(server); + } + + [Fact] + public async Task MapA2AHttpJson_WithAgentSelector_SelectsAgentFromRouteValueAsync() + { + // Arrange + await using WebApplication app = await StartAsync( + services => + { + services.AddAIAgent("alpha", (_, name) => new TurnCountingAgent(name), ServiceLifetime.Scoped) + .WithInMemorySessionStore(withIsolation: false); + services.AddAIAgent("beta", (_, name) => new TurnCountingAgent(name), ServiceLifetime.Transient) + .WithInMemorySessionStore(withIsolation: false); + }, + endpoints => endpoints.MapA2AHttpJson("/agents/{agentId}/a2a", SelectAgentFromRoute)); + using HttpClient httpClient = app.GetTestClient(); + using A2AHttpJsonClient alphaClient = new(new Uri(httpClient.BaseAddress!, "/agents/alpha/a2a"), httpClient); + using A2AHttpJsonClient betaClient = new(new Uri(httpClient.BaseAddress!, "/agents/beta/a2a"), httpClient); + + // Act + SendMessageResponse firstAlphaResponse = await alphaClient.SendMessageAsync(CreateRequest("ctx-1")); + SendMessageResponse secondAlphaResponse = await alphaClient.SendMessageAsync(CreateRequest("ctx-1")); + SendMessageResponse betaResponse = await betaClient.SendMessageAsync(CreateRequest("ctx-1")); + + // Assert + Assert.Equal("alpha turn 1", GetText(firstAlphaResponse)); + Assert.Equal("alpha turn 2", GetText(secondAlphaResponse)); + Assert.Equal("beta turn 1", GetText(betaResponse)); + } + + [Fact] + public async Task MapA2AJsonRpc_WithAgentSelector_PersistsSessionsAcrossNewAgentInstancesAsync() + { + // Arrange + await using WebApplication app = await StartAsync( + services => services.AddKeyedSingleton("dynamic", new InMemoryAgentSessionStore()), + endpoints => endpoints.MapA2AJsonRpc("/a2a", _ => new TurnCountingAgent("dynamic"))); + using HttpClient httpClient = app.GetTestClient(); + using A2AClient client = new(new Uri(httpClient.BaseAddress!, "/a2a"), httpClient); + + // Act + SendMessageResponse firstResponse = await client.SendMessageAsync(CreateRequest("ctx-1")); + SendMessageResponse secondResponse = await client.SendMessageAsync(CreateRequest("ctx-1")); + SendMessageResponse otherContextResponse = await client.SendMessageAsync(CreateRequest("ctx-2")); + + // Assert + Assert.Equal("dynamic turn 1", GetText(firstResponse)); + Assert.Equal("dynamic turn 2", GetText(secondResponse)); + Assert.Equal("dynamic turn 1", GetText(otherContextResponse)); + } + + [Fact] + public async Task MapA2AHttpJson_WithAgentSelector_StreamsSelectedAgentAsync() + { + // Arrange + await using WebApplication app = await StartAsync( + _ => { }, + endpoints => endpoints.MapA2AHttpJson("/a2a", _ => new TurnCountingAgent("streaming"))); + using HttpClient httpClient = app.GetTestClient(); + using A2AHttpJsonClient client = new(new Uri(httpClient.BaseAddress!, "/a2a"), httpClient); + + // Act + List responses = []; + await foreach (StreamResponse response in client.SendStreamingMessageAsync(CreateRequest("ctx-1"))) + { + responses.Add(response); + } + + // Assert + Message message = Assert.IsType(Assert.Single(responses).Message); + Assert.Equal("streaming turn 1", string.Concat(message.Parts.Select(part => part.Text))); + } + + [Fact] + public async Task MapA2AHttpJson_WithAgentSelector_SharesTaskStateAcrossRequestsAsync() + { + // Arrange + await using WebApplication app = await StartAsync( + _ => { }, + endpoints => endpoints.MapA2AHttpJson( + "/a2a", + _ => new TurnCountingAgent("tasks"), + options => options.AgentRunMode = AgentRunMode.ReturnTask)); + using HttpClient httpClient = app.GetTestClient(); + using A2AHttpJsonClient client = new(new Uri(httpClient.BaseAddress!, "/a2a"), httpClient); + + // Act + SendMessageResponse response = await client.SendMessageAsync(CreateRequest("ctx-1")); + AgentTask task = await client.GetTaskAsync(new GetTaskRequest { Id = response.Task!.Id }); + + // Assert + Assert.Equal(response.Task.Id, task.Id); + Assert.Equal(TaskState.Completed, task.Status.State); + } + + [Fact] + public async Task MapA2AJsonRpc_WithAgentSelectorReturningTask_RejectsReturnImmediatelyAsync() + { + // Arrange + await using WebApplication app = await StartAsync( + _ => { }, + endpoints => endpoints.MapA2AJsonRpc( + "/a2a", + _ => new TurnCountingAgent("tasks"), + options => options.AgentRunMode = AgentRunMode.ReturnTask)); + using HttpClient httpClient = app.GetTestClient(); + using A2AClient client = new(new Uri(httpClient.BaseAddress!, "/a2a"), httpClient); + + // Act + A2AException exception = await Assert.ThrowsAsync( + () => client.SendMessageAsync(CreateRequest("ctx-1", returnImmediately: true))); + + // Assert + Assert.Equal(A2AErrorCode.UnsupportedOperation, exception.ErrorCode); + } + + [Fact] + public async Task MapA2AHttpJson_WithAgentSelectorReturningMessage_IgnoresReturnImmediatelyAsync() + { + // Arrange + await using WebApplication app = await StartAsync( + _ => { }, + endpoints => endpoints.MapA2AHttpJson("/a2a", _ => new TurnCountingAgent("messages"))); + using HttpClient httpClient = app.GetTestClient(); + using A2AHttpJsonClient client = new(new Uri(httpClient.BaseAddress!, "/a2a"), httpClient); + + // Act + SendMessageResponse response = await client.SendMessageAsync(CreateRequest("ctx-1", returnImmediately: true)); + + // Assert + Assert.Equal("messages turn 1", GetText(response)); + } + + [Fact] + public async Task MapA2AHttpJson_WithAgentSelectorReturningNull_ReturnsNotFoundAsync() + { + // Arrange + int selections = 0; + await using WebApplication app = await StartAsync( + _ => { }, + endpoints => endpoints.MapA2AHttpJson("/agents/{agentId}/a2a", _ => + { + Interlocked.Increment(ref selections); + return null; + })); + using HttpClient httpClient = app.GetTestClient(); + + // Act + using HttpResponseMessage response = await httpClient.GetAsync(new Uri("/agents/unknown/a2a/tasks/task-1", UriKind.Relative)); + + // Assert + Assert.Equal(HttpStatusCode.NotFound, response.StatusCode); + Assert.Equal(1, selections); + } + + [Fact] + public async Task MapA2AJsonRpc_WithAgentSelectorReturningUnnamedAgent_FailsRequestAsync() + { + // Arrange + await using WebApplication app = await StartAsync( + _ => { }, + endpoints => endpoints.MapA2AJsonRpc("/a2a", _ => new TurnCountingAgent(name: null))); + using HttpClient httpClient = app.GetTestClient(); + using A2AClient client = new(new Uri(httpClient.BaseAddress!, "/a2a"), httpClient); + + // Act & Assert + InvalidOperationException exception = await Assert.ThrowsAsync(() => client.SendMessageAsync(CreateRequest("ctx-1"))); + Assert.Contains("stable name", exception.Message, StringComparison.Ordinal); + } + + [Fact] + public void MapA2AHttpJson_WithNullAgentSelector_ThrowsArgumentNullException() + { + // Arrange + Mock endpoints = new(); + + // Act & Assert + Assert.Throws(() => + endpoints.Object.MapA2AHttpJson("/a2a", (Func)null!)); + } + + private static AIAgent? SelectAgentFromRoute(HttpContext context) => + context.GetRouteValue("agentId") is string agentId + ? context.RequestServices.GetKeyedService(agentId) + : null; + + private static async Task StartAsync(Action configureServices, Action mapEndpoints) + { + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + builder.WebHost.UseTestServer(); + builder.Host.UseDefaultServiceProvider(options => options.ValidateScopes = true); + configureServices(builder.Services); + + WebApplication app = builder.Build(); + mapEndpoints(app); + await app.StartAsync(); + return app; + } + + private static SendMessageRequest CreateRequest(string contextId, bool returnImmediately = false) => new() + { + Message = new Message + { + MessageId = Guid.NewGuid().ToString("N"), + ContextId = contextId, + Role = Role.User, + Parts = [Part.FromText("hi")] + }, + Configuration = returnImmediately ? new SendMessageConfiguration { ReturnImmediately = true } : null + }; + + private static string? GetText(SendMessageResponse response) + { + Message message = Assert.IsType(response.Message); + return string.Concat(message.Parts.Select(part => part.Text)); + } + + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")] + private sealed class RequestDependency; + + private sealed class TurnCountingAgent(string? name) : AIAgent + { + public override string? Name => name; + + protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => + new(new TurnCountingSession()); + + protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => + new(session.StateBag.Serialize()); + + protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => + new(new TurnCountingSession(AgentSessionStateBag.Deserialize(serializedState))); + + protected override Task RunCoreAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + int turn = 1; + if (session is not null) + { + if (session.StateBag.TryGetValue("turn", out string? previousTurn)) + { + turn = int.Parse(previousTurn!, CultureInfo.InvariantCulture) + 1; + } + + session.StateBag.SetValue("turn", turn.ToString(CultureInfo.InvariantCulture)); + } + + await Task.Yield(); + yield return new AgentResponseUpdate + { + MessageId = "message-1", + Role = ChatRole.Assistant, + Contents = [new TextContent($"{name} turn {turn}")] + }; + } + } + + private sealed class TurnCountingSession : AgentSession + { + public TurnCountingSession() + { + } + + public TurnCountingSession(AgentSessionStateBag stateBag) + : base(stateBag) + { + } + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs index 1b16307e9ea..a66234f7804 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs @@ -206,7 +206,9 @@ public async Task AddA2AServer_WithCustomAgentSessionStore_ResolvesSuccessfullyA await using var provider = services.BuildServiceProvider(); var server = provider.GetKeyedService(AgentName); Assert.NotNull(server); - mockSessionStore.Verify(s => s.GetService(typeof(IsolationKeyScopedAgentSessionStore), null), Times.Once); + + // The session store is resolved for each A2A operation rather than when the server is created. + mockSessionStore.Verify(s => s.GetService(typeof(IsolationKeyScopedAgentSessionStore), null), Times.Never); } /// From b3033049ec9544f8a47378b2c900aa945cfb1473 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Mon, 19 Oct 2026 09:44:15 +0000 Subject: [PATCH 2/6] .NET: Keep A2A selector request services alive until the agent run completes A2AServer keeps streaming runs going after the response stream ends early, for example when the client disconnects, while the selected agent and its session store belong to the HTTP request. The selector request handler now waits for the agent run before the request ends and rejects runs that would start after it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../A2AEndpointRouteBuilderExtensions.cs | 4 +- .../SelectedAgentA2ARequestHandler.cs | 95 ++++++++++++- .../A2ARequestScopedAgentTests.cs | 131 ++++++++++++++++++ 3 files changed, 222 insertions(+), 8 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs index 975a7d6e1dc..a3ec1230705 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs @@ -120,7 +120,9 @@ public static IEndpointConventionBuilder MapA2AHttpJson(this IEndpointRouteBuild /// agent must therefore have a non-empty name. The session store is resolved from /// , so the run cannot outlive the request: when the configured /// returns a task, non-streaming requests with ReturnImmediately enabled are - /// rejected with . + /// rejected with , and a streaming request that ends early, for + /// example because the client disconnects, completes only after the agent run finishes. The run is not canceled, so + /// its task state is still persisted. /// /// /// Selection is not an authorization boundary. See diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs index d29d81ec256..10937b82640 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs @@ -57,19 +57,49 @@ public void Select(AIAgent agent, IServiceProvider requestServices) } /// - public Task SendMessageAsync(SendMessageRequest request, CancellationToken cancellationToken = default) + public async Task SendMessageAsync(SendMessageRequest request, CancellationToken cancellationToken = default) { Selection selection = this.GetSelection(); - selection.BindAgent(request.Message); - return selection.A2AServer.SendMessageAsync(request, cancellationToken); + RequestBoundAgentHandler handler = selection.BindAgent(request.Message); + try + { + return await selection.A2AServer.SendMessageAsync(request, cancellationToken).ConfigureAwait(false); + } + finally + { + await handler.ReleaseAsync().ConfigureAwait(false); + } } /// public IAsyncEnumerable SendStreamingMessageAsync(SendMessageRequest request, CancellationToken cancellationToken = default) { + // The selection is only visible in the endpoint filter's asynchronous flow, so bind it before the + // response stream is enumerated. Selection selection = this.GetSelection(); - selection.BindAgent(request.Message); - return selection.A2AServer.SendStreamingMessageAsync(request, cancellationToken); + RequestBoundAgentHandler handler = selection.BindAgent(request.Message); + return StreamAsync(selection.A2AServer, handler, request, cancellationToken); + } + + private static async IAsyncEnumerable StreamAsync( + A2AServer server, + RequestBoundAgentHandler handler, + SendMessageRequest request, + [EnumeratorCancellation] CancellationToken cancellationToken) + { + try + { + await foreach (StreamResponse response in server.SendStreamingMessageAsync(request, cancellationToken).ConfigureAwait(false)) + { + yield return response; + } + } + finally + { + // A2AServer keeps a streaming run going after the stream ends early, for example when the client + // disconnects, but the selected agent's services belong to this request. + await handler.ReleaseAsync().ConfigureAwait(false); + } } /// @@ -120,10 +150,61 @@ private sealed class Selection( { public A2AServer A2AServer => server.A2AServer; - public void BindAgent(Message message) + public RequestBoundAgentHandler BindAgent(Message message) { AIHostAgent hostAgent = A2AServerServiceCollectionExtensions.CreateHostAgent(requestServices, agent, agentName, sessionStorageIdentity: agentName); - server.Bind(message, new A2AAgentHandler(hostAgent, runMode, allowReturnImmediately: false)); + RequestBoundAgentHandler handler = new(new A2AAgentHandler(hostAgent, runMode, allowReturnImmediately: false)); + server.Bind(message, handler); + return handler; + } + } + + /// + /// Runs the selected agent only while the request that owns its services is active. + /// + /// + /// The request calls before it ends, which waits for a run in progress and rejects a run + /// that A2AServer starts in the background after the request has ended. + /// + private sealed class RequestBoundAgentHandler(IAgentHandler innerHandler) : IAgentHandler + { + private readonly object _gate = new(); + private Task? _execution; + private bool _released; + + public async Task ExecuteAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) + { + TaskCompletionSource completion = new(TaskCreationOptions.RunContinuationsAsynchronously); + lock (this._gate) + { + if (this._released) + { + throw new OperationCanceledException("The request that selected the agent ended before the agent run started."); + } + + this._execution = completion.Task; + } + + try + { + await innerHandler.ExecuteAsync(context, eventQueue, cancellationToken).ConfigureAwait(false); + } + finally + { + completion.SetResult(); + } + } + + public Task CancelAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) + => innerHandler.CancelAsync(context, eventQueue, cancellationToken); + + public Task ReleaseAsync() + { + lock (this._gate) + { + this._released = true; + return this._execution ?? Task.CompletedTask; + } } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs index 05afb712c60..59450610122 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs @@ -28,6 +28,8 @@ namespace Microsoft.Agents.AI.Hosting.A2A.UnitTests; /// public sealed class A2ARequestScopedAgentTests { + private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(10); + [Fact] public async Task AddA2AServer_WithScopedAgent_ResolvesAgentForEachOperationAsync() { @@ -289,6 +291,54 @@ public async Task MapA2AJsonRpc_WithAgentSelectorReturningUnnamedAgent_FailsRequ Assert.Contains("stable name", exception.Message, StringComparison.Ordinal); } + [Fact] + public async Task MapA2AHttpJson_WithAgentSelector_KeepsRequestServicesUntilStreamingRunCompletesAfterDisconnectAsync() + { + // Arrange + DisconnectProbe probe = new(); + await using WebApplication app = await StartAsync( + services => + { + services.AddScoped(); + services.AddKeyedScoped("streaming", (_, _) => new TrackingSessionStore(probe)); + }, + endpoints => + { + endpoints.Use(async (context, next) => + { + try + { + await next(context); + } + finally + { + probe.RequestCompleted.TrySetResult(); + } + }); + endpoints.MapA2AHttpJson("/a2a", context => new BlockingAgent(context.RequestServices.GetRequiredService(), probe)); + }); + using HttpClient httpClient = app.GetTestClient(); + using A2AHttpJsonClient client = new(new Uri(httpClient.BaseAddress!, "/a2a"), httpClient); + using CancellationTokenSource disconnect = new(); + + // Act + Task streaming = ConsumeAsync(client.SendStreamingMessageAsync(CreateRequest("ctx-1"), disconnect.Token)); + await Task.WhenAny(probe.AgentStarted.Task, streaming).WaitAsync(s_timeout); + Assert.False(streaming.IsCompleted); + disconnect.Cancel(); + Task first = await Task.WhenAny(probe.RequestCompleted.Task, Task.Delay(TimeSpan.FromMilliseconds(500))); + bool requestCompletedBeforeRun = first == probe.RequestCompleted.Task; + probe.ReleaseAgent.SetResult(); + await probe.SessionSaved.Task.WaitAsync(s_timeout); + await probe.RequestCompleted.Task.WaitAsync(s_timeout); + + // Assert + Assert.False(requestCompletedBeforeRun); + Assert.False(probe.DependencyDisposedWhenRunFinished); + Assert.False(probe.StoreDisposedWhenSessionSaved); + await Assert.ThrowsAnyAsync(() => streaming.WaitAsync(s_timeout)); + } + [Fact] public void MapA2AHttpJson_WithNullAgentSelector_ThrowsArgumentNullException() { @@ -336,6 +386,87 @@ private static async Task StartAsync(Action return string.Concat(message.Parts.Select(part => part.Text)); } + private static async Task ConsumeAsync(IAsyncEnumerable responses) + { + await foreach (StreamResponse _ in responses) + { + } + } + + private sealed class DisconnectProbe + { + public TaskCompletionSource AgentStarted { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource ReleaseAgent { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource SessionSaved { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource RequestCompleted { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public bool? DependencyDisposedWhenRunFinished { get; set; } + + public bool? StoreDisposedWhenSessionSaved { get; set; } + } + + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")] + private sealed class DisposableRequestDependency : IDisposable + { + public bool IsDisposed { get; private set; } + + public void Dispose() => this.IsDisposed = true; + } + + private sealed class TrackingSessionStore(DisconnectProbe probe) : AgentSessionStore, IDisposable + { + private bool _disposed; + + public override ValueTask SaveSessionAsync(AIAgent agent, AgentSessionStoreKey key, AgentSession session, CancellationToken cancellationToken = default) + { + probe.StoreDisposedWhenSessionSaved = this._disposed; + probe.SessionSaved.TrySetResult(); + return default; + } + + public override ValueTask GetSessionAsync(AIAgent agent, AgentSessionStoreKey key, CancellationToken cancellationToken = default) => + new((AgentSession?)null); + + public void Dispose() => this._disposed = true; + } + + private sealed class BlockingAgent(DisposableRequestDependency dependency, DisconnectProbe probe) : AIAgent + { + public override string? Name => "streaming"; + + protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => + new(new TurnCountingSession()); + + protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => + new(session.StateBag.Serialize()); + + protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => + new(new TurnCountingSession(AgentSessionStateBag.Deserialize(serializedState))); + + protected override Task RunCoreAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + probe.AgentStarted.TrySetResult(); + await probe.ReleaseAgent.Task.ConfigureAwait(false); + probe.DependencyDisposedWhenRunFinished = dependency.IsDisposed; + yield return new AgentResponseUpdate + { + MessageId = "message-1", + Role = ChatRole.Assistant, + Contents = [new TextContent("done")] + }; + } + } + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")] private sealed class RequestDependency; From 9c042e2388b01687065e65a6da37f7db706e9bcb Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Mon, 19 Oct 2026 10:01:00 +0000 Subject: [PATCH 3/6] .NET: Resolve the A2A selector task isolation key from each request The task store of the server cached for a selected agent name resolved AgentIsolationKeyProvider from the root provider, so scoped providers failed scope validation or were captured for every request. The cached task store now asks the provider resolved from the current request's services for each operation, and the selector endpoint filter executes the endpoint result in its own flow so streaming operations see the selection. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../A2AEndpointRouteBuilderExtensions.cs | 12 +++- .../SelectedAgentA2ARequestHandler.cs | 50 ++++++++++++++--- .../A2AServerServiceCollectionExtensions.cs | 20 +++++-- .../A2ARequestScopedAgentTests.cs | 55 +++++++++++++++++++ 4 files changed, 125 insertions(+), 12 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs index a3ec1230705..26c803bad02 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs @@ -282,7 +282,17 @@ private static void AddAgentSelection( } requestHandler.Select(agent, httpContext.RequestServices); - return await next(invocationContext).ConfigureAwait(false); + object? result = await next(invocationContext).ConfigureAwait(false); + + // Streaming results enumerate the request handler when they execute, which would otherwise happen after + // this filter returns and outside the asynchronous flow that carries the selection. + if (result is IResult httpResult) + { + await httpResult.ExecuteAsync(httpContext).ConfigureAwait(false); + return Results.Empty; + } + + return result; }); } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs index 10937b82640..b06b7e6f45d 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs @@ -44,16 +44,22 @@ public SelectedAgentA2ARequestHandler(IServiceProvider applicationServices, A2AS /// Makes the target of the A2A operations dispatched in the current asynchronous flow. /// /// The agent selected for the request. - /// The request services that own the agent's session services. + /// The request services that own the agent's session and isolation services. public void Select(AIAgent agent, IServiceProvider requestServices) { string agentName = HostedAgentResolution.GetSelectedAgentName(agent); SelectedAgentServer server = this._servers.GetOrAdd( agentName, - static (name, state) => new SelectedAgentServer(state._applicationServices, name, state._options), + static (name, owner) => new SelectedAgentServer(owner, name), this); - this._selection.Value = new Selection(agent, agentName, requestServices, server, this._options?.AgentRunMode ?? AgentRunMode.ReturnMessage); + this._selection.Value = new Selection( + agent, + agentName, + requestServices, + requestServices.GetService(), + server, + this._options?.AgentRunMode ?? AgentRunMode.ReturnMessage); } /// @@ -74,8 +80,6 @@ public async Task SendMessageAsync(SendMessageRequest reque /// public IAsyncEnumerable SendStreamingMessageAsync(SendMessageRequest request, CancellationToken cancellationToken = default) { - // The selection is only visible in the endpoint filter's asynchronous flow, so bind it before the - // response stream is enumerated. Selection selection = this.GetSelection(); RequestBoundAgentHandler handler = selection.BindAgent(request.Message); return StreamAsync(selection.A2AServer, handler, request, cancellationToken); @@ -145,11 +149,14 @@ private sealed class Selection( AIAgent agent, string agentName, IServiceProvider requestServices, + AgentIsolationKeyProvider? isolationKeyProvider, SelectedAgentServer server, AgentRunMode runMode) { public A2AServer A2AServer => server.A2AServer; + public AgentIsolationKeyProvider? IsolationKeyProvider => isolationKeyProvider; + public RequestBoundAgentHandler BindAgent(Message message) { AIHostAgent hostAgent = A2AServerServiceCollectionExtensions.CreateHostAgent(requestServices, agent, agentName, sessionStorageIdentity: agentName); @@ -216,9 +223,17 @@ private sealed class SelectedAgentServer : IAgentHandler // A2AServer passes the incoming message to the handler, including when it runs the handler in the background. private readonly ConditionalWeakTable _handlers = new(); - public SelectedAgentServer(IServiceProvider applicationServices, string agentName, A2AServerRegistrationOptions? options) + public SelectedAgentServer(SelectedAgentA2ARequestHandler owner, string agentName) { - this.A2AServer = A2AServerServiceCollectionExtensions.CreateA2AServer(applicationServices, agentName, this, options); + // The server is shared by every request that selects the agent name, so its task store must not capture + // request-scoped isolation services. Each operation uses the provider of the request that performs it. + this.A2AServer = A2AServerServiceCollectionExtensions.CreateA2AServer( + owner._applicationServices, + agentName, + this, + owner._options, + new SelectedRequestIsolationKeyProvider(owner), + requireIsolationKey: false); } public A2AServer A2AServer { get; } @@ -238,4 +253,25 @@ public Task ExecuteAsync(RequestContext context, AgentEventQueue eventQueue, Can public Task CancelAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) => A2AAgentHandler.CancelTaskAsync(context, eventQueue, cancellationToken); } + + /// + /// Supplies the isolation key of the request that performs the current operation. + /// + /// + /// A key is required only when the request services provide an , matching + /// the strict isolation applied to the session store. + /// + private sealed class SelectedRequestIsolationKeyProvider(SelectedAgentA2ARequestHandler owner) : AgentIsolationKeyProvider + { + public override async ValueTask GetIsolationKeyAsync(CancellationToken cancellationToken = default) + { + if (owner.GetSelection().IsolationKeyProvider is not { } provider) + { + return null; + } + + return await provider.GetIsolationKeyAsync(cancellationToken).ConfigureAwait(false) + ?? throw new InvalidOperationException("Agent isolation key is required but was not provided by the configured AgentIsolationKeyProvider."); + } + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs index c466e394de4..8c9783fc464 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs @@ -227,7 +227,8 @@ private static A2AServer CreateA2AServer( sessionStorageIdentity, options?.AgentRunMode ?? AgentRunMode.ReturnMessage); - return CreateA2AServer(serviceProvider, agentName, agentHandler, options); + var isolationKeyProvider = serviceProvider.GetService(); + return CreateA2AServer(serviceProvider, agentName, agentHandler, options, isolationKeyProvider, requireIsolationKey: isolationKeyProvider != null); } /// @@ -237,17 +238,28 @@ private static A2AServer CreateA2AServer( /// The name that keys the agent's hosted services. /// The handler that executes A2A operations. /// The optional registration options. + /// + /// The provider that partitions the default task store, or when tasks are not isolated. + /// + /// + /// to reject task store operations for which returns no key. + /// /// The created server. - internal static A2AServer CreateA2AServer(IServiceProvider serviceProvider, string agentName, IAgentHandler agentHandler, A2AServerRegistrationOptions? options) + internal static A2AServer CreateA2AServer( + IServiceProvider serviceProvider, + string agentName, + IAgentHandler agentHandler, + A2AServerRegistrationOptions? options, + AgentIsolationKeyProvider? isolationKeyProvider, + bool requireIsolationKey) { - var isolationKeyProvider = serviceProvider.GetService(); var loggerFactory = serviceProvider.GetService() ?? NullLoggerFactory.Instance; ITaskStore taskStore = serviceProvider.GetKeyedService(agentName) ?? new InMemoryTaskStore(); // Wrap the task store with isolation key scoping, same as the session store. if (taskStore is not IsolationKeyScopedTaskStore) { - taskStore = new IsolationKeyScopedTaskStore(taskStore, isolationKeyProvider, strict: isolationKeyProvider != null); + taskStore = new IsolationKeyScopedTaskStore(taskStore, isolationKeyProvider, strict: requireIsolationKey); } return new A2AServer( diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs index 59450610122..b623f68e03b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs @@ -339,6 +339,46 @@ public async Task MapA2AHttpJson_WithAgentSelector_KeepsRequestServicesUntilStre await Assert.ThrowsAnyAsync(() => streaming.WaitAsync(s_timeout)); } + [Fact] + public async Task MapA2AHttpJson_WithAgentSelectorAndScopedIsolationKeyProvider_IsolatesTasksPerRequestAsync() + { + // Arrange + await using WebApplication app = await StartAsync( + services => services.AddScoped(), + endpoints => + { + endpoints.Use((context, next) => + { + var provider = (RequestIsolationKeyProvider)context.RequestServices.GetRequiredService(); + provider.Key = context.Request.Headers["X-User"].ToString(); + return next(context); + }); + endpoints.MapA2AHttpJson("/a2a", _ => new TurnCountingAgent("tasks"), options => options.AgentRunMode = AgentRunMode.ReturnTask); + }); + using HttpClient aliceHttpClient = CreateUserHttpClient(app, "alice"); + using HttpClient bobHttpClient = CreateUserHttpClient(app, "bob"); + using A2AHttpJsonClient alice = new(new Uri(aliceHttpClient.BaseAddress!, "/a2a"), aliceHttpClient); + using A2AHttpJsonClient bob = new(new Uri(bobHttpClient.BaseAddress!, "/a2a"), bobHttpClient); + + // Act + SendMessageResponse sent = await alice.SendMessageAsync(CreateRequest("ctx-1")); + List streamed = []; + await foreach (StreamResponse response in alice.SendStreamingMessageAsync(CreateRequest("ctx-2"))) + { + streamed.Add(response); + } + + string streamedTaskId = streamed.Select(response => response.Task?.Id ?? response.StatusUpdate?.TaskId).First(id => id is not null)!; + AgentTask sentTask = await alice.GetTaskAsync(new GetTaskRequest { Id = sent.Task!.Id }); + AgentTask streamedTask = await alice.GetTaskAsync(new GetTaskRequest { Id = streamedTaskId }); + A2AException otherUserError = await Assert.ThrowsAsync(() => bob.GetTaskAsync(new GetTaskRequest { Id = sent.Task.Id })); + + // Assert + Assert.Equal(TaskState.Completed, sentTask.Status.State); + Assert.Equal(TaskState.Completed, streamedTask.Status.State); + Assert.Equal(A2AErrorCode.TaskNotFound, otherUserError.ErrorCode); + } + [Fact] public void MapA2AHttpJson_WithNullAgentSelector_ThrowsArgumentNullException() { @@ -393,6 +433,21 @@ private static async Task ConsumeAsync(IAsyncEnumerable response } } + private static HttpClient CreateUserHttpClient(WebApplication app, string user) + { + HttpClient httpClient = app.GetTestClient(); + httpClient.DefaultRequestHeaders.Add("X-User", user); + return httpClient; + } + + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")] + private sealed class RequestIsolationKeyProvider : AgentIsolationKeyProvider + { + public string? Key { get; set; } + + public override ValueTask GetIsolationKeyAsync(CancellationToken cancellationToken = default) => new(this.Key); + } + private sealed class DisconnectProbe { public TaskCompletionSource AgentStarted { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); From c5887a31b7677dd15c8dead25d0977de57070595 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Mon, 19 Oct 2026 10:33:09 +0000 Subject: [PATCH 4/6] .NET: Resolve the A2A isolation key from the request for registered servers Servers registered with AddA2AServer resolved AgentIsolationKeyProvider from the root provider for the task store and from a new operation scope for the session store, so scoped providers failed scope validation or did not see the request. The A2A endpoint extensions now resolve the isolation key once per request from its services and make it ambient for the request flow, including operations that continue in the background. The task and session stores of both registered and selector servers use that key, falling back to a new service scope outside A2A requests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../A2AEndpointRouteBuilderExtensions.cs | 37 +++++++++-- .../SelectedAgentA2ARequestHandler.cs | 40 +++--------- .../A2ARequestIsolation.cs | 61 +++++++++++++++++++ .../A2ARequestIsolationKeyProvider.cs | 35 +++++++++++ .../A2AServerServiceCollectionExtensions.cs | 39 ++++++------ .../ScopedA2AAgentHandler.cs | 13 +++- .../A2ARequestScopedAgentTests.cs | 60 ++++++++++++++++-- 7 files changed, 223 insertions(+), 62 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2ARequestIsolation.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2ARequestIsolationKeyProvider.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs index 26c803bad02..12ae9f1431d 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs @@ -77,6 +77,9 @@ public static IEndpointConventionBuilder MapA2AHttpJson(this IEndpointRouteBuild /// task isolation remains necessary without session persistence. For claims-based isolation, register /// AddHttpContextAccessor() and UseClaimsBasedAgentIsolation(...) from /// Microsoft.Agents.AI.Hosting.AspNetCore, using a claim that uniquely identifies the caller. + /// The is resolved from once + /// per request, so it can be registered with any lifetime, and its key is reused by the operations that the + /// request continues in the background. /// Protect each enabled protocol binding, not just one of the HTTP+JSON and JSON-RPC endpoints. /// public static IEndpointConventionBuilder MapA2AHttpJson(this IEndpointRouteBuilder endpoints, string agentName, string path) @@ -91,6 +94,7 @@ public static IEndpointConventionBuilder MapA2AHttpJson(this IEndpointRouteBuild $"Call services.AddA2AServer(\"{agentName}\") or agentBuilder.AddA2AServer() during service registration to register one."); IEndpointConventionBuilder endpoint = endpoints.MapHttpA2A(a2aServer, path); + AddRequestFlow(endpoint); MarkFeatureUsed(); return endpoint; } @@ -213,6 +217,7 @@ public static IEndpointConventionBuilder MapA2AJsonRpc(this IEndpointRouteBuilde $"Call services.AddA2AServer(\"{agentName}\") or agentBuilder.AddA2AServer() during service registration to register one."); IEndpointConventionBuilder endpoint = endpoints.MapA2A(a2aServer, path); + AddRequestFlow(endpoint); MarkFeatureUsed(); return endpoint; } @@ -271,21 +276,45 @@ private static void AddAgentSelection( SelectedAgentA2ARequestHandler requestHandler, Func agentSelector) { - // The A2A endpoints call the request handler without request context, so select the agent for the + AddRequestFlow(endpoint, httpContext => + { + if (agentSelector(httpContext) is not { } agent) + { + return false; + } + + requestHandler.Select(agent, httpContext.RequestServices); + return true; + }); + } + + /// + /// Runs the A2A endpoints in an asynchronous flow that carries the request's isolation key and, when + /// is specified, the agent selected for the request. + /// + /// The A2A endpoints. + /// + /// An optional delegate that selects the agent for the request, returning when no agent matches. + /// + private static void AddRequestFlow(IEndpointConventionBuilder endpoint, Func? selectAgent = null) + { + // The A2A endpoints call the request handler without request context, so record the request state for the // asynchronous flow of the request before the endpoint runs. endpoint.AddEndpointFilter(async (invocationContext, next) => { HttpContext httpContext = invocationContext.HttpContext; - if (agentSelector(httpContext) is not { } agent) + if (selectAgent?.Invoke(httpContext) == false) { return Results.NotFound(); } - requestHandler.Select(agent, httpContext.RequestServices); + // A2A servers are shared by all requests, so resolve the key from the request services while they are + // alive. Operations that continue in the background reuse it. + A2ARequestIsolation.Current = await A2ARequestIsolation.ResolveAsync(httpContext.RequestServices, httpContext.RequestAborted).ConfigureAwait(false); object? result = await next(invocationContext).ConfigureAwait(false); // Streaming results enumerate the request handler when they execute, which would otherwise happen after - // this filter returns and outside the asynchronous flow that carries the selection. + // this filter returns and outside the asynchronous flow that carries the request state. if (result is IResult httpResult) { await httpResult.ExecuteAsync(httpContext).ConfigureAwait(false); diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs index b06b7e6f45d..dd93d1cdd6f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs @@ -27,6 +27,7 @@ internal sealed class SelectedAgentA2ARequestHandler : IA2ARequestHandler private readonly AsyncLocal _selection = new(); private readonly ConcurrentDictionary _servers = new(StringComparer.Ordinal); private readonly IServiceProvider _applicationServices; + private readonly A2ARequestIsolationKeyProvider _isolationKeyProvider; private readonly A2AServerRegistrationOptions? _options; /// @@ -37,6 +38,7 @@ internal sealed class SelectedAgentA2ARequestHandler : IA2ARequestHandler public SelectedAgentA2ARequestHandler(IServiceProvider applicationServices, A2AServerRegistrationOptions? options) { this._applicationServices = applicationServices; + this._isolationKeyProvider = new A2ARequestIsolationKeyProvider(applicationServices.GetRequiredService()); this._options = options; } @@ -44,7 +46,7 @@ public SelectedAgentA2ARequestHandler(IServiceProvider applicationServices, A2AS /// Makes the target of the A2A operations dispatched in the current asynchronous flow. /// /// The agent selected for the request. - /// The request services that own the agent's session and isolation services. + /// The request services that own the agent's session services. public void Select(AIAgent agent, IServiceProvider requestServices) { string agentName = HostedAgentResolution.GetSelectedAgentName(agent); @@ -57,7 +59,7 @@ public void Select(AIAgent agent, IServiceProvider requestServices) agent, agentName, requestServices, - requestServices.GetService(), + this._isolationKeyProvider, server, this._options?.AgentRunMode ?? AgentRunMode.ReturnMessage); } @@ -149,17 +151,15 @@ private sealed class Selection( AIAgent agent, string agentName, IServiceProvider requestServices, - AgentIsolationKeyProvider? isolationKeyProvider, + A2ARequestIsolationKeyProvider isolationKeyProvider, SelectedAgentServer server, AgentRunMode runMode) { public A2AServer A2AServer => server.A2AServer; - public AgentIsolationKeyProvider? IsolationKeyProvider => isolationKeyProvider; - public RequestBoundAgentHandler BindAgent(Message message) { - AIHostAgent hostAgent = A2AServerServiceCollectionExtensions.CreateHostAgent(requestServices, agent, agentName, sessionStorageIdentity: agentName); + AIHostAgent hostAgent = A2AServerServiceCollectionExtensions.CreateHostAgent(requestServices, agent, agentName, sessionStorageIdentity: agentName, isolationKeyProvider); RequestBoundAgentHandler handler = new(new A2AAgentHandler(hostAgent, runMode, allowReturnImmediately: false)); server.Bind(message, handler); return handler; @@ -225,15 +225,12 @@ private sealed class SelectedAgentServer : IAgentHandler public SelectedAgentServer(SelectedAgentA2ARequestHandler owner, string agentName) { - // The server is shared by every request that selects the agent name, so its task store must not capture - // request-scoped isolation services. Each operation uses the provider of the request that performs it. this.A2AServer = A2AServerServiceCollectionExtensions.CreateA2AServer( owner._applicationServices, agentName, this, - owner._options, - new SelectedRequestIsolationKeyProvider(owner), - requireIsolationKey: false); + owner._isolationKeyProvider, + owner._options); } public A2AServer A2AServer { get; } @@ -253,25 +250,4 @@ public Task ExecuteAsync(RequestContext context, AgentEventQueue eventQueue, Can public Task CancelAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken) => A2AAgentHandler.CancelTaskAsync(context, eventQueue, cancellationToken); } - - /// - /// Supplies the isolation key of the request that performs the current operation. - /// - /// - /// A key is required only when the request services provide an , matching - /// the strict isolation applied to the session store. - /// - private sealed class SelectedRequestIsolationKeyProvider(SelectedAgentA2ARequestHandler owner) : AgentIsolationKeyProvider - { - public override async ValueTask GetIsolationKeyAsync(CancellationToken cancellationToken = default) - { - if (owner.GetSelection().IsolationKeyProvider is not { } provider) - { - return null; - } - - return await provider.GetIsolationKeyAsync(cancellationToken).ConfigureAwait(false) - ?? throw new InvalidOperationException("Agent isolation key is required but was not provided by the configured AgentIsolationKeyProvider."); - } - } } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2ARequestIsolation.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2ARequestIsolation.cs new file mode 100644 index 00000000000..c23f14d5e20 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2ARequestIsolation.cs @@ -0,0 +1,61 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.DependencyInjection; + +namespace Microsoft.Agents.AI.Hosting.A2A; + +/// +/// The isolation key of the request that performs the A2A operations in the current asynchronous flow. +/// +/// +/// The key is resolved once while the request services are alive, so operations that A2A continues in the +/// background after the request ends reuse it instead of resolving request-scoped services. +/// +internal sealed class A2ARequestIsolation +{ + private static readonly AsyncLocal s_current = new(); + + private readonly string? _key; + private readonly bool _isRequired; + + private A2ARequestIsolation(string? key, bool isRequired) + { + this._key = key; + this._isRequired = isRequired; + } + + /// + /// Gets or sets the isolation of the request that performs the operations in the current asynchronous flow. + /// + public static A2ARequestIsolation? Current + { + get => s_current.Value; + set => s_current.Value = value; + } + + /// + /// Resolves the isolation key from the registered in . + /// + /// The services of the request or operation. + /// A to cancel the operation. + /// The resolved isolation, which requires a key when a provider is registered. + public static async ValueTask ResolveAsync(IServiceProvider services, CancellationToken cancellationToken) + { + AgentIsolationKeyProvider? provider = services.GetService(); + string? key = provider is null ? null : await provider.GetIsolationKeyAsync(cancellationToken).ConfigureAwait(false); + return new A2ARequestIsolation(key, isRequired: provider is not null); + } + + /// + /// Gets the isolation key. + /// + /// The isolation key, or when no provider is registered. + /// A provider is registered but did not supply a non-empty key. + public string? GetKey() => + this._isRequired && string.IsNullOrWhiteSpace(this._key) + ? throw new InvalidOperationException("Agent isolation key is required but was not provided by the configured AgentIsolationKeyProvider.") + : this._key; +} diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2ARequestIsolationKeyProvider.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2ARequestIsolationKeyProvider.cs new file mode 100644 index 00000000000..63997a73894 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2ARequestIsolationKeyProvider.cs @@ -0,0 +1,35 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.DependencyInjection; + +namespace Microsoft.Agents.AI.Hosting.A2A; + +/// +/// Supplies the isolation key of the request that performs the current A2A operation to the task and session stores. +/// +/// +/// A2A servers are shared by all requests, so they must not capture an +/// that can be scoped. Within a request mapped by the A2A endpoint extensions, the key resolved for the request is +/// used. Otherwise, for example when an A2AServer is invoked directly, the provider is resolved from a new +/// service scope for each operation. +/// +internal sealed class A2ARequestIsolationKeyProvider(IServiceScopeFactory scopeFactory) : AgentIsolationKeyProvider +{ + /// + public override async ValueTask GetIsolationKeyAsync(CancellationToken cancellationToken = default) + { + A2ARequestIsolation? isolation = A2ARequestIsolation.Current; + if (isolation is null) + { + AsyncServiceScope scope = scopeFactory.CreateAsyncScope(); + await using (scope.ConfigureAwait(false)) + { + isolation = await A2ARequestIsolation.ResolveAsync(scope.ServiceProvider, cancellationToken).ConfigureAwait(false); + } + } + + return isolation.GetKey(); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs index 8c9783fc464..b8730549ded 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AServerServiceCollectionExtensions.cs @@ -220,15 +220,17 @@ private static A2AServer CreateA2AServer( { // The server owns task state shared across requests, while the agent and its session services are // resolved for each operation so that scoped and transient registrations keep their lifetimes. + var scopeFactory = serviceProvider.GetRequiredService(); + var isolationKeyProvider = new A2ARequestIsolationKeyProvider(scopeFactory); var agentHandler = new ScopedA2AAgentHandler( - serviceProvider.GetRequiredService(), + scopeFactory, agentName, agent, sessionStorageIdentity, + isolationKeyProvider, options?.AgentRunMode ?? AgentRunMode.ReturnMessage); - var isolationKeyProvider = serviceProvider.GetService(); - return CreateA2AServer(serviceProvider, agentName, agentHandler, options, isolationKeyProvider, requireIsolationKey: isolationKeyProvider != null); + return CreateA2AServer(serviceProvider, agentName, agentHandler, isolationKeyProvider, options); } /// @@ -237,29 +239,24 @@ private static A2AServer CreateA2AServer( /// The application service provider. /// The name that keys the agent's hosted services. /// The handler that executes A2A operations. + /// The provider that supplies the isolation key of the current request. /// The optional registration options. - /// - /// The provider that partitions the default task store, or when tasks are not isolated. - /// - /// - /// to reject task store operations for which returns no key. - /// /// The created server. internal static A2AServer CreateA2AServer( IServiceProvider serviceProvider, string agentName, IAgentHandler agentHandler, - A2AServerRegistrationOptions? options, - AgentIsolationKeyProvider? isolationKeyProvider, - bool requireIsolationKey) + A2ARequestIsolationKeyProvider isolationKeyProvider, + A2AServerRegistrationOptions? options) { var loggerFactory = serviceProvider.GetService() ?? NullLoggerFactory.Instance; ITaskStore taskStore = serviceProvider.GetKeyedService(agentName) ?? new InMemoryTaskStore(); - // Wrap the task store with isolation key scoping, same as the session store. + // Wrap the task store with isolation key scoping, same as the session store. The provider enforces + // that a key is supplied when an AgentIsolationKeyProvider is registered. if (taskStore is not IsolationKeyScopedTaskStore) { - taskStore = new IsolationKeyScopedTaskStore(taskStore, isolationKeyProvider, strict: requireIsolationKey); + taskStore = new IsolationKeyScopedTaskStore(taskStore, isolationKeyProvider, strict: false); } return new A2AServer( @@ -279,17 +276,23 @@ internal static A2AServer CreateA2AServer( /// /// The stable logical identity that partitions persisted sessions, or to use the agent instance identity. /// + /// The provider that supplies the isolation key of the current request. /// The hosting wrapper for . - internal static AIHostAgent CreateHostAgent(IServiceProvider services, AIAgent agent, string agentName, string? sessionStorageIdentity) + internal static AIHostAgent CreateHostAgent( + IServiceProvider services, + AIAgent agent, + string agentName, + string? sessionStorageIdentity, + A2ARequestIsolationKeyProvider isolationKeyProvider) { - var isolationKeyProvider = services.GetService(); var agentSessionStore = services.GetKeyedService(agentName); - // Ensure that we have an IsolationKeyScopedAgentSessionStore registered. + // Ensure that we have an IsolationKeyScopedAgentSessionStore registered. The provider enforces that a key + // is supplied when an AgentIsolationKeyProvider is registered. if (agentSessionStore?.GetService() is null) { agentSessionStore ??= new NoopAgentSessionStore(); - agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = isolationKeyProvider != null }); + agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = false }); } return new AIHostAgent(agent, agentSessionStore, sessionStorageIdentity); diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ScopedA2AAgentHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ScopedA2AAgentHandler.cs index 49352b2f8bf..ee73a2a5ed1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ScopedA2AAgentHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/ScopedA2AAgentHandler.cs @@ -25,6 +25,7 @@ internal sealed class ScopedA2AAgentHandler : IAgentHandler private readonly string _agentName; private readonly AIAgent? _agent; private readonly string? _sessionStorageIdentity; + private readonly A2ARequestIsolationKeyProvider _isolationKeyProvider; private readonly AgentRunMode _runMode; /// @@ -36,13 +37,21 @@ internal sealed class ScopedA2AAgentHandler : IAgentHandler /// /// The stable logical identity that partitions persisted sessions, or to use the agent instance identity. /// + /// The provider that supplies the isolation key of the request that performs each operation. /// Controls which A2A artifact the agent response is returned as. - public ScopedA2AAgentHandler(IServiceScopeFactory scopeFactory, string agentName, AIAgent? agent, string? sessionStorageIdentity, AgentRunMode runMode) + public ScopedA2AAgentHandler( + IServiceScopeFactory scopeFactory, + string agentName, + AIAgent? agent, + string? sessionStorageIdentity, + A2ARequestIsolationKeyProvider isolationKeyProvider, + AgentRunMode runMode) { this._scopeFactory = scopeFactory; this._agentName = agentName; this._agent = agent; this._sessionStorageIdentity = sessionStorageIdentity; + this._isolationKeyProvider = isolationKeyProvider; this._runMode = runMode; } @@ -78,7 +87,7 @@ private IAgentHandler CreateHandler(IServiceProvider services) } AIAgent agent = this._agent ?? services.GetRequiredKeyedService(this._agentName); - AIHostAgent hostAgent = A2AServerServiceCollectionExtensions.CreateHostAgent(services, agent, this._agentName, this._sessionStorageIdentity); + AIHostAgent hostAgent = A2AServerServiceCollectionExtensions.CreateHostAgent(services, agent, this._agentName, this._sessionStorageIdentity, this._isolationKeyProvider); return new A2AAgentHandler(hostAgent, this._runMode); } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs index b623f68e03b..f12aaac9613 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs @@ -339,6 +339,48 @@ public async Task MapA2AHttpJson_WithAgentSelector_KeepsRequestServicesUntilStre await Assert.ThrowsAnyAsync(() => streaming.WaitAsync(s_timeout)); } + [Fact] + public async Task MapA2AHttpJson_WithRegisteredServerAndScopedIsolationKeyProvider_IsolatesTasksAndSessionsPerRequestAsync() + { + // Arrange + await using WebApplication app = await StartAsync( + services => + { + services.AddScoped(); + services.AddAIAgent("assistant", (_, name) => new TurnCountingAgent(name), ServiceLifetime.Scoped) + .WithInMemorySessionStore(withIsolation: false); + services.AddA2AServer("assistant", options => options.AgentRunMode = AgentRunMode.ReturnTask); + }, + endpoints => + { + UseRequestIsolationKeyFromHeader(endpoints); + endpoints.MapA2AHttpJson("assistant", "/a2a"); + }); + using HttpClient aliceHttpClient = CreateUserHttpClient(app, "alice"); + using HttpClient bobHttpClient = CreateUserHttpClient(app, "bob"); + using A2AHttpJsonClient alice = new(new Uri(aliceHttpClient.BaseAddress!, "/a2a"), aliceHttpClient); + using A2AHttpJsonClient bob = new(new Uri(bobHttpClient.BaseAddress!, "/a2a"), bobHttpClient); + + // Act + SendMessageResponse aliceFirst = await alice.SendMessageAsync(CreateRequest("ctx-1")); + List aliceStreamed = []; + await foreach (StreamResponse response in alice.SendStreamingMessageAsync(CreateRequest("ctx-1"))) + { + aliceStreamed.Add(response); + } + + SendMessageResponse bobFirst = await bob.SendMessageAsync(CreateRequest("ctx-1")); + string aliceStreamedTaskId = aliceStreamed.Select(response => response.Task?.Id ?? response.StatusUpdate?.TaskId).First(id => id is not null)!; + AgentTask aliceStreamedTask = await alice.GetTaskAsync(new GetTaskRequest { Id = aliceStreamedTaskId }); + A2AException otherUserError = await Assert.ThrowsAsync(() => bob.GetTaskAsync(new GetTaskRequest { Id = aliceFirst.Task!.Id })); + + // Assert + Assert.Equal("assistant turn 1", GetArtifactText(aliceFirst.Task!)); + Assert.Equal("assistant turn 2", GetArtifactText(aliceStreamedTask)); + Assert.Equal("assistant turn 1", GetArtifactText(bobFirst.Task!)); + Assert.Equal(A2AErrorCode.TaskNotFound, otherUserError.ErrorCode); + } + [Fact] public async Task MapA2AHttpJson_WithAgentSelectorAndScopedIsolationKeyProvider_IsolatesTasksPerRequestAsync() { @@ -347,12 +389,7 @@ public async Task MapA2AHttpJson_WithAgentSelectorAndScopedIsolationKeyProvider_ services => services.AddScoped(), endpoints => { - endpoints.Use((context, next) => - { - var provider = (RequestIsolationKeyProvider)context.RequestServices.GetRequiredService(); - provider.Key = context.Request.Headers["X-User"].ToString(); - return next(context); - }); + UseRequestIsolationKeyFromHeader(endpoints); endpoints.MapA2AHttpJson("/a2a", _ => new TurnCountingAgent("tasks"), options => options.AgentRunMode = AgentRunMode.ReturnTask); }); using HttpClient aliceHttpClient = CreateUserHttpClient(app, "alice"); @@ -440,6 +477,17 @@ private static HttpClient CreateUserHttpClient(WebApplication app, string user) return httpClient; } + private static void UseRequestIsolationKeyFromHeader(WebApplication app) => + app.Use((context, next) => + { + var provider = (RequestIsolationKeyProvider)context.RequestServices.GetRequiredService(); + provider.Key = context.Request.Headers["X-User"].ToString(); + return next(context); + }); + + private static string GetArtifactText(AgentTask task) => + string.Concat(task.Artifacts!.SelectMany(artifact => artifact.Parts).Select(part => part.Text)); + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")] private sealed class RequestIsolationKeyProvider : AgentIsolationKeyProvider { From 4e70ad2daaab4357c0a3d1c978f4e365bd822236 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Mon, 19 Oct 2026 10:53:46 +0000 Subject: [PATCH 5/6] .NET: Dispose A2A selector servers when the application stops The A2A servers cached for selected agents were created manually, so nothing disposed them and application shutdown could dispose their dependencies while detached operations were still running. The selector request handler now owns and disposes these servers, and it is disposed when the application has stopped, after in-flight requests and before the service provider is disposed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../A2AEndpointRouteBuilderExtensions.cs | 20 ++++- .../SelectedAgentA2ARequestHandler.cs | 49 +++++++++-- .../A2ARequestScopedAgentTests.cs | 86 +++++++++++++++---- 3 files changed, 128 insertions(+), 27 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs index 12ae9f1431d..1f776204508 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs @@ -10,6 +10,7 @@ using Microsoft.AspNetCore.Http; using Microsoft.AspNetCore.Routing; using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Hosting; using Microsoft.Shared.DiagnosticIds; namespace Microsoft.AspNetCore.Builder; @@ -126,7 +127,8 @@ public static IEndpointConventionBuilder MapA2AHttpJson(this IEndpointRouteBuild /// returns a task, non-streaming requests with ReturnImmediately enabled are /// rejected with , and a streaming request that ends early, for /// example because the client disconnects, completes only after the agent run finishes. The run is not canceled, so - /// its task state is still persisted. + /// its task state is still persisted. When the application stops, operations that are still running for selected + /// agents are canceled and awaited. /// /// /// Selection is not an authorization boundary. See @@ -268,7 +270,21 @@ private static SelectedAgentA2ARequestHandler CreateSelectedAgentRequestHandler( configureOptions(options); } - return new SelectedAgentA2ARequestHandler(endpoints.ServiceProvider, options); + var requestHandler = new SelectedAgentA2ARequestHandler(endpoints.ServiceProvider, options); + + // The handler is created after the service provider is built, so the host cannot own it. Dispose it once the + // server has stopped processing requests and before the service provider is disposed, so the operations its + // A2A servers still run are canceled and joined while their dependencies are alive. + if (endpoints.ServiceProvider.GetService() is { } lifetime) + { +#pragma warning disable VSTHRD002 // Lifetime callbacks are synchronous, and the host must wait for the servers to stop. + _ = lifetime.ApplicationStopped.Register( + static state => ((SelectedAgentA2ARequestHandler)state!).DisposeAsync().AsTask().GetAwaiter().GetResult(), + requestHandler); +#pragma warning restore VSTHRD002 + } + + return requestHandler; } private static void AddAgentSelection( diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs index dd93d1cdd6f..588ecdc36ef 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs @@ -1,7 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. using System; -using System.Collections.Concurrent; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; using System.Runtime.CompilerServices; @@ -19,16 +18,19 @@ namespace Microsoft.Agents.AI.Hosting.A2A; /// /// The A2A endpoints invoke this handler without request context, so the endpoint filter records the selected agent /// for the asynchronous flow of the request before the endpoint runs. Each selected agent name owns an -/// whose task state is shared by all requests that select that name. +/// whose task state is shared by all requests that select that name. The handler owns these +/// servers, so disposing it cancels and joins the operations they still run. /// [Experimental(DiagnosticIds.Experiments.AIResponseContinuations)] -internal sealed class SelectedAgentA2ARequestHandler : IA2ARequestHandler +internal sealed class SelectedAgentA2ARequestHandler : IA2ARequestHandler, IAsyncDisposable { private readonly AsyncLocal _selection = new(); - private readonly ConcurrentDictionary _servers = new(StringComparer.Ordinal); + private readonly object _serversGate = new(); + private readonly Dictionary _servers = new(StringComparer.Ordinal); private readonly IServiceProvider _applicationServices; private readonly A2ARequestIsolationKeyProvider _isolationKeyProvider; private readonly A2AServerRegistrationOptions? _options; + private Task? _disposal; /// /// Initializes a new instance of the class. @@ -47,13 +49,20 @@ public SelectedAgentA2ARequestHandler(IServiceProvider applicationServices, A2AS /// /// The agent selected for the request. /// The request services that own the agent's session services. + /// The handler has been disposed. public void Select(AIAgent agent, IServiceProvider requestServices) { string agentName = HostedAgentResolution.GetSelectedAgentName(agent); - SelectedAgentServer server = this._servers.GetOrAdd( - agentName, - static (name, owner) => new SelectedAgentServer(owner, name), - this); + SelectedAgentServer? server; + lock (this._serversGate) + { + ObjectDisposedException.ThrowIf(this._disposal is not null, this); + if (!this._servers.TryGetValue(agentName, out server)) + { + server = new SelectedAgentServer(this, agentName); + this._servers.Add(agentName, server); + } + } this._selection.Value = new Selection( agent, @@ -64,6 +73,30 @@ public void Select(AIAgent agent, IServiceProvider requestServices) this._options?.AgentRunMode ?? AgentRunMode.ReturnMessage); } + /// + /// Disposes the servers of all selected agents, canceling and joining the operations they still run. + /// + /// A task that represents the asynchronous dispose operation. + public ValueTask DisposeAsync() + { + lock (this._serversGate) + { + this._disposal ??= DisposeServersAsync([.. this._servers.Values]); + return new(this._disposal); + } + + static async Task DisposeServersAsync(SelectedAgentServer[] servers) + { + List disposals = new(servers.Length); + foreach (SelectedAgentServer server in servers) + { + disposals.Add(server.A2AServer.DisposeAsync().AsTask()); + } + + await Task.WhenAll(disposals).ConfigureAwait(false); + } + } + /// public async Task SendMessageAsync(SendMessageRequest request, CancellationToken cancellationToken = default) { diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs index f12aaac9613..c3348e6a5e6 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs @@ -416,6 +416,30 @@ public async Task MapA2AHttpJson_WithAgentSelectorAndScopedIsolationKeyProvider_ Assert.Equal(A2AErrorCode.TaskNotFound, otherUserError.ErrorCode); } + [Fact] + public async Task MapA2AHttpJson_WithAgentSelector_CancelsAndJoinsRunsWhenApplicationStopsAsync() + { + // Arrange + ShutdownProbe probe = new(); + await using WebApplication app = await StartAsync( + _ => { }, + endpoints => endpoints.MapA2AHttpJson("/a2a", _ => new WaitForCancellationAgent(probe))); + using HttpClient httpClient = app.GetTestClient(); + using A2AHttpJsonClient client = new(new Uri(httpClient.BaseAddress!, "/a2a"), httpClient); + Task streaming = ConsumeAsync(client.SendStreamingMessageAsync(CreateRequest("ctx-1"))); + await Task.WhenAny(probe.AgentStarted.Task, streaming).WaitAsync(s_timeout); + Assert.False(streaming.IsCompleted); + + // Act + await app.StopAsync().WaitAsync(s_timeout); + bool runFinishedWhenStopped = probe.RunFinished.Task.IsCompleted; + + // Assert + Assert.True(runFinishedWhenStopped); + Assert.True(probe.RunCanceled); + Assert.Same(streaming, await Task.WhenAny(streaming, Task.Delay(s_timeout))); + } + [Fact] public void MapA2AHttpJson_WithNullAgentSelector_ThrowsArgumentNullException() { @@ -536,22 +560,10 @@ public override ValueTask SaveSessionAsync(AIAgent agent, AgentSessionStoreKey k public void Dispose() => this._disposed = true; } - private sealed class BlockingAgent(DisposableRequestDependency dependency, DisconnectProbe probe) : AIAgent + private sealed class BlockingAgent(DisposableRequestDependency dependency, DisconnectProbe probe) : StreamingTestAgent { public override string? Name => "streaming"; - protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => - new(new TurnCountingSession()); - - protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => - new(session.StateBag.Serialize()); - - protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => - new(new TurnCountingSession(AgentSessionStateBag.Deserialize(serializedState))); - - protected override Task RunCoreAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) => - throw new NotSupportedException(); - protected override async IAsyncEnumerable RunCoreStreamingAsync( IEnumerable messages, AgentSession? session = null, @@ -570,13 +582,45 @@ protected override async IAsyncEnumerable RunCoreStreamingA } } - [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")] - private sealed class RequestDependency; + private sealed class ShutdownProbe + { + public TaskCompletionSource AgentStarted { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource RunFinished { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public bool RunCanceled { get; set; } + } - private sealed class TurnCountingAgent(string? name) : AIAgent + private sealed class WaitForCancellationAgent(ShutdownProbe probe) : StreamingTestAgent { - public override string? Name => name; + public override string? Name => "shutdown"; + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + probe.AgentStarted.TrySetResult(); + try + { + await Task.Delay(Timeout.Infinite, cancellationToken).ConfigureAwait(false); + } + finally + { + probe.RunCanceled = cancellationToken.IsCancellationRequested; + probe.RunFinished.TrySetResult(); + } + + yield break; + } + } + + /// + /// An agent that persists its session state bag and only supports streaming runs. + /// + private abstract class StreamingTestAgent : AIAgent + { protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => new(new TurnCountingSession()); @@ -588,6 +632,14 @@ protected override ValueTask DeserializeSessionCoreAsync(JsonEleme protected override Task RunCoreAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) => throw new NotSupportedException(); + } + + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")] + private sealed class RequestDependency; + + private sealed class TurnCountingAgent(string? name) : StreamingTestAgent + { + public override string? Name => name; protected override async IAsyncEnumerable RunCoreStreamingAsync( IEnumerable messages, From 6d73d581a23aabf244013ecb4653c519f0e4c1b9 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Mon, 19 Oct 2026 11:07:52 +0000 Subject: [PATCH 6/6] .NET: Reject ReturnImmediately task continuations on A2A selector endpoints The ReturnImmediately guard of the selector endpoints only covered new messages, so continuing an existing task with ReturnImmediately invoked the selected agent with a run that would outlive the request. Such continuations are now rejected with UnsupportedOperation before A2AServer starts their execution, which keeps the continued task and its history unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../A2AEndpointRouteBuilderExtensions.cs | 12 ++++---- .../SelectedAgentA2ARequestHandler.cs | 8 +++++ .../A2AAgentHandler.cs | 18 +++++++++--- .../A2ARequestScopedAgentTests.cs | 29 +++++++++++++++++++ 4 files changed, 57 insertions(+), 10 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs index 1f776204508..ceaf76d67eb 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/A2AEndpointRouteBuilderExtensions.cs @@ -123,12 +123,12 @@ public static IEndpointConventionBuilder MapA2AHttpJson(this IEndpointRouteBuild /// identifies the A2A task state shared by the requests that select it through this endpoint, so conversations /// and tasks continue across requests even when the selector returns a new agent instance each time. The selected /// agent must therefore have a non-empty name. The session store is resolved from - /// , so the run cannot outlive the request: when the configured - /// returns a task, non-streaming requests with ReturnImmediately enabled are - /// rejected with , and a streaming request that ends early, for - /// example because the client disconnects, completes only after the agent run finishes. The run is not canceled, so - /// its task state is still persisted. When the application stops, operations that are still running for selected - /// agents are canceled and awaited. + /// , so the run cannot outlive the request: non-streaming requests with + /// ReturnImmediately enabled are rejected with when they + /// continue an existing task or when the configured returns a task, and a streaming + /// request that ends early, for example because the client disconnects, completes only after the agent run + /// finishes. The run is not canceled, so its task state is still persisted. When the application stops, operations + /// that are still running for selected agents are canceled and awaited. /// /// /// Selection is not an authorization boundary. See diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs index 588ecdc36ef..846225ef3df 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/SelectedAgentA2ARequestHandler.cs @@ -100,6 +100,14 @@ static async Task DisposeServersAsync(SelectedAgentServer[] servers) /// public async Task SendMessageAsync(SendMessageRequest request, CancellationToken cancellationToken = default) { + // A continuation always returns a task, so with ReturnImmediately its run would outlive the request. Reject it + // before A2AServer starts the execution, which would otherwise append the message to the task history and fail + // the task when the agent handler rejects the run. + if (request.Configuration?.ReturnImmediately is true && !string.IsNullOrEmpty(request.Message.TaskId)) + { + throw A2AAgentHandler.CreateReturnImmediatelyNotSupportedException(); + } + Selection selection = this.GetSelection(); RequestBoundAgentHandler handler = selection.BindAgent(request.Message); try diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs index b3ec933be67..aef30d3cfb0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/A2AAgentHandler.cs @@ -36,8 +36,10 @@ internal sealed class A2AAgentHandler : IAgentHandler /// The hosted agent that provides the execution logic. /// Controls which A2A artifact the agent response is returned as. /// - /// to reject non-streaming requests with ReturnImmediately = true when + /// to reject new non-streaming messages with ReturnImmediately = true when /// returns a task, because the run would continue after the request completes. + /// Task continuations always return a task, so callers must reject them before the server starts their execution; + /// rejecting them here would fail the task being continued. /// public A2AAgentHandler( AIHostAgent hostAgent, @@ -90,6 +92,16 @@ internal static async Task CancelTaskAsync(RequestContext context, AgentEventQue await taskUpdater.CancelAsync(cancellationToken: cancellationToken).ConfigureAwait(false); } + /// + /// Creates the error returned for non-streaming requests with ReturnImmediately = true that would return a + /// task while the agent is selected for each request. + /// + /// An with . + internal static A2AException CreateReturnImmediatelyNotSupportedException() => + new( + "ReturnImmediately is not supported when the agent is selected for each request, because the selected agent and its services are bound to the HTTP request. Send the message with ReturnImmediately disabled.", + A2AErrorCode.UnsupportedOperation); + /// /// Runs the agent for a new message and emits the response events, shared by the streaming and non-streaming endpoints. /// @@ -159,9 +171,7 @@ private async Task HandleNewMessageAsync(RequestContext context, AgentEventQueue if (returnTask && !aggregateTaskUpdates && !context.StreamingResponse && !this._allowReturnImmediately) { - throw new A2AException( - "ReturnImmediately is not supported when the agent is selected for each request, because the selected agent and its services are bound to the HTTP request. Send the message with ReturnImmediately disabled.", - A2AErrorCode.UnsupportedOperation); + throw CreateReturnImmediatelyNotSupportedException(); } var options = CreateRunOptions(context); diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs index c3348e6a5e6..b8672df2e7f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2ARequestScopedAgentTests.cs @@ -237,6 +237,35 @@ public async Task MapA2AJsonRpc_WithAgentSelectorReturningTask_RejectsReturnImme Assert.Equal(A2AErrorCode.UnsupportedOperation, exception.ErrorCode); } + [Fact] + public async Task MapA2AHttpJson_WithAgentSelectorContinuingTask_RejectsReturnImmediatelyWithoutChangingTaskAsync() + { + // Arrange + InMemoryTaskStore taskStore = new(); + await taskStore.SaveTaskAsync("task-1", new AgentTask + { + Id = "task-1", + ContextId = "ctx-1", + Status = new global::A2A.TaskStatus { State = TaskState.InputRequired, Timestamp = DateTimeOffset.UtcNow } + }); + await using WebApplication app = await StartAsync( + services => services.AddKeyedSingleton("tasks", taskStore), + endpoints => endpoints.MapA2AHttpJson("/a2a", _ => new TurnCountingAgent("tasks"))); + using HttpClient httpClient = app.GetTestClient(); + using A2AHttpJsonClient client = new(new Uri(httpClient.BaseAddress!, "/a2a"), httpClient); + SendMessageRequest request = CreateRequest("ctx-1", returnImmediately: true); + request.Message.TaskId = "task-1"; + + // Act + A2AException exception = await Assert.ThrowsAsync(() => client.SendMessageAsync(request)); + AgentTask task = await client.GetTaskAsync(new GetTaskRequest { Id = "task-1" }); + + // Assert + Assert.Equal(A2AErrorCode.UnsupportedOperation, exception.ErrorCode); + Assert.Equal(TaskState.InputRequired, task.Status.State); + Assert.Empty(task.History ?? []); + } + [Fact] public async Task MapA2AHttpJson_WithAgentSelectorReturningMessage_IgnoresReturnImmediatelyAsync() {