.NET: Add MapAGUI overload that resolves the agent per request via a factory

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;
@@ -110,66 +110,143 @@ public static IEndpointConventionBuilder MapAGUIServer(
         ArgumentNullException.ThrowIfNull(endpoints);
         ArgumentNullException.ThrowIfNull(aiAgent);
 
-        var agentSessionStore = endpoints.ServiceProvider.GetKeyedService<AgentSessionStore>(aiAgent.Name);
+        var hostAgent = BuildHostAgent(endpoints.ServiceProvider, aiAgent, aiAgent.Name);
+
+        return endpoints.MapPost(pattern, (
+            [FromBody] RunAgentInput? input,
+            [FromServices] IOptions<Microsoft.AspNetCore.Http.Json.JsonOptions> jsonOptions,
+            HttpContext context,
+            CancellationToken cancellationToken)
+            => HandleRunAsync(input, hostAgent, jsonOptions, context, cancellationToken));
+    }
+
+    /// <summary>
+    /// Maps an AG-UI agent endpoint that resolves the agent <strong>per request</strong> via a factory
+    /// delegate, rather than capturing a single agent instance at startup.
+    /// </summary>
+    /// <param name="endpoints">The endpoint route builder.</param>
+    /// <param name="agentName">The logical agent name passed to the factory and used as the
+    /// <see cref="AgentSessionStore"/> resolution key and the agent's stable session identity.</param>
+    /// <param name="pattern">The URL pattern for the endpoint.</param>
+    /// <param name="createAgentDelegate">A factory invoked once per request with the request's
+    /// <see cref="IServiceProvider"/> and the <paramref name="agentName"/>, returning the
+    /// <see cref="AIAgent"/> that handles that request.</param>
+    /// <returns>An <see cref="IEndpointConventionBuilder"/> for the mapped endpoint.</returns>
+    /// <remarks>
+    /// <para>
+    /// Use this overload when the agent must be built with request-scoped state — per-request
+    /// authentication, scoped tool/MCP sessions, or conversation-scoped configuration — rather than a
+    /// process-wide singleton resolved once at startup (as the other <c>MapAGUIServer</c> overloads do).
+    /// The factory receives the request's <see cref="IServiceProvider"/>
+    /// (<see cref="HttpContext.RequestServices"/>), so scoped services resolve correctly within the request.
+    /// </para>
+    /// <para>
+    /// The keyed <see cref="AgentSessionStore"/> lookup and the <c>ThreadId</c> trust model are identical
+    /// to <see cref="MapAGUIServer(IEndpointRouteBuilder, string, AIAgent)"/>; the store is resolved per
+    /// request from <see cref="HttpContext.RequestServices"/> using <paramref name="agentName"/> as the key.
+    /// </para>
+    /// <para>
+    /// Because the agent is created per request, this overload assigns it a <strong>stable session
+    /// identity</strong> equal to <paramref name="agentName"/> so that <see cref="AgentSessionStore"/>
+    /// lookups — keyed by <c>(agent.Id, conversationId)</c> — stay consistent across requests and the
+    /// AG-UI <c>ThreadId</c> continues the same persisted session on every turn. Without this, each
+    /// request's agent would report a different per-instance <see cref="AIAgent.Id"/> and the previously
+    /// saved session would never be found. Any id set on the agent returned by
+    /// <paramref name="createAgentDelegate"/> is therefore not used for session keying.
+    /// </para>
+    /// </remarks>
+    public static IEndpointConventionBuilder MapAGUIServer(
+        this IEndpointRouteBuilder endpoints,
+        string agentName,
+        [StringSyntax("route")] string pattern,
+        Func<IServiceProvider, string, AIAgent> createAgentDelegate)
+    {
+        ArgumentNullException.ThrowIfNull(endpoints);
+        ArgumentNullException.ThrowIfNull(agentName);
+        ArgumentNullException.ThrowIfNull(createAgentDelegate);
+
+        return endpoints.MapPost(pattern, (
+            [FromBody] RunAgentInput? input,
+            [FromServices] IOptions<Microsoft.AspNetCore.Http.Json.JsonOptions> jsonOptions,
+            HttpContext context,
+            CancellationToken cancellationToken) =>
+        {
+            var created = createAgentDelegate(context.RequestServices, agentName)
+                ?? throw new InvalidOperationException($"The agent factory for '{agentName}' returned null.");
+
+            // The session store keys persisted sessions by (agent.Id, conversationId), and AIAgent.Id
+            // defaults to a per-instance value. Since this overload builds a fresh agent every request,
+            // wrap it so the id is the stable agentName; otherwise each request would key a different
+            // session and AG-UI thread continuation would silently reset. See StableIdentityAIAgent.
+            var stableAgent = new StableIdentityAIAgent(created, agentName);
+            var hostAgent = BuildHostAgent(context.RequestServices, stableAgent, agentName);
+            return HandleRunAsync(input, hostAgent, jsonOptions, context, cancellationToken);
+        });
+    }
+
+    private static AIHostAgent BuildHostAgent(IServiceProvider services, AIAgent agent, string? sessionStoreKey)
+    {
+        var agentSessionStore = services.GetKeyedService<AgentSessionStore>(sessionStoreKey);
 
         // Ensure that we have an IsolationKeyScopedAgentSessionStore registered.
-        var isolationKeyProvider = endpoints.ServiceProvider.GetService<SessionIsolationKeyProvider>();
+        var isolationKeyProvider = services.GetService<SessionIsolationKeyProvider>();
         if (agentSessionStore?.GetService<IsolationKeyScopedAgentSessionStore>() is null)
         {
             agentSessionStore ??= new NoopAgentSessionStore();
             agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = isolationKeyProvider != null });
         }
 
-        var hostAgent = new AIHostAgent(aiAgent, agentSessionStore);
+        return new AIHostAgent(agent, agentSessionStore);
+    }
 
-        return endpoints.MapPost(pattern, async (
-            [FromBody] RunAgentInput? input,
-            [FromServices] IOptions<Microsoft.AspNetCore.Http.Json.JsonOptions> jsonOptions,
-            HttpContext context,
-            CancellationToken cancellationToken) =>
+    private static async Task<IResult> HandleRunAsync(
+        RunAgentInput? input,
+        AIHostAgent hostAgent,
+        IOptions<Microsoft.AspNetCore.Http.Json.JsonOptions> jsonOptions,
+        HttpContext context,
+        CancellationToken cancellationToken)
+    {
+        if (input is null)
         {
-            if (input is null)
-            {
-                return Results.BadRequest();
-            }
+            return Results.BadRequest();
+        }
 
-            var jsonSerializerOptions = jsonOptions.Value.SerializerOptions;
-            var streamOptions = context.GetEndpoint()?.Metadata.GetMetadata<AGUIStreamOptions>()
-                ?? context.RequestServices.GetService<IOptions<AGUIStreamOptions>>()?.Value;
+        var jsonSerializerOptions = jsonOptions.Value.SerializerOptions;
+        var streamOptions = context.GetEndpoint()?.Metadata.GetMetadata<AGUIStreamOptions>()
+            ?? context.RequestServices.GetService<IOptions<AGUIStreamOptions>>()?.Value;
 
-            var ctx = input.ToChatRequestContext(jsonSerializerOptions, streamOptions);
+        var ctx = input.ToChatRequestContext(jsonSerializerOptions, streamOptions);
 
-            // AG-UI continuation is keyed by thread id. When the client does not supply one, generate a
-            // stable id and write it back onto the input so the persisted session, the RUN_STARTED /
-            // RUN_FINISHED events, and any continuation the client sends back all agree on the same id.
-            var threadId = string.IsNullOrWhiteSpace(ctx.Input.ThreadId) ? Guid.NewGuid().ToString("N") : ctx.Input.ThreadId;
-            ctx.Input.ThreadId = threadId;
+        // AG-UI continuation is keyed by thread id. When the client does not supply one, generate a
+        // stable id and write it back onto the input so the persisted session, the RUN_STARTED /
+        // RUN_FINISHED events, and any continuation the client sends back all agree on the same id.
+        var threadId = string.IsNullOrWhiteSpace(ctx.Input.ThreadId) ? Guid.NewGuid().ToString("N") : ctx.Input.ThreadId;
+        ctx.Input.ThreadId = threadId;
 
-            var session = await hostAgent.GetOrCreateSessionAsync(threadId, cancellationToken).ConfigureAwait(false);
+        var session = await hostAgent.GetOrCreateSessionAsync(threadId, cancellationToken).ConfigureAwait(false);
 
-            var events = hostAgent
-                .RunStreamingAsync(
-                    ctx.Messages,
-                    session: session,
-                    options: new ChatClientAgentRunOptions { ChatOptions = ctx.ChatOptions },
-                    cancellationToken: cancellationToken)
-                .AsChatResponseUpdatesAsync()
-                .AsAGUIEventStreamAsync(ctx, cancellationToken);
+        var events = hostAgent
+            .RunStreamingAsync(
+                ctx.Messages,
+                session: session,
+                options: new ChatClientAgentRunOptions { ChatOptions = ctx.ChatOptions },
+                cancellationToken: cancellationToken)
+            .AsChatResponseUpdatesAsync()
+            .AsAGUIEventStreamAsync(ctx, cancellationToken);
 
-            // Wrap the event stream to save the session after streaming completes.
-            var eventsWithSessionSave = SaveSessionAfterStreamingAsync(events, hostAgent, threadId, session, cancellationToken);
+        // Wrap the event stream to save the session after streaming completes.
+        var eventsWithSessionSave = SaveSessionAfterStreamingAsync(events, hostAgent, threadId, session, cancellationToken);
 
 #if NET10_0_OR_GREATER
-            // On net10+ the framework provides first-class SSE result that flows through the
-            // configured ASP.NET Core JsonSerializerOptions (which AddAGUIServer() augments with
-            // AGUIJsonSerializerContext via the resolver chain).
-            return TypedResults.ServerSentEvents(eventsWithSessionSave);
+        // On net10+ the framework provides first-class SSE result that flows through the
+        // configured ASP.NET Core JsonSerializerOptions (which AddAGUIServer() augments with
+        // AGUIJsonSerializerContext via the resolver chain).
+        return TypedResults.ServerSentEvents(eventsWithSessionSave);
 #else
-            // On older TFMs we ship a small polyfill that emulates TypedResults.ServerSentEvents.
-            var sseLogger = context.RequestServices.GetRequiredService<ILogger<AGUIServerSentEventsResult>>();
-            return new AGUIServerSentEventsResult(eventsWithSessionSave, sseLogger);
+        // On older TFMs we ship a small polyfill that emulates TypedResults.ServerSentEvents.
+        var sseLogger = context.RequestServices.GetRequiredService<ILogger<AGUIServerSentEventsResult>>();
+        return new AGUIServerSentEventsResult(eventsWithSessionSave, sseLogger);
 #endif
-        });
     }
 
     private static async IAsyncEnumerable<BaseEvent> SaveSessionAfterStreamingAsync(

dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/StableIdentityAIAgent.cs

@@ -0,0 +1,36 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+namespace Microsoft.Agents.AI.Hosting.AGUI.AspNetCore;
+
+/// <summary>
+/// Wraps a per-request <see cref="AIAgent"/> so that its <see cref="AIAgent.Id"/> reports a
+/// caller-supplied stable value instead of the per-instance identifier the base type generates.
+/// </summary>
+/// <remarks>
+/// <see cref="AgentSessionStore"/> implementations key persisted sessions by the pair
+/// <c>(agent.Id, conversationId)</c>, and <see cref="AIAgent.Id"/> defaults to a value that is unique
+/// per agent instance. The per-request factory overload of <c>MapAGUIServer</c> builds a fresh agent on
+/// every request, so without a stable identity each request would compute a different session key and
+/// the previously persisted session would never be found — silently resetting AG-UI thread
+/// continuation every turn. Wrapping the per-request agent so its id is the logical agent name keeps
+/// the session key constant across requests, matching the startup-capture overloads where a single
+/// agent instance is reused.
+/// </remarks>
+internal sealed class StableIdentityAIAgent : DelegatingAIAgent
+{
+    private readonly string stableId;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="StableIdentityAIAgent"/> class.
+    /// </summary>
+    /// <param name="innerAgent">The per-request agent to delegate all behavior to.</param>
+    /// <param name="stableId">The stable identifier to report from <see cref="AIAgent.Id"/>.</param>
+    public StableIdentityAIAgent(AIAgent innerAgent, string stableId)
+        : base(innerAgent)
+    {
+        this.stableId = stableId;
+    }
+
+    /// <inheritdoc/>
+    protected override string? IdCore => this.stableId;
+}

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/MapAGUIFactoryDelegateTests.cs

@@ -0,0 +1,115 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Net.Http;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+using AGUI.Client;
+using FluentAssertions;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.TestHost;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests;
+
+/// <summary>
+/// Integration tests for the per-request factory-delegate overload of
+/// <c>MapAGUIServer(endpoints, agentName, pattern, Func&lt;IServiceProvider, string, AIAgent&gt;)</c>.
+/// Unlike the startup-capture overloads, the factory is invoked once per request from the request's
+/// <see cref="IServiceProvider"/>.
+/// </summary>
+public sealed class MapAGUIFactoryDelegateTests : IAsyncDisposable
+{
+    private WebApplication? _app;
+    private HttpClient? _client;
+
+    [Fact]
+    public async Task MapAGUIServer_WithFactoryDelegate_InvokesFactoryPerRequest_AndStreamsAsync()
+    {
+        // Arrange - map the endpoint with a factory that records how many times it is invoked.
+        int factoryInvocations = 0;
+        await this.SetupTestServerWithFactoryAsync((_, name) =>
+        {
+            Interlocked.Increment(ref factoryInvocations);
+            return new FakeSessionAgent(name);
+        });
+
+        var chatClient = new AGUIChatClient(new(this._client!, ""));
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        AgentSession session = await agent.CreateSessionAsync();
+
+        // Act - two turns => two HTTP requests.
+        List<AgentResponseUpdate> firstTurn = [];
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "First")], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            firstTurn.Add(update);
+        }
+
+        List<AgentResponseUpdate> secondTurn = [];
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Second")], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            secondTurn.Add(update);
+        }
+
+        // Assert - the factory ran once per request (not captured once at startup), and the agent streamed.
+        factoryInvocations.Should().Be(2, "the factory delegate is invoked per request");
+        firstTurn.Should().NotBeEmpty();
+        firstTurn.ToAgentResponse().Messages[0].Text.Should().Contain("Hello from session agent");
+    }
+
+    [Fact]
+    public async Task MapAGUIServer_WithFactoryDelegate_WhenFactoryReturnsNull_FailsTheRequestAsync()
+    {
+        // Arrange - a factory that returns null should surface a clear failure when a request arrives.
+        await this.SetupTestServerWithFactoryAsync((_, _) => null!);
+
+        const string Json = """
+            {"threadId":"t1","runId":"r1","messages":[{"id":"m1","role":"user","content":"hi"}],"tools":[],"context":[],"state":{}}
+            """;
+        using StringContent content = new(Json, Encoding.UTF8, "application/json");
+
+        // Act
+        Func<Task> act = async () =>
+        {
+            using HttpResponseMessage response = await this._client!.PostAsync((Uri?)null, content);
+            response.EnsureSuccessStatusCode();
+        };
+
+        // Assert - the factory returning null surfaces a clear InvalidOperationException naming the agent.
+        (await act.Should().ThrowAsync<InvalidOperationException>())
+            .WithMessage("*factory for 'factory-agent' returned null*");
+    }
+
+    private async Task SetupTestServerWithFactoryAsync(Func<IServiceProvider, string, AIAgent> factory)
+    {
+        WebApplicationBuilder builder = WebApplication.CreateBuilder();
+        builder.WebHost.UseTestServer();
+        builder.Services.AddAGUIServer();
+
+        this._app = builder.Build();
+
+        // Per-request factory overload - no keyed AIAgent registration required.
+        this._app.MapAGUIServer("factory-agent", "/agent", factory);
+
+        await this._app.StartAsync();
+
+        TestServer testServer = this._app.Services.GetRequiredService<IServer>() as TestServer
+            ?? throw new InvalidOperationException("TestServer not found");
+
+        this._client = testServer.CreateClient();
+        this._client.BaseAddress = new Uri("http://localhost/agent");
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        this._client?.Dispose();
+        if (this._app != null)
+        {
+            await this._app.DisposeAsync();
+        }
+    }
+}