.NET: Update AgentThread, MessageStores and Context Providers to deserialize via constructor. (#779)

* Update AgentThread, MessageStores and Context Providers to deserialize via constructor.

* Fix pr comment.

* Add additional validation for AgentThread deserialization

* Update WorkflowMessageStore desreialize to improve error checking.

* Reduce allocations in InMemoryChatMessageStore

---------

Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>

dotnet/src/Microsoft.Agents.Workflows/Specialized/AIAgentHostExecutor.cs

@@ -77,8 +77,7 @@ internal class AIAgentHostExecutor : Executor
         JsonElement? threadValue = await context.ReadStateAsync<JsonElement?>(ThreadStateKey).ConfigureAwait(false);
         if (threadValue.HasValue)
         {
-            this._thread = await this._agent.DeserializeThreadAsync(threadValue.Value, cancellationToken: cancellation)
-                                            .ConfigureAwait(false);
+            this._thread = this._agent.DeserializeThread(threadValue.Value, cancellationToken: cancellation);
         }
 
         JsonElement? messagesValue = await context.ReadStateAsync<JsonElement?>(PendingMessagesStateKey).ConfigureAwait(false);

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

@@ -6,6 +6,7 @@ using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
 using System.Runtime.CompilerServices;
+using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -47,6 +48,9 @@ internal class WorkflowHostAgent : AIAgent
 
     public override AgentThread GetNewThread() => new WorkflowThread(this.Id, this.Name, this.GenerateNewId());
 
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        => new WorkflowThread(serializedThread, jsonSerializerOptions);
+
     private async
     IAsyncEnumerable<AgentRunResponseUpdate> InvokeStageAsync(
         WorkflowThread conversation,

dotnet/src/Microsoft.Agents.Workflows/WorkflowMessageStore.cs

@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System;
 using System.Collections.Generic;
 using System.Text.Json;
 using System.Threading;
@@ -14,6 +15,30 @@ internal class WorkflowMessageStore : IChatMessageStore
     private int _bookmark = 0;
     private readonly List<ChatMessage> _chatMessages = new();
 
+    public WorkflowMessageStore()
+    {
+    }
+
+    public WorkflowMessageStore(JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
+    {
+        if (serializedStoreState.ValueKind != JsonValueKind.Object)
+        {
+            throw new ArgumentException("The provided JsonElement must be a json object", nameof(serializedStoreState));
+        }
+
+        StoreState? state =
+            JsonSerializer.Deserialize(
+                serializedStoreState,
+                AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState;
+
+        if (state?.Messages is not null)
+        {
+            this._chatMessages.AddRange(state.Messages);
+        }
+
+        this._bookmark = state?.Bookmark ?? 0;
+    }
+
     internal class StoreState
     {
         public int Bookmark { get; set; }
@@ -50,31 +75,6 @@ internal class WorkflowMessageStore : IChatMessageStore
         this._bookmark = this._chatMessages.Count;
     }
 
-    public ValueTask DeserializeStateAsync(JsonElement? serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-    {
-        if (serializedStoreState is null)
-        {
-            return default;
-        }
-
-        object? maybeState =
-            JsonSerializer.Deserialize(
-                serializedStoreState.Value,
-                AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState)));
-
-        if (maybeState is not StoreState state)
-        {
-            throw new JsonException("Invalid state format for WorkflowMessageStore.");
-        }
-
-        this._chatMessages.Clear();
-        this._chatMessages.AddRange(state.Messages);
-
-        this._bookmark = state.Bookmark;
-
-        return default;
-    }
-
     public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
     {
         StoreState state = new()

dotnet/src/Microsoft.Agents.Workflows/WorkflowThread.cs

@@ -25,6 +25,11 @@ internal class WorkflowThread : AgentThread
         this._workflowName = workflowName;
     }
 
+    public WorkflowThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    {
+        throw new NotImplementedException("Pending Checkpointing work.");
+    }
+
     public string RunId { get; }
     public int Halts { get; } = 0;
 
@@ -35,11 +40,6 @@ internal class WorkflowThread : AgentThread
         throw new NotImplementedException("Pending Checkpointing work.");
     }
 
-    protected override Task DeserializeAsync(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-    {
-        throw new NotImplementedException("Pending Checkpointing work.");
-    }
-
     public AgentRunResponseUpdate CreateUpdate(params AIContent[] parts)
     {
         Throw.IfNullOrEmpty(parts);

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/AIAgent.cs

@@ -90,12 +90,8 @@ public abstract class AIAgent
     /// <param name="jsonSerializerOptions">Optional <see cref="JsonSerializerOptions"/> to use for deserializing the thread state.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns>The deserialized <see cref="AgentThread"/> instance.</returns>
-    public async ValueTask<AgentThread> DeserializeThreadAsync(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-    {
-        var thread = this.GetNewThread();
-        await thread.DeserializeAsync(serializedThread, jsonSerializerOptions, cancellationToken).ConfigureAwait(false);
-        return thread;
-    }
+    public virtual AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        => new(serializedThread, jsonSerializerOptions);
 
     /// <summary>
     /// Run the agent with no message assuming that all required instructions are already provided to the agent or on the thread.

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/AgentThread.cs

@@ -26,6 +26,46 @@ public class AgentThread
     {
     }
 
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AgentThread"/> class from serialized state.
+    /// </summary>
+    /// <param name="serializedThreadState">A <see cref="JsonElement"/> representing the serialized state of the thread.</param>
+    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
+    /// <param name="chatMessageStoreFactory">An optional factory function to create a custom <see cref="IChatMessageStore"/>.</param>
+    /// <param name="aiContextProviderFactory">An optional factory function to create a custom <see cref="AIContextProvider"/>.</param>
+    public AgentThread(
+        JsonElement serializedThreadState,
+        JsonSerializerOptions? jsonSerializerOptions = null,
+        Func<JsonElement, JsonSerializerOptions?, IChatMessageStore>? chatMessageStoreFactory = null,
+        Func<JsonElement, JsonSerializerOptions?, AIContextProvider>? aiContextProviderFactory = null)
+    {
+        if (serializedThreadState.ValueKind != JsonValueKind.Object)
+        {
+            throw new ArgumentException("The serialized thread state must be a JSON object.", nameof(serializedThreadState));
+        }
+
+        var state = JsonSerializer.Deserialize(
+            serializedThreadState,
+            AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ThreadState))) as ThreadState;
+
+        this.AIContextProvider = aiContextProviderFactory?.Invoke(state?.AIContextProviderState ?? default, jsonSerializerOptions);
+
+        if (state?.ConversationId is string threadId)
+        {
+            this.ConversationId = threadId;
+
+            // Since we have an ID, we should not have a chat message store and we can return here.
+            return;
+        }
+
+        this._messageStore = chatMessageStoreFactory?.Invoke(state?.StoreState ?? default, jsonSerializerOptions);
+        if (this._messageStore is null)
+        {
+            // If we didn't get a custom store, create an in-memory one.
+            this._messageStore = new InMemoryChatMessageStore(state?.StoreState ?? default, jsonSerializerOptions);
+        }
+    }
+
     /// <summary>
     /// Gets or sets the ID of the underlying service thread to support cases where the chat history is stored by the agent service.
     /// </summary>
@@ -175,47 +215,6 @@ public class AgentThread
         }
     }
 
-    /// <summary>
-    /// Deserializes the state contained in the provided <see cref="JsonElement"/> into the properties on this thread.
-    /// </summary>
-    /// <param name="serializedThread">A <see cref="JsonElement"/> representing the state of the thread.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A <see cref="ValueTask"/> that completes when the state has been deserialized.</returns>
-    protected internal virtual async Task DeserializeAsync(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-    {
-        var state = JsonSerializer.Deserialize(
-            serializedThread,
-            AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ThreadState))) as ThreadState;
-
-        if (state?.ConversationId is string threadId)
-        {
-            this.ConversationId = threadId;
-
-            // Since we have an ID, we should not have a chat message store and we can return here.
-            return;
-        }
-
-        if (state?.AIContextProviderState.HasValue is true && this.AIContextProvider is not null)
-        {
-            await this.AIContextProvider.DeserializeAsync(state.AIContextProviderState.Value, jsonSerializerOptions, cancellationToken).ConfigureAwait(false);
-        }
-
-        // If we don't have any IChatMessageStore state return here.
-        if (state?.StoreState is null || state?.StoreState.Value.ValueKind is JsonValueKind.Undefined or JsonValueKind.Null)
-        {
-            return;
-        }
-
-        if (this._messageStore is null)
-        {
-            // If we don't have a chat message store yet, create an in-memory one.
-            this._messageStore = new InMemoryChatMessageStore();
-        }
-
-        await this._messageStore.DeserializeStateAsync(state!.StoreState.Value, jsonSerializerOptions, cancellationToken).ConfigureAwait(false);
-    }
-
     internal sealed class ThreadState
     {
         public string? ConversationId { get; set; }

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/DelegatingAIAgent.cs

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Shared.Diagnostics;
@@ -52,6 +53,10 @@ public class DelegatingAIAgent : AIAgent
     /// <inheritdoc />
     public override AgentThread GetNewThread() => this.InnerAgent.GetNewThread();
 
+    /// <inheritdoc />
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        => this.InnerAgent.DeserializeThread(serializedThread, jsonSerializerOptions, cancellationToken);
+
     /// <inheritdoc />
     public override Task<AgentRunResponse> RunAsync(
         IEnumerable<ChatMessage> messages,

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/IChatMessageStore.cs

@@ -44,28 +44,11 @@ public interface IChatMessageStore
     /// <returns>An async task.</returns>
     Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default);
 
-    /// <summary>
-    /// Deserializes the state contained in the provided <see cref="JsonElement"/> into the properties on this store.
-    /// </summary>
-    /// <param name="serializedStoreState">A <see cref="JsonElement"/> representing the state of the store.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A <see cref="ValueTask"/> that completes when the state has been deserialized.</returns>
-    /// <remarks>
-    /// This method, together with <see cref="SerializeStateAsync(JsonSerializerOptions?, CancellationToken)"/> can be used to save and load messages from a persistent store
-    /// if this store only has messages in memory.
-    /// </remarks>
-    ValueTask DeserializeStateAsync(JsonElement? serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, 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>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
-    /// <remarks>
-    /// This method, together with <see cref="DeserializeStateAsync(JsonElement?, JsonSerializerOptions?, CancellationToken)"/> can be used to save and load messages from a persistent store
-    /// if this store only has messages in memory.
-    /// </remarks>
     ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default);
 }

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/InMemoryChatMessageStore.cs

@@ -17,12 +17,23 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
 {
     private readonly IChatReducer? _chatReducer;
     private readonly ChatReducerTriggerEvent _reducerTriggerEvent;
-    private List<ChatMessage> _messages = new();
+    private List<ChatMessage> _messages;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="InMemoryChatMessageStore"/> class.
     /// </summary>
     public InMemoryChatMessageStore()
+    {
+        this._messages = new();
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="InMemoryChatMessageStore"/> class, with an existing state from a serialized JSON element.
+    /// </summary>
+    /// <param name="serializedStoreState">A <see cref="JsonElement"/> representing the serialized state of the store.</param>
+    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
+    public InMemoryChatMessageStore(JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
+        : this(null, serializedStoreState, jsonSerializerOptions, ChatReducerTriggerEvent.BeforeMessagesRetrieval)
     {
     }
 
@@ -32,9 +43,40 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
     /// <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="reducerTriggerEvent">The event that should trigger the reducer invocation.</param>
     public InMemoryChatMessageStore(IChatReducer chatReducer, ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent.BeforeMessagesRetrieval)
+        : this(chatReducer, default, null, reducerTriggerEvent)
     {
-        this._chatReducer = Throw.IfNull(chatReducer);
+        Throw.IfNull(chatReducer);
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="InMemoryChatMessageStore"/> 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="serializedStoreState">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="reducerTriggerEvent">The event that should trigger the reducer invocation.</param>
+    public InMemoryChatMessageStore(IChatReducer? chatReducer, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent.BeforeMessagesRetrieval)
+    {
+        this._chatReducer = chatReducer;
         this._reducerTriggerEvent = reducerTriggerEvent;
+
+        if (serializedStoreState.ValueKind != JsonValueKind.Object)
+        {
+            this._messages = new();
+            return;
+        }
+
+        var state = JsonSerializer.Deserialize(
+            serializedStoreState,
+            AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState;
+
+        if (state?.Messages is { Count: > 0 } messages)
+        {
+            this._messages = messages;
+            return;
+        }
+
+        this._messages = new();
     }
 
     /// <summary>
@@ -84,26 +126,6 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
         return this._messages;
     }
 
-    /// <inheritdoc />
-    public ValueTask DeserializeStateAsync(JsonElement? serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-    {
-        if (serializedStoreState is null)
-        {
-            return default;
-        }
-
-        var state = JsonSerializer.Deserialize(
-            serializedStoreState.Value,
-            AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState;
-
-        if (state?.Messages is { Count: > 0 } messages)
-        {
-            this._messages.AddRange(messages);
-        }
-
-        return default;
-    }
-
     /// <inheritdoc />
     public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
     {
@@ -157,7 +179,7 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
 
     internal sealed class StoreState
     {
-        public IList<ChatMessage> Messages { get; set; } = new List<ChatMessage>();
+        public List<ChatMessage> Messages { get; set; } = new List<ChatMessage>();
     }
 
     /// <summary>

dotnet/src/Microsoft.Extensions.AI.Agents.Hosting/AgentActor.cs

@@ -39,7 +39,7 @@ internal sealed class AgentActor(
             if (threadResult.Value is { } threadJson)
             {
                 // Deserialize the thread state if it exists
-                this._thread = await agent.DeserializeThreadAsync(threadJson, cancellationToken: cancellationToken).ConfigureAwait(false);
+                this._thread = agent.DeserializeThread(threadJson, cancellationToken: cancellationToken);
                 hasExistingThread = true;
             }
         }

dotnet/src/Microsoft.Extensions.AI.Agents.Hosting/AgentProxy.cs

@@ -36,6 +36,10 @@ public sealed class AgentProxy : AIAgent
     /// <inheritdoc/>
     public override AgentThread GetNewThread() => new AgentProxyThread();
 
+    /// <inheritdoc/>
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        => new AgentProxyThread(serializedThread, jsonSerializerOptions);
+
     /// <summary>
     /// Gets a thread by its <see cref="AgentThread.ConversationId"/>.
     /// </summary>

dotnet/src/Microsoft.Extensions.AI.Agents.Hosting/AgentProxyThread.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Text.Json;
 using System.Text.RegularExpressions;
 using Microsoft.Shared.Diagnostics;
 
@@ -55,6 +56,16 @@ internal sealed partial class AgentProxyThread : AgentThread
     {
     }
 
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AgentProxyThread"/> class from serialized state.
+    /// </summary>
+    /// <param name="serializedThreadState">A <see cref="JsonElement"/> representing the serialized state of the thread.</param>
+    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
+    public AgentProxyThread(JsonElement serializedThreadState, JsonSerializerOptions? jsonSerializerOptions = null)
+        : base(serializedThreadState, jsonSerializerOptions)
+    {
+    }
+
     internal static string CreateId() => Guid.NewGuid().ToString("N");
 
     /// <summary>

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

@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System.Text.Json;
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.AI.Agents;
 using Microsoft.Extensions.Logging;
@@ -83,6 +84,10 @@ public class OpenAIChatClientAgent : AIAgent
     public sealed override AgentThread GetNewThread()
         => this._chatClientAgent.GetNewThread();
 
+    /// <inheritdoc/>
+    public sealed override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        => this._chatClientAgent.DeserializeThread(serializedThread, jsonSerializerOptions, cancellationToken);
+
     /// <inheritdoc/>
     public sealed override Task<AgentRunResponse> RunAsync(
         IEnumerable<Microsoft.Extensions.AI.ChatMessage> messages,

dotnet/src/Microsoft.Extensions.AI.Agents/ChatCompletion/ChatClientAgent.cs

@@ -4,6 +4,7 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Runtime.CompilerServices;
+using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
@@ -241,12 +242,16 @@ public sealed class ChatClientAgent : AIAgent
     {
         var thread = new AgentThread
         {
-            MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke(),
-            AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke()
+            MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke(default, null),
+            AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke(default, null)
         };
         return thread;
     }
 
+    /// <inheritdoc/>
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+        => new(serializedThread, jsonSerializerOptions, this._agentOptions?.ChatMessageStoreFactory, this._agentOptions?.AIContextProviderFactory);
+
     #region Private
 
     /// <summary>
@@ -493,7 +498,7 @@ public sealed class ChatClientAgent : AIAgent
             // If the service doesn't use service side thread storage (i.e. we got no id back from invocation), and
             // the thread has no MessageStore yet, and we have a custom messages store, we should update the thread
             // with the custom MessageStore so that it has somewhere to store the chat history.
-            thread.MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke();
+            thread.MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke(default, null);
         }
     }
 

dotnet/src/Microsoft.Extensions.AI.Agents/ChatCompletion/ChatClientAgentOptions.cs

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Text.Json;
 
 namespace Microsoft.Extensions.AI.Agents;
 
@@ -78,14 +79,14 @@ public class ChatClientAgentOptions
     /// Gets or sets a factory function to create an instance of <see cref="IChatMessageStore"/>
     /// which will be used to store chat messages for this agent.
     /// </summary>
-    public Func<IChatMessageStore>? ChatMessageStoreFactory { get; set; }
+    public Func<JsonElement, JsonSerializerOptions?, IChatMessageStore>? ChatMessageStoreFactory { 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.
     /// </summary>
-    public Func<AIContextProvider>? AIContextProviderFactory { get; set; }
+    public Func<JsonElement, JsonSerializerOptions?, AIContextProvider>? AIContextProviderFactory { get; set; }
 
     /// <summary>
     /// Gets or sets a value indicating whether to use the provided <see cref="IChatClient"/> instance as is,