.Net: Add support for 3rd party thread storage and thread serialization (#203)

* Add thread storage and serialization POC * Switch to using JsonElement and add unit tests * Add additional unit tests. * Exclude private debugger properties from CodeCoverage. * Rename IChatMessagesStorable to IChatMessageStore * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Improve xml doc. * Update the message storing thread to always use external store for both local and remote storage. * Remove threadid from the IChatMessageStore interface, since the store should own the thread id itself, if it requires one. * Switch GetMessages to IEnumerable * Address pr comments. * Make jsonserializer options default consistent on DeserializeThreadAsync * Move message storing thread functionality into AgentThread and simplify AgentThread behavior. * Remove embedding generation from VectorStore chat history sample. * Remove unecessary code and fix formatting. * Make GetNewThread and DeserializeThread virtual with default implementations. Remove unsued json utilities. * Fix formatting * Remove problem test. * Add more unit tests * Remove unused using clause. * Address pr feedback. * Address PR comments. * Make InMemory store internal * Switch InMemoryChatMessageStore to implement IList instead of inheriting from List. * Rename store deserialize param. * Update serialization based on PR comments. * Remove confusing comment. * Address Deserialization PR comments in the same way as Serialization * Add State to IChatMessageStore Serialize and Deserialize names. Make Thread Deserialize internal. Make AgentThread type switching fobidden. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2025-08-05 17:24:25 +00:00
parent c1d306ec95
commit ff3e13c2aa
38 changed files with 1143 additions and 1427 deletions
@@ -70,13 +70,13 @@ public abstract partial class OrchestratingAgent : AIAgent

        if (thread is not null)
        {
-            if (thread is not IMessagesRetrievableThread retrievableThread)
+            if (thread.MessageStore is null)
            {
-                throw new InvalidOperationException($"The thread type '{thread.GetType().Name}' is not supported by this agent. Use {nameof(GetNewThread)} to create a thread when needed.");
+                throw new InvalidOperationException("An agent service managed thread is not supported by this agent.");
            }

            List<ChatMessage> messagesList = [];
-            await foreach (var threadMessage in retrievableThread.GetMessagesAsync(cancellationToken).ConfigureAwait(false))
+            await foreach (var threadMessage in thread.GetMessagesAsync(cancellationToken).ConfigureAwait(false))
            {
                messagesList.Add(threadMessage);
            }
@@ -101,9 +101,6 @@ public abstract partial class OrchestratingAgent : AIAgent
        }
    }

-    /// <inheritdoc />
-    public sealed override AgentThread GetNewThread() => new ChatClientAgentThread();
-
    /// <summary>
    /// Initiates processing of the orchestration.
    /// </summary>
@@ -207,10 +204,6 @@ public abstract partial class OrchestratingAgent : AIAgent
        return response;
    }

-    /// <inheritdoc />
-    protected sealed override TThreadType ValidateOrCreateThreadType<TThreadType>(AgentThread? thread, Func<TThreadType> constructThread) =>
-        base.ValidateOrCreateThreadType(thread, constructThread);
-
    /// <summary>Writes the specified checkpoint state to the runtime.</summary>
    /// <param name="state">The state to persist.</param>
    /// <param name="context">The context for the orchestrating operation.</param>
@@ -2,6 +2,7 @@

 using System;
 using System.Collections.Generic;
+using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Shared.Diagnostics;
@@ -50,7 +51,21 @@ public abstract class AIAgent
    /// If the thread needs to be created via a service call it would be created on first use.
    /// </para>
    /// </remarks>
-    public abstract AgentThread GetNewThread();
+    public virtual AgentThread GetNewThread() => new();
+
+    /// <summary>
+    /// Deserialize the thread from JSON.
+    /// </summary>
+    /// <param name="serializedThread">The <see cref="JsonElement"/> representing the thread state.</param>
+    /// <param name="jsonSerializerOptions">Optional <see cref="JsonSerializerOptions"/> to use for deserializing the thread state.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>The deserialized <see cref="AgentThread"/> instance.</returns>
+    public async ValueTask<AgentThread> DeserializeThreadAsync(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        var thread = this.GetNewThread();
+        await thread.DeserializeAsync(serializedThread, jsonSerializerOptions, cancellationToken).ConfigureAwait(false);
+        return thread;
+    }

    /// <summary>
    /// Run the agent with no message assuming that all required instructions are already provided to the agent or on the thread.
@@ -192,30 +207,6 @@ public abstract class AIAgent
        AgentRunOptions? options = null,
        CancellationToken cancellationToken = default);

-    /// <summary>
-    /// Checks that the thread is of the expected type, or if null, creates the default thread type.
-    /// </summary>
-    /// <typeparam name="TThreadType">The expected type of the thead.</typeparam>
-    /// <param name="thread">The thread to create if it's null and validate its type if not null.</param>
-    /// <param name="constructThread">A callback to use to construct the thread if it's null.</param>
-    /// <returns>An async task that completes once all update are complete.</returns>
-    protected virtual TThreadType ValidateOrCreateThreadType<TThreadType>(
-        AgentThread? thread,
-        Func<TThreadType> constructThread)
-        where TThreadType : AgentThread
-    {
-        Throw.IfNull(constructThread);
-
-        thread ??= constructThread();
-
-        if (thread is not TThreadType concreteThreadType)
-        {
-            throw new NotSupportedException($"{this.GetType().Name} currently only supports agent threads of type {typeof(TThreadType).Name}.");
-        }
-
-        return concreteThreadType;
-    }
-
    /// <summary>
    /// Notfiy the given thread that new messages are available.
    /// </summary>
@@ -57,7 +57,8 @@ public static partial class AgentAbstractionsJsonUtilities
    [JsonSerializable(typeof(AgentRunResponse[]))]
    [JsonSerializable(typeof(AgentRunResponseUpdate))]
    [JsonSerializable(typeof(AgentRunResponseUpdate[]))]
-    [JsonSerializable(typeof(AgentThread))]
+    [JsonSerializable(typeof(AgentThread.ThreadState))]
+    [JsonSerializable(typeof(InMemoryChatMessageStore.StoreState))]

    [ExcludeFromCodeCoverage]
    internal sealed partial class JsonContext : JsonSerializerContext;
@@ -141,9 +141,11 @@ public class AgentRunResponseUpdate

    /// <summary>Gets a <see cref="AIContent"/> object to display in the debugger display.</summary>
    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+    [ExcludeFromCodeCoverage]
    private AIContent? ContentForDebuggerDisplay => this._contents is { Count: > 0 } ? this._contents[0] : null;

    /// <summary>Gets an indication for the debugger display of whether there's more content.</summary>
    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+    [ExcludeFromCodeCoverage]
    private string EllipsesForDebuggerDisplay => this._contents is { Count: > 1 } ? ", ..." : string.Empty;
 }
@@ -2,8 +2,12 @@

 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
+using System.Runtime.CompilerServices;
+using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Shared.Diagnostics;

 namespace Microsoft.Extensions.AI.Agents;

@@ -13,21 +17,114 @@ namespace Microsoft.Extensions.AI.Agents;
 /// </summary>
 public class AgentThread
 {
+    private string? _conversationId;
+    private IChatMessageStore? _messageStore;
+
    /// <summary>
-    /// Gets or sets the id of the current thread.
+    /// Initializes a new instance of the <see cref="AgentThread"/> class.
+    /// </summary>
+    public AgentThread()
+    {
+    }
+
+    /// <summary>
+    /// Gets or sets the id of the current thread to support cases where the thread is owned by the agent service.
    /// </summary>
    /// <remarks>
    /// <para>
-    /// This id may be null if the thread has no id, or
-    /// if it represents a service-owned thread but the service
-    /// has not yet been called to create the thread.
+    /// Note that either <see cref="ConversationId"/> or <see cref="MessageStore "/> may be set, but not both.
+    /// If <see cref="MessageStore "/> is not null, and <see cref="ConversationId"/> is set, <see cref="MessageStore "/>
+    /// will be reverted to null, and vice versa.
    /// </para>
    /// <para>
-    /// The id may also change over time where the <see cref="AgentThread"/>
-    /// is a proxy to a service owned thread that forks on each agent invocation.
+    /// This property may be null in the following cases:
+    /// <list type="bullet">
+    /// <item>The thread stores messages via the <see cref="IChatMessageStore"/> and not in the agent service.</item>
+    /// <item>This thread object is new and a server managed thread has not yet been created in the agent service.</item>
+    /// </list>
+    /// </para>
+    /// <para>
+    /// The id may also change over time where the the id is pointing at a
+    /// agent service managed thread, and the default behavior of a service is
+    /// to fork the thread with each iteration.
    /// </para>
    /// </remarks>
-    public string? Id { get; set; }
+    public string? ConversationId
+    {
+        get { return this._conversationId; }
+        set
+        {
+            if (string.IsNullOrWhiteSpace(this._conversationId) && string.IsNullOrWhiteSpace(value))
+            {
+                return;
+            }
+
+            if (this._messageStore is not null)
+            {
+                // If we have a message store already, we shouldn't switch the thread to use a conversation id
+                // since it means that the thread contents will essentially be deleted, and the thread will not work
+                // with the original agent anymore.
+                throw new InvalidOperationException("Only the ConversationId or MessageStore may be set, but not both and switching from one to another is not supported.");
+            }
+
+            this._conversationId = Throw.IfNullOrWhitespace(value);
+        }
+    }
+
+    /// <summary>
+    /// Gets or sets the <see cref="IChatMessageStore"/> used by this thread, for cases where messages should be stored in a custom location.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// Note that either <see cref="ConversationId"/> or <see cref="MessageStore "/> may be set, but not both.
+    /// If <see cref="ConversationId"/> is not null, and <see cref="MessageStore "/> is set, <see cref="ConversationId"/>
+    /// will be reverted to null, and vice versa.
+    /// </para>
+    /// <para>
+    /// This property may be null in the following cases:
+    /// <list type="bullet">
+    /// <item>The thread stores messages in the agent service and just has an id to the remove thread, instead of in an <see cref="IChatMessageStore"/>.</item>
+    /// <item>This thread object is new it is not yet clear whether it will be backed by a server managed thread or an <see cref="IChatMessageStore"/>.</item>
+    /// </list>
+    /// </para>
+    /// </remarks>
+    public IChatMessageStore? MessageStore
+    {
+        get { return this._messageStore; }
+        set
+        {
+            if (this._messageStore is null && value is null)
+            {
+                return;
+            }
+
+            if (!string.IsNullOrWhiteSpace(this._conversationId))
+            {
+                // If we have a conversation id already, we shouldn't switch the thread to use a message store
+                // since it means that the thread will not work with the original agent anymore.
+                throw new InvalidOperationException("Only the ConversationId or MessageStore may be set, but not both and switching from one to another is not supported.");
+            }
+
+            this._messageStore = Throw.IfNull(value);
+        }
+    }
+
+    /// <summary>
+    /// Retrieves any messages stored in the <see cref="IChatMessageStore"/> of the thread, otherwise returns an empty collection.
+    /// </summary>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>The messages from the <see cref="IChatMessageStore"/> in ascending chronological order, with the oldest message first.</returns>
+    public virtual async IAsyncEnumerable<ChatMessage> GetMessagesAsync([EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        if (this._messageStore is not null)
+        {
+            var messages = await this._messageStore!.GetMessagesAsync(cancellationToken).ConfigureAwait(false);
+            foreach (var message in messages)
+            {
+                yield return message;
+            }
+        }
+    }

    /// <summary>
    /// This method is called when new messages have been contributed to the chat by any participant.
@@ -39,8 +136,92 @@ public class AgentThread
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
    /// <returns>A task that completes when the context has been updated.</returns>
    /// <exception cref="InvalidOperationException">The thread has been deleted.</exception>
-    protected internal virtual Task OnNewMessagesAsync(IReadOnlyCollection<ChatMessage> newMessages, CancellationToken cancellationToken = default)
+    protected internal virtual async Task OnNewMessagesAsync(IReadOnlyCollection<ChatMessage> newMessages, CancellationToken cancellationToken = default)
    {
-        return Task.CompletedTask;
+        switch (this)
+        {
+            case { ConversationId: not null }:
+                // If the thread messages are stored in the service
+                // there is nothing to do here, since invoking the
+                // service should already update the thread.
+                break;
+
+            case { MessageStore: null }:
+                // If there is no conversation id, and no store we can createa a default in memory store and add messages to it.
+                this._messageStore = new InMemoryChatMessageStore();
+                await this._messageStore!.AddMessagesAsync(newMessages, cancellationToken).ConfigureAwait(false);
+                break;
+
+            case { MessageStore: not null }:
+                // If a store has been provided, we need to add the messages to the store.
+                await this._messageStore!.AddMessagesAsync(newMessages, cancellationToken).ConfigureAwait(false);
+                break;
+
+            default:
+                throw new UnreachableException();
+        }
+    }
+
+    /// <summary>
+    /// Deserializes the state contained in the provided <see cref="JsonElement"/> into the properties on this thread.
+    /// </summary>
+    /// <param name="serializedThread">A <see cref="JsonElement"/> representing the state of the thread.</param>
+    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    protected internal virtual async Task DeserializeAsync(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        var state = JsonSerializer.Deserialize(
+            serializedThread,
+            AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ThreadState))) as ThreadState;
+
+        if (state?.ConversationId is string threadId)
+        {
+            this.ConversationId = threadId;
+
+            // Since we have an ID, we should not have a chat message store and we can return here.
+            return;
+        }
+
+        // If we don't have any IChatMessageStore state return here.
+        if (state?.StoreState is null || state?.StoreState?.ValueKind is JsonValueKind.Undefined or JsonValueKind.Null)
+        {
+            return;
+        }
+
+        if (this._messageStore is null)
+        {
+            // If we don't have a chat message store yet, create an in-memory one.
+            this._messageStore = new InMemoryChatMessageStore();
+        }
+
+        await this._messageStore.DeserializeStateAsync(state!.StoreState.Value, jsonSerializerOptions, cancellationToken).ConfigureAwait(false);
+    }
+
+    /// <summary>
+    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
+    /// </summary>
+    /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
+    public virtual async Task<JsonElement> SerializeAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        var storeState = this._messageStore is null ?
+            (JsonElement?)null :
+            await this._messageStore.SerializeStateAsync(jsonSerializerOptions, cancellationToken).ConfigureAwait(false);
+
+        var state = new ThreadState
+        {
+            ConversationId = this.ConversationId,
+            StoreState = storeState
+        };
+
+        return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ThreadState)));
+    }
+
+    internal class ThreadState
+    {
+        public string? ConversationId { get; set; }
+
+        public JsonElement? StoreState { get; set; }
    }
 }
@@ -0,0 +1,70 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Collections.Generic;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Extensions.AI.Agents;
+
+/// <summary>
+/// Defines methods for storing and retrieving chat messages associated with a specific thread.
+/// </summary>
+/// <remarks>
+/// Implementations of this interface are responsible for managing the storage of chat messages,
+/// including handling large volumes of data by truncating or summarizing messages as necessary.
+/// </remarks>
+public interface IChatMessageStore
+{
+    /// <summary>
+    /// Gets all the messages from the store that should be used for the next agent invocation.
+    /// </summary>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>A collection of chat messages.</returns>
+    /// <remarks>
+    /// <para>
+    /// Messages are returned in ascending chronological order, with the oldest message first.
+    /// </para>
+    /// <para>
+    /// If the messages stored in the store become very large, it is up to the store to
+    /// truncate, summarize or otherwise limit the number of messages returned.
+    /// </para>
+    /// <para>
+    /// When using implementations of <see cref="IChatMessageStore"/>, a new one should be created for each thread
+    /// since they may contain state that is specific to a thread.
+    /// </para>
+    /// </remarks>
+    Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Adds messages to the store.
+    /// </summary>
+    /// <param name="messages">The messages to add.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>An async task.</returns>
+    Task AddMessagesAsync(IReadOnlyCollection<ChatMessage> messages, CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Deserializes the state contained in the provided <see cref="JsonElement"/> into the properties on this store.
+    /// </summary>
+    /// <param name="serializedStoreState">A <see cref="JsonElement"/> representing the state of the store.</param>
+    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <remarks>
+    /// This method, together with <see cref="SerializeStateAsync(JsonSerializerOptions?, CancellationToken)"/> can be used to save and load messages from a persistent store
+    /// if this store only has messages in memory.
+    /// </remarks>
+    ValueTask DeserializeStateAsync(JsonElement? serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
+    /// </summary>
+    /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
+    /// <remarks>
+    /// This method, together with <see cref="DeserializeStateAsync(JsonElement?, JsonSerializerOptions?, CancellationToken)"/> can be used to save and load messages from a persistent store
+    /// if this store only has messages in memory.
+    /// </remarks>
+    ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default);
+}
@@ -1,36 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Collections.Generic;
-using System.Threading;
-
-namespace Microsoft.Extensions.AI.Agents;
-
-/// <summary>
-/// An interface for agent threads that allow retrieval of messages in the thread for agent invocation.
-/// </summary>
-/// <remarks>
-/// <para>
-/// Some agents need to be invoked with all relevant chat history messages in order to produce a result, while some must be invoked
-/// with the id of a server side thread that contains the chat history.
-/// </para>
-/// <para>
-/// This interface can be implemented by all thread types that support the case where the agent is invoked with the chat history.
-/// Implementations must consider the size of the messages provided, so that they do not exceed the maximum size of the context window
-/// of the agent they are used with. Where appropriate, implementations should truncate or summarize messages so that the size of messages
-/// are constrained.
-/// </para>
-/// </remarks>
-public interface IMessagesRetrievableThread
-{
-    /// <summary>
-    /// Asynchronously retrieves all messages to be used for the agent invocation.
-    /// </summary>
-    /// <remarks>
-    /// Messages are returned in ascending chronological order.
-    /// </remarks>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>The messages in the thread.</returns>
-    /// <exception cref="InvalidOperationException">The thread has been deleted.</exception>
-    IAsyncEnumerable<ChatMessage> GetMessagesAsync(CancellationToken cancellationToken = default);
-}
@@ -0,0 +1,121 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Collections;
+using System.Collections.Generic;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Extensions.AI.Agents;
+
+/// <summary>
+/// Represents an in-memory store for chat messages associated with a specific thread.
+/// </summary>
+internal class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageStore
+{
+    private readonly List<ChatMessage> _messages = new();
+
+    /// <inheritdoc />
+    public int Count => this._messages.Count;
+
+    /// <inheritdoc />
+    public bool IsReadOnly => ((IList)this._messages).IsReadOnly;
+
+    /// <inheritdoc />
+    public ChatMessage this[int index]
+    {
+        get => this._messages[index];
+        set => this._messages[index] = value;
+    }
+
+    /// <inheritdoc />
+    public Task AddMessagesAsync(IReadOnlyCollection<ChatMessage> messages, CancellationToken cancellationToken)
+    {
+        _ = Throw.IfNull(messages);
+        this._messages.AddRange(messages);
+        return Task.CompletedTask;
+    }
+
+    /// <inheritdoc />
+    public Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken)
+    {
+        return Task.FromResult<IEnumerable<ChatMessage>>(this._messages);
+    }
+
+    /// <inheritdoc />
+    public ValueTask DeserializeStateAsync(JsonElement? serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        if (serializedStoreState is null)
+        {
+            return new ValueTask();
+        }
+
+        var state = JsonSerializer.Deserialize(
+            serializedStoreState.Value,
+            AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState;
+
+        if (state?.Messages is { Count: > 0 } messages)
+        {
+            this._messages.AddRange(messages);
+        }
+
+        return new ValueTask();
+    }
+
+    /// <inheritdoc />
+    public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    {
+        StoreState state = new()
+        {
+            Messages = this._messages,
+        };
+
+        return new ValueTask<JsonElement?>(JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))));
+    }
+
+    /// <inheritdoc />
+    public int IndexOf(ChatMessage item)
+        => this._messages.IndexOf(item);
+
+    /// <inheritdoc />
+    public void Insert(int index, ChatMessage item)
+        => this._messages.Insert(index, item);
+
+    /// <inheritdoc />
+    public void RemoveAt(int index)
+        => this._messages.RemoveAt(index);
+
+    /// <inheritdoc />
+    public void Add(ChatMessage item)
+        => this._messages.Add(item);
+
+    /// <inheritdoc />
+    public void Clear()
+        => this._messages.Clear();
+
+    /// <inheritdoc />
+    public bool Contains(ChatMessage item)
+        => this._messages.Contains(item);
+
+    /// <inheritdoc />
+    public void CopyTo(ChatMessage[] array, int arrayIndex)
+        => this._messages.CopyTo(array, arrayIndex);
+
+    /// <inheritdoc />
+    public bool Remove(ChatMessage item)
+        => this._messages.Remove(item);
+
+    /// <inheritdoc />
+    public IEnumerator<ChatMessage> GetEnumerator()
+        => this._messages.GetEnumerator();
+
+    /// <inheritdoc />
+    IEnumerator IEnumerable.GetEnumerator()
+        => this.GetEnumerator();
+
+    internal class StoreState
+    {
+        public IList<ChatMessage> Messages { get; set; } = new List<ChatMessage>();
+    }
+}
@@ -10,6 +10,7 @@

  <PropertyGroup>
    <InjectSharedThrow>true</InjectSharedThrow>
+    <InjectDiagnosticClassesOnLegacy>true</InjectDiagnosticClassesOnLegacy>
    <InjectTrimAttributesOnLegacy>true</InjectTrimAttributesOnLegacy>
  </PropertyGroup>

@@ -37,12 +37,6 @@ public class CopilotStudioAgent : AIAgent
        this._logger = (loggerFactory ?? NullLoggerFactory.Instance).CreateLogger<CopilotStudioAgent>();
    }

-    /// <inheritdoc/>
-    public override AgentThread GetNewThread()
-    {
-        return new CopilotStudioAgentThread();
-    }
-
    /// <inheritdoc/>
    public override async Task<AgentRunResponse> RunAsync(
        IReadOnlyCollection<ChatMessage> messages,
@@ -54,12 +48,12 @@ public class CopilotStudioAgent : AIAgent

        // Ensure that we have a valid thread to work with.
        // If the thread ID is null, we need to start a new conversation and set the thread ID accordingly.
-        CopilotStudioAgentThread copilotStudioAgentThread = base.ValidateOrCreateThreadType(thread, () => new CopilotStudioAgentThread());
-        copilotStudioAgentThread.Id ??= await this.StartNewConversationAsync(cancellationToken).ConfigureAwait(false);
+        thread ??= this.GetNewThread();
+        thread.ConversationId ??= await this.StartNewConversationAsync(cancellationToken).ConfigureAwait(false);

        // Invoke the Copilot Studio agent with the provided messages.
        string question = string.Join("\n", messages.Select(m => m.Text));
-        var responseMessages = ActivityProcessor.ProcessActivityAsync(this.Client.AskQuestionAsync(question, copilotStudioAgentThread.Id, cancellationToken), streaming: false, this._logger);
+        var responseMessages = ActivityProcessor.ProcessActivityAsync(this.Client.AskQuestionAsync(question, thread.ConversationId, cancellationToken), streaming: false, this._logger);
        var responseMessagesList = new List<ChatMessage>();
        await foreach (var message in responseMessages.ConfigureAwait(false))
        {
@@ -87,12 +81,12 @@ public class CopilotStudioAgent : AIAgent

        // Ensure that we have a valid thread to work with.
        // If the thread ID is null, we need to start a new conversation and set the thread ID accordingly.
-        CopilotStudioAgentThread copilotStudioAgentThread = base.ValidateOrCreateThreadType(thread, () => new CopilotStudioAgentThread());
-        copilotStudioAgentThread.Id ??= await this.StartNewConversationAsync(cancellationToken).ConfigureAwait(false);
+        thread ??= this.GetNewThread();
+        thread.ConversationId ??= await this.StartNewConversationAsync(cancellationToken).ConfigureAwait(false);

        // Invoke the Copilot Studio agent with the provided messages.
        string question = string.Join("\n", messages.Select(m => m.Text));
-        var responseMessages = ActivityProcessor.ProcessActivityAsync(this.Client.AskQuestionAsync(question, copilotStudioAgentThread.Id, cancellationToken), streaming: true, this._logger);
+        var responseMessages = ActivityProcessor.ProcessActivityAsync(this.Client.AskQuestionAsync(question, thread.ConversationId, cancellationToken), streaming: true, this._logger);

        // Enumerate the response messages
        await foreach (ChatMessage message in responseMessages.ConfigureAwait(false))
@@ -1,8 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-namespace Microsoft.Extensions.AI.Agents.CopilotStudio;
-
-/// <summary>
-/// Represents a thread for interacting with a Copilot Studio agent.
-/// </summary>
-public class CopilotStudioAgentThread : AgentThread;
@@ -1,20 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System.Collections.Generic;
-using System.Text.Json;
-using System.Text.Json.Serialization;
-
-namespace Microsoft.Extensions.AI.Agents;
-
-/// <summary>
-/// Source-generated JSON type information for use by all Agents implementations.
-/// </summary>
-[JsonSourceGenerationOptions(
-    JsonSerializerDefaults.Web,
-    UseStringEnumConverter = true,
-    DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
-    WriteIndented = false)]
-[JsonSerializable(typeof(ChatMessage))]
-[JsonSerializable(typeof(List<ChatMessage>))]
-[JsonSerializable(typeof(ChatClientAgentThread))]
-internal sealed partial class AgentsJsonContext : JsonSerializerContext;
@@ -100,7 +100,7 @@ public sealed class ChatClientAgent : AIAgent
    {
        Throw.IfNull(messages);

-        (ChatClientAgentThread chatClientThread, ChatOptions? chatOptions, List<ChatMessage> threadMessages) =
+        (AgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> threadMessages) =
            await this.PrepareThreadAndMessagesAsync(thread, messages, options, cancellationToken).ConfigureAwait(false);

        var agentName = this.GetLoggingAgentName();
@@ -113,10 +113,10 @@ public sealed class ChatClientAgent : AIAgent

        // We can derive the type of supported thread from whether we have a conversation id,
        // so let's update it and set the conversation id for the service thread case.
-        this.UpdateThreadWithTypeAndConversationId(chatClientThread, chatResponse.ConversationId);
+        this.UpdateThreadWithTypeAndConversationId(safeThread, chatResponse.ConversationId);

        // Only notify the thread of new messages if the chatResponse was successful to avoid inconsistent messages state in the thread.
-        await this.NotifyThreadOfNewMessagesAsync(chatClientThread, messages, cancellationToken).ConfigureAwait(false);
+        await this.NotifyThreadOfNewMessagesAsync(safeThread, messages, cancellationToken).ConfigureAwait(false);

        // Ensure that the author name is set for each message in the response.
        foreach (ChatMessage chatResponseMessage in chatResponse.Messages)
@@ -127,7 +127,7 @@ public sealed class ChatClientAgent : AIAgent
        // Convert the chat response messages to a valid IReadOnlyCollection for notification signatures below.
        var chatResponseMessages = chatResponse.Messages as IReadOnlyCollection<ChatMessage> ?? [.. chatResponse.Messages];

-        await this.NotifyThreadOfNewMessagesAsync(chatClientThread, chatResponseMessages, cancellationToken).ConfigureAwait(false);
+        await this.NotifyThreadOfNewMessagesAsync(safeThread, chatResponseMessages, cancellationToken).ConfigureAwait(false);

        return new(chatResponse) { AgentId = this.Id };
    }
@@ -141,7 +141,7 @@ public sealed class ChatClientAgent : AIAgent
    {
        var inputMessages = Throw.IfNull(messages);

-        (ChatClientAgentThread chatClientThread, ChatOptions? chatOptions, List<ChatMessage> threadMessages) =
+        (AgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> threadMessages) =
            await this.PrepareThreadAndMessagesAsync(thread, inputMessages, options, cancellationToken).ConfigureAwait(false);

        int messageCount = threadMessages.Count;
@@ -177,16 +177,20 @@ public sealed class ChatClientAgent : AIAgent

        // We can derive the type of supported thread from whether we have a conversation id,
        // so let's update it and set the conversation id for the service thread case.
-        this.UpdateThreadWithTypeAndConversationId(chatClientThread, chatResponse.ConversationId);
+        this.UpdateThreadWithTypeAndConversationId(safeThread, chatResponse.ConversationId);

        // To avoid inconsistent state we only notify the thread of the input messages if no error occurs after the initial request.
-        await this.NotifyThreadOfNewMessagesAsync(chatClientThread, inputMessages, cancellationToken).ConfigureAwait(false);
+        await this.NotifyThreadOfNewMessagesAsync(safeThread, inputMessages, cancellationToken).ConfigureAwait(false);

-        await this.NotifyThreadOfNewMessagesAsync(chatClientThread, chatResponseMessages, cancellationToken).ConfigureAwait(false);
+        await this.NotifyThreadOfNewMessagesAsync(safeThread, chatResponseMessages, cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc/>
-    public override AgentThread GetNewThread() => new ChatClientAgentThread();
+    public override AgentThread GetNewThread()
+    {
+        var thread = new AgentThread() { MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke() };
+        return thread;
+    }

    #region Private

@@ -312,7 +316,7 @@ public sealed class ChatClientAgent : AIAgent
    /// <param name="runOptions">Optional parameters for agent invocation.</param>
    /// <param name="cancellationToken">The cancellation token.</param>
    /// <returns>A tuple containing the thread, chat options, and thread messages.</returns>
-    private async Task<(ChatClientAgentThread, ChatOptions?, List<ChatMessage>)> PrepareThreadAndMessagesAsync(
+    private async Task<(AgentThread, ChatOptions?, List<ChatMessage>)> PrepareThreadAndMessagesAsync(
        AgentThread? thread,
        IReadOnlyCollection<ChatMessage> inputMessages,
        AgentRunOptions? runOptions,
@@ -320,16 +324,13 @@ public sealed class ChatClientAgent : AIAgent
    {
        ChatOptions? chatOptions = this.CreateConfiguredChatOptions(runOptions);

-        var chatClientThread = this.ValidateOrCreateThreadType<ChatClientAgentThread>(thread, () => new());
+        thread ??= this.GetNewThread();

        // Add any existing messages from the thread to the messages to be sent to the chat client.
        List<ChatMessage> threadMessages = [];
-        if (chatClientThread is IMessagesRetrievableThread messagesRetrievableThread)
+        await foreach (ChatMessage message in thread.GetMessagesAsync(cancellationToken).ConfigureAwait(false))
        {
-            await foreach (ChatMessage message in messagesRetrievableThread.GetMessagesAsync(cancellationToken).ConfigureAwait(false))
-            {
-                threadMessages.Add(message);
-            }
+            threadMessages.Add(message);
        }

        // Update the messages with agent instructions.
@@ -340,39 +341,43 @@ public sealed class ChatClientAgent : AIAgent

        // If a user provided two different thread ids, via the thread object and options, we should throw
        // since we don't know which one to use.
-        if (!string.IsNullOrWhiteSpace(chatClientThread.Id) && !string.IsNullOrWhiteSpace(chatOptions?.ConversationId) && chatClientThread.Id != chatOptions.ConversationId)
+        if (!string.IsNullOrWhiteSpace(thread.ConversationId) && !string.IsNullOrWhiteSpace(chatOptions?.ConversationId) && thread.ConversationId != chatOptions.ConversationId)
        {
            throw new InvalidOperationException(
                $"The {nameof(chatOptions.ConversationId)} provided via {nameof(Microsoft.Extensions.AI.ChatOptions)} is different to the id of the provided {nameof(AgentThread)}. Only one thread id can be used for a run.");
        }

        // Only clone and update ChatOptions if we have an id on the thread and we don't have the same one already in ChatOptions.
-        if (!string.IsNullOrWhiteSpace(chatClientThread.Id) && chatClientThread.Id != chatOptions?.ConversationId)
+        if (!string.IsNullOrWhiteSpace(thread.ConversationId) && thread.ConversationId != chatOptions?.ConversationId)
        {
            chatOptions ??= new();
-            chatOptions.ConversationId = chatClientThread.Id;
+            chatOptions.ConversationId = thread.ConversationId;
        }

-        return (chatClientThread, chatOptions, threadMessages);
+        return (thread, chatOptions, threadMessages);
    }

-    private void UpdateThreadWithTypeAndConversationId(ChatClientAgentThread chatClientThread, string? responseConversationId)
+    private void UpdateThreadWithTypeAndConversationId(AgentThread thread, string? responseConversationId)
    {
-        // Set the thread's storage location, the first time that we use it.
-        chatClientThread.StorageLocation ??= string.IsNullOrWhiteSpace(responseConversationId)
-            ? ChatClientAgentThreadType.InMemoryMessages
-            : ChatClientAgentThreadType.ConversationId;
-
-        // If we got a conversation id back from the chat client, it means that the service supports server side thread storage
-        // so we should capture the id and update the thread with the new id.
-        if (chatClientThread.StorageLocation == ChatClientAgentThreadType.ConversationId)
+        if (string.IsNullOrWhiteSpace(responseConversationId) && !string.IsNullOrWhiteSpace(thread.ConversationId))
        {
-            if (string.IsNullOrWhiteSpace(responseConversationId))
-            {
-                throw new InvalidOperationException("Service did not return a valid conversation id when using a service managed thread.");
-            }
+            // We were passed a thread that is service managed, but we got no conversation id back from the chat client,
+            // meaning the service doesn't support service managed threads, so the thread cannot be used with this service.
+            throw new InvalidOperationException("Service did not return a valid conversation id when using a service managed thread.");
+        }

-            chatClientThread.Id = responseConversationId;
+        if (!string.IsNullOrWhiteSpace(responseConversationId))
+        {
+            // If we got a conversation id back from the chat client, it means that the service supports server side thread storage
+            // so we should update the thread with the new id.
+            thread.ConversationId = responseConversationId;
+        }
+        else if (thread.MessageStore is null)
+        {
+            // If the service doesn't use service side thread storage (i.e. we got no id back from invocation), and
+            // the thread has no MessageStore yet, and we have a custom messages store, we should update the thread
+            // with the custom MessageStore so that it has somewhere to store the chat history.
+            thread.MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke();
        }
    }

@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.

+using System;
 using System.Collections.Generic;

 namespace Microsoft.Extensions.AI.Agents;
@@ -72,6 +73,12 @@ public class ChatClientAgentOptions
    /// </summary>
    public ChatOptions? ChatOptions { get; set; }

+    /// <summary>
+    /// Gets or sets a factory function to create an instance of <see cref="IChatMessageStore"/>
+    /// which will be used to store chat messages for this agent.
+    /// </summary>
+    public Func<IChatMessageStore>? ChatMessageStoreFactory { get; set; } = null;
+
    /// <summary>
    /// Creates a new instance of <see cref="ChatClientAgentOptions"/> with the same values as this instance.
    /// </summary>
@@ -82,6 +89,7 @@ public class ChatClientAgentOptions
            Name = this.Name,
            Instructions = this.Instructions,
            Description = this.Description,
-            ChatOptions = this.ChatOptions?.Clone()
+            ChatOptions = this.ChatOptions?.Clone(),
+            ChatMessageStoreFactory = this.ChatMessageStoreFactory
        };
 }
@@ -1,185 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Collections.Generic;
-using System.ComponentModel;
-using System.Diagnostics;
-using System.Runtime.CompilerServices;
-using System.Text.Json;
-using System.Text.Json.Serialization;
-using System.Threading;
-using System.Threading.Tasks;
-using Microsoft.Shared.Diagnostics;
-
-namespace Microsoft.Extensions.AI.Agents;
-
-/// <summary>
-/// Chat client agent thread.
-/// </summary>
-[JsonConverter(typeof(Converter))]
-public sealed class ChatClientAgentThread : AgentThread, IMessagesRetrievableThread
-{
-    private readonly List<ChatMessage> _chatMessages = [];
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="ChatClientAgentThread"/> class.
-    /// </summary>
-    public ChatClientAgentThread()
-    {
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="ChatClientAgentThread"/> class.
-    /// </summary>
-    /// <param name="id">The id of an existing server side thread to continue.</param>
-    /// <remarks>
-    /// This constructor creates a <see cref="ChatClientAgentThread"/> that supports in-service message storage.
-    /// </remarks>
-    public ChatClientAgentThread(string id)
-    {
-        Throw.IfNullOrWhitespace(id);
-
-        this.Id = id;
-        this.StorageLocation = ChatClientAgentThreadType.ConversationId;
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="ChatClientAgentThread"/> class.
-    /// </summary>
-    /// <param name="messages">A set of initial messages to seed the thread with.</param>
-    /// <remarks>
-    /// This constructor creates a <see cref="ChatClientAgentThread"/> that supports local in-memory message storage.
-    /// </remarks>
-    public ChatClientAgentThread(IEnumerable<ChatMessage> messages)
-    {
-        Throw.IfNull(messages);
-
-        this._chatMessages.AddRange(messages);
-        this.StorageLocation = ChatClientAgentThreadType.InMemoryMessages;
-    }
-
-    /// <summary>
-    /// Gets the location of the thread contents.
-    /// </summary>
-    internal ChatClientAgentThreadType? StorageLocation { get; set; }
-
-#pragma warning disable CS1998 // Async method lacks 'await' operators and will run synchronously
-    /// <inheritdoc/>
-    public async IAsyncEnumerable<ChatMessage> GetMessagesAsync([EnumeratorCancellation] CancellationToken cancellationToken = default)
-    {
-        foreach (var message in this._chatMessages)
-        {
-            yield return message;
-        }
-    }
-
-#pragma warning restore CS1998 // Async method lacks 'await' operators and will run synchronously
-
-    /// <inheritdoc/>
-    protected override Task OnNewMessagesAsync(IReadOnlyCollection<ChatMessage> newMessages, CancellationToken cancellationToken = default)
-    {
-        switch (this.StorageLocation)
-        {
-            case ChatClientAgentThreadType.InMemoryMessages:
-                this._chatMessages.AddRange(newMessages);
-                break;
-            case ChatClientAgentThreadType.ConversationId:
-                // If the thread messages are stored in the service
-                // there is nothing to do here, since invoking the
-                // service should already update the thread.
-                break;
-            default:
-                throw new UnreachableException();
-        }
-
-        return Task.CompletedTask;
-    }
-
-    /// <summary>
-    /// Provides a <see cref="JsonConverter"/> for <see cref="ChatClientAgentThread"/> objects.
-    /// </summary>
-    [EditorBrowsable(EditorBrowsableState.Never)]
-    public sealed class Converter : JsonConverter<ChatClientAgentThread>
-    {
-        /// <inheritdoc/>
-        public override ChatClientAgentThread? Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
-        {
-            if (reader.TokenType != JsonTokenType.StartObject)
-            {
-                throw new JsonException("Expected StartObject token");
-            }
-
-            using var doc = JsonDocument.ParseValue(ref reader);
-            var root = doc.RootElement;
-
-            // Extract properties from JSON
-            string? id = null;
-            if (root.TryGetProperty("id", out var idProperty))
-            {
-                id = idProperty.GetString();
-            }
-
-            List<ChatMessage>? messages = null;
-            if (root.TryGetProperty("messages", out var messagesProperty))
-            {
-                if (messagesProperty.ValueKind == JsonValueKind.Array)
-                {
-                    messages = [];
-                    foreach (var messageElement in messagesProperty.EnumerateArray())
-                    {
-                        var message = messageElement.Deserialize(options.GetTypeInfo<ChatMessage>(AgentsJsonContext.Default));
-                        if (message != null)
-                        {
-                            messages.Add(message);
-                        }
-                    }
-                }
-            }
-
-            // Create the appropriate instance based on available data
-            // StorageLocation will be set automatically by the constructors
-            ChatClientAgentThread thread;
-            if (messages?.Count > 0)
-            {
-                thread = new ChatClientAgentThread(messages);
-            }
-            else if (!string.IsNullOrWhiteSpace(id))
-            {
-                thread = new ChatClientAgentThread(id);
-            }
-            else
-            {
-                thread = new ChatClientAgentThread();
-            }
-
-            // Override Id if it was explicitly set in JSON (for cases where messages exist but ID is also provided)
-            if (id != null)
-            {
-                thread.Id = id;
-            }
-
-            return thread;
-        }
-
-        /// <inheritdoc/>
-        public override void Write(Utf8JsonWriter writer, ChatClientAgentThread value, JsonSerializerOptions options)
-        {
-            writer.WriteStartObject();
-
-            // Write base properties
-            if (value.Id != null)
-            {
-                writer.WriteString("id", value.Id);
-            }
-
-            // Write messages if in memory storage (StorageLocation is determined by presence of messages vs ID)
-            if (value.StorageLocation == ChatClientAgentThreadType.InMemoryMessages)
-            {
-                writer.WritePropertyName("messages");
-                JsonSerializer.Serialize(writer, value._chatMessages, options.GetTypeInfo<List<ChatMessage>>(AgentsJsonContext.Default));
-            }
-
-            writer.WriteEndObject();
-        }
-    }
-}
@@ -1,19 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-namespace Microsoft.Extensions.AI.Agents;
-
-/// <summary>
-/// Defines the different supported storage locations for <see cref="ChatClientAgentThread"/>.
-/// </summary>
-internal enum ChatClientAgentThreadType
-{
-    /// <summary>
-    /// Messages are stored in memory inside the thread object.
-    /// </summary>
-    InMemoryMessages,
-
-    /// <summary>
-    /// Messages are stored in the service and the thread object just has an id reference the service storage.
-    /// </summary>
-    ConversationId
-}
@@ -210,9 +210,9 @@ public sealed class OpenTelemetryAgent : AIAgent, IDisposable
                }

                // Add conversation ID if thread is available (following gen_ai.conversation.id convention)
-                if (!string.IsNullOrWhiteSpace(thread?.Id))
+                if (!string.IsNullOrWhiteSpace(thread?.ConversationId))
                {
-                    _ = activity.AddTag(AgentOpenTelemetryConsts.GenAI.ConversationId, thread.Id);
+                    _ = activity.AddTag(AgentOpenTelemetryConsts.GenAI.ConversationId, thread.ConversationId);
                }

                // Add instructions if available (for ChatClientAgent)