.NET: Update AGUI service to support session storage (#5193)

* Update AGUI service to support session storage * Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Address PR comments --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-04-13 18:03:51 +00:00
parent 76fe7319e0
commit b1fb63eb81
6 changed files with 509 additions and 7 deletions
@@ -1,9 +1,12 @@
 // Copyright (c) Microsoft. All rights reserved.

+using System;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.Shared;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Http;
@@ -21,6 +24,42 @@ namespace Microsoft.Agents.AI.Hosting.AGUI.AspNetCore;
 /// </summary>
 public static class AGUIEndpointRouteBuilderExtensions
 {
+    /// <summary>
+    /// Maps an AG-UI agent endpoint using an agent registered in dependency injection via <see cref="IHostedAgentBuilder"/>.
+    /// </summary>
+    /// <param name="endpoints">The endpoint route builder.</param>
+    /// <param name="agentBuilder">The hosted agent builder that identifies the agent registration.</param>
+    /// <param name="pattern">The URL pattern for the endpoint.</param>
+    /// <returns>An <see cref="IEndpointConventionBuilder"/> for the mapped endpoint.</returns>
+    public static IEndpointConventionBuilder MapAGUI(
+        this IEndpointRouteBuilder endpoints,
+        IHostedAgentBuilder agentBuilder,
+        [StringSyntax("route")] string pattern)
+    {
+        ArgumentNullException.ThrowIfNull(endpoints);
+        ArgumentNullException.ThrowIfNull(agentBuilder);
+        return endpoints.MapAGUI(agentBuilder.Name, pattern);
+    }
+
+    /// <summary>
+    /// Maps an AG-UI agent endpoint using a named agent registered in dependency injection.
+    /// </summary>
+    /// <param name="endpoints">The endpoint route builder.</param>
+    /// <param name="agentName">The name of the keyed agent registration to resolve from dependency injection.</param>
+    /// <param name="pattern">The URL pattern for the endpoint.</param>
+    /// <returns>An <see cref="IEndpointConventionBuilder"/> for the mapped endpoint.</returns>
+    public static IEndpointConventionBuilder MapAGUI(
+        this IEndpointRouteBuilder endpoints,
+        string agentName,
+        [StringSyntax("route")] string pattern)
+    {
+        ArgumentNullException.ThrowIfNull(endpoints);
+        ArgumentNullException.ThrowIfNull(agentName);
+
+        var agent = endpoints.ServiceProvider.GetRequiredKeyedService<AIAgent>(agentName);
+        return endpoints.MapAGUI(pattern, agent);
+    }
+
    /// <summary>
    /// Maps an AG-UI agent endpoint.
    /// </summary>
@@ -28,11 +67,24 @@ public static class AGUIEndpointRouteBuilderExtensions
    /// <param name="pattern">The URL pattern for the endpoint.</param>
    /// <param name="aiAgent">The agent instance.</param>
    /// <returns>An <see cref="IEndpointConventionBuilder"/> for the mapped endpoint.</returns>
+    /// <remarks>
+    /// <para>
+    /// If an <see cref="AgentSessionStore"/> is registered in dependency injection keyed by the agent's name,
+    /// it will be used to persist conversation sessions across requests using the AG-UI thread ID as the
+    /// conversation identifier. If no session store is registered, sessions are ephemeral (not persisted).
+    /// </para>
+    /// </remarks>
    public static IEndpointConventionBuilder MapAGUI(
        this IEndpointRouteBuilder endpoints,
        [StringSyntax("route")] string pattern,
        AIAgent aiAgent)
    {
+        ArgumentNullException.ThrowIfNull(endpoints);
+        ArgumentNullException.ThrowIfNull(aiAgent);
+
+        var agentSessionStore = endpoints.ServiceProvider.GetKeyedService<AgentSessionStore>(aiAgent.Name);
+        var hostAgent = new AIHostAgent(aiAgent, agentSessionStore ?? new NoopAgentSessionStore());
+
        return endpoints.MapPost(pattern, async ([FromBody] RunAgentInput? input, HttpContext context, CancellationToken cancellationToken) =>
        {
            if (input is null)
@@ -63,21 +115,43 @@ public static class AGUIEndpointRouteBuilderExtensions
                }
            };

+            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 = aiAgent.RunStreamingAsync(
+            var events = hostAgent.RunStreamingAsync(
                messages,
+                session: session,
                options: runOptions,
                cancellationToken: cancellationToken)
                .AsChatResponseUpdatesAsync()
                .FilterServerToolsFromMixedToolInvocationsAsync(clientTools, cancellationToken)
                .AsAGUIEventStreamAsync(
-                    input.ThreadId,
+                    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<ILogger<AGUIServerSentEventsResult>>();
-            return new AGUIServerSentEventsResult(events, sseLogger);
+            return new AGUIServerSentEventsResult(eventsWithSessionSave, sseLogger);
        });
    }
+
+    private static async IAsyncEnumerable<BaseEvent> SaveSessionAfterStreamingAsync(
+        IAsyncEnumerable<BaseEvent> events,
+        AIHostAgent hostAgent,
+        string threadId,
+        AgentSession session,
+        [EnumeratorCancellation] CancellationToken cancellationToken)
+    {
+        await foreach (BaseEvent evt in events.ConfigureAwait(false))
+        {
+            yield return evt;
+        }
+
+        await hostAgent.SaveSessionAsync(threadId, session, cancellationToken).ConfigureAwait(false);
+    }
 }
@@ -19,6 +19,7 @@

  <ItemGroup>
    <ProjectReference Include="..\Microsoft.Agents.AI\Microsoft.Agents.AI.csproj" />
+    <ProjectReference Include="..\Microsoft.Agents.AI.Hosting\Microsoft.Agents.AI.Hosting.csproj" />
  </ItemGroup>

  <ItemGroup>