.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/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();
     }