Add a public StateKey property to providers (#3810)

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

@@ -75,12 +75,13 @@ namespace SampleApp
     /// </summary>
     internal sealed class VectorChatHistoryProvider : ChatHistoryProvider
     {
-        private const string DefaultStateBagKey = "VectorChatHistoryProvider.State";
-
         private readonly VectorStore _vectorStore;
         private readonly Func<AgentSession?, State> _stateInitializer;
         private readonly string _stateKey;
 
+        /// <inheritdoc />
+        public override string StateKey => this._stateKey;
+
         public VectorChatHistoryProvider(
             VectorStore vectorStore,
             Func<AgentSession?, State>? stateInitializer = null,
@@ -88,7 +89,7 @@ namespace SampleApp
         {
             this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
             this._stateInitializer = stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N")));
-            this._stateKey = stateKey ?? DefaultStateBagKey;
+            this._stateKey = stateKey ?? base.StateKey;
         }
 
         public string GetSessionDbKey(AgentSession session)

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

@@ -79,13 +79,11 @@ namespace SampleApp
     /// </summary>
     internal sealed class TodoListAIContextProvider : AIContextProvider
     {
-        private const string StateKey = nameof(TodoListAIContextProvider);
-
         private static List<string> GetTodoItems(AgentSession? session)
-            => session?.StateBag.GetValue<List<string>>(StateKey) ?? new List<string>();
+            => session?.StateBag.GetValue<List<string>>(nameof(TodoListAIContextProvider)) ?? new List<string>();
 
         private static void SetTodoItems(AgentSession? session, List<string> items)
-            => session?.StateBag.SetValue(StateKey, items);
+            => session?.StateBag.SetValue(nameof(TodoListAIContextProvider), items);
 
         protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
         {

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

@@ -50,6 +50,16 @@ public abstract class AIContextProvider
         this._sourceName = sourceName;
     }
 
+    /// <summary>
+    /// Gets the key 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.
+    /// </remarks>
+    public virtual string StateKey => this.GetType().Name;
+
     /// <summary>
     /// Called at the start of agent invocation to provide additional context.
     /// </summary>

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

@@ -59,6 +59,16 @@ public abstract class ChatHistoryProvider
         this._sourceName = sourceName;
     }
 
+    /// <summary>
+    /// Gets the key 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.
+    /// </remarks>
+    public virtual string StateKey => this.GetType().Name;
+
     /// <summary>
     /// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation.
     /// </summary>

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

@@ -47,6 +47,9 @@ public sealed class ChatHistoryProviderMessageFilter : ChatHistoryProvider
         this._invokedMessagesFilter = invokedMessagesFilter;
     }
 
+    /// <inheritdoc />
+    public override string StateKey => this._innerProvider.StateKey;
+
     /// <inheritdoc />
     protected override async ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {

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

@@ -27,8 +27,6 @@ namespace Microsoft.Agents.AI;
 /// </remarks>
 public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
 {
-    private const string DefaultStateBagKey = "InMemoryChatHistoryProvider.State";
-
     private readonly string _stateKey;
     private readonly Func<AgentSession?, State> _stateInitializer;
     private readonly JsonSerializerOptions _jsonSerializerOptions;
@@ -45,10 +43,13 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
         this._stateInitializer = options?.StateInitializer ?? (_ => new State());
         this.ChatReducer = options?.ChatReducer;
         this.ReducerTriggerEvent = options?.ReducerTriggerEvent ?? InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval;
-        this._stateKey = options?.StateKey ?? DefaultStateBagKey;
+        this._stateKey = options?.StateKey ?? base.StateKey;
         this._jsonSerializerOptions = options?.JsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
     }
 
+    /// <inheritdoc />
+    public override string StateKey => this._stateKey;
+
     /// <summary>
     /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied.
     /// </summary>

dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs

@@ -21,8 +21,6 @@ namespace Microsoft.Agents.AI;
 [RequiresDynamicCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with NativeAOT.")]
 public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
 {
-    private const string DefaultStateBagKey = "CosmosChatHistoryProvider.State";
-
     private readonly CosmosClient _cosmosClient;
     private readonly Container _container;
     private readonly bool _ownsClient;
@@ -45,6 +43,9 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         return options;
     }
 
+    /// <inheritdoc />
+    public override string StateKey => this._stateKey;
+
     /// <summary>
     /// Gets or sets the maximum number of messages to return in a single query batch.
     /// Default is 100 for optimal performance.
@@ -105,7 +106,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         this._container = this._cosmosClient.GetContainer(databaseId, containerId);
         this._stateInitializer = Throw.IfNull(stateInitializer);
         this._ownsClient = ownsClient;
-        this._stateKey = stateKey ?? DefaultStateBagKey;
+        this._stateKey = stateKey ?? base.StateKey;
     }
 
     /// <summary>

dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs

@@ -25,7 +25,6 @@ namespace Microsoft.Agents.AI.Mem0;
 public sealed class Mem0Provider : AIContextProvider
 {
     private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:";
-    private const string DefaultStateBagKey = "Mem0Provider.State";
 
     private readonly string _contextPrompt;
     private readonly bool _enableSensitiveTelemetryData;
@@ -67,9 +66,12 @@ public sealed class Mem0Provider : AIContextProvider
 
         this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt;
         this._enableSensitiveTelemetryData = options?.EnableSensitiveTelemetryData ?? false;
-        this._stateKey = options?.StateKey ?? DefaultStateBagKey;
+        this._stateKey = options?.StateKey ?? base.StateKey;
     }
 
+    /// <inheritdoc />
+    public override string StateKey => this._stateKey;
+
     /// <summary>
     /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present.
     /// </summary>

dotnet/src/Microsoft.Agents.AI.Mem0/Mem0ProviderOptions.cs

@@ -22,6 +22,6 @@ public sealed class Mem0ProviderOptions
     /// <summary>
     /// Gets or sets the key used to store the provider state in the session's <see cref="AgentSessionStateBag"/>.
     /// </summary>
-    /// <value>Defaults to "Mem0Provider.State".</value>
+    /// <value>Defaults to the provider's type name.</value>
     public string? StateKey { get; set; }
 }

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

@@ -11,7 +11,6 @@ namespace Microsoft.Agents.AI.Workflows;
 
 internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
 {
-    private const string DefaultStateBagKey = "WorkflowChatHistoryProvider.State";
     private readonly JsonSerializerOptions _jsonSerializerOptions;
 
     /// <summary>
@@ -35,7 +34,7 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
 
     private StoreState GetOrInitializeState(AgentSession? session)
     {
-        if (session?.StateBag.TryGetValue<StoreState>(DefaultStateBagKey, out var state, this._jsonSerializerOptions) is true && state is not null)
+        if (session?.StateBag.TryGetValue<StoreState>(this.StateKey, out var state, this._jsonSerializerOptions) is true && state is not null)
         {
             return state;
         }
@@ -43,7 +42,7 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
         state = new();
         if (session is not null)
         {
-            session.StateBag.SetValue(DefaultStateBagKey, state, this._jsonSerializerOptions);
+            session.StateBag.SetValue(this.StateKey, state, this._jsonSerializerOptions);
         }
 
         return state;

dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs

@@ -40,7 +40,6 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
     private const int DefaultMaxResults = 3;
     private const string DefaultFunctionToolName = "Search";
     private const string DefaultFunctionToolDescription = "Allows searching for related previous chat history to help answer the user question.";
-    private const string DefaultStateBagKey = "ChatHistoryMemoryProvider.State";
 
 #pragma warning disable CA2213 // VectorStore is not owned by this class - caller is responsible for disposal
     private readonly VectorStore _vectorStore;
@@ -86,7 +85,7 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
         this._contextPrompt = options.ContextPrompt ?? DefaultContextPrompt;
         this._enableSensitiveTelemetryData = options.EnableSensitiveTelemetryData;
         this._searchTime = options.SearchTime;
-        this._stateKey = options.StateKey ?? DefaultStateBagKey;
+        this._stateKey = options.StateKey ?? base.StateKey;
         this._logger = loggerFactory?.CreateLogger<ChatHistoryMemoryProvider>();
         this._toolName = options.FunctionToolName ?? DefaultFunctionToolName;
         this._toolDescription = options.FunctionToolDescription ?? DefaultFunctionToolDescription;
@@ -113,6 +112,9 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
         this._collection = this._vectorStore.GetDynamicCollection(Throw.IfNullOrWhitespace(collectionName), definition);
     }
 
+    /// <inheritdoc />
+    public override string StateKey => this._stateKey;
+
     /// <summary>
     /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present.
     /// </summary>

dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderOptions.cs

@@ -48,7 +48,7 @@ public sealed class ChatHistoryMemoryProviderOptions
     /// Gets or sets the key used to store provider state in the <see cref="AgentSession.StateBag"/>.
     /// </summary>
     /// <value>
-    /// Defaults to "ChatHistoryMemoryProvider.State". Override this if you need multiple
+    /// Defaults to the provider's type name. Override this if you need multiple
     /// <see cref="ChatHistoryMemoryProvider"/> instances with separate state in the same session.
     /// </value>
     public string? StateKey { get; set; }

dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs

@@ -38,7 +38,6 @@ public sealed class TextSearchProvider : AIContextProvider
     private const string DefaultPluginSearchFunctionDescription = "Allows searching for additional information to help answer the user question.";
     private const string DefaultContextPrompt = "## Additional Context\nConsider the following information from source documents when responding to the user:";
     private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available.";
-    private const string DefaultStateBagKey = "TextSearchProvider.RecentMessagesText";
 
     private readonly Func<string, CancellationToken, Task<IEnumerable<TextSearchResult>>> _searchAsync;
     private readonly ILogger<TextSearchProvider>? _logger;
@@ -71,7 +70,7 @@ public sealed class TextSearchProvider : AIContextProvider
         this._searchTime = options?.SearchTime ?? TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke;
         this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt;
         this._citationsPrompt = options?.CitationsPrompt ?? DefaultCitationsPrompt;
-        this._stateKey = options?.StateKey ?? DefaultStateBagKey;
+        this._stateKey = options?.StateKey ?? base.StateKey;
         this._contextFormatter = options?.ContextFormatter;
 
         // Create the on-demand search tool (only used if behavior is OnDemandFunctionCalling)
@@ -84,6 +83,9 @@ public sealed class TextSearchProvider : AIContextProvider
         ];
     }
 
+    /// <inheritdoc />
+    public override string StateKey => this._stateKey;
+
     /// <inheritdoc />
     protected override async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {

dotnet/src/Microsoft.Agents.AI/TextSearchProviderOptions.cs

@@ -63,7 +63,7 @@ public sealed class TextSearchProviderOptions
     /// Gets or sets the key used to store provider state in the <see cref="AgentSession.StateBag"/>.
     /// </summary>
     /// <value>
-    /// Defaults to "TextSearchProvider.RecentMessagesText". Override this if you need multiple
+    /// Defaults to the provider's type name. Override this if you need multiple
     /// <see cref="TextSearchProvider"/> instances with separate state in the same session.
     /// </value>
     public string? StateKey { get; set; }

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

@@ -52,6 +52,20 @@ public sealed class ChatHistoryProviderMessageFilterTests
         Assert.NotNull(filter);
     }
 
+    [Fact]
+    public void StateKey_DelegatesToInnerProvider()
+    {
+        // Arrange
+        var innerProviderMock = new Mock<ChatHistoryProvider>();
+        innerProviderMock.Setup(p => p.StateKey).Returns("inner-state-key");
+
+        // Act
+        var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, x => x);
+
+        // Assert
+        Assert.Equal("inner-state-key", filter.StateKey);
+    }
+
     [Fact]
     public async Task InvokingAsync_WithNoOpFilters_ReturnsInnerProviderMessagesAsync()
     {

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

@@ -42,6 +42,26 @@ public class InMemoryChatHistoryProviderTests
         Assert.Equal(InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded, provider.ReducerTriggerEvent);
     }
 
+    [Fact]
+    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    {
+        // Arrange & Act
+        var provider = new InMemoryChatHistoryProvider();
+
+        // Assert
+        Assert.Equal("InMemoryChatHistoryProvider", provider.StateKey);
+    }
+
+    [Fact]
+    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    {
+        // Arrange & Act
+        var provider = new InMemoryChatHistoryProvider(new() { StateKey = "custom-key" });
+
+        // Assert
+        Assert.Equal("custom-key", provider.StateKey);
+    }
+
     [Fact]
     public async Task InvokedAsyncAddsMessagesAsync()
     {

dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs

@@ -148,6 +148,35 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
 
     #region Constructor Tests
 
+    [SkippableFact]
+    [Trait("Category", "CosmosDB")]
+    public void StateKey_ReturnsDefaultKey_WhenNoStateKeyProvided()
+    {
+        // Arrange & Act
+        this.SkipIfEmulatorNotAvailable();
+
+        using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId,
+            _ => new CosmosChatHistoryProvider.State("test-conversation"));
+
+        // Assert
+        Assert.Equal("CosmosChatHistoryProvider", provider.StateKey);
+    }
+
+    [SkippableFact]
+    [Trait("Category", "CosmosDB")]
+    public void StateKey_ReturnsCustomKey_WhenSetViaConstructor()
+    {
+        // Arrange & Act
+        this.SkipIfEmulatorNotAvailable();
+
+        using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId,
+            _ => new CosmosChatHistoryProvider.State("test-conversation"),
+            stateKey: "custom-key");
+
+        // Assert
+        Assert.Equal("custom-key", provider.StateKey);
+    }
+
     [SkippableFact]
     [Trait("Category", "CosmosDB")]
     public void Constructor_WithConnectionString_ShouldCreateInstance()

dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs

@@ -66,6 +66,29 @@ public sealed class Mem0ProviderTests : IDisposable
         Assert.Contains("stateInitializer", ex.Message);
     }
 
+    [Fact]
+    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    {
+        // Arrange & Act
+        var provider = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(new Mem0ProviderScope { ThreadId = "tid" }));
+
+        // Assert
+        Assert.Equal("Mem0Provider", provider.StateKey);
+    }
+
+    [Fact]
+    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    {
+        // Arrange & Act
+        var provider = new Mem0Provider(
+            this._httpClient,
+            _ => new Mem0Provider.State(new Mem0ProviderScope { ThreadId = "tid" }),
+            new Mem0ProviderOptions { StateKey = "custom-key" });
+
+        // Assert
+        Assert.Equal("custom-key", provider.StateKey);
+    }
+
     [Fact]
     public async Task InvokingAsync_PerformsSearch_AndReturnsContextMessageAsync()
     {

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

@@ -48,7 +48,7 @@ public class ChatClientAgentSessionTests
         var json = JsonSerializer.Deserialize("""
             {
                 "stateBag": {
-                    "InMemoryChatHistoryProvider.State": {
+                    "InMemoryChatHistoryProvider": {
                         "messages": [{"authorName": "testAuthor"}]
                     }
                 }
@@ -164,7 +164,7 @@ public class ChatClientAgentSessionTests
         // Messages should be stored in the stateBag
         Assert.True(json.TryGetProperty("stateBag", out var stateBagProperty));
         Assert.Equal(JsonValueKind.Object, stateBagProperty.ValueKind);
-        Assert.True(stateBagProperty.TryGetProperty("InMemoryChatHistoryProvider.State", out var providerStateProperty));
+        Assert.True(stateBagProperty.TryGetProperty("InMemoryChatHistoryProvider", out var providerStateProperty));
         Assert.Equal(JsonValueKind.Object, providerStateProperty.ValueKind);
         Assert.True(providerStateProperty.TryGetProperty("messages", out var messagesProperty));
         Assert.Equal(JsonValueKind.Array, messagesProperty.ValueKind);

dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs

@@ -38,6 +38,28 @@ public sealed class TextSearchProviderTests
             .Returns(true);
     }
 
+    [Fact]
+    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    {
+        // Arrange & Act
+        var provider = new TextSearchProvider((_, _) => Task.FromResult<IEnumerable<TextSearchProvider.TextSearchResult>>([]));
+
+        // Assert
+        Assert.Equal("TextSearchProvider", provider.StateKey);
+    }
+
+    [Fact]
+    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    {
+        // Arrange & Act
+        var provider = new TextSearchProvider(
+            (_, _) => Task.FromResult<IEnumerable<TextSearchProvider.TextSearchResult>>([]),
+            new TextSearchProviderOptions { StateKey = "custom-key" });
+
+        // Assert
+        Assert.Equal("custom-key", provider.StateKey);
+    }
+
     [Theory]
     [InlineData(null, null, true)]
     [InlineData("Custom context prompt", "Custom citations prompt", false)]
@@ -523,7 +545,7 @@ public sealed class TextSearchProviderTests
 
         // Assert - State should be in the session's StateBag
         var stateBagSerialized = session.StateBag.Serialize();
-        Assert.True(stateBagSerialized.TryGetProperty("TextSearchProvider.RecentMessagesText", out var stateProperty));
+        Assert.True(stateBagSerialized.TryGetProperty("TextSearchProvider", out var stateProperty));
         Assert.True(stateProperty.TryGetProperty("recentMessagesText", out var recentProperty));
         Assert.Equal(JsonValueKind.Array, recentProperty.ValueKind);
         var list = recentProperty.EnumerateArray().Select(e => e.GetString()).ToList();

dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs

@@ -55,6 +55,35 @@ public class ChatHistoryMemoryProviderTests
             .Returns(this._vectorStoreCollectionMock.Object);
     }
 
+    [Fact]
+    public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+    {
+        // Arrange & Act
+        var provider = new ChatHistoryMemoryProvider(
+            this._vectorStoreMock.Object,
+            TestCollectionName,
+            1,
+            _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }));
+
+        // Assert
+        Assert.Equal("ChatHistoryMemoryProvider", provider.StateKey);
+    }
+
+    [Fact]
+    public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+    {
+        // Arrange & Act
+        var provider = new ChatHistoryMemoryProvider(
+            this._vectorStoreMock.Object,
+            TestCollectionName,
+            1,
+            _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }),
+            new ChatHistoryMemoryProviderOptions { StateKey = "custom-key" });
+
+        // Assert
+        Assert.Equal("custom-key", provider.StateKey);
+    }
+
     [Fact]
     public void Constructor_Throws_ForNullVectorStore()
     {