.NET: [BREAKING] Fix issue where AIContextProvider messages were not added to MessageStores (#1788)

* Fix issue where AIContextProvider messages were not added to MessageStores

* Fix typos

* Update XML docs to reduce ambiguity.

* Update AIContext XML docs

* Fix merge issue

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

@@ -11,7 +11,7 @@ namespace Microsoft.Agents.AI;
 /// <remarks>
 /// <para>
 /// <see cref="AIContext"/> serves as a container for contextual information that <see cref="AIContextProvider"/> instances
-/// can supply to enhance AI model interactions. This context is combined across multiple providers and merged with
+/// can supply to enhance AI model interactions. This context is merged with
 /// the agent's base configuration before being passed to the underlying AI model.
 /// </para>
 /// <para>
@@ -24,7 +24,7 @@ namespace Microsoft.Agents.AI;
 /// </list>
 /// </para>
 /// <para>
-/// Context information is transient by default and applies only to the current invocation, though messages
+/// Context information is transient by default and applies only to the current invocation, however messages
 /// added through the <see cref="Messages"/> property will be permanently incorporated into the conversation history.
 /// </para>
 /// </remarks>

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

@@ -18,25 +18,25 @@ namespace Microsoft.Agents.AI;
 /// An AI context provider is a component that participates in the agent invocation lifecycle by:
 /// <list type="bullet">
 /// <item><description>Listening to changes in conversations</description></item>
-/// <item><description>Providing additional context to AI models or agents before invocation</description></item>
+/// <item><description>Providing additional context to agents during invocation</description></item>
 /// <item><description>Supplying additional function tools for enhanced capabilities</description></item>
 /// <item><description>Processing invocation results for state management or learning</description></item>
 /// </list>
 /// </para>
 /// <para>
-/// Context providers operate through a two-phase lifecycle: they are called before invocation via
-/// <see cref="InvokingAsync"/> to provide context, and optionally called after invocation via
+/// Context providers operate through a two-phase lifecycle: they are called at the start of invocation via
+/// <see cref="InvokingAsync"/> to provide context, and optionally called at the end of invocation via
 /// <see cref="InvokedAsync"/> to process results.
 /// </para>
 /// </remarks>
 public abstract class AIContextProvider
 {
     /// <summary>
-    /// Called immediately before an AI model or agent is invoked to provide additional context.
+    /// Called at the start of agent invocation to provide additional context.
     /// </summary>
-    /// <param name="context">Contains the request context including the messages that will be sent to the AI model or agent.</param>
+    /// <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 the <see cref="AIContext"/> with additional context to be provided to the AI model or agent.</returns>
+    /// <returns>A task that represents the asynchronous operation. The task result contains the <see cref="AIContext"/> with additional context to be used by the agent during this invocation.</returns>
     /// <remarks>
     /// <para>
     /// Implementers can load any additional context required at this time, such as:
@@ -47,14 +47,11 @@ public abstract class AIContextProvider
     /// <item><description>Injecting contextual messages from conversation history</description></item>
     /// </list>
     /// </para>
-    /// <para>
-    /// The returned context will be combined with context from other providers before being passed to the AI model or agent.
-    /// </para>
     /// </remarks>
     public abstract ValueTask<AIContext> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default);
 
     /// <summary>
-    /// Called immediately after an AI model or agent has been invoked to process the results.
+    /// Called at the end of the agent invocation to process the invocation results.
     /// </summary>
     /// <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>
@@ -123,8 +120,8 @@ public abstract class AIContextProvider
     /// Contains the context information provided to <see cref="InvokingAsync(InvokingContext, CancellationToken)"/>.
     /// </summary>
     /// <remarks>
-    /// This class provides context about the upcoming AI model or agent invocation, including the messages
-    /// that will be sent. Context providers can use this information to determine what additional context
+    /// This class provides context about the invocation before the underlying AI model is invoked, including the messages
+    /// that will be used. Context providers can use this information to determine what additional context
     /// should be provided for the invocation.
     /// </remarks>
     public class InvokingContext
@@ -132,7 +129,7 @@ public abstract class AIContextProvider
         /// <summary>
         /// Initializes a new instance of the <see cref="InvokingContext"/> class with the specified request messages.
         /// </summary>
-        /// <param name="requestMessages">The messages to be sent to the AI model or agent for this invocation.</param>
+        /// <param name="requestMessages">The 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)
         {
@@ -140,11 +137,10 @@ public abstract class AIContextProvider
         }
 
         /// <summary>
-        /// Gets the messages that will be sent to the AI model or agent for this invocation.
+        /// 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 the conversation history
-        /// and new messages that will be processed by the AI model or agent.
+        /// A collection of <see cref="ChatMessage"/> instances representing new messages that were provided by the caller.
         /// </value>
         public IEnumerable<ChatMessage> RequestMessages { get; }
     }
@@ -153,8 +149,8 @@ public abstract class AIContextProvider
     /// Contains the context information provided to <see cref="InvokedAsync(InvokedContext, CancellationToken)"/>.
     /// </summary>
     /// <remarks>
-    /// This class provides context about a completed AI model or agent invocation, including both the
-    /// request messages that were sent and the response messages that were generated. It also indicates
+    /// 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 class InvokedContext
@@ -162,30 +158,41 @@ public abstract class AIContextProvider
         /// <summary>
         /// Initializes a new instance of the <see cref="InvokedContext"/> class with the specified request messages.
         /// </summary>
-        /// <param name="requestMessages">The messages that were sent to the AI model or agent for this invocation.</param>
+        /// <param name="requestMessages">The caller provided messages that were used by the agent for this invocation.</param>
+        /// <param name="aiContextProviderMessages">The messages provided by the <see cref="AIContextProvider"/> for this invocation, if any.</param>
         /// <exception cref="ArgumentNullException"><paramref name="requestMessages"/> is <see langword="null"/>.</exception>
-        public InvokedContext(IEnumerable<ChatMessage> requestMessages)
+        public InvokedContext(IEnumerable<ChatMessage> requestMessages, IEnumerable<ChatMessage>? aiContextProviderMessages)
         {
             this.RequestMessages = requestMessages ?? throw new ArgumentNullException(nameof(requestMessages));
+            this.AIContextProviderMessages = aiContextProviderMessages;
         }
 
         /// <summary>
-        /// Gets the messages that were sent to the AI model or agent for this invocation.
+        /// Gets the caller provided messages that were used by the agent for this invocation.
         /// </summary>
         /// <value>
-        /// A collection of <see cref="ChatMessage"/> instances representing the conversation history
-        /// and new messages that were processed by the AI model or agent.
+        /// A collection of <see cref="ChatMessage"/> instances representing new messages that were provided by the caller.
+        /// This does not include any <see cref="AIContextProvider"/> supplied messages.
         /// </value>
         public IEnumerable<ChatMessage> RequestMessages { get; }
 
         /// <summary>
-        /// Gets the collection of response messages generated by the AI model or agent if the invocation succeeded.
+        /// Gets the messages provided by the <see cref="AIContextProvider"/> for this invocation, if any.
         /// </summary>
         /// <value>
-        /// A collection of <see cref="ChatMessage"/> instances representing the response from the AI model or agent,
+        /// 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; }
+
+        /// <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; init; }
+        public IEnumerable<ChatMessage>? ResponseMessages { get; set; }
 
         /// <summary>
         /// Gets the <see cref="Exception"/> that was thrown during the invocation, if the invocation failed.
@@ -193,6 +200,6 @@ public abstract class AIContextProvider
         /// <value>
         /// The exception that caused the invocation to fail, or <see langword="null"/> if the invocation succeeded.
         /// </value>
-        public Exception? InvokeException { get; init; }
+        public Exception? InvokeException { get; set; }
     }
 }

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

@@ -161,6 +161,11 @@ public sealed class Mem0Provider : AIContextProvider
     /// <inheritdoc />
     public override async ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
+        if (context.InvokeException is not null)
+        {
+            return; // Do not update memory on failed invocations.
+        }
+
         // Persist request and response messages after invocation.
         await this.PersistMessagesAsync(context.RequestMessages, cancellationToken).ConfigureAwait(false);
 

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

@@ -201,11 +201,9 @@ public sealed partial class ChatClientAgent : AIAgent
     {
         var inputMessages = Throw.IfNull(messages) as IReadOnlyCollection<ChatMessage> ?? messages.ToList();
 
-        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> threadMessages) =
+        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> inputMessagesForChatClient, IList<ChatMessage>? aiContextProviderMessages) =
             await this.PrepareThreadAndMessagesAsync(thread, inputMessages, options, cancellationToken).ConfigureAwait(false);
 
-        int messageCount = threadMessages.Count;
-
         var chatClient = this.ChatClient;
 
         chatClient = ApplyRunOptionsTransformations(options, chatClient);
@@ -221,11 +219,11 @@ public sealed partial class ChatClientAgent : AIAgent
         try
         {
             // Using the enumerator to ensure we consider the case where no updates are returned for notification.
-            responseUpdatesEnumerator = chatClient.GetStreamingResponseAsync(threadMessages, chatOptions, cancellationToken).GetAsyncEnumerator(cancellationToken);
+            responseUpdatesEnumerator = chatClient.GetStreamingResponseAsync(inputMessagesForChatClient, chatOptions, cancellationToken).GetAsyncEnumerator(cancellationToken);
         }
         catch (Exception ex)
         {
-            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             throw;
         }
 
@@ -239,7 +237,7 @@ public sealed partial class ChatClientAgent : AIAgent
         }
         catch (Exception ex)
         {
-            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             throw;
         }
 
@@ -260,7 +258,7 @@ public sealed partial class ChatClientAgent : AIAgent
             }
             catch (Exception ex)
             {
-                await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, cancellationToken).ConfigureAwait(false);
+                await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
                 throw;
             }
         }
@@ -272,10 +270,10 @@ 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 NotifyThreadOfNewMessagesAsync(safeThread, inputMessages.Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);
+        await NotifyThreadOfNewMessagesAsync(safeThread, inputMessages.Concat(aiContextProviderMessages ?? []).Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);
 
         // Notify the AIContextProvider of all new messages.
-        await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
+        await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
     }
 
     /// <inheritdoc/>
@@ -348,7 +346,7 @@ public sealed partial class ChatClientAgent : AIAgent
     {
         var inputMessages = Throw.IfNull(messages) as IReadOnlyCollection<ChatMessage> ?? messages.ToList();
 
-        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> threadMessages) =
+        (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List<ChatMessage> inputMessagesForChatClient, IList<ChatMessage>? aiContextProviderMessages) =
             await this.PrepareThreadAndMessagesAsync(thread, inputMessages, options, cancellationToken).ConfigureAwait(false);
 
         var chatClient = this.ChatClient;
@@ -363,11 +361,11 @@ public sealed partial class ChatClientAgent : AIAgent
         TChatClientResponse chatResponse;
         try
         {
-            chatResponse = await chatClientRunFunc.Invoke(chatClient, threadMessages, chatOptions, cancellationToken).ConfigureAwait(false);
+            chatResponse = await chatClientRunFunc.Invoke(chatClient, inputMessagesForChatClient, chatOptions, cancellationToken).ConfigureAwait(false);
         }
         catch (Exception ex)
         {
-            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, cancellationToken).ConfigureAwait(false);
+            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false);
             throw;
         }
 
@@ -384,10 +382,10 @@ 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 NotifyThreadOfNewMessagesAsync(safeThread, inputMessages.Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);
+        await NotifyThreadOfNewMessagesAsync(safeThread, inputMessages.Concat(aiContextProviderMessages ?? []).Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);
 
         // Notify the AIContextProvider of all new messages.
-        await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
+        await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
 
         var agentResponse = agentResponseFactoryFunc(chatResponse);
 
@@ -399,11 +397,16 @@ public sealed partial class ChatClientAgent : AIAgent
     /// <summary>
     /// Notify the <see cref="AIContextProvider"/> when an agent run succeeded, if there is an <see cref="AIContextProvider"/>.
     /// </summary>
-    private static async Task NotifyAIContextProviderOfSuccessAsync(ChatClientAgentThread thread, IEnumerable<ChatMessage> inputMessages, IEnumerable<ChatMessage> responseMessages, CancellationToken cancellationToken)
+    private static async Task NotifyAIContextProviderOfSuccessAsync(
+        ChatClientAgentThread thread,
+        IEnumerable<ChatMessage> inputMessages,
+        IList<ChatMessage>? aiContextProviderMessages,
+        IEnumerable<ChatMessage> responseMessages,
+        CancellationToken cancellationToken)
     {
         if (thread.AIContextProvider is not null)
         {
-            await thread.AIContextProvider.InvokedAsync(new(inputMessages) { ResponseMessages = responseMessages },
+            await thread.AIContextProvider.InvokedAsync(new(inputMessages, aiContextProviderMessages) { ResponseMessages = responseMessages },
                 cancellationToken).ConfigureAwait(false);
         }
     }
@@ -411,11 +414,16 @@ public sealed partial class ChatClientAgent : AIAgent
     /// <summary>
     /// Notify the <see cref="AIContextProvider"/> of any failure during an agent run, if there is an <see cref="AIContextProvider"/>.
     /// </summary>
-    private static async Task NotifyAIContextProviderOfFailureAsync(ChatClientAgentThread thread, Exception ex, IEnumerable<ChatMessage> inputMessages, CancellationToken cancellationToken)
+    private static async Task NotifyAIContextProviderOfFailureAsync(
+        ChatClientAgentThread thread,
+        Exception ex,
+        IEnumerable<ChatMessage> inputMessages,
+        IList<ChatMessage>? aiContextProviderMessages,
+        CancellationToken cancellationToken)
     {
         if (thread.AIContextProvider is not null)
         {
-            await thread.AIContextProvider.InvokedAsync(new(inputMessages) { InvokeException = ex },
+            await thread.AIContextProvider.InvokedAsync(new(inputMessages, aiContextProviderMessages) { InvokeException = ex },
                 cancellationToken).ConfigureAwait(false);
         }
     }
@@ -556,7 +564,7 @@ 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> ThreadMessages)> PrepareThreadAndMessagesAsync(
+    private async Task<(ChatClientAgentThread AgentThread, ChatOptions? ChatOptions, List<ChatMessage> InputMessagesForChatClient, IList<ChatMessage>? AIContextProviderMessages)> PrepareThreadAndMessagesAsync(
         AgentThread? thread,
         IEnumerable<ChatMessage> inputMessages,
         AgentRunOptions? runOptions,
@@ -582,7 +590,8 @@ public sealed partial class ChatClientAgent : AIAgent
         {
             throw new InvalidOperationException("Input messages are not allowed when continuing a background response using a continuation token.");
         }
-        List<ChatMessage> threadMessages = [];
+        List<ChatMessage> inputMessagesForChatClient = [];
+        IList<ChatMessage>? aiContextProviderMessages = null;
 
         // Populate the thread messages only if we are not continuing an existing response as it's not allowed
         if (chatOptions?.ContinuationToken is null)
@@ -590,7 +599,7 @@ 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)
             {
-                threadMessages.AddRange(await typedThread.MessageStore.GetMessagesAsync(cancellationToken).ConfigureAwait(false));
+                inputMessagesForChatClient.AddRange(await typedThread.MessageStore.GetMessagesAsync(cancellationToken).ConfigureAwait(false));
             }
 
             // If we have an AIContextProvider, we should get context from it, and update our
@@ -601,7 +610,8 @@ public sealed partial class ChatClientAgent : AIAgent
                 var aiContext = await typedThread.AIContextProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
                 if (aiContext.Messages is { Count: > 0 })
                 {
-                    threadMessages.AddRange(aiContext.Messages);
+                    inputMessagesForChatClient.AddRange(aiContext.Messages);
+                    aiContextProviderMessages = aiContext.Messages;
                 }
 
                 if (aiContext.Tools is { Count: > 0 })
@@ -622,7 +632,7 @@ public sealed partial class ChatClientAgent : AIAgent
             }
 
             // Add the input messages to the end of thread messages.
-            threadMessages.AddRange(inputMessages);
+            inputMessagesForChatClient.AddRange(inputMessages);
         }
 
         // If a user provided two different thread ids, via the thread object and options, we should throw
@@ -649,7 +659,7 @@ public sealed partial class ChatClientAgent : AIAgent
             chatOptions.ConversationId = typedThread.ConversationId;
         }
 
-        return (typedThread, chatOptions, threadMessages);
+        return (typedThread, chatOptions, inputMessagesForChatClient, aiContextProviderMessages);
     }
 
     private void UpdateThreadWithTypeAndConversationId(ChatClientAgentThread thread, string? responseConversationId)

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

@@ -176,6 +176,11 @@ public sealed class TextSearchProvider : AIContextProvider
             return default; // Memory disabled.
         }
 
+        if (context.InvokeException is not null)
+        {
+            return default; // Do not update memory on failed invocations.
+        }
+
         var messagesText = context.RequestMessages
             .Concat(context.ResponseMessages ?? [])
             .Where(m => (m.Role == ChatRole.User || m.Role == ChatRole.Assistant) && !string.IsNullOrWhiteSpace(m.Text))

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

@@ -16,7 +16,7 @@ public class AIContextProviderTests
     {
         var provider = new TestAIContextProvider();
         var messages = new ReadOnlyCollection<ChatMessage>([]);
-        var task = provider.InvokedAsync(new(messages));
+        var task = provider.InvokedAsync(new(messages, aiContextProviderMessages: null));
         Assert.Equal(default, task);
     }
 
@@ -37,7 +37,7 @@ public class AIContextProviderTests
     [Fact]
     public void InvokedContext_Constructor_ThrowsForNullMessages()
     {
-        Assert.Throws<ArgumentNullException>(() => new AIContextProvider.InvokedContext(null!));
+        Assert.Throws<ArgumentNullException>(() => new AIContextProvider.InvokedContext(null!, aiContextProviderMessages: null));
     }
 
     #region GetService Method Tests

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

@@ -53,7 +53,7 @@ public sealed class Mem0ProviderTests : IDisposable
         Assert.DoesNotContain("Caoimhe", ctxBefore.Messages?[0].Text ?? string.Empty);
 
         // Act
-        await sut.InvokedAsync(new AIContextProvider.InvokedContext(new[] { input }));
+        await sut.InvokedAsync(new AIContextProvider.InvokedContext(new[] { input }, aiContextProviderMessages: null));
         var ctxAfterAdding = await GetContextWithRetryAsync(sut, question);
         await sut.ClearStoredMemoriesAsync();
         var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(new[] { question }));
@@ -77,7 +77,7 @@ public sealed class Mem0ProviderTests : IDisposable
         Assert.DoesNotContain("Caoimhe", ctxBefore.Messages?[0].Text ?? string.Empty);
 
         // Act
-        await sut.InvokedAsync(new AIContextProvider.InvokedContext(new[] { assistantIntro }));
+        await sut.InvokedAsync(new AIContextProvider.InvokedContext(new[] { assistantIntro }, aiContextProviderMessages: null));
         var ctxAfterAdding = await GetContextWithRetryAsync(sut, question);
         await sut.ClearStoredMemoriesAsync();
         var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(new[] { question }));
@@ -105,7 +105,7 @@ public sealed class Mem0ProviderTests : IDisposable
         Assert.DoesNotContain("Caoimhe", ctxBefore2.Messages?[0].Text ?? string.Empty);
 
         // Act
-        await sut1.InvokedAsync(new AIContextProvider.InvokedContext(new[] { assistantIntro }));
+        await sut1.InvokedAsync(new AIContextProvider.InvokedContext(new[] { assistantIntro }, aiContextProviderMessages: null));
         var ctxAfterAdding1 = await GetContextWithRetryAsync(sut1, question);
         var ctxAfterAdding2 = await GetContextWithRetryAsync(sut2, question);
 

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

@@ -123,8 +123,7 @@ public sealed class Mem0ProviderTests : IDisposable
         };
 
         // Act
-        await sut.InvokedAsync(new AIContextProvider.InvokedContext(requestMessages)); // persists request messages
-        await sut.InvokedAsync(new AIContextProvider.InvokedContext(responseMessages)); // persists assistant
+        await sut.InvokedAsync(new AIContextProvider.InvokedContext(requestMessages, aiContextProviderMessages: null) { ResponseMessages = responseMessages });
 
         // Assert
         var memoryPosts = this._handler.Requests.Where(r => r.RequestMessage.RequestUri!.AbsolutePath == "/v1/memories/" && r.RequestMessage.Method == HttpMethod.Post).ToList();
@@ -136,6 +135,27 @@ public sealed class Mem0ProviderTests : IDisposable
         Assert.DoesNotContain(memoryPosts, r => ContainsOrdinal(r.RequestBody, "Tool text"));
     }
 
+    [Fact]
+    public async Task InvokedAsync_PersistsNothingForFailedRequestAsync()
+    {
+        // Arrange
+        var options = new Mem0ProviderOptions { ApplicationId = "a", AgentId = "b", ThreadId = "c", UserId = "d" };
+        var sut = new Mem0Provider(this._httpClient, options);
+
+        var requestMessages = new List<ChatMessage>
+        {
+            new(ChatRole.User, "User text"),
+            new(ChatRole.System, "System text"),
+            new(ChatRole.Tool, "Tool text should be ignored")
+        };
+
+        // Act
+        await sut.InvokedAsync(new AIContextProvider.InvokedContext(requestMessages, aiContextProviderMessages: null) { ResponseMessages = null, InvokeException = new InvalidOperationException("Request Failed") });
+
+        // Assert
+        Assert.Empty(this._handler.Requests);
+    }
+
     [Fact]
     public async Task ClearStoredMemoriesAsync_SendsDeleteWithQueryAsync()
     {

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

@@ -220,10 +220,10 @@ public partial class ChatClientAgentTests
     }
 
     /// <summary>
-    /// Verify that RunAsync works with existing thread and retrieves messages from IMessagesRetrievableThread.
+    /// Verify that RunAsync works with existing thread and can retreive messages if the thread has a MessageStore.
     /// </summary>
     [Fact]
-    public async Task RunAsyncRetrievesMessagesFromThreadWhenThreadImplementsIMessagesRetrievableThreadAsync()
+    public async Task RunAsyncRetrievesMessagesFromThreadWhenThreadStoresMessagesThreadAsync()
     {
         // Arrange
         Mock<IChatClient> mockService = new();
@@ -449,7 +449,10 @@ public partial class ChatClientAgentTests
         await agent.RunAsync([new(ChatRole.User, "test")], thread);
 
         // Assert
-        Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
+        var messageStore = Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
+        Assert.Equal(2, messageStore.Count);
+        Assert.Equal("test", messageStore[0].Text);
+        Assert.Equal("response", messageStore[1].Text);
         mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Once);
     }
 
@@ -492,6 +495,7 @@ public partial class ChatClientAgentTests
         // Arrange
         ChatMessage[] requestMessages = [new(ChatRole.User, "user message")];
         ChatMessage[] responseMessages = [new(ChatRole.Assistant, "response")];
+        ChatMessage[] aiContextProviderMessages = [new(ChatRole.System, "context provider message")];
         Mock<IChatClient> mockService = new();
         List<ChatMessage> capturedMessages = [];
         string capturedInstructions = string.Empty;
@@ -517,7 +521,7 @@ public partial class ChatClientAgentTests
             .Setup(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(new AIContext
             {
-                Messages = [new(ChatRole.System, "context provider message")],
+                Messages = aiContextProviderMessages,
                 Instructions = "context provider instructions",
                 Tools = [AIFunctionFactory.Create(() => { }, "context provider function")]
             });
@@ -528,7 +532,8 @@ public partial class ChatClientAgentTests
         ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = _ => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
 
         // Act
-        await agent.RunAsync(requestMessages);
+        var thread = agent.GetNewThread() as ChatClientAgentThread;
+        await agent.RunAsync(requestMessages, thread);
 
         // Assert
         // Should contain: base instructions, context message, user message, base function, context function
@@ -541,8 +546,20 @@ public partial class ChatClientAgentTests
         Assert.Equal(2, capturedTools.Count);
         Assert.Contains(capturedTools, t => t.Name == "base function");
         Assert.Contains(capturedTools, t => t.Name == "context provider function");
+
+        // Verify that the thread was updated with the input, ai context and response messages
+        var messageStore = Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
+        Assert.Equal(3, messageStore.Count);
+        Assert.Equal("user message", messageStore[0].Text);
+        Assert.Equal("context provider message", messageStore[1].Text);
+        Assert.Equal("response", messageStore[2].Text);
+
         mockProvider.Verify(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()), Times.Once);
-        mockProvider.Verify(p => p.InvokedAsync(It.Is<AIContextProvider.InvokedContext>(x => x.RequestMessages == requestMessages && x.ResponseMessages == responseMessages && x.InvokeException == null), It.IsAny<CancellationToken>()), Times.Once);
+        mockProvider.Verify(p => p.InvokedAsync(It.Is<AIContextProvider.InvokedContext>(x =>
+            x.RequestMessages == requestMessages &&
+            x.AIContextProviderMessages == aiContextProviderMessages &&
+            x.ResponseMessages == responseMessages &&
+            x.InvokeException == null), It.IsAny<CancellationToken>()), Times.Once);
     }
 
     /// <summary>
@@ -554,6 +571,7 @@ public partial class ChatClientAgentTests
         // Arrange
         ChatMessage[] requestMessages = [new(ChatRole.User, "user message")];
         ChatMessage[] responseMessages = [new(ChatRole.Assistant, "response")];
+        ChatMessage[] aiContextProviderMessages = [new(ChatRole.System, "context provider message")];
         Mock<IChatClient> mockService = new();
         mockService
             .Setup(s => s.GetResponseAsync(
@@ -565,7 +583,10 @@ public partial class ChatClientAgentTests
         var mockProvider = new Mock<AIContextProvider>();
         mockProvider
             .Setup(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(new AIContext());
+            .ReturnsAsync(new AIContext
+            {
+                Messages = aiContextProviderMessages,
+            });
         mockProvider
             .Setup(p => p.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
             .Returns(new ValueTask());
@@ -577,7 +598,11 @@ public partial class ChatClientAgentTests
 
         // Assert
         mockProvider.Verify(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()), Times.Once);
-        mockProvider.Verify(p => p.InvokedAsync(It.Is<AIContextProvider.InvokedContext>(x => x.RequestMessages == requestMessages && x.ResponseMessages == null && x.InvokeException is InvalidOperationException), It.IsAny<CancellationToken>()), Times.Once);
+        mockProvider.Verify(p => p.InvokedAsync(It.Is<AIContextProvider.InvokedContext>(x =>
+            x.RequestMessages == requestMessages &&
+            x.AIContextProviderMessages == aiContextProviderMessages &&
+            x.ResponseMessages == null &&
+            x.InvokeException is InvalidOperationException), It.IsAny<CancellationToken>()), Times.Once);
     }
 
     /// <summary>
@@ -1879,7 +1904,10 @@ public partial class ChatClientAgentTests
         await agent.RunStreamingAsync([new(ChatRole.User, "test")], thread).ToListAsync();
 
         // Assert
-        Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
+        var messageStore = Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
+        Assert.Equal(2, messageStore.Count);
+        Assert.Equal("test", messageStore[0].Text);
+        Assert.Equal("what?", messageStore[1].Text);
         mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Once);
     }
 
@@ -1918,6 +1946,130 @@ public partial class ChatClientAgentTests
         mockFactory.Verify(f => f(It.IsAny<ChatClientAgentOptions.ChatMessageStoreFactoryContext>()), Times.Never);
     }
 
+    /// <summary>
+    /// Verify that RunStreamingAsync invokes any provided AIContextProvider and uses the result.
+    /// </summary>
+    [Fact]
+    public async Task RunStreamingAsyncInvokesAIContextProviderAndUsesResultAsync()
+    {
+        // Arrange
+        ChatMessage[] requestMessages = [new(ChatRole.User, "user message")];
+        ChatResponseUpdate[] responseUpdates = [new(ChatRole.Assistant, "response")];
+        ChatMessage[] aiContextProviderMessages = [new(ChatRole.System, "context provider message")];
+        Mock<IChatClient> mockService = new();
+        List<ChatMessage> capturedMessages = [];
+        string capturedInstructions = string.Empty;
+        List<AITool> capturedTools = [];
+        mockService
+            .Setup(s => s.GetStreamingResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>()))
+            .Callback<IEnumerable<ChatMessage>, ChatOptions, CancellationToken>((msgs, opts, ct) =>
+            {
+                capturedMessages.AddRange(msgs);
+                capturedInstructions = opts.Instructions ?? string.Empty;
+                if (opts.Tools is not null)
+                {
+                    capturedTools.AddRange(opts.Tools);
+                }
+            })
+            .Returns(ToAsyncEnumerableAsync(responseUpdates));
+
+        var mockProvider = new Mock<AIContextProvider>();
+        mockProvider
+            .Setup(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new AIContext
+            {
+                Messages = aiContextProviderMessages,
+                Instructions = "context provider instructions",
+                Tools = [AIFunctionFactory.Create(() => { }, "context provider function")]
+            });
+        mockProvider
+            .Setup(p => p.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
+            .Returns(new ValueTask());
+
+        ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = _ => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
+
+        // Act
+        var thread = agent.GetNewThread() as ChatClientAgentThread;
+        var updates = agent.RunStreamingAsync(requestMessages, thread);
+        _ = await updates.ToAgentRunResponseAsync();
+
+        // Assert
+        // Should contain: base instructions, context message, user message, base function, context function
+        Assert.Equal(2, capturedMessages.Count);
+        Assert.Equal("base instructions\ncontext provider instructions", capturedInstructions);
+        Assert.Equal("context provider message", capturedMessages[0].Text);
+        Assert.Equal(ChatRole.System, capturedMessages[0].Role);
+        Assert.Equal("user message", capturedMessages[1].Text);
+        Assert.Equal(ChatRole.User, capturedMessages[1].Role);
+        Assert.Equal(2, capturedTools.Count);
+        Assert.Contains(capturedTools, t => t.Name == "base function");
+        Assert.Contains(capturedTools, t => t.Name == "context provider function");
+
+        // Verify that the thread was updated with the input, ai context and response messages
+        var messageStore = Assert.IsType<InMemoryChatMessageStore>(thread!.MessageStore);
+        Assert.Equal(3, messageStore.Count);
+        Assert.Equal("user message", messageStore[0].Text);
+        Assert.Equal("context provider message", messageStore[1].Text);
+        Assert.Equal("response", messageStore[2].Text);
+
+        mockProvider.Verify(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()), Times.Once);
+        mockProvider.Verify(p => p.InvokedAsync(It.Is<AIContextProvider.InvokedContext>(x =>
+            x.RequestMessages == requestMessages &&
+            x.AIContextProviderMessages == aiContextProviderMessages &&
+            x.ResponseMessages!.Count() == 1 &&
+            x.ResponseMessages!.ElementAt(0).Text == "response" &&
+            x.InvokeException == null), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that RunStreamingAsync invokes any provided AIContextProvider when the downstream GetStreamingResponse call fails.
+    /// </summary>
+    [Fact]
+    public async Task RunStreamingAsyncInvokesAIContextProviderWhenGetResponseFailsAsync()
+    {
+        // Arrange
+        ChatMessage[] requestMessages = [new(ChatRole.User, "user message")];
+        ChatMessage[] aiContextProviderMessages = [new(ChatRole.System, "context provider message")];
+        Mock<IChatClient> mockService = new();
+        mockService
+            .Setup(s => s.GetStreamingResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions>(),
+                It.IsAny<CancellationToken>()))
+            .Throws(new InvalidOperationException("downstream failure"));
+
+        var mockProvider = new Mock<AIContextProvider>();
+        mockProvider
+            .Setup(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new AIContext
+            {
+                Messages = aiContextProviderMessages,
+            });
+        mockProvider
+            .Setup(p => p.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
+            .Returns(new ValueTask());
+
+        ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = _ => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
+
+        // Act
+        await Assert.ThrowsAsync<InvalidOperationException>(async () =>
+        {
+            var updates = agent.RunStreamingAsync(requestMessages);
+            await updates.ToAgentRunResponseAsync();
+        });
+
+        // Assert
+        mockProvider.Verify(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()), Times.Once);
+        mockProvider.Verify(p => p.InvokedAsync(It.Is<AIContextProvider.InvokedContext>(x =>
+            x.RequestMessages == requestMessages &&
+            x.AIContextProviderMessages == aiContextProviderMessages &&
+            x.ResponseMessages == null &&
+            x.InvokeException is InvalidOperationException), It.IsAny<CancellationToken>()), Times.Once);
+    }
+
     #endregion
 
     #region GetNewThread Tests

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

@@ -289,6 +289,45 @@ public sealed class TextSearchProviderTests
 
     #region Recent Message Memory Tests
 
+    [Fact]
+    public async Task InvokingAsync_WithPreviousFailedRequest_ShouldNotIncludeFailedRequestInputInSearchInputAsync()
+    {
+        // Arrange
+        var options = new TextSearchProviderOptions
+        {
+            SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke,
+            RecentMessageMemoryLimit = 3
+        };
+        string? capturedInput = null;
+        Task<IEnumerable<TextSearchProvider.TextSearchResult>> SearchDelegateAsync(string input, CancellationToken ct)
+        {
+            capturedInput = input;
+            return Task.FromResult<IEnumerable<TextSearchProvider.TextSearchResult>>([]); // No results needed.
+        }
+        var provider = new TextSearchProvider(SearchDelegateAsync, options);
+
+        // Populate memory with more messages than the limit (A,B,C,D) -> should retain B,C,D
+        var initialMessages = new[]
+        {
+            new ChatMessage(ChatRole.User, "A"),
+            new ChatMessage(ChatRole.Assistant, "B"),
+            new ChatMessage(ChatRole.User, "C"),
+            new ChatMessage(ChatRole.Assistant, "D"),
+        };
+        await provider.InvokedAsync(new(initialMessages, aiContextProviderMessages: null) { InvokeException = new InvalidOperationException("Request Failed") });
+
+        var invokingContext = new AIContextProvider.InvokingContext(new[]
+        {
+            new ChatMessage(ChatRole.User, "E")
+        });
+
+        // Act
+        await provider.InvokingAsync(invokingContext, CancellationToken.None);
+
+        // Assert
+        Assert.Equal("E", capturedInput); // Only the messages from the current request, since previous failed request should not be stored.
+    }
+
     [Fact]
     public async Task InvokingAsync_WithRecentMessageMemory_ShouldIncludeStoredMessagesInSearchInputAsync()
     {
@@ -314,7 +353,7 @@ public sealed class TextSearchProviderTests
             new ChatMessage(ChatRole.User, "C"),
             new ChatMessage(ChatRole.Assistant, "D"),
         };
-        await provider.InvokedAsync(new(initialMessages));
+        await provider.InvokedAsync(new(initialMessages, aiContextProviderMessages: null));
 
         var invokingContext = new AIContextProvider.InvokingContext(new[]
         {
@@ -350,7 +389,7 @@ public sealed class TextSearchProviderTests
         {
             new ChatMessage(ChatRole.User, "A"),
             new ChatMessage(ChatRole.Assistant, "B"),
-        }));
+        }, aiContextProviderMessages: null));
 
         // Second memory update (C,D,E)
         await provider.InvokedAsync(new(new[]
@@ -358,7 +397,7 @@ public sealed class TextSearchProviderTests
             new ChatMessage(ChatRole.User, "C"),
             new ChatMessage(ChatRole.Assistant, "D"),
             new ChatMessage(ChatRole.User, "E"),
-        }));
+        }, aiContextProviderMessages: null));
 
         var invokingContext = new AIContextProvider.InvokingContext(new[] { new ChatMessage(ChatRole.User, "F") });
 
@@ -410,7 +449,7 @@ public sealed class TextSearchProviderTests
         };
 
         // Act
-        await provider.InvokedAsync(new(messages)); // Populate recent memory.
+        await provider.InvokedAsync(new(messages, aiContextProviderMessages: null)); // Populate recent memory.
         var state = provider.Serialize();
 
         // Assert
@@ -438,7 +477,7 @@ public sealed class TextSearchProviderTests
             new ChatMessage(ChatRole.User, "C"),
             new ChatMessage(ChatRole.Assistant, "D"),
         };
-        await provider.InvokedAsync(new(messages));
+        await provider.InvokedAsync(new(messages, aiContextProviderMessages: null));
 
         // Act
         var state = provider.Serialize();
@@ -478,7 +517,7 @@ public sealed class TextSearchProviderTests
             new ChatMessage(ChatRole.Assistant, "L4"),
             new ChatMessage(ChatRole.User, "L5"),
         };
-        await initialProvider.InvokedAsync(new(messages));
+        await initialProvider.InvokedAsync(new(messages, aiContextProviderMessages: null));
         var state = initialProvider.Serialize();
 
         string? capturedInput = null;