.NET: API to manage AgentThreads in hosting scenarios

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

@@ -25,7 +25,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")
+    .AddInMemoryThreadStore();
 
 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 NoopAgentThreadStore());
 
+        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.GetOrCreateThreadAsync(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,73 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+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 : DelegatingAIAgent
+{
+    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)
+        : base(innerAgent)
+    {
+        this._threadStore = Throw.IfNull(threadStore);
+    }
+
+    /// <summary>
+    /// Gets an existing agent thread for the specified conversation, or creates a new one if none exists.
+    /// </summary>
+    /// <param name="conversationId">The unique identifier of the conversation for which to retrieve or create the agent thread. Cannot be null,
+    /// empty, or consist only of white-space characters.</param>
+    /// <param name="cancellationToken">A cancellation token that can be used to cancel the asynchronous operation.</param>
+    /// <returns>A task that represents the asynchronous operation. The task result contains the agent thread associated with the
+    /// specified conversation. If no thread exists, a new thread is created and returned.</returns>
+    public ValueTask<AgentThread> GetOrCreateThreadAsync(string conversationId, CancellationToken cancellationToken = default)
+    {
+        _ = Throw.IfNullOrWhitespace(conversationId);
+
+        return this._threadStore.GetThreadAsync(this.InnerAgent, conversationId, cancellationToken);
+    }
+
+    /// <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 ValueTask SaveThreadAsync(string conversationId, AgentThread thread, CancellationToken cancellationToken = default)
+    {
+        _ = Throw.IfNullOrWhitespace(conversationId);
+        _ = Throw.IfNull(thread);
+
+        return this._threadStore.SaveThreadAsync(this.InnerAgent, conversationId, thread, cancellationToken);
+    }
+}

dotnet/src/Microsoft.Agents.AI.Hosting/HostedAgentBuilderExtensions.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 HostedAgentBuilderExtensions
+{
+    /// <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 IHostedAgentBuilder AddInMemoryThreadStore(this IHostedAgentBuilder 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 IHostedAgentBuilder AddThreadStore(this IHostedAgentBuilder 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 IHostedAgentBuilder AddThreadStore(this IHostedAgentBuilder 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,46 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+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="agent">The agent that owns this thread.</param>
+    /// <param name="conversationId">The unique identifier for the conversation/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>
+    ValueTask SaveThreadAsync(
+        AIAgent agent,
+        string conversationId,
+        AgentThread thread,
+        CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Retrieves a serialized agent thread from persistent storage.
+    /// </summary>
+    /// <param name="agent">The agent that owns this thread.</param>
+    /// <param name="conversationId">The unique identifier for the conversation/thread to retrieve.</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<AgentThread> GetThreadAsync(
+        AIAgent agent,
+        string conversationId,
+        CancellationToken cancellationToken = default);
+}

dotnet/src/Microsoft.Agents.AI.Hosting/Local/InMemoryAgentThreadStore.cs

@@ -0,0 +1,54 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Collections.Concurrent;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Agents.AI.Hosting;
+
+/// <summary>
+/// Provides an in-memory implementation of <see cref="IAgentThreadStore"/> for development and testing scenarios.
+/// </summary>
+/// <remarks>
+/// <para>
+/// This implementation stores threads in memory using a concurrent dictionary and is suitable for:
+/// <list type="bullet">
+/// <item><description>Single-instance development scenarios</description></item>
+/// <item><description>Testing and prototyping</description></item>
+/// <item><description>Scenarios where thread persistence across restarts is not required</description></item>
+/// </list>
+/// </para>
+/// <para>
+/// <strong>Warning:</strong> All stored threads will be lost when the application restarts.
+/// For production use with multiple instances or persistence across restarts, use a durable storage implementation
+/// such as Redis, SQL Server, or Azure Cosmos DB.
+/// </para>
+/// </remarks>
+public sealed class InMemoryAgentThreadStore : IAgentThreadStore
+{
+    private readonly ConcurrentDictionary<string, JsonElement> _threads = new();
+
+    /// <inheritdoc/>
+    public ValueTask SaveThreadAsync(AIAgent agent, string conversationId, AgentThread thread, CancellationToken cancellationToken = default)
+    {
+        var key = GetKey(conversationId, agent.Id);
+        this._threads[key] = thread.Serialize();
+        return default;
+    }
+
+    /// <inheritdoc/>
+    public ValueTask<AgentThread> GetThreadAsync(AIAgent agent, string conversationId, CancellationToken cancellationToken = default)
+    {
+        var key = GetKey(conversationId, agent.Id);
+        JsonElement? threadContent = this._threads.TryGetValue(key, out var existingThread) ? existingThread : null;
+
+        return threadContent switch
+        {
+            null => new ValueTask<AgentThread>(agent.GetNewThread()),
+            _ => new ValueTask<AgentThread>(agent.DeserializeThread(threadContent.Value)),
+        };
+    }
+
+    private static string GetKey(string conversationId, string agentId) => $"{agentId}:{conversationId}";
+}

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

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Agents.AI.Hosting;
+
+/// <summary>
+/// This store implementation does not have any store under the hood and operates with empty threads.
+/// It is the "noop" store, and could be used if you are keeping the thread contents on the client side for example.
+/// </summary>
+public sealed class NoopAgentThreadStore : IAgentThreadStore
+{
+    /// <inheritdoc/>
+    public ValueTask SaveThreadAsync(AIAgent agent, string conversationId, AgentThread thread, CancellationToken cancellationToken = default)
+    {
+        return new ValueTask();
+    }
+
+    /// <inheritdoc/>
+    public ValueTask<AgentThread> GetThreadAsync(AIAgent agent, string conversationId, CancellationToken cancellationToken = default)
+    {
+        return new ValueTask<AgentThread>(agent.GetNewThread());
+    }
+}