.NET: [BREAKING] Refactor ChatMessageStore methods to be similar to AIContextProvider and add filtering support (#2604)

* Refactor ChatMessageStore methods to be similar to AIContextProvider

* Fix file encoding

* Ensure that AIContextProvider messages area also persisted.

* Update formatting and seal context classes

* Improve formatting

* Remove optional messages from constructor and add unit test

* Add ChatMessageStore filtering via a decorator

* Update sample and cosmos message store to store AIContextProvider messages in right order. Fix unit tests.

* Update Workflowmessage store to use aicontext provider messages.

* Apply suggestions from code review

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

* Apply suggestions from code review

Co-authored-by: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com>

* Improve xml docs messaging

* Address code review comments.

* Also notify message store on failure

---------

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

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

@@ -32,8 +32,9 @@ namespace Microsoft.Agents.AI;
 public abstract class ChatMessageStore
 {
     /// <summary>
-    /// Asynchronously retrieves all messages from the store that should be provided as context for the next agent invocation.
+    /// Called at the start of agent invocation to retrieve all messages from the store that should be provided as context for the next agent invocation.
     /// </summary>
+    /// <param name="context">Contains the request context including the caller provided messages that will be used by the agent for this invocation.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns>
     /// A task that represents the asynchronous operation. The task result contains a collection of <see cref="ChatMessage"/>
@@ -59,20 +60,19 @@ public abstract class ChatMessageStore
     /// and context management.
     /// </para>
     /// </remarks>
-    public abstract Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default);
+    public abstract ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default);
 
     /// <summary>
-    /// Asynchronously adds new messages to the store.
+    /// Called at the end of the agent invocation to add new messages to the store.
     /// </summary>
-    /// <param name="messages">The collection of chat messages to add to the store.</param>
+    /// <param name="context">Contains the invocation context including request messages, response messages, and any exception that occurred.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns>A task that represents the asynchronous add operation.</returns>
-    /// <exception cref="ArgumentNullException"><paramref name="messages"/> is <see langword="null"/>.</exception>
     /// <remarks>
     /// <para>
     /// Messages should be added in the order they were generated to maintain proper chronological sequence.
     /// The store is responsible for preserving message ordering and ensuring that subsequent calls to
-    /// <see cref="GetMessagesAsync"/> return messages in the correct chronological order.
+    /// <see cref="InvokingAsync"/> return messages in the correct chronological order.
     /// </para>
     /// <para>
     /// Implementations may perform additional processing during message addition, such as:
@@ -83,8 +83,12 @@ public abstract class ChatMessageStore
     /// <item><description>Updating indices or search capabilities</description></item>
     /// </list>
     /// </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>
     /// </remarks>
-    public abstract Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default);
+    public abstract ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default);
 
     /// <summary>
     /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
@@ -121,4 +125,100 @@ public abstract class ChatMessageStore
     /// </remarks>
     public TService? GetService<TService>(object? serviceKey = null)
         => this.GetService(typeof(TService), serviceKey) is TService service ? service : default;
+
+    /// <summary>
+    /// Contains the context information provided to <see cref="InvokingAsync(InvokingContext, CancellationToken)"/>.
+    /// </summary>
+    /// <remarks>
+    /// This class provides context about the invocation before the messages are retrieved from the store,
+    /// including the new messages that will be used. Stores can use this information to determine what
+    /// messages should be retrieved for the invocation.
+    /// </remarks>
+    public sealed class InvokingContext
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="InvokingContext"/> class with the specified request messages.
+        /// </summary>
+        /// <param name="requestMessages">The new messages to be used by the agent for this invocation.</param>
+        /// <exception cref="ArgumentNullException"><paramref name="requestMessages"/> is <see langword="null"/>.</exception>
+        public InvokingContext(IEnumerable<ChatMessage> requestMessages)
+        {
+            this.RequestMessages = requestMessages ?? throw new ArgumentNullException(nameof(requestMessages));
+        }
+
+        /// <summary>
+        /// Gets the caller provided messages that will be used by the agent for this invocation.
+        /// </summary>
+        /// <value>
+        /// A collection of <see cref="ChatMessage"/> instances representing new messages that were provided by the caller.
+        /// </value>
+        public IEnumerable<ChatMessage> RequestMessages { get; }
+    }
+
+    /// <summary>
+    /// Contains the context information provided to <see cref="InvokedAsync(InvokedContext, CancellationToken)"/>.
+    /// </summary>
+    /// <remarks>
+    /// This class provides context about a completed agent invocation, including both the
+    /// request messages that were used and the response messages that were generated. It also indicates
+    /// whether the invocation succeeded or failed.
+    /// </remarks>
+    public sealed class InvokedContext
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="InvokedContext"/> class with the specified request messages.
+        /// </summary>
+        /// <param name="requestMessages">The caller provided messages that were used by the agent for this invocation.</param>
+        /// <param name="chatMessageStoreMessages">The messages retrieved from the <see cref="ChatMessageStore"/> for this invocation.</param>
+        /// <exception cref="ArgumentNullException"><paramref name="requestMessages"/> is <see langword="null"/>.</exception>
+        public InvokedContext(IEnumerable<ChatMessage> requestMessages, IEnumerable<ChatMessage> chatMessageStoreMessages)
+        {
+            this.RequestMessages = Throw.IfNull(requestMessages);
+            this.ChatMessageStoreMessages = chatMessageStoreMessages;
+        }
+
+        /// <summary>
+        /// Gets the caller provided messages that were used by the agent for this invocation.
+        /// </summary>
+        /// <value>
+        /// A collection of <see cref="ChatMessage"/> instances representing new messages that were provided by the caller.
+        /// This does not include any <see cref="ChatMessageStore"/> supplied messages.
+        /// </value>
+        public IEnumerable<ChatMessage> RequestMessages { get; }
+
+        /// <summary>
+        /// Gets the messages retrieved from the <see cref="ChatMessageStore"/> for this invocation, if any.
+        /// </summary>
+        /// <value>
+        /// A collection of <see cref="ChatMessage"/> instances that were retrieved from the <see cref="ChatMessageStore"/>,
+        /// and were used by the agent as part of the invocation.
+        /// </value>
+        public IEnumerable<ChatMessage> ChatMessageStoreMessages { get; }
+
+        /// <summary>
+        /// Gets or sets the messages provided by the <see cref="AIContextProvider"/> for this invocation, if any.
+        /// </summary>
+        /// <value>
+        /// A collection of <see cref="ChatMessage"/> instances that were provided by the <see cref="AIContextProvider"/>,
+        /// and were used by the agent as part of the invocation.
+        /// </value>
+        public IEnumerable<ChatMessage>? AIContextProviderMessages { get; set; }
+
+        /// <summary>
+        /// Gets the collection of response messages generated during this invocation if the invocation succeeded.
+        /// </summary>
+        /// <value>
+        /// A collection of <see cref="ChatMessage"/> instances representing the response,
+        /// or <see langword="null"/> if the invocation failed or did not produce response messages.
+        /// </value>
+        public IEnumerable<ChatMessage>? ResponseMessages { get; set; }
+
+        /// <summary>
+        /// Gets the <see cref="Exception"/> that was thrown during the invocation, if the invocation failed.
+        /// </summary>
+        /// <value>
+        /// The exception that caused the invocation to fail, or <see langword="null"/> if the invocation succeeded.
+        /// </value>
+        public Exception? InvokeException { get; set; }
+    }
 }

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

@@ -0,0 +1,50 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using Microsoft.Extensions.AI;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Contains extension methods for the <see cref="ChatMessageStore"/> class.
+/// </summary>
+public static class ChatMessageStoreExtensions
+{
+    /// <summary>
+    /// Adds message filtering to an existing store, so that messages passed to the store and messages produced by the store
+    /// can be filtered, updated or replaced.
+    /// </summary>
+    /// <param name="store">The store to add the message filter to.</param>
+    /// <param name="invokingMessagesFilter">An optional filter function to apply to messages produced by the store. If null, no filter is applied at this
+    /// stage.</param>
+    /// <param name="invokedMessagesFilter">An optional filter function to apply to the invoked context messages before they are passed to the store. If null, no
+    /// filter is applied at this stage.</param>
+    /// <returns>The <see cref="ChatMessageStore"/> with filtering applied.</returns>
+    public static ChatMessageStore WithMessageFilters(
+        this ChatMessageStore store,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? invokingMessagesFilter = null,
+        Func<ChatMessageStore.InvokedContext, ChatMessageStore.InvokedContext>? invokedMessagesFilter = null)
+    {
+        return new ChatMessageStoreMessageFilter(
+            innerChatMessageStore: store,
+            invokingMessagesFilter: invokingMessagesFilter,
+            invokedMessagesFilter: invokedMessagesFilter);
+    }
+
+    /// <summary>
+    /// Decorates the provided chat message store so that it does not store messages produced by any <see cref="AIContextProvider"/>.
+    /// </summary>
+    /// <param name="store">The store to add the message filter to.</param>
+    /// <returns>A new <see cref="ChatMessageStore"/> instance that filters out <see cref="AIContextProvider"/> messages so they do not get stored.</returns>
+    public static ChatMessageStore WithAIContextProviderMessageRemoval(this ChatMessageStore store)
+    {
+        return new ChatMessageStoreMessageFilter(
+            innerChatMessageStore: store,
+            invokedMessagesFilter: (ctx) =>
+            {
+                ctx.AIContextProviderMessages = null;
+                return ctx;
+            });
+    }
+}

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

@@ -0,0 +1,74 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// A <see cref="ChatMessageStore"/> decorator that allows filtering the messages
+/// passed into and out of an inner <see cref="ChatMessageStore"/>.
+/// </summary>
+public sealed class ChatMessageStoreMessageFilter : ChatMessageStore
+{
+    private readonly ChatMessageStore _innerChatMessageStore;
+    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? _invokingMessagesFilter;
+    private readonly Func<InvokedContext, InvokedContext>? _invokedMessagesFilter;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ChatMessageStoreMessageFilter"/> class.
+    /// </summary>
+    /// <remarks>Use this constructor to customize how messages are filtered before and after invocation by
+    /// providing appropriate filter functions. If no filters are provided, the message store operates without
+    /// additional filtering.</remarks>
+    /// <param name="innerChatMessageStore">The underlying chat message store to be wrapped. Cannot be null.</param>
+    /// <param name="invokingMessagesFilter">An optional filter function to apply to messages before they are invoked. If null, no filter is applied at this
+    /// stage.</param>
+    /// <param name="invokedMessagesFilter">An optional filter function to apply to the invocation context after messages have been invoked. If null, no
+    /// filter is applied at this stage.</param>
+    /// <exception cref="ArgumentNullException">Thrown if innerChatMessageStore is null.</exception>
+    public ChatMessageStoreMessageFilter(
+        ChatMessageStore innerChatMessageStore,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? invokingMessagesFilter = null,
+        Func<InvokedContext, InvokedContext>? invokedMessagesFilter = null)
+    {
+        this._innerChatMessageStore = Throw.IfNull(innerChatMessageStore);
+
+        if (invokingMessagesFilter == null && invokedMessagesFilter == null)
+        {
+            throw new ArgumentException("At least one filter function, invokingMessagesFilter or invokedMessagesFilter, must be provided.");
+        }
+
+        this._invokingMessagesFilter = invokingMessagesFilter;
+        this._invokedMessagesFilter = invokedMessagesFilter;
+    }
+
+    /// <inheritdoc />
+    public override async ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    {
+        var messages = await this._innerChatMessageStore.InvokingAsync(context, cancellationToken).ConfigureAwait(false);
+        return this._invokingMessagesFilter != null ? this._invokingMessagesFilter(messages) : messages;
+    }
+
+    /// <inheritdoc />
+    public override ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default)
+    {
+        if (this._invokedMessagesFilter != null)
+        {
+            context = this._invokedMessagesFilter(context);
+        }
+
+        return this._innerChatMessageStore.InvokedAsync(context, cancellationToken);
+    }
+
+    /// <inheritdoc />
+    public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
+    {
+        return this._innerChatMessageStore.Serialize(jsonSerializerOptions);
+    }
+}

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

@@ -134,21 +134,10 @@ public sealed class InMemoryChatMessageStore : ChatMessageStore, IList<ChatMessa
     }
 
     /// <inheritdoc />
-    public override async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
+    public override async ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {
-        _ = Throw.IfNull(messages);
+        _ = Throw.IfNull(context);
 
-        this._messages.AddRange(messages);
-
-        if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null)
-        {
-            this._messages = (await this.ChatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList();
-        }
-    }
-
-    /// <inheritdoc />
-    public override async Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default)
-    {
         if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null)
         {
             this._messages = (await this.ChatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList();
@@ -157,6 +146,26 @@ public sealed class InMemoryChatMessageStore : ChatMessageStore, IList<ChatMessa
         return this._messages;
     }
 
+    /// <inheritdoc />
+    public override async ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default)
+    {
+        _ = Throw.IfNull(context);
+
+        if (context.InvokeException is not null)
+        {
+            return;
+        }
+
+        // Add request, AI context provider, and response messages to the store
+        var allNewMessages = context.RequestMessages.Concat(context.AIContextProviderMessages ?? []).Concat(context.ResponseMessages ?? []);
+        this._messages.AddRange(allNewMessages);
+
+        if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null)
+        {
+            this._messages = (await this.ChatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList();
+        }
+    }
+
     /// <inheritdoc />
     public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
     {
@@ -221,7 +230,7 @@ public sealed class InMemoryChatMessageStore : ChatMessageStore, IList<ChatMessa
     {
         /// <summary>
         /// Trigger the reducer when a new message is added.
-        /// <see cref="AddMessagesAsync(IEnumerable{ChatMessage}, CancellationToken)"/> will only complete when reducer processing is done.
+        /// <see cref="InvokedAsync(InvokedContext, CancellationToken)"/> will only complete when reducer processing is done.
         /// </summary>
         AfterMessageAdded,
 

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

@@ -287,7 +287,7 @@ public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable
     }
 
     /// <inheritdoc />
-    public override async Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default)
+    public override async ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {
 #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
         if (this._disposed)
@@ -347,11 +347,14 @@ public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable
     }
 
     /// <inheritdoc />
-    public override async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
+    public override async ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
-        if (messages is null)
+        Throw.IfNull(context);
+
+        if (context.InvokeException is not null)
         {
-            throw new ArgumentNullException(nameof(messages));
+            // Do not store messages if there was an exception during invocation
+            return;
         }
 
 #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
@@ -361,7 +364,7 @@ public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable
         }
 #pragma warning restore CA1513
 
-        var messageList = messages as IReadOnlyCollection<ChatMessage> ?? messages.ToList();
+        var messageList = context.RequestMessages.Concat(context.AIContextProviderMessages ?? []).Concat(context.ResponseMessages ?? []).ToList();
         if (messageList.Count == 0)
         {
             return;
@@ -381,7 +384,7 @@ public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable
     /// <summary>
     /// Adds multiple messages using transactional batch operations for atomicity.
     /// </summary>
-    private async Task AddMessagesInBatchAsync(IReadOnlyCollection<ChatMessage> messages, CancellationToken cancellationToken)
+    private async Task AddMessagesInBatchAsync(List<ChatMessage> messages, CancellationToken cancellationToken)
     {
         var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds();
 

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

@@ -66,7 +66,7 @@ internal sealed class WorkflowHostAgent : AIAgent
     public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
         => new WorkflowThread(this._workflow, serializedThread, this._executionEnvironment, this._checkpointManager, jsonSerializerOptions);
 
-    private async ValueTask<WorkflowThread> UpdateThreadAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, CancellationToken cancellationToken = default)
+    private ValueTask<WorkflowThread> UpdateThreadAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, CancellationToken cancellationToken = default)
     {
         thread ??= this.GetNewThread();
 
@@ -75,8 +75,10 @@ internal sealed class WorkflowHostAgent : AIAgent
             throw new ArgumentException($"Incompatible thread type: {thread.GetType()} (expecting {typeof(WorkflowThread)})", nameof(thread));
         }
 
-        await workflowThread.MessageStore.AddMessagesAsync(messages, cancellationToken).ConfigureAwait(false);
-        return workflowThread;
+        // For workflow threads, messages are added directly via the internal AddMessages method
+        // The MessageStore methods are used for agent invocation scenarios
+        workflowThread.MessageStore.AddMessages(messages);
+        return new ValueTask<WorkflowThread>(workflowThread);
     }
 
     protected override async

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

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System.Collections.Generic;
+using System.Linq;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
@@ -45,15 +46,22 @@ internal sealed class WorkflowMessageStore : ChatMessageStore
 
     internal void AddMessages(params IEnumerable<ChatMessage> messages) => this._chatMessages.AddRange(messages);
 
-    public override Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
+    public override ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
+        => new(this._chatMessages.AsReadOnly());
+
+    public override ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
-        this._chatMessages.AddRange(messages);
+        if (context.InvokeException is not null)
+        {
+            return default;
+        }
 
-        return Task.CompletedTask;
+        var allNewMessages = context.RequestMessages.Concat(context.AIContextProviderMessages ?? []).Concat(context.ResponseMessages ?? []);
+        this._chatMessages.AddRange(allNewMessages);
+
+        return default;
     }
 
-    public override Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default) => Task.FromResult<IEnumerable<ChatMessage>>(this._chatMessages.AsReadOnly());
-
     public IEnumerable<ChatMessage> GetFromBookmark()
     {
         for (int i = this._bookmark; i < this._chatMessages.Count; i++)

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

@@ -201,7 +201,7 @@ public sealed partial class ChatClientAgent : AIAgent
     {
         var inputMessages = Throw.IfNull(messages) as IReadOnlyCollection<ChatMessage> ?? messages.ToList();
 
-        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> inputMessagesForChatClient, IList<ChatMessage>? aiContextProviderMessages) =
+        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> inputMessagesForChatClient, IList<ChatMessage>? aiContextProviderMessages, IList<ChatMessage>? chatMessageStoreMessages) =
             await this.PrepareThreadAndMessagesAsync(thread, inputMessages, options, cancellationToken).ConfigureAwait(false);
 
         ValidateStreamResumptionAllowed(chatOptions?.ContinuationToken, safeThread);
@@ -225,6 +225,7 @@ public sealed partial class ChatClientAgent : AIAgent
         }
         catch (Exception ex)
         {
+            await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             throw;
         }
@@ -239,6 +240,7 @@ public sealed partial class ChatClientAgent : AIAgent
         }
         catch (Exception ex)
         {
+            await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             throw;
         }
@@ -260,6 +262,7 @@ public sealed partial class ChatClientAgent : AIAgent
             }
             catch (Exception ex)
             {
+                await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
                 await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
                 throw;
             }
@@ -272,7 +275,7 @@ public sealed partial class ChatClientAgent : AIAgent
         this.UpdateThreadWithTypeAndConversationId(safeThread, chatResponse.ConversationId);
 
         // To avoid inconsistent state we only notify the thread of the input messages if no error occurs after the initial request.
-        await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages.Concat(aiContextProviderMessages ?? []).Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);
+        await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
 
         // Notify the AIContextProvider of all new messages.
         await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
@@ -379,7 +382,7 @@ public sealed partial class ChatClientAgent : AIAgent
     {
         var inputMessages = Throw.IfNull(messages) as IReadOnlyCollection<ChatMessage> ?? messages.ToList();
 
-        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> inputMessagesForChatClient, IList<ChatMessage>? aiContextProviderMessages) =
+        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> inputMessagesForChatClient, IList<ChatMessage>? aiContextProviderMessages, IList<ChatMessage>? chatMessageStoreMessages) =
             await this.PrepareThreadAndMessagesAsync(thread, inputMessages, options, cancellationToken).ConfigureAwait(false);
 
         var chatClient = this.ChatClient;
@@ -398,6 +401,7 @@ public sealed partial class ChatClientAgent : AIAgent
         }
         catch (Exception ex)
         {
+            await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             throw;
         }
@@ -415,7 +419,7 @@ public sealed partial class ChatClientAgent : AIAgent
         }
 
         // Only notify the thread of new messages if the chatResponse was successful to avoid inconsistent message state in the thread.
-        await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages.Concat(aiContextProviderMessages ?? []).Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);
+        await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
 
         // Notify the AIContextProvider of all new messages.
         await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
@@ -603,7 +607,14 @@ public sealed partial class ChatClientAgent : AIAgent
     /// <param name="runOptions">Optional parameters for agent invocation.</param>
     /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
     /// <returns>A tuple containing the thread, chat options, and thread messages.</returns>
-    private async Task<(ChatClientAgentThread AgentThread, ChatOptions? ChatOptions, List<ChatMessage> InputMessagesForChatClient, IList<ChatMessage>? AIContextProviderMessages)> PrepareThreadAndMessagesAsync(
+    private async Task
+        <(
+            ChatClientAgentThread AgentThread,
+            ChatOptions? ChatOptions,
+            List<ChatMessage> InputMessagesForChatClient,
+            IList<ChatMessage>? AIContextProviderMessages,
+            IList<ChatMessage>? ChatMessageStoreMessages
+        )> PrepareThreadAndMessagesAsync(
         AgentThread? thread,
         IEnumerable<ChatMessage> inputMessages,
         AgentRunOptions? runOptions,
@@ -637,6 +648,7 @@ public sealed partial class ChatClientAgent : AIAgent
 
         List<ChatMessage> inputMessagesForChatClient = [];
         IList<ChatMessage>? aiContextProviderMessages = null;
+        IList<ChatMessage>? chatMessageStoreMessages = null;
 
         // Populate the thread messages only if we are not continuing an existing response as it's not allowed
         if (chatOptions?.ContinuationToken is null)
@@ -644,7 +656,10 @@ public sealed partial class ChatClientAgent : AIAgent
             // Add any existing messages from the thread to the messages to be sent to the chat client.
             if (typedThread.MessageStore is not null)
             {
-                inputMessagesForChatClient.AddRange(await typedThread.MessageStore.GetMessagesAsync(cancellationToken).ConfigureAwait(false));
+                var invokingContext = new ChatMessageStore.InvokingContext(inputMessages);
+                var storeMessages = await typedThread.MessageStore.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
+                inputMessagesForChatClient.AddRange(storeMessages);
+                chatMessageStoreMessages = storeMessages as IList<ChatMessage> ?? storeMessages.ToList();
             }
 
             // If we have an AIContextProvider, we should get context from it, and update our
@@ -698,7 +713,7 @@ public sealed partial class ChatClientAgent : AIAgent
             chatOptions.ConversationId = typedThread.ConversationId;
         }
 
-        return (typedThread, chatOptions, inputMessagesForChatClient, aiContextProviderMessages);
+        return (typedThread, chatOptions, inputMessagesForChatClient, aiContextProviderMessages, chatMessageStoreMessages);
     }
 
     private void UpdateThreadWithTypeAndConversationId(ChatClientAgentThread thread, string? responseConversationId)
@@ -725,7 +740,13 @@ public sealed partial class ChatClientAgent : AIAgent
         }
     }
 
-    private static Task NotifyMessageStoreOfNewMessagesAsync(ChatClientAgentThread thread, IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken)
+    private static Task NotifyMessageStoreOfFailureAsync(
+        ChatClientAgentThread thread,
+        Exception ex,
+        IEnumerable<ChatMessage> requestMessages,
+        IEnumerable<ChatMessage>? chatMessageStoreMessages,
+        IEnumerable<ChatMessage>? aiContextProviderMessages,
+        CancellationToken cancellationToken)
     {
         var messageStore = thread.MessageStore;
 
@@ -733,7 +754,38 @@ public sealed partial class ChatClientAgent : AIAgent
         // If we don't have one, it means that the chat history is service managed and the underlying service is responsible for storing messages.
         if (messageStore is not null)
         {
-            return messageStore.AddMessagesAsync(newMessages, cancellationToken);
+            var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages!)
+            {
+                AIContextProviderMessages = aiContextProviderMessages,
+                InvokeException = ex
+            };
+
+            return messageStore.InvokedAsync(invokedContext, cancellationToken).AsTask();
+        }
+
+        return Task.CompletedTask;
+    }
+
+    private static Task NotifyMessageStoreOfNewMessagesAsync(
+        ChatClientAgentThread thread,
+        IEnumerable<ChatMessage> requestMessages,
+        IEnumerable<ChatMessage>? chatMessageStoreMessages,
+        IEnumerable<ChatMessage>? aiContextProviderMessages,
+        IEnumerable<ChatMessage> responseMessages,
+        CancellationToken cancellationToken)
+    {
+        var messageStore = thread.MessageStore;
+
+        // Only notify the message store 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.
+        if (messageStore is not null)
+        {
+            var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages!)
+            {
+                AIContextProviderMessages = aiContextProviderMessages,
+                ResponseMessages = responseMessages
+            };
+            return messageStore.InvokedAsync(invokedContext, cancellationToken).AsTask();
         }
 
         return Task.CompletedTask;