.NET: Add per run / thread feature collection support and improved custom ChatMessageStore support (#2345)

* Add the ability to override services on an agent per run.

* Remove Run from AgentFeatureCollection name.

* Adding features param to GetNewThread.

* Move feature collection.

* Add features to DeserializeThread

* Remove servicecollection based option

* Add feature collection unit tests and fix bug identified in code review.

* Add more unit tests for DelegatingAIAgent and AgentRunOptions

* Fix formatting.

* Address PR comments.

* Switch to dedicated ConversationIdAgentFeature and improve 3rd party storage samples.

* Fix bug in sample.

dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs

@@ -28,10 +28,10 @@ namespace SampleApp
     {
         public override string? Name => "UpperCaseParrotAgent";
 
-        public override AgentThread GetNewThread()
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
             => new CustomAgentThread();
 
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
             => new CustomAgentThread(serializedThread, jsonSerializerOptions);
 
         public override async Task<AgentRunResponse> RunAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
@@ -39,11 +39,16 @@ namespace SampleApp
             // Create a thread if the user didn't supply one.
             thread ??= this.GetNewThread();
 
+            if (thread is not CustomAgentThread typedThread)
+            {
+                throw new ArgumentException($"The provided thread is not of type {nameof(CustomAgentThread)}.", nameof(thread));
+            }
+
             // Clone the input messages and turn them into response messages with upper case text.
             List<ChatMessage> responseMessages = CloneAndToUpperCase(messages, this.DisplayName).ToList();
 
             // Notify the thread of the input and output messages.
-            await NotifyThreadOfNewMessagesAsync(thread, messages.Concat(responseMessages), cancellationToken);
+            await typedThread.MessageStore.AddMessagesAsync(messages.Concat(responseMessages), cancellationToken);
 
             return new AgentRunResponse
             {
@@ -58,11 +63,16 @@ namespace SampleApp
             // Create a thread if the user didn't supply one.
             thread ??= this.GetNewThread();
 
+            if (thread is not CustomAgentThread typedThread)
+            {
+                throw new ArgumentException($"The provided thread is not of type {nameof(CustomAgentThread)}.", nameof(thread));
+            }
+
             // Clone the input messages and turn them into response messages with upper case text.
             List<ChatMessage> responseMessages = CloneAndToUpperCase(messages, this.DisplayName).ToList();
 
             // Notify the thread of the input and output messages.
-            await NotifyThreadOfNewMessagesAsync(thread, messages.Concat(responseMessages), cancellationToken);
+            await typedThread.MessageStore.AddMessagesAsync(messages.Concat(responseMessages), cancellationToken);
 
             foreach (var message in responseMessages)
             {

dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs

@@ -21,53 +21,200 @@ var deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT
 // Replace this with a vector store implementation of your choice if you want to persist the chat history to disk.
 VectorStore vectorStore = new InMemoryVectorStore();
 
-// Create the agent
-AIAgent agent = new AzureOpenAIClient(
-    new Uri(endpoint),
-    new AzureCliCredential())
-    .GetChatClient(deploymentName)
-    .CreateAIAgent(new ChatClientAgentOptions
-    {
-        Instructions = "You are good at telling jokes.",
-        Name = "Joker",
-        ChatMessageStoreFactory = ctx =>
+// Execute various samples showing how to use a custom ChatMessageStore with an agent.
+await CustomChatMessageStore_UsingFactory_Async();
+await CustomChatMessageStore_UsingFactoryAndExistingExternalId_Async();
+await CustomChatMessageStore_PerThread_Async();
+await CustomChatMessageStore_PerRun_Async();
+
+// Here we can see how to create a custom ChatMessageStore using a factory method
+// provided to the agent via the ChatMessageStoreFactory option.
+// This allows us to use a custom chat message store, where the consumer of the agent
+// doesn't need to know anything about the storage mechanism used.
+async Task CustomChatMessageStore_UsingFactory_Async()
+{
+    Console.WriteLine("\n--- With Factory ---\n");
+
+    // Create the agent
+    AIAgent agent = new AzureOpenAIClient(
+        new Uri(endpoint),
+        new AzureCliCredential())
+        // Use a service that doesn't require storage of chat history in the service itself.
+        .GetChatClient(deploymentName)
+        .CreateAIAgent(new ChatClientAgentOptions
         {
-            // Create a new chat message store for this agent that stores the messages in a vector store.
-            // Each thread must get its own copy of the VectorChatMessageStore, since the store
-            // also contains the id that the thread is stored under.
-            return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions);
-        }
-    });
+            Instructions = "You are good at telling jokes.",
+            Name = "Joker",
+            ChatMessageStoreFactory = ctx =>
+            {
+                // Create a new chat message store for this agent that stores the messages in a vector store.
+                // Each thread must get its own copy of the VectorChatMessageStore, since the store
+                // also contains the id that the thread is stored under.
+                return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions, ctx.Features);
+            }
+        });
 
-// Start a new thread for the agent conversation.
-AgentThread thread = agent.GetNewThread();
+    // Start a new thread for the agent conversation.
+    AgentThread thread = agent.GetNewThread();
 
-// Run the agent with the thread that stores conversation history in the vector store.
-Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
+    // Run the agent with the thread that stores conversation history in the vector store.
+    Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
 
-// Serialize the thread state, so it can be stored for later use.
-// Since the chat history is stored in the vector store, the serialized thread
-// only contains the guid that the messages are stored under in the vector store.
-JsonElement serializedThread = thread.Serialize();
+    // Serialize the thread state, so it can be stored for later use.
+    // Since the chat history is stored in the vector store, the serialized thread
+    // only contains the guid that the messages are stored under in the vector store.
+    JsonElement serializedThread = thread.Serialize();
 
-Console.WriteLine("\n--- Serialized thread ---\n");
-Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
+    Console.WriteLine("\n--- Serialized thread ---\n");
+    Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
 
-// The serialized thread can now be saved to a database, file, or any other storage mechanism
-// and loaded again later.
+    // The serialized thread can now be saved to a database, file, or any other storage mechanism
+    // and loaded again later.
 
-// Deserialize the thread state after loading from storage.
-AgentThread resumedThread = agent.DeserializeThread(serializedThread);
+    // Deserialize the thread state after loading from storage.
+    AgentThread resumedThread = agent.DeserializeThread(serializedThread);
 
-// Run the agent with the thread that stores conversation history in the vector store a second time.
-Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
+    // Run the agent with the thread that stores conversation history in the vector store a second time.
+    Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
+}
 
-// We can access the VectorChatMessageStore via the thread's GetService method if we need to read the key under which threads are stored.
-var messageStore = resumedThread.GetService<VectorChatMessageStore>()!;
-Console.WriteLine($"\nThread is stored in vector store under key: {messageStore.ThreadDbKey}");
+// Here we can see how to create a custom ChatMessageStore using a factory method
+// provided to the agent via the ChatMessageStoreFactory option.
+// It also shows how we can pass a custom storage id at runtime to the message store using
+// the VectorChatMessageStoreThreadDbKeyFeature.
+// Note that not all agents or chat message stores may support this feature.
+async Task CustomChatMessageStore_UsingFactoryAndExistingExternalId_Async()
+{
+    Console.WriteLine("\n--- With Factory and Existing External ID ---\n");
+
+    // Create the agent
+    AIAgent agent = new AzureOpenAIClient(
+        new Uri(endpoint),
+        new AzureCliCredential())
+        // Use a service that doesn't require storage of chat history in the service itself.
+        .GetChatClient(deploymentName)
+        .CreateAIAgent(new ChatClientAgentOptions
+        {
+            Instructions = "You are good at telling jokes.",
+            Name = "Joker",
+            ChatMessageStoreFactory = ctx =>
+            {
+                // Create a new chat message store for this agent that stores the messages in a vector store.
+                // Each thread must get its own copy of the VectorChatMessageStore, since the store
+                // also contains the id that the thread is stored under.
+                return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions, ctx.Features);
+            }
+        });
+
+    // Start a new thread for the agent conversation.
+    AgentThread thread = agent.GetNewThread();
+
+    // Run the agent with the thread that stores conversation history in the vector store.
+    Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
+
+    // We can access the VectorChatMessageStore via the thread's GetService method if we need to read the key under which threads are stored.
+    var messageStoreFromFactory = thread.GetService<VectorChatMessageStore>()!;
+    Console.WriteLine($"\nThread is stored in vector store under key: {messageStoreFromFactory.ThreadDbKey}");
+
+    // It's possible to create a new thread that uses the same chat message store id by providing
+    // the VectorChatMessageStoreThreadDbKeyFeature in the feature collection when creating the new thread.
+    AgentFeatureCollection features = new();
+    features.Set(new VectorChatMessageStoreThreadDbKeyFeature(messageStoreFromFactory.ThreadDbKey!));
+    AgentThread resumedThread = agent.GetNewThread(features);
+
+    // Run the agent with the thread that stores conversation history in the vector store.
+    Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
+}
+
+// Here we can see how to create a custom ChatMessageStore and pass it to the thread
+// when creating a new thread.
+async Task CustomChatMessageStore_PerThread_Async()
+{
+    Console.WriteLine("\n--- Per Thread ---\n");
+
+    // We can also create an agent without a factory that provides a ChatMessageStore.
+    AIAgent agent = new AzureOpenAIClient(
+        new Uri(endpoint),
+        new AzureCliCredential())
+        // Use a service that doesn't require storage of chat history in the service itself.
+        .GetChatClient(deploymentName)
+        .CreateAIAgent(new ChatClientAgentOptions
+        {
+            Instructions = "You are good at telling jokes.",
+            Name = "Joker"
+        });
+
+    // Instead of using a factory on the agent to create the ChatMessageStore, we can
+    // create a VectorChatMessageStore ourselves and register it in a feature collection.
+    // We can then pass the feature collection when creating a new thread.
+    // We also have the opportunity here to pass any id that we want for storing the chat history in the vector store.
+    VectorChatMessageStore perThreadMessageStore = new(vectorStore, "chat-history-1");
+    AgentFeatureCollection features = new();
+    features.Set<ChatMessageStore>(perThreadMessageStore);
+    AgentThread thread = agent.GetNewThread(features);
+
+    Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
+
+    // When serializing this thread, we'll see that it has the id from the message store stored in its state.
+    JsonElement serializedThread = thread.Serialize();
+
+    Console.WriteLine("\n--- Serialized thread ---\n");
+    Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
+}
+
+// Here we can see how to create a custom ChatMessageStore for a single run using the Features option
+// passed when we run the agent.
+// Note that if the agent doesn't support a chat message store, it would be ignored.
+async Task CustomChatMessageStore_PerRun_Async()
+{
+    Console.WriteLine("\n--- Per Run ---\n");
+
+    // We can also create an agent without a factory that provides a ChatMessageStore.
+    AIAgent agent = new AzureOpenAIClient(
+        new Uri(endpoint),
+        new AzureCliCredential())
+        // Use a service that doesn't require storage of chat history in the service itself.
+        .GetChatClient(deploymentName)
+        .CreateAIAgent(new ChatClientAgentOptions
+        {
+            Instructions = "You are good at telling jokes.",
+            Name = "Joker"
+        });
+
+    // Start a new thread for the agent conversation.
+    AgentThread thread = agent.GetNewThread();
+
+    // Instead of using a factory on the agent to create the ChatMessageStore, we can
+    // create a VectorChatMessageStore ourselves and register it in a feature collection.
+    // We can then pass the feature collection to the agent when running it by using the Features option.
+    // The message store would only be used for the run that it's passed to.
+    // If the agent doesn't support a message store, it would be ignored.
+    // We also have the opportunity here to pass any id that we want for storing the chat history in the vector store.
+    VectorChatMessageStore perRunMessageStore = new(vectorStore, "chat-history-1");
+    AgentFeatureCollection features = new();
+    features.Set<ChatMessageStore>(perRunMessageStore);
+
+    Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread, options: new AgentRunOptions() { Features = features }));
+
+    // When serializing this thread, we'll see that it has no messagestore state, since the messagestore was not attached to the thread,
+    // but just provided for the single run. Note that, depending on the circumstances, the thread may still contain other state, e.g. Memories,
+    // if an AIContextProvider is attached which adds memory to an agent.
+    JsonElement serializedThread = thread.Serialize();
+
+    Console.WriteLine("\n--- Serialized thread ---\n");
+    Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
+}
 
 namespace SampleApp
 {
+    /// <summary>
+    /// A feature that allows providing the thread database key for the <see cref="VectorChatMessageStore"/>.
+    /// </summary>
+    internal sealed class VectorChatMessageStoreThreadDbKeyFeature(string threadDbKey)
+    {
+        public string ThreadDbKey { get; } = threadDbKey;
+    }
+
     /// <summary>
     /// A sample implementation of <see cref="ChatMessageStore"/> that stores chat messages in a vector store.
     /// </summary>
@@ -75,29 +222,35 @@ namespace SampleApp
     {
         private readonly VectorStore _vectorStore;
 
-        public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
+        public VectorChatMessageStore(VectorStore vectorStore, string threadDbKey)
+        {
+            this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
+            this.ThreadDbKey = threadDbKey ?? throw new ArgumentNullException(nameof(threadDbKey));
+        }
+
+        public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? features = null)
         {
             this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
 
-            if (serializedStoreState.ValueKind is JsonValueKind.String)
-            {
-                // Here we can deserialize the thread id so that we can access the same messages as before the suspension.
-                this.ThreadDbKey = serializedStoreState.Deserialize<string>();
-            }
+            // Here we can deserialize the thread id so that we can access the same messages as before the suspension, or if
+            // a user provided a ConversationIdAgentFeature in the features collection, we can use that
+            // or finally we can generate one ourselves.
+            this.ThreadDbKey = serializedStoreState.ValueKind is JsonValueKind.String
+                ? serializedStoreState.Deserialize<string>()
+                : features?.Get<VectorChatMessageStoreThreadDbKeyFeature>()?.ThreadDbKey
+                ?? Guid.NewGuid().ToString("N");
         }
 
-        public string? ThreadDbKey { get; private set; }
+        public string? ThreadDbKey { get; }
 
         public override async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
         {
-            this.ThreadDbKey ??= Guid.NewGuid().ToString("N");
-
             var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
             await collection.EnsureCollectionExistsAsync(cancellationToken);
 
             await collection.UpsertAsync(messages.Select(x => new ChatHistoryItem()
             {
-                Key = this.ThreadDbKey + x.MessageId,
+                Key = this.ThreadDbKey + (string.IsNullOrWhiteSpace(x.MessageId) ? Guid.NewGuid().ToString("N") : x.MessageId),
                 Timestamp = DateTimeOffset.UtcNow,
                 ThreadId = this.ThreadDbKey,
                 SerializedMessage = JsonSerializer.Serialize(x),

dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs

@@ -54,8 +54,8 @@ internal sealed class A2AAgent : AIAgent
     }
 
     /// <inheritdoc/>
-    public sealed override AgentThread GetNewThread()
-        => new A2AAgentThread();
+    public sealed override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
+        => new A2AAgentThread() { ContextId = featureCollection?.Get<ConversationIdAgentFeature>()?.ConversationId };
 
     /// <summary>
     /// Get a new <see cref="AgentThread"/> instance using an existing context id, to continue that conversation.
@@ -66,7 +66,7 @@ internal sealed class A2AAgent : AIAgent
         => new A2AAgentThread() { ContextId = contextId };
 
     /// <inheritdoc/>
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
         => new A2AAgentThread(serializedThread, jsonSerializerOptions);
 
     /// <inheritdoc/>

dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgent.cs

@@ -108,6 +108,7 @@ public abstract class AIAgent
     /// <summary>
     /// Creates a new conversation thread that is compatible with this agent.
     /// </summary>
+    /// <param name="featureCollection">An optional feature collection to override or provide additional context or capabilities to the thread where the thread supports these features.</param>
     /// <returns>A new <see cref="AgentThread"/> instance ready for use with this agent.</returns>
     /// <remarks>
     /// <para>
@@ -121,13 +122,14 @@ public abstract class AIAgent
     /// may be deferred until first use to optimize performance.
     /// </para>
     /// </remarks>
-    public abstract AgentThread GetNewThread();
+    public abstract AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null);
 
     /// <summary>
     /// Deserializes an agent thread from its JSON serialized representation.
     /// </summary>
     /// <param name="serializedThread">A <see cref="JsonElement"/> containing the serialized thread state.</param>
     /// <param name="jsonSerializerOptions">Optional settings to customize the deserialization process.</param>
+    /// <param name="featureCollection">An optional feature collection to override or provide additional context or capabilities to the thread where the thread supports these features.</param>
     /// <returns>A restored <see cref="AgentThread"/> instance with the state from <paramref name="serializedThread"/>.</returns>
     /// <exception cref="ArgumentException">The <paramref name="serializedThread"/> is not in the expected format.</exception>
     /// <exception cref="JsonException">The serialized data is invalid or cannot be deserialized.</exception>
@@ -136,7 +138,7 @@ public abstract class AIAgent
     /// allowing conversations to resume across application restarts or be migrated between
     /// different agent instances.
     /// </remarks>
-    public abstract AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null);
+    public abstract AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null);
 
     /// <summary>
     /// Run the agent with no message assuming that all required instructions are already provided to the agent or on the thread.
@@ -328,28 +330,4 @@ public abstract class AIAgent
         AgentThread? thread = null,
         AgentRunOptions? options = null,
         CancellationToken cancellationToken = default);
-
-    /// <summary>
-    /// Notifies the specified thread about new messages that have been added to the conversation.
-    /// </summary>
-    /// <param name="thread">The conversation thread to notify about the new messages.</param>
-    /// <param name="messages">The collection of new messages to report to the thread.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A task that represents the asynchronous notification operation.</returns>
-    /// <exception cref="ArgumentNullException"><paramref name="thread"/> or <paramref name="messages"/> is <see langword="null"/>.</exception>
-    /// <remarks>
-    /// <para>
-    /// This method ensures that conversation threads are kept informed about message additions, which
-    /// is important for threads that manage their own state, memory components, or derived context.
-    /// While all agent implementations should notify their threads, the specific actions taken by
-    /// each thread type may vary.
-    /// </para>
-    /// </remarks>
-    protected static async Task NotifyThreadOfNewMessagesAsync(AgentThread thread, IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
-    {
-        _ = Throw.IfNull(thread);
-        _ = Throw.IfNull(messages);
-
-        await thread.MessagesReceivedAsync(messages, cancellationToken).ConfigureAwait(false);
-    }
 }

dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRunOptions.cs

@@ -34,6 +34,7 @@ public class AgentRunOptions
         this.ContinuationToken = options.ContinuationToken;
         this.AllowBackgroundResponses = options.AllowBackgroundResponses;
         this.AdditionalProperties = options.AdditionalProperties?.Clone();
+        this.Features = options.Features;
     }
 
     /// <summary>
@@ -90,4 +91,9 @@ public class AgentRunOptions
     /// preserving implementation-specific details or extending the options with custom data.
     /// </remarks>
     public AdditionalPropertiesDictionary? AdditionalProperties { get; set; }
+
+    /// <summary>
+    /// Gets or sets the collection of features provided by the caller and middleware for this run.
+    /// </summary>
+    public IAgentFeatureCollection? Features { get; set; }
 }

dotnet/src/Microsoft.Agents.AI.Abstractions/AgentThread.cs

@@ -1,11 +1,7 @@
 // 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;
@@ -30,8 +26,8 @@ namespace Microsoft.Agents.AI;
 /// <item><description>Chat history reduction, e.g. where messages needs to be summarized or truncated to reduce the size.</description></item>
 /// </list>
 /// An <see cref="AgentThread"/> is always constructed by an <see cref="AIAgent"/> so that the <see cref="AIAgent"/>
-/// can attach any necessary behaviors to the <see cref="AgentThread"/>. See the <see cref="AIAgent.GetNewThread()"/>
-/// and <see cref="AIAgent.DeserializeThread(JsonElement, JsonSerializerOptions?)"/> methods for more information.
+/// can attach any necessary behaviors to the <see cref="AgentThread"/>. See the <see cref="AIAgent.GetNewThread(Microsoft.Agents.AI.IAgentFeatureCollection?)"/>
+/// and <see cref="AIAgent.DeserializeThread(JsonElement, JsonSerializerOptions?, Microsoft.Agents.AI.IAgentFeatureCollection?)"/> methods for more information.
 /// </para>
 /// <para>
 /// Because of these behaviors, an <see cref="AgentThread"/> may not be reusable across different agents, since each agent
@@ -41,13 +37,13 @@ namespace Microsoft.Agents.AI;
 /// To support conversations that may need to survive application restarts or separate service requests, an <see cref="AgentThread"/> can be serialized
 /// and deserialized, so that it can be saved in a persistent store.
 /// The <see cref="AgentThread"/> provides the <see cref="Serialize(JsonSerializerOptions?)"/> method to serialize the thread to a
-/// <see cref="JsonElement"/> and the <see cref="AIAgent.DeserializeThread(JsonElement, JsonSerializerOptions?)"/> method
+/// <see cref="JsonElement"/> and the <see cref="AIAgent.DeserializeThread(JsonElement, JsonSerializerOptions?, Microsoft.Agents.AI.IAgentFeatureCollection?)"/> method
 /// can be used to deserialize the thread.
 /// </para>
 /// </remarks>
 /// <seealso cref="AIAgent"/>
-/// <seealso cref="AIAgent.GetNewThread()"/>
-/// <seealso cref="AIAgent.DeserializeThread(JsonElement, JsonSerializerOptions?)"/>
+/// <seealso cref="AIAgent.GetNewThread(Microsoft.Agents.AI.IAgentFeatureCollection?)"/>
+/// <seealso cref="AIAgent.DeserializeThread(JsonElement, JsonSerializerOptions?, Microsoft.Agents.AI.IAgentFeatureCollection?)"/>
 public abstract class AgentThread
 {
     /// <summary>
@@ -65,19 +61,6 @@ public abstract class AgentThread
     public virtual JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
         => default;
 
-    /// <summary>
-    /// This method is called when new messages have been contributed to the chat by any participant.
-    /// </summary>
-    /// <remarks>
-    /// Inheritors can use this method to update their context based on the new message.
-    /// </remarks>
-    /// <param name="newMessages">The new messages.</param>
-    /// <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 MessagesReceivedAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken = default)
-        => Task.CompletedTask;
-
     /// <summary>Asks the <see cref="AgentThread"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
     /// <param name="serviceType">The type of object being requested.</param>
     /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>

dotnet/src/Microsoft.Agents.AI.Abstractions/DelegatingAIAgent.cs

@@ -74,11 +74,11 @@ public class DelegatingAIAgent : AIAgent
     }
 
     /// <inheritdoc />
-    public override AgentThread GetNewThread() => this.InnerAgent.GetNewThread();
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null) => this.InnerAgent.GetNewThread(featureCollection);
 
     /// <inheritdoc />
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
-        => this.InnerAgent.DeserializeThread(serializedThread, jsonSerializerOptions);
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
+        => this.InnerAgent.DeserializeThread(serializedThread, jsonSerializerOptions, featureCollection);
 
     /// <inheritdoc />
     public override Task<AgentRunResponse> RunAsync(

dotnet/src/Microsoft.Agents.AI.Abstractions/Features/AgentFeatureCollection.cs

@@ -0,0 +1,192 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Linq;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Agents.AI;
+
+#pragma warning disable CA1043 // Use Integral Or String Argument For Indexers
+
+/// <summary>
+/// Default implementation for <see cref="IAgentFeatureCollection"/>.
+/// </summary>
+[DebuggerDisplay("Count = {GetCount()}")]
+[DebuggerTypeProxy(typeof(FeatureCollectionDebugView))]
+public class AgentFeatureCollection : IAgentFeatureCollection
+{
+    private static readonly KeyComparer s_featureKeyComparer = new();
+    private readonly IAgentFeatureCollection? _defaults;
+    private readonly int _initialCapacity;
+    private Dictionary<Type, object>? _features;
+    private volatile int _containerRevision;
+
+    /// <summary>
+    /// Initializes a new instance of <see cref="AgentFeatureCollection"/>.
+    /// </summary>
+    public AgentFeatureCollection()
+    {
+    }
+
+    /// <summary>
+    /// Initializes a new instance of <see cref="AgentFeatureCollection"/> with the specified initial capacity.
+    /// </summary>
+    /// <param name="initialCapacity">The initial number of elements that the collection can contain.</param>
+    /// <exception cref="ArgumentOutOfRangeException"><paramref name="initialCapacity"/> is less than 0</exception>
+    public AgentFeatureCollection(int initialCapacity)
+    {
+        Throw.IfLessThan(initialCapacity, 0);
+
+        this._initialCapacity = initialCapacity;
+    }
+
+    /// <summary>
+    /// Initializes a new instance of <see cref="AgentFeatureCollection"/> with the specified defaults.
+    /// </summary>
+    /// <param name="defaults">The feature defaults.</param>
+    public AgentFeatureCollection(IAgentFeatureCollection defaults)
+    {
+        this._defaults = defaults;
+    }
+
+    /// <inheritdoc />
+    public virtual int Revision
+    {
+        get { return this._containerRevision + (this._defaults?.Revision ?? 0); }
+    }
+
+    /// <inheritdoc />
+    public bool IsReadOnly { get { return false; } }
+
+    /// <inheritdoc />
+    public object? this[Type key]
+    {
+        get
+        {
+            Throw.IfNull(key);
+
+            return this._features != null && this._features.TryGetValue(key, out var result) ? result : this._defaults?[key];
+        }
+        set
+        {
+            Throw.IfNull(key);
+
+            if (value == null)
+            {
+                if (this._features?.Remove(key) is true)
+                {
+                    this._containerRevision++;
+                }
+                return;
+            }
+
+            if (this._features == null)
+            {
+                this._features = new Dictionary<Type, object>(this._initialCapacity);
+            }
+            this._features[key] = value;
+            this._containerRevision++;
+        }
+    }
+
+    IEnumerator IEnumerable.GetEnumerator()
+    {
+        return this.GetEnumerator();
+    }
+
+    /// <inheritdoc />
+    public IEnumerator<KeyValuePair<Type, object>> GetEnumerator()
+    {
+        if (this._features != null)
+        {
+            foreach (var pair in this._features)
+            {
+                yield return pair;
+            }
+        }
+
+        if (this._defaults != null)
+        {
+            // Don't return features masked by the wrapper.
+            foreach (var pair in this._features == null ? this._defaults : this._defaults.Except(this._features, s_featureKeyComparer))
+            {
+                yield return pair;
+            }
+        }
+    }
+
+    /// <inheritdoc />
+    public TFeature? Get<TFeature>()
+    {
+        if (typeof(TFeature).IsValueType)
+        {
+            var feature = this[typeof(TFeature)];
+            if (feature is null && Nullable.GetUnderlyingType(typeof(TFeature)) is null)
+            {
+                throw new InvalidOperationException(
+                    $"{typeof(TFeature).FullName} does not exist in the feature collection " +
+                    $"and because it is a struct the method can't return null. Use 'AgentFeatureCollection[typeof({typeof(TFeature).FullName})] is not null' to check if the feature exists.");
+            }
+            return (TFeature?)feature;
+        }
+        return (TFeature?)this[typeof(TFeature)];
+    }
+
+    /// <inheritdoc />
+    public void Set<TFeature>(TFeature? instance)
+    {
+        this[typeof(TFeature)] = instance;
+    }
+
+    // Used by the debugger. Count over enumerable is required to get the correct value.
+    private int GetCount() => this.Count();
+
+    private sealed class KeyComparer : IEqualityComparer<KeyValuePair<Type, object>>
+    {
+        public bool Equals(KeyValuePair<Type, object> x, KeyValuePair<Type, object> y)
+        {
+            return x.Key.Equals(y.Key);
+        }
+
+        public int GetHashCode(KeyValuePair<Type, object> obj)
+        {
+            return obj.Key.GetHashCode();
+        }
+    }
+
+    private sealed class FeatureCollectionDebugView(AgentFeatureCollection features)
+    {
+        private readonly AgentFeatureCollection _features = features;
+
+        [DebuggerBrowsable(DebuggerBrowsableState.RootHidden)]
+        public DictionaryItemDebugView<Type, object>[] Items => this._features.Select(pair => new DictionaryItemDebugView<Type, object>(pair)).ToArray();
+    }
+
+    /// <summary>
+    /// Defines a key/value pair for displaying an item of a dictionary by a debugger.
+    /// </summary>
+    [DebuggerDisplay("{Value}", Name = "[{Key}]")]
+    internal readonly struct DictionaryItemDebugView<TKey, TValue>
+    {
+        public DictionaryItemDebugView(TKey key, TValue value)
+        {
+            this.Key = key;
+            this.Value = value;
+        }
+
+        public DictionaryItemDebugView(KeyValuePair<TKey, TValue> keyValue)
+        {
+            this.Key = keyValue.Key;
+            this.Value = keyValue.Value;
+        }
+
+        [DebuggerBrowsable(DebuggerBrowsableState.Collapsed)]
+        public TKey Key { get; }
+
+        [DebuggerBrowsable(DebuggerBrowsableState.Collapsed)]
+        public TValue Value { get; }
+    }
+}

dotnet/src/Microsoft.Agents.AI.Abstractions/Features/ConversationIdAgentFeature.cs

@@ -0,0 +1,29 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// An agent feature that allows providing a conversation identifier.
+/// </summary>
+/// <remarks>
+/// This feature allows a user to provide a specific identifier for chat history whether stored in the underlying AI service or stored in a 3rd party store.
+/// </remarks>
+public class ConversationIdAgentFeature
+{
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ConversationIdAgentFeature"/> class with the specified thread
+    /// identifier.
+    /// </summary>
+    /// <param name="conversationId">The unique identifier of the thread required by the underlying AI service or 3rd party store. Cannot be <see langword="null"/> or empty.</param>
+    public ConversationIdAgentFeature(string conversationId)
+    {
+        this.ConversationId = Throw.IfNullOrWhitespace(conversationId);
+    }
+
+    /// <summary>
+    /// Gets the conversation identifier.
+    /// </summary>
+    public string ConversationId { get; }
+}

dotnet/src/Microsoft.Agents.AI.Abstractions/Features/IAgentFeatureCollection.cs

@@ -0,0 +1,46 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+
+namespace Microsoft.Agents.AI;
+
+#pragma warning disable CA1043 // Use Integral Or String Argument For Indexers
+#pragma warning disable CA1716 // Identifiers should not match keywords
+
+/// <summary>
+/// Represents a collection of Agent features.
+/// </summary>
+public interface IAgentFeatureCollection : IEnumerable<KeyValuePair<Type, object>>
+{
+    /// <summary>
+    /// Indicates if the collection can be modified.
+    /// </summary>
+    bool IsReadOnly { get; }
+
+    /// <summary>
+    /// Incremented for each modification and can be used to verify cached results.
+    /// </summary>
+    int Revision { get; }
+
+    /// <summary>
+    /// Gets or sets a given feature. Setting a null value removes the feature.
+    /// </summary>
+    /// <param name="key"></param>
+    /// <returns>The requested feature, or null if it is not present.</returns>
+    object? this[Type key] { get; set; }
+
+    /// <summary>
+    /// Retrieves the requested feature from the collection.
+    /// </summary>
+    /// <typeparam name="TFeature">The feature key.</typeparam>
+    /// <returns>The requested feature, or null if it is not present.</returns>
+    TFeature? Get<TFeature>();
+
+    /// <summary>
+    /// Sets the given feature in the collection.
+    /// </summary>
+    /// <typeparam name="TFeature">The feature key.</typeparam>
+    /// <param name="instance">The feature value.</param>
+    void Set<TFeature>(TFeature? instance);
+}

dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryAgentThread.cs

@@ -4,8 +4,6 @@ using System;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Text.Json;
-using System.Threading;
-using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
 
 namespace Microsoft.Agents.AI;
@@ -116,10 +114,6 @@ public abstract class InMemoryAgentThread : AgentThread
     public override object? GetService(Type serviceType, object? serviceKey = null) =>
         base.GetService(serviceType, serviceKey) ?? this.MessageStore?.GetService(serviceType, serviceKey);
 
-    /// <inheritdoc />
-    protected internal override Task MessagesReceivedAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken = default)
-        => this.MessageStore.AddMessagesAsync(newMessages, cancellationToken);
-
     [DebuggerBrowsable(DebuggerBrowsableState.Never)]
     private string DebuggerDisplay => $"Count = {this.MessageStore.Count}";
 

dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs

@@ -42,8 +42,8 @@ public class CopilotStudioAgent : AIAgent
     }
 
     /// <inheritdoc/>
-    public sealed override AgentThread GetNewThread()
-        => new CopilotStudioAgentThread();
+    public sealed override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
+        => new CopilotStudioAgentThread() { ConversationId = featureCollection?.Get<ConversationIdAgentFeature>()?.ConversationId };
 
     /// <summary>
     /// Get a new <see cref="AgentThread"/> instance using an existing conversation id, to continue that conversation.
@@ -54,7 +54,7 @@ public class CopilotStudioAgent : AIAgent
         => new CopilotStudioAgentThread() { ConversationId = conversationId };
 
     /// <inheritdoc/>
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
         => new CopilotStudioAgentThread(serializedThread, jsonSerializerOptions);
 
     /// <inheritdoc/>

dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs

@@ -33,21 +33,17 @@ public sealed class DurableAIAgent : AIAgent
     /// Creates a new agent thread for this agent using a random session ID.
     /// </summary>
     /// <returns>A new agent thread.</returns>
-    public override AgentThread GetNewThread()
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
     {
         AgentSessionId sessionId = this._context.NewAgentSessionId(this._agentName);
         return new DurableAgentThread(sessionId);
     }
 
-    /// <summary>
-    /// Deserializes an agent thread from JSON.
-    /// </summary>
-    /// <param name="serializedThread">The serialized thread data.</param>
-    /// <param name="jsonSerializerOptions">Optional JSON serializer options.</param>
-    /// <returns>The deserialized agent thread.</returns>
+    /// <inheritdoc/>
     public override AgentThread DeserializeThread(
         JsonElement serializedThread,
-        JsonSerializerOptions? jsonSerializerOptions = null)
+        JsonSerializerOptions? jsonSerializerOptions = null,
+        IAgentFeatureCollection? featureCollection = null)
     {
         return DurableAgentThread.Deserialize(serializedThread, jsonSerializerOptions);
     }

dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs

@@ -13,12 +13,13 @@ internal class DurableAIAgentProxy(string name, IDurableAgentClient agentClient)
 
     public override AgentThread DeserializeThread(
         JsonElement serializedThread,
-        JsonSerializerOptions? jsonSerializerOptions = null)
+        JsonSerializerOptions? jsonSerializerOptions = null,
+        IAgentFeatureCollection? featureCollection = null)
     {
         return DurableAgentThread.Deserialize(serializedThread, jsonSerializerOptions);
     }
 
-    public override AgentThread GetNewThread()
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
     {
         return new DurableAgentThread(AgentSessionId.WithRandomKey(this.Name!));
     }

dotnet/src/Microsoft.Agents.AI.Purview/PurviewAgent.cs

@@ -30,15 +30,15 @@ internal class PurviewAgent : AIAgent, IDisposable
     }
 
     /// <inheritdoc/>
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
     {
-        return this._innerAgent.DeserializeThread(serializedThread, jsonSerializerOptions);
+        return this._innerAgent.DeserializeThread(serializedThread, jsonSerializerOptions, featureCollection);
     }
 
     /// <inheritdoc/>
-    public override AgentThread GetNewThread()
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
     {
-        return this._innerAgent.GetNewThread();
+        return this._innerAgent.GetNewThread(featureCollection);
     }
 
     /// <inheritdoc/>

dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs

@@ -61,9 +61,9 @@ internal sealed class WorkflowHostAgent : AIAgent
         protocol.ThrowIfNotChatProtocol();
     }
 
-    public override AgentThread GetNewThread() => new WorkflowThread(this._workflow, this.GenerateNewId(), this._executionEnvironment, this._checkpointManager);
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null) => new WorkflowThread(this._workflow, this.GenerateNewId(), this._executionEnvironment, this._checkpointManager);
 
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
         => new WorkflowThread(this._workflow, serializedThread, this._executionEnvironment, this._checkpointManager, jsonSerializerOptions);
 
     private async ValueTask<WorkflowThread> UpdateThreadAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, CancellationToken cancellationToken = default)

dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowThread.cs

@@ -68,9 +68,6 @@ internal sealed class WorkflowThread : AgentThread
 
     public CheckpointInfo? LastCheckpoint { get; set; }
 
-    protected override Task MessagesReceivedAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken = default)
-        => this.MessageStore.AddMessagesAsync(newMessages, cancellationToken);
-
     public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
     {
         JsonMarshaller marshaller = new(jsonSerializerOptions);

dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs

@@ -270,7 +270,7 @@ public sealed partial class ChatClientAgent : AIAgent
         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 NotifyThreadOfNewMessagesAsync(safeThread, inputMessages.Concat(aiContextProviderMessages ?? []).Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);
+        await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages.Concat(aiContextProviderMessages ?? []).Concat(chatResponse.Messages), options, cancellationToken).ConfigureAwait(false);
 
         // Notify the AIContextProvider of all new messages.
         await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
@@ -286,10 +286,16 @@ public sealed partial class ChatClientAgent : AIAgent
         : this.ChatClient.GetService(serviceType, serviceKey));
 
     /// <inheritdoc/>
-    public override AgentThread GetNewThread()
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
         => new ChatClientAgentThread
         {
-            AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke(new() { SerializedState = default, JsonSerializerOptions = null })
+            ConversationId = featureCollection?.Get<ConversationIdAgentFeature>()?.ConversationId,
+            MessageStore =
+                featureCollection?.Get<ChatMessageStore>()
+                ?? this._agentOptions?.ChatMessageStoreFactory?.Invoke(new() { SerializedState = default, Features = featureCollection, JsonSerializerOptions = null }),
+            AIContextProvider =
+                featureCollection?.Get<AIContextProvider>()
+                ?? this._agentOptions?.AIContextProviderFactory?.Invoke(new() { SerializedState = default, Features = featureCollection, JsonSerializerOptions = null })
         };
 
     /// <summary>
@@ -316,16 +322,52 @@ public sealed partial class ChatClientAgent : AIAgent
             AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke(new() { SerializedState = default, JsonSerializerOptions = null })
         };
 
-    /// <inheritdoc/>
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-        Func<JsonElement, JsonSerializerOptions?, ChatMessageStore>? chatMessageStoreFactory = this._agentOptions?.ChatMessageStoreFactory is null ?
-            null :
-            (jse, jso) => this._agentOptions.ChatMessageStoreFactory.Invoke(new() { SerializedState = jse, JsonSerializerOptions = jso });
+    /// <summary>
+    /// Creates a new agent thread instance using an existing <see cref="ChatMessageStore"/> to continue a conversation.
+    /// </summary>
+    /// <param name="chatMessageStore">The chat history of the existing conversation to continue.</param>
+    /// <returns>
+    /// A new <see cref="AgentThread"/> instance configured to work with the provided <paramref name="chatMessageStore"/>.
+    /// </returns>
+    /// <remarks>
+    /// <para>
+    /// This method creates threads that do not support server-side conversation storage.
+    /// Some AI services require server-side conversation storage to function properly, and creating a thread
+    /// with a <see cref="ChatMessageStore"/> may not be compatible with these services.
+    /// </para>
+    /// <para>
+    /// Where a service requires server-side conversation storage, use <see cref="GetNewThread(string)"/>.
+    /// </para>
+    /// <para>
+    /// If the agent detects, during the first run, that the underlying AI service requires server-side conversation storage,
+    /// the thread will throw an exception to indicate that it cannot continue using the provided <see cref="ChatMessageStore"/>.
+    /// </para>
+    /// </remarks>
+    public AgentThread GetNewThread(ChatMessageStore chatMessageStore)
+        => new ChatClientAgentThread()
+        {
+            MessageStore = chatMessageStore,
+            AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke(new() { SerializedState = default, JsonSerializerOptions = null })
+        };
 
-        Func<JsonElement, JsonSerializerOptions?, AIContextProvider>? aiContextProviderFactory = this._agentOptions?.AIContextProviderFactory is null ?
-            null :
-            (jse, jso) => this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = jse, JsonSerializerOptions = jso });
+    /// <inheritdoc/>
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
+    {
+        var chatMessageStoreFeature = featureCollection?.Get<ChatMessageStore>();
+        Func<JsonElement, JsonSerializerOptions?, ChatMessageStore>? chatMessageStoreFactory =
+            chatMessageStoreFeature is not null
+                ? (jse, jso) => chatMessageStoreFeature
+                : this._agentOptions?.ChatMessageStoreFactory is not null
+                    ? (jse, jso) => this._agentOptions.ChatMessageStoreFactory.Invoke(new() { SerializedState = jse, Features = featureCollection, JsonSerializerOptions = jso })
+                    : null;
+
+        var aiContextProviderFeature = featureCollection?.Get<AIContextProvider>();
+        Func<JsonElement, JsonSerializerOptions?, AIContextProvider>? aiContextProviderFactory =
+            aiContextProviderFeature is not null
+                ? (jse, jso) => aiContextProviderFeature
+                : this._agentOptions?.AIContextProviderFactory is not null
+                    ? (jse, jso) => this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = jse, Features = featureCollection, JsonSerializerOptions = jso })
+                    : null;
 
         return new ChatClientAgentThread(
             serializedThread,
@@ -384,7 +426,7 @@ public sealed partial class ChatClientAgent : AIAgent
         }
 
         // Only notify the thread of new messages if the chatResponse was successful to avoid inconsistent message state in the thread.
-        await NotifyThreadOfNewMessagesAsync(safeThread, inputMessages.Concat(aiContextProviderMessages ?? []).Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);
+        await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages.Concat(aiContextProviderMessages ?? []).Concat(chatResponse.Messages), options, cancellationToken).ConfigureAwait(false);
 
         // Notify the AIContextProvider of all new messages.
         await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
@@ -598,10 +640,19 @@ public sealed partial class ChatClientAgent : AIAgent
         // Populate the thread messages only if we are not continuing an existing response as it's not allowed
         if (chatOptions?.ContinuationToken is null)
         {
-            // Add any existing messages from the thread to the messages to be sent to the chat client.
-            if (typedThread.MessageStore is not null)
+            var messageStore = typedThread.MessageStore;
+
+            // If the caller provided an override message store via run options, we should use that instead of the message store
+            // on the thread.
+            if (runOptions?.Features?.Get<ChatMessageStore>() is ChatMessageStore chatMessageStoreFeature)
             {
-                inputMessagesForChatClient.AddRange(await typedThread.MessageStore.GetMessagesAsync(cancellationToken).ConfigureAwait(false));
+                messageStore = chatMessageStoreFeature;
+            }
+
+            // Add any existing messages from the thread to the messages to be sent to the chat client.
+            if (messageStore is not null)
+            {
+                inputMessagesForChatClient.AddRange(await messageStore.GetMessagesAsync(cancellationToken).ConfigureAwait(false));
             }
 
             // If we have an AIContextProvider, we should get context from it, and update our
@@ -684,10 +735,31 @@ public sealed partial class ChatClientAgent : AIAgent
             // 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(new() { SerializedState = default, JsonSerializerOptions = null });
+            thread.MessageStore ??= this._agentOptions?.ChatMessageStoreFactory?.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }) ?? new InMemoryChatMessageStore();
         }
     }
 
+    private static Task NotifyMessageStoreOfNewMessagesAsync(ChatClientAgentThread thread, IEnumerable<ChatMessage> newMessages, AgentRunOptions? runOptions, CancellationToken cancellationToken)
+    {
+        var messageStore = thread.MessageStore;
+
+        // If the caller provided an override message store via run options, we should use that instead of the message store
+        // on the thread.
+        if (runOptions?.Features?.Get<ChatMessageStore>() is ChatMessageStore chatMessageStoreFeature)
+        {
+            messageStore = chatMessageStoreFeature;
+        }
+
+        // Only notify the message store if we have one.
+        // If we don't have one, it means that the chat history is service managed and the underlying service is responsible for storing messages.
+        if (messageStore is not null)
+        {
+            return messageStore.AddMessagesAsync(newMessages, cancellationToken);
+        }
+
+        return Task.CompletedTask;
+    }
+
     private string GetLoggingAgentName() => this.Name ?? "UnnamedAgent";
     #endregion
 }

dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs

@@ -128,6 +128,11 @@ public class ChatClientAgentOptions
         /// Gets or sets the JSON serialization options to use when deserializing the <see cref="SerializedState"/>.
         /// </summary>
         public JsonSerializerOptions? JsonSerializerOptions { get; set; }
+
+        /// <summary>
+        /// Gets or sets the collection of features provided by the caller and middleware.
+        /// </summary>
+        public IAgentFeatureCollection? Features { get; set; }
     }
 
     /// <summary>
@@ -145,5 +150,10 @@ public class ChatClientAgentOptions
         /// Gets or sets the JSON serialization options to use when deserializing the <see cref="SerializedState"/>.
         /// </summary>
         public JsonSerializerOptions? JsonSerializerOptions { get; set; }
+
+        /// <summary>
+        /// Gets or sets the collection of features provided by the caller and middleware.
+        /// </summary>
+        public IAgentFeatureCollection? Features { get; set; }
     }
 }

dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentThread.cs

@@ -1,12 +1,8 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
-using System.Collections.Generic;
 using System.Diagnostics;
 using System.Text.Json;
-using System.Threading;
-using System.Threading.Tasks;
-using Microsoft.Extensions.AI;
 using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Agents.AI;
@@ -182,33 +178,6 @@ public class ChatClientAgentThread : AgentThread
             ?? this.AIContextProvider?.GetService(serviceType, serviceKey)
             ?? this.MessageStore?.GetService(serviceType, serviceKey);
 
-    /// <inheritdoc />
-    protected override async Task MessagesReceivedAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken = default)
-    {
-        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();
-        }
-    }
-
     [DebuggerBrowsable(DebuggerBrowsableState.Never)]
     private string DebuggerDisplay =>
         this._conversationId is { } conversationId ? $"ConversationId = {conversationId}" :

dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs

@@ -73,6 +73,24 @@ public sealed class A2AAgentTests : IDisposable
         Assert.Equal(agent.Id, agent.DisplayName);
     }
 
+    [Fact]
+    public void GetNewThread_WithStringFeature_UsesItForContextId()
+    {
+        // Arrange
+        var contextIdFeature = new ConversationIdAgentFeature("feature-context-id");
+        var agentWithFeature = new A2AAgent(this._a2aClient);
+
+        // Act
+        var features = new AgentFeatureCollection();
+        features.Set(contextIdFeature);
+        var thread = agentWithFeature.GetNewThread(features);
+
+        // Assert
+        Assert.IsType<A2AAgentThread>(thread);
+        var a2aThread = (A2AAgentThread)thread;
+        Assert.Equal(contextIdFeature.ConversationId, a2aThread.ContextId);
+    }
+
     [Fact]
     public async Task RunAsync_AllowsNonUserRoleMessagesAsync()
     {

dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIAgentTests.cs

@@ -8,7 +8,6 @@ using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
 using Moq;
-using Moq.Protected;
 
 namespace Microsoft.Agents.AI.Abstractions.UnitTests;
 
@@ -222,21 +221,6 @@ public class AIAgentTests
         Assert.Equal(id, agent.Id);
     }
 
-    [Fact]
-    public async Task NotifyThreadOfNewMessagesNotifiesThreadAsync()
-    {
-        var cancellationToken = default(CancellationToken);
-
-        var messages = new[] { new ChatMessage(ChatRole.User, "msg1"), new ChatMessage(ChatRole.User, "msg2") };
-
-        var threadMock = new Mock<TestAgentThread> { CallBase = true };
-        threadMock.SetupAllProperties();
-
-        await MockAgent.NotifyThreadOfNewMessagesAsync(threadMock.Object, messages, cancellationToken);
-
-        threadMock.Protected().Verify("MessagesReceivedAsync", Times.Once(), messages, cancellationToken);
-    }
-
     #region GetService Method Tests
 
     /// <summary>
@@ -360,13 +344,10 @@ public class AIAgentTests
 
     private sealed class MockAgent : AIAgent
     {
-        public static new Task NotifyThreadOfNewMessagesAsync(AgentThread thread, IEnumerable<ChatMessage> messages, CancellationToken cancellationToken) =>
-            AIAgent.NotifyThreadOfNewMessagesAsync(thread, messages, cancellationToken);
-
-        public override AgentThread GetNewThread()
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
             => throw new NotImplementedException();
 
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
             => throw new NotImplementedException();
 
         public override Task<AgentRunResponse> RunAsync(

dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentFeatureCollectionTests.cs

@@ -0,0 +1,119 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+
+namespace Microsoft.Agents.AI.Abstractions.UnitTests;
+
+/// <summary>
+/// Contains unit tests for the <see cref="AgentFeatureCollection"/> class.
+/// </summary>
+public class AgentFeatureCollectionTests
+{
+    [Fact]
+    public void AddedInterfaceIsReturned()
+    {
+        var interfaces = new AgentFeatureCollection();
+        var thing = new Thing();
+
+        interfaces[typeof(IThing)] = thing;
+
+        var thing2 = interfaces[typeof(IThing)];
+        Assert.Equal(thing2, thing);
+    }
+
+    [Fact]
+    public void IndexerAlsoAddsItems()
+    {
+        var interfaces = new AgentFeatureCollection();
+        var thing = new Thing();
+
+        interfaces[typeof(IThing)] = thing;
+
+        Assert.Equal(interfaces[typeof(IThing)], thing);
+    }
+
+    [Fact]
+    public void SetNullValueRemoves()
+    {
+        var interfaces = new AgentFeatureCollection();
+        var thing = new Thing();
+
+        interfaces[typeof(IThing)] = thing;
+        Assert.Equal(interfaces[typeof(IThing)], thing);
+
+        interfaces[typeof(IThing)] = null;
+
+        var thing2 = interfaces[typeof(IThing)];
+        Assert.Null(thing2);
+    }
+
+    [Fact]
+    public void GetMissingStructFeatureThrows()
+    {
+        var interfaces = new AgentFeatureCollection();
+
+        var ex = Assert.Throws<InvalidOperationException>(() => interfaces.Get<int>());
+        Assert.Equal("System.Int32 does not exist in the feature collection and because it is a struct the method can't return null. Use 'AgentFeatureCollection[typeof(System.Int32)] is not null' to check if the feature exists.", ex.Message);
+    }
+
+    [Fact]
+    public void GetMissingFeatureReturnsNull()
+    {
+        var interfaces = new AgentFeatureCollection();
+
+        Assert.Null(interfaces.Get<Thing>());
+    }
+
+    [Fact]
+    public void GetStructFeature()
+    {
+        var interfaces = new AgentFeatureCollection();
+        const int Value = 20;
+        interfaces.Set(Value);
+
+        Assert.Equal(Value, interfaces.Get<int>());
+    }
+
+    [Fact]
+    public void GetNullableStructFeatureWhenSetWithNonNullableStruct()
+    {
+        var interfaces = new AgentFeatureCollection();
+        const int Value = 20;
+        interfaces.Set(Value);
+
+        Assert.Null(interfaces.Get<int?>());
+    }
+
+    [Fact]
+    public void GetNullableStructFeatureWhenSetWithNullableStruct()
+    {
+        var interfaces = new AgentFeatureCollection();
+        const int Value = 20;
+        interfaces.Set<int?>(Value);
+
+        Assert.Equal(Value, interfaces.Get<int?>());
+    }
+
+    [Fact]
+    public void GetFeature()
+    {
+        var interfaces = new AgentFeatureCollection();
+        var thing = new Thing();
+        interfaces.Set(thing);
+
+        Assert.Equal(thing, interfaces.Get<Thing>());
+    }
+
+    private interface IThing
+    {
+        string Hello();
+    }
+
+    private sealed class Thing : IThing
+    {
+        public string Hello()
+        {
+            return "World";
+        }
+    }
+}

dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRunOptionsTests.cs

@@ -23,7 +23,8 @@ public class AgentRunOptionsTests
             {
                 ["key1"] = "value1",
                 ["key2"] = 42
-            }
+            },
+            Features = new AgentFeatureCollection()
         };
 
         // Act
@@ -37,6 +38,7 @@ public class AgentRunOptionsTests
         Assert.NotSame(options.AdditionalProperties, clone.AdditionalProperties);
         Assert.Equal("value1", clone.AdditionalProperties["key1"]);
         Assert.Equal(42, clone.AdditionalProperties["key2"]);
+        Assert.Same(options.Features, clone.Features);
     }
 
     [Fact]

dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentThreadTests.cs

@@ -1,8 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
-using System.Collections.Generic;
-using Microsoft.Extensions.AI;
 
 #pragma warning disable CA1861 // Avoid constant arrays as arguments
 
@@ -21,15 +19,6 @@ public class AgentThreadTests
         Assert.Equal(default, result);
     }
 
-    [Fact]
-    public void MessagesReceivedAsync_ReturnsCompletedTask()
-    {
-        var thread = new TestAgentThread();
-        var messages = new List<ChatMessage> { new(ChatRole.User, "hello") };
-        var result = thread.MessagesReceivedAsync(messages);
-        Assert.True(result.IsCompleted);
-    }
-
     #region GetService Method Tests
 
     /// <summary>

dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/DelegatingAIAgentTests.cs

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -34,7 +35,12 @@ public class DelegatingAIAgentTests
         this._innerAgentMock.Setup(x => x.Id).Returns("test-agent-id");
         this._innerAgentMock.Setup(x => x.Name).Returns("Test Agent");
         this._innerAgentMock.Setup(x => x.Description).Returns("Test Description");
-        this._innerAgentMock.Setup(x => x.GetNewThread()).Returns(this._testThread);
+        this._innerAgentMock.Setup(x => x.GetNewThread(It.IsAny<IAgentFeatureCollection?>())).Returns(this._testThread);
+        this._innerAgentMock.Setup(x => x.DeserializeThread(
+            It.IsAny<JsonElement>(),
+            It.IsAny<JsonSerializerOptions?>(),
+            It.IsAny<IAgentFeatureCollection?>()))
+            .Returns(this._testThread);
 
         this._innerAgentMock
             .Setup(x => x.RunAsync(
@@ -135,11 +141,29 @@ public class DelegatingAIAgentTests
     public void GetNewThread_DelegatesToInnerAgent()
     {
         // Act
-        var thread = this._delegatingAgent.GetNewThread();
+        var featureCollection = new AgentFeatureCollection();
+        var thread = this._delegatingAgent.GetNewThread(featureCollection);
 
         // Assert
         Assert.Same(this._testThread, thread);
-        this._innerAgentMock.Verify(x => x.GetNewThread(), Times.Once);
+        this._innerAgentMock.Verify(x => x.GetNewThread(featureCollection), Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that DeserializeThread delegates to inner agent.
+    /// </summary>
+    [Fact]
+    public void DeserializeThread_DelegatesToInnerAgent()
+    {
+        // Act
+        var featureCollection = new AgentFeatureCollection();
+        var jsonElement = new JsonElement();
+        var jso = new JsonSerializerOptions();
+        var thread = this._delegatingAgent.DeserializeThread(jsonElement, jso, featureCollection);
+
+        // Assert
+        Assert.Same(this._testThread, thread);
+        this._innerAgentMock.Verify(x => x.DeserializeThread(jsonElement, jso, featureCollection), Times.Once);
     }
 
     /// <summary>

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs

@@ -289,12 +289,12 @@ internal sealed class FakeChatClientAgent : AIAgent
 
     public override string? Description => this._description;
 
-    public override AgentThread GetNewThread()
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
     {
         return new FakeInMemoryAgentThread();
     }
 
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
     {
         return new FakeInMemoryAgentThread(serializedThread, jsonSerializerOptions);
     }
@@ -366,12 +366,12 @@ internal sealed class FakeMultiMessageAgent : AIAgent
 
     public override string? Description => this._description;
 
-    public override AgentThread GetNewThread()
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
     {
         return new FakeInMemoryAgentThread();
     }
 
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
     {
         return new FakeInMemoryAgentThread(serializedThread, jsonSerializerOptions);
     }

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs

@@ -417,9 +417,9 @@ internal sealed class FakeStateAgent : AIAgent
         await Task.CompletedTask;
     }
 
-    public override AgentThread GetNewThread() => new FakeInMemoryAgentThread();
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null) => new FakeInMemoryAgentThread();
 
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
     {
         return new FakeInMemoryAgentThread(serializedThread, jsonSerializerOptions);
     }

dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs

@@ -425,9 +425,9 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests
 
         public override string? Description => "Agent that produces multiple text chunks";
 
-        public override AgentThread GetNewThread() => new TestInMemoryAgentThread();
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null) => new TestInMemoryAgentThread();
 
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null) =>
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null) =>
             new TestInMemoryAgentThread(serializedThread, jsonSerializerOptions);
 
         public override Task<AgentRunResponse> RunAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
@@ -514,9 +514,9 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests
 
         public override string? Description => "Test agent";
 
-        public override AgentThread GetNewThread() => new TestInMemoryAgentThread();
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null) => new TestInMemoryAgentThread();
 
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null) =>
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null) =>
             new TestInMemoryAgentThread(serializedThread, jsonSerializerOptions);
 
         public override Task<AgentRunResponse> RunAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)

dotnet/tests/Microsoft.Agents.AI.Hosting.AzureFunctions.UnitTests/TestAgent.cs

@@ -11,11 +11,12 @@ internal sealed class TestAgent(string name, string description) : AIAgent
 
     public override string? Description => description;
 
-    public override AgentThread GetNewThread() => new DummyAgentThread();
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null) => new DummyAgentThread();
 
     public override AgentThread DeserializeThread(
         JsonElement serializedThread,
-        JsonSerializerOptions? jsonSerializerOptions = null) => new DummyAgentThread();
+        JsonSerializerOptions? jsonSerializerOptions = null,
+        IAgentFeatureCollection? featureCollection = null) => new DummyAgentThread();
 
     public override Task<AgentRunResponse> RunAsync(
         IEnumerable<ChatMessage> messages,

dotnet/tests/Microsoft.Agents.AI.UnitTests/AgentExtensionsTests.cs

@@ -324,10 +324,10 @@ public class AgentExtensionsTests
             this._exceptionToThrow = exceptionToThrow;
         }
 
-        public override AgentThread GetNewThread()
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
             => throw new NotImplementedException();
 
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
             => throw new NotImplementedException();
 
         public override string? Name { get; }

dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs

@@ -426,10 +426,10 @@ public partial class ChatClientAgentTests
     }
 
     /// <summary>
-    /// Verify that RunAsync uses the ChatMessageStore factory when the chat client returns no conversation id.
+    /// Verify that RunAsync uses the default InMemoryChatMessageStore when the chat client returns no conversation id.
     /// </summary>
     [Fact]
-    public async Task RunAsyncUsesChatMessageStoreWhenNoConversationIdReturnedByChatClientAsync()
+    public async Task RunAsyncUsesDefaultInMemoryChatMessageStoreWhenNoConversationIdReturnedByChatClientAsync()
     {
         // Arrange
         Mock<IChatClient> mockService = new();
@@ -438,12 +438,9 @@ public partial class ChatClientAgentTests
                 It.IsAny<IEnumerable<ChatMessage>>(),
                 It.IsAny<ChatOptions>(),
                 It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
-        Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, ChatMessageStore>> mockFactory = new();
-        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>())).Returns(new InMemoryChatMessageStore());
         ChatClientAgent agent = new(mockService.Object, options: new()
         {
             Instructions = "test instructions",
-            ChatMessageStoreFactory = mockFactory.Object
         });
 
         // Act
@@ -455,14 +452,82 @@ public partial class ChatClientAgentTests
         Assert.Equal(2, messageStore.Count);
         Assert.Equal("test", messageStore[0].Text);
         Assert.Equal("response", messageStore[1].Text);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync uses the ChatMessageStore factory when the chat client returns no conversation id.
+    /// </summary>
+    [Fact]
+    public async Task RunAsyncUsesChatMessageStoreFactoryWhenProvidedAndNoConversationIdReturnedByChatClientAsync()
+    {
+        // Arrange
+        Mock<IChatClient> mockService = new();
+        mockService.Setup(
+            s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
+
+        Mock<ChatMessageStore> mockChatMessageStore = new();
+
+        Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, ChatMessageStore>> mockFactory = new();
+        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>())).Returns(mockChatMessageStore.Object);
+
+        ChatClientAgent agent = new(mockService.Object, options: new()
+        {
+            Instructions = "test instructions",
+            ChatMessageStoreFactory = mockFactory.Object
+        });
+
+        // Act
+        ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread;
+        await agent.RunAsync([new(ChatRole.User, "test")], thread);
+
+        // Assert
+        Assert.IsType<ChatMessageStore>(thread!.MessageStore, exactMatch: false);
+        mockChatMessageStore.Verify(s => s.AddMessagesAsync(It.Is<IEnumerable<ChatMessage>>(x => x.Count() == 2), It.IsAny<CancellationToken>()), Times.Once);
         mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Once);
     }
 
     /// <summary>
-    /// Verify that RunAsync doesn't use the ChatMessageStore factory when the chat client returns a conversation id.
+    /// Verify that RunAsync uses the ChatMessageStore provided via run params when the chat client returns no conversation id.
     /// </summary>
     [Fact]
-    public async Task RunAsyncIgnoresChatMessageStoreWhenConversationIdReturnedByChatClientAsync()
+    public async Task RunAsyncUsesChatMessageStoreWhenProvidedViaFeaturesAndNoConversationIdReturnedByChatClientAsync()
+    {
+        // Arrange
+        Mock<IChatClient> mockService = new();
+        mockService.Setup(
+            s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
+
+        Mock<ChatMessageStore> mockChatMessageStore = new();
+
+        ChatClientAgent agent = new(mockService.Object, options: new()
+        {
+            Instructions = "test instructions",
+        });
+
+        AgentFeatureCollection features = new();
+        features.Set(mockChatMessageStore.Object);
+
+        // Act
+        ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread;
+        await agent.RunAsync([new(ChatRole.User, "test")], thread, options: new AgentRunOptions() { Features = features });
+
+        // Assert
+        Assert.IsType<ChatMessageStore>(thread!.MessageStore, exactMatch: false);
+        mockChatMessageStore.Verify(s => s.GetMessagesAsync(It.IsAny<CancellationToken>()), Times.Once);
+        mockChatMessageStore.Verify(s => s.AddMessagesAsync(It.Is<IEnumerable<ChatMessage>>(x => x.Count() == 2), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync throws when a ChatMessageStore Factory is provided but when the chat client returns a conversation id.
+    /// </summary>
+    [Fact]
+    public async Task RunAsyncThrowsWhenChatMessageStoreFactoryProvidedAndConversationIdReturnedByChatClientAsync()
     {
         // Arrange
         Mock<IChatClient> mockService = new();
@@ -479,13 +544,10 @@ public partial class ChatClientAgentTests
             ChatMessageStoreFactory = mockFactory.Object
         });
 
-        // Act
+        // Act & Assert
         ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread;
-        await agent.RunAsync([new(ChatRole.User, "test")], thread);
-
-        // Assert
-        Assert.Equal("ConvId", thread!.ConversationId);
-        mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Never);
+        var exception = await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync([new(ChatRole.User, "test")], thread));
+        Assert.Equal("Only the ConversationId or MessageStore may be set, but not both and switching from one to another is not supported.", exception.Message);
     }
 
     /// <summary>
@@ -1914,10 +1976,10 @@ public partial class ChatClientAgentTests
     }
 
     /// <summary>
-    /// Verify that RunStreamingAsync doesn't use the ChatMessageStore factory when the chat client returns a conversation id.
+    /// Verify that RunStreamingAsync throws when a ChatMessageStore factory is provided and the chat client returns a conversation id.
     /// </summary>
     [Fact]
-    public async Task RunStreamingAsyncIgnoresChatMessageStoreWhenConversationIdReturnedByChatClientAsync()
+    public async Task RunStreamingAsyncThrowsWhenChatMessageStoreFactoryProvidedAndConversationIdReturnedByChatClientAsync()
     {
         // Arrange
         Mock<IChatClient> mockService = new();
@@ -1939,13 +2001,10 @@ public partial class ChatClientAgentTests
             ChatMessageStoreFactory = mockFactory.Object
         });
 
-        // Act
+        // Act & Assert
         ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread;
-        await agent.RunStreamingAsync([new(ChatRole.User, "test")], thread).ToListAsync();
-
-        // Assert
-        Assert.Equal("ConvId", thread!.ConversationId);
-        mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Never);
+        var exception = await Assert.ThrowsAsync<InvalidOperationException>(async () => await agent.RunStreamingAsync([new(ChatRole.User, "test")], thread).ToListAsync());
+        Assert.Equal("Only the ConversationId or MessageStore may be set, but not both and switching from one to another is not supported.", exception.Message);
     }
 
     /// <summary>
@@ -2074,37 +2133,6 @@ public partial class ChatClientAgentTests
 
     #endregion
 
-    #region GetNewThread Tests
-
-    [Fact]
-    public void GetNewThreadUsesAIContextProviderFactoryIfProvided()
-    {
-        // Arrange
-        var mockChatClient = new Mock<IChatClient>();
-        var mockContextProvider = new Mock<AIContextProvider>();
-        var factoryCalled = false;
-        var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
-        {
-            Instructions = "Test instructions",
-            AIContextProviderFactory = _ =>
-            {
-                factoryCalled = true;
-                return mockContextProvider.Object;
-            }
-        });
-
-        // Act
-        var thread = agent.GetNewThread();
-
-        // Assert
-        Assert.True(factoryCalled, "AIContextProviderFactory was not called.");
-        Assert.IsType<ChatClientAgentThread>(thread);
-        var typedThread = (ChatClientAgentThread)thread;
-        Assert.Same(mockContextProvider.Object, typedThread.AIContextProvider);
-    }
-
-    #endregion
-
     #region Background Responses Tests
 
     [Theory]

dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentThreadTests.cs

@@ -4,7 +4,6 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text.Json;
-using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
 using Moq;
@@ -91,50 +90,6 @@ public class ChatClientAgentThreadTests
 
     #endregion Constructor and Property Tests
 
-    #region OnNewMessagesAsync Tests
-
-    [Fact]
-    public async Task OnNewMessagesAsyncDoesNothingWhenAgentServiceIdAsync()
-    {
-        // Arrange
-        var thread = new ChatClientAgentThread { ConversationId = "thread-123" };
-        var messages = new List<ChatMessage>
-        {
-            new(ChatRole.User, "Hello"),
-            new(ChatRole.Assistant, "Hi there!")
-        };
-        var agent = new MessageSendingAgent();
-
-        // Act
-        await agent.SendMessagesAsync(thread, messages, CancellationToken.None);
-        Assert.Equal("thread-123", thread.ConversationId);
-        Assert.Null(thread.MessageStore);
-    }
-
-    [Fact]
-    public async Task OnNewMessagesAsyncAddsMessagesToStoreAsync()
-    {
-        // Arrange
-        var store = new InMemoryChatMessageStore();
-        var thread = new ChatClientAgentThread { MessageStore = store };
-        var messages = new List<ChatMessage>
-        {
-            new(ChatRole.User, "Hello"),
-            new(ChatRole.Assistant, "Hi there!")
-        };
-        var agent = new MessageSendingAgent();
-
-        // Act
-        await agent.SendMessagesAsync(thread, messages, CancellationToken.None);
-
-        // Assert
-        Assert.Equal(2, store.Count);
-        Assert.Equal("Hello", store[0].Text);
-        Assert.Equal("Hi there!", store[1].Text);
-    }
-
-    #endregion OnNewMessagesAsync Tests
-
     #region Deserialize Tests
 
     [Fact]
@@ -372,22 +327,4 @@ public class ChatClientAgentThreadTests
     }
 
     #endregion
-
-    private sealed class MessageSendingAgent : AIAgent
-    {
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
-            => throw new NotImplementedException();
-
-        public override AgentThread GetNewThread()
-            => throw new NotImplementedException();
-
-        public override Task<AgentRunResponse> RunAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
-            => throw new NotImplementedException();
-
-        public override IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
-            => throw new NotImplementedException();
-
-        public Task SendMessagesAsync(AgentThread thread, IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
-            => NotifyThreadOfNewMessagesAsync(thread, messages, cancellationToken);
-    }
 }

dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_DeserializeThreadTests.cs

@@ -0,0 +1,143 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Text.Json;
+using Microsoft.Extensions.AI;
+using Moq;
+
+namespace Microsoft.Agents.AI.UnitTests.ChatClient;
+
+/// <summary>
+/// Contains unit tests for the ChatClientAgent.DeserializeThread methods.
+/// </summary>
+public class ChatClientAgent_DeserializeThreadTests
+{
+    [Fact]
+    public void DeserializeThread_UsesAIContextProviderFactory_IfProvided()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockContextProvider = new Mock<AIContextProvider>();
+        var factoryCalled = false;
+        var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
+        {
+            Instructions = "Test instructions",
+            AIContextProviderFactory = _ =>
+            {
+                factoryCalled = true;
+                return mockContextProvider.Object;
+            }
+        });
+
+        var json = JsonSerializer.Deserialize("""
+            {
+                "aiContextProviderState": ["CP1"]
+            }
+            """, TestJsonSerializerContext.Default.JsonElement);
+
+        // Act
+        var thread = agent.DeserializeThread(json);
+
+        // Assert
+        Assert.True(factoryCalled, "AIContextProviderFactory was not called.");
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Same(mockContextProvider.Object, typedThread.AIContextProvider);
+    }
+
+    [Fact]
+    public void DeserializeThread_UsesChatMessageStoreFactory_IfProvided()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockMessageStore = new Mock<ChatMessageStore>();
+        var factoryCalled = false;
+        var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
+        {
+            Instructions = "Test instructions",
+            ChatMessageStoreFactory = _ =>
+            {
+                factoryCalled = true;
+                return mockMessageStore.Object;
+            }
+        });
+
+        var json = JsonSerializer.Deserialize("""
+            {
+                "storeState": { }
+            }
+            """, TestJsonSerializerContext.Default.JsonElement);
+
+        // Act
+        var thread = agent.DeserializeThread(json);
+
+        // Assert
+        Assert.True(factoryCalled, "ChatMessageStoreFactory was not called.");
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Same(mockMessageStore.Object, typedThread.MessageStore);
+    }
+
+    [Fact]
+    public void DeserializeThread_UsesChatMessageStore_FromFeatureOverload()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockMessageStore = new Mock<ChatMessageStore>();
+        var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
+        {
+            Instructions = "Test instructions",
+            ChatMessageStoreFactory = _ =>
+            {
+                Assert.Fail("ChatMessageStoreFactory should not have been called.");
+                return null!;
+            }
+        });
+
+        var json = JsonSerializer.Deserialize("""
+            {
+            }
+            """, TestJsonSerializerContext.Default.JsonElement);
+
+        // Act
+        var agentFeatures = new AgentFeatureCollection();
+        agentFeatures.Set(mockMessageStore.Object);
+        var thread = agent.DeserializeThread(json, null, agentFeatures);
+
+        // Assert
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Same(mockMessageStore.Object, typedThread.MessageStore);
+    }
+
+    [Fact]
+    public void DeserializeThread_UsesAIContextProvider_FromFeatureOverload()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockContextProvider = new Mock<AIContextProvider>();
+        var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
+        {
+            Instructions = "Test instructions",
+            AIContextProviderFactory = _ =>
+            {
+                Assert.Fail("AIContextProviderFactory should not have been called.");
+                return null!;
+            }
+        });
+
+        var json = JsonSerializer.Deserialize("""
+            {
+            }
+            """, TestJsonSerializerContext.Default.JsonElement);
+
+        // Act
+        var agentFeatures = new AgentFeatureCollection();
+        agentFeatures.Set(mockContextProvider.Object);
+        var thread = agent.DeserializeThread(json, null, agentFeatures);
+
+        // Assert
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Same(mockContextProvider.Object, typedThread.AIContextProvider);
+    }
+}

dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_GetNewThreadTests.cs

@@ -0,0 +1,176 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using Microsoft.Extensions.AI;
+using Moq;
+
+namespace Microsoft.Agents.AI.UnitTests.ChatClient;
+
+/// <summary>
+/// Contains unit tests for the ChatClientAgent.GetNewThread methods.
+/// </summary>
+public class ChatClientAgent_GetNewThreadTests
+{
+    [Fact]
+    public void GetNewThread_UsesAIContextProviderFactory_IfProvided()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockContextProvider = new Mock<AIContextProvider>();
+        var factoryCalled = false;
+        var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
+        {
+            Instructions = "Test instructions",
+            AIContextProviderFactory = _ =>
+            {
+                factoryCalled = true;
+                return mockContextProvider.Object;
+            }
+        });
+
+        // Act
+        var thread = agent.GetNewThread();
+
+        // Assert
+        Assert.True(factoryCalled, "AIContextProviderFactory was not called.");
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Same(mockContextProvider.Object, typedThread.AIContextProvider);
+    }
+
+    [Fact]
+    public void GetNewThread_UsesChatMessageStoreFactory_IfProvided()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockMessageStore = new Mock<ChatMessageStore>();
+        var factoryCalled = false;
+        var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
+        {
+            Instructions = "Test instructions",
+            ChatMessageStoreFactory = _ =>
+            {
+                factoryCalled = true;
+                return mockMessageStore.Object;
+            }
+        });
+
+        // Act
+        var thread = agent.GetNewThread();
+
+        // Assert
+        Assert.True(factoryCalled, "ChatMessageStoreFactory was not called.");
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Same(mockMessageStore.Object, typedThread.MessageStore);
+    }
+
+    [Fact]
+    public void GetNewThread_UsesChatMessageStore_FromTypedOverload()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockMessageStore = new Mock<ChatMessageStore>();
+        var agent = new ChatClientAgent(mockChatClient.Object);
+
+        // Act
+        var thread = agent.GetNewThread(mockMessageStore.Object);
+
+        // Assert
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Same(mockMessageStore.Object, typedThread.MessageStore);
+    }
+
+    [Fact]
+    public void GetNewThread_UsesConversationId_FromTypedOverload()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        const string TestConversationId = "test_conversation_id";
+        var agent = new ChatClientAgent(mockChatClient.Object);
+
+        // Act
+        var thread = agent.GetNewThread(TestConversationId);
+
+        // Assert
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Equal(TestConversationId, typedThread.ConversationId);
+    }
+
+    [Fact]
+    public void GetNewThread_UsesConversationId_FromFeatureOverload()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var testConversationId = new ConversationIdAgentFeature("test_conversation_id");
+        var agent = new ChatClientAgent(mockChatClient.Object);
+
+        // Act
+        var agentFeatures = new AgentFeatureCollection();
+        agentFeatures.Set(testConversationId);
+        var thread = agent.GetNewThread(agentFeatures);
+
+        // Assert
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Equal(testConversationId.ConversationId, typedThread.ConversationId);
+    }
+
+    [Fact]
+    public void GetNewThread_UsesChatMessageStore_FromFeatureOverload()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockMessageStore = new Mock<ChatMessageStore>();
+        var agent = new ChatClientAgent(mockChatClient.Object);
+
+        // Act
+        var agentFeatures = new AgentFeatureCollection();
+        agentFeatures.Set(mockMessageStore.Object);
+        var thread = agent.GetNewThread(agentFeatures);
+
+        // Assert
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Same(mockMessageStore.Object, typedThread.MessageStore);
+    }
+
+    [Fact]
+    public void GetNewThread_UsesAIContextProvider_FromFeatureOverload()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockContextProvider = new Mock<AIContextProvider>();
+        var agent = new ChatClientAgent(mockChatClient.Object);
+
+        // Act
+        var agentFeatures = new AgentFeatureCollection();
+        agentFeatures.Set(mockContextProvider.Object);
+        var thread = agent.GetNewThread(agentFeatures);
+
+        // Assert
+        Assert.IsType<ChatClientAgentThread>(thread);
+        var typedThread = (ChatClientAgentThread)thread;
+        Assert.Same(mockContextProvider.Object, typedThread.AIContextProvider);
+    }
+
+    [Fact]
+    public void GetNewThread_Throws_IfBothConversationIdAndMessageStoreAreSet()
+    {
+        // Arrange
+        var mockChatClient = new Mock<IChatClient>();
+        var mockMessageStore = new Mock<ChatMessageStore>();
+        var testConversationId = new ConversationIdAgentFeature("test_conversation_id");
+        var agent = new ChatClientAgent(mockChatClient.Object);
+
+        // Act & Assert
+        var agentFeatures = new AgentFeatureCollection();
+        agentFeatures.Set(mockMessageStore.Object);
+        agentFeatures.Set(testConversationId);
+
+        var exception = Assert.Throws<InvalidOperationException>(() => agent.GetNewThread(agentFeatures));
+        Assert.Equal("Only the ConversationId or MessageStore may be set, but not both and switching from one to another is not supported.", exception.Message);
+    }
+}

dotnet/tests/Microsoft.Agents.AI.UnitTests/TestAIAgent.cs

@@ -24,10 +24,10 @@ internal sealed class TestAIAgent : AIAgent
 
     public override string? Description => this.DescriptionFunc?.Invoke() ?? base.Description;
 
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null) =>
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null) =>
         this.DeserializeThreadFunc(serializedThread, jsonSerializerOptions);
 
-    public override AgentThread GetNewThread() =>
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null) =>
         this.GetNewThreadFunc();
 
     public override Task<AgentRunResponse> RunAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) =>

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/AgentWorkflowBuilderTests.cs

@@ -135,10 +135,10 @@ public class AgentWorkflowBuilderTests
     {
         public override string Name => name;
 
-        public override AgentThread GetNewThread()
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
             => new DoubleEchoAgentThread();
 
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
             => new DoubleEchoAgentThread();
 
         public override Task<AgentRunResponse> RunAsync(

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessExecutionTests.cs

@@ -146,10 +146,12 @@ public class InProcessExecutionTests
 
         public override string Name => this._name;
 
-        public override AgentThread GetNewThread() => new SimpleTestAgentThread();
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null) => new SimpleTestAgentThread();
 
-        public override AgentThread DeserializeThread(System.Text.Json.JsonElement serializedThread,
-            System.Text.Json.JsonSerializerOptions? jsonSerializerOptions = null) => new SimpleTestAgentThread();
+        public override AgentThread DeserializeThread(
+            System.Text.Json.JsonElement serializedThread,
+            System.Text.Json.JsonSerializerOptions? jsonSerializerOptions = null,
+            IAgentFeatureCollection? featureCollection = null) => new SimpleTestAgentThread();
 
         public override Task<AgentRunResponse> RunAsync(
             IEnumerable<ChatMessage> messages,

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RepresentationTests.cs

@@ -24,10 +24,10 @@ public class RepresentationTests
 
     private sealed class TestAgent : AIAgent
     {
-        public override AgentThread GetNewThread()
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
             => throw new NotImplementedException();
 
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
             => throw new NotImplementedException();
 
         public override Task<AgentRunResponse> RunAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) =>

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/06_GroupChat_Workflow.cs

@@ -60,10 +60,10 @@ internal sealed class HelloAgent(string id = nameof(HelloAgent)) : AIAgent
     public override string Id => id;
     public override string? Name => id;
 
-    public override AgentThread GetNewThread()
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
         => new HelloAgentThread();
 
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
         => new HelloAgentThread();
 
     public override async Task<AgentRunResponse> RunAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/SpecializedExecutorSmokeTests.cs

@@ -51,10 +51,10 @@ public class SpecializedExecutorSmokeTests
             return result;
         }
 
-        public override AgentThread GetNewThread()
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
             => new TestAgentThread();
 
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
             => new TestAgentThread();
 
         public static TestAIAgent FromStrings(params string[] messages) =>

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestEchoAgent.cs

@@ -16,12 +16,12 @@ internal class TestEchoAgent(string? id = null, string? name = null, string? pre
     public override string Id => id ?? base.Id;
     public override string? Name => name ?? base.Name;
 
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
     {
         return JsonSerializer.Deserialize<EchoAgentThread>(serializedThread, jsonSerializerOptions) ?? this.GetNewThread();
     }
 
-    public override AgentThread GetNewThread()
+    public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
     {
         return new EchoAgentThread();
     }