diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs index 0d4c390bbb8..04883f48c00 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; @@ -121,52 +121,130 @@ public static IEndpointConventionBuilder MapAGUI( return Results.BadRequest(); } - var jsonOptions = context.RequestServices.GetRequiredService>(); - var jsonSerializerOptions = jsonOptions.Value.SerializerOptions; + return await ExecuteAgentRequestAsync(hostAgent, input, context, cancellationToken).ConfigureAwait(false); + }); + } + + /// + /// Maps an AG-UI agent endpoint using a factory delegate for per-request agent resolution. + /// This enables dynamic, multi-tenant agent hosting where the agent is selected based on + /// route parameters, request headers, claims, or other per-request information. + /// + /// The endpoint route builder. + /// The URL pattern for the endpoint (e.g., "/agents/{agentId}"). + /// + /// A factory delegate that resolves an for each request. + /// The delegate receives the current and a . + /// Return null to produce a 404 Not Found response. + /// + /// An for the mapped endpoint. + /// + /// + /// Unlike the static overload, + /// this method does not capture the agent at startup. Instead, the agent and its + /// are resolved per-request from the factory delegate and + /// respectively. This fixes the singleton-capture + /// issue where scoped or transient session stores were inadvertently captured at startup. + /// + /// + /// Trust model. See remarks on + /// for session isolation guidance. + /// + /// + public static IEndpointConventionBuilder MapAGUI( + this IEndpointRouteBuilder endpoints, + [StringSyntax("route")] string pattern, + Func> agentFactory) + { + ArgumentNullException.ThrowIfNull(endpoints); + ArgumentNullException.ThrowIfNull(agentFactory); + + return endpoints.MapPost(pattern, async ([FromBody] RunAgentInput? input, HttpContext context, CancellationToken cancellationToken) => + { + if (input is null) + { + return Results.BadRequest(); + } + + // Resolve agent per-request via factory delegate + var aiAgent = await agentFactory(context, cancellationToken).ConfigureAwait(false); + if (aiAgent is null) + { + return Results.NotFound(); + } + + // Resolve session store per-request from the request's DI scope (not app-level) + var agentSessionStore = context.RequestServices.GetKeyedService(aiAgent.Name); + + // Ensure that we have an IsolationKeyScopedAgentSessionStore registered. + var isolationKeyProvider = context.RequestServices.GetService(); + if (agentSessionStore?.GetService() is null) + { + agentSessionStore ??= new NoopAgentSessionStore(); + agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = isolationKeyProvider != null }); + } + + var hostAgent = new AIHostAgent(aiAgent, agentSessionStore); + + return await ExecuteAgentRequestAsync(hostAgent, input, context, cancellationToken).ConfigureAwait(false); + }); + } - var messages = input.Messages.AsChatMessages(jsonSerializerOptions); - var clientTools = input.Tools?.AsAITools().ToList(); + /// + /// Shared execution pipeline for AG-UI agent requests. Converts the input to chat messages, + /// runs the agent, and returns an SSE result with session persistence. + /// + private static async Task ExecuteAgentRequestAsync( + AIHostAgent hostAgent, + RunAgentInput input, + HttpContext context, + CancellationToken cancellationToken) + { + var jsonOptions = context.RequestServices.GetRequiredService>(); + var jsonSerializerOptions = jsonOptions.Value.SerializerOptions; + + var messages = input.Messages.AsChatMessages(jsonSerializerOptions); + var clientTools = input.Tools?.AsAITools().ToList(); - // Create run options with AG-UI context in AdditionalProperties - var runOptions = new ChatClientAgentRunOptions + // Create run options with AG-UI context in AdditionalProperties + var runOptions = new ChatClientAgentRunOptions + { + ChatOptions = new ChatOptions { - ChatOptions = new ChatOptions + Tools = clientTools, + AdditionalProperties = new AdditionalPropertiesDictionary { - Tools = clientTools, - AdditionalProperties = new AdditionalPropertiesDictionary - { - ["ag_ui_state"] = input.State, - ["ag_ui_context"] = input.Context?.Select(c => new KeyValuePair(c.Description, c.Value)).ToArray(), - ["ag_ui_forwarded_properties"] = input.ForwardedProperties, - ["ag_ui_thread_id"] = input.ThreadId, - ["ag_ui_run_id"] = input.RunId - } + ["ag_ui_state"] = input.State, + ["ag_ui_context"] = input.Context?.Select(c => new KeyValuePair(c.Description, c.Value)).ToArray(), + ["ag_ui_forwarded_properties"] = input.ForwardedProperties, + ["ag_ui_thread_id"] = input.ThreadId, + ["ag_ui_run_id"] = input.RunId } - }; - - var threadId = string.IsNullOrWhiteSpace(input.ThreadId) ? Guid.NewGuid().ToString("N") : input.ThreadId; - var session = await hostAgent.GetOrCreateSessionAsync(threadId, cancellationToken).ConfigureAwait(false); - - // Run the agent and convert to AG-UI events - var events = hostAgent.RunStreamingAsync( - messages, - session: session, - options: runOptions, - cancellationToken: cancellationToken) - .AsChatResponseUpdatesAsync() - .FilterServerToolsFromMixedToolInvocationsAsync(clientTools, cancellationToken) - .AsAGUIEventStreamAsync( - threadId, - input.RunId, - jsonSerializerOptions, - cancellationToken); - - // Wrap the event stream to save the session after streaming completes - var eventsWithSessionSave = SaveSessionAfterStreamingAsync(events, hostAgent, threadId, session, cancellationToken); - - var sseLogger = context.RequestServices.GetRequiredService>(); - return new AGUIServerSentEventsResult(eventsWithSessionSave, sseLogger); - }); + } + }; + + var threadId = string.IsNullOrWhiteSpace(input.ThreadId) ? Guid.NewGuid().ToString("N") : input.ThreadId; + var session = await hostAgent.GetOrCreateSessionAsync(threadId, cancellationToken).ConfigureAwait(false); + + // Run the agent and convert to AG-UI events + var events = hostAgent.RunStreamingAsync( + messages, + session: session, + options: runOptions, + cancellationToken: cancellationToken) + .AsChatResponseUpdatesAsync() + .FilterServerToolsFromMixedToolInvocationsAsync(clientTools, cancellationToken) + .AsAGUIEventStreamAsync( + threadId, + input.RunId, + jsonSerializerOptions, + cancellationToken); + + // Wrap the event stream to save the session after streaming completes + var eventsWithSessionSave = SaveSessionAfterStreamingAsync(events, hostAgent, threadId, session, cancellationToken); + + var sseLogger = context.RequestServices.GetRequiredService>(); + return new AGUIServerSentEventsResult(eventsWithSessionSave, sseLogger); } private static async IAsyncEnumerable SaveSessionAfterStreamingAsync( diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs index 248629b3924..66cb48ee1ca 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; @@ -747,4 +747,55 @@ protected override async IAsyncEnumerable RunCoreStreamingA yield return new AgentResponseUpdate(new ChatResponseUpdate(ChatRole.Assistant, "Test response")); } } + + #region Factory Delegate Overload Tests + + [Fact] + public void MapAGUI_WithFactoryDelegate_MapsEndpoint_AtSpecifiedPattern() + { + // Arrange + Mock endpointsMock = new(); + Mock serviceProviderMock = new(); + serviceProviderMock.As(); + + endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object); + endpointsMock.Setup(e => e.DataSources).Returns([]); + + const string Pattern = "/agents/{agentId}"; + + // Act + IEndpointConventionBuilder? result = endpointsMock.Object.MapAGUI( + Pattern, + (HttpContext context, CancellationToken ct) => new ValueTask(new TestAgent())); + + // Assert + Assert.NotNull(result); + } + + [Fact] + public void MapAGUI_WithNullFactory_ThrowsArgumentNullException() + { + // Arrange + Mock endpointsMock = new(); + Mock serviceProviderMock = new(); + serviceProviderMock.As(); + endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object); + + // Act & Assert + Assert.Throws(() => + endpointsMock.Object.MapAGUI("/agents/{agentId}", (Func>)null!)); + } + + [Fact] + public void MapAGUI_WithFactoryDelegate_AndNullEndpoints_ThrowsArgumentNullException() + { + // Act & Assert + Assert.Throws(() => + AGUIEndpointRouteBuilderExtensions.MapAGUI( + null!, + "/agents/{agentId}", + (HttpContext context, CancellationToken ct) => new ValueTask(new TestAgent()))); + } + + #endregion }