From 90493aeb89dd14afeaf88f77d93a3ed4630a8e7b Mon Sep 17 00:00:00 2001
From: westey <164392973+westey-m@users.noreply.github.com>
Date: Mon, 2 Mar 2026 17:00:16 +0000
Subject: [PATCH 1/3] Change *Provider StateKey to list of StateKeys

---
 .../01-get-started/04_memory/Program.cs       |  2 +-
 .../Program.cs                                |  2 +-
 .../AIContextProvider.cs                      | 11 ++---
 .../ChatHistoryProvider.cs                    | 11 ++---
 .../InMemoryChatHistoryProvider.cs            |  2 +-
 .../CosmosChatHistoryProvider.cs              |  2 +-
 .../FoundryMemoryProvider.cs                  |  2 +-
 .../Microsoft.Agents.AI.Mem0/Mem0Provider.cs  |  2 +-
 .../WorkflowChatHistoryProvider.cs            |  2 +-
 .../ChatClient/ChatClientAgent.cs             | 40 +++++++++++++------
 .../Memory/ChatHistoryMemoryProvider.cs       |  2 +-
 .../Microsoft.Agents.AI/TextSearchProvider.cs |  2 +-
 .../InMemoryChatHistoryProviderTests.cs       | 10 +++--
 .../CosmosChatHistoryProviderTests.cs         | 10 +++--
 .../Mem0ProviderTests.cs                      | 12 +++---
 .../AIContextProviderChatClientTests.cs       |  2 +-
 .../ChatClient/ChatClientAgentTests.cs        | 29 ++++++++------
 ...hatClientAgent_BackgroundResponsesTests.cs | 16 ++++----
 ...tClientAgent_ChatHistoryManagementTests.cs |  4 ++
 .../Data/TextSearchProviderTests.cs           | 10 +++--
 .../Memory/ChatHistoryMemoryProviderTests.cs  | 10 +++--
 21 files changed, 109 insertions(+), 74 deletions(-)

diff --git a/dotnet/samples/01-get-started/04_memory/Program.cs b/dotnet/samples/01-get-started/04_memory/Program.cs
index fa6940f5fd..c25dec2f17 100644
--- a/dotnet/samples/01-get-started/04_memory/Program.cs
+++ b/dotnet/samples/01-get-started/04_memory/Program.cs
@@ -100,7 +100,7 @@ public UserInfoMemory(IChatClient chatClient, Func<AgentSession?, UserInfo>? sta
             this._chatClient = chatClient;
         }
 
-        public override string StateKey => this._sessionState.StateKey;
+        public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
 
         public UserInfo GetUserInfo(AgentSession session)
             => this._sessionState.GetOrInitializeState(session);
diff --git a/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs b/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs
index cbcf14157e..8fd16cdd94 100644
--- a/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs
+++ b/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs
@@ -93,7 +93,7 @@ public VectorChatHistoryProvider(
             this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
         }
 
-        public override string StateKey => this._sessionState.StateKey;
+        public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
 
         public string GetSessionDbKey(AgentSession session)
             => this._sessionState.GetOrInitializeState(session).SessionDbKey;
diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs
index 7ac4eed18c..c588ca0aed 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs
@@ -58,14 +58,15 @@ protected AIContextProvider(
     protected Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> StoreInputMessageFilter { get; }
 
     /// <summary>
-    /// Gets the key used to store the provider state in the <see cref="AgentSession.StateBag"/>.
+    /// Gets the set of keys used to store the provider state in the <see cref="AgentSession.StateBag"/>.
     /// </summary>
     /// <remarks>
-    /// The default value is the name of the concrete type (e.g. <c>"TextSearchProvider"</c>).
-    /// Implementations may override this to provide a custom key, for example when multiple
-    /// instances of the same provider type are used in the same session.
+    /// The default value is a single-element set containing the name of the concrete type (e.g. <c>"TextSearchProvider"</c>).
+    /// Implementations may override this to provide custom keys, for example when multiple
+    /// instances of the same provider type are used in the same session, or when a provider
+    /// stores state under more than one key.
     /// </remarks>
-    public virtual string StateKey => this.GetType().Name;
+    public virtual IReadOnlyList<string> StateKeys => [this.GetType().Name];
 
     /// <summary>
     /// Called at the start of agent invocation to provide additional context.
diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs
index ad3f3aacfb..bd53107107 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs
@@ -60,14 +60,15 @@ protected ChatHistoryProvider(
     }
 
     /// <summary>
-    /// Gets the key used to store the provider state in the <see cref="AgentSession.StateBag"/>.
+    /// Gets the set of keys used to store the provider state in the <see cref="AgentSession.StateBag"/>.
     /// </summary>
     /// <remarks>
-    /// The default value is the name of the concrete type (e.g. <c>"InMemoryChatHistoryProvider"</c>).
-    /// Implementations may override this to provide a custom key, for example when multiple
-    /// instances of the same provider type are used in the same session.
+    /// The default value is a single-element set containing the name of the concrete type (e.g. <c>"InMemoryChatHistoryProvider"</c>).
+    /// Implementations may override this to provide custom keys, for example when multiple
+    /// instances of the same provider type are used in the same session, or when a provider
+    /// stores state under more than one key.
     /// </remarks>
-    public virtual string StateKey => this.GetType().Name;
+    public virtual IReadOnlyList<string> StateKeys => [this.GetType().Name];
 
     /// <summary>
     /// Called at the start of agent invocation to provide messages for the next agent invocation.
diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs
index 12e935b23e..6ecedcf5f4 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs
@@ -49,7 +49,7 @@ public InMemoryChatHistoryProvider(InMemoryChatHistoryProviderOptions? options =
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
 
     /// <summary>
     /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied.
diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
index f1670fbb84..df4b5e35e8 100644
--- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
@@ -112,7 +112,7 @@ public CosmosChatHistoryProvider(
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
 
     /// <summary>
     /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using a connection string.
diff --git a/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs
index 9ffeda3fb5..5e5e05eb72 100644
--- a/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs
@@ -82,7 +82,7 @@ public FoundryMemoryProvider(
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
 
     private static Func<AgentSession?, State> ValidateStateInitializer(Func<AgentSession?, State> stateInitializer) =>
         session =>
diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs
index 1924bc0da2..341cba326e 100644
--- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs
@@ -72,7 +72,7 @@ public Mem0Provider(HttpClient httpClient, Func<AgentSession?, State> stateIniti
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
 
     private static Func<AgentSession?, State> ValidateStateInitializer(Func<AgentSession?, State> stateInitializer) =>
         session =>
diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs
index b9d5f3ae49..a66152343a 100644
--- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs
@@ -31,7 +31,7 @@ public WorkflowChatHistoryProvider(JsonSerializerOptions? jsonSerializerOptions
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
 
     internal sealed class StoreState
     {
diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs
index d52ea52e43..7db4eff6d8 100644
--- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs
+++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs
@@ -112,7 +112,7 @@ public ChatClientAgent(IChatClient chatClient, ChatClientAgentOptions? options,
         this.ChatHistoryProvider = options?.ChatHistoryProvider ?? new InMemoryChatHistoryProvider();
         this.AIContextProviders = this._agentOptions?.AIContextProviders as IReadOnlyList<AIContextProvider> ?? this._agentOptions?.AIContextProviders?.ToList();
 
-        // Validate that no two providers share the same StateKey, since they would overwrite each other's state in the session.
+        // Validate that no two providers share any StateKeys, since they would overwrite each other's state in the session.
         this._aiContextProviderStateKeys = ValidateAndCollectStateKeys(this._agentOptions?.AIContextProviders, this.ChatHistoryProvider);
 
         this._logger = (loggerFactory ?? chatClient.GetService<ILoggerFactory>() ?? NullLoggerFactory.Instance).CreateLogger<ChatClientAgent>();
@@ -824,11 +824,17 @@ private Task NotifyChatHistoryProviderOfNewMessagesAsync(
                     $"Only {nameof(ChatClientAgentSession.ConversationId)} or {nameof(this.ChatHistoryProvider)} may be used, but not both. The current {nameof(ChatClientAgentSession)} has a {nameof(ChatClientAgentSession.ConversationId)} indicating server-side chat history management, but an override {nameof(this.ChatHistoryProvider)} was provided via {nameof(AgentRunOptions.AdditionalProperties)}.");
             }
 
-            // Validate that the override provider's StateKey does not clash with any AIContextProvider's StateKey.
-            if (overrideProvider is not null && this._aiContextProviderStateKeys.Contains(overrideProvider.StateKey))
+            // Validate that the override provider's StateKeys do not clash with any AIContextProvider's StateKeys.
+            if (overrideProvider is not null)
             {
-                throw new InvalidOperationException(
-                    $"The ChatHistoryProvider '{overrideProvider.GetType().Name}' uses the state key '{overrideProvider.StateKey}' which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state.");
+                foreach (var key in overrideProvider.StateKeys)
+                {
+                    if (this._aiContextProviderStateKeys.Contains(key))
+                    {
+                        throw new InvalidOperationException(
+                            $"The ChatHistoryProvider '{overrideProvider.GetType().Name}' uses state key '{key}' which is already used by one of the configured AIContextProviders. Each provider must use unique state keys to avoid overwriting each other's state.");
+                    }
+                }
             }
 
             provider = overrideProvider;
@@ -879,7 +885,7 @@ private static List<ChatResponseUpdate> GetResponseUpdates(ChatClientAgentContin
     private string GetLoggingAgentName() => this.Name ?? "UnnamedAgent";
 
     /// <summary>
-    /// Validates that all configured providers have unique <see cref="AIContextProvider.StateKey"/> values
+    /// Validates that all configured providers have unique <see cref="AIContextProvider.StateKeys"/> values
     /// and returns a <see cref="HashSet{T}"/> of the AIContextProvider state keys.
     /// </summary>
     private static HashSet<string> ValidateAndCollectStateKeys(IEnumerable<AIContextProvider>? aiContextProviders, ChatHistoryProvider? chatHistoryProvider)
@@ -890,10 +896,13 @@ private static HashSet<string> ValidateAndCollectStateKeys(IEnumerable<AIContext
         {
             foreach (var provider in aiContextProviders)
             {
-                if (!stateKeys.Add(provider.StateKey))
+                foreach (var key in provider.StateKeys)
                 {
-                    throw new InvalidOperationException(
-                        $"Multiple providers use the same state key '{provider.StateKey}'. Each provider must use a unique state key to avoid overwriting each other's state.");
+                    if (!stateKeys.Add(key))
+                    {
+                        throw new InvalidOperationException(
+                            $"Multiple providers use the same state key '{key}'. Each provider must use a unique state key to avoid overwriting each other's state.");
+                    }
                 }
             }
         }
@@ -905,11 +914,16 @@ private static HashSet<string> ValidateAndCollectStateKeys(IEnumerable<AIContext
                 $"The default {nameof(InMemoryChatHistoryProvider)} uses the state key '{nameof(InMemoryChatHistoryProvider)}', which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state. To resolve this, either configure a different state key for the AIContextProvider that is using '{nameof(InMemoryChatHistoryProvider)}' as its state key, or provide a custom ChatHistoryProvider with a unique state key.");
         }
 
-        if (chatHistoryProvider is not null
-            && stateKeys.Contains(chatHistoryProvider.StateKey))
+        if (chatHistoryProvider is not null)
         {
-            throw new InvalidOperationException(
-                $"The ChatHistoryProvider '{chatHistoryProvider.GetType().Name}' uses the state key '{chatHistoryProvider.StateKey}' which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state. To resolve this, either configure a different state key for the AIContextProvider that is using '{chatHistoryProvider.StateKey}' as its state key, or reconfigure the custom ChatHistoryProvider with a unique state key.");
+            foreach (var key in chatHistoryProvider.StateKeys)
+            {
+                if (stateKeys.Contains(key))
+                {
+                    throw new InvalidOperationException(
+                        $"The ChatHistoryProvider '{chatHistoryProvider.GetType().Name}' uses state key '{key}' which is already used by one of the configured AIContextProviders. Each provider must use unique state keys to avoid overwriting each other's state. To resolve this, either configure different state keys for the AIContextProvider that shares keys with the ChatHistoryProvider, or reconfigure the custom ChatHistoryProvider with unique state keys.");
+                }
+            }
         }
 
         return stateKeys;
diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs
index 7905db74b8..59fb1161d9 100644
--- a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs
@@ -128,7 +128,7 @@ public ChatHistoryMemoryProvider(
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
 
     /// <inheritdoc />
     protected override async ValueTask<AIContext> ProvideAIContextAsync(AIContextProvider.InvokingContext context, CancellationToken cancellationToken = default)
diff --git a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs
index dd62b0eb9b..6dd9e9571a 100644
--- a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs
@@ -88,7 +88,7 @@ public TextSearchProvider(
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._sessionState.StateKey;
+    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
 
     /// <inheritdoc />
     protected override async ValueTask<AIContext> ProvideAIContextAsync(AIContextProvider.InvokingContext context, CancellationToken cancellationToken = default)
diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs
index ebe1131ab7..93c9854eae 100644
--- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs
@@ -43,23 +43,25 @@ public void Constructor_Arguments_SetOnPropertiesCorrectly()
     }
 
     [Fact]
-    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
     {
         // Arrange & Act
         var provider = new InMemoryChatHistoryProvider();
 
         // Assert
-        Assert.Equal("InMemoryChatHistoryProvider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("InMemoryChatHistoryProvider", provider.StateKeys);
     }
 
     [Fact]
-    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
     {
         // Arrange & Act
         var provider = new InMemoryChatHistoryProvider(new() { StateKey = "custom-key" });
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [Fact]
diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs
index a790b19cdd..059639d899 100644
--- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs
@@ -150,7 +150,7 @@ private void SkipIfEmulatorNotAvailable()
 
     [SkippableFact]
     [Trait("Category", "CosmosDB")]
-    public void StateKey_ReturnsDefaultKey_WhenNoStateKeyProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoStateKeyProvided()
     {
         // Arrange & Act
         this.SkipIfEmulatorNotAvailable();
@@ -159,12 +159,13 @@ public void StateKey_ReturnsDefaultKey_WhenNoStateKeyProvided()
             _ => new CosmosChatHistoryProvider.State("test-conversation"));
 
         // Assert
-        Assert.Equal("CosmosChatHistoryProvider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("CosmosChatHistoryProvider", provider.StateKeys);
     }
 
     [SkippableFact]
     [Trait("Category", "CosmosDB")]
-    public void StateKey_ReturnsCustomKey_WhenSetViaConstructor()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaConstructor()
     {
         // Arrange & Act
         this.SkipIfEmulatorNotAvailable();
@@ -174,7 +175,8 @@ public void StateKey_ReturnsCustomKey_WhenSetViaConstructor()
             stateKey: "custom-key");
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [SkippableFact]
diff --git a/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs
index 02e18f324e..122627da76 100644
--- a/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs
@@ -67,17 +67,18 @@ public void Constructor_Throws_WhenStateInitializerIsNull()
     }
 
     [Fact]
-    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
     {
         // Arrange & Act
         var provider = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(new Mem0ProviderScope { ThreadId = "tid" }));
 
         // Assert
-        Assert.Equal("Mem0Provider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("Mem0Provider", provider.StateKeys);
     }
 
     [Fact]
-    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
     {
         // Arrange & Act
         var provider = new Mem0Provider(
@@ -86,7 +87,8 @@ public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
             new Mem0ProviderOptions { StateKey = "custom-key" });
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [Fact]
@@ -419,7 +421,7 @@ public async Task StateInitializer_IsCalledOnceAndStoredInStateBagAsync()
     }
 
     [Fact]
-    public async Task StateKey_CanBeConfiguredViaOptionsAsync()
+    public async Task StateKeys_CanBeConfiguredViaOptionsAsync()
     {
         // Arrange
         this._handler.EnqueueJsonResponse("[]");
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs
index 51eb4be3ab..69fe0e1f22 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs
@@ -389,7 +389,7 @@ private sealed class TestAIContextProvider : AIContextProvider
 
         public InvokedContext? LastInvokedContext { get; private set; }
 
-        public override string StateKey => this._stateKey;
+        public override IReadOnlyList<string> StateKeys => [this._stateKey];
 
         public TestAIContextProvider(
             string stateKey,
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
index 12446c89c0..1f3a2d941c 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
@@ -105,8 +105,8 @@ public void Constructor_ThrowsWhenChatHistoryProviderStateKeyClashesWithAIContex
                 ChatHistoryProvider = historyProvider
             }));
 
-        Assert.Contains("SharedKey", ex.Message);
-        Assert.Contains(nameof(ChatHistoryProvider), ex.Message);
+        Assert.Contains("ChatHistoryProvider", ex.Message);
+        Assert.Contains("state key 'SharedKey'", ex.Message);
     }
 
     /// <summary>
@@ -159,11 +159,11 @@ public async Task RunAsync_ThrowsWhenOverrideChatHistoryProviderStateKeyClashesW
         var ex = await Assert.ThrowsAsync<InvalidOperationException>(() =>
             agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties }));
 
-        Assert.Contains("SharedKey", ex.Message);
+        Assert.Contains("state key 'SharedKey'", ex.Message);
     }
 
     /// <summary>
-    /// Verify that RunAsync succeeds when an override ChatHistoryProvider uses the same StateKey as the default ChatHistoryProvider.
+    /// Verify that RunAsync succeeds when an override ChatHistoryProvider uses the same StateKeys as the default ChatHistoryProvider.
     /// </summary>
     [Fact]
     public async Task RunAsync_SucceedsWhenOverrideChatHistoryProviderSharesKeyWithDefaultAsync()
@@ -489,6 +489,7 @@ public async Task RunAsyncInvokesAIContextProviderAndUsesResultAsync()
             .ReturnsAsync(new ChatResponse(responseMessages));
 
         var mockProvider = new Mock<AIContextProvider>(null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -560,6 +561,7 @@ public async Task RunAsyncInvokesAIContextProviderWhenGetResponseFailsAsync()
             .Throws(new InvalidOperationException("downstream failure"));
 
         var mockProvider = new Mock<AIContextProvider>(null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -618,6 +620,7 @@ public async Task RunAsyncInvokesAIContextProviderAndSucceedsWithEmptyAIContextA
             .ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
 
         var mockProvider = new Mock<AIContextProvider>(null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -678,7 +681,7 @@ public async Task RunAsyncInvokesMultipleAIContextProvidersInOrderAsync()
 
         // Provider 1: adds a system message and a tool
         var mockProvider1 = new Mock<AIContextProvider>(null, null);
-        mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockProvider1.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockProvider1
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -697,7 +700,7 @@ public async Task RunAsyncInvokesMultipleAIContextProvidersInOrderAsync()
         // Provider 2: adds another system message and verifies it receives accumulated context from provider 1
         AIContext? provider2ReceivedContext = null;
         var mockProvider2 = new Mock<AIContextProvider>(null, null);
-        mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
+        mockProvider2.SetupGet(p => p.StateKeys).Returns(["Provider2"]);
         mockProvider2
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -785,7 +788,7 @@ public async Task RunAsyncInvokesMultipleAIContextProvidersOnFailureAsync()
             .ThrowsAsync(new InvalidOperationException("downstream failure"));
 
         var mockProvider1 = new Mock<AIContextProvider>(null, null);
-        mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockProvider1.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockProvider1
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -802,7 +805,7 @@ public async Task RunAsyncInvokesMultipleAIContextProvidersOnFailureAsync()
             .Returns(new ValueTask());
 
         var mockProvider2 = new Mock<AIContextProvider>(null, null);
-        mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
+        mockProvider2.SetupGet(p => p.StateKeys).Returns(["Provider2"]);
         mockProvider2
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -870,7 +873,7 @@ public async Task RunStreamingAsyncInvokesMultipleAIContextProvidersAsync()
             .Returns(ToAsyncEnumerableAsync(responseUpdates));
 
         var mockProvider1 = new Mock<AIContextProvider>(null, null);
-        mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockProvider1.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockProvider1
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -887,7 +890,7 @@ public async Task RunStreamingAsyncInvokesMultipleAIContextProvidersAsync()
             .Returns(new ValueTask());
 
         var mockProvider2 = new Mock<AIContextProvider>(null, null);
-        mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
+        mockProvider2.SetupGet(p => p.StateKeys).Returns(["Provider2"]);
         mockProvider2
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -1829,6 +1832,7 @@ public async Task RunStreamingAsyncInvokesAIContextProviderAndUsesResultAsync()
             .Returns(ToAsyncEnumerableAsync(responseUpdates));
 
         var mockProvider = new Mock<AIContextProvider>(null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -1908,6 +1912,7 @@ public async Task RunStreamingAsyncInvokesAIContextProviderWhenGetResponseFailsA
             .Throws(new InvalidOperationException("downstream failure"));
 
         var mockProvider = new Mock<AIContextProvider>(null, null);
+        mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -1965,7 +1970,7 @@ private sealed partial class JsonContext2 : JsonSerializerContext;
 
     private sealed class TestAIContextProvider(string stateKey) : AIContextProvider
     {
-        public override string StateKey => stateKey;
+        public override IReadOnlyList<string> StateKeys => [stateKey];
 
         protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
             => new(context.AIContext);
@@ -1973,7 +1978,7 @@ protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext contex
 
     private sealed class TestChatHistoryProvider(string stateKey) : ChatHistoryProvider
     {
-        public override string StateKey => stateKey;
+        public override IReadOnlyList<string> StateKeys => [stateKey];
 
         protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
             => new(context.RequestMessages);
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs
index 64835f2b2f..52266a32c2 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs
@@ -339,7 +339,7 @@ public async Task RunAsync_WhenContinuationTokenProvided_SkipsSessionMessagePopu
 
         // Create a mock chat history provider that would normally provide messages
         var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null);
-        mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -347,7 +347,7 @@ public async Task RunAsync_WhenContinuationTokenProvided_SkipsSessionMessagePopu
 
         // Create a mock AI context provider that would normally provide context
         var mockContextProvider = new Mock<AIContextProvider>(null, null);
-        mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockContextProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -408,7 +408,7 @@ public async Task RunStreamingAsync_WhenContinuationTokenProvided_SkipsSessionMe
 
         // Create a mock chat history provider that would normally provide messages
         var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null);
-        mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -416,7 +416,7 @@ public async Task RunStreamingAsync_WhenContinuationTokenProvided_SkipsSessionMe
 
         // Create a mock AI context provider that would normally provide context
         var mockContextProvider = new Mock<AIContextProvider>(null, null);
-        mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockContextProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -639,7 +639,7 @@ public async Task RunStreamingAsync_WhenResumingStreaming_UsesUpdatesFromInitial
 
         List<ChatMessage> capturedMessagesAddedToProvider = [];
         var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null);
-        mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -648,7 +648,7 @@ public async Task RunStreamingAsync_WhenResumingStreaming_UsesUpdatesFromInitial
 
         AIContextProvider.InvokedContext? capturedInvokedContext = null;
         var mockContextProvider = new Mock<AIContextProvider>(null, null);
-        mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockContextProvider
             .Protected()
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -703,7 +703,7 @@ public async Task RunStreamingAsync_WhenResumingStreaming_UsesInputMessagesFromI
 
         List<ChatMessage> capturedMessagesAddedToProvider = [];
         var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null);
-        mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -712,7 +712,7 @@ public async Task RunStreamingAsync_WhenResumingStreaming_UsesInputMessagesFromI
 
         AIContextProvider.InvokedContext? capturedInvokedContext = null;
         var mockContextProvider = new Mock<AIContextProvider>(null, null);
-        mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+        mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
         mockContextProvider
             .Protected()
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs
index 4d8326269a..06d7362f7e 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs
@@ -186,6 +186,7 @@ public async Task RunAsync_UsesChatHistoryProvider_WhenProvidedAndNoConversation
                 It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
 
         Mock<ChatHistoryProvider> mockChatHistoryProvider = new(null, null);
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -241,6 +242,7 @@ public async Task RunAsync_NotifiesChatHistoryProvider_OnFailureAsync()
                 It.IsAny<CancellationToken>())).Throws(new InvalidOperationException("Test Error"));
 
         Mock<ChatHistoryProvider> mockChatHistoryProvider = new(null, null);
+        mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -430,6 +432,7 @@ public async Task RunAsync_UsesOverrideChatHistoryProvider_WhenProvidedViaAdditi
 
         // Arrange a chat history provider to override the factory provided one.
         Mock<ChatHistoryProvider> mockOverrideChatHistoryProvider = new(null, null);
+        mockOverrideChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
         mockOverrideChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -443,6 +446,7 @@ public async Task RunAsync_UsesOverrideChatHistoryProvider_WhenProvidedViaAdditi
         // Arrange a chat history provider to provide to the agent at construction time.
         // This one shouldn't be used since it is being overridden.
         Mock<ChatHistoryProvider> mockAgentOptionsChatHistoryProvider = new(null, null);
+        mockAgentOptionsChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
         mockAgentOptionsChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs
index 46c56fc483..c2b0be06ad 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs
@@ -39,17 +39,18 @@ public TextSearchProviderTests()
     }
 
     [Fact]
-    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
     {
         // Arrange & Act
         var provider = new TextSearchProvider((_, _) => Task.FromResult<IEnumerable<TextSearchProvider.TextSearchResult>>([]));
 
         // Assert
-        Assert.Equal("TextSearchProvider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("TextSearchProvider", provider.StateKeys);
     }
 
     [Fact]
-    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
     {
         // Arrange & Act
         var provider = new TextSearchProvider(
@@ -57,7 +58,8 @@ public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
             new TextSearchProviderOptions { StateKey = "custom-key" });
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [Theory]
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs
index ff5d709202..1e20c74bcb 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs
@@ -56,7 +56,7 @@ public ChatHistoryMemoryProviderTests()
     }
 
     [Fact]
-    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
     {
         // Arrange & Act
         var provider = new ChatHistoryMemoryProvider(
@@ -66,11 +66,12 @@ public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
             _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }));
 
         // Assert
-        Assert.Equal("ChatHistoryMemoryProvider", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("ChatHistoryMemoryProvider", provider.StateKeys);
     }
 
     [Fact]
-    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
     {
         // Arrange & Act
         var provider = new ChatHistoryMemoryProvider(
@@ -81,7 +82,8 @@ public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
             new ChatHistoryMemoryProviderOptions { StateKey = "custom-key" });
 
         // Assert
-        Assert.Equal("custom-key", provider.StateKey);
+        Assert.Single(provider.StateKeys);
+        Assert.Contains("custom-key", provider.StateKeys);
     }
 
     [Fact]

From ffc59ad5824d45ef4f004bc6d204778deaa224c9 Mon Sep 17 00:00:00 2001
From: westey <164392973+westey-m@users.noreply.github.com>
Date: Mon, 2 Mar 2026 17:32:49 +0000
Subject: [PATCH 2/3] Add more statekey validation tests

---
 .../ChatClient/ChatClientAgentTests.cs        | 115 ++++++++++++++++++
 1 file changed, 115 insertions(+)

diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
index 1f3a2d941c..6bcc0a8acd 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
@@ -192,6 +192,102 @@ public async Task RunAsync_SucceedsWhenOverrideChatHistoryProviderSharesKeyWithD
         await agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties });
     }
 
+    /// <summary>
+    /// Verify that the constructor throws when two multi-key AIContextProviders have an overlapping key.
+    /// </summary>
+    [Fact]
+    public void Constructor_ThrowsWhenMultiKeyAIContextProvidersOverlap()
+    {
+        // Arrange
+        var chatClient = new Mock<IChatClient>().Object;
+        var provider1 = new MultiKeyTestAIContextProvider("Key1", "SharedKey");
+        var provider2 = new MultiKeyTestAIContextProvider("Key2", "SharedKey");
+
+        // Act & Assert
+        var ex = Assert.Throws<InvalidOperationException>(() =>
+            new ChatClientAgent(chatClient, options: new()
+            {
+                AIContextProviders = [provider1, provider2]
+            }));
+
+        Assert.Contains("state key 'SharedKey'", ex.Message);
+    }
+
+    /// <summary>
+    /// Verify that the constructor throws when a multi-key ChatHistoryProvider has an overlapping key with an AIContextProvider.
+    /// </summary>
+    [Fact]
+    public void Constructor_ThrowsWhenMultiKeyChatHistoryProviderOverlapsWithAIContextProvider()
+    {
+        // Arrange
+        var chatClient = new Mock<IChatClient>().Object;
+        var contextProvider = new MultiKeyTestAIContextProvider("Key1", "SharedKey");
+        var historyProvider = new MultiKeyTestChatHistoryProvider("Key2", "SharedKey");
+
+        // Act & Assert
+        var ex = Assert.Throws<InvalidOperationException>(() =>
+            new ChatClientAgent(chatClient, options: new()
+            {
+                AIContextProviders = [contextProvider],
+                ChatHistoryProvider = historyProvider
+            }));
+
+        Assert.Contains("state key 'SharedKey'", ex.Message);
+    }
+
+    /// <summary>
+    /// Verify that the constructor succeeds when multi-key providers have no overlapping keys.
+    /// </summary>
+    [Fact]
+    public void Constructor_SucceedsWithMultiKeyProvidersWithUniqueKeys()
+    {
+        // Arrange
+        var chatClient = new Mock<IChatClient>().Object;
+        var contextProvider1 = new MultiKeyTestAIContextProvider("Key1", "Key2");
+        var contextProvider2 = new MultiKeyTestAIContextProvider("Key3", "Key4");
+        var historyProvider = new MultiKeyTestChatHistoryProvider("Key5", "Key6");
+
+        // Act & Assert - should not throw
+        _ = new ChatClientAgent(chatClient, options: new()
+        {
+            AIContextProviders = [contextProvider1, contextProvider2],
+            ChatHistoryProvider = historyProvider
+        });
+    }
+
+    /// <summary>
+    /// Verify that RunAsync throws when a multi-key override ChatHistoryProvider has an overlapping key with an AIContextProvider.
+    /// </summary>
+    [Fact]
+    public async Task RunAsync_ThrowsWhenMultiKeyOverrideChatHistoryProviderClashesWithAIContextProviderAsync()
+    {
+        // 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")]));
+
+        var contextProvider = new MultiKeyTestAIContextProvider("Key1", "SharedKey");
+        var overrideHistoryProvider = new MultiKeyTestChatHistoryProvider("Key2", "SharedKey");
+
+        ChatClientAgent agent = new(mockService.Object, options: new()
+        {
+            AIContextProviders = [contextProvider]
+        });
+
+        // Act & Assert
+        ChatClientAgentSession? session = await agent.CreateSessionAsync() as ChatClientAgentSession;
+        AdditionalPropertiesDictionary additionalProperties = new();
+        additionalProperties.Add<ChatHistoryProvider>(overrideHistoryProvider);
+
+        var ex = await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties }));
+
+        Assert.Contains("state key 'SharedKey'", ex.Message);
+    }
+
     #endregion
 
     #region RunAsync Tests
@@ -1976,6 +2072,14 @@ protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext contex
             => new(context.AIContext);
     }
 
+    private sealed class MultiKeyTestAIContextProvider(params string[] stateKeys) : AIContextProvider
+    {
+        public override IReadOnlyList<string> StateKeys => stateKeys;
+
+        protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+            => new(context.AIContext);
+    }
+
     private sealed class TestChatHistoryProvider(string stateKey) : ChatHistoryProvider
     {
         public override IReadOnlyList<string> StateKeys => [stateKey];
@@ -1986,4 +2090,15 @@ protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(Invokin
         protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
             => default;
     }
+
+    private sealed class MultiKeyTestChatHistoryProvider(params string[] stateKeys) : ChatHistoryProvider
+    {
+        public override IReadOnlyList<string> StateKeys => stateKeys;
+
+        protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+            => new(context.RequestMessages);
+
+        protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+            => default;
+    }
 }

From ea8553bad5696e4fda9b590969573fe51e96635c Mon Sep 17 00:00:00 2001
From: westey <164392973+westey-m@users.noreply.github.com>
Date: Mon, 2 Mar 2026 18:34:59 +0000
Subject: [PATCH 3/3] Address PR comments

---
 dotnet/samples/01-get-started/04_memory/Program.cs        | 3 ++-
 .../Agent_Step04_3rdPartyChatHistoryStorage/Program.cs    | 3 ++-
 .../Microsoft.Agents.AI.Abstractions/AIContextProvider.cs | 4 +++-
 .../ChatHistoryProvider.cs                                | 3 ++-
 .../InMemoryChatHistoryProvider.cs                        | 3 ++-
 .../CosmosChatHistoryProvider.cs                          | 3 ++-
 .../FoundryMemoryProvider.cs                              | 3 ++-
 dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs       | 3 ++-
 .../WorkflowChatHistoryProvider.cs                        | 3 ++-
 .../Memory/ChatHistoryMemoryProvider.cs                   | 3 ++-
 dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs      | 3 ++-
 .../AIContextProviderChatClientTests.cs                   | 6 +++---
 .../ChatClient/ChatClientAgentTests.cs                    | 8 ++++++--
 13 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/dotnet/samples/01-get-started/04_memory/Program.cs b/dotnet/samples/01-get-started/04_memory/Program.cs
index c25dec2f17..e07a84ad4b 100644
--- a/dotnet/samples/01-get-started/04_memory/Program.cs
+++ b/dotnet/samples/01-get-started/04_memory/Program.cs
@@ -89,6 +89,7 @@ namespace SampleApp
     internal sealed class UserInfoMemory : AIContextProvider
     {
         private readonly ProviderSessionState<UserInfo> _sessionState;
+        private IReadOnlyList<string>? _stateKeys;
         private readonly IChatClient _chatClient;
 
         public UserInfoMemory(IChatClient chatClient, Func<AgentSession?, UserInfo>? stateInitializer = null)
@@ -100,7 +101,7 @@ public UserInfoMemory(IChatClient chatClient, Func<AgentSession?, UserInfo>? sta
             this._chatClient = chatClient;
         }
 
-        public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
+        public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
         public UserInfo GetUserInfo(AgentSession session)
             => this._sessionState.GetOrInitializeState(session);
diff --git a/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs b/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs
index 8fd16cdd94..578d8d51a6 100644
--- a/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs
+++ b/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs
@@ -79,6 +79,7 @@ namespace SampleApp
     internal sealed class VectorChatHistoryProvider : ChatHistoryProvider
     {
         private readonly ProviderSessionState<State> _sessionState;
+        private IReadOnlyList<string>? _stateKeys;
         private readonly VectorStore _vectorStore;
 
         public VectorChatHistoryProvider(
@@ -93,7 +94,7 @@ public VectorChatHistoryProvider(
             this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
         }
 
-        public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
+        public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
         public string GetSessionDbKey(AgentSession session)
             => this._sessionState.GetOrInitializeState(session).SessionDbKey;
diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs
index c588ca0aed..09f8393f62 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs
@@ -34,6 +34,8 @@ public abstract class AIContextProvider
     private static IEnumerable<ChatMessage> DefaultExternalOnlyFilter(IEnumerable<ChatMessage> messages)
         => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External);
 
+    private IReadOnlyList<string>? _stateKeys;
+
     /// <summary>
     /// Initializes a new instance of the <see cref="AIContextProvider"/> class.
     /// </summary>
@@ -66,7 +68,7 @@ protected AIContextProvider(
     /// instances of the same provider type are used in the same session, or when a provider
     /// stores state under more than one key.
     /// </remarks>
-    public virtual IReadOnlyList<string> StateKeys => [this.GetType().Name];
+    public virtual IReadOnlyList<string> StateKeys => this._stateKeys ??= [this.GetType().Name];
 
     /// <summary>
     /// Called at the start of agent invocation to provide additional context.
diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs
index bd53107107..b8d3cc83b5 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs
@@ -43,6 +43,7 @@ public abstract class ChatHistoryProvider
     private static IEnumerable<ChatMessage> DefaultExcludeChatHistoryFilter(IEnumerable<ChatMessage> messages)
         => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory);
 
+    private IReadOnlyList<string>? _stateKeys;
     private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? _provideOutputMessageFilter;
     private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _storeInputMessageFilter;
 
@@ -68,7 +69,7 @@ protected ChatHistoryProvider(
     /// instances of the same provider type are used in the same session, or when a provider
     /// stores state under more than one key.
     /// </remarks>
-    public virtual IReadOnlyList<string> StateKeys => [this.GetType().Name];
+    public virtual IReadOnlyList<string> StateKeys => this._stateKeys ??= [this.GetType().Name];
 
     /// <summary>
     /// Called at the start of agent invocation to provide messages for the next agent invocation.
diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs
index 6ecedcf5f4..eb1ace8a1a 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs
@@ -27,6 +27,7 @@ namespace Microsoft.Agents.AI;
 public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
 {
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="InMemoryChatHistoryProvider"/> class.
@@ -49,7 +50,7 @@ public InMemoryChatHistoryProvider(InMemoryChatHistoryProviderOptions? options =
     }
 
     /// <inheritdoc />
-    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     /// <summary>
     /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied.
diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
index df4b5e35e8..6a1b8ed798 100644
--- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
@@ -22,6 +22,7 @@ namespace Microsoft.Agents.AI;
 public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
 {
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
     private readonly CosmosClient _cosmosClient;
     private readonly Container _container;
     private readonly bool _ownsClient;
@@ -112,7 +113,7 @@ public CosmosChatHistoryProvider(
     }
 
     /// <inheritdoc />
-    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     /// <summary>
     /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using a connection string.
diff --git a/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs
index 5e5e05eb72..6086b0c641 100644
--- a/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs
@@ -32,6 +32,7 @@ public sealed class FoundryMemoryProvider : AIContextProvider
     private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:";
 
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
     private readonly string _contextPrompt;
     private readonly string _memoryStoreName;
     private readonly int _maxMemories;
@@ -82,7 +83,7 @@ public FoundryMemoryProvider(
     }
 
     /// <inheritdoc />
-    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     private static Func<AgentSession?, State> ValidateStateInitializer(Func<AgentSession?, State> stateInitializer) =>
         session =>
diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs
index 341cba326e..9f809305d1 100644
--- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs
@@ -27,6 +27,7 @@ public sealed class Mem0Provider : MessageAIContextProvider
     private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:";
 
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
     private readonly string _contextPrompt;
     private readonly bool _enableSensitiveTelemetryData;
 
@@ -72,7 +73,7 @@ public Mem0Provider(HttpClient httpClient, Func<AgentSession?, State> stateIniti
     }
 
     /// <inheritdoc />
-    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     private static Func<AgentSession?, State> ValidateStateInitializer(Func<AgentSession?, State> stateInitializer) =>
         session =>
diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs
index a66152343a..ff52e86e58 100644
--- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs
@@ -12,6 +12,7 @@ namespace Microsoft.Agents.AI.Workflows;
 internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
 {
     private readonly ProviderSessionState<StoreState> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="WorkflowChatHistoryProvider"/> class.
@@ -31,7 +32,7 @@ public WorkflowChatHistoryProvider(JsonSerializerOptions? jsonSerializerOptions
     }
 
     /// <inheritdoc />
-    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     internal sealed class StoreState
     {
diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs
index 59fb1161d9..a0f50b4323 100644
--- a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs
@@ -54,6 +54,7 @@ public sealed class ChatHistoryMemoryProvider : MessageAIContextProvider, IDispo
     private const string ContentEmbeddingField = "ContentEmbedding";
 
     private readonly ProviderSessionState<State> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
 
 #pragma warning disable CA2213 // VectorStore is not owned by this class - caller is responsible for disposal
     private readonly VectorStore _vectorStore;
@@ -128,7 +129,7 @@ public ChatHistoryMemoryProvider(
     }
 
     /// <inheritdoc />
-    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     /// <inheritdoc />
     protected override async ValueTask<AIContext> ProvideAIContextAsync(AIContextProvider.InvokingContext context, CancellationToken cancellationToken = default)
diff --git a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs
index 6dd9e9571a..845a154bdf 100644
--- a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs
@@ -40,6 +40,7 @@ public sealed class TextSearchProvider : MessageAIContextProvider
     private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available.";
 
     private readonly ProviderSessionState<TextSearchProviderState> _sessionState;
+    private IReadOnlyList<string>? _stateKeys;
     private readonly Func<string, CancellationToken, Task<IEnumerable<TextSearchResult>>> _searchAsync;
     private readonly ILogger<TextSearchProvider>? _logger;
     private readonly AITool[] _tools;
@@ -88,7 +89,7 @@ public TextSearchProvider(
     }
 
     /// <inheritdoc />
-    public override IReadOnlyList<string> StateKeys => [this._sessionState.StateKey];
+    public override IReadOnlyList<string> StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
 
     /// <inheritdoc />
     protected override async ValueTask<AIContext> ProvideAIContextAsync(AIContextProvider.InvokingContext context, CancellationToken cancellationToken = default)
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs
index 69fe0e1f22..3b06bbb772 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs
@@ -380,7 +380,7 @@ private static async IAsyncEnumerable<ChatResponseUpdate> ToAsyncEnumerableAsync
     /// </summary>
     private sealed class TestAIContextProvider : AIContextProvider
     {
-        private readonly string _stateKey;
+        private readonly IReadOnlyList<string> _stateKeys;
         private readonly IEnumerable<ChatMessage> _provideMessages;
         private readonly string? _provideInstructions;
         private readonly IEnumerable<AITool>? _provideTools;
@@ -389,7 +389,7 @@ private sealed class TestAIContextProvider : AIContextProvider
 
         public InvokedContext? LastInvokedContext { get; private set; }
 
-        public override IReadOnlyList<string> StateKeys => [this._stateKey];
+        public override IReadOnlyList<string> StateKeys => this._stateKeys;
 
         public TestAIContextProvider(
             string stateKey,
@@ -397,7 +397,7 @@ public TestAIContextProvider(
             string? provideInstructions = null,
             IEnumerable<AITool>? provideTools = null)
         {
-            this._stateKey = stateKey;
+            this._stateKeys = [stateKey];
             this._provideMessages = provideMessages ?? [];
             this._provideInstructions = provideInstructions;
             this._provideTools = provideTools;
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
index 6bcc0a8acd..0a706bbf64 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
@@ -2066,7 +2066,9 @@ private sealed partial class JsonContext2 : JsonSerializerContext;
 
     private sealed class TestAIContextProvider(string stateKey) : AIContextProvider
     {
-        public override IReadOnlyList<string> StateKeys => [stateKey];
+        private readonly IReadOnlyList<string> _stateKeys = [stateKey];
+
+        public override IReadOnlyList<string> StateKeys => this._stateKeys;
 
         protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
             => new(context.AIContext);
@@ -2082,7 +2084,9 @@ protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext contex
 
     private sealed class TestChatHistoryProvider(string stateKey) : ChatHistoryProvider
     {
-        public override IReadOnlyList<string> StateKeys => [stateKey];
+        private readonly IReadOnlyList<string> _stateKeys = [stateKey];
+
+        public override IReadOnlyList<string> StateKeys => this._stateKeys;
 
         protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
             => new(context.RequestMessages);