.NET: [BREAKING] Add session statebag to use for state storage instead of inside providers (#3737)

* Add a StateBag to AgentSession and pass Agent and AgentSession to AIContextProvider and ChatHistoryProviders

* Convert all AIContextProviders to use the statebag

* Update InMemoryChatHistoryProvider to use StateBag

* Update Comsos and Workflow ChatHistoryProviders

* Update 3rd party chat history storage sample.

* Remove serialize method from providers

* Replacing provider factories with properties

* Remove Providers from Session and flatten state bag serialization

* Update samples to use getservice on agent

* Updated additional session types to serialize statebag

* Fix regression

* Address PR comments

* Address PR comments.

* Fix formatting

* Fix unit tests

* Remove InMemoryAgentSession since it is not required anymore.

* Address PR comments

* Convert sessions for A2AAgent, ChatClientAgent, CopilotStudioAgent and GithubCopilotAgent to use regular json serialization.

* Fix durable agent session jso usgae

* Add jso to InMemory and Workflow ChatHistoryProviders

* Update InMemoryChatHistoryProvider to use an options class for it's many optional settings.

* Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Address PR feedback

* Fix verification bug.

* Improve state bag thread safety

* Address PR comments and fix unit tests

* Address PR comments

* Fix unit test

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

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

@@ -6,6 +6,7 @@
 
 using System.Runtime.CompilerServices;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 using Microsoft.Agents.AI;
 using Microsoft.Extensions.AI;
 using SampleApp;
@@ -28,6 +29,8 @@ namespace SampleApp
     {
         public override string? Name => "UpperCaseParrotAgent";
 
+        public readonly ChatHistoryProvider ChatHistoryProvider = new InMemoryChatHistoryProvider();
+
         protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default)
             => new(new CustomAgentSession());
 
@@ -38,11 +41,11 @@ namespace SampleApp
                 throw new ArgumentException($"The provided session is not of type {nameof(CustomAgentSession)}.", nameof(session));
             }
 
-            return typedSession.Serialize(jsonSerializerOptions);
+            return JsonSerializer.SerializeToElement(typedSession, jsonSerializerOptions);
         }
 
         protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-            => new(new CustomAgentSession(serializedState, jsonSerializerOptions));
+            => new(serializedState.Deserialize<CustomAgentSession>(jsonSerializerOptions)!);
 
         protected override async Task<AgentResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
         {
@@ -56,7 +59,7 @@ namespace SampleApp
 
             // Get existing messages from the store
             var invokingContext = new ChatHistoryProvider.InvokingContext(this, session, messages);
-            var storeMessages = await typedSession.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken);
+            var storeMessages = await this.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken);
 
             // Clone the input messages and turn them into response messages with upper case text.
             List<ChatMessage> responseMessages = CloneAndToUpperCase(messages, this.Name).ToList();
@@ -66,7 +69,7 @@ namespace SampleApp
             {
                 ResponseMessages = responseMessages
             };
-            await typedSession.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken);
+            await this.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken);
 
             return new AgentResponse
             {
@@ -88,7 +91,7 @@ namespace SampleApp
 
             // Get existing messages from the store
             var invokingContext = new ChatHistoryProvider.InvokingContext(this, session, messages);
-            var storeMessages = await typedSession.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken);
+            var storeMessages = await this.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken);
 
             // Clone the input messages and turn them into response messages with upper case text.
             List<ChatMessage> responseMessages = CloneAndToUpperCase(messages, this.Name).ToList();
@@ -98,7 +101,7 @@ namespace SampleApp
             {
                 ResponseMessages = responseMessages
             };
-            await typedSession.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken);
+            await this.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken);
 
             foreach (var message in responseMessages)
             {
@@ -140,15 +143,16 @@ namespace SampleApp
         /// <summary>
         /// A session type for our custom agent that only supports in memory storage of messages.
         /// </summary>
-        internal sealed class CustomAgentSession : InMemoryAgentSession
+        internal sealed class CustomAgentSession : AgentSession
         {
-            internal CustomAgentSession() { }
+            internal CustomAgentSession()
+            {
+            }
 
-            internal CustomAgentSession(JsonElement serializedSessionState, JsonSerializerOptions? jsonSerializerOptions = null)
-                : base(serializedSessionState, jsonSerializerOptions) { }
-
-            internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-                => base.Serialize(jsonSerializerOptions);
+            [JsonConstructor]
+            internal CustomAgentSession(AgentSessionStateBag stateBag) : base(stateBag)
+            {
+            }
         }
     }
 }

dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step01_ChatHistoryMemory/Program.cs

@@ -34,16 +34,21 @@ AIAgent agent = new AzureOpenAIClient(
     {
         ChatOptions = new() { Instructions = "You are good at telling jokes." },
         Name = "Joker",
-        AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new ChatHistoryMemoryProvider(
+        AIContextProvider = new ChatHistoryMemoryProvider(
             vectorStore,
             collectionName: "chathistory",
             vectorDimensions: 3072,
-            // Configure the scope values under which chat messages will be stored.
-            // In this case, we are using a fixed user ID and a unique session ID for each new session.
-            storageScope: new() { UserId = "UID1", SessionId = Guid.NewGuid().ToString() },
-            // Configure the scope which would be used to search for relevant prior messages.
-            // In this case, we are searching for any messages for the user across all sessions.
-            searchScope: new() { UserId = "UID1" }))
+            // Callback to configure the initial state of the ChatHistoryMemoryProvider.
+            // The ChatHistoryMemoryProvider stores its state in the AgentSession and this callback
+            // will be called whenever the ChatHistoryMemoryProvider cannot find existing state in the session,
+            // typically the first time it is used with a new session.
+            session => new ChatHistoryMemoryProvider.State(
+                // Configure the scope values under which chat messages will be stored.
+                // In this case, we are using a fixed user ID and a unique session ID for each new session.
+                storageScope: new() { UserId = "UID1", SessionId = Guid.NewGuid().ToString() },
+                // Configure the scope which would be used to search for relevant prior messages.
+                // In this case, we are searching for any messages for the user across all sessions.
+                searchScope: new() { UserId = "UID1" }))
     });
 
 // Start a new session for the agent conversation.

dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step02_MemoryUsingMem0/Program.cs

@@ -31,20 +31,21 @@ AIAgent agent = new AzureOpenAIClient(
     .AsAIAgent(new ChatClientAgentOptions()
     {
         ChatOptions = new() { Instructions = "You are a friendly travel assistant. Use known memories about the user when responding, and do not invent details." },
-        AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(ctx.SerializedState.ValueKind is not JsonValueKind.Null and not JsonValueKind.Undefined
-            // If each session should have its own Mem0 scope, you can create a new id per session here:
-            // ? new Mem0Provider(mem0HttpClient, new Mem0ProviderScope() { ThreadId = Guid.NewGuid().ToString() })
-            // In this case we are storing memories scoped by application and user instead so that memories are retained across threads.
-            ? new Mem0Provider(mem0HttpClient, new Mem0ProviderScope() { ApplicationId = "getting-started-agents", UserId = "sample-user" })
-            // For cases where we are restoring from serialized state:
-            : new Mem0Provider(mem0HttpClient, ctx.SerializedState, ctx.JsonSerializerOptions))
+        // The stateInitializer can be used to customize the Mem0 scope per session and it will be called each time a session
+        // is encountered by the Mem0Provider that does not already have Mem0Provider state stored on the session.
+        // If each session should have its own Mem0 scope, you can create a new id per session via the stateInitializer, e.g.:
+        // new Mem0Provider(mem0HttpClient, stateInitializer: _ => new(new Mem0ProviderScope() { ThreadId = Guid.NewGuid().ToString() }))
+        // In our case we are storing memories scoped by application and user instead so that memories are retained across threads.
+        AIContextProvider = new Mem0Provider(mem0HttpClient, stateInitializer: _ => new(new Mem0ProviderScope() { ApplicationId = "getting-started-agents", UserId = "sample-user" }))
     });
 
 AgentSession session = await agent.CreateSessionAsync();
 
 // Clear any existing memories for this scope to demonstrate fresh behavior.
-Mem0Provider mem0Provider = session.GetService<Mem0Provider>()!;
-await mem0Provider.ClearStoredMemoriesAsync();
+// Note that the ClearStoredMemoriesAsync method will clear memories
+// using the scope stored in the session, or provided via the stateInitializer.
+Mem0Provider mem0Provider = agent.GetService<Mem0Provider>()!;
+await mem0Provider.ClearStoredMemoriesAsync(session);
 
 Console.WriteLine(await agent.RunAsync("Hi there! My name is Taylor and I'm planning a hiking trip to Patagonia in November.", session));
 Console.WriteLine(await agent.RunAsync("I'm travelling with my sister and we love finding scenic viewpoints.", session));

dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs

@@ -33,7 +33,7 @@ ChatClient chatClient = new AzureOpenAIClient(
 AIAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions()
 {
     ChatOptions = new() { Instructions = "You are a friendly assistant. Always address the user by their name." },
-    AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new UserInfoMemory(chatClient.AsIChatClient(), ctx.SerializedState, ctx.JsonSerializerOptions))
+    AIContextProvider = new UserInfoMemory(chatClient.AsIChatClient())
 });
 
 // Create a new session for the conversation.
@@ -55,10 +55,10 @@ Console.WriteLine("\n>> Use deserialized session with previously created memorie
 var deserializedSession = await agent.DeserializeSessionAsync(sesionElement);
 Console.WriteLine(await agent.RunAsync("What is my name and age?", deserializedSession));
 
-Console.WriteLine("\n>> Read memories from memory component\n");
+Console.WriteLine("\n>> Read memories using memory component\n");
 
-// It's possible to access the memory component via the session's GetService method.
-var userInfo = deserializedSession.GetService<UserInfoMemory>()?.UserInfo;
+// It's possible to access the memory component via the agent's GetService method.
+var userInfo = agent.GetService<UserInfoMemory>()?.GetUserInfo(deserializedSession);
 
 // Output the user info that was captured by the memory component.
 Console.WriteLine($"MEMORY - User Name: {userInfo?.UserName}");
@@ -66,12 +66,12 @@ Console.WriteLine($"MEMORY - User Age: {userInfo?.UserAge}");
 
 Console.WriteLine("\n>> Use new session with previously created memories\n");
 
-// It is also possible to set the memories in a memory component on an individual session.
+// It is also possible to set the memories using a memory component on an individual session.
 // This is useful if we want to start a new session, but have it share the same memories as a previous session.
 var newSession = await agent.CreateSessionAsync();
-if (userInfo is not null && newSession.GetService<UserInfoMemory>() is UserInfoMemory newSessionMemory)
+if (userInfo is not null && agent.GetService<UserInfoMemory>() is UserInfoMemory newSessionMemory)
 {
-    newSessionMemory.UserInfo = userInfo;
+    newSessionMemory.SetUserInfo(newSession, userInfo);
 }
 
 // Invoke the agent and output the text result.
@@ -86,28 +86,27 @@ namespace SampleApp
     internal sealed class UserInfoMemory : AIContextProvider
     {
         private readonly IChatClient _chatClient;
+        private readonly Func<AgentSession?, UserInfo> _stateInitializer;
 
-        public UserInfoMemory(IChatClient chatClient, UserInfo? userInfo = null)
+        public UserInfoMemory(IChatClient chatClient, Func<AgentSession?, UserInfo>? stateInitializer = null)
         {
             this._chatClient = chatClient;
-            this.UserInfo = userInfo ?? new UserInfo();
+            this._stateInitializer = stateInitializer ?? (_ => new UserInfo());
         }
 
-        public UserInfoMemory(IChatClient chatClient, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
-        {
-            this._chatClient = chatClient;
+        public UserInfo GetUserInfo(AgentSession session)
+            => session.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory)) ?? new UserInfo();
 
-            this.UserInfo = serializedState.ValueKind == JsonValueKind.Object ?
-                serializedState.Deserialize<UserInfo>(jsonSerializerOptions)! :
-                new UserInfo();
-        }
-
-        public UserInfo UserInfo { get; set; }
+        public void SetUserInfo(AgentSession session, UserInfo userInfo)
+            => session.StateBag.SetValue(nameof(UserInfoMemory), userInfo);
 
         protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
         {
+            var userInfo = context.Session?.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory))
+                ?? this._stateInitializer.Invoke(context.Session);
+
             // Try and extract the user name and age from the message if we don't have it already and it's a user message.
-            if ((this.UserInfo.UserName is null || this.UserInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User))
+            if ((userInfo.UserName is null || userInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User))
             {
                 var result = await this._chatClient.GetResponseAsync<UserInfo>(
                     context.RequestMessages,
@@ -117,36 +116,36 @@ namespace SampleApp
                     },
                     cancellationToken: cancellationToken);
 
-                this.UserInfo.UserName ??= result.Result.UserName;
-                this.UserInfo.UserAge ??= result.Result.UserAge;
+                userInfo.UserName ??= result.Result.UserName;
+                userInfo.UserAge ??= result.Result.UserAge;
             }
+
+            context.Session?.StateBag.SetValue(nameof(UserInfoMemory), userInfo);
         }
 
         protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
         {
+            var userInfo = context.Session?.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory))
+                ?? this._stateInitializer.Invoke(context.Session);
+
             StringBuilder instructions = new();
 
             // If we don't already know the user's name and age, add instructions to ask for them, otherwise just provide what we have to the context.
             instructions
                 .AppendLine(
-                    this.UserInfo.UserName is null ?
+                    userInfo.UserName is null ?
                         "Ask the user for their name and politely decline to answer any questions until they provide it." :
-                        $"The user's name is {this.UserInfo.UserName}.")
+                        $"The user's name is {userInfo.UserName}.")
                 .AppendLine(
-                    this.UserInfo.UserAge is null ?
+                    userInfo.UserAge is null ?
                         "Ask the user for their age and politely decline to answer any questions until they provide it." :
-                        $"The user's age is {this.UserInfo.UserAge}.");
+                        $"The user's age is {userInfo.UserAge}.");
 
             return new ValueTask<AIContext>(new AIContext
             {
                 Instructions = instructions.ToString()
             });
         }
-
-        public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-        {
-            return JsonSerializer.SerializeToElement(this.UserInfo, jsonSerializerOptions);
-        }
     }
 
     internal sealed class UserInfo

dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs

@@ -62,12 +62,12 @@ AIAgent agent = azureOpenAIClient
     .AsAIAgent(new ChatClientAgentOptions
     {
         ChatOptions = new() { Instructions = "You are a helpful support specialist for Contoso Outdoors. Answer questions using the provided context and cite the source document when available." },
-        AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new TextSearchProvider(SearchAdapter, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions)),
+        AIContextProvider = new TextSearchProvider(SearchAdapter, textSearchOptions),
         // Since we are using ChatCompletion which stores chat history locally, we can also add a message removal policy
         // that removes messages produced by the TextSearchProvider before they are added to the chat history, so that
         // we don't bloat chat history with all the search result messages.
-        ChatHistoryProviderFactory = (ctx, ct) => new ValueTask<ChatHistoryProvider>(new InMemoryChatHistoryProvider(ctx.SerializedState, ctx.JsonSerializerOptions)
-            .WithAIContextProviderMessageRemoval()),
+        ChatHistoryProvider = new InMemoryChatHistoryProvider()
+            .WithAIContextProviderMessageRemoval(),
     });
 
 AgentSession session = await agent.CreateSessionAsync();

dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step02_CustomVectorStoreRAG/Program.cs

@@ -71,7 +71,7 @@ AIAgent agent = azureOpenAIClient
     .AsAIAgent(new ChatClientAgentOptions
     {
         ChatOptions = new() { Instructions = "You are a helpful support specialist for the Microsoft Agent Framework. Answer questions using the provided context and cite the source document when available. Keep responses brief." },
-        AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new TextSearchProvider(SearchAdapter, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions))
+        AIContextProvider = new TextSearchProvider(SearchAdapter, textSearchOptions)
     });
 
 AgentSession session = await agent.CreateSessionAsync();

dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step03_CustomRAGDataSource/Program.cs

@@ -29,7 +29,7 @@ AIAgent agent = new AzureOpenAIClient(
     .AsAIAgent(new ChatClientAgentOptions
     {
         ChatOptions = new() { Instructions = "You are a helpful support specialist for Contoso Outdoors. Answer questions using the provided context and cite the source document when available." },
-        AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new TextSearchProvider(MockSearchAsync, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions))
+        AIContextProvider = new TextSearchProvider(MockSearchAsync, textSearchOptions)
     });
 
 AgentSession session = await agent.CreateSessionAsync();

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

@@ -3,7 +3,7 @@
 #pragma warning disable CA1869 // Cache and reuse 'JsonSerializerOptions' instances
 
 // This sample shows how to create and use a simple AI agent with custom ChatHistoryProvider that stores chat history in a custom storage location.
-// The state of the custom ChatHistoryProvider (SessionDbKey) is stored with the agent session, so that when the session is resumed later,
+// The state of the custom ChatHistoryProvider (SessionDbKey) is stored in the AgentSession's StateBag, so that when the session is resumed later,
 // the chat history can be retrieved from the custom storage location.
 
 using System.Text.Json;
@@ -33,11 +33,8 @@ AIAgent agent = new AzureOpenAIClient(
     {
         ChatOptions = new() { Instructions = "You are good at telling jokes." },
         Name = "Joker",
-        ChatHistoryProviderFactory = (ctx, ct) => new ValueTask<ChatHistoryProvider>(
-            // Create a new ChatHistoryProvider for this agent that stores chat history in a vector store.
-            // Each session must get its own copy of the VectorChatHistoryProvider, since the provider
-            // also contains the id that the chat history is stored under.
-            new VectorChatHistoryProvider(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions))
+        // Create a new ChatHistoryProvider for this agent that stores chat history in a vector store.
+        ChatHistoryProvider = new VectorChatHistoryProvider(vectorStore)
     });
 
 // Start a new session for the agent conversation.
@@ -63,46 +60,70 @@ AgentSession resumedSession = await agent.DeserializeSessionAsync(serializedSess
 // Run the agent with the session that stores chat history in the vector store a second time.
 Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedSession));
 
-// We can access the VectorChatHistoryProvider via the session's GetService method if we need to read the key under which chat history is stored.
-var chatHistoryProvider = resumedSession.GetService<VectorChatHistoryProvider>()!;
-Console.WriteLine($"\nSession is stored in vector store under key: {chatHistoryProvider.SessionDbKey}");
+// We can access the VectorChatHistoryProvider via the agent's GetService method
+// if we need to read the key under which chat history is stored. The key is stored
+// in the session state, and therefore we need to provide the session when reading it.
+var chatHistoryProvider = agent.GetService<VectorChatHistoryProvider>()!;
+Console.WriteLine($"\nSession is stored in vector store under key: {chatHistoryProvider.GetSessionDbKey(resumedSession)}");
 
 namespace SampleApp
 {
     /// <summary>
     /// A sample implementation of <see cref="ChatHistoryProvider"/> that stores chat history in a vector store.
+    /// State (the session DB key) is stored in the <see cref="AgentSession.StateBag"/> so it roundtrips
+    /// automatically with session serialization.
     /// </summary>
     internal sealed class VectorChatHistoryProvider : ChatHistoryProvider
     {
-        private readonly VectorStore _vectorStore;
+        private const string DefaultStateBagKey = "VectorChatHistoryProvider.State";
 
-        public VectorChatHistoryProvider(VectorStore vectorStore, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
+        private readonly VectorStore _vectorStore;
+        private readonly Func<AgentSession?, State> _stateInitializer;
+        private readonly string _stateKey;
+
+        public VectorChatHistoryProvider(
+            VectorStore vectorStore,
+            Func<AgentSession?, State>? stateInitializer = null,
+            string? stateKey = null)
         {
             this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
-
-            if (serializedState.ValueKind is JsonValueKind.String)
-            {
-                // Here we can deserialize the session id so that we can access the same messages as before the suspension.
-                this.SessionDbKey = serializedState.Deserialize<string>();
-            }
+            this._stateInitializer = stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N")));
+            this._stateKey = stateKey ?? DefaultStateBagKey;
         }
 
-        public string? SessionDbKey { get; private set; }
+        public string GetSessionDbKey(AgentSession session)
+            => this.GetOrInitializeState(session).SessionDbKey;
+
+        private State GetOrInitializeState(AgentSession? session)
+        {
+            if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state) is true && state is not null)
+            {
+                return state;
+            }
+
+            state = this._stateInitializer(session);
+            if (session is not null)
+            {
+                session.StateBag.SetValue(this._stateKey, state);
+            }
+
+            return state;
+        }
 
         protected override async ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
         {
+            var state = this.GetOrInitializeState(context.Session);
             var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
             await collection.EnsureCollectionExistsAsync(cancellationToken);
 
             var records = await collection
                 .GetAsync(
-                    x => x.SessionId == this.SessionDbKey, 10,
+                    x => x.SessionId == state.SessionDbKey, 10,
                     new() { OrderBy = x => x.Descending(y => y.Timestamp) },
                     cancellationToken)
                 .ToListAsync(cancellationToken);
 
-            var messages = records.ConvertAll(x => JsonSerializer.Deserialize<ChatMessage>(x.SerializedMessage!)!)
-;
+            var messages = records.ConvertAll(x => JsonSerializer.Deserialize<ChatMessage>(x.SerializedMessage!)!);
             messages.Reverse();
             return messages;
         }
@@ -115,7 +136,7 @@ namespace SampleApp
                 return;
             }
 
-            this.SessionDbKey ??= Guid.NewGuid().ToString("N");
+            var state = this.GetOrInitializeState(context.Session);
 
             var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
             await collection.EnsureCollectionExistsAsync(cancellationToken);
@@ -126,17 +147,26 @@ namespace SampleApp
 
             await collection.UpsertAsync(allNewMessages.Select(x => new ChatHistoryItem()
             {
-                Key = this.SessionDbKey + x.MessageId,
+                Key = state.SessionDbKey + x.MessageId,
                 Timestamp = DateTimeOffset.UtcNow,
-                SessionId = this.SessionDbKey,
+                SessionId = state.SessionDbKey,
                 SerializedMessage = JsonSerializer.Serialize(x),
                 MessageText = x.Text
             }), cancellationToken);
         }
 
-        public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) =>
-            // We have to serialize the session id, so that on deserialization we can retrieve the messages using the same session id.
-            JsonSerializer.SerializeToElement(this.SessionDbKey);
+        /// <summary>
+        /// Represents the per-session state stored in the <see cref="AgentSession.StateBag"/>.
+        /// </summary>
+        public sealed class State
+        {
+            public State(string sessionDbKey)
+            {
+                this.SessionDbKey = sessionDbKey ?? throw new ArgumentNullException(nameof(sessionDbKey));
+            }
+
+            public string SessionDbKey { get; }
+        }
 
         /// <summary>
         /// The data structure used to store chat history items in the vector store.

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

@@ -24,7 +24,7 @@ AIAgent agent = new AzureOpenAIClient(
     {
         ChatOptions = new() { Instructions = "You are good at telling jokes." },
         Name = "Joker",
-        ChatHistoryProviderFactory = (ctx, ct) => new ValueTask<ChatHistoryProvider>(new InMemoryChatHistoryProvider(new MessageCountingChatReducer(2), ctx.SerializedState, ctx.JsonSerializerOptions))
+        ChatHistoryProvider = new InMemoryChatHistoryProvider(new() { ChatReducer = new MessageCountingChatReducer(2) })
     });
 
 AgentSession session = await agent.CreateSessionAsync();
@@ -33,7 +33,10 @@ AgentSession session = await agent.CreateSessionAsync();
 Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", session));
 
 // Get the chat history to see how many messages are stored.
-IList<ChatMessage>? chatHistory = session.GetService<IList<ChatMessage>>();
+// We can use the ChatHistoryProvider, that is also used by the agent, to read the
+// chat history from the session state, and see how the reducer is affecting the stored messages.
+var provider = agent.GetService<InMemoryChatHistoryProvider>();
+List<ChatMessage>? chatHistory = provider?.GetMessages(session);
 Console.WriteLine($"\nChat history has {chatHistory?.Count} messages.\n");
 
 // Invoke the agent a few more times.

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

@@ -7,7 +7,6 @@
 
 #pragma warning disable CA1869 // Cache and reuse 'JsonSerializerOptions' instances
 
-using System.ComponentModel;
 using System.Text;
 using System.Text.Json;
 using Azure.AI.OpenAI;
@@ -45,16 +44,16 @@ AIAgent agent = new AzureOpenAIClient(
         You manage a TODO list for the user. When the user has completed one of the tasks it can be removed from the TODO list. Only provide the list of TODO items if asked.
         You remind users of upcoming calendar events when the user interacts with you.
         """ },
-        ChatHistoryProviderFactory = (ctx, ct) => new ValueTask<ChatHistoryProvider>(new InMemoryChatHistoryProvider()
+        ChatHistoryProvider = new InMemoryChatHistoryProvider()
             // Use WithAIContextProviderMessageRemoval, so that we don't store the messages from the AI context provider in the chat history.
             // You may want to store these messages, depending on their content and your requirements.
-            .WithAIContextProviderMessageRemoval()),
+            .WithAIContextProviderMessageRemoval(),
         // Add an AI context provider that maintains a todo list for the agent and one that provides upcoming calendar entries.
         // Wrap these in an AI context provider that aggregates the other two.
-        AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new AggregatingAIContextProvider([
-            AggregatingAIContextProvider.CreateFactory((jsonElement, jsonSerializerOptions) => new TodoListAIContextProvider(jsonElement, jsonSerializerOptions)),
-            AggregatingAIContextProvider.CreateFactory((_, _) => new CalendarSearchAIContextProvider(loadNextThreeCalendarEvents))
-        ], ctx.SerializedState, ctx.JsonSerializerOptions)),
+        AIContextProvider = new AggregatingAIContextProvider([
+            new TodoListAIContextProvider(),
+            new CalendarSearchAIContextProvider(loadNextThreeCalendarEvents)
+        ]),
     });
 
 // Invoke the agent and output the text result.
@@ -80,51 +79,62 @@ namespace SampleApp
     /// </summary>
     internal sealed class TodoListAIContextProvider : AIContextProvider
     {
-        private readonly List<string> _todoItems = new();
+        private const string StateKey = nameof(TodoListAIContextProvider);
 
-        public TodoListAIContextProvider(JsonElement jsonElement, JsonSerializerOptions? jsonSerializerOptions = null)
-        {
-            // Only try and restore the state if we got an array, since any other json would be invalid or undefined/null meaning
-            // it's the first time we are running.
-            if (jsonElement.ValueKind == JsonValueKind.Array)
-            {
-                this._todoItems = JsonSerializer.Deserialize<List<string>>(jsonElement.GetRawText(), jsonSerializerOptions) ?? new List<string>();
-            }
-        }
+        private static List<string> GetTodoItems(AgentSession? session)
+            => session?.StateBag.GetValue<List<string>>(StateKey) ?? new List<string>();
+
+        private static void SetTodoItems(AgentSession? session, List<string> items)
+            => session?.StateBag.SetValue(StateKey, items);
 
         protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
         {
+            var todoItems = GetTodoItems(context.Session);
+
             StringBuilder outputMessageBuilder = new();
             outputMessageBuilder.AppendLine("Your todo list contains the following items:");
 
-            if (this._todoItems.Count == 0)
+            if (todoItems.Count == 0)
             {
                 outputMessageBuilder.AppendLine("  (no items)");
             }
             else
             {
-                for (int i = 0; i < this._todoItems.Count; i++)
+                for (int i = 0; i < todoItems.Count; i++)
                 {
-                    outputMessageBuilder.AppendLine($"{i}. {this._todoItems[i]}");
+                    outputMessageBuilder.AppendLine($"{i}. {todoItems[i]}");
                 }
             }
 
             return new ValueTask<AIContext>(new AIContext
             {
-                Tools = [AIFunctionFactory.Create(this.AddTodoItem), AIFunctionFactory.Create(this.RemoveTodoItem)],
+                Tools =
+                [
+                    AIFunctionFactory.Create((string item) => AddTodoItem(context.Session, item), "AddTodoItem", "Adds an item to the todo list."),
+                    AIFunctionFactory.Create((int index) => RemoveTodoItem(context.Session, index), "RemoveTodoItem", "Removes an item from the todo list. Index is zero based.")
+                ],
                 Messages = [new MEAI.ChatMessage(ChatRole.User, outputMessageBuilder.ToString())]
             });
         }
 
-        [Description("Adds an item to the todo list. Index is zero based.")]
-        private void RemoveTodoItem(int index) =>
-            this._todoItems.RemoveAt(index);
+        private static void RemoveTodoItem(AgentSession? session, int index)
+        {
+            var items = GetTodoItems(session);
+            items.RemoveAt(index);
+            SetTodoItems(session, items);
+        }
 
-        private void AddTodoItem(string item) =>
-            this._todoItems.Add(string.IsNullOrWhiteSpace(item) ? throw new ArgumentException("Item must have a value") : item);
+        private static void AddTodoItem(AgentSession? session, string item)
+        {
+            if (string.IsNullOrWhiteSpace(item))
+            {
+                throw new ArgumentException("Item must have a value");
+            }
 
-        public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) =>
-            JsonSerializer.SerializeToElement(this._todoItems, jsonSerializerOptions);
+            var items = GetTodoItems(session);
+            items.Add(item);
+            SetTodoItems(session, items);
+        }
     }
 
     /// <summary>
@@ -155,28 +165,15 @@ namespace SampleApp
 
     /// <summary>
     /// An <see cref="AIContextProvider"/> which aggregates multiple AI context providers into one.
-    /// Serialized state for the different providers are stored under their type name.
     /// Tools and messages from all providers are combined, and instructions are concatenated.
     /// </summary>
     internal sealed class AggregatingAIContextProvider : AIContextProvider
     {
-        private readonly List<AIContextProvider> _providers = new();
+        private readonly List<AIContextProvider> _providers;
 
-        public AggregatingAIContextProvider(ProviderFactory[] providerFactories, JsonElement jsonElement, JsonSerializerOptions? jsonSerializerOptions)
+        public AggregatingAIContextProvider(List<AIContextProvider> providers)
         {
-            // We received a json object, so let's check if it has some previously serialized state that we can use.
-            if (jsonElement.ValueKind == JsonValueKind.Object)
-            {
-                this._providers = providerFactories
-                    .Select(factory => factory.FactoryMethod(jsonElement.TryGetProperty(factory.ProviderType.Name, out var prop) ? prop : default, jsonSerializerOptions))
-                    .ToList();
-                return;
-            }
-
-            // We didn't receive any valid json, so we can just construct fresh providers.
-            this._providers = providerFactories
-                .Select(factory => factory.FactoryMethod(default, jsonSerializerOptions))
-                .ToList();
+            this._providers = providers;
         }
 
         protected override async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
@@ -193,36 +190,5 @@ namespace SampleApp
                 Instructions = string.Join("\n", results.Select(r => r.Instructions).Where(s => !string.IsNullOrEmpty(s)))
             };
         }
-
-        public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-        {
-            Dictionary<string, JsonElement> elements = new();
-            foreach (var provider in this._providers)
-            {
-                JsonElement element = provider.Serialize(jsonSerializerOptions);
-
-                // Don't try to store state for any providers that aren't producing any.
-                if (element.ValueKind != JsonValueKind.Undefined && element.ValueKind != JsonValueKind.Null)
-                {
-                    elements[provider.GetType().Name] = element;
-                }
-            }
-
-            return JsonSerializer.SerializeToElement(elements, jsonSerializerOptions);
-        }
-
-        public static ProviderFactory CreateFactory<TProviderType>(Func<JsonElement, JsonSerializerOptions?, TProviderType> factoryMethod)
-            where TProviderType : AIContextProvider => new()
-            {
-                FactoryMethod = (jsonElement, jsonSerializerOptions) => factoryMethod(jsonElement, jsonSerializerOptions),
-                ProviderType = typeof(TProviderType)
-            };
-
-        public readonly struct ProviderFactory
-        {
-            public Func<JsonElement, JsonSerializerOptions?, AIContextProvider> FactoryMethod { get; init; }
-
-            public Type ProviderType { get; init; }
-        }
     }
 }