.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/src/Microsoft.Agents.AI.A2A/A2AAgent.cs

@@ -80,7 +80,7 @@ public sealed class A2AAgent : AIAgent
 
     /// <inheritdoc/>
     protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-        => new(new A2AAgentSession(serializedState, jsonSerializerOptions));
+        => new(A2AAgentSession.Deserialize(serializedState, jsonSerializerOptions));
 
     /// <inheritdoc/>
     protected override async Task<AgentResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)

dotnet/src/Microsoft.Agents.AI.A2A/A2AAgentSession.cs

@@ -1,66 +1,61 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Diagnostics;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 
 namespace Microsoft.Agents.AI.A2A;
 
 /// <summary>
 /// Session for A2A based agents.
 /// </summary>
+[DebuggerDisplay("{DebuggerDisplay,nq}")]
 public sealed class A2AAgentSession : AgentSession
 {
     internal A2AAgentSession()
     {
     }
 
-    internal A2AAgentSession(JsonElement serializedSessionState, JsonSerializerOptions? jsonSerializerOptions = null)
+    [JsonConstructor]
+    internal A2AAgentSession(string? contextId, string? taskId, AgentSessionStateBag? stateBag) : base(stateBag ?? new())
     {
-        if (serializedSessionState.ValueKind != JsonValueKind.Object)
-        {
-            throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedSessionState));
-        }
-
-        var state = serializedSessionState.Deserialize(
-            A2AJsonUtilities.DefaultOptions.GetTypeInfo(typeof(A2AAgentSessionState))) as A2AAgentSessionState;
-
-        if (state?.ContextId is string contextId)
-        {
-            this.ContextId = contextId;
-        }
-
-        if (state?.TaskId is string taskId)
-        {
-            this.TaskId = taskId;
-        }
+        this.ContextId = contextId;
+        this.TaskId = taskId;
     }
 
     /// <summary>
     /// Gets the ID for the current conversation with the A2A agent.
     /// </summary>
+    [JsonPropertyName("contextId")]
     public string? ContextId { get; internal set; }
 
     /// <summary>
     /// Gets the ID for the task the agent is currently working on.
     /// </summary>
+    [JsonPropertyName("taskId")]
     public string? TaskId { get; internal set; }
 
     /// <inheritdoc/>
     internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        var state = new A2AAgentSessionState
-        {
-            ContextId = this.ContextId,
-            TaskId = this.TaskId
-        };
-
-        return JsonSerializer.SerializeToElement(state, A2AJsonUtilities.DefaultOptions.GetTypeInfo(typeof(A2AAgentSessionState)));
+        var jso = jsonSerializerOptions ?? A2AJsonUtilities.DefaultOptions;
+        return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(A2AAgentSession)));
     }
 
-    internal sealed class A2AAgentSessionState
+    internal static A2AAgentSession Deserialize(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        public string? ContextId { get; set; }
+        if (serializedState.ValueKind != JsonValueKind.Object)
+        {
+            throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState));
+        }
 
-        public string? TaskId { get; set; }
+        var jso = jsonSerializerOptions ?? A2AJsonUtilities.DefaultOptions;
+        return serializedState.Deserialize(jso.GetTypeInfo(typeof(A2AAgentSession))) as A2AAgentSession
+            ?? new A2AAgentSession();
     }
+
+    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+    private string DebuggerDisplay =>
+        $"ContextId = {this.ContextId}, TaskId = {this.TaskId}, StateBag Count = {this.StateBag.Count}";
 }

dotnet/src/Microsoft.Agents.AI.A2A/A2AJsonUtilities.cs

@@ -74,7 +74,7 @@ public static partial class A2AJsonUtilities
         NumberHandling = JsonNumberHandling.AllowReadingFromString)]
 
     // A2A agent types
-    [JsonSerializable(typeof(A2AAgentSession.A2AAgentSessionState))]
+    [JsonSerializable(typeof(A2AAgentSession))]
     [ExcludeFromCodeCoverage]
     private sealed partial class JsonContext : JsonSerializerContext;
 }

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

@@ -3,7 +3,6 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -130,13 +129,18 @@ public abstract class AIContextProvider
     /// <para>
     /// Implementers can use the request and response messages in the provided <paramref name="context"/> to:
     /// <list type="bullet">
-    /// <item><description>Update internal state based on conversation outcomes</description></item>
+    /// <item><description>Update state based on conversation outcomes</description></item>
     /// <item><description>Extract and store memories or preferences from user messages</description></item>
     /// <item><description>Log or audit conversation details</description></item>
     /// <item><description>Perform cleanup or finalization tasks</description></item>
     /// </list>
     /// </para>
     /// <para>
+    /// The <see cref="AIContextProvider"/> is passed a reference to the <see cref="AgentSession"/> via <see cref="InvokingContext"/> and <see cref="InvokedContext"/>
+    /// allowing it to store state in the <see cref="AgentSession.StateBag"/>. Since an <see cref="AIContextProvider"/> is used with many different sessions, it should
+    /// not store any session-specific information within its own instance fields. Instead, any session-specific state should be stored in the associated <see cref="AgentSession.StateBag"/>.
+    /// </para>
+    /// <para>
     /// This method is called regardless of whether the invocation succeeded or failed.
     /// To check if the invocation was successful, inspect the <see cref="InvokedContext.InvokeException"/> property.
     /// </para>
@@ -168,18 +172,6 @@ public abstract class AIContextProvider
     protected virtual ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
         => default;
 
-    /// <summary>
-    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
-    /// </summary>
-    /// <param name="jsonSerializerOptions">The JSON serialization options to use for the serialization process.</param>
-    /// <returns>A <see cref="JsonElement"/> representation of the object's state, or a default <see cref="JsonElement"/> if the provider has no serializable state.</returns>
-    /// <remarks>
-    /// The default implementation returns a default <see cref="JsonElement"/>. Override this method if the provider
-    /// maintains state that should be preserved across sessions or distributed scenarios.
-    /// </remarks>
-    public virtual JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-        => default;
-
     /// <summary>Asks the <see cref="AIContextProvider"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
     /// <param name="serviceType">The type of object being requested.</param>
     /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>

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

@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System.Collections.Concurrent;
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Encodings.Web;
 using System.Text.Json;
@@ -80,9 +81,9 @@ public static partial class AgentAbstractionsJsonUtilities
     [JsonSerializable(typeof(AgentResponse[]))]
     [JsonSerializable(typeof(AgentResponseUpdate))]
     [JsonSerializable(typeof(AgentResponseUpdate[]))]
-    [JsonSerializable(typeof(ServiceIdAgentSession.ServiceIdAgentSessionState))]
-    [JsonSerializable(typeof(InMemoryAgentSession.InMemoryAgentSessionState))]
     [JsonSerializable(typeof(InMemoryChatHistoryProvider.State))]
+    [JsonSerializable(typeof(AgentSessionStateBag))]
+    [JsonSerializable(typeof(ConcurrentDictionary<string, AgentSessionStateBagValue>))]
 
     [ExcludeFromCodeCoverage]
     private sealed partial class JsonContext : JsonSerializerContext;

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

@@ -1,7 +1,9 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Diagnostics;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Agents.AI;
@@ -44,6 +46,7 @@ namespace Microsoft.Agents.AI;
 /// <seealso cref="AIAgent"/>
 /// <seealso cref="AIAgent.CreateSessionAsync(System.Threading.CancellationToken)"/>
 /// <seealso cref="AIAgent.DeserializeSessionAsync(JsonElement, JsonSerializerOptions?, System.Threading.CancellationToken)"/>
+[DebuggerDisplay("{DebuggerDisplay,nq}")]
 public abstract class AgentSession
 {
     /// <summary>
@@ -53,6 +56,20 @@ public abstract class AgentSession
     {
     }
 
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AgentSession"/> class.
+    /// </summary>
+    protected AgentSession(AgentSessionStateBag stateBag)
+    {
+        this.StateBag = Throw.IfNull(stateBag);
+    }
+
+    /// <summary>
+    /// Gets any arbitrary state associated with this session.
+    /// </summary>
+    [JsonPropertyName("stateBag")]
+    public AgentSessionStateBag StateBag { get; protected set; } = new();
+
     /// <summary>Asks the <see cref="AgentSession"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
     /// <param name="serviceType">The type of object being requested.</param>
     /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
@@ -82,4 +99,7 @@ public abstract class AgentSession
     /// </remarks>
     public TService? GetService<TService>(object? serviceKey = null)
         => this.GetService(typeof(TService), serviceKey) is TService service ? service : default;
+
+    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+    private string DebuggerDisplay => $"StateBag Count = {this.StateBag.Count}";
 }

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

@@ -0,0 +1,145 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Concurrent;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Provides a thread-safe key-value store for managing session-scoped state with support for type-safe access and JSON
+/// serialization options.
+/// </summary>
+/// <remarks>
+/// SessionState enables storing and retrieving objects associated with a session using string keys.
+/// Values can be accessed in a type-safe manner and are serialized or deserialized using configurable JSON serializer
+/// options. This class is designed for concurrent access and is safe to use across multiple threads.
+/// </remarks>
+[JsonConverter(typeof(AgentSessionStateBagJsonConverter))]
+public class AgentSessionStateBag
+{
+    private readonly ConcurrentDictionary<string, AgentSessionStateBagValue> _state;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AgentSessionStateBag"/> class.
+    /// </summary>
+    public AgentSessionStateBag()
+    {
+        this._state = new ConcurrentDictionary<string, AgentSessionStateBagValue>();
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AgentSessionStateBag"/> class.
+    /// </summary>
+    /// <param name="state">The initial state dictionary.</param>
+    internal AgentSessionStateBag(ConcurrentDictionary<string, AgentSessionStateBagValue>? state)
+    {
+        this._state = state ?? new ConcurrentDictionary<string, AgentSessionStateBagValue>();
+    }
+
+    /// <summary>
+    /// Gets the number of key-value pairs contained in the session state.
+    /// </summary>
+    public int Count => this._state.Count;
+
+    /// <summary>
+    /// Tries to get a value from the session state.
+    /// </summary>
+    /// <typeparam name="T">The type of the value to retrieve.</typeparam>
+    /// <param name="key">The key from which to retrieve the value.</param>
+    /// <param name="value">The value if found and convertible to the required type; otherwise, null.</param>
+    /// <param name="jsonSerializerOptions">The JSON serializer options to use for serializing/deserializing the value.</param>
+    /// <returns><see langword="true"/> if the value was successfully retrieved, <see langword="false"/> otherwise.</returns>
+    public bool TryGetValue<T>(string key, out T? value, JsonSerializerOptions? jsonSerializerOptions = null)
+        where T : class
+    {
+        _ = Throw.IfNullOrWhitespace(key);
+        var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
+
+        if (this._state.TryGetValue(key, out var stateValue))
+        {
+            return stateValue.TryReadDeserializedValue(out value, jso);
+        }
+
+        value = null;
+        return false;
+    }
+
+    /// <summary>
+    /// Gets a value from the session state.
+    /// </summary>
+    /// <typeparam name="T">The type of value to get.</typeparam>
+    /// <param name="key">The key from which to retrieve the value.</param>
+    /// <param name="jsonSerializerOptions">The JSON serializer options to use for serializing/deserialing the value.</param>
+    /// <returns>The retrieved value or null if not found.</returns>
+    /// <exception cref="InvalidOperationException">The value could not be deserialized into the required type.</exception>
+    public T? GetValue<T>(string key, JsonSerializerOptions? jsonSerializerOptions = null)
+        where T : class
+    {
+        _ = Throw.IfNullOrWhitespace(key);
+        var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
+
+        if (this._state.TryGetValue(key, out var stateValue))
+        {
+            return stateValue.ReadDeserializedValue<T>(jso);
+        }
+
+        return null;
+    }
+
+    /// <summary>
+    /// Sets a value in the session state.
+    /// </summary>
+    /// <typeparam name="T">The type of the value to set.</typeparam>
+    /// <param name="key">The key to store the value under.</param>
+    /// <param name="value">The value to set.</param>
+    /// <param name="jsonSerializerOptions">The JSON serializer options to use for serializing the value.</param>
+    public void SetValue<T>(string key, T? value, JsonSerializerOptions? jsonSerializerOptions = null)
+        where T : class
+    {
+        _ = Throw.IfNullOrWhitespace(key);
+        var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
+
+        var stateValue = this._state.GetOrAdd(key, _ =>
+            new AgentSessionStateBagValue(value, typeof(T), jso));
+
+        stateValue.SetDeserialized(value, typeof(T), jso);
+    }
+
+    /// <summary>
+    /// Tries to remove a value from the session state.
+    /// </summary>
+    /// <param name="key">The key of the value to remove.</param>
+    /// <returns><see langword="true"/> if the value was successfully removed; otherwise, <see langword="false"/>.</returns>
+    public bool TryRemoveValue(string key)
+        => this._state.TryRemove(Throw.IfNullOrWhitespace(key), out _);
+
+    /// <summary>
+    /// Serializes all session state values to a JSON object.
+    /// </summary>
+    /// <returns>A <see cref="JsonElement"/> representing the serialized session state.</returns>
+    /// <exception cref="InvalidOperationException">Thrown when a session state value is not properly initialized.</exception>
+    public JsonElement Serialize()
+    {
+        return JsonSerializer.SerializeToElement(this._state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ConcurrentDictionary<string, AgentSessionStateBagValue>)));
+    }
+
+    /// <summary>
+    /// Deserializes a JSON object into an <see cref="AgentSessionStateBag"/> instance.
+    /// </summary>
+    /// <param name="jsonElement">The element to deserialize.</param>
+    /// <returns>The deserialized <see cref="AgentSessionStateBag"/>.</returns>
+    public static AgentSessionStateBag Deserialize(JsonElement jsonElement)
+    {
+        if (jsonElement.ValueKind is JsonValueKind.Undefined or JsonValueKind.Null)
+        {
+            return new AgentSessionStateBag();
+        }
+
+        return new AgentSessionStateBag(
+            jsonElement.Deserialize(AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ConcurrentDictionary<string, AgentSessionStateBagValue>))) as ConcurrentDictionary<string, AgentSessionStateBagValue>
+            ?? new ConcurrentDictionary<string, AgentSessionStateBagValue>());
+    }
+}

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

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Custom JSON converter for <see cref="AgentSessionStateBag"/> that serializes and deserializes
+/// the internal dictionary contents rather than the container object's public properties.
+/// </summary>
+public sealed class AgentSessionStateBagJsonConverter : JsonConverter<AgentSessionStateBag>
+{
+    /// <inheritdoc/>
+    public override AgentSessionStateBag Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
+    {
+        var element = JsonElement.ParseValue(ref reader);
+        return AgentSessionStateBag.Deserialize(element);
+    }
+
+    /// <inheritdoc/>
+    public override void Write(Utf8JsonWriter writer, AgentSessionStateBag value, JsonSerializerOptions options)
+    {
+        var element = value.Serialize();
+        element.WriteTo(writer);
+    }
+}

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

@@ -0,0 +1,170 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Used to store a value in session state.
+/// </summary>
+[JsonConverter(typeof(AgentSessionStateBagValueJsonConverter))]
+internal class AgentSessionStateBagValue
+{
+    private readonly object _lock = new();
+    private DeserializedCache? _cache;
+    private JsonElement _jsonValue;
+
+    /// <summary>
+    /// Initializes a new instance of the SessionStateValue class with the specified value.
+    /// </summary>
+    /// <param name="jsonValue">The serialized value to associate with the session state.</param>
+    public AgentSessionStateBagValue(JsonElement jsonValue)
+    {
+        this.JsonValue = jsonValue;
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the SessionStateValue class with the specified value.
+    /// </summary>
+    /// <param name="deserializedValue">The value to associate with the session state. Can be any object, including null.</param>
+    /// <param name="valueType">The type of the value.</param>
+    /// <param name="jsonSerializerOptions">The JSON serializer options to use for serializing the value.</param>
+    public AgentSessionStateBagValue(object? deserializedValue, Type valueType, JsonSerializerOptions jsonSerializerOptions)
+    {
+        this._cache = new DeserializedCache(deserializedValue, valueType, jsonSerializerOptions);
+    }
+
+    /// <summary>
+    /// Gets or sets the value associated with this instance.
+    /// </summary>
+    public JsonElement JsonValue
+    {
+        get
+        {
+            lock (this._lock)
+            {
+                // We are assuming here that JsonValue will only be read when the object is being serialized,
+                // which means that we will only call SerializeToElement when serializing and therefore it's
+                // OK to serialize on each read if the cache is set.
+                if (this._cache is { } cache)
+                {
+                    this._jsonValue = JsonSerializer.SerializeToElement(cache.Value, cache.Options.GetTypeInfo(cache.ValueType));
+                }
+
+                return this._jsonValue;
+            }
+        }
+        set
+        {
+            lock (this._lock)
+            {
+                this._jsonValue = value;
+                this._cache = null;
+            }
+        }
+    }
+
+    /// <summary>
+    /// Tries to read the deserialized value of this session state value.
+    /// Returns false if the value could not be deserialized into the required type, or if the value is undefined.
+    /// Returns true and sets the out parameter to null if the value is null.
+    /// </summary>
+    public bool TryReadDeserializedValue<T>(out T? value, JsonSerializerOptions? jsonSerializerOptions = null)
+        where T : class
+    {
+        var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
+
+        lock (this._lock)
+        {
+            if (this._cache is { } cache)
+            {
+                value = cache.Value as T;
+                return true;
+            }
+
+            switch (this._jsonValue)
+            {
+                case JsonElement jsonElement when jsonElement.ValueKind == JsonValueKind.Undefined:
+                    value = null;
+                    return false;
+                case JsonElement jsonElement when jsonElement.ValueKind == JsonValueKind.Null:
+                    value = null;
+                    return true;
+                default:
+                    T? result = this._jsonValue.Deserialize(jso.GetTypeInfo(typeof(T))) as T;
+                    if (result is null)
+                    {
+                        value = null;
+                        return false;
+                    }
+
+                    this._cache = new DeserializedCache(result, typeof(T), jso);
+
+                    value = result;
+                    return true;
+            }
+        }
+    }
+
+    /// <summary>
+    /// Reads the deserialized value of this session state value, throwing an exception if the value could not be deserialized into the required type or is undefined.
+    /// </summary>
+    public T? ReadDeserializedValue<T>(JsonSerializerOptions? jsonSerializerOptions = null)
+        where T : class
+    {
+        var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
+
+        lock (this._lock)
+        {
+            if (this._cache is { } cache)
+            {
+                return cache.Value as T;
+            }
+
+            switch (this._jsonValue)
+            {
+                case JsonElement jsonElement when jsonElement.ValueKind == JsonValueKind.Null || jsonElement.ValueKind == JsonValueKind.Undefined:
+                    return null;
+                default:
+                    T? result = this._jsonValue.Deserialize(jso.GetTypeInfo(typeof(T))) as T;
+                    if (result is null)
+                    {
+                        throw new InvalidOperationException($"Failed to deserialize session state value to type {typeof(T).FullName}.");
+                    }
+
+                    this._cache = new DeserializedCache(result, typeof(T), jso);
+                    return result;
+            }
+        }
+    }
+
+    /// <summary>
+    /// Sets the deserialized value of this session state value, updating the cache accordingly.
+    /// This does not update the JsonValue directly; the JsonValue will be updated on the next read or when the object is serialized.
+    /// </summary>
+    public void SetDeserialized(object? deserializedValue, Type valueType, JsonSerializerOptions jsonSerializerOptions)
+    {
+        lock (this._lock)
+        {
+            this._cache = new DeserializedCache(deserializedValue, valueType, jsonSerializerOptions);
+        }
+    }
+
+    private readonly struct DeserializedCache
+    {
+        public DeserializedCache(object? value, Type valueType, JsonSerializerOptions options)
+        {
+            this.Value = value;
+            this.ValueType = valueType;
+            this.Options = options;
+        }
+
+        public object? Value { get; }
+
+        public Type ValueType { get; }
+
+        public JsonSerializerOptions Options { get; }
+    }
+}

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

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Custom JSON converter for <see cref="AgentSessionStateBagValue"/> that serializes and deserializes
+/// the <see cref="AgentSessionStateBagValue.JsonValue"/> directly rather than wrapping it in a container object.
+/// </summary>
+internal sealed class AgentSessionStateBagValueJsonConverter : JsonConverter<AgentSessionStateBagValue>
+{
+    /// <inheritdoc/>
+    public override AgentSessionStateBagValue Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
+    {
+        var element = JsonElement.ParseValue(ref reader);
+        return new AgentSessionStateBagValue(element);
+    }
+
+    /// <inheritdoc/>
+    public override void Write(Utf8JsonWriter writer, AgentSessionStateBagValue value, JsonSerializerOptions options)
+    {
+        value.JsonValue.WriteTo(writer);
+    }
+}

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

@@ -3,7 +3,6 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -27,10 +26,14 @@ namespace Microsoft.Agents.AI;
 /// <item><description>Storing chat messages with proper ordering and metadata preservation</description></item>
 /// <item><description>Retrieving messages in chronological order for agent context</description></item>
 /// <item><description>Managing storage limits through truncation, summarization, or other strategies</description></item>
-/// <item><description>Supporting serialization for thread persistence and migration</description></item>
 /// </list>
 /// </para>
 /// <para>
+/// The <see cref="ChatHistoryProvider"/> is passed a reference to the <see cref="AgentSession"/> via <see cref="InvokingContext"/> and <see cref="InvokedContext"/>
+/// allowing it to store state in the <see cref="AgentSession.StateBag"/>. Since a <see cref="ChatHistoryProvider"/> is used with many different sessions, it should
+/// not store any session-specific information within its own instance fields. Instead, any session-specific state should be stored in the associated <see cref="AgentSession.StateBag"/>.
+/// </para>
+/// <para>
 /// A <see cref="ChatHistoryProvider"/> is only relevant for scenarios where the underlying AI service that the agent is using
 /// does not use in-service chat history storage.
 /// </para>
@@ -80,10 +83,6 @@ public abstract class ChatHistoryProvider
     /// <item><description>Archiving old messages while keeping active conversation context</description></item>
     /// </list>
     /// </para>
-    /// <para>
-    /// Each <see cref="ChatHistoryProvider"/> instance should be associated with a single <see cref="AgentSession"/> to ensure proper message isolation
-    /// and context management.
-    /// </para>
     /// </remarks>
     public async ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {
@@ -198,13 +197,6 @@ public abstract class ChatHistoryProvider
     /// </remarks>
     protected abstract ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default);
 
-    /// <summary>
-    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
-    /// </summary>
-    /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
-    /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
-    public abstract JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null);
-
     /// <summary>Asks the <see cref="ChatHistoryProvider"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
     /// <param name="serviceType">The type of object being requested.</param>
     /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>

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

@@ -2,7 +2,6 @@
 
 using System;
 using System.Collections.Generic;
-using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -65,10 +64,4 @@ public sealed class ChatHistoryProviderMessageFilter : ChatHistoryProvider
 
         return this._innerProvider.InvokedAsync(context, cancellationToken);
     }
-
-    /// <inheritdoc />
-    public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-        return this._innerProvider.Serialize(jsonSerializerOptions);
-    }
 }

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

@@ -1,124 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Collections.Generic;
-using System.Diagnostics;
-using System.Text.Json;
-using Microsoft.Extensions.AI;
-
-namespace Microsoft.Agents.AI;
-
-/// <summary>
-/// Provides an abstract base class for an <see cref="AgentSession"/> that maintain all chat history in local memory.
-/// </summary>
-/// <remarks>
-/// <para>
-/// <see cref="InMemoryAgentSession"/> is designed for scenarios where chat history should be stored locally
-/// rather than in external services or databases. This approach provides high performance and simplicity while
-/// maintaining full control over the conversation data.
-/// </para>
-/// <para>
-/// In-memory threads do not persist conversation data across application restarts
-/// unless explicitly serialized and restored.
-/// </para>
-/// </remarks>
-[DebuggerDisplay("{DebuggerDisplay,nq}")]
-public abstract class InMemoryAgentSession : AgentSession
-{
-    /// <summary>
-    /// Initializes a new instance of the <see cref="InMemoryAgentSession"/> class.
-    /// </summary>
-    /// <param name="chatHistoryProvider">
-    /// An optional <see cref="InMemoryChatHistoryProvider"/> instance to use for storing chat messages.
-    /// If <see langword="null"/>, a new empty <see cref="InMemoryChatHistoryProvider"/> will be created.
-    /// </param>
-    /// <remarks>
-    /// This constructor allows sharing of <see cref="ChatHistoryProvider"/> between sessions or providing pre-configured
-    /// <see cref="ChatHistoryProvider"/> with specific reduction or processing logic.
-    /// </remarks>
-    protected InMemoryAgentSession(InMemoryChatHistoryProvider? chatHistoryProvider = null)
-    {
-        this.ChatHistoryProvider = chatHistoryProvider ?? [];
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="InMemoryAgentSession"/> class.
-    /// </summary>
-    /// <param name="messages">The initial messages to populate the conversation history.</param>
-    /// <exception cref="ArgumentNullException"><paramref name="messages"/> is <see langword="null"/>.</exception>
-    /// <remarks>
-    /// This constructor is useful for initializing sessions with existing conversation history or
-    /// for migrating conversations from other storage systems.
-    /// </remarks>
-    protected InMemoryAgentSession(IEnumerable<ChatMessage> messages)
-    {
-        this.ChatHistoryProvider = [.. messages];
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="InMemoryAgentSession"/> class from previously serialized state.
-    /// </summary>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized state of the session.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="chatHistoryProviderFactory">
-    /// Optional factory function to create the <see cref="InMemoryChatHistoryProvider"/> from its serialized state.
-    /// If not provided, a default factory will be used that creates a basic <see cref="InMemoryChatHistoryProvider"/>.
-    /// </param>
-    /// <exception cref="ArgumentException">The <paramref name="serializedState"/> is not a JSON object.</exception>
-    /// <exception cref="JsonException">The <paramref name="serializedState"/> is invalid or cannot be deserialized to the expected type.</exception>
-    /// <remarks>
-    /// This constructor enables restoration of in-memory threads from previously saved state, allowing
-    /// conversations to be resumed across application restarts or migrated between different instances.
-    /// </remarks>
-    protected InMemoryAgentSession(
-        JsonElement serializedState,
-        JsonSerializerOptions? jsonSerializerOptions = null,
-        Func<JsonElement, JsonSerializerOptions?, InMemoryChatHistoryProvider>? chatHistoryProviderFactory = null)
-    {
-        if (serializedState.ValueKind != JsonValueKind.Object)
-        {
-            throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState));
-        }
-
-        var state = serializedState.Deserialize(
-            AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(InMemoryAgentSessionState))) as InMemoryAgentSessionState;
-
-        this.ChatHistoryProvider =
-            chatHistoryProviderFactory?.Invoke(state?.ChatHistoryProviderState ?? default, jsonSerializerOptions) ??
-            new InMemoryChatHistoryProvider(state?.ChatHistoryProviderState ?? default, jsonSerializerOptions);
-    }
-
-    /// <summary>
-    /// Gets or sets the <see cref="InMemoryChatHistoryProvider"/> used by this thread.
-    /// </summary>
-    public InMemoryChatHistoryProvider ChatHistoryProvider { get; }
-
-    /// <summary>
-    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
-    /// </summary>
-    /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
-    /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
-    protected internal virtual JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-        var chatHistoryProviderState = this.ChatHistoryProvider.Serialize(jsonSerializerOptions);
-
-        var state = new InMemoryAgentSessionState
-        {
-            ChatHistoryProviderState = chatHistoryProviderState,
-        };
-
-        return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(InMemoryAgentSessionState)));
-    }
-
-    /// <inheritdoc/>
-    public override object? GetService(Type serviceType, object? serviceKey = null) =>
-        base.GetService(serviceType, serviceKey) ?? this.ChatHistoryProvider?.GetService(serviceType, serviceKey);
-
-    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
-    private string DebuggerDisplay => $"Count = {this.ChatHistoryProvider.Count}";
-
-    internal sealed class InMemoryAgentSessionState
-    {
-        public JsonElement? ChatHistoryProviderState { get; set; }
-    }
-}

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

@@ -1,11 +1,10 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
-using System.Collections;
 using System.Collections.Generic;
-using System.Diagnostics;
 using System.Linq;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -14,99 +13,40 @@ using Microsoft.Shared.Diagnostics;
 namespace Microsoft.Agents.AI;
 
 /// <summary>
-/// Provides an in-memory implementation of <see cref="ChatHistoryProvider"/> with support for message reduction and collection semantics.
+/// Provides an in-memory implementation of <see cref="ChatHistoryProvider"/> with support for message reduction.
 /// </summary>
 /// <remarks>
 /// <para>
-/// <see cref="InMemoryChatHistoryProvider"/> stores chat messages entirely in local memory, providing fast access and manipulation
-/// capabilities. It implements both <see cref="ChatHistoryProvider"/> for agent integration and <see cref="IList{ChatMessage}"/>
-/// for direct collection manipulation.
+/// <see cref="InMemoryChatHistoryProvider"/> stores chat messages in the <see cref="AgentSession.StateBag"/>,
+/// providing fast access and manipulation capabilities integrated with session state management.
 /// </para>
 /// <para>
 /// This <see cref="ChatHistoryProvider"/> maintains all messages in memory. For long-running conversations or high-volume scenarios, consider using
 /// message reduction strategies or alternative storage implementations.
 /// </para>
 /// </remarks>
-[DebuggerDisplay("Count = {Count}")]
-[DebuggerTypeProxy(typeof(DebugView))]
-public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider, IList<ChatMessage>, IReadOnlyList<ChatMessage>
+public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
 {
-    private List<ChatMessage> _messages;
+    private const string DefaultStateBagKey = "InMemoryChatHistoryProvider.State";
+
+    private readonly string _stateKey;
+    private readonly Func<AgentSession?, State> _stateInitializer;
+    private readonly JsonSerializerOptions _jsonSerializerOptions;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="InMemoryChatHistoryProvider"/> class.
     /// </summary>
-    /// <remarks>
-    /// This constructor creates a basic in-memory <see cref="ChatHistoryProvider"/> without message reduction capabilities.
-    /// Messages will be stored exactly as added without any automatic processing or reduction.
-    /// </remarks>
-    public InMemoryChatHistoryProvider()
-    {
-        this._messages = [];
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="InMemoryChatHistoryProvider"/> class from previously serialized state.
-    /// </summary>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized state of the provider.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <exception cref="ArgumentException">The <paramref name="serializedState"/> is not a valid JSON object or cannot be deserialized.</exception>
-    /// <remarks>
-    /// This constructor enables restoration of messages from previously saved state, allowing
-    /// conversation history to be preserved across application restarts or migrated between instances.
-    /// </remarks>
-    public InMemoryChatHistoryProvider(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
-        : this(null, serializedState, jsonSerializerOptions, ChatReducerTriggerEvent.BeforeMessagesRetrieval)
-    {
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="InMemoryChatHistoryProvider"/> class.
-    /// </summary>
-    /// <param name="chatReducer">
-    /// A <see cref="IChatReducer"/> instance used to process, reduce, or optimize chat messages.
-    /// This can be used to implement strategies like message summarization, truncation, or cleanup.
+    /// <param name="options">
+    /// Optional configuration options that control the provider's behavior, including state initialization,
+    /// message reduction, and serialization settings. If <see langword="null"/>, default settings will be used.
     /// </param>
-    /// <param name="reducerTriggerEvent">
-    /// Specifies when the message reducer should be invoked. The default is <see cref="ChatReducerTriggerEvent.BeforeMessagesRetrieval"/>,
-    /// which applies reduction logic when messages are retrieved for agent consumption.
-    /// </param>
-    /// <exception cref="ArgumentNullException"><paramref name="chatReducer"/> is <see langword="null"/>.</exception>
-    /// <remarks>
-    /// Message reducers enable automatic management of message storage by implementing strategies to
-    /// keep memory usage under control while preserving important conversation context.
-    /// </remarks>
-    public InMemoryChatHistoryProvider(IChatReducer chatReducer, ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent.BeforeMessagesRetrieval)
-        : this(chatReducer, default, null, reducerTriggerEvent)
+    public InMemoryChatHistoryProvider(InMemoryChatHistoryProviderOptions? options = null)
     {
-        Throw.IfNull(chatReducer);
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="InMemoryChatHistoryProvider"/> class, with an existing state from a serialized JSON element.
-    /// </summary>
-    /// <param name="chatReducer">An optional <see cref="IChatReducer"/> instance used to process or reduce chat messages. If null, no reduction logic will be applied.</param>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized state of the provider.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="reducerTriggerEvent">The event that should trigger the reducer invocation.</param>
-    public InMemoryChatHistoryProvider(IChatReducer? chatReducer, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent.BeforeMessagesRetrieval)
-    {
-        this.ChatReducer = chatReducer;
-        this.ReducerTriggerEvent = reducerTriggerEvent;
-
-        if (serializedState.ValueKind is JsonValueKind.Object)
-        {
-            var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
-            var state = serializedState.Deserialize(
-                jso.GetTypeInfo(typeof(State))) as State;
-            if (state?.Messages is { } messages)
-            {
-                this._messages = messages;
-                return;
-            }
-        }
-
-        this._messages = [];
+        this._stateInitializer = options?.StateInitializer ?? (_ => new State());
+        this.ChatReducer = options?.ChatReducer;
+        this.ReducerTriggerEvent = options?.ReducerTriggerEvent ?? InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval;
+        this._stateKey = options?.StateKey ?? DefaultStateBagKey;
+        this._jsonSerializerOptions = options?.JsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
     }
 
     /// <summary>
@@ -117,19 +57,49 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider, IList<Cha
     /// <summary>
     /// Gets the event that triggers the reducer invocation in this provider.
     /// </summary>
-    public ChatReducerTriggerEvent ReducerTriggerEvent { get; }
+    public InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent ReducerTriggerEvent { get; }
 
-    /// <inheritdoc />
-    public int Count => this._messages.Count;
+    /// <summary>
+    /// Gets the chat messages stored for the specified session.
+    /// </summary>
+    /// <param name="session">The agent session containing the state.</param>
+    /// <returns>A list of chat messages, or an empty list if no state is found.</returns>
+    public List<ChatMessage> GetMessages(AgentSession? session)
+        => this.GetOrInitializeState(session).Messages;
 
-    /// <inheritdoc />
-    public bool IsReadOnly => ((IList)this._messages).IsReadOnly;
-
-    /// <inheritdoc />
-    public ChatMessage this[int index]
+    /// <summary>
+    /// Sets the chat messages for the specified session.
+    /// </summary>
+    /// <param name="session">The agent session containing the state.</param>
+    /// <param name="messages">The messages to store.</param>
+    /// <exception cref="ArgumentNullException"><paramref name="messages"/> is <see langword="null"/>.</exception>
+    public void SetMessages(AgentSession? session, List<ChatMessage> messages)
     {
-        get => this._messages[index];
-        set => this._messages[index] = value;
+        _ = Throw.IfNull(messages);
+
+        var state = this.GetOrInitializeState(session);
+        state.Messages = messages;
+    }
+
+    /// <summary>
+    /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present.
+    /// </summary>
+    /// <param name="session">The agent session containing the StateBag.</param>
+    /// <returns>The provider state, or null if no session is available.</returns>
+    private State GetOrInitializeState(AgentSession? session)
+    {
+        if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null)
+        {
+            return state;
+        }
+
+        state = this._stateInitializer(session);
+        if (session is not null)
+        {
+            session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions);
+        }
+
+        return state;
     }
 
     /// <inheritdoc />
@@ -137,12 +107,14 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider, IList<Cha
     {
         _ = Throw.IfNull(context);
 
-        if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null)
+        var state = this.GetOrInitializeState(context.Session);
+
+        if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null)
         {
-            this._messages = (await this.ChatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList();
+            state.Messages = (await this.ChatReducer.ReduceAsync(state.Messages, cancellationToken).ConfigureAwait(false)).ToList();
         }
 
-        return this._messages;
+        return state.Messages;
     }
 
     /// <inheritdoc />
@@ -155,94 +127,27 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider, IList<Cha
             return;
         }
 
+        var state = this.GetOrInitializeState(context.Session);
+
         // Add request and response messages to the provider
         var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []);
-        this._messages.AddRange(allNewMessages);
+        state.Messages.AddRange(allNewMessages);
 
-        if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null)
+        if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null)
         {
-            this._messages = (await this.ChatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList();
+            state.Messages = (await this.ChatReducer.ReduceAsync(state.Messages, cancellationToken).ConfigureAwait(false)).ToList();
         }
     }
 
-    /// <inheritdoc />
-    public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-        State state = new()
-        {
-            Messages = this._messages,
-        };
-
-        var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
-        return JsonSerializer.SerializeToElement(state, jso.GetTypeInfo(typeof(State)));
-    }
-
-    /// <inheritdoc />
-    public int IndexOf(ChatMessage item)
-        => this._messages.IndexOf(item);
-
-    /// <inheritdoc />
-    public void Insert(int index, ChatMessage item)
-        => this._messages.Insert(index, item);
-
-    /// <inheritdoc />
-    public void RemoveAt(int index)
-        => this._messages.RemoveAt(index);
-
-    /// <inheritdoc />
-    public void Add(ChatMessage item)
-        => this._messages.Add(item);
-
-    /// <inheritdoc />
-    public void Clear()
-        => this._messages.Clear();
-
-    /// <inheritdoc />
-    public bool Contains(ChatMessage item)
-        => this._messages.Contains(item);
-
-    /// <inheritdoc />
-    public void CopyTo(ChatMessage[] array, int arrayIndex)
-        => this._messages.CopyTo(array, arrayIndex);
-
-    /// <inheritdoc />
-    public bool Remove(ChatMessage item)
-        => this._messages.Remove(item);
-
-    /// <inheritdoc />
-    public IEnumerator<ChatMessage> GetEnumerator()
-        => this._messages.GetEnumerator();
-
-    /// <inheritdoc />
-    IEnumerator IEnumerable.GetEnumerator()
-        => this.GetEnumerator();
-
-    internal sealed class State
+    /// <summary>
+    /// Represents the state of a <see cref="InMemoryChatHistoryProvider"/> stored in the <see cref="AgentSession.StateBag"/>.
+    /// </summary>
+    public sealed class State
     {
+        /// <summary>
+        /// Gets or sets the list of chat messages.
+        /// </summary>
+        [JsonPropertyName("messages")]
         public List<ChatMessage> Messages { get; set; } = [];
     }
-
-    /// <summary>
-    /// Defines the events that can trigger a reducer in the <see cref="InMemoryChatHistoryProvider"/>.
-    /// </summary>
-    public enum ChatReducerTriggerEvent
-    {
-        /// <summary>
-        /// Trigger the reducer when a new message is added.
-        /// <see cref="InvokedCoreAsync(InvokedContext, CancellationToken)"/> will only complete when reducer processing is done.
-        /// </summary>
-        AfterMessageAdded,
-
-        /// <summary>
-        /// Trigger the reducer before messages are retrieved from the provider.
-        /// The reducer will process the messages before they are returned to the caller.
-        /// </summary>
-        BeforeMessagesRetrieval
-    }
-
-    private sealed class DebugView(InMemoryChatHistoryProvider provider)
-    {
-        [DebuggerBrowsable(DebuggerBrowsableState.RootHidden)]
-        public ChatMessage[] Items => provider._messages.ToArray();
-    }
 }

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

@@ -0,0 +1,67 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Text.Json;
+using Microsoft.Extensions.AI;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Represents configuration options for <see cref="InMemoryChatHistoryProvider"/>.
+/// </summary>
+public sealed class InMemoryChatHistoryProviderOptions
+{
+    /// <summary>
+    /// Gets or sets an optional delegate that initializes the provider state on the first invocation.
+    /// If <see langword="null"/>, a default initializer that creates an empty state will be used.
+    /// </summary>
+    public Func<AgentSession?, InMemoryChatHistoryProvider.State>? StateInitializer { get; set; }
+
+    /// <summary>
+    /// Gets or sets an optional <see cref="IChatReducer"/> instance used to process, reduce, or optimize chat messages.
+    /// This can be used to implement strategies like message summarization, truncation, or cleanup.
+    /// </summary>
+    public IChatReducer? ChatReducer { get; set; }
+
+    /// <summary>
+    /// Gets or sets when the message reducer should be invoked.
+    /// The default is <see cref="ChatReducerTriggerEvent.BeforeMessagesRetrieval"/>,
+    /// which applies reduction logic when messages are retrieved for agent consumption.
+    /// </summary>
+    /// <remarks>
+    /// Message reducers enable automatic management of message storage by implementing strategies to
+    /// keep memory usage under control while preserving important conversation context.
+    /// </remarks>
+    public ChatReducerTriggerEvent ReducerTriggerEvent { get; set; } = ChatReducerTriggerEvent.BeforeMessagesRetrieval;
+
+    /// <summary>
+    /// Gets or sets an optional key to use for storing the state in the <see cref="AgentSession.StateBag"/>.
+    /// If <see langword="null"/>, a default key will be used.
+    /// </summary>
+    public string? StateKey { get; set; }
+
+    /// <summary>
+    /// Gets or sets optional JSON serializer options for serializing the state of this provider.
+    /// This is valuable for cases like when the chat history contains custom <see cref="AIContent"/> types
+    /// and source generated serializers are required, or Native AOT / Trimming is required.
+    /// </summary>
+    public JsonSerializerOptions? JsonSerializerOptions { get; set; }
+
+    /// <summary>
+    /// Defines the events that can trigger a reducer in the <see cref="InMemoryChatHistoryProvider"/>.
+    /// </summary>
+    public enum ChatReducerTriggerEvent
+    {
+        /// <summary>
+        /// Trigger the reducer when a new message is added.
+        /// <see cref="AIContextProvider.InvokedAsync"/> will only complete when reducer processing is done.
+        /// </summary>
+        AfterMessageAdded,
+
+        /// <summary>
+        /// Trigger the reducer before messages are retrieved from the provider.
+        /// The reducer will process the messages before they are returned to the caller.
+        /// </summary>
+        BeforeMessagesRetrieval
+    }
+}

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

@@ -1,104 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System;
-using System.Diagnostics;
-using System.Text.Json;
-using Microsoft.Shared.Diagnostics;
-
-namespace Microsoft.Agents.AI;
-
-/// <summary>
-/// Provides a base class for agent sessions that store conversation state remotely in a service and maintain only an identifier reference locally.
-/// </summary>
-/// <remarks>
-/// This class is designed for scenarios where conversation state is managed by an external service (such as a cloud-based AI service)
-/// rather than being stored locally. The session maintains only the service identifier needed to reference the remote conversation state.
-/// </remarks>
-[DebuggerDisplay("ServiceSessionId = {ServiceSessionId}")]
-public abstract class ServiceIdAgentSession : AgentSession
-{
-    /// <summary>
-    /// Initializes a new instance of the <see cref="ServiceIdAgentSession"/> class without a service session identifier.
-    /// </summary>
-    /// <remarks>
-    /// When using this constructor, the <see cref="ServiceSessionId"/> will be <see langword="null"/> initially
-    /// and should be set by derived classes when the remote conversation is created.
-    /// </remarks>
-    protected ServiceIdAgentSession()
-    {
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="ServiceIdAgentSession"/> class with the specified service session identifier.
-    /// </summary>
-    /// <param name="serviceSessionId">The unique identifier that references the conversation state stored in the remote service.</param>
-    /// <exception cref="ArgumentNullException"><paramref name="serviceSessionId"/> is <see langword="null"/>.</exception>
-    /// <exception cref="ArgumentException"><paramref name="serviceSessionId"/> is empty or contains only whitespace.</exception>
-    protected ServiceIdAgentSession(string serviceSessionId)
-    {
-        this.ServiceSessionId = Throw.IfNullOrEmpty(serviceSessionId);
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="ServiceIdAgentSession"/> class from previously serialized state.
-    /// </summary>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized state of the session.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <exception cref="ArgumentException">The <paramref name="serializedState"/> is not a JSON object.</exception>
-    /// <exception cref="JsonException">The <paramref name="serializedState"/> is invalid or cannot be deserialized to the expected type.</exception>
-    /// <remarks>
-    /// This constructor enables restoration of a service-backed session from serialized state, typically used
-    /// when deserializing session information that was previously saved or transmitted across application boundaries.
-    /// </remarks>
-    protected ServiceIdAgentSession(
-        JsonElement serializedState,
-        JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-        if (serializedState.ValueKind != JsonValueKind.Object)
-        {
-            throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState));
-        }
-
-        var state = serializedState.Deserialize(
-            AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ServiceIdAgentSessionState))) as ServiceIdAgentSessionState;
-
-        if (state?.ServiceSessionId is string serviceSessionId)
-        {
-            this.ServiceSessionId = serviceSessionId;
-        }
-    }
-
-    /// <summary>
-    /// Gets or sets the unique identifier that references the conversation state stored in the remote service.
-    /// </summary>
-    /// <value>
-    /// A string identifier that uniquely identifies the conversation within the remote service,
-    /// or <see langword="null"/> if no remote conversation has been established yet.
-    /// </value>
-    /// <remarks>
-    /// This identifier is used by derived classes to reference the remote conversation state when making
-    /// API calls to the backing service. The exact format and meaning of this identifier depends on the
-    /// specific service implementation.
-    /// </remarks>
-    protected string? ServiceSessionId { get; set; }
-
-    /// <summary>
-    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
-    /// </summary>
-    /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
-    /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
-    protected internal virtual JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-        var state = new ServiceIdAgentSessionState
-        {
-            ServiceSessionId = this.ServiceSessionId,
-        };
-
-        return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ServiceIdAgentSessionState)));
-    }
-
-    internal sealed class ServiceIdAgentSessionState
-    {
-        public string? ServiceSessionId { get; set; }
-    }
-}

dotnet/src/Microsoft.Agents.AI.AzureAI.Persistent/PersistentAgentsClientExtensions.cs

@@ -191,8 +191,8 @@ public static class PersistentAgentsClientExtensions
             Name = options.Name ?? persistentAgentMetadata.Name,
             Description = options.Description ?? persistentAgentMetadata.Description,
             ChatOptions = options.ChatOptions,
-            AIContextProviderFactory = options.AIContextProviderFactory,
-            ChatHistoryProviderFactory = options.ChatHistoryProviderFactory,
+            AIContextProvider = options.AIContextProvider,
+            ChatHistoryProvider = options.ChatHistoryProvider,
             UseProvidedChatClientAsIs = options.UseProvidedChatClientAsIs
         };
 

dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs

@@ -589,8 +589,8 @@ public static partial class AzureAIProjectChatClientExtensions
         var agentOptions = CreateChatClientAgentOptions(agentVersion, options?.ChatOptions, requireInvocableTools);
         if (options is not null)
         {
-            agentOptions.AIContextProviderFactory = options.AIContextProviderFactory;
-            agentOptions.ChatHistoryProviderFactory = options.ChatHistoryProviderFactory;
+            agentOptions.AIContextProvider = options.AIContextProvider;
+            agentOptions.ChatHistoryProvider = options.ChatHistoryProvider;
             agentOptions.UseProvidedChatClientAsIs = options.UseProvidedChatClientAsIs;
         }
 

dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs

@@ -68,7 +68,7 @@ public class CopilotStudioAgent : AIAgent
 
     /// <inheritdoc/>
     protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-        => new(new CopilotStudioAgentSession(serializedState, jsonSerializerOptions));
+        => new(CopilotStudioAgentSession.Deserialize(serializedState, jsonSerializerOptions));
 
     /// <inheritdoc/>
     protected override async Task<AgentResponse> RunCoreAsync(

dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgentSession.cs

@@ -1,36 +1,58 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System;
+using System.Diagnostics;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 
 namespace Microsoft.Agents.AI.CopilotStudio;
 
 /// <summary>
 /// Session for CopilotStudio based agents.
 /// </summary>
-public sealed class CopilotStudioAgentSession : ServiceIdAgentSession
+[DebuggerDisplay("{DebuggerDisplay,nq}")]
+public sealed class CopilotStudioAgentSession : AgentSession
 {
     internal CopilotStudioAgentSession()
     {
     }
 
-    internal CopilotStudioAgentSession(JsonElement serializedSessionState, JsonSerializerOptions? jsonSerializerOptions = null) : base(serializedSessionState, jsonSerializerOptions)
+    [JsonConstructor]
+    internal CopilotStudioAgentSession(string? conversationId, AgentSessionStateBag? stateBag) : base(stateBag ?? new())
     {
+        this.ConversationId = conversationId;
     }
 
     /// <summary>
     /// Gets the ID for the current conversation with the Copilot Studio agent.
     /// </summary>
-    public string? ConversationId
-    {
-        get { return this.ServiceSessionId; }
-        internal set { this.ServiceSessionId = value; }
-    }
+    [JsonPropertyName("serviceSessionId")]
+    public string? ConversationId { get; internal set; }
 
     /// <summary>
     /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
     /// </summary>
     /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
     /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
-    internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-        => base.Serialize(jsonSerializerOptions);
+    internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
+    {
+        var jso = jsonSerializerOptions ?? CopilotStudioJsonUtilities.DefaultOptions;
+        return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(CopilotStudioAgentSession)));
+    }
+
+    internal static CopilotStudioAgentSession Deserialize(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
+    {
+        if (serializedState.ValueKind != JsonValueKind.Object)
+        {
+            throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState));
+        }
+
+        var jso = jsonSerializerOptions ?? CopilotStudioJsonUtilities.DefaultOptions;
+        return serializedState.Deserialize(jso.GetTypeInfo(typeof(CopilotStudioAgentSession))) as CopilotStudioAgentSession
+            ?? new CopilotStudioAgentSession();
+    }
+
+    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+    private string DebuggerDisplay =>
+        $"ConversationId = {this.ConversationId}, StateBag Count = {this.StateBag.Count}";
 }

dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioJsonUtilities.cs

@@ -0,0 +1,48 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Diagnostics.CodeAnalysis;
+using System.Text.Encodings.Web;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+namespace Microsoft.Agents.AI.CopilotStudio;
+
+/// <summary>
+/// Provides utility methods and configurations for JSON serialization operations within the Copilot Studio agent implementation.
+/// </summary>
+internal static partial class CopilotStudioJsonUtilities
+{
+    /// <summary>
+    /// Gets the default <see cref="JsonSerializerOptions"/> instance used for JSON serialization operations.
+    /// </summary>
+    public static JsonSerializerOptions DefaultOptions { get; } = CreateDefaultOptions();
+
+    /// <summary>
+    /// Creates and configures the default JSON serialization options.
+    /// </summary>
+    /// <returns>The configured options.</returns>
+    private static JsonSerializerOptions CreateDefaultOptions()
+    {
+        // Copy the configuration from the source generated context.
+        JsonSerializerOptions options = new(JsonContext.Default.Options)
+        {
+            Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping,
+        };
+
+        // Chain in the resolvers from both AgentAbstractionsJsonUtilities and our source generated context.
+        options.TypeInfoResolverChain.Clear();
+        options.TypeInfoResolverChain.Add(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver!);
+        options.TypeInfoResolverChain.Add(JsonContext.Default.Options.TypeInfoResolver!);
+
+        options.MakeReadOnly();
+        return options;
+    }
+
+    [JsonSourceGenerationOptions(JsonSerializerDefaults.Web,
+        UseStringEnumConverter = true,
+        DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
+        NumberHandling = JsonNumberHandling.AllowReadingFromString)]
+    [JsonSerializable(typeof(CopilotStudioAgentSession))]
+    [ExcludeFromCodeCoverage]
+    private sealed partial class JsonContext : JsonSerializerContext;
+}

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

@@ -21,17 +21,15 @@ 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;
+    private readonly string _stateKey;
+    private readonly Func<AgentSession?, State> _stateInitializer;
     private bool _disposed;
 
-    // Hierarchical partition key support
-    private readonly string? _tenantId;
-    private readonly string? _userId;
-    private readonly PartitionKey _partitionKey;
-    private readonly bool _useHierarchicalPartitioning;
-
     /// <summary>
     /// Cached JSON serializer options for .NET 9.0 compatibility.
     /// </summary>
@@ -72,11 +70,6 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// </summary>
     public int? MessageTtlSeconds { get; set; } = 86400;
 
-    /// <summary>
-    /// Gets the conversation ID associated with this provider.
-    /// </summary>
-    public string ConversationId { get; init; }
-
     /// <summary>
     /// Gets the database ID associated with this provider.
     /// </summary>
@@ -88,36 +81,31 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     public string ContainerId { get; init; }
 
     /// <summary>
-    /// Internal primary constructor used by all public constructors.
+    /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class.
     /// </summary>
     /// <param name="cosmosClient">The <see cref="CosmosClient"/> instance to use for Cosmos DB operations.</param>
     /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
     /// <param name="containerId">The identifier of the Cosmos DB container.</param>
-    /// <param name="conversationId">The unique identifier for this conversation thread.</param>
+    /// <param name="stateInitializer">A delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId).</param>
     /// <param name="ownsClient">Whether this instance owns the CosmosClient and should dispose it.</param>
-    /// <param name="tenantId">Optional tenant identifier for hierarchical partitioning.</param>
-    /// <param name="userId">Optional user identifier for hierarchical partitioning.</param>
-    internal CosmosChatHistoryProvider(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId, bool ownsClient, string? tenantId = null, string? userId = null)
+    /// <param name="stateKey">An optional key to use for storing the state in the <see cref="AgentSession.StateBag"/>.</param>
+    /// <exception cref="ArgumentNullException">Thrown when <paramref name="cosmosClient"/> or <paramref name="stateInitializer"/> is <see langword="null"/>.</exception>
+    /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
+    public CosmosChatHistoryProvider(
+        CosmosClient cosmosClient,
+        string databaseId,
+        string containerId,
+        Func<AgentSession?, State> stateInitializer,
+        bool ownsClient = false,
+        string? stateKey = null)
     {
         this._cosmosClient = Throw.IfNull(cosmosClient);
-        this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId));
-        this.ConversationId = Throw.IfNullOrWhitespace(conversationId);
-        this.DatabaseId = databaseId;
-        this.ContainerId = containerId;
+        this.DatabaseId = Throw.IfNullOrWhitespace(databaseId);
+        this.ContainerId = Throw.IfNullOrWhitespace(containerId);
+        this._container = this._cosmosClient.GetContainer(databaseId, containerId);
+        this._stateInitializer = Throw.IfNull(stateInitializer);
         this._ownsClient = ownsClient;
-
-        // Initialize partitioning mode
-        this._tenantId = tenantId;
-        this._userId = userId;
-        this._useHierarchicalPartitioning = tenantId != null && userId != null;
-
-        this._partitionKey = this._useHierarchicalPartitioning
-            ? new PartitionKeyBuilder()
-                .Add(tenantId!)
-                .Add(userId!)
-                .Add(conversationId)
-                .Build()
-            : new PartitionKey(conversationId);
+        this._stateKey = stateKey ?? DefaultStateBagKey;
     }
 
     /// <summary>
@@ -126,24 +114,17 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// <param name="connectionString">The Cosmos DB connection string.</param>
     /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
     /// <param name="containerId">The identifier of the Cosmos DB container.</param>
+    /// <param name="stateInitializer">A delegate that initializes the provider state on the first invocation.</param>
+    /// <param name="stateKey">An optional key to use for storing the state in the <see cref="AgentSession.StateBag"/>.</param>
     /// <exception cref="ArgumentNullException">Thrown when any required parameter is null.</exception>
     /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
-    public CosmosChatHistoryProvider(string connectionString, string databaseId, string containerId)
-        : this(connectionString, databaseId, containerId, Guid.NewGuid().ToString("N"))
-    {
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using a connection string.
-    /// </summary>
-    /// <param name="connectionString">The Cosmos DB connection string.</param>
-    /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
-    /// <param name="containerId">The identifier of the Cosmos DB container.</param>
-    /// <param name="conversationId">The unique identifier for this conversation thread.</param>
-    /// <exception cref="ArgumentNullException">Thrown when any required parameter is null.</exception>
-    /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
-    public CosmosChatHistoryProvider(string connectionString, string databaseId, string containerId, string conversationId)
-        : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, conversationId, ownsClient: true)
+    public CosmosChatHistoryProvider(
+        string connectionString,
+        string databaseId,
+        string containerId,
+        Func<AgentSession?, State> stateInitializer,
+        string? stateKey = null)
+        : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey)
     {
     }
 
@@ -154,136 +135,63 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// <param name="tokenCredential">The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).</param>
     /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
     /// <param name="containerId">The identifier of the Cosmos DB container.</param>
+    /// <param name="stateInitializer">A delegate that initializes the provider state on the first invocation.</param>
+    /// <param name="stateKey">An optional key to use for storing the state in the <see cref="AgentSession.StateBag"/>.</param>
     /// <exception cref="ArgumentNullException">Thrown when any required parameter is null.</exception>
     /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
-    public CosmosChatHistoryProvider(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId)
-        : this(accountEndpoint, tokenCredential, databaseId, containerId, Guid.NewGuid().ToString("N"))
+    public CosmosChatHistoryProvider(
+        string accountEndpoint,
+        TokenCredential tokenCredential,
+        string databaseId,
+        string containerId,
+        Func<AgentSession?, State> stateInitializer,
+        string? stateKey = null)
+        : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey)
     {
     }
 
     /// <summary>
-    /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using a TokenCredential for authentication.
+    /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present.
     /// </summary>
-    /// <param name="accountEndpoint">The Cosmos DB account endpoint URI.</param>
-    /// <param name="tokenCredential">The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).</param>
-    /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
-    /// <param name="containerId">The identifier of the Cosmos DB container.</param>
-    /// <param name="conversationId">The unique identifier for this conversation thread.</param>
-    /// <exception cref="ArgumentNullException">Thrown when any required parameter is null.</exception>
-    /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
-    public CosmosChatHistoryProvider(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string conversationId)
-        : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, conversationId, ownsClient: true)
+    /// <param name="session">The agent session containing the StateBag.</param>
+    /// <returns>The provider state, or null if no session is available.</returns>
+    private State GetOrInitializeState(AgentSession? session)
     {
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using an existing <see cref="CosmosClient"/>.
-    /// </summary>
-    /// <param name="cosmosClient">The <see cref="CosmosClient"/> instance to use for Cosmos DB operations.</param>
-    /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
-    /// <param name="containerId">The identifier of the Cosmos DB container.</param>
-    /// <exception cref="ArgumentNullException">Thrown when <paramref name="cosmosClient"/> is null.</exception>
-    /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
-    public CosmosChatHistoryProvider(CosmosClient cosmosClient, string databaseId, string containerId)
-        : this(cosmosClient, databaseId, containerId, Guid.NewGuid().ToString("N"))
-    {
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using an existing <see cref="CosmosClient"/>.
-    /// </summary>
-    /// <param name="cosmosClient">The <see cref="CosmosClient"/> instance to use for Cosmos DB operations.</param>
-    /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
-    /// <param name="containerId">The identifier of the Cosmos DB container.</param>
-    /// <param name="conversationId">The unique identifier for this conversation thread.</param>
-    /// <exception cref="ArgumentNullException">Thrown when <paramref name="cosmosClient"/> is null.</exception>
-    /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
-    public CosmosChatHistoryProvider(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId)
-        : this(cosmosClient, databaseId, containerId, conversationId, ownsClient: false)
-    {
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using a connection string with hierarchical partition keys.
-    /// </summary>
-    /// <param name="connectionString">The Cosmos DB connection string.</param>
-    /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
-    /// <param name="containerId">The identifier of the Cosmos DB container.</param>
-    /// <param name="tenantId">The tenant identifier for hierarchical partitioning.</param>
-    /// <param name="userId">The user identifier for hierarchical partitioning.</param>
-    /// <param name="sessionId">The session identifier for hierarchical partitioning.</param>
-    /// <exception cref="ArgumentNullException">Thrown when any required parameter is null.</exception>
-    /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
-    public CosmosChatHistoryProvider(string connectionString, string databaseId, string containerId, string tenantId, string userId, string sessionId)
-        : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId))
-    {
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using a TokenCredential for authentication with hierarchical partition keys.
-    /// </summary>
-    /// <param name="accountEndpoint">The Cosmos DB account endpoint URI.</param>
-    /// <param name="tokenCredential">The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).</param>
-    /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
-    /// <param name="containerId">The identifier of the Cosmos DB container.</param>
-    /// <param name="tenantId">The tenant identifier for hierarchical partitioning.</param>
-    /// <param name="userId">The user identifier for hierarchical partitioning.</param>
-    /// <param name="sessionId">The session identifier for hierarchical partitioning.</param>
-    /// <exception cref="ArgumentNullException">Thrown when any required parameter is null.</exception>
-    /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
-    public CosmosChatHistoryProvider(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string tenantId, string userId, string sessionId)
-        : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId))
-    {
-    }
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class using an existing <see cref="CosmosClient"/> with hierarchical partition keys.
-    /// </summary>
-    /// <param name="cosmosClient">The <see cref="CosmosClient"/> instance to use for Cosmos DB operations.</param>
-    /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
-    /// <param name="containerId">The identifier of the Cosmos DB container.</param>
-    /// <param name="tenantId">The tenant identifier for hierarchical partitioning.</param>
-    /// <param name="userId">The user identifier for hierarchical partitioning.</param>
-    /// <param name="sessionId">The session identifier for hierarchical partitioning.</param>
-    /// <exception cref="ArgumentNullException">Thrown when <paramref name="cosmosClient"/> is null.</exception>
-    /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
-    public CosmosChatHistoryProvider(CosmosClient cosmosClient, string databaseId, string containerId, string tenantId, string userId, string sessionId)
-        : this(cosmosClient, databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: false, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId))
-    {
-    }
-
-    /// <summary>
-    /// Creates a new instance of the <see cref="CosmosChatHistoryProvider"/> class from previously serialized state.
-    /// </summary>
-    /// <param name="cosmosClient">The <see cref="CosmosClient"/> instance to use for Cosmos DB operations.</param>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized state of the provider.</param>
-    /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
-    /// <param name="containerId">The identifier of the Cosmos DB container.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <returns>A new instance of <see cref="CosmosChatHistoryProvider"/> initialized from the serialized state.</returns>
-    /// <exception cref="ArgumentNullException">Thrown when <paramref name="cosmosClient"/> is null.</exception>
-    /// <exception cref="ArgumentException">Thrown when the serialized state cannot be deserialized.</exception>
-    public static CosmosChatHistoryProvider CreateFromSerializedState(CosmosClient cosmosClient, JsonElement serializedState, string databaseId, string containerId, JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-        Throw.IfNull(cosmosClient);
-        Throw.IfNullOrWhitespace(databaseId);
-        Throw.IfNullOrWhitespace(containerId);
-
-        if (serializedState.ValueKind is not JsonValueKind.Object)
+        if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state, AgentAbstractionsJsonUtilities.DefaultOptions) is true && state is not null)
         {
-            throw new ArgumentException("Invalid serialized state", nameof(serializedState));
+            return state;
         }
 
-        var state = serializedState.Deserialize<State>(jsonSerializerOptions);
-        if (state?.ConversationIdentifier is not { } conversationId)
+        state = this._stateInitializer(session);
+        if (session is not null)
         {
-            throw new ArgumentException("Invalid serialized state", nameof(serializedState));
+            session.StateBag.SetValue(this._stateKey, state, AgentAbstractionsJsonUtilities.DefaultOptions);
         }
 
-        // Use the internal constructor with all parameters to ensure partition key logic is centralized
-        return state.UseHierarchicalPartitioning && state.TenantId != null && state.UserId != null
-            ? new CosmosChatHistoryProvider(cosmosClient, databaseId, containerId, conversationId, ownsClient: false, state.TenantId, state.UserId)
-            : new CosmosChatHistoryProvider(cosmosClient, databaseId, containerId, conversationId, ownsClient: false);
+        return state;
+    }
+
+    /// <summary>
+    /// Determines whether hierarchical partitioning should be used based on the state.
+    /// </summary>
+    private static bool UseHierarchicalPartitioning(State state) =>
+        state.TenantId is not null && state.UserId is not null;
+
+    /// <summary>
+    /// Builds the partition key from the state.
+    /// </summary>
+    private static PartitionKey BuildPartitionKey(State state)
+    {
+        if (UseHierarchicalPartitioning(state))
+        {
+            return new PartitionKeyBuilder()
+                .Add(state.TenantId)
+                .Add(state.UserId)
+                .Add(state.ConversationId)
+                .Build();
+        }
+
+        return new PartitionKey(state.ConversationId);
     }
 
     /// <inheritdoc />
@@ -296,15 +204,20 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         }
 #pragma warning restore CA1513
 
+        _ = Throw.IfNull(context);
+
+        var state = this.GetOrInitializeState(context.Session);
+        var partitionKey = BuildPartitionKey(state);
+
         // Fetch most recent messages in descending order when limit is set, then reverse to ascending
         var orderDirection = this.MaxMessagesToRetrieve.HasValue ? "DESC" : "ASC";
         var query = new QueryDefinition($"SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp {orderDirection}")
-            .WithParameter("@conversationId", this.ConversationId)
+            .WithParameter("@conversationId", state.ConversationId)
             .WithParameter("@type", "ChatMessage");
 
         var iterator = this._container.GetItemQueryIterator<CosmosMessageDocument>(query, requestOptions: new QueryRequestOptions
         {
-            PartitionKey = this._partitionKey,
+            PartitionKey = partitionKey,
             MaxItemCount = this.MaxItemCount // Configurable query performance
         });
 
@@ -364,27 +277,30 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         }
 #pragma warning restore CA1513
 
+        var state = this.GetOrInitializeState(context.Session);
         var messageList = context.RequestMessages.Concat(context.ResponseMessages ?? []).ToList();
         if (messageList.Count == 0)
         {
             return;
         }
 
+        var partitionKey = BuildPartitionKey(state);
+
         // Use transactional batch for atomic operations
         if (messageList.Count > 1)
         {
-            await this.AddMessagesInBatchAsync(messageList, cancellationToken).ConfigureAwait(false);
+            await this.AddMessagesInBatchAsync(partitionKey, state, messageList, cancellationToken).ConfigureAwait(false);
         }
         else
         {
-            await this.AddSingleMessageAsync(messageList.First(), cancellationToken).ConfigureAwait(false);
+            await this.AddSingleMessageAsync(partitionKey, state, messageList.First(), cancellationToken).ConfigureAwait(false);
         }
     }
 
     /// <summary>
     /// Adds multiple messages using transactional batch operations for atomicity.
     /// </summary>
-    private async Task AddMessagesInBatchAsync(List<ChatMessage> messages, CancellationToken cancellationToken)
+    private async Task AddMessagesInBatchAsync(PartitionKey partitionKey, State state, List<ChatMessage> messages, CancellationToken cancellationToken)
     {
         var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds();
 
@@ -392,7 +308,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         for (int i = 0; i < messages.Count; i += this.MaxBatchSize)
         {
             var batchMessages = messages.Skip(i).Take(this.MaxBatchSize).ToList();
-            await this.ExecuteBatchOperationAsync(batchMessages, currentTimestamp, cancellationToken).ConfigureAwait(false);
+            await this.ExecuteBatchOperationAsync(partitionKey, state, batchMessages, currentTimestamp, cancellationToken).ConfigureAwait(false);
         }
     }
 
@@ -400,13 +316,13 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// Executes a single batch operation with enhanced error handling.
     /// Cosmos SDK handles throttling (429) retries automatically.
     /// </summary>
-    private async Task ExecuteBatchOperationAsync(List<ChatMessage> messages, long timestamp, CancellationToken cancellationToken)
+    private async Task ExecuteBatchOperationAsync(PartitionKey partitionKey, State state, List<ChatMessage> messages, long timestamp, CancellationToken cancellationToken)
     {
         // Create all documents upfront for validation and batch operation
         var documents = new List<CosmosMessageDocument>(messages.Count);
         foreach (var message in messages)
         {
-            documents.Add(this.CreateMessageDocument(message, timestamp));
+            documents.Add(this.CreateMessageDocument(state, message, timestamp));
         }
 
         // Defensive check: Verify all messages share the same partition key values
@@ -414,7 +330,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         // In simple partitioning, this means same conversationId
         if (documents.Count > 0)
         {
-            if (this._useHierarchicalPartitioning)
+            if (UseHierarchicalPartitioning(state))
             {
                 // Verify all documents have matching hierarchical partition key components
                 var firstDoc = documents[0];
@@ -436,7 +352,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
 
         // All messages in this store share the same partition key by design
         // Transactional batches require all items to share the same partition key
-        var batch = this._container.CreateTransactionalBatch(this._partitionKey);
+        var batch = this._container.CreateTransactionalBatch(partitionKey);
 
         foreach (var document in documents)
         {
@@ -457,7 +373,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
             if (messages.Count == 1)
             {
                 // Can't split further, use single operation
-                await this.AddSingleMessageAsync(messages[0], cancellationToken).ConfigureAwait(false);
+                await this.AddSingleMessageAsync(partitionKey, state, messages[0], cancellationToken).ConfigureAwait(false);
                 return;
             }
 
@@ -466,21 +382,21 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
             var firstHalf = messages.Take(midpoint).ToList();
             var secondHalf = messages.Skip(midpoint).ToList();
 
-            await this.ExecuteBatchOperationAsync(firstHalf, timestamp, cancellationToken).ConfigureAwait(false);
-            await this.ExecuteBatchOperationAsync(secondHalf, timestamp, cancellationToken).ConfigureAwait(false);
+            await this.ExecuteBatchOperationAsync(partitionKey, state, firstHalf, timestamp, cancellationToken).ConfigureAwait(false);
+            await this.ExecuteBatchOperationAsync(partitionKey, state, secondHalf, timestamp, cancellationToken).ConfigureAwait(false);
         }
     }
 
     /// <summary>
     /// Adds a single message to the store.
     /// </summary>
-    private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken cancellationToken)
+    private async Task AddSingleMessageAsync(PartitionKey partitionKey, State state, ChatMessage message, CancellationToken cancellationToken)
     {
-        var document = this.CreateMessageDocument(message, DateTimeOffset.UtcNow.ToUnixTimeSeconds());
+        var document = this.CreateMessageDocument(state, message, DateTimeOffset.UtcNow.ToUnixTimeSeconds());
 
         try
         {
-            await this._container.CreateItemAsync(document, this._partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false);
+            await this._container.CreateItemAsync(document, partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false);
         }
         catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge)
         {
@@ -495,12 +411,14 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// <summary>
     /// Creates a message document with enhanced metadata.
     /// </summary>
-    private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long timestamp)
+    private CosmosMessageDocument CreateMessageDocument(State state, ChatMessage message, long timestamp)
     {
+        var useHierarchical = UseHierarchicalPartitioning(state);
+
         return new CosmosMessageDocument
         {
             Id = Guid.NewGuid().ToString(),
-            ConversationId = this.ConversationId,
+            ConversationId = state.ConversationId,
             Timestamp = timestamp,
             MessageId = message.MessageId,
             Role = message.Role.Value,
@@ -508,41 +426,20 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
             Type = "ChatMessage", // Type discriminator
             Ttl = this.MessageTtlSeconds, // Configurable TTL
             // Include hierarchical metadata when using hierarchical partitioning
-            TenantId = this._useHierarchicalPartitioning ? this._tenantId : null,
-            UserId = this._useHierarchicalPartitioning ? this._userId : null,
-            SessionId = this._useHierarchicalPartitioning ? this.ConversationId : null
+            TenantId = useHierarchical ? state.TenantId : null,
+            UserId = useHierarchical ? state.UserId : null,
+            SessionId = useHierarchical ? state.ConversationId : null
         };
     }
 
-    /// <inheritdoc />
-    public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
-        if (this._disposed)
-        {
-            throw new ObjectDisposedException(this.GetType().FullName);
-        }
-#pragma warning restore CA1513
-
-        var state = new State
-        {
-            ConversationIdentifier = this.ConversationId,
-            TenantId = this._tenantId,
-            UserId = this._userId,
-            UseHierarchicalPartitioning = this._useHierarchicalPartitioning
-        };
-
-        var options = jsonSerializerOptions ?? s_defaultJsonOptions;
-        return JsonSerializer.SerializeToElement(state, options);
-    }
-
     /// <summary>
     /// Gets the count of messages in this conversation.
     /// This is an additional utility method beyond the base contract.
     /// </summary>
+    /// <param name="session">The agent session to get state from.</param>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>The number of messages in the conversation.</returns>
-    public async Task<int> GetMessageCountAsync(CancellationToken cancellationToken = default)
+    public async Task<int> GetMessageCountAsync(AgentSession? session, CancellationToken cancellationToken = default)
     {
 #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
         if (this._disposed)
@@ -551,14 +448,17 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         }
 #pragma warning restore CA1513
 
+        var state = this.GetOrInitializeState(session);
+        var partitionKey = BuildPartitionKey(state);
+
         // Efficient count query
         var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.Type = @type")
-            .WithParameter("@conversationId", this.ConversationId)
+            .WithParameter("@conversationId", state.ConversationId)
             .WithParameter("@type", "ChatMessage");
 
         var iterator = this._container.GetItemQueryIterator<int>(query, requestOptions: new QueryRequestOptions
         {
-            PartitionKey = this._partitionKey
+            PartitionKey = partitionKey
         });
 
         // COUNT queries always return a result
@@ -570,9 +470,10 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// Deletes all messages in this conversation.
     /// This is an additional utility method beyond the base contract.
     /// </summary>
+    /// <param name="session">The agent session to get state from.</param>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>The number of messages deleted.</returns>
-    public async Task<int> ClearMessagesAsync(CancellationToken cancellationToken = default)
+    public async Task<int> ClearMessagesAsync(AgentSession? session, CancellationToken cancellationToken = default)
     {
 #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
         if (this._disposed)
@@ -581,14 +482,17 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         }
 #pragma warning restore CA1513
 
+        var state = this.GetOrInitializeState(session);
+        var partitionKey = BuildPartitionKey(state);
+
         // Batch delete for efficiency
         var query = new QueryDefinition("SELECT VALUE c.id FROM c WHERE c.conversationId = @conversationId AND c.Type = @type")
-            .WithParameter("@conversationId", this.ConversationId)
+            .WithParameter("@conversationId", state.ConversationId)
             .WithParameter("@type", "ChatMessage");
 
         var iterator = this._container.GetItemQueryIterator<string>(query, requestOptions: new QueryRequestOptions
         {
-            PartitionKey = this._partitionKey,
+            PartitionKey = partitionKey,
             MaxItemCount = this.MaxItemCount
         });
 
@@ -597,7 +501,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         while (iterator.HasMoreResults)
         {
             var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false);
-            var batch = this._container.CreateTransactionalBatch(this._partitionKey);
+            var batch = this._container.CreateTransactionalBatch(partitionKey);
             var batchItemCount = 0;
 
             foreach (var itemId in response)
@@ -632,12 +536,38 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         }
     }
 
-    private sealed class State
+    /// <summary>
+    /// Represents the per-session state of a <see cref="CosmosChatHistoryProvider"/> stored in the <see cref="AgentSession.StateBag"/>.
+    /// </summary>
+    public sealed class State
     {
-        public string ConversationIdentifier { get; set; } = string.Empty;
-        public string? TenantId { get; set; }
-        public string? UserId { get; set; }
-        public bool UseHierarchicalPartitioning { get; set; }
+        /// <summary>
+        /// Initializes a new instance of the <see cref="State"/> class.
+        /// </summary>
+        /// <param name="conversationId">The unique identifier for this conversation thread.</param>
+        /// <param name="tenantId">Optional tenant identifier for hierarchical partitioning.</param>
+        /// <param name="userId">Optional user identifier for hierarchical partitioning.</param>
+        public State(string conversationId, string? tenantId = null, string? userId = null)
+        {
+            this.ConversationId = Throw.IfNullOrWhitespace(conversationId);
+            this.TenantId = tenantId;
+            this.UserId = userId;
+        }
+
+        /// <summary>
+        /// Gets the conversation ID associated with this state.
+        /// </summary>
+        public string ConversationId { get; }
+
+        /// <summary>
+        /// Gets the tenant identifier for hierarchical partitioning, if any.
+        /// </summary>
+        public string? TenantId { get; }
+
+        /// <summary>
+        /// Gets the user identifier for hierarchical partitioning, if any.
+        /// </summary>
+        public string? UserId { get; }
     }
 
     /// <summary>

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

@@ -2,7 +2,6 @@
 
 using System;
 using System.Diagnostics.CodeAnalysis;
-using System.Threading.Tasks;
 using Azure.Core;
 using Microsoft.Azure.Cosmos;
 
@@ -13,6 +12,9 @@ namespace Microsoft.Agents.AI;
 /// </summary>
 public static class CosmosDBChatExtensions
 {
+    private static readonly Func<AgentSession?, CosmosChatHistoryProvider.State> s_defaultStateInitializer =
+        _ => new CosmosChatHistoryProvider.State(Guid.NewGuid().ToString("N"));
+
     /// <summary>
     /// Configures the agent to use Cosmos DB for message storage with connection string authentication.
     /// </summary>
@@ -20,6 +22,7 @@ public static class CosmosDBChatExtensions
     /// <param name="connectionString">The Cosmos DB connection string.</param>
     /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
     /// <param name="containerId">The identifier of the Cosmos DB container.</param>
+    /// <param name="stateInitializer">An optional delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId). When not provided, a new conversation ID is generated automatically.</param>
     /// <returns>The configured <see cref="ChatClientAgentOptions"/>.</returns>
     /// <exception cref="ArgumentNullException">Thrown when <paramref name="options"/> is null.</exception>
     /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
@@ -29,14 +32,16 @@ public static class CosmosDBChatExtensions
         this ChatClientAgentOptions options,
         string connectionString,
         string databaseId,
-        string containerId)
+        string containerId,
+        Func<AgentSession?, CosmosChatHistoryProvider.State>? stateInitializer = null)
     {
         if (options is null)
         {
             throw new ArgumentNullException(nameof(options));
         }
 
-        options.ChatHistoryProviderFactory = (context, ct) => new ValueTask<ChatHistoryProvider>(new CosmosChatHistoryProvider(connectionString, databaseId, containerId));
+        options.ChatHistoryProvider =
+            new CosmosChatHistoryProvider(connectionString, databaseId, containerId, stateInitializer ?? s_defaultStateInitializer);
         return options;
     }
 
@@ -48,6 +53,7 @@ public static class CosmosDBChatExtensions
     /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
     /// <param name="containerId">The identifier of the Cosmos DB container.</param>
     /// <param name="tokenCredential">The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).</param>
+    /// <param name="stateInitializer">An optional delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId). When not provided, a new conversation ID is generated automatically.</param>
     /// <returns>The configured <see cref="ChatClientAgentOptions"/>.</returns>
     /// <exception cref="ArgumentNullException">Thrown when <paramref name="options"/> or <paramref name="tokenCredential"/> is null.</exception>
     /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
@@ -58,7 +64,8 @@ public static class CosmosDBChatExtensions
         string accountEndpoint,
         string databaseId,
         string containerId,
-        TokenCredential tokenCredential)
+        TokenCredential tokenCredential,
+        Func<AgentSession?, CosmosChatHistoryProvider.State>? stateInitializer = null)
     {
         if (options is null)
         {
@@ -70,7 +77,8 @@ public static class CosmosDBChatExtensions
             throw new ArgumentNullException(nameof(tokenCredential));
         }
 
-        options.ChatHistoryProviderFactory = (context, ct) => new ValueTask<ChatHistoryProvider>(new CosmosChatHistoryProvider(accountEndpoint, tokenCredential, databaseId, containerId));
+        options.ChatHistoryProvider =
+            new CosmosChatHistoryProvider(accountEndpoint, tokenCredential, databaseId, containerId, stateInitializer ?? s_defaultStateInitializer);
         return options;
     }
 
@@ -81,6 +89,7 @@ public static class CosmosDBChatExtensions
     /// <param name="cosmosClient">The <see cref="CosmosClient"/> instance to use for Cosmos DB operations.</param>
     /// <param name="databaseId">The identifier of the Cosmos DB database.</param>
     /// <param name="containerId">The identifier of the Cosmos DB container.</param>
+    /// <param name="stateInitializer">An optional delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId). When not provided, a new conversation ID is generated automatically.</param>
     /// <returns>The configured <see cref="ChatClientAgentOptions"/>.</returns>
     /// <exception cref="ArgumentNullException">Thrown when any required parameter is null.</exception>
     /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
@@ -90,14 +99,16 @@ public static class CosmosDBChatExtensions
         this ChatClientAgentOptions options,
         CosmosClient cosmosClient,
         string databaseId,
-        string containerId)
+        string containerId,
+        Func<AgentSession?, CosmosChatHistoryProvider.State>? stateInitializer = null)
     {
         if (options is null)
         {
             throw new ArgumentNullException(nameof(options));
         }
 
-        options.ChatHistoryProviderFactory = (context, ct) => new ValueTask<ChatHistoryProvider>(new CosmosChatHistoryProvider(cosmosClient, databaseId, containerId));
+        options.ChatHistoryProvider =
+            new CosmosChatHistoryProvider(cosmosClient, databaseId, containerId, stateInitializer ?? s_defaultStateInitializer);
         return options;
     }
 }

dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentSession.cs

@@ -7,17 +7,22 @@ using System.Text.Json.Serialization;
 namespace Microsoft.Agents.AI.DurableTask;
 
 /// <summary>
-/// An agent thread implementation for durable agents.
+/// An <see cref="AgentSession"/> implementation for durable agents.
 /// </summary>
-[DebuggerDisplay("{SessionId}")]
+[DebuggerDisplay("{DebuggerDisplay,nq}")]
 public sealed class DurableAgentSession : AgentSession
 {
-    [JsonConstructor]
     internal DurableAgentSession(AgentSessionId sessionId)
     {
         this.SessionId = sessionId;
     }
 
+    [JsonConstructor]
+    internal DurableAgentSession(AgentSessionId sessionId, AgentSessionStateBag stateBag) : base(stateBag)
+    {
+        this.SessionId = sessionId;
+    }
+
     /// <summary>
     /// Gets the agent session ID.
     /// </summary>
@@ -28,9 +33,8 @@ public sealed class DurableAgentSession : AgentSession
     /// <inheritdoc/>
     internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        return JsonSerializer.SerializeToElement(
-            this,
-            DurableAgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(DurableAgentSession)));
+        var jso = jsonSerializerOptions ?? DurableAgentJsonUtilities.DefaultOptions;
+        return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(DurableAgentSession)));
     }
 
     /// <summary>
@@ -49,7 +53,11 @@ public sealed class DurableAgentSession : AgentSession
 
         string sessionIdString = sessionIdElement.GetString() ?? throw new JsonException("sessionId property is null.");
         AgentSessionId sessionId = AgentSessionId.Parse(sessionIdString);
-        return new DurableAgentSession(sessionId);
+        AgentSessionStateBag stateBag = serializedSession.TryGetProperty("stateBag", out JsonElement stateBagElement)
+            ? AgentSessionStateBag.Deserialize(stateBagElement)
+            : new AgentSessionStateBag();
+
+        return new DurableAgentSession(sessionId, stateBag);
     }
 
     /// <inheritdoc/>
@@ -68,4 +76,8 @@ public sealed class DurableAgentSession : AgentSession
     {
         return this.SessionId.ToString();
     }
+
+    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+    private string DebuggerDisplay =>
+        $"SessionId = {this.SessionId}, StateBag Count = {this.StateBag.Count}";
 }

dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgent.cs

@@ -115,7 +115,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
         JsonElement serializedState,
         JsonSerializerOptions? jsonSerializerOptions = null,
         CancellationToken cancellationToken = default)
-        => new(new GitHubCopilotAgentSession(serializedState, jsonSerializerOptions));
+        => new(GitHubCopilotAgentSession.Deserialize(serializedState, jsonSerializerOptions));
 
     /// <inheritdoc/>
     protected override Task<AgentResponse> RunCoreAsync(

dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgentSession.cs

@@ -1,17 +1,22 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System;
+using System.Diagnostics;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 
 namespace Microsoft.Agents.AI.GitHub.Copilot;
 
 /// <summary>
 /// Represents a session for a GitHub Copilot agent conversation.
 /// </summary>
+[DebuggerDisplay("{DebuggerDisplay,nq}")]
 public sealed class GitHubCopilotAgentSession : AgentSession
 {
     /// <summary>
     /// Gets or sets the session ID for the GitHub Copilot conversation.
     /// </summary>
+    [JsonPropertyName("sessionId")]
     public string? SessionId { get; internal set; }
 
     /// <summary>
@@ -21,35 +26,32 @@ public sealed class GitHubCopilotAgentSession : AgentSession
     {
     }
 
-    /// <summary>
-    /// Initializes a new instance of the <see cref="GitHubCopilotAgentSession"/> class from serialized data.
-    /// </summary>
-    /// <param name="serializedThread">The serialized thread data.</param>
-    /// <param name="jsonSerializerOptions">Optional JSON serialization options.</param>
-    internal GitHubCopilotAgentSession(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    [JsonConstructor]
+    internal GitHubCopilotAgentSession(string? sessionId, AgentSessionStateBag? stateBag) : base(stateBag ?? new())
     {
-        // The JSON serialization uses camelCase
-        if (serializedThread.TryGetProperty("sessionId", out JsonElement sessionIdElement))
-        {
-            this.SessionId = sessionIdElement.GetString();
-        }
+        this.SessionId = sessionId;
     }
 
     /// <inheritdoc/>
     internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        State state = new()
-        {
-            SessionId = this.SessionId
-        };
-
-        return JsonSerializer.SerializeToElement(
-            state,
-            GitHubCopilotJsonUtilities.DefaultOptions.GetTypeInfo(typeof(State)));
+        var jso = jsonSerializerOptions ?? GitHubCopilotJsonUtilities.DefaultOptions;
+        return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(GitHubCopilotAgentSession)));
     }
 
-    internal sealed class State
+    internal static GitHubCopilotAgentSession Deserialize(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        public string? SessionId { get; set; }
+        if (serializedState.ValueKind != JsonValueKind.Object)
+        {
+            throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState));
+        }
+
+        var jso = jsonSerializerOptions ?? GitHubCopilotJsonUtilities.DefaultOptions;
+        return serializedState.Deserialize(jso.GetTypeInfo(typeof(GitHubCopilotAgentSession))) as GitHubCopilotAgentSession
+            ?? new GitHubCopilotAgentSession();
     }
+
+    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+    private string DebuggerDisplay =>
+        $"SessionId = {this.SessionId}, StateBag Count = {this.StateBag.Count}";
 }

dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotJsonUtilities.cs

@@ -42,7 +42,7 @@ internal static partial class GitHubCopilotJsonUtilities
         UseStringEnumConverter = true,
         DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
         NumberHandling = JsonNumberHandling.AllowReadingFromString)]
-    [JsonSerializable(typeof(GitHubCopilotAgentSession.State))]
+    [JsonSerializable(typeof(GitHubCopilotAgentSession))]
     [ExcludeFromCodeCoverage]
     private sealed partial class JsonContext : JsonSerializerContext;
 }

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

@@ -65,7 +65,7 @@ public static partial class Mem0JsonUtilities
         NumberHandling = JsonNumberHandling.AllowReadingFromString)]
 
     // Agent abstraction types
-    [JsonSerializable(typeof(Mem0Provider.Mem0State))]
+    [JsonSerializable(typeof(Mem0Provider.State))]
 
     [ExcludeFromCodeCoverage]
     internal sealed partial class JsonContext : JsonSerializerContext;

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

@@ -4,7 +4,6 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Net.Http;
-using System.Text.Json;
 using System.Text.Json.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
@@ -26,24 +25,24 @@ 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;
+    private readonly string _stateKey;
+    private readonly Func<AgentSession?, State> _stateInitializer;
 
     private readonly Mem0Client _client;
     private readonly ILogger<Mem0Provider>? _logger;
 
-    private readonly Mem0ProviderScope _storageScope;
-    private readonly Mem0ProviderScope _searchScope;
-
     /// <summary>
     /// Initializes a new instance of the <see cref="Mem0Provider"/> class.
     /// </summary>
     /// <param name="httpClient">Configured <see cref="HttpClient"/> (base address + auth).</param>
-    /// <param name="storageScope">Optional values to scope the memory storage with.</param>
-    /// <param name="searchScope">Optional values to scope the memory search with. Defaults to <paramref name="storageScope"/> if not provided.</param>
+    /// <param name="stateInitializer">A delegate that initializes the provider state on the first invocation, providing the storage and search scopes.</param>
     /// <param name="options">Provider options.</param>
     /// <param name="loggerFactory">Optional logger factory.</param>
+    /// <exception cref="ArgumentNullException">Thrown when <paramref name="httpClient"/> or <paramref name="stateInitializer"/> is <see langword="null"/>.</exception>
     /// <remarks>
     /// The base address of the required mem0 service, and any authentication headers, should be set on the <paramref name="httpClient"/>
     /// already, when passed as a parameter here. E.g.:
@@ -51,83 +50,55 @@ public sealed class Mem0Provider : AIContextProvider
     /// using var httpClient = new HttpClient();
     /// httpClient.BaseAddress = new Uri("https://api.mem0.ai");
     /// httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Token", "&lt;Your APIKey&gt;");
-    /// new Mem0AIContextProvider(httpClient);
+    /// new Mem0Provider(httpClient);
     /// </code>
     /// </remarks>
-    public Mem0Provider(HttpClient httpClient, Mem0ProviderScope storageScope, Mem0ProviderScope? searchScope = null, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null)
+    public Mem0Provider(HttpClient httpClient, Func<AgentSession?, State> stateInitializer, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null)
     {
+        Throw.IfNull(httpClient);
         if (string.IsNullOrWhiteSpace(httpClient.BaseAddress?.AbsoluteUri))
         {
             throw new ArgumentException("The HttpClient BaseAddress must be set for Mem0 operations.", nameof(httpClient));
         }
 
+        this._stateInitializer = Throw.IfNull(stateInitializer);
         this._logger = loggerFactory?.CreateLogger<Mem0Provider>();
         this._client = new Mem0Client(httpClient);
 
         this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt;
         this._enableSensitiveTelemetryData = options?.EnableSensitiveTelemetryData ?? false;
-        this._storageScope = new Mem0ProviderScope(Throw.IfNull(storageScope));
-        this._searchScope = searchScope ?? storageScope;
-
-        if (string.IsNullOrWhiteSpace(this._storageScope.ApplicationId)
-            && string.IsNullOrWhiteSpace(this._storageScope.AgentId)
-            && string.IsNullOrWhiteSpace(this._storageScope.ThreadId)
-            && string.IsNullOrWhiteSpace(this._storageScope.UserId))
-        {
-            throw new ArgumentException("At least one of ApplicationId, AgentId, ThreadId, or UserId must be provided for the storage scope.");
-        }
-
-        if (string.IsNullOrWhiteSpace(this._searchScope.ApplicationId)
-            && string.IsNullOrWhiteSpace(this._searchScope.AgentId)
-            && string.IsNullOrWhiteSpace(this._searchScope.ThreadId)
-            && string.IsNullOrWhiteSpace(this._searchScope.UserId))
-        {
-            throw new ArgumentException("At least one of ApplicationId, AgentId, ThreadId, or UserId must be provided for the search scope.");
-        }
+        this._stateKey = options?.StateKey ?? DefaultStateBagKey;
     }
 
     /// <summary>
-    /// Initializes a new instance of the <see cref="Mem0Provider"/> class, with existing state from a serialized JSON element.
+    /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present.
     /// </summary>
-    /// <param name="httpClient">Configured <see cref="HttpClient"/> (base address + auth).</param>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized state of the store.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="options">Provider options.</param>
-    /// <param name="loggerFactory">Optional logger factory.</param>
-    /// <exception cref="ArgumentException"></exception>
-    /// <remarks>
-    /// The base address of the required mem0 service, and any authentication headers, should be set on the <paramref name="httpClient"/>
-    /// already, when passed as a parameter here. E.g.:
-    /// <code>
-    /// using var httpClient = new HttpClient();
-    /// httpClient.BaseAddress = new Uri("https://api.mem0.ai");
-    /// httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Token", "&lt;Your APIKey&gt;");
-    /// new Mem0AIContextProvider(httpClient, state);
-    /// </code>
-    /// </remarks>
-    public Mem0Provider(HttpClient httpClient, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null)
+    /// <param name="session">The agent session containing the StateBag.</param>
+    /// <returns>The provider state, or null if no session is available.</returns>
+    private State? GetOrInitializeState(AgentSession? session)
     {
-        if (string.IsNullOrWhiteSpace(httpClient.BaseAddress?.AbsoluteUri))
+        if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state, Mem0JsonUtilities.DefaultOptions) is true && state is not null)
         {
-            throw new ArgumentException("The HttpClient BaseAddress must be set for Mem0 operations.", nameof(httpClient));
+            return state;
         }
 
-        this._logger = loggerFactory?.CreateLogger<Mem0Provider>();
-        this._client = new Mem0Client(httpClient);
+        state = this._stateInitializer(session);
 
-        this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt;
-        this._enableSensitiveTelemetryData = options?.EnableSensitiveTelemetryData ?? false;
-
-        var jso = jsonSerializerOptions ?? Mem0JsonUtilities.DefaultOptions;
-        var state = serializedState.Deserialize(jso.GetTypeInfo(typeof(Mem0State))) as Mem0State;
-
-        if (state == null || state.StorageScope == null || state.SearchScope == null)
+        if (state is null
+            || state.StorageScope is null
+            || (state.StorageScope.AgentId is null && state.StorageScope.ThreadId is null && state.StorageScope.UserId is null && state.StorageScope.ApplicationId is null)
+            || state.SearchScope is null
+            || (state.SearchScope.AgentId is null && state.SearchScope.ThreadId is null && state.SearchScope.UserId is null && state.SearchScope.ApplicationId is null))
         {
-            throw new InvalidOperationException("The Mem0Provider state did not contain the required scope properties.");
+            throw new InvalidOperationException("State initializer must return a non-null state with valid storage and search scopes, where at lest one scoping parameter is set for each.");
         }
 
-        this._storageScope = state.StorageScope;
-        this._searchScope = state.SearchScope;
+        if (session is not null)
+        {
+            session.StateBag.SetValue(this._stateKey, state, Mem0JsonUtilities.DefaultOptions);
+        }
+
+        return state;
     }
 
     /// <inheritdoc />
@@ -135,6 +106,9 @@ public sealed class Mem0Provider : AIContextProvider
     {
         Throw.IfNull(context);
 
+        var state = this.GetOrInitializeState(context.Session);
+        var searchScope = state?.SearchScope ?? new Mem0ProviderScope();
+
         string queryText = string.Join(
             Environment.NewLine,
             context.RequestMessages
@@ -145,10 +119,10 @@ public sealed class Mem0Provider : AIContextProvider
         try
         {
             var memories = (await this._client.SearchAsync(
-                this._searchScope.ApplicationId,
-                this._searchScope.AgentId,
-                this._searchScope.ThreadId,
-                this._searchScope.UserId,
+                searchScope.ApplicationId,
+                searchScope.AgentId,
+                searchScope.ThreadId,
+                searchScope.UserId,
                 queryText,
                 cancellationToken).ConfigureAwait(false)).ToList();
 
@@ -161,10 +135,10 @@ public sealed class Mem0Provider : AIContextProvider
                 this._logger.LogInformation(
                     "Mem0AIContextProvider: Retrieved {Count} memories. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.",
                     memories.Count,
-                    this._searchScope.ApplicationId,
-                    this._searchScope.AgentId,
-                    this._searchScope.ThreadId,
-                    this.SanitizeLogData(this._searchScope.UserId));
+                    searchScope.ApplicationId,
+                    searchScope.AgentId,
+                    searchScope.ThreadId,
+                    this.SanitizeLogData(searchScope.UserId));
 
                 if (outputMessageText is not null && this._logger.IsEnabled(LogLevel.Trace))
                 {
@@ -172,10 +146,10 @@ public sealed class Mem0Provider : AIContextProvider
                         "Mem0AIContextProvider: Search Results\nInput:{Input}\nOutput:{MessageText}\nApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.",
                         this.SanitizeLogData(queryText),
                         this.SanitizeLogData(outputMessageText),
-                        this._searchScope.ApplicationId,
-                        this._searchScope.AgentId,
-                        this._searchScope.ThreadId,
-                        this.SanitizeLogData(this._searchScope.UserId));
+                        searchScope.ApplicationId,
+                        searchScope.AgentId,
+                        searchScope.ThreadId,
+                        this.SanitizeLogData(searchScope.UserId));
                 }
             }
 
@@ -195,10 +169,10 @@ public sealed class Mem0Provider : AIContextProvider
                 this._logger.LogError(
                     ex,
                     "Mem0AIContextProvider: Failed to search Mem0 for memories due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.",
-                    this._searchScope.ApplicationId,
-                    this._searchScope.AgentId,
-                    this._searchScope.ThreadId,
-                    this.SanitizeLogData(this._searchScope.UserId));
+                    searchScope.ApplicationId,
+                    searchScope.AgentId,
+                    searchScope.ThreadId,
+                    this.SanitizeLogData(searchScope.UserId));
             }
             return new AIContext();
         }
@@ -212,10 +186,14 @@ public sealed class Mem0Provider : AIContextProvider
             return; // Do not update memory on failed invocations.
         }
 
+        var state = this.GetOrInitializeState(context.Session);
+        var storageScope = state?.StorageScope ?? new Mem0ProviderScope();
+
         try
         {
             // Persist request and response messages after invocation.
             await this.PersistMessagesAsync(
+                storageScope,
                 context.RequestMessages
                     .Where(m => m.GetAgentRequestMessageSource() == AgentRequestMessageSourceType.External)
                     .Concat(context.ResponseMessages ?? []),
@@ -228,36 +206,39 @@ public sealed class Mem0Provider : AIContextProvider
                 this._logger.LogError(
                     ex,
                     "Mem0AIContextProvider: Failed to send messages to Mem0 due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.",
-                    this._storageScope.ApplicationId,
-                    this._storageScope.AgentId,
-                    this._storageScope.ThreadId,
-                    this.SanitizeLogData(this._storageScope.UserId));
+                    storageScope.ApplicationId,
+                    storageScope.AgentId,
+                    storageScope.ThreadId,
+                    this.SanitizeLogData(storageScope.UserId));
             }
         }
     }
 
     /// <summary>
-    /// Clears stored memories for the configured scopes.
+    /// Clears stored memories for the specified scope.
     /// </summary>
+    /// <param name="session">The session containing the scope state to clear memories for.</param>
     /// <param name="cancellationToken">Cancellation token.</param>
-    public Task ClearStoredMemoriesAsync(CancellationToken cancellationToken = default) =>
-        this._client.ClearMemoryAsync(
-            this._storageScope.ApplicationId,
-            this._storageScope.AgentId,
-            this._storageScope.ThreadId,
-            this._storageScope.UserId,
-            cancellationToken);
-
-    /// <inheritdoc />
-    public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
+    public Task ClearStoredMemoriesAsync(AgentSession session, CancellationToken cancellationToken = default)
     {
-        var state = new Mem0State(this._storageScope, this._searchScope);
+        Throw.IfNull(session);
+        var state = this.GetOrInitializeState(session);
+        var storageScope = state?.StorageScope;
 
-        var jso = jsonSerializerOptions ?? Mem0JsonUtilities.DefaultOptions;
-        return JsonSerializer.SerializeToElement(state, jso.GetTypeInfo(typeof(Mem0State)));
+        if (storageScope is null)
+        {
+            return Task.CompletedTask; // Nothing to clear if there is no state.
+        }
+
+        return this._client.ClearMemoryAsync(
+            storageScope.ApplicationId,
+            storageScope.AgentId,
+            storageScope.ThreadId,
+            storageScope.UserId,
+            cancellationToken);
     }
 
-    private async Task PersistMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
+    private async Task PersistMessagesAsync(Mem0ProviderScope storageScope, IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
     {
         foreach (var message in messages)
         {
@@ -277,27 +258,42 @@ public sealed class Mem0Provider : AIContextProvider
             }
 
             await this._client.CreateMemoryAsync(
-                this._storageScope.ApplicationId,
-                this._storageScope.AgentId,
-                this._storageScope.ThreadId,
-                this._storageScope.UserId,
+                storageScope.ApplicationId,
+                storageScope.AgentId,
+                storageScope.ThreadId,
+                storageScope.UserId,
                 message.Text,
                 message.Role.Value,
                 cancellationToken).ConfigureAwait(false);
         }
     }
 
-    internal sealed class Mem0State
+    /// <summary>
+    /// Represents the state of a <see cref="Mem0Provider"/> stored in the <see cref="AgentSession.StateBag"/>.
+    /// </summary>
+    public sealed class State
     {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="State"/> class with the specified storage and search scopes.
+        /// </summary>
+        /// <param name="storageScope">The scope to use when storing memories.</param>
+        /// <param name="searchScope">The scope to use when searching for memories. If null, the storage scope will be used for searching as well.</param>
         [JsonConstructor]
-        public Mem0State(Mem0ProviderScope storageScope, Mem0ProviderScope searchScope)
+        public State(Mem0ProviderScope storageScope, Mem0ProviderScope? searchScope = null)
         {
-            this.StorageScope = storageScope;
-            this.SearchScope = searchScope;
+            this.StorageScope = Throw.IfNull(storageScope);
+            this.SearchScope = searchScope ?? storageScope;
         }
 
-        public Mem0ProviderScope StorageScope { get; set; }
-        public Mem0ProviderScope SearchScope { get; set; }
+        /// <summary>
+        /// Gets the scope used when storing memories.
+        /// </summary>
+        public Mem0ProviderScope StorageScope { get; }
+
+        /// <summary>
+        /// Gets the scope used when searching memories.
+        /// </summary>
+        public Mem0ProviderScope SearchScope { get; }
     }
 
     private string? SanitizeLogData(string? data) => this._enableSensitiveTelemetryData ? data : "<redacted>";

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

@@ -18,4 +18,10 @@ public sealed class Mem0ProviderOptions
     /// </summary>
     /// <value>Defaults to <see langword="false"/>.</value>
     public bool EnableSensitiveTelemetryData { get; set; }
+
+    /// <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>
+    public string? StateKey { get; set; }
 }

dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIAssistantClientExtensions.cs

@@ -204,8 +204,8 @@ public static class OpenAIAssistantClientExtensions
             Name = options.Name ?? assistantMetadata.Name,
             Description = options.Description ?? assistantMetadata.Description,
             ChatOptions = options.ChatOptions,
-            AIContextProviderFactory = options.AIContextProviderFactory,
-            ChatHistoryProviderFactory = options.ChatHistoryProviderFactory,
+            AIContextProvider = options.AIContextProvider,
+            ChatHistoryProvider = options.ChatHistoryProvider,
             UseProvidedChatClientAsIs = options.UseProvidedChatClientAsIs
         };
 

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

@@ -6,48 +6,54 @@ using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
-using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Agents.AI.Workflows;
 
 internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
 {
-    private int _bookmark;
-    private readonly List<ChatMessage> _chatMessages = [];
+    private const string DefaultStateBagKey = "WorkflowChatHistoryProvider.State";
+    private readonly JsonSerializerOptions _jsonSerializerOptions;
 
-    public WorkflowChatHistoryProvider()
+    /// <summary>
+    /// Initializes a new instance of the <see cref="WorkflowChatHistoryProvider"/> class.
+    /// </summary>
+    /// <param name="jsonSerializerOptions">
+    /// Optional JSON serializer options for serializing the state of this provider.
+    /// This is valuable for cases like when the chat history contains custom <see cref="AIContent"/> types
+    /// and source generated serializers are required, or Native AOT / Trimming is required.
+    /// </param>
+    public WorkflowChatHistoryProvider(JsonSerializerOptions? jsonSerializerOptions = null)
     {
-    }
-
-    public WorkflowChatHistoryProvider(StoreState state)
-    {
-        this.ImportStoreState(Throw.IfNull(state));
-    }
-
-    private void ImportStoreState(StoreState state, bool clearMessages = false)
-    {
-        if (clearMessages)
-        {
-            this._chatMessages.Clear();
-        }
-
-        if (state?.Messages is not null)
-        {
-            this._chatMessages.AddRange(state.Messages);
-        }
-        this._bookmark = state?.Bookmark ?? 0;
+        this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
     }
 
     internal sealed class StoreState
     {
         public int Bookmark { get; set; }
-        public IList<ChatMessage> Messages { get; set; } = [];
+        public List<ChatMessage> Messages { get; set; } = [];
     }
 
-    internal void AddMessages(params IEnumerable<ChatMessage> messages) => this._chatMessages.AddRange(messages);
+    private StoreState GetOrInitializeState(AgentSession? session)
+    {
+        if (session?.StateBag.TryGetValue<StoreState>(DefaultStateBagKey, out var state, this._jsonSerializerOptions) is true && state is not null)
+        {
+            return state;
+        }
+
+        state = new();
+        if (session is not null)
+        {
+            session.StateBag.SetValue(DefaultStateBagKey, state, this._jsonSerializerOptions);
+        }
+
+        return state;
+    }
+
+    internal void AddMessages(AgentSession session, params IEnumerable<ChatMessage> messages)
+        => this.GetOrInitializeState(session).Messages.AddRange(messages);
 
     protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
-        => new(this._chatMessages.AsReadOnly());
+        => new(this.GetOrInitializeState(context.Session).Messages.AsReadOnly());
 
     protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
@@ -57,28 +63,24 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
         }
 
         var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []);
-        this._chatMessages.AddRange(allNewMessages);
+        this.GetOrInitializeState(context.Session).Messages.AddRange(allNewMessages);
 
         return default;
     }
 
-    public IEnumerable<ChatMessage> GetFromBookmark()
+    public IEnumerable<ChatMessage> GetFromBookmark(AgentSession session)
     {
-        for (int i = this._bookmark; i < this._chatMessages.Count; i++)
+        var state = this.GetOrInitializeState(session);
+
+        for (int i = state.Bookmark; i < state.Messages.Count; i++)
         {
-            yield return this._chatMessages[i];
+            yield return state.Messages[i];
         }
     }
 
-    public void UpdateBookmark() => this._bookmark = this._chatMessages.Count;
-
-    public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
+    public void UpdateBookmark(AgentSession session)
     {
-        StoreState state = this.ExportStoreState();
-
-        return JsonSerializer.SerializeToElement(state,
-            WorkflowsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState)));
+        var state = this.GetOrInitializeState(session);
+        state.Bookmark = state.Messages.Count;
     }
-
-    internal StoreState ExportStoreState() => new() { Bookmark = this._bookmark, Messages = this._chatMessages };
 }

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

@@ -94,7 +94,7 @@ internal sealed class WorkflowHostAgent : AIAgent
 
         // For workflow threads, messages are added directly via the internal AddMessages method
         // The MessageStore methods are used for agent invocation scenarios
-        workflowSession.ChatHistoryProvider.AddMessages(messages);
+        workflowSession.ChatHistoryProvider.AddMessages(session, messages);
         return workflowSession;
     }
 

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

@@ -70,7 +70,8 @@ internal sealed class WorkflowSession : AgentSession
 
         this.RunId = sessionState.RunId;
         this.LastCheckpoint = sessionState.LastCheckpoint;
-        this.ChatHistoryProvider = new WorkflowChatHistoryProvider(sessionState.ChatHistoryProviderState);
+        this.ChatHistoryProvider = new WorkflowChatHistoryProvider();
+        this.StateBag = sessionState.StateBag;
     }
 
     public CheckpointInfo? LastCheckpoint { get; set; }
@@ -81,8 +82,8 @@ internal sealed class WorkflowSession : AgentSession
         SessionState info = new(
             this.RunId,
             this.LastCheckpoint,
-            this.ChatHistoryProvider.ExportStoreState(),
-            this._inMemoryCheckpointManager);
+            this._inMemoryCheckpointManager,
+            this.StateBag);
 
         return marshaller.Marshal(info);
     }
@@ -100,7 +101,7 @@ internal sealed class WorkflowSession : AgentSession
             RawRepresentation = raw
         };
 
-        this.ChatHistoryProvider.AddMessages(update.ToChatMessage());
+        this.ChatHistoryProvider.AddMessages(this, update.ToChatMessage());
 
         return update;
     }
@@ -117,7 +118,7 @@ internal sealed class WorkflowSession : AgentSession
             RawRepresentation = raw
         };
 
-        this.ChatHistoryProvider.AddMessages(update.ToChatMessage());
+        this.ChatHistoryProvider.AddMessages(this, update.ToChatMessage());
 
         return update;
     }
@@ -156,7 +157,7 @@ internal sealed class WorkflowSession : AgentSession
         try
         {
             this.LastResponseId = Guid.NewGuid().ToString("N");
-            List<ChatMessage> messages = this.ChatHistoryProvider.GetFromBookmark().ToList();
+            List<ChatMessage> messages = this.ChatHistoryProvider.GetFromBookmark(this).ToList();
 
 #pragma warning disable CA2007 // Analyzer misfiring and not seeing .ConfigureAwait(false) below.
             await using Checkpointed<StreamingRun> checkpointed =
@@ -240,7 +241,7 @@ internal sealed class WorkflowSession : AgentSession
         finally
         {
             // Do we want to try to undo the step, and not update the bookmark?
-            this.ChatHistoryProvider.UpdateBookmark();
+            this.ChatHistoryProvider.UpdateBookmark(this);
         }
     }
 
@@ -254,12 +255,12 @@ internal sealed class WorkflowSession : AgentSession
     internal sealed class SessionState(
         string runId,
         CheckpointInfo? lastCheckpoint,
-        WorkflowChatHistoryProvider.StoreState chatHistoryProviderState,
-        InMemoryCheckpointManager? checkpointManager = null)
+        InMemoryCheckpointManager? checkpointManager = null,
+        AgentSessionStateBag? stateBag = null)
     {
         public string RunId { get; } = runId;
         public CheckpointInfo? LastCheckpoint { get; } = lastCheckpoint;
-        public WorkflowChatHistoryProvider.StoreState ChatHistoryProviderState { get; } = chatHistoryProviderState;
         public InMemoryCheckpointManager? CheckpointManager { get; } = checkpointManager;
+        public AgentSessionStateBag StateBag { get; } = stateBag ?? new();
     }
 }

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

@@ -65,9 +65,9 @@ internal static partial class AgentJsonUtilities
         NumberHandling = JsonNumberHandling.AllowReadingFromString)]
 
     // Agent abstraction types
-    [JsonSerializable(typeof(ChatClientAgentSession.SessionState))]
+    [JsonSerializable(typeof(ChatClientAgentSession))]
     [JsonSerializable(typeof(TextSearchProvider.TextSearchProviderState))]
-    [JsonSerializable(typeof(ChatHistoryMemoryProvider.ChatHistoryMemoryProviderState))]
+    [JsonSerializable(typeof(ChatHistoryMemoryProvider.State))]
 
     [ExcludeFromCodeCoverage]
     internal sealed partial class JsonContext : JsonSerializerContext;

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

@@ -105,6 +105,11 @@ public sealed partial class ChatClientAgent : AIAgent
         // If the user has not opted out of using our default decorators, we wrap the chat client.
         this.ChatClient = options?.UseProvidedChatClientAsIs is true ? chatClient : chatClient.WithDefaultAgentMiddleware(options, services);
 
+        // Use the ChatHistoryProvider from options if provided.
+        // If one was not provided, and we later find out that the underlying service does not manage chat history server-side,
+        // we will use the default InMemoryChatHistoryProvider at that time.
+        this.ChatHistoryProvider = options?.ChatHistoryProvider;
+
         this._logger = (loggerFactory ?? chatClient.GetService<ILoggerFactory>() ?? NullLoggerFactory.Instance).CreateLogger<ChatClientAgent>();
     }
 
@@ -120,6 +125,14 @@ public sealed partial class ChatClientAgent : AIAgent
     /// </remarks>
     public IChatClient ChatClient { get; }
 
+    /// <summary>
+    /// Gets the <see cref="ChatHistoryProvider"/> used by this agent, to support cases where the chat history is not stored by the agent service.
+    /// </summary>
+    /// <remarks>
+    /// This property may be null in case the agent stores messages in the underlying agent service.
+    /// </remarks>
+    public ChatHistoryProvider? ChatHistoryProvider { get; private set; }
+
     /// <inheritdoc/>
     protected override string? IdCore => this._agentOptions?.Id;
 
@@ -282,7 +295,7 @@ public sealed partial class ChatClientAgent : AIAgent
 
         // We can derive the type of supported session from whether we have a conversation id,
         // so let's update it and set the conversation id for the service session case.
-        await this.UpdateSessionWithTypeAndConversationIdAsync(safeSession, chatResponse.ConversationId, cancellationToken).ConfigureAwait(false);
+        this.UpdateSessionConversationId(safeSession, chatResponse.ConversationId, cancellationToken);
 
         // To avoid inconsistent state we only notify the session of the input messages if no error occurs after the initial request.
         await this.NotifyChatHistoryProviderOfNewMessagesAsync(safeSession, GetInputMessages(inputMessagesForProviders, continuationToken), chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false);
@@ -298,24 +311,14 @@ public sealed partial class ChatClientAgent : AIAgent
         : serviceType == typeof(IChatClient) ? this.ChatClient
         : serviceType == typeof(ChatOptions) ? this._agentOptions?.ChatOptions
         : serviceType == typeof(ChatClientAgentOptions) ? this._agentOptions
-        : this.ChatClient.GetService(serviceType, serviceKey));
+        : this._agentOptions?.AIContextProvider?.GetService(serviceType, serviceKey)
+        ?? this.ChatHistoryProvider?.GetService(serviceType, serviceKey)
+        ?? this.ChatClient.GetService(serviceType, serviceKey));
 
     /// <inheritdoc/>
-    protected override async ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default)
+    protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default)
     {
-        ChatHistoryProvider? chatHistoryProvider = this._agentOptions?.ChatHistoryProviderFactory is not null
-            ? await this._agentOptions.ChatHistoryProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false)
-            : null;
-
-        AIContextProvider? contextProvider = this._agentOptions?.AIContextProviderFactory is not null
-            ? await this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false)
-            : null;
-
-        return new ChatClientAgentSession
-        {
-            ChatHistoryProvider = chatHistoryProvider,
-            AIContextProvider = contextProvider
-        };
+        return new(new ChatClientAgentSession());
     }
 
     /// <summary>
@@ -336,52 +339,12 @@ public sealed partial class ChatClientAgent : AIAgent
     /// instances that support server-side conversation storage through their underlying <see cref="IChatClient"/>.
     /// </para>
     /// </remarks>
-    public async ValueTask<AgentSession> CreateSessionAsync(string conversationId, CancellationToken cancellationToken = default)
+    public ValueTask<AgentSession> CreateSessionAsync(string conversationId, CancellationToken cancellationToken = default)
     {
-        AIContextProvider? contextProvider = this._agentOptions?.AIContextProviderFactory is not null
-            ? await this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false)
-            : null;
-
-        return new ChatClientAgentSession()
+        return new(new ChatClientAgentSession()
         {
             ConversationId = conversationId,
-            AIContextProvider = contextProvider
-        };
-    }
-
-    /// <summary>
-    /// Creates a new agent session instance using an existing <see cref="ChatHistoryProvider"/> to continue a conversation.
-    /// </summary>
-    /// <param name="chatHistoryProvider">The <see cref="ChatHistoryProvider"/> instance to use for managing the conversation's message history.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
-    /// <returns>
-    /// A value task representing the asynchronous operation. The task result contains a new <see cref="AgentSession"/> instance configured to work with the provided <paramref name="chatHistoryProvider"/>.
-    /// </returns>
-    /// <remarks>
-    /// <para>
-    /// This method creates threads that do not support server-side conversation storage.
-    /// Some AI services require server-side conversation storage to function properly, and creating a session
-    /// with a <see cref="ChatHistoryProvider"/> may not be compatible with these services.
-    /// </para>
-    /// <para>
-    /// Where a service requires server-side conversation storage, use <see cref="CreateSessionAsync(string, CancellationToken)"/>.
-    /// </para>
-    /// <para>
-    /// If the agent detects, during the first run, that the underlying AI service requires server-side conversation storage,
-    /// the session will throw an exception to indicate that it cannot continue using the provided <see cref="ChatHistoryProvider"/>.
-    /// </para>
-    /// </remarks>
-    public async ValueTask<AgentSession> CreateSessionAsync(ChatHistoryProvider chatHistoryProvider, CancellationToken cancellationToken = default)
-    {
-        AIContextProvider? contextProvider = this._agentOptions?.AIContextProviderFactory is not null
-            ? await this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false)
-            : null;
-
-        return new ChatClientAgentSession()
-        {
-            ChatHistoryProvider = Throw.IfNull(chatHistoryProvider),
-            AIContextProvider = contextProvider
-        };
+        });
     }
 
     /// <inheritdoc/>
@@ -398,22 +361,9 @@ public sealed partial class ChatClientAgent : AIAgent
     }
 
     /// <inheritdoc/>
-    protected override async ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
     {
-        Func<JsonElement, JsonSerializerOptions?, CancellationToken, ValueTask<ChatHistoryProvider>>? chatHistoryProviderFactory = this._agentOptions?.ChatHistoryProviderFactory is null ?
-            null :
-            (jse, jso, ct) => this._agentOptions.ChatHistoryProviderFactory.Invoke(new() { SerializedState = jse, JsonSerializerOptions = jso }, ct);
-
-        Func<JsonElement, JsonSerializerOptions?, CancellationToken, ValueTask<AIContextProvider>>? aiContextProviderFactory = this._agentOptions?.AIContextProviderFactory is null ?
-            null :
-            (jse, jso, ct) => this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = jse, JsonSerializerOptions = jso }, ct);
-
-        return await ChatClientAgentSession.DeserializeAsync(
-            serializedState,
-            jsonSerializerOptions,
-            chatHistoryProviderFactory,
-            aiContextProviderFactory,
-            cancellationToken).ConfigureAwait(false);
+        return new(ChatClientAgentSession.Deserialize(serializedState, jsonSerializerOptions));
     }
 
     #region Private
@@ -462,7 +412,7 @@ public sealed partial class ChatClientAgent : AIAgent
 
         // We can derive the type of supported session from whether we have a conversation id,
         // so let's update it and set the conversation id for the service session case.
-        await this.UpdateSessionWithTypeAndConversationIdAsync(safeSession, chatResponse.ConversationId, cancellationToken).ConfigureAwait(false);
+        this.UpdateSessionConversationId(safeSession, chatResponse.ConversationId, cancellationToken);
 
         // Ensure that the author name is set for each message in the response.
         foreach (ChatMessage chatResponseMessage in chatResponse.Messages)
@@ -492,9 +442,9 @@ public sealed partial class ChatClientAgent : AIAgent
         IEnumerable<ChatMessage> responseMessages,
         CancellationToken cancellationToken)
     {
-        if (session.AIContextProvider is not null)
+        if (this._agentOptions?.AIContextProvider is { } contextProvider)
         {
-            await session.AIContextProvider.InvokedAsync(new(this, session, inputMessages) { ResponseMessages = responseMessages },
+            await contextProvider.InvokedAsync(new(this, session, inputMessages) { ResponseMessages = responseMessages },
                 cancellationToken).ConfigureAwait(false);
         }
     }
@@ -508,9 +458,9 @@ public sealed partial class ChatClientAgent : AIAgent
         IEnumerable<ChatMessage> inputMessages,
         CancellationToken cancellationToken)
     {
-        if (session.AIContextProvider is not null)
+        if (this._agentOptions?.AIContextProvider is { } contextProvider)
         {
-            await session.AIContextProvider.InvokedAsync(new(this, session, inputMessages) { InvokeException = ex },
+            await contextProvider.InvokedAsync(new(this, session, inputMessages) { InvokeException = ex },
                 cancellationToken).ConfigureAwait(false);
         }
     }
@@ -715,7 +665,7 @@ public sealed partial class ChatClientAgent : AIAgent
         // Populate the session messages only if we are not continuing an existing response as it's not allowed
         if (chatOptions?.ContinuationToken is null)
         {
-            ChatHistoryProvider? chatHistoryProvider = ResolveChatHistoryProvider(typedSession, chatOptions);
+            ChatHistoryProvider? chatHistoryProvider = this.ResolveChatHistoryProvider(chatOptions, typedSession);
 
             // Add any existing messages from the session to the messages to be sent to the chat client.
             if (chatHistoryProvider is not null)
@@ -731,10 +681,10 @@ public sealed partial class ChatClientAgent : AIAgent
 
             // If we have an AIContextProvider, we should get context from it, and update our
             // messages and options with the additional context.
-            if (typedSession.AIContextProvider is not null)
+            if (this._agentOptions?.AIContextProvider is { } aiContextProvider)
             {
                 var invokingContext = new AIContextProvider.InvokingContext(this, typedSession, inputMessages);
-                var aiContext = await typedSession.AIContextProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
+                var aiContext = await aiContextProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
                 if (aiContext.Messages is { Count: > 0 })
                 {
                     inputMessagesForProviders.AddRange(aiContext.Messages);
@@ -780,7 +730,7 @@ public sealed partial class ChatClientAgent : AIAgent
         return (typedSession, chatOptions, inputMessagesForProviders, inputMessagesForChatClient, continuationToken);
     }
 
-    private async Task UpdateSessionWithTypeAndConversationIdAsync(ChatClientAgentSession session, string? responseConversationId, CancellationToken cancellationToken)
+    private void UpdateSessionConversationId(ChatClientAgentSession session, string? responseConversationId, CancellationToken cancellationToken)
     {
         if (string.IsNullOrWhiteSpace(responseConversationId) && !string.IsNullOrWhiteSpace(session.ConversationId))
         {
@@ -791,6 +741,14 @@ public sealed partial class ChatClientAgent : AIAgent
 
         if (!string.IsNullOrWhiteSpace(responseConversationId))
         {
+            if (this.ChatHistoryProvider is not null)
+            {
+                // The agent has a ChatHistoryProvider configured, but the service returned a conversation id,
+                // meaning the service manages chat history server-side. Both cannot be used simultaneously.
+                throw new InvalidOperationException(
+                    $"Only {nameof(ChatClientAgentSession.ConversationId)} or {nameof(this.ChatHistoryProvider)} may be used, but not both. The service returned a conversation id indicating server-side chat history management, but the agent has a {nameof(this.ChatHistoryProvider)} configured.");
+            }
+
             // If we got a conversation id back from the chat client, it means that the service supports server side session storage
             // so we should update the session with the new id.
             session.ConversationId = responseConversationId;
@@ -798,11 +756,9 @@ public sealed partial class ChatClientAgent : AIAgent
         else
         {
             // If the service doesn't use service side chat history storage (i.e. we got no id back from invocation), and
-            // the session has no ChatHistoryProvider yet, we should update the session with the custom ChatHistoryProvider or
-            // default InMemoryChatHistoryProvider so that it has somewhere to store the chat history.
-            session.ChatHistoryProvider ??= this._agentOptions?.ChatHistoryProviderFactory is not null
-                ? await this._agentOptions.ChatHistoryProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false)
-                : new InMemoryChatHistoryProvider();
+            // the agent has no ChatHistoryProvider yet, we should use the default InMemoryChatHistoryProvider so that
+            // we have somewhere to store the chat history.
+            this.ChatHistoryProvider ??= new InMemoryChatHistoryProvider();
         }
     }
 
@@ -813,7 +769,7 @@ public sealed partial class ChatClientAgent : AIAgent
         ChatOptions? chatOptions,
         CancellationToken cancellationToken)
     {
-        ChatHistoryProvider? provider = ResolveChatHistoryProvider(session, chatOptions);
+        ChatHistoryProvider? provider = this.ResolveChatHistoryProvider(chatOptions, session);
 
         // Only notify the provider if we have one.
         // If we don't have one, it means that the chat history is service managed and the underlying service is responsible for storing messages.
@@ -837,7 +793,7 @@ public sealed partial class ChatClientAgent : AIAgent
         ChatOptions? chatOptions,
         CancellationToken cancellationToken)
     {
-        ChatHistoryProvider? provider = ResolveChatHistoryProvider(session, chatOptions);
+        ChatHistoryProvider? provider = this.ResolveChatHistoryProvider(chatOptions, session);
 
         // Only notify the provider if we have one.
         // If we don't have one, it means that the chat history is service managed and the underlying service is responsible for storing messages.
@@ -853,13 +809,25 @@ public sealed partial class ChatClientAgent : AIAgent
         return Task.CompletedTask;
     }
 
-    private static ChatHistoryProvider? ResolveChatHistoryProvider(ChatClientAgentSession session, ChatOptions? chatOptions)
+    private ChatHistoryProvider? ResolveChatHistoryProvider(ChatOptions? chatOptions, ChatClientAgentSession session)
     {
-        ChatHistoryProvider? provider = session.ChatHistoryProvider;
+        ChatHistoryProvider? provider = this.ChatHistoryProvider;
 
-        // If someone provided an override ChatHistoryProvider via AdditionalProperties, we should use that instead of the one on the session.
+        if (session.ConversationId is not null && provider is not null)
+        {
+            throw new InvalidOperationException(
+                $"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 the agent has a {nameof(this.ChatHistoryProvider)} configured.");
+        }
+
+        // If someone provided an override ChatHistoryProvider via AdditionalProperties, we should use that instead.
         if (chatOptions?.AdditionalProperties?.TryGetValue(out ChatHistoryProvider? overrideProvider) is true)
         {
+            if (session.ConversationId is not null && overrideProvider is not null)
+            {
+                throw new InvalidOperationException(
+                    $"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)}.");
+            }
+
             provider = overrideProvider;
         }
 

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

@@ -1,9 +1,5 @@
 // Copyright (c) Microsoft. All rights reserved.
 
-using System;
-using System.Text.Json;
-using System.Threading;
-using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
 
 namespace Microsoft.Agents.AI;
@@ -39,17 +35,14 @@ public sealed class ChatClientAgentOptions
     public ChatOptions? ChatOptions { get; set; }
 
     /// <summary>
-    /// Gets or sets a factory function to create an instance of <see cref="ChatHistoryProvider"/>
-    /// which will be used to provide chat history for this agent.
+    /// Gets or sets the <see cref="ChatHistoryProvider"/> instance to use for providing chat history for this agent.
     /// </summary>
-    public Func<ChatHistoryProviderFactoryContext, CancellationToken, ValueTask<ChatHistoryProvider>>? ChatHistoryProviderFactory { get; set; }
+    public ChatHistoryProvider? ChatHistoryProvider { get; set; }
 
     /// <summary>
-    /// Gets or sets a factory function to create an instance of <see cref="AIContextProvider"/>
-    /// which will be used to create a context provider for each new thread, and can then
-    /// provide additional context for each agent run.
+    /// Gets or sets the <see cref="AIContextProvider"/> instance to use for providing additional context for each agent run.
     /// </summary>
-    public Func<AIContextProviderFactoryContext, CancellationToken, ValueTask<AIContextProvider>>? AIContextProviderFactory { get; set; }
+    public AIContextProvider? AIContextProvider { get; set; }
 
     /// <summary>
     /// Gets or sets a value indicating whether to use the provided <see cref="IChatClient"/> instance as is,
@@ -75,41 +68,7 @@ public sealed class ChatClientAgentOptions
             Name = this.Name,
             Description = this.Description,
             ChatOptions = this.ChatOptions?.Clone(),
-            ChatHistoryProviderFactory = this.ChatHistoryProviderFactory,
-            AIContextProviderFactory = this.AIContextProviderFactory,
+            ChatHistoryProvider = this.ChatHistoryProvider,
+            AIContextProvider = this.AIContextProvider,
         };
-
-    /// <summary>
-    /// Context object passed to the <see cref="AIContextProviderFactory"/> to create a new instance of <see cref="AIContextProvider"/>.
-    /// </summary>
-    public sealed class AIContextProviderFactoryContext
-    {
-        /// <summary>
-        /// Gets or sets the serialized state of the <see cref="AIContextProvider"/>, if any.
-        /// </summary>
-        /// <value><see langword="default"/> if there is no state, e.g. when the <see cref="AIContextProvider"/> is first created.</value>
-        public JsonElement SerializedState { get; set; }
-
-        /// <summary>
-        /// Gets or sets the JSON serialization options to use when deserializing the <see cref="SerializedState"/>.
-        /// </summary>
-        public JsonSerializerOptions? JsonSerializerOptions { get; set; }
-    }
-
-    /// <summary>
-    /// Context object passed to the <see cref="ChatHistoryProviderFactory"/> to create a new instance of <see cref="ChatHistoryProvider"/>.
-    /// </summary>
-    public sealed class ChatHistoryProviderFactoryContext
-    {
-        /// <summary>
-        /// Gets or sets the serialized state of the <see cref="ChatHistoryProvider"/>, if any.
-        /// </summary>
-        /// <value><see langword="default"/> if there is no state, e.g. when the <see cref="ChatHistoryProvider"/> is first created.</value>
-        public JsonElement SerializedState { get; set; }
-
-        /// <summary>
-        /// Gets or sets the JSON serialization options to use when deserializing the <see cref="SerializedState"/>.
-        /// </summary>
-        public JsonSerializerOptions? JsonSerializerOptions { get; set; }
-    }
 }

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

@@ -3,8 +3,7 @@
 using System;
 using System.Diagnostics;
 using System.Text.Json;
-using System.Threading;
-using System.Threading.Tasks;
+using System.Text.Json.Serialization;
 using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Agents.AI;
@@ -15,8 +14,6 @@ namespace Microsoft.Agents.AI;
 [DebuggerDisplay("{DebuggerDisplay,nq}")]
 public sealed class ChatClientAgentSession : AgentSession
 {
-    private ChatHistoryProvider? _chatHistoryProvider;
-
     /// <summary>
     /// Initializes a new instance of the <see cref="ChatClientAgentSession"/> class.
     /// </summary>
@@ -24,29 +21,30 @@ public sealed class ChatClientAgentSession : AgentSession
     {
     }
 
+    [JsonConstructor]
+    internal ChatClientAgentSession(string? conversationId, AgentSessionStateBag? stateBag) : base(stateBag ?? new())
+    {
+        this.ConversationId = conversationId;
+    }
+
     /// <summary>
-    /// Gets or sets the ID of the underlying service thread to support cases where the chat history is stored by the agent service.
+    /// Gets or sets the ID of the underlying service chat history to support cases where the chat history is stored by the agent service.
     /// </summary>
     /// <remarks>
     /// <para>
-    /// Note that either <see cref="ConversationId"/> or <see cref="ChatHistoryProvider "/> may be set, but not both.
-    /// If <see cref="ChatHistoryProvider "/> is not null, setting <see cref="ConversationId"/> will throw an
-    /// <see cref="InvalidOperationException "/> exception.
-    /// </para>
-    /// <para>
     /// This property may be null in the following cases:
     /// <list type="bullet">
-    /// <item><description>The thread stores messages via the <see cref="AI.ChatHistoryProvider"/> and not in the agent service.</description></item>
-    /// <item><description>This thread object is new and a server managed thread has not yet been created in the agent service.</description></item>
+    /// <item><description>The agent stores messages via a <see cref="ChatHistoryProvider"/> and not in the agent service.</description></item>
+    /// <item><description>This session object is new and server managed chat history has not yet been created in the agent service.</description></item>
     /// </list>
     /// </para>
     /// <para>
-    /// The id may also change over time where the id is pointing at a
-    /// agent service managed thread, and the default behavior of a service is
-    /// to fork the thread with each iteration.
+    /// The id may also change over time where the id is pointing at
+    /// agent service managed chat history, and the default behavior of a service is
+    /// to fork the chat history with each iteration.
     /// </para>
     /// </remarks>
-    /// <exception cref="InvalidOperationException">Attempted to set a conversation ID but a <see cref="ChatHistoryProvider"/> is already set.</exception>
+    [JsonPropertyName("conversationId")]
     public string? ConversationId
     {
         get;
@@ -57,149 +55,37 @@ public sealed class ChatClientAgentSession : AgentSession
                 return;
             }
 
-            if (this._chatHistoryProvider is not null)
-            {
-                // If we have a ChatHistoryProvider already, we shouldn't switch the session to use a conversation id
-                // since it means that the session contents will essentially be deleted, and the session will not work
-                // with the original agent anymore.
-                throw new InvalidOperationException("Only the ConversationId or ChatHistoryProvider may be set, but not both and switching from one to another is not supported.");
-            }
-
             field = Throw.IfNullOrWhitespace(value);
         }
     }
 
-    /// <summary>
-    /// Gets or sets the <see cref="AI.ChatHistoryProvider"/> used by this thread, for cases where messages should be stored in a custom location.
-    /// </summary>
-    /// <remarks>
-    /// <para>
-    /// Note that either <see cref="ConversationId"/> or <see cref="ChatHistoryProvider "/> may be set, but not both.
-    /// If <see cref="ConversationId"/> is not null, and <see cref="ChatHistoryProvider "/> is set, <see cref="ConversationId"/>
-    /// will be reverted to null, and vice versa.
-    /// </para>
-    /// <para>
-    /// This property may be null in the following cases:
-    /// <list type="bullet">
-    /// <item><description>The thread stores messages in the agent service and just has an id to the remove thread, instead of in an <see cref="AI.ChatHistoryProvider"/>.</description></item>
-    /// <item><description>This thread object is new it is not yet clear whether it will be backed by a server managed thread or an <see cref="AI.ChatHistoryProvider"/>.</description></item>
-    /// </list>
-    /// </para>
-    /// </remarks>
-    public ChatHistoryProvider? ChatHistoryProvider
-    {
-        get => this._chatHistoryProvider;
-        internal set
-        {
-            if (this._chatHistoryProvider is null && value is null)
-            {
-                return;
-            }
-
-            if (!string.IsNullOrWhiteSpace(this.ConversationId))
-            {
-                // If we have a conversation id already, we shouldn't switch the session to use a ChatHistoryProvider
-                // since it means that the session will not work with the original agent anymore.
-                throw new InvalidOperationException("Only the ConversationId or ChatHistoryProvider may be set, but not both and switching from one to another is not supported.");
-            }
-
-            this._chatHistoryProvider = Throw.IfNull(value);
-        }
-    }
-
-    /// <summary>
-    /// Gets or sets the <see cref="AIContextProvider"/> used by this thread to provide additional context to the AI model before each invocation.
-    /// </summary>
-    public AIContextProvider? AIContextProvider { get; internal set; }
-
     /// <summary>
     /// Creates a new instance of the <see cref="ChatClientAgentSession"/> class from previously serialized state.
     /// </summary>
     /// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized state of the session.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="chatHistoryProviderFactory">
-    /// An optional factory function to create a custom <see cref="AI.ChatHistoryProvider"/> from its serialized state.
-    /// If not provided, the default <see cref="InMemoryChatHistoryProvider"/> will be used.
-    /// </param>
-    /// <param name="aiContextProviderFactory">
-    /// An optional factory function to create a custom <see cref="AIContextProvider"/> from its serialized state.
-    /// If not provided, no context provider will be configured.
-    /// </param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
-    /// <returns>A task representing the asynchronous operation. The task result contains the deserialized <see cref="ChatClientAgentSession"/>.</returns>
-    internal static async Task<ChatClientAgentSession> DeserializeAsync(
-        JsonElement serializedState,
-        JsonSerializerOptions? jsonSerializerOptions = null,
-        Func<JsonElement, JsonSerializerOptions?, CancellationToken, ValueTask<ChatHistoryProvider>>? chatHistoryProviderFactory = null,
-        Func<JsonElement, JsonSerializerOptions?, CancellationToken, ValueTask<AIContextProvider>>? aiContextProviderFactory = null,
-        CancellationToken cancellationToken = default)
+    /// <param name="jsonSerializerOptions">Optional JSON serialization options to use instead of the default options.</param>
+    /// <returns>The deserialized <see cref="ChatClientAgentSession"/>.</returns>
+    internal static ChatClientAgentSession Deserialize(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
     {
         if (serializedState.ValueKind != JsonValueKind.Object)
         {
             throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState));
         }
 
-        var state = serializedState.Deserialize(
-            AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(SessionState))) as SessionState;
-
-        var session = new ChatClientAgentSession();
-
-        session.AIContextProvider = aiContextProviderFactory is not null
-            ? await aiContextProviderFactory.Invoke(state?.AIContextProviderState ?? default, jsonSerializerOptions, cancellationToken).ConfigureAwait(false)
-            : null;
-
-        if (state?.ConversationId is string sessionId)
-        {
-            session.ConversationId = sessionId;
-
-            // Since we have an ID, we should not have a ChatHistoryProvider and we can return here.
-            return session;
-        }
-
-        session._chatHistoryProvider =
-            chatHistoryProviderFactory is not null
-                ? await chatHistoryProviderFactory.Invoke(state?.ChatHistoryProviderState ?? default, jsonSerializerOptions, cancellationToken).ConfigureAwait(false)
-                : new InMemoryChatHistoryProvider(state?.ChatHistoryProviderState ?? default, jsonSerializerOptions); // default to an in-memory ChatHistoryProvider
-
-        return session;
+        var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
+        return serializedState.Deserialize(jso.GetTypeInfo(typeof(ChatClientAgentSession))) as ChatClientAgentSession
+            ?? new ChatClientAgentSession();
     }
 
     /// <inheritdoc/>
     internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        JsonElement? chatHistoryProviderState = this._chatHistoryProvider?.Serialize(jsonSerializerOptions);
-
-        JsonElement? aiContextProviderState = this.AIContextProvider?.Serialize(jsonSerializerOptions);
-
-        var state = new SessionState
-        {
-            ConversationId = this.ConversationId,
-            ChatHistoryProviderState = chatHistoryProviderState is { ValueKind: not JsonValueKind.Undefined } ? chatHistoryProviderState : null,
-            AIContextProviderState = aiContextProviderState is { ValueKind: not JsonValueKind.Undefined } ? aiContextProviderState : null,
-        };
-
-        return JsonSerializer.SerializeToElement(state, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(SessionState)));
+        var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
+        return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(ChatClientAgentSession)));
     }
 
-    /// <inheritdoc/>
-    public override object? GetService(Type serviceType, object? serviceKey = null) =>
-        base.GetService(serviceType, serviceKey)
-            ?? this.AIContextProvider?.GetService(serviceType, serviceKey)
-            ?? this.ChatHistoryProvider?.GetService(serviceType, serviceKey);
-
     [DebuggerBrowsable(DebuggerBrowsableState.Never)]
     private string DebuggerDisplay =>
-        this.ConversationId is { } conversationId ? $"ConversationId = {conversationId}" :
-        this._chatHistoryProvider is InMemoryChatHistoryProvider inMemoryChatHistoryProvider ? $"Count = {inMemoryChatHistoryProvider.Count}" :
-        this._chatHistoryProvider is { } chatHistoryProvider ? $"ChatHistoryProvider = {chatHistoryProvider.GetType().Name}" :
-        "Count = 0";
-
-    internal sealed class SessionState
-    {
-        public string? ConversationId { get; set; }
-
-        public JsonElement? ChatHistoryProviderState { get; set; }
-
-        public JsonElement? AIContextProviderState { get; set; }
-    }
+        this.ConversationId is { } conversationId ? $"ConversationId = {conversationId}, StateBag Count = {this.StateBag.Count}" :
+        $"StateBag Count = {this.StateBag.Count}";
 }

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

@@ -4,7 +4,6 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Linq.Expressions;
-using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -41,18 +40,21 @@ 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;
+#pragma warning restore CA2213
     private readonly VectorStoreCollection<object, Dictionary<string, object?>> _collection;
     private readonly int _maxResults;
     private readonly string _contextPrompt;
     private readonly bool _enableSensitiveTelemetryData;
     private readonly ChatHistoryMemoryProviderOptions.SearchBehavior _searchTime;
-    private readonly AITool[] _tools;
+    private readonly string _toolName;
+    private readonly string _toolDescription;
     private readonly ILogger<ChatHistoryMemoryProvider>? _logger;
-
-    private readonly ChatHistoryMemoryProviderScope _storageScope;
-    private readonly ChatHistoryMemoryProviderScope _searchScope;
+    private readonly string _stateKey;
+    private readonly Func<AgentSession?, State> _stateInitializer;
 
     private bool _collectionInitialized;
     private readonly SemaphoreSlim _initializationLock = new(1, 1);
@@ -64,93 +66,30 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
     /// <param name="vectorStore">The vector store to use for storing and retrieving chat history.</param>
     /// <param name="collectionName">The name of the collection for storing chat history in the vector store.</param>
     /// <param name="vectorDimensions">The number of dimensions to use for the chat history vector store embeddings.</param>
-    /// <param name="storageScope">Optional values to scope the chat history storage with.</param>
-    /// <param name="searchScope">Optional values to scope the chat history search with. Where values are null, no filtering is done using those values. Defaults to <paramref name="storageScope"/> if not provided.</param>
+    /// <param name="stateInitializer">A delegate that initializes the provider state on the first invocation, providing the storage and search scopes.</param>
     /// <param name="options">Optional configuration options.</param>
     /// <param name="loggerFactory">Optional logger factory.</param>
-    /// <exception cref="ArgumentNullException">Thrown when <paramref name="vectorStore"/> is <see langword="null"/>.</exception>
+    /// <exception cref="ArgumentNullException">Thrown when <paramref name="vectorStore"/> or <paramref name="stateInitializer"/> is <see langword="null"/>.</exception>
     public ChatHistoryMemoryProvider(
         VectorStore vectorStore,
         string collectionName,
         int vectorDimensions,
-        ChatHistoryMemoryProviderScope storageScope,
-        ChatHistoryMemoryProviderScope? searchScope = null,
+        Func<AgentSession?, State> stateInitializer,
         ChatHistoryMemoryProviderOptions? options = null,
         ILoggerFactory? loggerFactory = null)
-        : this(
-            vectorStore,
-            collectionName,
-            vectorDimensions,
-            new ChatHistoryMemoryProviderState
-            {
-                StorageScope = new(Throw.IfNull(storageScope)),
-                SearchScope = searchScope ?? new(storageScope),
-            },
-            options,
-            loggerFactory)
     {
-    }
+        this._vectorStore = Throw.IfNull(vectorStore);
+        this._stateInitializer = Throw.IfNull(stateInitializer);
 
-    /// <summary>
-    /// Initializes a new instance of the <see cref="ChatHistoryMemoryProvider"/> class from previously serialized state.
-    /// </summary>
-    /// <param name="vectorStore">The vector store to use for storing and retrieving chat history.</param>
-    /// <param name="collectionName">The name of the collection for storing chat history in the vector store.</param>
-    /// <param name="vectorDimensions">The number of dimensions to use for the chat history vector store embeddings.</param>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized state of the provider.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="options">Optional configuration options.</param>
-    /// <param name="loggerFactory">Optional logger factory.</param>
-    public ChatHistoryMemoryProvider(
-        VectorStore vectorStore,
-        string collectionName,
-        int vectorDimensions,
-        JsonElement serializedState,
-        JsonSerializerOptions? jsonSerializerOptions = null,
-        ChatHistoryMemoryProviderOptions? options = null,
-        ILoggerFactory? loggerFactory = null)
-        : this(
-            vectorStore,
-            collectionName,
-            vectorDimensions,
-            DeserializeState(serializedState, jsonSerializerOptions),
-            options,
-            loggerFactory)
-    {
-    }
-
-    private ChatHistoryMemoryProvider(
-        VectorStore vectorStore,
-        string collectionName,
-        int vectorDimensions,
-        ChatHistoryMemoryProviderState? state = null,
-        ChatHistoryMemoryProviderOptions? options = null,
-        ILoggerFactory? loggerFactory = null)
-    {
-        this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
         options ??= new ChatHistoryMemoryProviderOptions();
         this._maxResults = options.MaxResults.HasValue ? Throw.IfLessThanOrEqual(options.MaxResults.Value, 0) : DefaultMaxResults;
         this._contextPrompt = options.ContextPrompt ?? DefaultContextPrompt;
         this._enableSensitiveTelemetryData = options.EnableSensitiveTelemetryData;
         this._searchTime = options.SearchTime;
+        this._stateKey = options.StateKey ?? DefaultStateBagKey;
         this._logger = loggerFactory?.CreateLogger<ChatHistoryMemoryProvider>();
-
-        if (state == null || state.StorageScope == null || state.SearchScope == null)
-        {
-            throw new InvalidOperationException($"The {nameof(ChatHistoryMemoryProvider)} state did not contain the required scope properties.");
-        }
-
-        this._storageScope = state.StorageScope;
-        this._searchScope = state.SearchScope;
-
-        // Create on-demand search tool (only used when behavior is OnDemandFunctionCalling)
-        this._tools =
-        [
-            AIFunctionFactory.Create(
-                (Func<string, CancellationToken, Task<string>>)this.SearchTextAsync,
-                name: options.FunctionToolName ?? DefaultFunctionToolName,
-                description: options.FunctionToolDescription ?? DefaultFunctionToolDescription)
-        ];
+        this._toolName = options.FunctionToolName ?? DefaultFunctionToolName;
+        this._toolDescription = options.FunctionToolDescription ?? DefaultFunctionToolDescription;
 
         // Create a definition so that we can use the dimensions provided at runtime.
         var definition = new VectorStoreCollectionDefinition
@@ -174,15 +113,51 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
         this._collection = this._vectorStore.GetDynamicCollection(Throw.IfNullOrWhitespace(collectionName), definition);
     }
 
+    /// <summary>
+    /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present.
+    /// </summary>
+    /// <param name="session">The agent session containing the StateBag.</param>
+    /// <returns>The provider state, or null if no session is available.</returns>
+    private State? GetOrInitializeState(AgentSession? session)
+    {
+        if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state, AgentJsonUtilities.DefaultOptions) is true && state is not null)
+        {
+            return state;
+        }
+
+        state = this._stateInitializer(session);
+        if (state is not null && session is not null)
+        {
+            session.StateBag.SetValue(this._stateKey, state, AgentJsonUtilities.DefaultOptions);
+        }
+
+        return state;
+    }
+
     /// <inheritdoc />
     protected override async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {
         _ = Throw.IfNull(context);
 
+        var state = this.GetOrInitializeState(context.Session);
+        var searchScope = state?.SearchScope ?? new ChatHistoryMemoryProviderScope();
+
         if (this._searchTime == ChatHistoryMemoryProviderOptions.SearchBehavior.OnDemandFunctionCalling)
         {
+            Task<string> InlineSearchAsync(string userQuestion, CancellationToken ct)
+                => this.SearchTextAsync(userQuestion, searchScope, ct);
+
+            // Create on-demand search tool (only used when behavior is OnDemandFunctionCalling)
+            AITool[] tools =
+            [
+                AIFunctionFactory.Create(
+                    InlineSearchAsync,
+                    name: this._toolName,
+                    description: this._toolDescription)
+            ];
+
             // Expose search tool for on-demand invocation by the model
-            return new AIContext { Tools = this._tools };
+            return new AIContext { Tools = tools };
         }
 
         try
@@ -199,7 +174,7 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
             }
 
             // Search for relevant chat history
-            var contextText = await this.SearchTextAsync(requestText, cancellationToken).ConfigureAwait(false);
+            var contextText = await this.SearchTextAsync(requestText, searchScope, cancellationToken).ConfigureAwait(false);
 
             if (string.IsNullOrWhiteSpace(contextText))
             {
@@ -218,10 +193,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
                 this._logger.LogError(
                     ex,
                     "ChatHistoryMemoryProvider: Failed to search for chat history due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', SessionId: '{SessionId}', UserId: '{UserId}'.",
-                    this._searchScope.ApplicationId,
-                    this._searchScope.AgentId,
-                    this._searchScope.SessionId,
-                    this.SanitizeLogData(this._searchScope.UserId));
+                    searchScope.ApplicationId,
+                    searchScope.AgentId,
+                    searchScope.SessionId,
+                    this.SanitizeLogData(searchScope.UserId));
             }
 
             return new AIContext();
@@ -239,6 +214,9 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
             return;
         }
 
+        var state = this.GetOrInitializeState(context.Session);
+        var storageScope = state?.StorageScope ?? new ChatHistoryMemoryProviderScope();
+
         try
         {
             // Ensure the collection is initialized
@@ -253,10 +231,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
                     ["Role"] = message.Role.ToString(),
                     ["MessageId"] = message.MessageId,
                     ["AuthorName"] = message.AuthorName,
-                    ["ApplicationId"] = this._storageScope?.ApplicationId,
-                    ["AgentId"] = this._storageScope?.AgentId,
-                    ["UserId"] = this._storageScope?.UserId,
-                    ["SessionId"] = this._storageScope?.SessionId,
+                    ["ApplicationId"] = storageScope.ApplicationId,
+                    ["AgentId"] = storageScope.AgentId,
+                    ["UserId"] = storageScope.UserId,
+                    ["SessionId"] = storageScope.SessionId,
                     ["Content"] = message.Text,
                     ["CreatedAt"] = message.CreatedAt?.ToString("O") ?? DateTimeOffset.UtcNow.ToString("O"),
                     ["ContentEmbedding"] = message.Text,
@@ -275,10 +253,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
                 this._logger.LogError(
                     ex,
                     "ChatHistoryMemoryProvider: Failed to add messages to chat history vector store due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', SessionId: '{SessionId}', UserId: '{UserId}'.",
-                    this._searchScope.ApplicationId,
-                    this._searchScope.AgentId,
-                    this._searchScope.SessionId,
-                    this.SanitizeLogData(this._searchScope.UserId));
+                    storageScope.ApplicationId,
+                    storageScope.AgentId,
+                    storageScope.SessionId,
+                    this.SanitizeLogData(storageScope.UserId));
             }
         }
     }
@@ -287,16 +265,17 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
     /// Function callable by the AI model (when enabled) to perform an ad-hoc chat history search.
     /// </summary>
     /// <param name="userQuestion">The query text.</param>
+    /// <param name="searchScope">The scope to filter search results with.</param>
     /// <param name="cancellationToken">Cancellation token.</param>
     /// <returns>Formatted search results (may be empty).</returns>
-    internal async Task<string> SearchTextAsync(string userQuestion, CancellationToken cancellationToken = default)
+    private async Task<string> SearchTextAsync(string userQuestion, ChatHistoryMemoryProviderScope searchScope, CancellationToken cancellationToken = default)
     {
         if (string.IsNullOrWhiteSpace(userQuestion))
         {
             return string.Empty;
         }
 
-        var results = await this.SearchChatHistoryAsync(userQuestion, this._maxResults, cancellationToken).ConfigureAwait(false);
+        var results = await this.SearchChatHistoryAsync(userQuestion, searchScope, this._maxResults, cancellationToken).ConfigureAwait(false);
         if (!results.Any())
         {
             return string.Empty;
@@ -317,10 +296,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
                 "ChatHistoryMemoryProvider: Search Results\nInput:{Input}\nOutput:{MessageText}\n ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', SessionId: '{SessionId}', UserId: '{UserId}'.",
                 this.SanitizeLogData(userQuestion),
                 this.SanitizeLogData(formatted),
-                this._searchScope.ApplicationId,
-                this._searchScope.AgentId,
-                this._searchScope.SessionId,
-                this.SanitizeLogData(this._searchScope.UserId));
+                searchScope.ApplicationId,
+                searchScope.AgentId,
+                searchScope.SessionId,
+                this.SanitizeLogData(searchScope.UserId));
         }
 
         return formatted;
@@ -330,11 +309,13 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
     /// Searches for relevant chat history items based on the provided query text.
     /// </summary>
     /// <param name="queryText">The text to search for.</param>
+    /// <param name="searchScope">The scope to filter search results with.</param>
     /// <param name="top">The maximum number of results to return.</param>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>A list of relevant chat history items.</returns>
     private async Task<IEnumerable<Dictionary<string, object?>>> SearchChatHistoryAsync(
         string queryText,
+        ChatHistoryMemoryProviderScope searchScope,
         int top,
         CancellationToken cancellationToken = default)
     {
@@ -345,10 +326,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
 
         var collection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false);
 
-        string? applicationId = this._searchScope.ApplicationId;
-        string? agentId = this._searchScope.AgentId;
-        string? userId = this._searchScope.UserId;
-        string? sessionId = this._searchScope.SessionId;
+        string? applicationId = searchScope.ApplicationId;
+        string? agentId = searchScope.AgentId;
+        string? userId = searchScope.UserId;
+        string? sessionId = searchScope.SessionId;
 
         Expression<Func<Dictionary<string, object?>, bool>>? filter = null;
         if (applicationId != null)
@@ -401,10 +382,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
             this._logger.LogInformation(
                 "ChatHistoryMemoryProvider: Retrieved {Count} search results. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', SessionId: '{SessionId}', UserId: '{UserId}'.",
                 results.Count,
-                this._searchScope.ApplicationId,
-                this._searchScope.AgentId,
-                this._searchScope.SessionId,
-                this.SanitizeLogData(this._searchScope.UserId));
+                searchScope.ApplicationId,
+                searchScope.AgentId,
+                searchScope.SessionId,
+                this.SanitizeLogData(searchScope.UserId));
         }
 
         return results;
@@ -465,39 +446,32 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
         GC.SuppressFinalize(this);
     }
 
-    /// <summary>
-    /// Serializes the current provider state to a <see cref="JsonElement"/> including storage and search scopes.
-    /// </summary>
-    /// <param name="jsonSerializerOptions">Optional serializer options.</param>
-    /// <returns>Serialized provider state.</returns>
-    public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-        var state = new ChatHistoryMemoryProviderState
-        {
-            StorageScope = this._storageScope,
-            SearchScope = this._searchScope,
-        };
-
-        var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
-        return JsonSerializer.SerializeToElement(state, jso.GetTypeInfo(typeof(ChatHistoryMemoryProviderState)));
-    }
-
-    private static ChatHistoryMemoryProviderState? DeserializeState(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions)
-    {
-        if (serializedState.ValueKind != JsonValueKind.Object)
-        {
-            return null;
-        }
-
-        var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
-        return serializedState.Deserialize(jso.GetTypeInfo(typeof(ChatHistoryMemoryProviderState))) as ChatHistoryMemoryProviderState;
-    }
-
     private string? SanitizeLogData(string? data) => this._enableSensitiveTelemetryData ? data : "<redacted>";
 
-    internal sealed class ChatHistoryMemoryProviderState
+    /// <summary>
+    /// Represents the state of a <see cref="ChatHistoryMemoryProvider"/> stored in the <see cref="AgentSession.StateBag"/>.
+    /// </summary>
+    public sealed class State
     {
-        public ChatHistoryMemoryProviderScope? StorageScope { get; set; }
-        public ChatHistoryMemoryProviderScope? SearchScope { get; set; }
+        /// <summary>
+        /// Initializes a new instance of the <see cref="State"/> class with the specified storage and search scopes.
+        /// </summary>
+        /// <param name="storageScope">The scope to use when storing chat history messages.</param>
+        /// <param name="searchScope">The scope to use when searching for relevant chat history messages. If null, the storage scope will be used for searching as well.</param>
+        public State(ChatHistoryMemoryProviderScope storageScope, ChatHistoryMemoryProviderScope? searchScope = null)
+        {
+            this.StorageScope = Throw.IfNull(storageScope);
+            this.SearchScope = searchScope ?? storageScope;
+        }
+
+        /// <summary>
+        /// Gets or sets the scope used when storing chat history messages.
+        /// </summary>
+        public ChatHistoryMemoryProviderScope StorageScope { get; }
+
+        /// <summary>
+        /// Gets or sets the scope used when searching chat history messages.
+        /// </summary>
+        public ChatHistoryMemoryProviderScope SearchScope { get; }
     }
 }

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

@@ -44,6 +44,15 @@ public sealed class ChatHistoryMemoryProviderOptions
     /// <value>Defaults to <see langword="false"/>.</value>
     public bool EnableSensitiveTelemetryData { get; set; }
 
+    /// <summary>
+    /// 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
+    /// <see cref="ChatHistoryMemoryProvider"/> instances with separate state in the same session.
+    /// </value>
+    public string? StateKey { get; set; }
+
     /// <summary>
     /// Behavior choices for the provider.
     /// </summary>

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

@@ -4,7 +4,6 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
-using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -39,31 +38,28 @@ 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;
     private readonly AITool[] _tools;
-    private readonly Queue<string> _recentMessagesText;
     private readonly List<ChatRole> _recentMessageRolesIncluded;
     private readonly int _recentMessageMemoryLimit;
     private readonly TextSearchProviderOptions.TextSearchBehavior _searchTime;
     private readonly string _contextPrompt;
     private readonly string _citationsPrompt;
+    private readonly string _stateKey;
     private readonly Func<IList<TextSearchResult>, string>? _contextFormatter;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="TextSearchProvider"/> class.
     /// </summary>
     /// <param name="searchAsync">Delegate that executes the search logic. Must not be <see langword="null"/>.</param>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized provider state.</param>
-    /// <param name="jsonSerializerOptions">Optional serializer options (unused - source generated context is used).</param>
     /// <param name="options">Optional configuration options.</param>
     /// <param name="loggerFactory">Optional logger factory.</param>
     /// <exception cref="ArgumentNullException">Thrown when <paramref name="searchAsync"/> is <see langword="null"/>.</exception>
     public TextSearchProvider(
         Func<string, CancellationToken, Task<IEnumerable<TextSearchResult>>> searchAsync,
-        JsonElement serializedState,
-        JsonSerializerOptions? jsonSerializerOptions = null,
         TextSearchProviderOptions? options = null,
         ILoggerFactory? loggerFactory = null)
     {
@@ -75,27 +71,9 @@ 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._contextFormatter = options?.ContextFormatter;
 
-        // Restore recent messages from serialized state if provided
-        List<string>? restoredMessages = null;
-        if (serializedState.ValueKind is JsonValueKind.Null or JsonValueKind.Undefined)
-        {
-            this._recentMessagesText = new();
-        }
-        else
-        {
-            var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
-            var state = serializedState.Deserialize(jso.GetTypeInfo(typeof(TextSearchProviderState))) as TextSearchProviderState;
-            if (state?.RecentMessagesText is { Count: > 0 })
-            {
-                restoredMessages = state.RecentMessagesText;
-            }
-
-            // Restore recent messages respecting the limit (may truncate if limit changed afterwards).
-            this._recentMessagesText = restoredMessages is null ? new() : new(restoredMessages.Take(this._recentMessageMemoryLimit));
-        }
-
         // Create the on-demand search tool (only used if behavior is OnDemandFunctionCalling)
         this._tools =
         [
@@ -115,12 +93,16 @@ public sealed class TextSearchProvider : AIContextProvider
             return new AIContext { Tools = this._tools }; // No automatic message injection.
         }
 
+        // Retrieve recent messages from the session state bag.
+        var recentMessagesText = context.Session?.StateBag.GetValue<TextSearchProviderState>(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText
+            ?? [];
+
         // Aggregate text from memory + current request messages.
         var sbInput = new StringBuilder();
         var requestMessagesText = context.RequestMessages
             .Where(m => m.GetAgentRequestMessageSource() == AgentRequestMessageSourceType.External)
             .Where(x => !string.IsNullOrWhiteSpace(x?.Text)).Select(x => x.Text);
-        foreach (var messageText in this._recentMessagesText.Concat(requestMessagesText))
+        foreach (var messageText in recentMessagesText.Concat(requestMessagesText))
         {
             if (sbInput.Length > 0)
             {
@@ -176,12 +158,21 @@ public sealed class TextSearchProvider : AIContextProvider
             return default; // Memory disabled.
         }
 
+        if (context.Session is null)
+        {
+            return default; // No session to store state in.
+        }
+
         if (context.InvokeException is not null)
         {
             return default; // Do not update memory on failed invocations.
         }
 
-        var messagesText = context.RequestMessages
+        // Retrieve existing recent messages from the session state bag.
+        var recentMessagesText = context.Session.StateBag.GetValue<TextSearchProviderState>(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText
+            ?? [];
+
+        var newMessagesText = context.RequestMessages
             .Where(m => m.GetAgentRequestMessageSource() == AgentRequestMessageSourceType.External)
             .Concat(context.ResponseMessages ?? [])
             .Where(m =>
@@ -190,44 +181,23 @@ public sealed class TextSearchProvider : AIContextProvider
                 // Filter out any messages that were added by this class in InvokingAsync, since we don't want
                 // a feedback loop where previous search results are used to find new search results.
                 (m.AdditionalProperties == null || m.AdditionalProperties.TryGetValue("IsTextSearchProviderOutput", out bool isTextSearchProviderOutput) == false || !isTextSearchProviderOutput))
-            .Select(m => m.Text)
-            .ToList();
-        if (messagesText.Count > limit)
-        {
-            // If the current request/response exceeds the limit, only keep the most recent messages from it.
-            messagesText = messagesText.Skip(messagesText.Count - limit).ToList();
-        }
+            .Select(m => m.Text);
 
-        foreach (var message in messagesText)
-        {
-            this._recentMessagesText.Enqueue(message);
-        }
+        // Combine existing messages with new messages, then take the most recent up to the limit.
+        var allMessages = recentMessagesText.Concat(newMessagesText).ToList();
+        var updatedMessages = allMessages.Count > limit
+            ? allMessages.Skip(allMessages.Count - limit).ToList()
+            : allMessages;
 
-        while (this._recentMessagesText.Count > limit)
-        {
-            this._recentMessagesText.Dequeue();
-        }
+        // Store updated state back to the session state bag.
+        context.Session.StateBag.SetValue(
+            this._stateKey,
+            new TextSearchProviderState { RecentMessagesText = updatedMessages },
+            AgentJsonUtilities.DefaultOptions);
 
         return default;
     }
 
-    /// <summary>
-    /// Serializes the current provider state to a <see cref="JsonElement"/> containing any overridden prompts or descriptions.
-    /// </summary>
-    /// <param name="jsonSerializerOptions">Optional serializer options (ignored, source generated context is used).</param>
-    /// <returns>A <see cref="JsonElement"/> with overridden values, or default if nothing was overridden.</returns>
-    public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-    {
-        // Only persist values that differ from defaults plus recent memory configuration & messages.
-        TextSearchProviderState state = new();
-        if (this._recentMessageMemoryLimit > 0 && this._recentMessagesText.Count > 0)
-        {
-            state.RecentMessagesText = this._recentMessagesText.Take(this._recentMessageMemoryLimit).ToList();
-        }
-
-        return JsonSerializer.SerializeToElement(state, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(TextSearchProviderState)));
-    }
-
     /// <summary>
     /// Function callable by the AI model (when enabled) to perform an ad-hoc search.
     /// </summary>

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

@@ -59,6 +59,15 @@ public sealed class TextSearchProviderOptions
     /// </value>
     public int RecentMessageMemoryLimit { get; set; }
 
+    /// <summary>
+    /// 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
+    /// <see cref="TextSearchProvider"/> instances with separate state in the same session.
+    /// </value>
+    public string? StateKey { get; set; }
+
     /// <summary>
     /// Gets or sets the list of <see cref="ChatRole"/> types to filter recent messages to
     /// when deciding which recent messages to include when constructing the search input.