.NET: Honor hosted agent lifetimes and add request-time agent selection for AG-UI hosting

dotnet/eng/MSBuild/Shared.props

@@ -38,6 +38,9 @@
   <ItemGroup Condition="'$(InjectSharedOriginPinning)' == 'true'">
     <Compile Include="$(MSBuildThisFileDirectory)\..\..\src\Shared\OriginPinning\*.cs" LinkBase="Shared\OriginPinning" />
   </ItemGroup>
+  <ItemGroup Condition="'$(InjectSharedHostedAgents)' == 'true'">
+    <Compile Include="$(MSBuildThisFileDirectory)\..\..\src\Shared\HostedAgents\*.cs" LinkBase="Shared\HostedAgents" />
+  </ItemGroup>
   <ItemGroup Condition="'$(InjectSharedHttpHeaderValidation)' == 'true'">
     <Compile Include="$(MSBuildThisFileDirectory)\..\..\src\Shared\Http\*.cs" LinkBase="Shared\Http" />
   </ItemGroup>

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

@@ -40,8 +40,16 @@ public static class AGUIEndpointRouteBuilderExtensions
     /// <param name="pattern">The URL pattern for the endpoint.</param>
     /// <returns>An <see cref="IEndpointConventionBuilder"/> for the mapped endpoint.</returns>
     /// <remarks>
+    /// <para>
+    /// Singleton registrations are resolved once when the endpoint is mapped. Scoped and transient registrations are
+    /// resolved for each request as described in <see cref="MapAGUIServer(IEndpointRouteBuilder, string, string)"/>.
+    /// In both cases, the session store is resolved for each request and persisted sessions are partitioned by the
+    /// registration name.
+    /// </para>
+    /// <para>
     /// See <see cref="MapAGUIServer(IEndpointRouteBuilder, string, AIAgent)"/> for authentication,
     /// authorization, and caller-scoped session isolation requirements.
+    /// </para>
     /// </remarks>
     public static IEndpointConventionBuilder MapAGUIServer(
         this IEndpointRouteBuilder endpoints,
@@ -50,7 +58,16 @@ public static IEndpointConventionBuilder MapAGUIServer(
     {
         ArgumentNullException.ThrowIfNull(endpoints);
         ArgumentNullException.ThrowIfNull(agentBuilder);
-        return endpoints.MapAGUIServer(agentBuilder.Name, pattern);
+
+        if (agentBuilder.Lifetime != ServiceLifetime.Singleton)
+        {
+            return endpoints.MapAGUIServer(agentBuilder.Name, pattern);
+        }
+
+        string agentName = agentBuilder.Name;
+        var agent = endpoints.ServiceProvider.GetRequiredKeyedService<AIAgent>(agentName);
+        return MapAGUIServerCore(endpoints, pattern, context =>
+            CreateHostAgent(context.RequestServices, agent, agentName, sessionStorageIdentity: agentName));
     }
 
     /// <summary>
@@ -61,8 +78,17 @@ public static IEndpointConventionBuilder MapAGUIServer(
     /// <param name="pattern">The URL pattern for the endpoint.</param>
     /// <returns>An <see cref="IEndpointConventionBuilder"/> for the mapped endpoint.</returns>
     /// <remarks>
+    /// <para>
+    /// The keyed <see cref="AIAgent"/>, its keyed <see cref="AgentSessionStore"/>, and the
+    /// <see cref="AgentIsolationKeyProvider"/> are resolved from <see cref="HttpContext.RequestServices"/>
+    /// for each request, so the configured DI lifetime of the agent registration is honored.
+    /// Persisted sessions are partitioned by <paramref name="agentName"/>, which remains stable across
+    /// scoped or transient agent instances.
+    /// </para>
+    /// <para>
     /// See <see cref="MapAGUIServer(IEndpointRouteBuilder, string, AIAgent)"/> for authentication,
     /// authorization, and caller-scoped session isolation requirements.
+    /// </para>
     /// </remarks>
     public static IEndpointConventionBuilder MapAGUIServer(
         this IEndpointRouteBuilder endpoints,
@@ -72,8 +98,63 @@ public static IEndpointConventionBuilder MapAGUIServer(
         ArgumentNullException.ThrowIfNull(endpoints);
         ArgumentNullException.ThrowIfNull(agentName);
 
-        var agent = endpoints.ServiceProvider.GetRequiredKeyedService<AIAgent>(agentName);
-        return endpoints.MapAGUIServer(pattern, agent);
+        HostedAgentResolution.EnsureAgentRegistered(endpoints.ServiceProvider, agentName);
+
+        return MapAGUIServerCore(endpoints, pattern, context =>
+        {
+            var agent = context.RequestServices.GetRequiredKeyedService<AIAgent>(agentName);
+            return CreateHostAgent(context.RequestServices, agent, agentName, sessionStorageIdentity: agentName);
+        });
+    }
+
+    /// <summary>
+    /// Maps an AG-UI agent endpoint that selects the agent to invoke for each request.
+    /// </summary>
+    /// <param name="endpoints">The endpoint route builder.</param>
+    /// <param name="pattern">The URL pattern for the endpoint.</param>
+    /// <param name="agentSelector">
+    /// A delegate invoked once per request that returns the agent to invoke, or <see langword="null"/> when
+    /// no agent matches the request. The delegate can use route values, claims, headers, and
+    /// <see cref="HttpContext.RequestServices"/> to select or resolve the agent.
+    /// </param>
+    /// <returns>An <see cref="IEndpointConventionBuilder"/> for the mapped endpoint.</returns>
+    /// <remarks>
+    /// <para>
+    /// When <paramref name="agentSelector"/> returns <see langword="null"/>, the endpoint responds with
+    /// <c>404 Not Found</c> without loading session state.
+    /// </para>
+    /// <para>
+    /// The selected agent's <see cref="AIAgent.Name"/> is its stable logical identity: it is used to resolve the
+    /// keyed <see cref="AgentSessionStore"/> from <see cref="HttpContext.RequestServices"/> and to partition
+    /// persisted sessions, so conversations continue across requests even when the selector returns a new agent
+    /// instance each time. The selected agent must therefore have a non-empty name, and agents that must not
+    /// share persisted sessions must have different names.
+    /// </para>
+    /// <para>
+    /// Selection is not an authorization boundary. Enforce authorization with ASP.NET Core authorization
+    /// policies, or return <see langword="null"/> from <paramref name="agentSelector"/> for agents the caller must
+    /// not reach. See <see cref="MapAGUIServer(IEndpointRouteBuilder, string, AIAgent)"/> for authentication,
+    /// authorization, and caller-scoped session isolation requirements.
+    /// </para>
+    /// </remarks>
+    public static IEndpointConventionBuilder MapAGUIServer(
+        this IEndpointRouteBuilder endpoints,
+        [StringSyntax("route")] string pattern,
+        Func<HttpContext, AIAgent?> agentSelector)
+    {
+        ArgumentNullException.ThrowIfNull(endpoints);
+        ArgumentNullException.ThrowIfNull(agentSelector);
+
+        return MapAGUIServerCore(endpoints, pattern, context =>
+        {
+            if (agentSelector(context) is not { } agent)
+            {
+                return null;
+            }
+
+            string agentName = HostedAgentResolution.GetSelectedAgentName(agent);
+            return CreateHostAgent(context.RequestServices, agent, agentName, sessionStorageIdentity: agentName);
+        });
     }
 
     /// <summary>
@@ -88,6 +169,15 @@ public static IEndpointConventionBuilder MapAGUIServer(
     /// 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).
+    /// The supplied agent instance is used for every request, while the session store is resolved from
+    /// <see cref="HttpContext.RequestServices"/> for each request.
+    /// </para>
+    /// <para>
+    /// The agent's <see cref="AIAgent.Name"/> is its stable logical identity: it keys the session store and partitions
+    /// persisted sessions, as it does for the named, hosted, and selector mappings. The same agent exposed through
+    /// several mappings therefore continues the same conversations for a given thread ID, and so does any agent
+    /// instance with the same name that uses the same session store. Agents that must not share persisted sessions must
+    /// have different names. Sessions of an agent without a name are partitioned by its instance identity.
     /// </para>
     /// <para>
     /// <strong>Trust model.</strong> The AG-UI <c>RunAgentInput.ThreadId</c> arrives
@@ -125,18 +215,18 @@ public static IEndpointConventionBuilder MapAGUIServer(
         ArgumentNullException.ThrowIfNull(endpoints);
         ArgumentNullException.ThrowIfNull(aiAgent);
 
-        var agentSessionStore = endpoints.ServiceProvider.GetKeyedService<AgentSessionStore>(aiAgent.Name);
-
-        // Ensure that we have an IsolationKeyScopedAgentSessionStore registered.
-        var isolationKeyProvider = endpoints.ServiceProvider.GetService<AgentIsolationKeyProvider>();
-        if (agentSessionStore?.GetService<IsolationKeyScopedAgentSessionStore>() is null)
-        {
-            agentSessionStore ??= new NoopAgentSessionStore();
-            agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = isolationKeyProvider != null });
-        }
-
-        var hostAgent = new AIHostAgent(aiAgent, agentSessionStore);
+        // A named agent is identified by its name, like the named, hosted, and selector mappings, so the same agent
+        // exposed through several mappings continues the same conversations. Unnamed agents keep their instance identity.
+        string? sessionStorageIdentity = string.IsNullOrWhiteSpace(aiAgent.Name) ? null : aiAgent.Name;
+        return MapAGUIServerCore(endpoints, pattern, context =>
+            CreateHostAgent(context.RequestServices, aiAgent, aiAgent.Name, sessionStorageIdentity));
+    }
 
+    private static IEndpointConventionBuilder MapAGUIServerCore(
+        IEndpointRouteBuilder endpoints,
+        string pattern,
+        Func<HttpContext, AIHostAgent?> hostAgentFactory)
+    {
         IEndpointConventionBuilder endpoint = endpoints.MapPost(pattern, async (
             [FromBody] RunAgentInput? input,
             [FromServices] IOptions<Microsoft.AspNetCore.Http.Json.JsonOptions> jsonOptions,
@@ -148,6 +238,12 @@ public static IEndpointConventionBuilder MapAGUIServer(
                 return Results.BadRequest();
             }
 
+            // Resolve the agent before any session state is loaded so an unmatched selection never touches storage.
+            if (hostAgentFactory(context) is not { } hostAgent)
+            {
+                return Results.NotFound();
+            }
+
             var jsonSerializerOptions = jsonOptions.Value.SerializerOptions;
             var streamOptions = context.GetEndpoint()?.Metadata.GetMetadata<AGUIStreamOptions>()
                 ?? context.RequestServices.GetService<IOptions<AGUIStreamOptions>>()?.Value;
@@ -190,6 +286,25 @@ public static IEndpointConventionBuilder MapAGUIServer(
         return endpoint;
     }
 
+    private static AIHostAgent CreateHostAgent(
+        IServiceProvider services,
+        AIAgent agent,
+        string? sessionStoreKey,
+        string? sessionStorageIdentity)
+    {
+        var agentSessionStore = services.GetKeyedService<AgentSessionStore>(sessionStoreKey);
+
+        // Ensure that we have an IsolationKeyScopedAgentSessionStore registered.
+        var isolationKeyProvider = services.GetService<AgentIsolationKeyProvider>();
+        if (agentSessionStore?.GetService<IsolationKeyScopedAgentSessionStore>() is null)
+        {
+            agentSessionStore ??= new NoopAgentSessionStore();
+            agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = isolationKeyProvider != null });
+        }
+
+        return new AIHostAgent(agent, agentSessionStore, sessionStorageIdentity);
+    }
+
     private static void MarkFeatureUsed()
     {
 #pragma warning disable MAAI001

dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj

@@ -7,6 +7,7 @@
     <NoWarn>$(NoWarn);MAAI001</NoWarn>
     <InterceptorsNamespaces>$(InterceptorsNamespaces);Microsoft.AspNetCore.Http.Generated</InterceptorsNamespaces>
     <EnableRequestDelegateGenerator>true</EnableRequestDelegateGenerator>
+    <InjectSharedHostedAgents>true</InjectSharedHostedAgents>
   </PropertyGroup>
 
   <Import Project="$(RepoRoot)/dotnet/nuget/nuget-package.props" />

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;
+
+/// <summary>
+/// Provides validation shared by protocol hosting integrations that resolve agents for each request.
+/// </summary>
+internal static class HostedAgentResolution
+{
+    /// <summary>
+    /// Verifies that a keyed <see cref="AIAgent"/> registration exists without resolving it.
+    /// </summary>
+    /// <param name="services">The application service provider.</param>
+    /// <param name="agentName">The name of the keyed agent registration.</param>
+    /// <exception cref="InvalidOperationException">No agent is registered with <paramref name="agentName"/>.</exception>
+    /// <remarks>
+    /// 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.
+    /// </remarks>
+    public static void EnsureAgentRegistered(IServiceProvider services, string agentName)
+    {
+        if (services.GetService<IServiceProviderIsKeyedService>() 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.");
+        }
+    }
+
+    /// <summary>
+    /// Gets the stable logical identity of an agent returned by a request-time agent selector.
+    /// </summary>
+    /// <param name="agent">The selected agent.</param>
+    /// <returns>The agent name, which identifies the agent's hosted services and persisted sessions.</returns>
+    /// <exception cref="InvalidOperationException">The agent does not have a name.</exception>
+    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!;
+    }
+}