.NET: Allow overriding the ChatMessageStore to be used per agent run. (#3330)

* Allow overriding the ChatMessageStore to be used per agent run.

* Fix typos

* Fix Add and add TryAdd, Contains and Remove

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

@@ -0,0 +1,99 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Diagnostics.CodeAnalysis;
+using Microsoft.Extensions.AI;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Contains extension methods to allow storing and retrieving properties using the type name of the property as the key.
+/// </summary>
+public static class AdditionalPropertiesExtensions
+{
+    /// <summary>
+    /// Adds an additional property using the type name of the property as the key.
+    /// </summary>
+    /// <typeparam name="T">The type of the property to add.</typeparam>
+    /// <param name="additionalProperties">The dictionary of additional properties.</param>
+    /// <param name="value">The value to add.</param>
+    public static void Add<T>(this AdditionalPropertiesDictionary additionalProperties, T value)
+    {
+        _ = Throw.IfNull(additionalProperties);
+
+        additionalProperties.Add(typeof(T).FullName!, value);
+    }
+
+    /// <summary>
+    /// Attempts to add a property using the type name of the property as the key.
+    /// </summary>
+    /// <remarks>
+    /// This method uses the full name of the type parameter as the key. If the key already exists,
+    /// the value is not updated and the method returns <see langword="false"/>.
+    /// </remarks>
+    /// <typeparam name="T">The type of the property to add.</typeparam>
+    /// <param name="additionalProperties">The dictionary of additional properties.</param>
+    /// <param name="value">The value to add.</param>
+    /// <returns>
+    /// <see langword="true"/> if the value was added successfully; <see langword="false"/> if the key already exists.
+    /// </returns>
+    public static bool TryAdd<T>(this AdditionalPropertiesDictionary additionalProperties, T value)
+    {
+        _ = Throw.IfNull(additionalProperties);
+
+        return additionalProperties.TryAdd(typeof(T).FullName!, value);
+    }
+
+    /// <summary>
+    /// Attempts to retrieve a value from the additional properties dictionary using the type name of the property as the key.
+    /// </summary>
+    /// <remarks>
+    /// This method uses the full name of the type parameter as the key when searching the dictionary.
+    /// </remarks>
+    /// <typeparam name="T">The type of the property to be retrieved.</typeparam>
+    /// <param name="additionalProperties">The dictionary containing additional properties.</param>
+    /// <param name="value">
+    /// When this method returns, contains the value retrieved from the dictionary, if found and successfully converted to the requested type;
+    /// otherwise, the default value of <typeparamref name="T"/>.
+    /// </param>
+    /// <returns>
+    /// <see langword="true"/> if a non-<see langword="null"/> value was found
+    /// in the dictionary and converted to the requested type; otherwise, <see langword="false"/>.
+    /// </returns>
+    public static bool TryGetValue<T>(this AdditionalPropertiesDictionary additionalProperties, [NotNullWhen(true)] out T? value)
+    {
+        _ = Throw.IfNull(additionalProperties);
+
+        return additionalProperties.TryGetValue(typeof(T).FullName!, out value);
+    }
+
+    /// <summary>
+    /// Determines whether the additional properties dictionary contains a property with the name of the provided type as the key.
+    /// </summary>
+    /// <typeparam name="T">The type of the property to check for.</typeparam>
+    /// <param name="additionalProperties">The dictionary of additional properties.</param>
+    /// <returns>
+    /// <see langword="true"/> if the dictionary contains a property with the name of the provided type as the key; otherwise, <see langword="false"/>.
+    /// </returns>
+    public static bool Contains<T>(this AdditionalPropertiesDictionary additionalProperties)
+    {
+        _ = Throw.IfNull(additionalProperties);
+
+        return additionalProperties.ContainsKey(typeof(T).FullName!);
+    }
+
+    /// <summary>
+    /// Removes a property from the additional properties dictionary using the name of the provided type as the key.
+    /// </summary>
+    /// <typeparam name="T">The type of the property to remove.</typeparam>
+    /// <param name="additionalProperties">The dictionary of additional properties.</param>
+    /// <returns>
+    /// <see langword="true"/> if the property was successfully removed; otherwise, <see langword="false"/>.
+    /// </returns>
+    public static bool Remove<T>(this AdditionalPropertiesDictionary additionalProperties)
+    {
+        _ = Throw.IfNull(additionalProperties);
+
+        return additionalProperties.Remove(typeof(T).FullName!);
+    }
+}

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

@@ -171,10 +171,10 @@ public abstract class ChatMessageStore
         /// <param name="requestMessages">The caller provided messages that were used by the agent for this invocation.</param>
         /// <param name="chatMessageStoreMessages">The messages retrieved from the <see cref="ChatMessageStore"/> for this invocation.</param>
         /// <exception cref="ArgumentNullException"><paramref name="requestMessages"/> is <see langword="null"/>.</exception>
-        public InvokedContext(IEnumerable<ChatMessage> requestMessages, IEnumerable<ChatMessage> chatMessageStoreMessages)
+        public InvokedContext(IEnumerable<ChatMessage> requestMessages, IEnumerable<ChatMessage>? chatMessageStoreMessages)
         {
             this.RequestMessages = Throw.IfNull(requestMessages);
-            this.ChatMessageStoreMessages = Throw.IfNull(chatMessageStoreMessages);
+            this.ChatMessageStoreMessages = chatMessageStoreMessages;
         }
 
         /// <summary>
@@ -191,9 +191,9 @@ public abstract class ChatMessageStore
         /// </summary>
         /// <value>
         /// A collection of <see cref="ChatMessage"/> instances that were retrieved from the <see cref="ChatMessageStore"/>,
-        /// and were used by the agent as part of the invocation.
+        /// and were used by the agent as part of the invocation. May be null on the first run.
         /// </value>
-        public IEnumerable<ChatMessage> ChatMessageStoreMessages { get; set { field = Throw.IfNull(value); } }
+        public IEnumerable<ChatMessage>? ChatMessageStoreMessages { get; set; }
 
         /// <summary>
         /// Gets or sets the messages provided by the <see cref="AIContextProvider"/> for this invocation, if any.

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

@@ -231,7 +231,7 @@ public sealed partial class ChatClientAgent : AIAgent
         }
         catch (Exception ex)
         {
-            await NotifyMessageStoreOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyMessageStoreOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, chatOptions, cancellationToken).ConfigureAwait(false);
             await NotifyAIContextProviderOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             throw;
         }
@@ -246,7 +246,7 @@ public sealed partial class ChatClientAgent : AIAgent
         }
         catch (Exception ex)
         {
-            await NotifyMessageStoreOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyMessageStoreOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, chatOptions, cancellationToken).ConfigureAwait(false);
             await NotifyAIContextProviderOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             throw;
         }
@@ -273,7 +273,7 @@ public sealed partial class ChatClientAgent : AIAgent
             }
             catch (Exception ex)
             {
-                await NotifyMessageStoreOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+                await NotifyMessageStoreOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, chatOptions, cancellationToken).ConfigureAwait(false);
                 await NotifyAIContextProviderOfFailureAsync(safeThread, ex, GetInputMessages(inputMessages, continuationToken), aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
                 throw;
             }
@@ -286,7 +286,7 @@ public sealed partial class ChatClientAgent : AIAgent
         await this.UpdateThreadWithTypeAndConversationIdAsync(safeThread, chatResponse.ConversationId, cancellationToken).ConfigureAwait(false);
 
         // To avoid inconsistent state we only notify the thread of the input messages if no error occurs after the initial request.
-        await NotifyMessageStoreOfNewMessagesAsync(safeThread, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
+        await NotifyMessageStoreOfNewMessagesAsync(safeThread, GetInputMessages(inputMessages, continuationToken), chatMessageStoreMessages, aiContextProviderMessages, chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false);
 
         // Notify the AIContextProvider of all new messages.
         await NotifyAIContextProviderOfSuccessAsync(safeThread, GetInputMessages(inputMessages, continuationToken), aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
@@ -442,7 +442,7 @@ public sealed partial class ChatClientAgent : AIAgent
         }
         catch (Exception ex)
         {
-            await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, chatOptions, cancellationToken).ConfigureAwait(false);
             await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             throw;
         }
@@ -460,7 +460,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 NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
+        await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false);
 
         // Notify the AIContextProvider of all new messages.
         await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
@@ -703,16 +703,18 @@ public sealed partial class ChatClientAgent : AIAgent
 
         List<ChatMessage> inputMessagesForChatClient = [];
         IList<ChatMessage>? aiContextProviderMessages = null;
-        IList<ChatMessage>? chatMessageStoreMessages = [];
+        IList<ChatMessage>? chatMessageStoreMessages = null;
 
         // 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)
+            ChatMessageStore? chatMessageStore = ResolveChatMessageStore(typedThread, chatOptions);
+
+            // Add any existing messages from the chatMessageStore to the messages to be sent to the chat client.
+            if (chatMessageStore is not null)
             {
                 var invokingContext = new ChatMessageStore.InvokingContext(inputMessages);
-                var storeMessages = await typedThread.MessageStore.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
+                var storeMessages = await chatMessageStore.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
                 inputMessagesForChatClient.AddRange(storeMessages);
                 chatMessageStoreMessages = storeMessages as IList<ChatMessage> ?? storeMessages.ToList();
             }
@@ -803,21 +805,22 @@ public sealed partial class ChatClientAgent : AIAgent
         IEnumerable<ChatMessage> requestMessages,
         IEnumerable<ChatMessage>? chatMessageStoreMessages,
         IEnumerable<ChatMessage>? aiContextProviderMessages,
+        ChatOptions? chatOptions,
         CancellationToken cancellationToken)
     {
-        var messageStore = thread.MessageStore;
+        ChatMessageStore? chatMessageStore = ResolveChatMessageStore(thread, chatOptions);
 
         // 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)
+        if (chatMessageStore is not null)
         {
-            var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages!)
+            var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages)
             {
                 AIContextProviderMessages = aiContextProviderMessages,
                 InvokeException = ex
             };
 
-            return messageStore.InvokedAsync(invokedContext, cancellationToken).AsTask();
+            return chatMessageStore.InvokedAsync(invokedContext, cancellationToken).AsTask();
         }
 
         return Task.CompletedTask;
@@ -829,25 +832,39 @@ public sealed partial class ChatClientAgent : AIAgent
         IEnumerable<ChatMessage>? chatMessageStoreMessages,
         IEnumerable<ChatMessage>? aiContextProviderMessages,
         IEnumerable<ChatMessage> responseMessages,
+        ChatOptions? chatOptions,
         CancellationToken cancellationToken)
     {
-        var messageStore = thread.MessageStore;
+        ChatMessageStore? chatMessageStore = ResolveChatMessageStore(thread, chatOptions);
 
         // 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)
+        if (chatMessageStore is not null)
         {
-            var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages!)
+            var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages)
             {
                 AIContextProviderMessages = aiContextProviderMessages,
                 ResponseMessages = responseMessages
             };
-            return messageStore.InvokedAsync(invokedContext, cancellationToken).AsTask();
+            return chatMessageStore.InvokedAsync(invokedContext, cancellationToken).AsTask();
         }
 
         return Task.CompletedTask;
     }
 
+    private static ChatMessageStore? ResolveChatMessageStore(ChatClientAgentThread thread, ChatOptions? chatOptions)
+    {
+        ChatMessageStore? chatMessageStore = thread.MessageStore;
+
+        // If someone provided an override ChatMessageStore via AdditionalProperties, we should use that instead of the one on the thread.
+        if (chatOptions?.AdditionalProperties?.TryGetValue(out ChatMessageStore? overrideChatMessageStore) is true)
+        {
+            chatMessageStore = overrideChatMessageStore;
+        }
+
+        return chatMessageStore;
+    }
+
     private static ChatClientAgentContinuationToken? WrapContinuationToken(ResponseContinuationToken? continuationToken, IEnumerable<ChatMessage>? inputMessages = null, List<ChatResponseUpdate>? responseUpdates = null)
     {
         if (continuationToken is null)

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

@@ -0,0 +1,490 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using Microsoft.Extensions.AI;
+
+namespace Microsoft.Agents.AI.Abstractions.UnitTests;
+
+/// <summary>
+/// Contains tests for the <see cref="AdditionalPropertiesExtensions"/> class.
+/// </summary>
+public sealed class AdditionalPropertiesExtensionsTests
+{
+    #region Add Method Tests
+
+    [Fact]
+    public void Add_WithValidValue_StoresValueUsingTypeName()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass value = new() { Name = "Test" };
+
+        // Act
+        additionalProperties.Add(value);
+
+        // Assert
+        Assert.True(additionalProperties.ContainsKey(typeof(TestClass).FullName!));
+        Assert.Same(value, additionalProperties[typeof(TestClass).FullName!]);
+    }
+
+    [Fact]
+    public void Add_WithNullDictionary_ThrowsArgumentNullException()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary? additionalProperties = null;
+        TestClass value = new() { Name = "Test" };
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => additionalProperties!.Add(value));
+    }
+
+    [Fact]
+    public void Add_WithStringValue_StoresValueCorrectly()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        const string Value = "test string";
+
+        // Act
+        additionalProperties.Add(Value);
+
+        // Assert
+        Assert.True(additionalProperties.ContainsKey(typeof(string).FullName!));
+        Assert.Equal(Value, additionalProperties[typeof(string).FullName!]);
+    }
+
+    [Fact]
+    public void Add_WithIntValue_StoresValueCorrectly()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        const int Value = 42;
+
+        // Act
+        additionalProperties.Add(Value);
+
+        // Assert
+        Assert.True(additionalProperties.ContainsKey(typeof(int).FullName!));
+        Assert.Equal(Value, additionalProperties[typeof(int).FullName!]);
+    }
+
+    [Fact]
+    public void Add_ThrowsArgumentException_WhenSameTypeAddedTwice()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass firstValue = new() { Name = "First" };
+        TestClass secondValue = new() { Name = "Second" };
+        additionalProperties.Add(firstValue);
+
+        // Act & Assert
+        Assert.Throws<ArgumentException>(() => additionalProperties.Add(secondValue));
+    }
+
+    [Fact]
+    public void Add_WithMultipleDifferentTypes_StoresAllValues()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass testClassValue = new() { Name = "Test" };
+        AnotherTestClass anotherValue = new() { Id = 123 };
+        const string StringValue = "test";
+
+        // Act
+        additionalProperties.Add(testClassValue);
+        additionalProperties.Add(anotherValue);
+        additionalProperties.Add(StringValue);
+
+        // Assert
+        Assert.Equal(3, additionalProperties.Count);
+        Assert.Same(testClassValue, additionalProperties[typeof(TestClass).FullName!]);
+        Assert.Same(anotherValue, additionalProperties[typeof(AnotherTestClass).FullName!]);
+        Assert.Equal(StringValue, additionalProperties[typeof(string).FullName!]);
+    }
+
+    #endregion
+
+    #region TryAdd Method Tests
+
+    [Fact]
+    public void TryAdd_WithValidValue_ReturnsTrueAndStoresValue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass value = new() { Name = "Test" };
+
+        // Act
+        bool result = additionalProperties.TryAdd(value);
+
+        // Assert
+        Assert.True(result);
+        Assert.True(additionalProperties.ContainsKey(typeof(TestClass).FullName!));
+        Assert.Same(value, additionalProperties[typeof(TestClass).FullName!]);
+    }
+
+    [Fact]
+    public void TryAdd_WithNullDictionary_ThrowsArgumentNullException()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary? additionalProperties = null;
+        TestClass value = new() { Name = "Test" };
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => additionalProperties!.TryAdd(value));
+    }
+
+    [Fact]
+    public void TryAdd_WithExistingType_ReturnsFalseAndKeepsOriginalValue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass firstValue = new() { Name = "First" };
+        TestClass secondValue = new() { Name = "Second" };
+        additionalProperties.Add(firstValue);
+
+        // Act
+        bool result = additionalProperties.TryAdd(secondValue);
+
+        // Assert
+        Assert.False(result);
+        Assert.Single(additionalProperties);
+        Assert.Same(firstValue, additionalProperties[typeof(TestClass).FullName!]);
+    }
+
+    [Fact]
+    public void TryAdd_WithStringValue_ReturnsTrueAndStoresValue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        const string Value = "test string";
+
+        // Act
+        bool result = additionalProperties.TryAdd(Value);
+
+        // Assert
+        Assert.True(result);
+        Assert.True(additionalProperties.ContainsKey(typeof(string).FullName!));
+        Assert.Equal(Value, additionalProperties[typeof(string).FullName!]);
+    }
+
+    [Fact]
+    public void TryAdd_WithIntValue_ReturnsTrueAndStoresValue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        const int Value = 42;
+
+        // Act
+        bool result = additionalProperties.TryAdd(Value);
+
+        // Assert
+        Assert.True(result);
+        Assert.True(additionalProperties.ContainsKey(typeof(int).FullName!));
+        Assert.Equal(Value, additionalProperties[typeof(int).FullName!]);
+    }
+
+    [Fact]
+    public void TryAdd_WithMultipleDifferentTypes_StoresAllValues()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass testClassValue = new() { Name = "Test" };
+        AnotherTestClass anotherValue = new() { Id = 123 };
+        const string StringValue = "test";
+
+        // Act
+        bool result1 = additionalProperties.TryAdd(testClassValue);
+        bool result2 = additionalProperties.TryAdd(anotherValue);
+        bool result3 = additionalProperties.TryAdd(StringValue);
+
+        // Assert
+        Assert.True(result1);
+        Assert.True(result2);
+        Assert.True(result3);
+        Assert.Equal(3, additionalProperties.Count);
+        Assert.Same(testClassValue, additionalProperties[typeof(TestClass).FullName!]);
+        Assert.Same(anotherValue, additionalProperties[typeof(AnotherTestClass).FullName!]);
+        Assert.Equal(StringValue, additionalProperties[typeof(string).FullName!]);
+    }
+
+    #endregion
+
+    #region TryGetValue Method Tests
+
+    [Fact]
+    public void TryGetValue_WithExistingValue_ReturnsTrueAndValue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass expectedValue = new() { Name = "Test" };
+        additionalProperties.Add(expectedValue);
+
+        // Act
+        bool result = additionalProperties.TryGetValue(out TestClass? actualValue);
+
+        // Assert
+        Assert.True(result);
+        Assert.NotNull(actualValue);
+        Assert.Same(expectedValue, actualValue);
+    }
+
+    [Fact]
+    public void TryGetValue_WithNonExistingValue_ReturnsFalseAndNull()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+
+        // Act
+        bool result = additionalProperties.TryGetValue(out TestClass? actualValue);
+
+        // Assert
+        Assert.False(result);
+        Assert.Null(actualValue);
+    }
+
+    [Fact]
+    public void TryGetValue_WithNullDictionary_ThrowsArgumentNullException()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary? additionalProperties = null;
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => additionalProperties!.TryGetValue<TestClass>(out _));
+    }
+
+    [Fact]
+    public void TryGetValue_WithStringValue_ReturnsCorrectValue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        const string ExpectedValue = "test string";
+        additionalProperties.Add(ExpectedValue);
+
+        // Act
+        bool result = additionalProperties.TryGetValue(out string? actualValue);
+
+        // Assert
+        Assert.True(result);
+        Assert.Equal(ExpectedValue, actualValue);
+    }
+
+    [Fact]
+    public void TryGetValue_WithIntValue_ReturnsCorrectValue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        const int ExpectedValue = 42;
+        additionalProperties.Add(ExpectedValue);
+
+        // Act
+        bool result = additionalProperties.TryGetValue(out int actualValue);
+
+        // Assert
+        Assert.True(result);
+        Assert.Equal(ExpectedValue, actualValue);
+    }
+
+    [Fact]
+    public void TryGetValue_WithWrongType_ReturnsFalse()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass testValue = new() { Name = "Test" };
+        additionalProperties.Add(testValue);
+
+        // Act
+        bool result = additionalProperties.TryGetValue(out AnotherTestClass? actualValue);
+
+        // Assert
+        Assert.False(result);
+        Assert.Null(actualValue);
+    }
+
+    [Fact]
+    public void TryGetValue_AfterTryAddFails_ReturnsOriginalValue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass firstValue = new() { Name = "First" };
+        TestClass secondValue = new() { Name = "Second" };
+        additionalProperties.Add(firstValue);
+        additionalProperties.TryAdd(secondValue);
+
+        // Act
+        bool result = additionalProperties.TryGetValue(out TestClass? actualValue);
+
+        // Assert
+        Assert.Single(additionalProperties);
+        Assert.True(result);
+        Assert.Same(firstValue, actualValue);
+    }
+
+    #endregion
+
+    #region Contains Method Tests
+
+    [Fact]
+    public void Contains_WithExistingType_ReturnsTrue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass value = new() { Name = "Test" };
+        additionalProperties.Add(value);
+
+        // Act
+        bool result = additionalProperties.Contains<TestClass>();
+
+        // Assert
+        Assert.True(result);
+    }
+
+    [Fact]
+    public void Contains_WithNonExistingType_ReturnsFalse()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+
+        // Act
+        bool result = additionalProperties.Contains<TestClass>();
+
+        // Assert
+        Assert.False(result);
+    }
+
+    [Fact]
+    public void Contains_WithNullDictionary_ThrowsArgumentNullException()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary? additionalProperties = null;
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => additionalProperties!.Contains<TestClass>());
+    }
+
+    [Fact]
+    public void Contains_WithDifferentType_ReturnsFalse()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass value = new() { Name = "Test" };
+        additionalProperties.Add(value);
+
+        // Act
+        bool result = additionalProperties.Contains<AnotherTestClass>();
+
+        // Assert
+        Assert.False(result);
+    }
+
+    [Fact]
+    public void Contains_AfterRemove_ReturnsFalse()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass value = new() { Name = "Test" };
+        additionalProperties.Add(value);
+        additionalProperties.Remove<TestClass>();
+
+        // Act
+        bool result = additionalProperties.Contains<TestClass>();
+
+        // Assert
+        Assert.False(result);
+    }
+
+    #endregion
+
+    #region Remove Method Tests
+
+    [Fact]
+    public void Remove_WithExistingType_ReturnsTrueAndRemovesValue()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass value = new() { Name = "Test" };
+        additionalProperties.Add(value);
+
+        // Act
+        bool result = additionalProperties.Remove<TestClass>();
+
+        // Assert
+        Assert.True(result);
+        Assert.Empty(additionalProperties);
+    }
+
+    [Fact]
+    public void Remove_WithNonExistingType_ReturnsFalse()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+
+        // Act
+        bool result = additionalProperties.Remove<TestClass>();
+
+        // Assert
+        Assert.False(result);
+    }
+
+    [Fact]
+    public void Remove_WithNullDictionary_ThrowsArgumentNullException()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary? additionalProperties = null;
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => additionalProperties!.Remove<TestClass>());
+    }
+
+    [Fact]
+    public void Remove_OnlyRemovesSpecifiedType()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass testValue = new() { Name = "Test" };
+        AnotherTestClass anotherValue = new() { Id = 123 };
+        additionalProperties.Add(testValue);
+        additionalProperties.Add(anotherValue);
+
+        // Act
+        bool result = additionalProperties.Remove<TestClass>();
+
+        // Assert
+        Assert.True(result);
+        Assert.Single(additionalProperties);
+        Assert.False(additionalProperties.Contains<TestClass>());
+        Assert.True(additionalProperties.Contains<AnotherTestClass>());
+    }
+
+    [Fact]
+    public void Remove_CalledTwice_ReturnsFalseOnSecondCall()
+    {
+        // Arrange
+        AdditionalPropertiesDictionary additionalProperties = new();
+        TestClass value = new() { Name = "Test" };
+        additionalProperties.Add(value);
+
+        // Act
+        bool firstResult = additionalProperties.Remove<TestClass>();
+        bool secondResult = additionalProperties.Remove<TestClass>();
+
+        // Assert
+        Assert.True(firstResult);
+        Assert.False(secondResult);
+    }
+
+    #endregion
+
+    #region Test Helper Classes
+
+    private sealed class TestClass
+    {
+        public string Name { get; set; } = string.Empty;
+    }
+
+    private sealed class AnotherTestClass
+    {
+        public int Id { get; set; }
+    }
+
+    #endregion
+}

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

@@ -5,7 +5,7 @@ using System.Linq;
 using System.Text.Json;
 using Microsoft.Extensions.AI;
 
-namespace Microsoft.Agents.AI.UnitTests.ChatClient;
+namespace Microsoft.Agents.AI.UnitTests;
 
 public class ChatClientAgentContinuationTokenTests
 {

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

@@ -310,302 +310,6 @@ public partial class ChatClientAgentTests
         Assert.Empty(capturedMessages);
     }
 
-    /// <summary>
-    /// Verify that RunAsync does not throw when providing a thread with a ThreadId and a Conversationid
-    /// via ChatOptions and the two are the same.
-    /// </summary>
-    [Fact]
-    public async Task RunAsyncDoesNotThrowWhenSpecifyingTwoSameThreadIdsAsync()
-    {
-        // Arrange
-        var chatOptions = new ChatOptions { ConversationId = "ConvId" };
-        Mock<IChatClient> mockService = new();
-        mockService.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.Is<ChatOptions>(opts => opts.ConversationId == "ConvId"),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]) { ConversationId = "ConvId" });
-
-        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
-
-        ChatClientAgentThread thread = new() { ConversationId = "ConvId" };
-
-        // Act & Assert
-        var response = await agent.RunAsync([new(ChatRole.User, "test")], thread, options: new ChatClientAgentRunOptions(chatOptions));
-        Assert.NotNull(response);
-    }
-
-    /// <summary>
-    /// Verify that RunAsync throws when providing a thread with a ThreadId and a Conversationid
-    /// via ChatOptions and the two are different.
-    /// </summary>
-    [Fact]
-    public async Task RunAsyncThrowsWhenSpecifyingTwoDifferentThreadIdsAsync()
-    {
-        // Arrange
-        var chatOptions = new ChatOptions { ConversationId = "ConvId" };
-        Mock<IChatClient> mockService = new();
-
-        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
-
-        ChatClientAgentThread thread = new() { ConversationId = "ThreadId" };
-
-        // Act & Assert
-        await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync([new(ChatRole.User, "test")], thread, options: new ChatClientAgentRunOptions(chatOptions)));
-    }
-
-    /// <summary>
-    /// Verify that RunAsync clones the ChatOptions when providing a thread with a ThreadId and a ChatOptions.
-    /// </summary>
-    [Fact]
-    public async Task RunAsyncClonesChatOptionsToAddThreadIdAsync()
-    {
-        // Arrange
-        var chatOptions = new ChatOptions { MaxOutputTokens = 100 };
-        Mock<IChatClient> mockService = new();
-        mockService.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.Is<ChatOptions>(opts => opts.MaxOutputTokens == 100 && opts.ConversationId == "ConvId"),
-                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]) { ConversationId = "ConvId" });
-
-        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
-
-        ChatClientAgentThread thread = new() { ConversationId = "ConvId" };
-
-        // Act
-        await agent.RunAsync([new(ChatRole.User, "test")], thread, options: new ChatClientAgentRunOptions(chatOptions));
-
-        // Assert
-        Assert.Null(chatOptions.ConversationId);
-    }
-
-    /// <summary>
-    /// Verify that RunAsync throws if a thread is provided that uses a conversation id already, but the service does not return one on invoke.
-    /// </summary>
-    [Fact]
-    public async Task RunAsyncThrowsForMissingConversationIdWithConversationIdThreadAsync()
-    {
-        // 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")]));
-
-        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
-
-        ChatClientAgentThread thread = new() { ConversationId = "ConvId" };
-
-        // Act & Assert
-        await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync([new(ChatRole.User, "test")], thread));
-    }
-
-    /// <summary>
-    /// Verify that RunAsync sets the ConversationId on the thread when the service returns one.
-    /// </summary>
-    [Fact]
-    public async Task RunAsyncSetsConversationIdOnThreadWhenReturnedByChatClientAsync()
-    {
-        // 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")]) { ConversationId = "ConvId" });
-        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
-        ChatClientAgentThread thread = new();
-
-        // Act
-        await agent.RunAsync([new(ChatRole.User, "test")], thread);
-
-        // Assert
-        Assert.Equal("ConvId", thread.ConversationId);
-    }
-
-    /// <summary>
-    /// Verify that RunAsync uses the ChatMessageStore factory when the chat client returns no conversation id.
-    /// </summary>
-    [Fact]
-    public async Task RunAsyncUsesChatMessageStoreWhenNoConversationIdReturnedByChatClientAsync()
-    {
-        // 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<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, CancellationToken, ValueTask<ChatMessageStore>>> mockFactory = new();
-        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>())).ReturnsAsync(new InMemoryChatMessageStore());
-        ChatClientAgent agent = new(mockService.Object, options: new()
-        {
-            ChatOptions = new() { Instructions = "test instructions" },
-            ChatMessageStoreFactory = mockFactory.Object
-        });
-
-        // Act
-        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
-        await agent.RunAsync([new(ChatRole.User, "test")], thread);
-
-        // Assert
-        var messageStore = Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
-        Assert.Equal(2, messageStore.Count);
-        Assert.Equal("test", messageStore[0].Text);
-        Assert.Equal("response", messageStore[1].Text);
-        mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>()), Times.Once);
-    }
-
-    /// <summary>
-    /// Verify that RunAsync uses the default InMemoryChatMessageStore when the chat client returns no conversation id.
-    /// </summary>
-    [Fact]
-    public async Task RunAsyncUsesDefaultInMemoryChatMessageStoreWhenNoConversationIdReturnedByChatClientAsync()
-    {
-        // 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")]));
-        ChatClientAgent agent = new(mockService.Object, options: new()
-        {
-            ChatOptions = new() { Instructions = "test instructions" },
-        });
-
-        // Act
-        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
-        await agent.RunAsync([new(ChatRole.User, "test")], thread);
-
-        // Assert
-        var messageStore = Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
-        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();
-        mockChatMessageStore.Setup(s => s.InvokingAsync(
-            It.IsAny<ChatMessageStore.InvokingContext>(),
-            It.IsAny<CancellationToken>())).ReturnsAsync([new ChatMessage(ChatRole.User, "Existing Chat History")]);
-        mockChatMessageStore.Setup(s => s.InvokedAsync(
-            It.IsAny<ChatMessageStore.InvokedContext>(),
-            It.IsAny<CancellationToken>())).Returns(new ValueTask());
-
-        Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, CancellationToken, ValueTask<ChatMessageStore>>> mockFactory = new();
-        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>())).ReturnsAsync(mockChatMessageStore.Object);
-
-        ChatClientAgent agent = new(mockService.Object, options: new()
-        {
-            ChatOptions = new() { Instructions = "test instructions" },
-            ChatMessageStoreFactory = mockFactory.Object
-        });
-
-        // Act
-        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
-        await agent.RunAsync([new(ChatRole.User, "test")], thread);
-
-        // Assert
-        Assert.IsType<ChatMessageStore>(thread!.MessageStore, exactMatch: false);
-        mockService.Verify(
-            x => x.GetResponseAsync(
-                It.Is<IEnumerable<ChatMessage>>(msgs => msgs.Count() == 2 && msgs.Any(m => m.Text == "Existing Chat History") && msgs.Any(m => m.Text == "test")),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>()),
-            Times.Once);
-        mockChatMessageStore.Verify(s => s.InvokingAsync(
-            It.Is<ChatMessageStore.InvokingContext>(x => x.RequestMessages.Count() == 1),
-            It.IsAny<CancellationToken>()),
-            Times.Once);
-        mockChatMessageStore.Verify(s => s.InvokedAsync(
-            It.Is<ChatMessageStore.InvokedContext>(x => x.RequestMessages.Count() == 1 && x.ChatMessageStoreMessages.Count() == 1 && x.ResponseMessages!.Count() == 1),
-            It.IsAny<CancellationToken>()),
-            Times.Once);
-        mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>()), Times.Once);
-    }
-
-    /// <summary>
-    /// Verify that RunAsync notifies the ChatMessageStore on failure.
-    /// </summary>
-    [Fact]
-    public async Task RunAsyncNotifiesChatMessageStoreOnFailureAsync()
-    {
-        // Arrange
-        Mock<IChatClient> mockService = new();
-        mockService.Setup(
-            s => s.GetResponseAsync(
-                It.IsAny<IEnumerable<ChatMessage>>(),
-                It.IsAny<ChatOptions>(),
-                It.IsAny<CancellationToken>())).Throws(new InvalidOperationException("Test Error"));
-
-        Mock<ChatMessageStore> mockChatMessageStore = new();
-
-        Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, CancellationToken, ValueTask<ChatMessageStore>>> mockFactory = new();
-        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>())).ReturnsAsync(mockChatMessageStore.Object);
-
-        ChatClientAgent agent = new(mockService.Object, options: new()
-        {
-            ChatOptions = new() { Instructions = "test instructions" },
-            ChatMessageStoreFactory = mockFactory.Object
-        });
-
-        // Act
-        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
-        await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync([new(ChatRole.User, "test")], thread));
-
-        // Assert
-        Assert.IsType<ChatMessageStore>(thread!.MessageStore, exactMatch: false);
-        mockChatMessageStore.Verify(s => s.InvokedAsync(
-            It.Is<ChatMessageStore.InvokedContext>(x => x.RequestMessages.Count() == 1 && x.ResponseMessages == null && x.InvokeException!.Message == "Test Error"),
-            It.IsAny<CancellationToken>()),
-            Times.Once);
-        mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>()), Times.Once);
-    }
-
-    /// <summary>
-    /// Verify that RunAsync throws when a ChatMessageStore Factory is provided and the chat client returns a conversation id.
-    /// </summary>
-    [Fact]
-    public async Task RunAsyncThrowsWhenChatMessageStoreFactoryProvidedAndConversationIdReturnedByChatClientAsync()
-    {
-        // 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")]) { ConversationId = "ConvId" });
-        Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, CancellationToken, ValueTask<ChatMessageStore>>> mockFactory = new();
-        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>())).ReturnsAsync(new InMemoryChatMessageStore());
-        ChatClientAgent agent = new(mockService.Object, options: new()
-        {
-            ChatOptions = new() { Instructions = "test instructions" },
-            ChatMessageStoreFactory = mockFactory.Object
-        });
-
-        // Act & Assert
-        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
-        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>
     /// Verify that RunAsync invokes any provided AIContextProvider and uses the result.
     /// </summary>

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

@@ -0,0 +1,371 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Moq;
+using Xunit.Sdk;
+
+namespace Microsoft.Agents.AI.UnitTests;
+
+/// <summary>
+/// Contains unit tests that verify the chat history management functionality of the <see cref="ChatClientAgent"/> class,
+/// e.g. that it correctly reads and updates chat history in any available <see cref="ChatMessageStore"/> or that
+/// it uses conversation id correctly for service managed chat history.
+/// </summary>
+public class ChatClientAgent_ChatHistoryManagementTests
+{
+    #region ConversationId Tests
+
+    /// <summary>
+    /// Verify that RunAsync does not throw when providing a ConversationId via both AgentThread and
+    /// via ChatOptions and the two are the same.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_DoesNotThrow_WhenSpecifyingTwoSameConversationIdsAsync()
+    {
+        // Arrange
+        var chatOptions = new ChatOptions { ConversationId = "ConvId" };
+        Mock<IChatClient> mockService = new();
+        mockService.Setup(
+            s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.Is<ChatOptions>(opts => opts.ConversationId == "ConvId"),
+                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]) { ConversationId = "ConvId" });
+
+        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
+
+        ChatClientAgentThread thread = new() { ConversationId = "ConvId" };
+
+        // Act & Assert
+        var response = await agent.RunAsync([new(ChatRole.User, "test")], thread, options: new ChatClientAgentRunOptions(chatOptions));
+        Assert.NotNull(response);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync throws when providing a ConversationId via both AgentThread and
+    /// via ChatOptions and the two are different.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_Throws_WhenSpecifyingTwoDifferentConversationIdsAsync()
+    {
+        // Arrange
+        var chatOptions = new ChatOptions { ConversationId = "ConvId" };
+        Mock<IChatClient> mockService = new();
+
+        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
+
+        ChatClientAgentThread thread = new() { ConversationId = "ThreadId" };
+
+        // Act & Assert
+        await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync([new(ChatRole.User, "test")], thread, options: new ChatClientAgentRunOptions(chatOptions)));
+    }
+
+    /// <summary>
+    /// Verify that RunAsync clones the ChatOptions when providing a thread with a ConversationId and a ChatOptions.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_ClonesChatOptions_ToAddConversationIdAsync()
+    {
+        // Arrange
+        var chatOptions = new ChatOptions { MaxOutputTokens = 100 };
+        Mock<IChatClient> mockService = new();
+        mockService.Setup(
+            s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.Is<ChatOptions>(opts => opts.MaxOutputTokens == 100 && opts.ConversationId == "ConvId"),
+                It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]) { ConversationId = "ConvId" });
+
+        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
+
+        ChatClientAgentThread thread = new() { ConversationId = "ConvId" };
+
+        // Act
+        await agent.RunAsync([new(ChatRole.User, "test")], thread, options: new ChatClientAgentRunOptions(chatOptions));
+
+        // Assert
+        Assert.Null(chatOptions.ConversationId);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync throws if a thread is provided that uses a conversation id already, but the service does not return one on invoke.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_Throws_ForMissingConversationIdWithConversationIdThreadAsync()
+    {
+        // 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")]));
+
+        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
+
+        ChatClientAgentThread thread = new() { ConversationId = "ConvId" };
+
+        // Act & Assert
+        await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync([new(ChatRole.User, "test")], thread));
+    }
+
+    /// <summary>
+    /// Verify that RunAsync sets the ConversationId on the thread when the service returns one.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_SetsConversationIdOnThread_WhenReturnedByChatClientAsync()
+    {
+        // 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")]) { ConversationId = "ConvId" });
+        ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" } });
+        ChatClientAgentThread thread = new();
+
+        // Act
+        await agent.RunAsync([new(ChatRole.User, "test")], thread);
+
+        // Assert
+        Assert.Equal("ConvId", thread.ConversationId);
+    }
+
+    #endregion
+
+    #region ChatMessageStore Tests
+
+    /// <summary>
+    /// Verify that RunAsync uses the default InMemoryChatMessageStore when the chat client returns no conversation id.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_UsesDefaultInMemoryChatMessageStore_WhenNoConversationIdReturnedByChatClientAsync()
+    {
+        // 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")]));
+        ChatClientAgent agent = new(mockService.Object, options: new()
+        {
+            ChatOptions = new() { Instructions = "test instructions" },
+        });
+
+        // Act
+        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
+        await agent.RunAsync([new(ChatRole.User, "test")], thread);
+
+        // Assert
+        var messageStore = Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
+        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 RunAsync_UsesChatMessageStoreFactory_WhenProvidedAndNoConversationIdReturnedByChatClientAsync()
+    {
+        // 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();
+        mockChatMessageStore.Setup(s => s.InvokingAsync(
+            It.IsAny<ChatMessageStore.InvokingContext>(),
+            It.IsAny<CancellationToken>())).ReturnsAsync([new ChatMessage(ChatRole.User, "Existing Chat History")]);
+        mockChatMessageStore.Setup(s => s.InvokedAsync(
+            It.IsAny<ChatMessageStore.InvokedContext>(),
+            It.IsAny<CancellationToken>())).Returns(new ValueTask());
+
+        Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, CancellationToken, ValueTask<ChatMessageStore>>> mockFactory = new();
+        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>())).ReturnsAsync(mockChatMessageStore.Object);
+
+        ChatClientAgent agent = new(mockService.Object, options: new()
+        {
+            ChatOptions = new() { Instructions = "test instructions" },
+            ChatMessageStoreFactory = mockFactory.Object
+        });
+
+        // Act
+        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
+        await agent.RunAsync([new(ChatRole.User, "test")], thread);
+
+        // Assert
+        Assert.IsType<ChatMessageStore>(thread!.MessageStore, exactMatch: false);
+        mockService.Verify(
+            x => x.GetResponseAsync(
+                It.Is<IEnumerable<ChatMessage>>(msgs => msgs.Count() == 2 && msgs.Any(m => m.Text == "Existing Chat History") && msgs.Any(m => m.Text == "test")),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>()),
+            Times.Once);
+        mockChatMessageStore.Verify(s => s.InvokingAsync(
+            It.Is<ChatMessageStore.InvokingContext>(x => x.RequestMessages.Count() == 1),
+            It.IsAny<CancellationToken>()),
+            Times.Once);
+        mockChatMessageStore.Verify(s => s.InvokedAsync(
+            It.Is<ChatMessageStore.InvokedContext>(x => x.RequestMessages.Count() == 1 && x.ChatMessageStoreMessages != null && x.ChatMessageStoreMessages.Count() == 1 && x.ResponseMessages!.Count() == 1),
+            It.IsAny<CancellationToken>()),
+            Times.Once);
+        mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync notifies the ChatMessageStore on failure.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_NotifiesChatMessageStore_OnFailureAsync()
+    {
+        // Arrange
+        Mock<IChatClient> mockService = new();
+        mockService.Setup(
+            s => s.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>())).Throws(new InvalidOperationException("Test Error"));
+
+        Mock<ChatMessageStore> mockChatMessageStore = new();
+
+        Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, CancellationToken, ValueTask<ChatMessageStore>>> mockFactory = new();
+        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>())).ReturnsAsync(mockChatMessageStore.Object);
+
+        ChatClientAgent agent = new(mockService.Object, options: new()
+        {
+            ChatOptions = new() { Instructions = "test instructions" },
+            ChatMessageStoreFactory = mockFactory.Object
+        });
+
+        // Act
+        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
+        await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync([new(ChatRole.User, "test")], thread));
+
+        // Assert
+        Assert.IsType<ChatMessageStore>(thread!.MessageStore, exactMatch: false);
+        mockChatMessageStore.Verify(s => s.InvokedAsync(
+            It.Is<ChatMessageStore.InvokedContext>(x => x.RequestMessages.Count() == 1 && x.ResponseMessages == null && x.InvokeException!.Message == "Test Error"),
+            It.IsAny<CancellationToken>()),
+            Times.Once);
+        mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that RunAsync throws when a ChatMessageStore Factory is provided and the chat client returns a conversation id.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_Throws_WhenChatMessageStoreFactoryProvidedAndConversationIdReturnedByChatClientAsync()
+    {
+        // 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")]) { ConversationId = "ConvId" });
+        Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, CancellationToken, ValueTask<ChatMessageStore>>> mockFactory = new();
+        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>())).ReturnsAsync(new InMemoryChatMessageStore());
+        ChatClientAgent agent = new(mockService.Object, options: new()
+        {
+            ChatOptions = new() { Instructions = "test instructions" },
+            ChatMessageStoreFactory = mockFactory.Object
+        });
+
+        // Act & Assert
+        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
+        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);
+    }
+
+    #endregion
+
+    #region ChatMessageStore Override Tests
+
+    /// <summary>
+    /// Tests that RunAsync uses an override ChatMessageStore provided via AdditionalProperties instead of the store from a factory
+    /// if one is supplied.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_UsesOverrideChatMessageStore_WhenProvidedViaAdditionalPropertiesAsync()
+    {
+        // 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")]));
+
+        // Arrange a chat message store to override the factory provided one.
+        Mock<ChatMessageStore> mockOverrideChatMessageStore = new();
+        mockOverrideChatMessageStore.Setup(s => s.InvokingAsync(
+            It.IsAny<ChatMessageStore.InvokingContext>(),
+            It.IsAny<CancellationToken>())).ReturnsAsync([new ChatMessage(ChatRole.User, "Existing Chat History")]);
+        mockOverrideChatMessageStore.Setup(s => s.InvokedAsync(
+            It.IsAny<ChatMessageStore.InvokedContext>(),
+            It.IsAny<CancellationToken>())).Returns(new ValueTask());
+
+        // Arrange a chat message store to provide to the agent via a factory at construction time.
+        // This one shouldn't be used since it is being overridden.
+        Mock<ChatMessageStore> mockFactoryChatMessageStore = new();
+        mockFactoryChatMessageStore.Setup(s => s.InvokingAsync(
+            It.IsAny<ChatMessageStore.InvokingContext>(),
+            It.IsAny<CancellationToken>())).ThrowsAsync(FailException.ForFailure("Base ChatMessageStore shouldn't be used."));
+        mockFactoryChatMessageStore.Setup(s => s.InvokedAsync(
+            It.IsAny<ChatMessageStore.InvokedContext>(),
+            It.IsAny<CancellationToken>())).Throws(FailException.ForFailure("Base ChatMessageStore shouldn't be used."));
+
+        Mock<Func<ChatClientAgentOptions.ChatMessageStoreFactoryContext, CancellationToken, ValueTask<ChatMessageStore>>> mockFactory = new();
+        mockFactory.Setup(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>(), It.IsAny<CancellationToken>())).ReturnsAsync(mockFactoryChatMessageStore.Object);
+
+        ChatClientAgent agent = new(mockService.Object, options: new()
+        {
+            ChatOptions = new() { Instructions = "test instructions" },
+            ChatMessageStoreFactory = mockFactory.Object
+        });
+
+        // Act
+        ChatClientAgentThread? thread = await agent.GetNewThreadAsync() as ChatClientAgentThread;
+        var additionalProperties = new AdditionalPropertiesDictionary();
+        additionalProperties.Add(mockOverrideChatMessageStore.Object);
+        await agent.RunAsync([new(ChatRole.User, "test")], thread, options: new AgentRunOptions { AdditionalProperties = additionalProperties });
+
+        // Assert
+        Assert.Same(mockFactoryChatMessageStore.Object, thread!.MessageStore);
+        mockService.Verify(
+            x => x.GetResponseAsync(
+                It.Is<IEnumerable<ChatMessage>>(msgs => msgs.Count() == 2 && msgs.Any(m => m.Text == "Existing Chat History") && msgs.Any(m => m.Text == "test")),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>()),
+            Times.Once);
+        mockOverrideChatMessageStore.Verify(s => s.InvokingAsync(
+            It.Is<ChatMessageStore.InvokingContext>(x => x.RequestMessages.Count() == 1),
+            It.IsAny<CancellationToken>()),
+            Times.Once);
+        mockOverrideChatMessageStore.Verify(s => s.InvokedAsync(
+            It.Is<ChatMessageStore.InvokedContext>(x => x.RequestMessages.Count() == 1 && x.ChatMessageStoreMessages != null && x.ChatMessageStoreMessages.Count() == 1 && x.ResponseMessages!.Count() == 1),
+            It.IsAny<CancellationToken>()),
+            Times.Once);
+
+        mockFactoryChatMessageStore.Verify(s => s.InvokingAsync(
+            It.IsAny<ChatMessageStore.InvokingContext>(),
+            It.IsAny<CancellationToken>()),
+            Times.Never);
+        mockFactoryChatMessageStore.Verify(s => s.InvokedAsync(
+            It.IsAny<ChatMessageStore.InvokedContext>(),
+            It.IsAny<CancellationToken>()),
+            Times.Never);
+    }
+
+    #endregion
+}

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

@@ -5,7 +5,7 @@ using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
 using Moq;
 
-namespace Microsoft.Agents.AI.UnitTests.ChatClient;
+namespace Microsoft.Agents.AI.UnitTests;
 
 /// <summary>
 /// Contains unit tests for the ChatClientAgent.DeserializeThread methods.

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

@@ -4,7 +4,7 @@ using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
 using Moq;
 
-namespace Microsoft.Agents.AI.UnitTests.ChatClient;
+namespace Microsoft.Agents.AI.UnitTests;
 
 /// <summary>
 /// Contains unit tests for the ChatClientAgent.GetNewThreadAsync methods.