Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
westey-m wants to merge 2 commits into
base: main
Choose a base branch
from
+508 −7
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
29f2d11
Update AGUI service to support session storage
westey-m Apr 9, 2026
641bb78
Apply suggestion from @Copilot
westey-m Apr 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/Program.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
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;
Expand All @@ -13,11 +14,11 @@
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
Expand All @@ -27,7 +28,7 @@
new DefaultAzureCredential())
.GetChatClient(deploymentName)
.AsAIAgent(
name: "AGUIAssistant",
name: AgentName,
tools: [
AIFunctionFactory.Create(
() => DateTimeOffset.UtcNow,
Expand All @@ -48,7 +49,15 @@
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();
79 changes: 76 additions & 3 deletions dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Copy link

Copilot AI Apr 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

using System.Threading.Tasks; appears unused in this file and will trigger CS8019 (unnecessary using directive) in typical builds. Please remove it (and any other unused usings) to keep builds warning-free.

Suggested change
using System.Threading.Tasks;

Copilot uses AI. Check for mistakes.
using Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.Shared;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
Expand All @@ -21,18 +24,66 @@ 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)
{
Copy link

Copilot AI Apr 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The MapAGUI(IHostedAgentBuilder, ...) overload doesn't validate endpoints. If a caller passes null explicitly, this will throw a NullReferenceException when calling endpoints.MapAGUI(...) instead of an ArgumentNullException. Add ArgumentNullException.ThrowIfNull(endpoints); at the start for consistency with the other overloads.

Suggested change
{
{
ArgumentNullException.ThrowIfNull(endpoints);

Copilot uses AI. Check for mistakes.
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>
/// <param name="endpoints">The endpoint route builder.</param>
/// <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);
Copy link

Copilot AI Apr 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

aiAgent.Name is nullable. Passing a null/whitespace key into GetKeyedService<AgentSessionStore>(aiAgent.Name) can lead to surprising behavior (e.g., resolving an unkeyed store, or later failing when creating sessions). Consider guarding with string.IsNullOrWhiteSpace(aiAgent.Name) and treating that as "no store registered" (use NoopAgentSessionStore).

Suggested change
var agentSessionStore = endpoints.ServiceProvider.GetKeyedService<AgentSessionStore>(aiAgent.Name);
var agentSessionStore = string.IsNullOrWhiteSpace(aiAgent.Name)
? null
: endpoints.ServiceProvider.GetKeyedService<AgentSessionStore>(aiAgent.Name);

Copilot uses AI. Check for mistakes.
var hostAgent = new AIHostAgent(aiAgent, agentSessionStore ?? new NoopAgentSessionStore());

Comment on lines +84 to +86
Copy link

Copilot AI Apr 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The session store (and AIHostAgent) are resolved/constructed at endpoint-mapping time via endpoints.ServiceProvider. If a user registers an AgentSessionStore with Scoped/Transient lifetime (supported by WithSessionStore(..., lifetime)), resolving it from the root provider can throw ("Cannot resolve scoped service...") or capture the wrong lifetime. Resolve the store from HttpContext.RequestServices inside the request handler (or otherwise create the AIHostAgent per request) to respect DI lifetimes.

Copilot uses AI. Check for mistakes.
return endpoints.MapPost(pattern, async ([FromBody] RunAgentInput? input, HttpContext context, CancellationToken cancellationToken) =>
{
if (input is null)
Expand Down Expand Up @@ -63,21 +114,43 @@ public static IEndpointConventionBuilder MapAGUI(
}
};

var threadId = input.ThreadId ?? Guid.NewGuid().ToString("N");
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);
}
}
1 change: 1 addition & 0 deletions ...soft.Agents.AI.Hosting.AGUI.AspNetCore/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
1 change: 1 addition & 0 deletions ...Core.IntegrationTests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests.csproj
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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" />
Expand Down
Loading
Loading