.NET: introduce API to manage the thread in Hosting scenarios

dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs

@@ -26,7 +26,8 @@
     "pirate",
     instructions: "You are a pirate. Speak like a pirate",
     description: "An agent that speaks like a pirate.",
-    chatClientServiceKey: "chat-model");
+    chatClientServiceKey: "chat-model")
+    .WithInMemoryThreadStore();
 
 builder.AddAIAgent("knights-and-knaves", (sp, key) =>
 {

dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/WebApplicationExtensions.cs

@@ -22,9 +22,10 @@ public static class WebApplicationExtensions
     public static void MapA2A(this WebApplication app, string agentName, string path)
     {
         var agent = app.Services.GetRequiredKeyedService<AIAgent>(agentName);
+        var agentThreadStore = app.Services.GetKeyedService<IAgentThreadStore>(agentName);
         var loggerFactory = app.Services.GetRequiredService<ILoggerFactory>();
 
-        var taskManager = agent.MapA2A(loggerFactory: loggerFactory);
+        var taskManager = agent.MapA2A(loggerFactory: loggerFactory, agentThreadStore: agentThreadStore);
         app.MapA2A(taskManager, path);
     }
 
@@ -42,9 +43,10 @@ public static void MapA2A(
         AgentCard agentCard)
     {
         var agent = app.Services.GetRequiredKeyedService<AIAgent>(agentName);
+        var agentThreadStore = app.Services.GetKeyedService<IAgentThreadStore>(agentName);
         var loggerFactory = app.Services.GetRequiredService<ILoggerFactory>();
 
-        var taskManager = agent.MapA2A(agentCard: agentCard, loggerFactory: loggerFactory);
+        var taskManager = agent.MapA2A(agentCard: agentCard, loggerFactory: loggerFactory, agentThreadStore: agentThreadStore);
         app.MapA2A(taskManager, path);
     }
 

dotnet/src/Microsoft.Agents.AI.Hosting.A2A/AIAgentExtensions.cs

@@ -20,29 +20,37 @@ public static class AIAgentExtensions
     /// <param name="agent">Agent to attach A2A messaging processing capabilities to.</param>
     /// <param name="taskManager">Instance of <see cref="TaskManager"/> to configure for A2A messaging. New instance will be created if not passed.</param>
     /// <param name="loggerFactory">The logger factory to use for creating <see cref="ILogger"/> instances.</param>
+    /// <param name="agentThreadStore">The store to store thread contents and metadata.</param>
     /// <returns>The configured <see cref="TaskManager"/>.</returns>
     public static TaskManager MapA2A(
         this AIAgent agent,
         TaskManager? taskManager = null,
-        ILoggerFactory? loggerFactory = null)
+        ILoggerFactory? loggerFactory = null,
+        IAgentThreadStore? agentThreadStore = null)
     {
         ArgumentNullException.ThrowIfNull(agent);
         ArgumentNullException.ThrowIfNull(agent.Name);
 
-        taskManager ??= new();
+        var hostAgent = new AIHostAgent(
+            innerAgent: agent,
+            threadStore: agentThreadStore ?? new NoContextAgentThreadStore());
 
+        taskManager ??= new();
         taskManager.OnMessageReceived += OnMessageReceivedAsync;
-
         return taskManager;
 
         async Task<A2AResponse> OnMessageReceivedAsync(MessageSendParams messageSendParams, CancellationToken cancellationToken)
         {
-            var response = await agent.RunAsync(
+            var contextId = messageSendParams.Message.ContextId ?? Guid.NewGuid().ToString("N");
+            var thread = await hostAgent.RestoreThreadAsync(contextId, cancellationToken).ConfigureAwait(false);
+
+            var response = await hostAgent.RunAsync(
                 messageSendParams.ToChatMessages(),
+                thread: thread,
                 cancellationToken: cancellationToken).ConfigureAwait(false);
-            var contextId = messageSendParams.Message.ContextId ?? Guid.NewGuid().ToString("N");
-            var parts = response.Messages.ToParts();
 
+            await hostAgent.SaveThreadAsync(contextId, thread!, cancellationToken).ConfigureAwait(false);
+            var parts = response.Messages.ToParts();
             return new AgentMessage
             {
                 MessageId = response.ResponseId ?? Guid.NewGuid().ToString("N"),
@@ -60,14 +68,16 @@ async Task<A2AResponse> OnMessageReceivedAsync(MessageSendParams messageSendPara
     /// <param name="agentCard">The agent card to return on query.</param>
     /// <param name="taskManager">Instance of <see cref="TaskManager"/> to configure for A2A messaging. New instance will be created if not passed.</param>
     /// <param name="loggerFactory">The logger factory to use for creating <see cref="ILogger"/> instances.</param>
+    /// <param name="agentThreadStore">The store to store thread contents and metadata.</param>
     /// <returns>The configured <see cref="TaskManager"/>.</returns>
     public static TaskManager MapA2A(
         this AIAgent agent,
         AgentCard agentCard,
         TaskManager? taskManager = null,
-        ILoggerFactory? loggerFactory = null)
+        ILoggerFactory? loggerFactory = null,
+        IAgentThreadStore? agentThreadStore = null)
     {
-        taskManager = agent.MapA2A(taskManager, loggerFactory);
+        taskManager = agent.MapA2A(taskManager, loggerFactory, agentThreadStore);
 
         taskManager.OnAgentCardQuery += (context, query) =>
         {

dotnet/src/Microsoft.Agents.AI.Hosting/AIHostAgent.cs

@@ -0,0 +1,134 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Agents.AI.Hosting;
+
+/// <summary>
+/// Provides a hosting wrapper around an <see cref="AIAgent"/> that adds thread persistence capabilities
+/// for server-hosted scenarios where conversations need to be restored across requests.
+/// </summary>
+/// <remarks>
+/// <para>
+/// <see cref="AIHostAgent"/> wraps an existing agent implementation and adds the ability to
+/// persist and restore conversation threads using an <see cref="IAgentThreadStore"/>.
+/// </para>
+/// <para>
+/// This wrapper enables thread persistence without requiring type-specific knowledge of the thread type,
+/// as all thread operations work through the base <see cref="AgentThread"/> abstraction.
+/// </para>
+/// </remarks>
+public class AIHostAgent : AIAgent
+{
+    private readonly AIAgent _innerAgent;
+    private readonly IAgentThreadStore _threadStore;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AIHostAgent"/> class.
+    /// </summary>
+    /// <param name="innerAgent">The underlying agent implementation to wrap.</param>
+    /// <param name="threadStore">The thread store to use for persisting conversation state.</param>
+    /// <exception cref="ArgumentNullException">
+    /// <paramref name="innerAgent"/> or <paramref name="threadStore"/> is <see langword="null"/>.
+    /// </exception>
+    public AIHostAgent(AIAgent innerAgent, IAgentThreadStore threadStore)
+    {
+        this._innerAgent = Throw.IfNull(innerAgent);
+        this._threadStore = Throw.IfNull(threadStore);
+    }
+
+    /// <inheritdoc/>
+    public override string Id => this._innerAgent.Id;
+
+    /// <inheritdoc/>
+    public override string? Name => this._innerAgent.Name;
+
+    /// <inheritdoc/>
+    public override string? Description => this._innerAgent.Description;
+
+    /// <inheritdoc/>
+    public override object? GetService(Type serviceType, object? serviceKey = null)
+        => base.GetService(serviceType, serviceKey) ?? this._innerAgent.GetService(serviceType, serviceKey);
+
+    /// <inheritdoc/>
+    public override AgentThread GetNewThread() => this._innerAgent.GetNewThread();
+
+    /// <inheritdoc/>
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+        => this._innerAgent.DeserializeThread(serializedThread, jsonSerializerOptions);
+
+    /// <summary>
+    /// Restores a conversation thread from the thread store using the specified conversation ID.
+    /// </summary>
+    /// <param name="conversationId">The unique identifier of the conversation to restore.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
+    /// <returns>
+    /// A task that represents the asynchronous operation. The task result contains the restored thread,
+    /// or <see langword="null"/> if no thread with the given ID exists in the store.
+    /// </returns>
+    /// <exception cref="ArgumentException"><paramref name="conversationId"/> is null or whitespace.</exception>
+    public async ValueTask<AgentThread?> RestoreThreadAsync(
+        string conversationId,
+        CancellationToken cancellationToken = default)
+    {
+        _ = Throw.IfNullOrWhitespace(conversationId);
+
+        var serializedThread = await this._threadStore.GetOrCreateThreadAsync(
+            conversationId,
+            this.Id,
+            cancellationToken).ConfigureAwait(false);
+
+        if (serializedThread is null)
+        {
+            return null;
+        }
+
+        return this._innerAgent.DeserializeThread(serializedThread.Value);
+    }
+
+    /// <summary>
+    /// Persists a conversation thread to the thread store.
+    /// </summary>
+    /// <param name="conversationId">The unique identifier for the conversation.</param>
+    /// <param name="thread">The thread to persist.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
+    /// <returns>A task that represents the asynchronous save operation.</returns>
+    /// <exception cref="ArgumentException"><paramref name="conversationId"/> is null or whitespace.</exception>
+    /// <exception cref="ArgumentNullException"><paramref name="thread"/> is <see langword="null"/>.</exception>
+    public async ValueTask SaveThreadAsync(
+        string conversationId,
+        AgentThread thread,
+        CancellationToken cancellationToken = default)
+    {
+        _ = Throw.IfNullOrWhitespace(conversationId);
+        _ = Throw.IfNull(thread);
+
+        await this._threadStore.SaveThreadAsync(
+            conversationId,
+            this.Id,
+            thread,
+            cancellationToken).ConfigureAwait(false);
+    }
+
+    /// <inheritdoc/>
+    public override Task<AgentRunResponse> RunAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentThread? thread = null,
+        AgentRunOptions? options = null,
+        CancellationToken cancellationToken = default)
+        => this._innerAgent.RunAsync(messages, thread, options, cancellationToken);
+
+    /// <inheritdoc/>
+    public override IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(
+        IEnumerable<ChatMessage> messages,
+        AgentThread? thread = null,
+        AgentRunOptions? options = null,
+        CancellationToken cancellationToken = default)
+        => this._innerAgent.RunStreamingAsync(messages, thread, options, cancellationToken);
+}

dotnet/src/Microsoft.Agents.AI.Hosting/HostAgentBuilderExtensions.cs

@@ -0,0 +1,62 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Agents.AI.Hosting;
+
+/// <summary>
+/// Provides extension methods for configuring <see cref="AIAgent"/>.
+/// </summary>
+public static class HostAgentBuilderExtensions
+{
+    /// <summary>
+    /// Configures the host agent builder to use an in-memory thread store for agent thread management.
+    /// </summary>
+    /// <param name="builder">The host agent builder to configure with the in-memory thread store.</param>
+    /// <returns>The same <paramref name="builder"/> instance, configured to use an in-memory thread store.</returns>
+    public static IHostAgentBuilder WithInMemoryThreadStore(this IHostAgentBuilder builder)
+    {
+        builder.HostApplicationBuilder.Services.AddKeyedSingleton<IAgentThreadStore>(builder.Name, new InMemoryAgentThreadStore());
+        return builder;
+    }
+
+    /// <summary>
+    /// Registers the specified agent thread store with the host agent builder, enabling thread-specific storage for
+    /// agent operations.
+    /// </summary>
+    /// <param name="builder">The host agent builder to configure with the thread store. Cannot be null.</param>
+    /// <param name="store">The agent thread store instance to register. Cannot be null.</param>
+    /// <returns>The same host agent builder instance, allowing for method chaining.</returns>
+    public static IHostAgentBuilder WithThreadStore(this IHostAgentBuilder builder, IAgentThreadStore store)
+    {
+        builder.HostApplicationBuilder.Services.AddKeyedSingleton(builder.Name, store);
+        return builder;
+    }
+
+    /// <summary>
+    /// Configures the host agent builder to use a custom thread store implementation for agent threads.
+    /// </summary>
+    /// <param name="builder">The host agent builder to configure.</param>
+    /// <param name="createAgentThreadStore">A factory function that creates an agent thread store instance using the provided service provider and agent
+    /// name.</param>
+    /// <returns>The same host agent builder instance, enabling further configuration.</returns>
+    public static IHostAgentBuilder WithThreadStore(this IHostAgentBuilder builder, Func<IServiceProvider, string, IAgentThreadStore> createAgentThreadStore)
+    {
+        builder.HostApplicationBuilder.Services.AddKeyedSingleton(builder.Name, (sp, key) =>
+        {
+            Throw.IfNull(key);
+            var keyString = key as string;
+            Throw.IfNullOrEmpty(keyString);
+            var store = createAgentThreadStore(sp, keyString);
+            if (store is null)
+            {
+                throw new InvalidOperationException($"The agent thread store factory did not return a valid {nameof(IAgentThreadStore)} instance for key '{keyString}'.");
+            }
+
+            return store;
+        });
+        return builder;
+    }
+}

dotnet/src/Microsoft.Agents.AI.Hosting/IAgentThreadStore.cs

@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Agents.AI.Hosting;
+
+/// <summary>
+/// Defines the contract for storing and retrieving agent conversation threads.
+/// </summary>
+/// <remarks>
+/// Implementations of this interface enable persistent storage of conversation threads,
+/// allowing conversations to be resumed across HTTP requests, application restarts,
+/// or different service instances in hosted scenarios.
+/// </remarks>
+public interface IAgentThreadStore
+{
+    /// <summary>
+    /// Saves a serialized agent thread to persistent storage.
+    /// </summary>
+    /// <param name="conversationId">The unique identifier for the conversation/thread.</param>
+    /// <param name="agentId">The identifier of the agent that owns this thread.</param>
+    /// <param name="thread">The thread to save.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
+    /// <returns>A task that represents the asynchronous save operation.</returns>
+    /// <remarks>
+    /// The <paramref name="conversationId"/> is used as the primary key for thread lookup.
+    /// The <paramref name="agentId"/> provides additional scoping to prevent cross-agent thread access.
+    /// </remarks>
+    ValueTask SaveThreadAsync(
+        string conversationId,
+        string agentId,
+        AgentThread thread,
+        CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Retrieves a serialized agent thread from persistent storage.
+    /// </summary>
+    /// <param name="conversationId">The unique identifier for the conversation/thread to retrieve.</param>
+    /// <param name="agentId">The identifier of the agent that owns this thread.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
+    /// <returns>
+    /// A task that represents the asynchronous retrieval operation.
+    /// The task result contains the serialized thread state, or <see langword="null"/> if not found.
+    /// </returns>
+    ValueTask<JsonElement?> GetOrCreateThreadAsync(
+        string conversationId,
+        string agentId,
+        CancellationToken cancellationToken = default);
+}