.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 19:03:51 +01:00
parent 76fe7319e0
commit b1fb63eb81
6 changed files with 509 additions and 7 deletions
@@ -4,6 +4,7 @@ using System.ComponentModel;
 using AGUIServer;
 using Azure.AI.OpenAI;
 using Azure.Identity;
+using Microsoft.Agents.AI.Hosting;
 using Microsoft.Agents.AI.Hosting.AGUI.AspNetCore;
 using Microsoft.Extensions.AI;
 using OpenAI.Chat;
@@ -13,11 +14,11 @@ builder.Services.AddHttpClient().AddLogging();
 builder.Services.ConfigureHttpJsonOptions(options => options.SerializerOptions.TypeInfoResolverChain.Add(AGUIServerSerializerContext.Default));
 builder.Services.AddAGUI();

-WebApplication app = builder.Build();
-
 string endpoint = builder.Configuration["AZURE_OPENAI_ENDPOINT"] ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set.");
 string deploymentName = builder.Configuration["AZURE_OPENAI_DEPLOYMENT_NAME"] ?? throw new InvalidOperationException("AZURE_OPENAI_DEPLOYMENT_NAME is not set.");

+const string AgentName = "AGUIAssistant";
+
 // Create the AI agent with tools
 // WARNING: DefaultAzureCredential is convenient for development but requires careful consideration in production.
 // In production, consider using a specific credential (e.g., ManagedIdentityCredential) to avoid
@@ -27,7 +28,7 @@ var agent = new AzureOpenAIClient(
        new DefaultAzureCredential())
    .GetChatClient(deploymentName)
    .AsAIAgent(
-        name: "AGUIAssistant",
+        name: AgentName,
        tools: [
            AIFunctionFactory.Create(
                () => DateTimeOffset.UtcNow,
@@ -48,7 +49,15 @@ var agent = new AzureOpenAIClient(
                AGUIServerSerializerContext.Default.Options)
        ]);

+// Register the agent with the host and configure it to use an in-memory session store
+// so that conversation state is maintained across requests. In production, you may want to use a persistent session store.
+builder
+    .AddAIAgent(AgentName, (_, _) => agent)
+    .WithInMemorySessionStore();
+
+WebApplication app = builder.Build();
+
 // Map the AG-UI agent endpoint
-app.MapAGUI("/", agent);
+app.MapAGUI(AgentName, "/");

 await app.RunAsync();
@@ -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>
@@ -21,6 +21,7 @@

  <ItemGroup>
    <ProjectReference Include="..\..\src\Microsoft.Agents.AI\Microsoft.Agents.AI.csproj" />
+    <ProjectReference Include="..\..\src\Microsoft.Agents.AI.Hosting\Microsoft.Agents.AI.Hosting.csproj" />
    <ProjectReference Include="..\..\src\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj" />
    <ProjectReference Include="..\..\src\Microsoft.Agents.AI.AGUI\Microsoft.Agents.AI.AGUI.csproj" />
    <ProjectReference Include="..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
@@ -0,0 +1,226 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Net.Http;
+using System.Runtime.CompilerServices;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using System.Threading;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.Agents.AI.AGUI;
+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;
+
+public sealed class SessionPersistenceTests : IAsyncDisposable
+{
+    private WebApplication? _app;
+    private HttpClient? _client;
+
+    [Fact]
+    public async Task MultiTurnWithSessionStore_PersistsSessionAcrossRequestsAsync()
+    {
+        // Arrange - use hosting DI pattern with InMemorySessionStore.
+        // FakeSessionAgent tracks turn count in session StateBag so we can verify
+        // that state survives the serialization round-trip through the session store.
+        await this.SetupTestServerWithSessionStoreAsync();
+        var chatClient = new AGUIChatClient(this._client!, "", null);
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+
+        // Act - First turn
+        ChatMessage firstUserMessage = new(ChatRole.User, "First message");
+        List<AgentResponseUpdate> firstTurnUpdates = [];
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([firstUserMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            firstTurnUpdates.Add(update);
+        }
+
+        // Act - Second turn (same thread ID to test session persistence)
+        ChatMessage secondUserMessage = new(ChatRole.User, "Second message");
+        List<AgentResponseUpdate> secondTurnUpdates = [];
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([secondUserMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            secondTurnUpdates.Add(update);
+        }
+
+        // Assert - Verify turn count proves session state was persisted.
+        // If session persistence were broken, both turns would return "Turn 1"
+        // because a fresh session (with turn count 0) would be created each time.
+        AgentResponse firstResponse = firstTurnUpdates.ToAgentResponse();
+        firstResponse.Messages.Should().HaveCount(1);
+        firstResponse.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        firstResponse.Messages[0].Text.Should().Contain("Turn 1:");
+
+        AgentResponse secondResponse = secondTurnUpdates.ToAgentResponse();
+        secondResponse.Messages.Should().HaveCount(1);
+        secondResponse.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        secondResponse.Messages[0].Text.Should().Contain("Turn 2:");
+    }
+
+    [Fact]
+    public async Task MapAGUI_WithAgentName_StreamsResponseCorrectlyAsync()
+    {
+        // Arrange - use the MapAGUI(agentName, pattern) overload via hosting DI
+        await this.SetupTestServerWithSessionStoreAsync();
+        var chatClient = new AGUIChatClient(this._client!, "", null);
+        AIAgent agent = chatClient.AsAIAgent(instructions: null, name: "assistant", description: "Sample assistant", tools: []);
+        ChatClientAgentSession session = (ChatClientAgentSession)await agent.CreateSessionAsync();
+        ChatMessage userMessage = new(ChatRole.User, "hello");
+
+        List<AgentResponseUpdate> updates = [];
+
+        // Act
+        await foreach (AgentResponseUpdate update in agent.RunStreamingAsync([userMessage], session, new AgentRunOptions(), CancellationToken.None))
+        {
+            updates.Add(update);
+        }
+
+        // Assert
+        updates.Should().NotBeEmpty();
+        updates.Should().AllSatisfy(u => u.Role.Should().Be(ChatRole.Assistant));
+
+        AgentResponse response = updates.ToAgentResponse();
+        response.Messages.Should().HaveCount(1);
+        response.Messages[0].Role.Should().Be(ChatRole.Assistant);
+        response.Messages[0].Text.Should().Be("Turn 1: Hello from session agent!");
+    }
+
+    private async Task SetupTestServerWithSessionStoreAsync()
+    {
+        WebApplicationBuilder builder = WebApplication.CreateBuilder();
+        builder.WebHost.UseTestServer();
+
+        builder.Services.AddAGUI();
+
+        // Register agent using hosting DI pattern with InMemorySessionStore
+        builder.Services.AddAIAgent("session-test-agent", (_, name) => new FakeSessionAgent(name))
+            .WithInMemorySessionStore();
+
+        this._app = builder.Build();
+
+        // Use the agentName overload of MapAGUI
+        this._app.MapAGUI("session-test-agent", "/agent");
+
+        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();
+        }
+    }
+}
+
+[SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated via dependency injection")]
+internal sealed class FakeSessionAgent : AIAgent
+{
+    private readonly string _name;
+
+    public FakeSessionAgent(string name)
+    {
+        this._name = name;
+    }
+
+    protected override string? IdCore => this._name;
+
+    public override string? Name => this._name;
+
+    public override string? Description => "A fake agent with session support for testing";
+
+    protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default) =>
+        new(new FakeSessionAgentSession());
+
+    protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) =>
+        new(serializedState.Deserialize<FakeSessionAgentSession>(jsonSerializerOptions)!);
+
+    protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        if (session is not FakeSessionAgentSession fakeSession)
+        {
+            throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent.");
+        }
+
+        return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions));
+    }
+
+    protected override async Task<AgentResponse> RunCoreAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        CancellationToken cancellationToken = default)
+    {
+        List<AgentResponseUpdate> updates = [];
+        await foreach (AgentResponseUpdate update in this.RunStreamingAsync(messages, session, options, cancellationToken).ConfigureAwait(false))
+        {
+            updates.Add(update);
+        }
+
+        return updates.ToAgentResponse();
+    }
+
+    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentSession? session = null,
+        AgentRunOptions? options = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        // Track turn count in session state to enable persistence verification.
+        // If the session store works correctly, the turn count increments across requests.
+        int turnCount = 1;
+        if (session != null)
+        {
+            var counter = session.StateBag.GetValue<TurnCounter>("turnCounter");
+            turnCount = (counter?.Count ?? 0) + 1;
+            session.StateBag.SetValue("turnCounter", new TurnCounter { Count = turnCount });
+        }
+
+        string messageId = Guid.NewGuid().ToString("N");
+        string prefix = $"Turn {turnCount}: ";
+
+        foreach (string chunk in new[] { prefix, "Hello", " ", "from", " ", "session", " ", "agent", "!" })
+        {
+            yield return new AgentResponseUpdate
+            {
+                MessageId = messageId,
+                Role = ChatRole.Assistant,
+                Contents = [new TextContent(chunk)]
+            };
+
+            await Task.Yield();
+        }
+    }
+
+    internal sealed class TurnCounter
+    {
+        public int Count { get; set; }
+    }
+
+    private sealed class FakeSessionAgentSession : AgentSession
+    {
+        public FakeSessionAgentSession()
+        {
+        }
+
+        [JsonConstructor]
+        public FakeSessionAgentSession(AgentSessionStateBag stateBag) : base(stateBag)
+        {
+        }
+    }
+}
@@ -14,6 +14,7 @@ using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Routing;
 using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using Moq;
@@ -31,6 +32,7 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests
        // Arrange
        Mock<IEndpointRouteBuilder> endpointsMock = new();
        Mock<IServiceProvider> serviceProviderMock = new();
+        serviceProviderMock.As<IKeyedServiceProvider>();

        endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object);
        endpointsMock.Setup(e => e.DataSources).Returns([]);
@@ -45,6 +47,155 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests
        Assert.NotNull(result);
    }

+    [Fact]
+    public void MapAGUI_WithAgentName_ResolvesKeyedAgentFromDI()
+    {
+        // Arrange
+        Mock<IEndpointRouteBuilder> endpointsMock = new();
+        Mock<IServiceProvider> serviceProviderMock = new();
+        AIAgent agent = new NamedTestAgent();
+
+        serviceProviderMock.As<IKeyedServiceProvider>()
+            .Setup(sp => sp.GetRequiredKeyedService(typeof(AIAgent), "test-agent"))
+            .Returns(agent);
+
+        endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object);
+        endpointsMock.Setup(e => e.DataSources).Returns([]);
+
+        // Act
+        IEndpointConventionBuilder? result = endpointsMock.Object.MapAGUI("test-agent", "/api/agent");
+
+        // Assert
+        Assert.NotNull(result);
+        serviceProviderMock.As<IKeyedServiceProvider>()
+            .Verify(sp => sp.GetRequiredKeyedService(typeof(AIAgent), "test-agent"), Times.Once);
+    }
+
+    [Fact]
+    public void MapAGUI_WithHostedAgentBuilder_ResolvesAgentByBuilderName()
+    {
+        // Arrange
+        Mock<IEndpointRouteBuilder> endpointsMock = new();
+        Mock<IServiceProvider> serviceProviderMock = new();
+        Mock<IHostedAgentBuilder> agentBuilderMock = new();
+        AIAgent agent = new NamedTestAgent();
+
+        agentBuilderMock.Setup(b => b.Name).Returns("test-agent");
+
+        serviceProviderMock.As<IKeyedServiceProvider>()
+            .Setup(sp => sp.GetRequiredKeyedService(typeof(AIAgent), "test-agent"))
+            .Returns(agent);
+
+        endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object);
+        endpointsMock.Setup(e => e.DataSources).Returns([]);
+
+        // Act
+        IEndpointConventionBuilder? result = endpointsMock.Object.MapAGUI(agentBuilderMock.Object, "/api/agent");
+
+        // Assert
+        Assert.NotNull(result);
+        serviceProviderMock.As<IKeyedServiceProvider>()
+            .Verify(sp => sp.GetRequiredKeyedService(typeof(AIAgent), "test-agent"), Times.Once);
+    }
+
+    [Fact]
+    public void MapAGUI_WithAgent_ResolvesSessionStoreFromDI()
+    {
+        // Arrange
+        Mock<IEndpointRouteBuilder> endpointsMock = new();
+        Mock<IServiceProvider> serviceProviderMock = new();
+        Mock<AgentSessionStore> sessionStoreMock = new();
+        AIAgent agent = new NamedTestAgent();
+
+        serviceProviderMock.As<IKeyedServiceProvider>()
+            .Setup(sp => sp.GetKeyedService(typeof(AgentSessionStore), "test-agent"))
+            .Returns(sessionStoreMock.Object);
+
+        endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object);
+        endpointsMock.Setup(e => e.DataSources).Returns([]);
+
+        // Act
+        IEndpointConventionBuilder? result = endpointsMock.Object.MapAGUI("/api/agent", agent);
+
+        // Assert
+        Assert.NotNull(result);
+        serviceProviderMock.As<IKeyedServiceProvider>()
+            .Verify(sp => sp.GetKeyedService(typeof(AgentSessionStore), "test-agent"), Times.Once);
+    }
+
+    [Fact]
+    public void MapAGUI_WithoutSessionStore_FallsBackToNoopStore()
+    {
+        // Arrange
+        Mock<IEndpointRouteBuilder> endpointsMock = new();
+        Mock<IServiceProvider> serviceProviderMock = new();
+        AIAgent agent = new TestAgent();
+
+        // No session store registered - IKeyedServiceProvider returns null by default
+        serviceProviderMock.As<IKeyedServiceProvider>();
+
+        endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object);
+        endpointsMock.Setup(e => e.DataSources).Returns([]);
+
+        // Act - should not throw (falls back to NoopAgentSessionStore)
+        IEndpointConventionBuilder? result = endpointsMock.Object.MapAGUI("/api/agent", agent);
+
+        // Assert
+        Assert.NotNull(result);
+    }
+
+    [Fact]
+    public void MapAGUI_WithNullEndpoints_ThrowsArgumentNullException()
+    {
+        // Arrange
+        AIAgent agent = new TestAgent();
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() =>
+            AGUIEndpointRouteBuilderExtensions.MapAGUI(null!, "/api/agent", agent));
+    }
+
+    [Fact]
+    public void MapAGUI_WithNullAgent_ThrowsArgumentNullException()
+    {
+        // Arrange
+        Mock<IEndpointRouteBuilder> endpointsMock = new();
+        Mock<IServiceProvider> serviceProviderMock = new();
+        serviceProviderMock.As<IKeyedServiceProvider>();
+        endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object);
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() =>
+            endpointsMock.Object.MapAGUI("/api/agent", (AIAgent)null!));
+    }
+
+    [Fact]
+    public void MapAGUI_WithNullAgentName_ThrowsArgumentNullException()
+    {
+        // Arrange
+        Mock<IEndpointRouteBuilder> endpointsMock = new();
+        Mock<IServiceProvider> serviceProviderMock = new();
+        serviceProviderMock.As<IKeyedServiceProvider>();
+        endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object);
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() =>
+            endpointsMock.Object.MapAGUI((string)null!, "/api/agent"));
+    }
+
+    [Fact]
+    public void MapAGUI_WithNullAgentBuilder_ThrowsArgumentNullException()
+    {
+        // Arrange
+        Mock<IEndpointRouteBuilder> endpointsMock = new();
+        Mock<IServiceProvider> serviceProviderMock = new();
+        endpointsMock.Setup(e => e.ServiceProvider).Returns(serviceProviderMock.Object);
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() =>
+            endpointsMock.Object.MapAGUI((IHostedAgentBuilder)null!, "/api/agent"));
+    }
+
    [Fact]
    public async Task MapAGUIAgent_WithNullOrInvalidInput_Returns400BadRequestAsync()
    {
@@ -556,4 +707,44 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests
            yield return new AgentResponseUpdate(new ChatResponseUpdate(ChatRole.Assistant, "Test response"));
        }
    }
+
+    private sealed class NamedTestAgent : AIAgent
+    {
+        protected override string? IdCore => "test-agent";
+
+        public override string? Name => "test-agent";
+
+        public override string? Description => "Named test agent";
+
+        protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default) =>
+            new(new TestAgentSession());
+
+        protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) =>
+            new(serializedState.Deserialize<TestAgentSession>(jsonSerializerOptions)!);
+
+        protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        {
+            if (session is not TestAgentSession testSession)
+            {
+                throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(TestAgentSession)}' can be serialized by this agent.");
+            }
+
+            return new(JsonSerializer.SerializeToElement(testSession, jsonSerializerOptions));
+        }
+
+        protected override Task<AgentResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
+        {
+            throw new NotImplementedException();
+        }
+
+        protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
+            IEnumerable<ChatMessage> messages,
+            AgentSession? session = null,
+            AgentRunOptions? options = null,
+            [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            await Task.CompletedTask;
+            yield return new AgentResponseUpdate(new ChatResponseUpdate(ChatRole.Assistant, "Test response"));
+        }
+    }
 }