.NET: [BREAKING] Refactor providers to move common functionality to base (#3900)

* Move common functionality to provider base classes

Co-Authored-By: Copilot <175728472+Copilot@users.noreply.github.com>

* Address PR comments.

---------

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

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

@@ -86,27 +86,25 @@ namespace SampleApp
     /// <summary>
     /// Sample memory component that can remember a user's name and age.
     /// </summary>
-    internal sealed class UserInfoMemory : AIContextProvider
+    internal sealed class UserInfoMemory : AIContextProvider<UserInfo>
     {
         private readonly IChatClient _chatClient;
-        private readonly Func<AgentSession?, UserInfo> _stateInitializer;
 
         public UserInfoMemory(IChatClient chatClient, Func<AgentSession?, UserInfo>? stateInitializer = null)
+            : base(stateInitializer ?? (_ => new UserInfo()), null, null, null, null)
         {
             this._chatClient = chatClient;
-            this._stateInitializer = stateInitializer ?? (_ => new UserInfo());
         }
 
         public UserInfo GetUserInfo(AgentSession session)
-            => session.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory)) ?? new UserInfo();
+            => this.GetOrInitializeState(session);
 
         public void SetUserInfo(AgentSession session, UserInfo userInfo)
-            => session.StateBag.SetValue(nameof(UserInfoMemory), userInfo);
+            => this.SaveState(session, userInfo);
 
-        protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+        protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default)
         {
-            var userInfo = context.Session?.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory))
-                ?? this._stateInitializer.Invoke(context.Session);
+            var userInfo = this.GetOrInitializeState(context.Session);
 
             // Try and extract the user name and age from the message if we don't have it already and it's a user message.
             if ((userInfo.UserName is null || userInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User))
@@ -123,20 +121,14 @@ namespace SampleApp
                 userInfo.UserAge ??= result.Result.UserAge;
             }
 
-            context.Session?.StateBag.SetValue(nameof(UserInfoMemory), userInfo);
+            this.SaveState(context.Session, userInfo);
         }
 
-        protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+        protected override ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
         {
-            var inputContext = context.AIContext;
-            var userInfo = context.Session?.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory))
-                ?? this._stateInitializer.Invoke(context.Session);
+            var userInfo = this.GetOrInitializeState(context.Session);
 
             StringBuilder instructions = new();
-            if (!string.IsNullOrEmpty(inputContext.Instructions))
-            {
-                instructions.AppendLine(inputContext.Instructions);
-            }
 
             // If we don't already know the user's name and age, add instructions to ask for them, otherwise just provide what we have to the context.
             instructions
@@ -151,9 +143,7 @@ namespace SampleApp
 
             return new ValueTask<AIContext>(new AIContext
             {
-                Instructions = instructions.ToString(),
-                Messages = inputContext.Messages,
-                Tools = inputContext.Tools
+                Instructions = instructions.ToString()
             });
         }
     }

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

@@ -76,45 +76,23 @@ namespace SampleApp
     /// State (the session DB key) is stored in the <see cref="AgentSession.StateBag"/> so it roundtrips
     /// automatically with session serialization.
     /// </summary>
-    internal sealed class VectorChatHistoryProvider : ChatHistoryProvider
+    internal sealed class VectorChatHistoryProvider : ChatHistoryProvider<VectorChatHistoryProvider.State>
     {
         private readonly VectorStore _vectorStore;
-        private readonly Func<AgentSession?, State> _stateInitializer;
-        private readonly string _stateKey;
-
-        /// <inheritdoc />
-        public override string StateKey => this._stateKey;
 
         public VectorChatHistoryProvider(
             VectorStore vectorStore,
             Func<AgentSession?, State>? stateInitializer = null,
             string? stateKey = null)
+            : base(stateInitializer: stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N"))), stateKey: stateKey, jsonSerializerOptions: null, provideOutputMessageFilter: null, storeInputMessageFilter: null)
         {
             this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
-            this._stateInitializer = stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N")));
-            this._stateKey = stateKey ?? base.StateKey;
         }
 
         public string GetSessionDbKey(AgentSession session)
             => this.GetOrInitializeState(session).SessionDbKey;
 
-        private State GetOrInitializeState(AgentSession? session)
-        {
-            if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state) is true && state is not null)
-            {
-                return state;
-            }
-
-            state = this._stateInitializer(session);
-            if (session is not null)
-            {
-                session.StateBag.SetValue(this._stateKey, state);
-            }
-
-            return state;
-        }
-
-        protected override async ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+        protected override async ValueTask<IEnumerable<ChatMessage>> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default)
         {
             var state = this.GetOrInitializeState(context.Session);
             var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
@@ -129,29 +107,17 @@ namespace SampleApp
 
             var messages = records.ConvertAll(x => JsonSerializer.Deserialize<ChatMessage>(x.SerializedMessage!)!);
             messages.Reverse();
-            return messages
-                .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!))
-                .Concat(context.RequestMessages);
+            return messages;
         }
 
-        protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+        protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default)
         {
-            // Don't store messages if the request failed.
-            if (context.InvokeException is not null)
-            {
-                return;
-            }
-
             var state = this.GetOrInitializeState(context.Session);
 
             var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
             await collection.EnsureCollectionExistsAsync(cancellationToken);
 
-            // Add both request and response messages to the store, excluding messages that came from chat history.
-            // Optionally messages produced by the AIContextProvider can also be persisted (not shown).
-            var allNewMessages = context.RequestMessages
-                .Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory)
-                .Concat(context.ResponseMessages ?? []);
+            var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []);
 
             await collection.UpsertAsync(allNewMessages.Select(x => new ChatHistoryItem()
             {

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

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -10,7 +11,7 @@ using Microsoft.Shared.Diagnostics;
 namespace Microsoft.Agents.AI;
 
 /// <summary>
-/// Provides an abstract base class for components that enhance AI context management during agent invocations.
+/// Provides an abstract base class for components that enhance AI context during agent invocations.
 /// </summary>
 /// <remarks>
 /// <para>
@@ -30,6 +31,25 @@ namespace Microsoft.Agents.AI;
 /// </remarks>
 public abstract class AIContextProvider
 {
+    private static IEnumerable<ChatMessage> DefaultExternalOnlyFilter(IEnumerable<ChatMessage> messages)
+        => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External);
+
+    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _provideInputMessageFilter;
+    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _storeInputMessageFilter;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AIContextProvider"/> class.
+    /// </summary>
+    /// <param name="provideInputMessageFilter">An optional filter function to apply to input messages before providing context via <see cref="ProvideAIContextAsync"/>. If not set, defaults to including only <see cref="AgentRequestMessageSourceType.External"/> messages.</param>
+    /// <param name="storeInputMessageFilter">An optional filter function to apply to request messages before storing context via <see cref="StoreAIContextAsync"/>. If not set, defaults to including only <see cref="AgentRequestMessageSourceType.External"/> messages.</param>
+    protected AIContextProvider(
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideInputMessageFilter = null,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputMessageFilter = null)
+    {
+        this._provideInputMessageFilter = provideInputMessageFilter ?? DefaultExternalOnlyFilter;
+        this._storeInputMessageFilter = storeInputMessageFilter ?? DefaultExternalOnlyFilter;
+    }
+
     /// <summary>
     /// Gets the key used to store the provider state in the <see cref="AgentSession.StateBag"/>.
     /// </summary>
@@ -58,7 +78,7 @@ public abstract class AIContextProvider
     /// </para>
     /// </remarks>
     public ValueTask<AIContext> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
-        => this.InvokingCoreAsync(context, cancellationToken);
+        => this.InvokingCoreAsync(Throw.IfNull(context), cancellationToken);
 
     /// <summary>
     /// Called at the start of agent invocation to provide additional context.
@@ -76,8 +96,96 @@ public abstract class AIContextProvider
     /// <item><description>Injecting contextual messages from conversation history</description></item>
     /// </list>
     /// </para>
+    /// <para>
+    /// The default implementation of this method filters the input messages using the configured provide-input message filter
+    /// (which defaults to including only <see cref="AgentRequestMessageSourceType.External"/> messages),
+    /// then calls <see cref="ProvideAIContextAsync"/> to get additional context,
+    /// stamps any messages from the returned context with <see cref="AgentRequestMessageSourceType.AIContextProvider"/> source attribution,
+    /// and merges the returned context with the original (unfiltered) input context (concatenating instructions, messages, and tools).
+    /// For most scenarios, overriding <see cref="ProvideAIContextAsync"/> is sufficient to provide additional context,
+    /// while still benefiting from the default filtering, merging and source stamping behavior.
+    /// However, for scenarios that require more control over context filtering, merging or source stamping, overriding this method
+    /// allows you to directly control the full <see cref="AIContext"/> returned for the invocation.
+    /// </para>
     /// </remarks>
-    protected abstract ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default);
+    protected virtual async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    {
+        var inputContext = context.AIContext;
+
+        // Create a filtered context for ProvideAIContextAsync, filtering input messages
+        // to exclude non-external messages (e.g. chat history, other AI context provider messages).
+        var filteredContext = new InvokingContext(
+            context.Agent,
+            context.Session,
+            new AIContext
+            {
+                Instructions = inputContext.Instructions,
+                Messages = inputContext.Messages is not null ? this._provideInputMessageFilter(inputContext.Messages) : null,
+                Tools = inputContext.Tools
+            });
+
+        var provided = await this.ProvideAIContextAsync(filteredContext, cancellationToken).ConfigureAwait(false);
+
+        var mergedInstructions = (inputContext.Instructions, provided.Instructions) switch
+        {
+            (null, null) => null,
+            (string a, null) => a,
+            (null, string b) => b,
+            (string a, string b) => a + "\n" + b
+        };
+
+        var providedMessages = provided.Messages is not null
+            ? provided.Messages.Select(m => m.WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!))
+            : null;
+
+        var mergedMessages = (inputContext.Messages, providedMessages) switch
+        {
+            (null, null) => null,
+            (var a, null) => a,
+            (null, var b) => b,
+            (var a, var b) => a.Concat(b)
+        };
+
+        var mergedTools = (inputContext.Tools, provided.Tools) switch
+        {
+            (null, null) => null,
+            (var a, null) => a,
+            (null, var b) => b,
+            (var a, var b) => a.Concat(b)
+        };
+
+        return new AIContext
+        {
+            Instructions = mergedInstructions,
+            Messages = mergedMessages,
+            Tools = mergedTools
+        };
+    }
+
+    /// <summary>
+    /// When overridden in a derived class, provides additional AI context to be merged with the input context for the current invocation.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// This method is called from <see cref="InvokingCoreAsync"/>.
+    /// Note that <see cref="InvokingCoreAsync"/> can be overridden to directly control context merging and source stamping, in which case
+    /// it is up to the implementer to call this method as needed to retrieve the additional context.
+    /// </para>
+    /// <para>
+    /// In contrast with <see cref="InvokingCoreAsync"/>, this method only returns additional context to be merged with the input,
+    /// while <see cref="InvokingCoreAsync"/> is responsible for returning the full merged <see cref="AIContext"/> for the invocation.
+    /// </para>
+    /// </remarks>
+    /// <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 an <see cref="AIContext"/>
+    /// with additional context to be merged with the input context.
+    /// </returns>
+    protected virtual ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    {
+        return new ValueTask<AIContext>(new AIContext());
+    }
 
     /// <summary>
     /// Called at the end of the agent invocation to process the invocation results.
@@ -106,7 +214,7 @@ public abstract class AIContextProvider
     /// </para>
     /// </remarks>
     public ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default)
-        => this.InvokedCoreAsync(context, cancellationToken);
+        => this.InvokedCoreAsync(Throw.IfNull(context), cancellationToken);
 
     /// <summary>
     /// Called at the end of the agent invocation to process the invocation results.
@@ -128,9 +236,50 @@ public abstract class AIContextProvider
     /// 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>
+    /// <para>
+    /// The default implementation of this method skips execution for any invocation failures,
+    /// filters the request messages using the configured store-input message filter
+    /// (which defaults to including only <see cref="AgentRequestMessageSourceType.External"/> messages),
+    /// and calls <see cref="StoreAIContextAsync"/> to process the invocation results.
+    /// For most scenarios, overriding <see cref="StoreAIContextAsync"/> is sufficient to process invocation results,
+    /// while still benefiting from the default error handling and filtering behavior.
+    /// However, for scenarios that require more control over error handling or message filtering, overriding this method
+    /// allows you to directly control the processing of invocation results.
+    /// </para>
     /// </remarks>
     protected virtual ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
-        => default;
+    {
+        if (context.InvokeException is not null)
+        {
+            return default;
+        }
+
+        var subContext = new InvokedContext(context.Agent, context.Session, this._storeInputMessageFilter(context.RequestMessages), context.ResponseMessages!);
+        return this.StoreAIContextAsync(subContext, cancellationToken);
+    }
+
+    /// <summary>
+    /// When overridden in a derived class, processes invocation results at the end of the agent invocation.
+    /// </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>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    /// <remarks>
+    /// <para>
+    /// This method is called from <see cref="InvokedCoreAsync"/>.
+    /// Note that <see cref="InvokedCoreAsync"/> can be overridden to directly control error handling, in which case
+    /// it is up to the implementer to call this method as needed to process the invocation results.
+    /// </para>
+    /// <para>
+    /// In contrast with <see cref="InvokedCoreAsync"/>, this method only processes the invocation results,
+    /// while <see cref="InvokedCoreAsync"/> is also responsible for error handling.
+    /// </para>
+    /// <para>
+    /// The default implementation of <see cref="InvokedCoreAsync"/> only calls this method if the invocation succeeded.
+    /// </para>
+    /// </remarks>
+    protected virtual ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) =>
+        default;
 
     /// <summary>Asks the <see cref="AIContextProvider"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
     /// <param name="serviceType">The type of object being requested.</param>

dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs

@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Text.Json;
+using Microsoft.Extensions.AI;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Provides an abstract base class for components that enhance AI context during agent invocations with support for maintaining provider state of type <typeparamref name="TState"/>.
+/// </summary>
+/// <typeparam name="TState">The type of the state to be maintained by the context provider. Must be a reference type.</typeparam>
+/// <remarks>
+/// This class extends <see cref="AIContextProvider"/> by introducing a strongly-typed state management mechanism, allowing derived classes to maintain and persist custom state information across invocations.
+/// The state is stored in the session's StateBag using a configurable key and JSON serialization options, enabling seamless integration with the agent session lifecycle.
+/// </remarks>
+public abstract class AIContextProvider<TState> : AIContextProvider
+    where TState : class
+{
+    private readonly Func<AgentSession?, TState> _stateInitializer;
+    private readonly string _stateKey;
+    private readonly JsonSerializerOptions _jsonSerializerOptions;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AIContextProvider{TState}"/> class.
+    /// </summary>
+    /// <param name="stateInitializer">A function to initialize the state for the context provider.</param>
+    /// <param name="stateKey">The key used to store the state in the session's StateBag.</param>
+    /// <param name="jsonSerializerOptions">Options for JSON serialization and deserialization of the state.</param>
+    /// <param name="provideInputMessageFilter">An optional filter function to apply to input messages before providing context. If not set, defaults to including only <see cref="AgentRequestMessageSourceType.External"/> messages.</param>
+    /// <param name="storeInputMessageFilter">An optional filter function to apply to request messages before storing context. If not set, defaults to including only <see cref="AgentRequestMessageSourceType.External"/> messages.</param>
+    protected AIContextProvider(
+        Func<AgentSession?, TState> stateInitializer,
+        string? stateKey,
+        JsonSerializerOptions? jsonSerializerOptions,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideInputMessageFilter,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputMessageFilter)
+        : base(provideInputMessageFilter, storeInputMessageFilter)
+    {
+        this._stateInitializer = stateInitializer;
+        this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
+        this._stateKey = stateKey ?? this.GetType().Name;
+    }
+
+    /// <inheritdoc />
+    public override string StateKey => this._stateKey;
+
+    /// <summary>
+    /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present.
+    /// </summary>
+    /// <param name="session">The agent session containing the StateBag.</param>
+    /// <returns>The provider state.</returns>
+    protected virtual TState GetOrInitializeState(AgentSession? session)
+    {
+        if (session?.StateBag.TryGetValue<TState>(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null)
+        {
+            return state;
+        }
+
+        state = this._stateInitializer(session);
+        if (session is not null)
+        {
+            session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions);
+        }
+
+        return state;
+    }
+
+    /// <summary>
+    /// Saves the specified state to the session's StateBag using the configured state key and JSON serializer options.
+    /// If the session is null, this method does nothing.
+    /// </summary>
+    /// <remarks>
+    /// This method provides a convenient way for derived classes to persist state changes back to the session after processing.
+    /// It abstracts away the details of how state is stored in the session, allowing derived classes to focus on their specific logic.
+    /// </remarks>
+    /// <param name="session">The agent session containing the StateBag.</param>
+    /// <param name="state">The state to be saved.</param>
+    protected virtual void SaveState(AgentSession? session, TState state)
+    {
+        if (session is not null)
+        {
+            session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions);
+        }
+    }
+}

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

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -39,6 +40,25 @@ namespace Microsoft.Agents.AI;
 /// </remarks>
 public abstract class ChatHistoryProvider
 {
+    private static IEnumerable<ChatMessage> DefaultExcludeChatHistoryFilter(IEnumerable<ChatMessage> messages)
+        => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory);
+
+    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? _provideOutputMessageFilter;
+    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _storeInputMessageFilter;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ChatHistoryProvider"/> class.
+    /// </summary>
+    /// <param name="provideOutputMessageFilter">An optional filter function to apply to messages when retrieving them from the chat history.</param>
+    /// <param name="storeInputMessageFilter">An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type <see cref="AgentRequestMessageSourceType.ChatHistory"/>.</param>
+    protected ChatHistoryProvider(
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideOutputMessageFilter = null,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputMessageFilter = null)
+    {
+        this._provideOutputMessageFilter = provideOutputMessageFilter;
+        this._storeInputMessageFilter = storeInputMessageFilter ?? DefaultExcludeChatHistoryFilter;
+    }
+
     /// <summary>
     /// Gets the key used to store the provider state in the <see cref="AgentSession.StateBag"/>.
     /// </summary>
@@ -50,20 +70,16 @@ public abstract class ChatHistoryProvider
     public virtual string StateKey => this.GetType().Name;
 
     /// <summary>
-    /// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation.
+    /// Called at the start of agent invocation to provide messages 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"/>
-    /// instances in ascending chronological order (oldest first).
+    /// instances that will be used for the agent invocation.
     /// </returns>
     /// <remarks>
     /// <para>
-    /// Messages are returned in chronological order to maintain proper conversation flow and context for the agent.
-    /// The oldest messages appear first in the collection, followed by more recent messages.
-    /// </para>
-    /// <para>
     /// If the total message history becomes very large, implementations should apply appropriate strategies to manage
     /// storage constraints, such as:
     /// <list type="bullet">
@@ -75,23 +91,19 @@ public abstract class ChatHistoryProvider
     /// </para>
     /// </remarks>
     public ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
-        => this.InvokingCoreAsync(context, cancellationToken);
+        => this.InvokingCoreAsync(Throw.IfNull(context), cancellationToken);
 
     /// <summary>
-    /// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation.
+    /// Called at the start of agent invocation to provide messages 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"/>
-    /// instances in ascending chronological order (oldest first).
+    /// instances that will be used for the agent invocation.
     /// </returns>
     /// <remarks>
     /// <para>
-    /// Messages are returned in chronological order to maintain proper conversation flow and context for the agent.
-    /// The oldest messages appear first in the collection, followed by more recent messages.
-    /// </para>
-    /// <para>
     /// If the total message history becomes very large, implementations should apply appropriate strategies to manage
     /// storage constraints, such as:
     /// <list type="bullet">
@@ -102,11 +114,54 @@ public abstract class ChatHistoryProvider
     /// </list>
     /// </para>
     /// <para>
-    /// Each <see cref="ChatHistoryProvider"/> instance should be associated with a single <see cref="AgentSession"/> to ensure proper message isolation
-    /// and context management.
+    /// The default implementation of this method, calls <see cref="ProvideChatHistoryAsync"/> to get the chat history messages, applies the optional retrieval output filter,
+    /// and merges the returned messages with the caller provided messages (with chat history messages appearing first) before returning the full message list to be used for the invocation.
+    /// For most scenarios, overriding <see cref="ProvideChatHistoryAsync"/> is sufficient to return the desired chat history messages, while still benefiting from the default merging and filtering behavior.
+    /// However, for scenarios that require more control over message filtering, merging or source stamping, overriding this method allows you to directly control the full set of messages returned for the invocation.
     /// </para>
     /// </remarks>
-    protected abstract ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default);
+    protected virtual async ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    {
+        var output = await this.ProvideChatHistoryAsync(context, cancellationToken).ConfigureAwait(false);
+
+        if (this._provideOutputMessageFilter is not null)
+        {
+            output = this._provideOutputMessageFilter(output);
+        }
+
+        return output
+            .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!))
+            .Concat(context.RequestMessages);
+    }
+
+    /// <summary>
+    /// When overridden in a derived class, provides the chat history messages to be used for the current invocation.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// This method is called from <see cref="InvokingCoreAsync"/>.
+    /// Note that <see cref="InvokingCoreAsync"/> can be overridden to directly control message filtering, merging and source stamping, in which case
+    /// it is up to the implementer to call this method as needed to retrieve the unfiltered/unmerged chat history messages.
+    /// </para>
+    /// <para>
+    /// In contrast with <see cref="InvokingCoreAsync"/>, this method only returns additional messages to be added to the request,
+    /// while <see cref="InvokingCoreAsync"/> is responsible for returning the full set of messages to be used for the invocation (including caller provided messages).
+    /// </para>
+    /// <para>
+    /// Messages are returned in chronological order to maintain proper conversation flow and context for the agent.
+    /// The oldest messages appear first in the collection, followed by more recent messages.
+    /// </para>
+    /// </remarks>
+    /// <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"/>
+    /// instances in ascending chronological order (oldest first).
+    /// </returns>
+    protected virtual ValueTask<IEnumerable<ChatMessage>> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    {
+        return new ValueTask<IEnumerable<ChatMessage>>([]);
+    }
 
     /// <summary>
     /// Called at the end of the agent invocation to add new messages to the chat history.
@@ -134,7 +189,7 @@ public abstract class ChatHistoryProvider
     /// </para>
     /// </remarks>
     public ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) =>
-        this.InvokedCoreAsync(context, cancellationToken);
+        this.InvokedCoreAsync(Throw.IfNull(context), cancellationToken);
 
     /// <summary>
     /// Called at the end of the agent invocation to add new messages to the chat history.
@@ -160,8 +215,59 @@ public abstract class ChatHistoryProvider
     /// 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>
+    /// <para>
+    /// The default implementation of this method, skips execution for any invocation failures, filters messages using the optional storage input message filter
+    /// and calls <see cref="StoreChatHistoryAsync"/> to store new chat history messages.
+    /// For most scenarios, overriding <see cref="StoreChatHistoryAsync"/> is sufficient to store chat history messages, while still benefiting from the default error handling and filtering behavior.
+    /// However, for scenarios that require more control over error handling or message filtering, overriding this method allows you to directly control the messages that are stored for the invocation.
+    /// </para>
     /// </remarks>
-    protected abstract ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default);
+    protected virtual ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+    {
+        if (context.InvokeException is not null)
+        {
+            return default;
+        }
+
+        var subContext = new InvokedContext(context.Agent, context.Session, this._storeInputMessageFilter(context.RequestMessages), context.ResponseMessages!);
+        return this.StoreChatHistoryAsync(subContext, cancellationToken);
+    }
+
+    /// <summary>
+    /// When overridden in a derived class, adds new messages to the chat history at the end of the agent invocation.
+    /// </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>
+    /// <returns>A task that represents the asynchronous add operation.</returns>
+    /// <remarks>
+    /// <para>
+    /// Messages should be added in the order they were generated to maintain proper chronological sequence.
+    /// The <see cref="ChatHistoryProvider"/> is responsible for preserving message ordering and ensuring that subsequent calls to
+    /// <see cref="InvokingCoreAsync"/> return messages in the correct chronological order.
+    /// </para>
+    /// <para>
+    /// Implementations may perform additional processing during message addition, such as:
+    /// <list type="bullet">
+    /// <item><description>Validating message content and metadata</description></item>
+    /// <item><description>Applying storage optimizations or compression</description></item>
+    /// <item><description>Triggering background maintenance operations</description></item>
+    /// </list>
+    /// </para>
+    /// <para>
+    /// This method is called from <see cref="InvokedCoreAsync"/>.
+    /// Note that <see cref="InvokedCoreAsync"/> can be overridden to directly control message filtering and error handling, in which case
+    /// it is up to the implementer to call this method as needed to store messages.
+    /// </para>
+    /// <para>
+    /// In contrast with <see cref="InvokedCoreAsync"/>, this method only stores messages,
+    /// while <see cref="InvokedCoreAsync"/> is also responsible for messages filtering and error handling.
+    /// </para>
+    /// <para>
+    /// The default implementation of <see cref="InvokedCoreAsync"/> only calls this method if the invocation succeeded.
+    /// </para>
+    /// </remarks>
+    protected virtual ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) =>
+        default;
 
     /// <summary>Asks the <see cref="ChatHistoryProvider"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
     /// <param name="serviceType">The type of object being requested.</param>

dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs

@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Text.Json;
+using Microsoft.Extensions.AI;
+
+namespace Microsoft.Agents.AI;
+
+/// <summary>
+/// Provides an abstract base class for fetching chat messages from, and adding chat messages to, chat history for the purposes of agent execution with support for maintaining provider state of type <typeparamref name="TState"/>.
+/// </summary>
+/// <typeparam name="TState">The type of the state to be maintained by the chat history provider. Must be a reference type.</typeparam>
+/// <remarks>
+/// This class extends <see cref="ChatHistoryProvider"/> by introducing a strongly-typed state management mechanism, allowing derived classes to maintain and persist custom state information across invocations.
+/// The state is stored in the session's StateBag using a configurable key and JSON serialization options, enabling seamless integration with the agent session lifecycle.
+/// </remarks>
+public abstract class ChatHistoryProvider<TState> : ChatHistoryProvider
+    where TState : class
+{
+    private readonly Func<AgentSession?, TState> _stateInitializer;
+    private readonly string _stateKey;
+    private readonly JsonSerializerOptions _jsonSerializerOptions;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ChatHistoryProvider{TState}"/> class.
+    /// </summary>
+    /// <param name="stateInitializer">A function to initialize the state for the chat history provider.</param>
+    /// <param name="stateKey">The key used to store the state in the session's StateBag.</param>
+    /// <param name="jsonSerializerOptions">Options for JSON serialization and deserialization of the state.</param>
+    /// <param name="provideOutputMessageFilter">A filter function to apply to messages when retrieving them from the chat history.</param>
+    /// <param name="storeInputMessageFilter">A filter function to apply to messages before storing them in the chat history.</param>
+    protected ChatHistoryProvider(
+        Func<AgentSession?, TState> stateInitializer,
+        string? stateKey,
+        JsonSerializerOptions? jsonSerializerOptions,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideOutputMessageFilter,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputMessageFilter)
+        : base(provideOutputMessageFilter, storeInputMessageFilter)
+    {
+        this._stateInitializer = stateInitializer;
+        this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
+        this._stateKey = stateKey ?? this.GetType().Name;
+    }
+
+    /// <inheritdoc />
+    public override string StateKey => this._stateKey;
+
+    /// <summary>
+    /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present.
+    /// </summary>
+    /// <param name="session">The agent session containing the StateBag.</param>
+    /// <returns>The provider state, or null if no session is available.</returns>
+    protected virtual TState GetOrInitializeState(AgentSession? session)
+    {
+        if (session?.StateBag.TryGetValue<TState>(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null)
+        {
+            return state;
+        }
+
+        state = this._stateInitializer(session);
+        if (session is not null)
+        {
+            session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions);
+        }
+
+        return state;
+    }
+
+    /// <summary>
+    /// Saves the specified state to the session's StateBag using the configured state key and JSON serializer options.
+    /// If the session is null, this method does nothing.
+    /// </summary>
+    /// <remarks>
+    /// This method provides a convenient way for derived classes to persist state changes back to the session after processing.
+    /// It abstracts away the details of how state is stored in the session, allowing derived classes to focus on their specific logic.
+    /// </remarks>
+    /// <param name="session">The agent session containing the StateBag.</param>
+    /// <param name="state">The state to be saved.</param>
+    protected virtual void SaveState(AgentSession? session, TState state)
+    {
+        if (session is not null)
+        {
+            session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions);
+        }
+    }
+}

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

@@ -3,7 +3,6 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Text.Json;
 using System.Text.Json.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
@@ -25,17 +24,8 @@ namespace Microsoft.Agents.AI;
 /// message reduction strategies or alternative storage implementations.
 /// </para>
 /// </remarks>
-public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
+public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider<InMemoryChatHistoryProvider.State>
 {
-    private static IEnumerable<ChatMessage> DefaultExcludeChatHistoryFilter(IEnumerable<ChatMessage> messages)
-        => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory);
-
-    private readonly string _stateKey;
-    private readonly Func<AgentSession?, State> _stateInitializer;
-    private readonly JsonSerializerOptions _jsonSerializerOptions;
-    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _storageInputMessageFilter;
-    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? _retrievalOutputMessageFilter;
-
     /// <summary>
     /// Initializes a new instance of the <see cref="InMemoryChatHistoryProvider"/> class.
     /// </summary>
@@ -44,19 +34,17 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
     /// message reduction, and serialization settings. If <see langword="null"/>, default settings will be used.
     /// </param>
     public InMemoryChatHistoryProvider(InMemoryChatHistoryProviderOptions? options = null)
+        : base(
+            options?.StateInitializer ?? (_ => new State()),
+            options?.StateKey,
+            options?.JsonSerializerOptions,
+            options?.ProvideOutputMessageFilter,
+            options?.StorageInputMessageFilter)
     {
-        this._stateInitializer = options?.StateInitializer ?? (_ => new State());
         this.ChatReducer = options?.ChatReducer;
         this.ReducerTriggerEvent = options?.ReducerTriggerEvent ?? InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval;
-        this._stateKey = options?.StateKey ?? base.StateKey;
-        this._jsonSerializerOptions = options?.JsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
-        this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExcludeChatHistoryFilter;
-        this._retrievalOutputMessageFilter = options?.RetrievalOutputMessageFilter;
     }
 
-    /// <inheritdoc />
-    public override string StateKey => this._stateKey;
-
     /// <summary>
     /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied.
     /// </summary>
@@ -89,32 +77,9 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
         state.Messages = messages;
     }
 
-    /// <summary>
-    /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present.
-    /// </summary>
-    /// <param name="session">The agent session containing the StateBag.</param>
-    /// <returns>The provider state, or null if no session is available.</returns>
-    private State GetOrInitializeState(AgentSession? session)
-    {
-        if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null)
-        {
-            return state;
-        }
-
-        state = this._stateInitializer(session);
-        if (session is not null)
-        {
-            session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions);
-        }
-
-        return state;
-    }
-
     /// <inheritdoc />
-    protected override async ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    protected override async ValueTask<IEnumerable<ChatMessage>> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {
-        _ = Throw.IfNull(context);
-
         var state = this.GetOrInitializeState(context.Session);
 
         if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null)
@@ -122,30 +87,16 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
             state.Messages = (await this.ChatReducer.ReduceAsync(state.Messages, cancellationToken).ConfigureAwait(false)).ToList();
         }
 
-        IEnumerable<ChatMessage> output = state.Messages;
-        if (this._retrievalOutputMessageFilter is not null)
-        {
-            output = this._retrievalOutputMessageFilter(output);
-        }
-        return output
-            .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!))
-            .Concat(context.RequestMessages);
+        return state.Messages;
     }
 
     /// <inheritdoc />
-    protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+    protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
-        _ = Throw.IfNull(context);
-
-        if (context.InvokeException is not null)
-        {
-            return;
-        }
-
         var state = this.GetOrInitializeState(context.Session);
 
         // Add request and response messages to the provider
-        var allNewMessages = this._storageInputMessageFilter(context.RequestMessages).Concat(context.ResponseMessages ?? []);
+        var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []);
         state.Messages.AddRange(allNewMessages);
 
         if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null)

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

@@ -71,7 +71,7 @@ public sealed class InMemoryChatHistoryProviderOptions
     /// <value>
     /// When <see langword="null"/>, no filtering is applied to the output messages.
     /// </value>
-    public Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? RetrievalOutputMessageFilter { get; set; }
+    public Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? ProvideOutputMessageFilter { get; set; }
 
     /// <summary>
     /// Defines the events that can trigger a reducer in the <see cref="InMemoryChatHistoryProvider"/>.

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

@@ -19,16 +19,11 @@ namespace Microsoft.Agents.AI;
 /// </summary>
 [RequiresUnreferencedCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with trimming.")]
 [RequiresDynamicCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with NativeAOT.")]
-public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
+public sealed class CosmosChatHistoryProvider : ChatHistoryProvider<CosmosChatHistoryProvider.State>, IDisposable
 {
-    private static IEnumerable<ChatMessage> DefaultExcludeChatHistoryFilter(IEnumerable<ChatMessage> messages)
-        => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory);
-
     private readonly CosmosClient _cosmosClient;
     private readonly Container _container;
     private readonly bool _ownsClient;
-    private readonly string _stateKey;
-    private readonly Func<AgentSession?, State> _stateInitializer;
     private bool _disposed;
 
     /// <summary>
@@ -46,9 +41,6 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         return options;
     }
 
-    /// <inheritdoc />
-    public override string StateKey => this._stateKey;
-
     /// <summary>
     /// Gets or sets the maximum number of messages to return in a single query batch.
     /// Default is 100 for optimal performance.
@@ -84,25 +76,6 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// </summary>
     public string ContainerId { get; init; }
 
-    /// <summary>
-    /// A filter function applied to request messages before they are stored
-    /// during <see cref="ChatHistoryProvider.InvokedAsync"/>. The default filter excludes messages with the
-    /// <see cref="AgentRequestMessageSourceType.ChatHistory"/> source type.
-    /// </summary>
-    public Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> StorageInputMessageFilter { get; set { field = Throw.IfNull(value); } } = DefaultExcludeChatHistoryFilter;
-
-    /// <summary>
-    /// Gets or sets an optional filter function applied to messages produced by this provider
-    /// during <see cref="ChatHistoryProvider.InvokingAsync"/>.
-    /// </summary>
-    /// <remarks>
-    /// This filter is only applied to the messages that the provider itself produces (from its internal storage).
-    /// </remarks>
-    /// <value>
-    /// When <see langword="null"/>, no filtering is applied to the output messages.
-    /// </value>
-    public Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? RetrievalOutputMessageFilter { get; set; }
-
     /// <summary>
     /// Initializes a new instance of the <see cref="CosmosChatHistoryProvider"/> class.
     /// </summary>
@@ -112,6 +85,8 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// <param name="stateInitializer">A delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId).</param>
     /// <param name="ownsClient">Whether this instance owns the CosmosClient and should dispose it.</param>
     /// <param name="stateKey">An optional key to use for storing the state in the <see cref="AgentSession.StateBag"/>.</param>
+    /// <param name="provideOutputMessageFilter">An optional filter function to apply to messages when retrieving them from the chat history.</param>
+    /// <param name="storeInputMessageFilter">An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type <see cref="AgentRequestMessageSourceType.ChatHistory"/>.</param>
     /// <exception cref="ArgumentNullException">Thrown when <paramref name="cosmosClient"/> or <paramref name="stateInitializer"/> is <see langword="null"/>.</exception>
     /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
     public CosmosChatHistoryProvider(
@@ -120,15 +95,16 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         string containerId,
         Func<AgentSession?, State> stateInitializer,
         bool ownsClient = false,
-        string? stateKey = null)
+        string? stateKey = null,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideOutputMessageFilter = null,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputMessageFilter = null)
+        : base(Throw.IfNull(stateInitializer), stateKey, null, provideOutputMessageFilter, storeInputMessageFilter)
     {
         this._cosmosClient = Throw.IfNull(cosmosClient);
         this.DatabaseId = Throw.IfNullOrWhitespace(databaseId);
         this.ContainerId = Throw.IfNullOrWhitespace(containerId);
         this._container = this._cosmosClient.GetContainer(databaseId, containerId);
-        this._stateInitializer = Throw.IfNull(stateInitializer);
         this._ownsClient = ownsClient;
-        this._stateKey = stateKey ?? base.StateKey;
     }
 
     /// <summary>
@@ -139,6 +115,8 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// <param name="containerId">The identifier of the Cosmos DB container.</param>
     /// <param name="stateInitializer">A delegate that initializes the provider state on the first invocation.</param>
     /// <param name="stateKey">An optional key to use for storing the state in the <see cref="AgentSession.StateBag"/>.</param>
+    /// <param name="provideOutputMessageFilter">An optional filter function to apply to messages when retrieving them from the chat history.</param>
+    /// <param name="storeInputMessageFilter">An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type <see cref="AgentRequestMessageSourceType.ChatHistory"/>.</param>
     /// <exception cref="ArgumentNullException">Thrown when any required parameter is null.</exception>
     /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
     public CosmosChatHistoryProvider(
@@ -146,8 +124,10 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         string databaseId,
         string containerId,
         Func<AgentSession?, State> stateInitializer,
-        string? stateKey = null)
-        : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey)
+        string? stateKey = null,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideOutputMessageFilter = null,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputMessageFilter = null)
+        : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputMessageFilter)
     {
     }
 
@@ -160,6 +140,8 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     /// <param name="containerId">The identifier of the Cosmos DB container.</param>
     /// <param name="stateInitializer">A delegate that initializes the provider state on the first invocation.</param>
     /// <param name="stateKey">An optional key to use for storing the state in the <see cref="AgentSession.StateBag"/>.</param>
+    /// <param name="provideOutputMessageFilter">An optional filter function to apply to messages when retrieving them from the chat history.</param>
+    /// <param name="storeInputMessageFilter">An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type <see cref="AgentRequestMessageSourceType.ChatHistory"/>.</param>
     /// <exception cref="ArgumentNullException">Thrown when any required parameter is null.</exception>
     /// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
     public CosmosChatHistoryProvider(
@@ -168,32 +150,13 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         string databaseId,
         string containerId,
         Func<AgentSession?, State> stateInitializer,
-        string? stateKey = null)
-        : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey)
+        string? stateKey = null,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideOutputMessageFilter = null,
+        Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputMessageFilter = null)
+        : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputMessageFilter)
     {
     }
 
-    /// <summary>
-    /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present.
-    /// </summary>
-    /// <param name="session">The agent session containing the StateBag.</param>
-    /// <returns>The provider state, or null if no session is available.</returns>
-    private State GetOrInitializeState(AgentSession? session)
-    {
-        if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state, AgentAbstractionsJsonUtilities.DefaultOptions) is true && state is not null)
-        {
-            return state;
-        }
-
-        state = this._stateInitializer(session);
-        if (session is not null)
-        {
-            session.StateBag.SetValue(this._stateKey, state, AgentAbstractionsJsonUtilities.DefaultOptions);
-        }
-
-        return state;
-    }
-
     /// <summary>
     /// Determines whether hierarchical partitioning should be used based on the state.
     /// </summary>
@@ -218,7 +181,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
     }
 
     /// <inheritdoc />
-    protected override async ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    protected override async ValueTask<IEnumerable<ChatMessage>> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {
 #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
         if (this._disposed)
@@ -227,8 +190,6 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
         }
 #pragma warning restore CA1513
 
-        _ = Throw.IfNull(context);
-
         var state = this.GetOrInitializeState(context.Session);
         var partitionKey = BuildPartitionKey(state);
 
@@ -279,22 +240,12 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
             messages.Reverse();
         }
 
-        return (this.RetrievalOutputMessageFilter is not null ? this.RetrievalOutputMessageFilter(messages) : messages)
-            .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!))
-            .Concat(context.RequestMessages);
+        return messages;
     }
 
     /// <inheritdoc />
-    protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+    protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
-        Throw.IfNull(context);
-
-        if (context.InvokeException is not null)
-        {
-            // 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
         if (this._disposed)
         {
@@ -303,7 +254,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
 #pragma warning restore CA1513
 
         var state = this.GetOrInitializeState(context.Session);
-        var messageList = this.StorageInputMessageFilter(context.RequestMessages).Concat(context.ResponseMessages ?? []).ToList();
+        var messageList = context.RequestMessages.Concat(context.ResponseMessages ?? []).ToList();
         if (messageList.Count == 0)
         {
             return;

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

@@ -22,19 +22,12 @@ namespace Microsoft.Agents.AI.Mem0;
 /// for new invocations using a semantic search endpoint. Retrieved memories are injected as user messages
 /// to the model, prefixed by a configurable context prompt.
 /// </remarks>
-public sealed class Mem0Provider : AIContextProvider
+public sealed class Mem0Provider : AIContextProvider<Mem0Provider.State>
 {
     private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:";
 
-    private static IEnumerable<ChatMessage> DefaultExternalOnlyFilter(IEnumerable<ChatMessage> messages)
-        => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External);
-
     private readonly string _contextPrompt;
     private readonly bool _enableSensitiveTelemetryData;
-    private readonly string _stateKey;
-    private readonly Func<AgentSession?, State> _stateInitializer;
-    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _searchInputMessageFilter;
-    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _storageInputMessageFilter;
 
     private readonly Mem0Client _client;
     private readonly ILogger<Mem0Provider>? _logger;
@@ -58,6 +51,7 @@ public sealed class Mem0Provider : AIContextProvider
     /// </code>
     /// </remarks>
     public Mem0Provider(HttpClient httpClient, Func<AgentSession?, State> stateInitializer, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null)
+        : base(ValidateStateInitializer(Throw.IfNull(stateInitializer)), options?.StateKey, Mem0JsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter)
     {
         Throw.IfNull(httpClient);
         if (string.IsNullOrWhiteSpace(httpClient.BaseAddress?.AbsoluteUri))
@@ -65,63 +59,41 @@ public sealed class Mem0Provider : AIContextProvider
             throw new ArgumentException("The HttpClient BaseAddress must be set for Mem0 operations.", nameof(httpClient));
         }
 
-        this._stateInitializer = Throw.IfNull(stateInitializer);
         this._logger = loggerFactory?.CreateLogger<Mem0Provider>();
         this._client = new Mem0Client(httpClient);
 
         this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt;
         this._enableSensitiveTelemetryData = options?.EnableSensitiveTelemetryData ?? false;
-        this._stateKey = options?.StateKey ?? base.StateKey;
-        this._searchInputMessageFilter = options?.SearchInputMessageFilter ?? DefaultExternalOnlyFilter;
-        this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExternalOnlyFilter;
     }
 
-    /// <inheritdoc />
-    public override string StateKey => this._stateKey;
-
-    /// <summary>
-    /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present.
-    /// </summary>
-    /// <param name="session">The agent session containing the StateBag.</param>
-    /// <returns>The provider state, or null if no session is available.</returns>
-    private State? GetOrInitializeState(AgentSession? session)
-    {
-        if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state, Mem0JsonUtilities.DefaultOptions) is true && state is not null)
+    private static Func<AgentSession?, State> ValidateStateInitializer(Func<AgentSession?, State> stateInitializer) =>
+        session =>
         {
+            var state = stateInitializer(session);
+
+            if (state is null
+                || state.StorageScope is null
+                || (state.StorageScope.AgentId is null && state.StorageScope.ThreadId is null && state.StorageScope.UserId is null && state.StorageScope.ApplicationId is null)
+                || state.SearchScope is null
+                || (state.SearchScope.AgentId is null && state.SearchScope.ThreadId is null && state.SearchScope.UserId is null && state.SearchScope.ApplicationId is null))
+            {
+                throw new InvalidOperationException("State initializer must return a non-null state with valid storage and search scopes, where at least one scoping parameter is set for each.");
+            }
+
             return state;
-        }
-
-        state = this._stateInitializer(session);
-
-        if (state is null
-            || state.StorageScope is null
-            || (state.StorageScope.AgentId is null && state.StorageScope.ThreadId is null && state.StorageScope.UserId is null && state.StorageScope.ApplicationId is null)
-            || state.SearchScope is null
-            || (state.SearchScope.AgentId is null && state.SearchScope.ThreadId is null && state.SearchScope.UserId is null && state.SearchScope.ApplicationId is null))
-        {
-            throw new InvalidOperationException("State initializer must return a non-null state with valid storage and search scopes, where at lest one scoping parameter is set for each.");
-        }
-
-        if (session is not null)
-        {
-            session.StateBag.SetValue(this._stateKey, state, Mem0JsonUtilities.DefaultOptions);
-        }
-
-        return state;
-    }
+        };
 
     /// <inheritdoc />
-    protected override async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    protected override async ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {
         Throw.IfNull(context);
 
-        var inputContext = context.AIContext;
         var state = this.GetOrInitializeState(context.Session);
-        var searchScope = state?.SearchScope ?? new Mem0ProviderScope();
+        var searchScope = state.SearchScope;
 
         string queryText = string.Join(
             Environment.NewLine,
-                this._searchInputMessageFilter(inputContext.Messages ?? [])
+                (context.AIContext.Messages ?? [])
                 .Where(m => !string.IsNullOrWhiteSpace(m.Text))
                 .Select(m => m.Text));
 
@@ -138,9 +110,6 @@ public sealed class Mem0Provider : AIContextProvider
             var outputMessageText = memories.Count == 0
                 ? null
                 : $"{this._contextPrompt}\n{string.Join(Environment.NewLine, memories)}";
-            var outputMessage = memories.Count == 0
-                ? null
-                : new ChatMessage(ChatRole.User, outputMessageText!).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!);
 
             if (this._logger?.IsEnabled(LogLevel.Information) is true)
             {
@@ -167,11 +136,9 @@ public sealed class Mem0Provider : AIContextProvider
 
             return new AIContext
             {
-                Instructions = inputContext.Instructions,
-                Messages =
-                    (inputContext.Messages ?? [])
-                    .Concat(outputMessage is not null ? [outputMessage] : []),
-                Tools = inputContext.Tools
+                Messages = outputMessageText is not null
+                    ? [new ChatMessage(ChatRole.User, outputMessageText)]
+                    : null
             };
         }
         catch (ArgumentException)
@@ -190,27 +157,23 @@ public sealed class Mem0Provider : AIContextProvider
                     searchScope.ThreadId,
                     this.SanitizeLogData(searchScope.UserId));
             }
-            return inputContext;
+
+            return new AIContext();
         }
     }
 
     /// <inheritdoc />
-    protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+    protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
-        if (context.InvokeException is not null)
-        {
-            return; // Do not update memory on failed invocations.
-        }
-
         var state = this.GetOrInitializeState(context.Session);
-        var storageScope = state?.StorageScope ?? new Mem0ProviderScope();
+        var storageScope = state.StorageScope;
 
         try
         {
             // Persist request and response messages after invocation.
             await this.PersistMessagesAsync(
                 storageScope,
-                this._storageInputMessageFilter(context.RequestMessages)
+                context.RequestMessages
                     .Concat(context.ResponseMessages ?? []),
                 cancellationToken).ConfigureAwait(false);
         }
@@ -238,12 +201,7 @@ public sealed class Mem0Provider : AIContextProvider
     {
         Throw.IfNull(session);
         var state = this.GetOrInitializeState(session);
-        var storageScope = state?.StorageScope;
-
-        if (storageScope is null)
-        {
-            return Task.CompletedTask; // Nothing to clear if there is no state.
-        }
+        var storageScope = state.StorageScope;
 
         return this._client.ClearMemoryAsync(
             storageScope.ApplicationId,

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

@@ -9,10 +9,8 @@ using Microsoft.Extensions.AI;
 
 namespace Microsoft.Agents.AI.Workflows;
 
-internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
+internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider<WorkflowChatHistoryProvider.StoreState>
 {
-    private readonly JsonSerializerOptions _jsonSerializerOptions;
-
     /// <summary>
     /// Initializes a new instance of the <see cref="WorkflowChatHistoryProvider"/> class.
     /// </summary>
@@ -22,8 +20,8 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
     /// and source generated serializers are required, or Native AOT / Trimming is required.
     /// </param>
     public WorkflowChatHistoryProvider(JsonSerializerOptions? jsonSerializerOptions = null)
+        : base(stateInitializer: _ => new StoreState(), stateKey: null, jsonSerializerOptions: jsonSerializerOptions, provideOutputMessageFilter: null, storeInputMessageFilter: null)
     {
-        this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions;
     }
 
     internal sealed class StoreState
@@ -32,43 +30,16 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
         public List<ChatMessage> Messages { get; set; } = [];
     }
 
-    private StoreState GetOrInitializeState(AgentSession? session)
-    {
-        if (session?.StateBag.TryGetValue<StoreState>(this.StateKey, out var state, this._jsonSerializerOptions) is true && state is not null)
-        {
-            return state;
-        }
-
-        state = new();
-        if (session is not null)
-        {
-            session.StateBag.SetValue(this.StateKey, state, this._jsonSerializerOptions);
-        }
-
-        return state;
-    }
-
     internal void AddMessages(AgentSession session, params IEnumerable<ChatMessage> messages)
         => this.GetOrInitializeState(session).Messages.AddRange(messages);
 
-    protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
-        => new(this.GetOrInitializeState(context.Session)
-            .Messages
-            .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!))
-            .Concat(context.RequestMessages));
+    protected override ValueTask<IEnumerable<ChatMessage>> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default)
+        => new(this.GetOrInitializeState(context.Session).Messages);
 
-    protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+    protected override ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
-        if (context.InvokeException is not null)
-        {
-            return default;
-        }
-
-        var allNewMessages = context.RequestMessages
-            .Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory)
-            .Concat(context.ResponseMessages ?? []);
+        var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []);
         this.GetOrInitializeState(context.Session).Messages.AddRange(allNewMessages);
-
         return default;
     }
 

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

@@ -24,8 +24,8 @@ namespace Microsoft.Agents.AI;
 /// abstractions to work with any compatible vector store implementation.
 /// </para>
 /// <para>
-/// Messages are stored during the <see cref="InvokedCoreAsync"/> method and retrieved during the
-/// <see cref="InvokingCoreAsync"/> method using semantic similarity search.
+/// Messages are stored during the <see cref="StoreAIContextAsync"/> method and retrieved during the
+/// <see cref="ProvideAIContextAsync"/> method using semantic similarity search.
 /// </para>
 /// <para>
 /// Behavior is configurable through <see cref="ChatHistoryMemoryProviderOptions"/>. When
@@ -34,16 +34,13 @@ namespace Microsoft.Agents.AI;
 /// injecting them automatically on each invocation.
 /// </para>
 /// </remarks>
-public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
+public sealed class ChatHistoryMemoryProvider : AIContextProvider<ChatHistoryMemoryProvider.State>, IDisposable
 {
     private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:";
     private const int DefaultMaxResults = 3;
     private const string DefaultFunctionToolName = "Search";
     private const string DefaultFunctionToolDescription = "Allows searching for related previous chat history to help answer the user question.";
 
-    private static IEnumerable<ChatMessage> DefaultExternalOnlyFilter(IEnumerable<ChatMessage> messages)
-        => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External);
-
 #pragma warning disable CA2213 // VectorStore is not owned by this class - caller is responsible for disposal
     private readonly VectorStore _vectorStore;
 #pragma warning restore CA2213
@@ -55,10 +52,6 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
     private readonly string _toolName;
     private readonly string _toolDescription;
     private readonly ILogger<ChatHistoryMemoryProvider>? _logger;
-    private readonly string _stateKey;
-    private readonly Func<AgentSession?, State> _stateInitializer;
-    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _searchInputMessageFilter;
-    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _storageInputMessageFilter;
 
     private bool _collectionInitialized;
     private readonly SemaphoreSlim _initializationLock = new(1, 1);
@@ -81,21 +74,18 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
         Func<AgentSession?, State> stateInitializer,
         ChatHistoryMemoryProviderOptions? options = null,
         ILoggerFactory? loggerFactory = null)
+        : base(Throw.IfNull(stateInitializer), options?.StateKey, AgentJsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter)
     {
         this._vectorStore = Throw.IfNull(vectorStore);
-        this._stateInitializer = Throw.IfNull(stateInitializer);
 
         options ??= new ChatHistoryMemoryProviderOptions();
         this._maxResults = options.MaxResults.HasValue ? Throw.IfLessThanOrEqual(options.MaxResults.Value, 0) : DefaultMaxResults;
         this._contextPrompt = options.ContextPrompt ?? DefaultContextPrompt;
         this._enableSensitiveTelemetryData = options.EnableSensitiveTelemetryData;
         this._searchTime = options.SearchTime;
-        this._stateKey = options.StateKey ?? base.StateKey;
         this._logger = loggerFactory?.CreateLogger<ChatHistoryMemoryProvider>();
         this._toolName = options.FunctionToolName ?? DefaultFunctionToolName;
         this._toolDescription = options.FunctionToolDescription ?? DefaultFunctionToolDescription;
-        this._searchInputMessageFilter = options.SearchInputMessageFilter ?? DefaultExternalOnlyFilter;
-        this._storageInputMessageFilter = options.StorageInputMessageFilter ?? DefaultExternalOnlyFilter;
 
         // Create a definition so that we can use the dimensions provided at runtime.
         var definition = new VectorStoreCollectionDefinition
@@ -120,37 +110,12 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._stateKey;
-
-    /// <summary>
-    /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present.
-    /// </summary>
-    /// <param name="session">The agent session containing the StateBag.</param>
-    /// <returns>The provider state, or null if no session is available.</returns>
-    private State? GetOrInitializeState(AgentSession? session)
-    {
-        if (session?.StateBag.TryGetValue<State>(this._stateKey, out var state, AgentJsonUtilities.DefaultOptions) is true && state is not null)
-        {
-            return state;
-        }
-
-        state = this._stateInitializer(session);
-        if (state is not null && session is not null)
-        {
-            session.StateBag.SetValue(this._stateKey, state, AgentJsonUtilities.DefaultOptions);
-        }
-
-        return state;
-    }
-
-    /// <inheritdoc />
-    protected override async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    protected override async ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {
         _ = Throw.IfNull(context);
 
-        var inputContext = context.AIContext;
         var state = this.GetOrInitializeState(context.Session);
-        var searchScope = state?.SearchScope ?? new ChatHistoryMemoryProviderScope();
+        var searchScope = state.SearchScope;
 
         if (this._searchTime == ChatHistoryMemoryProviderOptions.SearchBehavior.OnDemandFunctionCalling)
         {
@@ -166,12 +131,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
                     description: this._toolDescription)
             ];
 
-            // Expose search tool for on-demand invocation by the model, accumulated with the input context
+            // Expose search tool for on-demand invocation by the model
             return new AIContext
             {
-                Instructions = inputContext.Instructions,
-                Messages = inputContext.Messages,
-                Tools = (inputContext.Tools ?? []).Concat(tools)
+                Tools = tools
             };
         }
 
@@ -179,13 +142,13 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
         {
             // Get the text from the current request messages
             var requestText = string.Join("\n",
-                this._searchInputMessageFilter(inputContext.Messages ?? [])
+                (context.AIContext.Messages ?? [])
                 .Where(m => m != null && !string.IsNullOrWhiteSpace(m.Text))
                 .Select(m => m.Text));
 
             if (string.IsNullOrWhiteSpace(requestText))
             {
-                return inputContext;
+                return new AIContext();
             }
 
             // Search for relevant chat history
@@ -193,19 +156,12 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
 
             if (string.IsNullOrWhiteSpace(contextText))
             {
-                return inputContext;
+                return new AIContext();
             }
 
             return new AIContext
             {
-                Instructions = inputContext.Instructions,
-                Messages =
-                    (inputContext.Messages ?? [])
-                    .Concat(
-                    [
-                        new ChatMessage(ChatRole.User, contextText).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!)
-                    ]),
-                Tools = inputContext.Tools
+                Messages = [new ChatMessage(ChatRole.User, contextText)]
             };
         }
         catch (Exception ex)
@@ -221,30 +177,24 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
                     this.SanitizeLogData(searchScope.UserId));
             }
 
-            return inputContext;
+            return new AIContext();
         }
     }
 
     /// <inheritdoc />
-    protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+    protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
         _ = Throw.IfNull(context);
 
-        // Only store if invocation was successful
-        if (context.InvokeException != null)
-        {
-            return;
-        }
-
         var state = this.GetOrInitializeState(context.Session);
-        var storageScope = state?.StorageScope ?? new ChatHistoryMemoryProviderScope();
+        var storageScope = state.StorageScope;
 
         try
         {
             // Ensure the collection is initialized
             var collection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false);
 
-            List<Dictionary<string, object?>> itemsToStore = this._storageInputMessageFilter(context.RequestMessages)
+            List<Dictionary<string, object?>> itemsToStore = context.RequestMessages
                 .Concat(context.ResponseMessages ?? [])
                 .Select(message => new Dictionary<string, object?>
                 {

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

@@ -32,16 +32,13 @@ namespace Microsoft.Agents.AI;
 /// multi-turn context to the retrieval layer without permanently altering the conversation history.
 /// </para>
 /// </remarks>
-public sealed class TextSearchProvider : AIContextProvider
+public sealed class TextSearchProvider : AIContextProvider<TextSearchProvider.TextSearchProviderState>
 {
     private const string DefaultPluginSearchFunctionName = "Search";
     private const string DefaultPluginSearchFunctionDescription = "Allows searching for additional information to help answer the user question.";
     private const string DefaultContextPrompt = "## Additional Context\nConsider the following information from source documents when responding to the user:";
     private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available.";
 
-    private static IEnumerable<ChatMessage> DefaultExternalOnlyFilter(IEnumerable<ChatMessage> messages)
-        => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External);
-
     private readonly Func<string, CancellationToken, Task<IEnumerable<TextSearchResult>>> _searchAsync;
     private readonly ILogger<TextSearchProvider>? _logger;
     private readonly AITool[] _tools;
@@ -50,10 +47,7 @@ public sealed class TextSearchProvider : AIContextProvider
     private readonly TextSearchProviderOptions.TextSearchBehavior _searchTime;
     private readonly string _contextPrompt;
     private readonly string _citationsPrompt;
-    private readonly string _stateKey;
     private readonly Func<IList<TextSearchResult>, string>? _contextFormatter;
-    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _searchInputMessageFilter;
-    private readonly Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>> _storageInputMessageFilter;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="TextSearchProvider"/> class.
@@ -66,6 +60,7 @@ public sealed class TextSearchProvider : AIContextProvider
         Func<string, CancellationToken, Task<IEnumerable<TextSearchResult>>> searchAsync,
         TextSearchProviderOptions? options = null,
         ILoggerFactory? loggerFactory = null)
+        : base(_ => new TextSearchProviderState(), options?.StateKey, AgentJsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter)
     {
         // Validate and assign parameters
         this._searchAsync = Throw.IfNull(searchAsync);
@@ -75,10 +70,7 @@ public sealed class TextSearchProvider : AIContextProvider
         this._searchTime = options?.SearchTime ?? TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke;
         this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt;
         this._citationsPrompt = options?.CitationsPrompt ?? DefaultCitationsPrompt;
-        this._stateKey = options?.StateKey ?? base.StateKey;
         this._contextFormatter = options?.ContextFormatter;
-        this._searchInputMessageFilter = options?.SearchInputMessageFilter ?? DefaultExternalOnlyFilter;
-        this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExternalOnlyFilter;
 
         // Create the on-demand search tool (only used if behavior is OnDemandFunctionCalling)
         this._tools =
@@ -91,32 +83,25 @@ public sealed class TextSearchProvider : AIContextProvider
     }
 
     /// <inheritdoc />
-    public override string StateKey => this._stateKey;
-
-    /// <inheritdoc />
-    protected override async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+    protected override async ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
     {
-        var inputContext = context.AIContext;
-
         if (this._searchTime != TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke)
         {
-            // Expose the search tool for on-demand invocation, accumulated with the input context.
+            // Expose the search tool for on-demand invocation.
             return new AIContext
             {
-                Instructions = inputContext.Instructions,
-                Messages = inputContext.Messages,
-                Tools = (inputContext.Tools ?? []).Concat(this._tools)
+                Tools = this._tools
             };
         }
 
-        // Retrieve recent messages from the session state bag.
-        var recentMessagesText = context.Session?.StateBag.GetValue<TextSearchProviderState>(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText
+        // Retrieve recent messages from the session state.
+        var recentMessagesText = this.GetOrInitializeState(context.Session).RecentMessagesText
             ?? [];
 
         // Aggregate text from memory + current request messages.
         var sbInput = new StringBuilder();
         var requestMessagesText =
-            this._searchInputMessageFilter(inputContext.Messages ?? [])
+            (context.AIContext.Messages ?? [])
             .Where(x => !string.IsNullOrWhiteSpace(x?.Text)).Select(x => x.Text);
         foreach (var messageText in recentMessagesText.Concat(requestMessagesText))
         {
@@ -142,7 +127,7 @@ public sealed class TextSearchProvider : AIContextProvider
 
             if (materialized.Count == 0)
             {
-                return inputContext;
+                return new AIContext();
             }
 
             // Format search results
@@ -155,25 +140,18 @@ public sealed class TextSearchProvider : AIContextProvider
 
             return new AIContext
             {
-                Instructions = inputContext.Instructions,
-                Messages =
-                    (inputContext.Messages ?? [])
-                    .Concat(
-                    [
-                        new ChatMessage(ChatRole.User, formatted).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!)
-                    ]),
-                Tools = inputContext.Tools
+                Messages = [new ChatMessage(ChatRole.User, formatted)]
             };
         }
         catch (Exception ex)
         {
             this._logger?.LogError(ex, "TextSearchProvider: Failed to search for data due to error");
-            return inputContext;
+            return new AIContext();
         }
     }
 
     /// <inheritdoc />
-    protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+    protected override ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default)
     {
         int limit = this._recentMessageMemoryLimit;
         if (limit <= 0)
@@ -186,16 +164,11 @@ public sealed class TextSearchProvider : AIContextProvider
             return default; // No session to store state in.
         }
 
-        if (context.InvokeException is not null)
-        {
-            return default; // Do not update memory on failed invocations.
-        }
-
-        // Retrieve existing recent messages from the session state bag.
-        var recentMessagesText = context.Session.StateBag.GetValue<TextSearchProviderState>(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText
+        // Retrieve existing recent messages from the session state.
+        var recentMessagesText = this.GetOrInitializeState(context.Session).RecentMessagesText
             ?? [];
 
-        var newMessagesText = this._storageInputMessageFilter(context.RequestMessages)
+        var newMessagesText = context.RequestMessages
             .Concat(context.ResponseMessages ?? [])
             .Where(m =>
                 this._recentMessageRolesIncluded.Contains(m.Role) &&
@@ -208,11 +181,10 @@ public sealed class TextSearchProvider : AIContextProvider
             ? allMessages.Skip(allMessages.Count - limit).ToList()
             : allMessages;
 
-        // Store updated state back to the session state bag.
-        context.Session.StateBag.SetValue(
-            this._stateKey,
-            new TextSearchProviderState { RecentMessagesText = updatedMessages },
-            AgentJsonUtilities.DefaultOptions);
+        // Store updated state back to the session.
+        this.SaveState(
+            context.Session,
+            new TextSearchProviderState { RecentMessagesText = updatedMessages });
 
         return default;
     }
@@ -311,8 +283,14 @@ public sealed class TextSearchProvider : AIContextProvider
         public object? RawRepresentation { get; set; }
     }
 
-    internal sealed class TextSearchProviderState
+    /// <summary>
+    /// Represents the per-session state of a <see cref="TextSearchProvider"/> stored in the <see cref="AgentSession.StateBag"/>.
+    /// </summary>
+    public sealed class TextSearchProviderState
     {
+        /// <summary>
+        /// Gets or sets the list of recent message texts retained for multi-turn search context.
+        /// </summary>
         public List<string>? RecentMessagesText { get; set; }
     }
 }

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

@@ -0,0 +1,199 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Moq;
+
+namespace Microsoft.Agents.AI.Abstractions.UnitTests;
+
+/// <summary>
+/// Contains tests for the <see cref="AIContextProvider{TState}"/> class.
+/// </summary>
+public class AIContextProviderTStateTests
+{
+    private static readonly AIAgent s_mockAgent = new Mock<AIAgent>().Object;
+
+    #region GetOrInitializeState Tests
+
+    [Fact]
+    public void GetOrInitializeState_InitializesFromStateInitializerOnFirstCall()
+    {
+        // Arrange
+        var expectedState = new TestState { Value = "initialized" };
+        var provider = new TestAIContextProvider(_ => expectedState);
+        var session = new TestAgentSession();
+
+        // Act
+        var state = provider.GetState(session);
+
+        // Assert
+        Assert.Same(expectedState, state);
+    }
+
+    [Fact]
+    public void GetOrInitializeState_ReturnsCachedStateFromStateBagOnSecondCall()
+    {
+        // Arrange
+        var callCount = 0;
+        var provider = new TestAIContextProvider(_ =>
+        {
+            callCount++;
+            return new TestState { Value = $"init-{callCount}" };
+        });
+        var session = new TestAgentSession();
+
+        // Act
+        var state1 = provider.GetState(session);
+        var state2 = provider.GetState(session);
+
+        // Assert - initializer called only once; second call reads from StateBag
+        Assert.Equal(1, callCount);
+        Assert.Equal("init-1", state1.Value);
+        Assert.Equal("init-1", state2.Value);
+    }
+
+    [Fact]
+    public void GetOrInitializeState_WorksWhenSessionIsNull()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider(_ => new TestState { Value = "no-session" });
+
+        // Act
+        var state = provider.GetState(null);
+
+        // Assert
+        Assert.Equal("no-session", state.Value);
+    }
+
+    [Fact]
+    public void GetOrInitializeState_ReInitializesWhenSessionIsNull()
+    {
+        // Arrange - without a session, state can't be cached in StateBag
+        var callCount = 0;
+        var provider = new TestAIContextProvider(_ =>
+        {
+            callCount++;
+            return new TestState { Value = $"init-{callCount}" };
+        });
+
+        // Act
+        provider.GetState(null);
+        provider.GetState(null);
+
+        // Assert - initializer called each time since there's no session to cache in
+        Assert.Equal(2, callCount);
+    }
+
+    #endregion
+
+    #region SaveState Tests
+
+    [Fact]
+    public void SaveState_SavesToStateBag()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider(_ => new TestState());
+        var session = new TestAgentSession();
+        var state = new TestState { Value = "saved" };
+
+        // Act
+        provider.DoSaveState(session, state);
+        var retrieved = provider.GetState(session);
+
+        // Assert
+        Assert.Equal("saved", retrieved.Value);
+    }
+
+    [Fact]
+    public void SaveState_NoOpWhenSessionIsNull()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider(_ => new TestState { Value = "default" });
+
+        // Act - should not throw
+        provider.DoSaveState(null, new TestState { Value = "saved" });
+
+        // Assert - no exception; can't verify further without a session
+    }
+
+    #endregion
+
+    #region StateKey Tests
+
+    [Fact]
+    public void StateKey_DefaultsToTypeName()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider(_ => new TestState());
+
+        // Act & Assert
+        Assert.Equal(nameof(TestAIContextProvider), provider.StateKey);
+    }
+
+    [Fact]
+    public void StateKey_UsesCustomKeyWhenProvided()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider(_ => new TestState(), stateKey: "custom-key");
+
+        // Act & Assert
+        Assert.Equal("custom-key", provider.StateKey);
+    }
+
+    #endregion
+
+    #region Integration Tests
+
+    [Fact]
+    public async Task InvokingCoreAsync_CanUseStateInProvideAIContextAsync()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider(_ => new TestState { Value = "state-value" });
+        var session = new TestAgentSession();
+        var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hi")] };
+        var context = new AIContextProvider.InvokingContext(s_mockAgent, session, inputContext);
+
+        // Act
+        var result = await provider.InvokingAsync(context);
+
+        // Assert - the provider uses state to produce context messages
+        var messages = result.Messages!.ToList();
+        Assert.Equal(2, messages.Count);
+        Assert.Contains("state-value", messages[1].Text);
+    }
+
+    #endregion
+
+    public sealed class TestState
+    {
+        public string Value { get; set; } = string.Empty;
+    }
+
+    private sealed class TestAIContextProvider : AIContextProvider<TestState>
+    {
+        public TestAIContextProvider(
+            Func<AgentSession?, TestState> stateInitializer,
+            string? stateKey = null)
+            : base(stateInitializer, stateKey, null, null, null)
+        {
+        }
+
+        public TestState GetState(AgentSession? session) => this.GetOrInitializeState(session);
+
+        public void DoSaveState(AgentSession? session, TestState state) => this.SaveState(session, state);
+
+        protected override ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
+        {
+            var state = this.GetOrInitializeState(context.Session);
+            return new(new AIContext
+            {
+                Messages = [new ChatMessage(ChatRole.System, $"Context from state: {state.Value}")]
+            });
+        }
+    }
+
+    private sealed class TestAgentSession : AgentSession;
+}

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

@@ -3,6 +3,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -337,9 +338,314 @@ public class AIContextProviderTests
 
     #endregion
 
+    #region InvokingAsync / InvokedAsync Null Check Tests
+
+    [Fact]
+    public async Task InvokingAsync_NullContext_ThrowsArgumentNullExceptionAsync()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider();
+
+        // Act & Assert
+        await Assert.ThrowsAsync<ArgumentNullException>(() => provider.InvokingAsync(null!).AsTask());
+    }
+
+    [Fact]
+    public async Task InvokedAsync_NullContext_ThrowsArgumentNullExceptionAsync()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider();
+
+        // Act & Assert
+        await Assert.ThrowsAsync<ArgumentNullException>(() => provider.InvokedAsync(null!).AsTask());
+    }
+
+    #endregion
+
+    #region InvokingCoreAsync Tests
+
+    [Fact]
+    public async Task InvokingCoreAsync_CallsProvideAIContextAndReturnsMergedContextAsync()
+    {
+        // Arrange
+        var providedMessages = new[] { new ChatMessage(ChatRole.System, "Context message") };
+        var provider = new TestAIContextProvider(provideContext: new AIContext { Messages = providedMessages });
+        var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "User input")] };
+        var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext);
+
+        // Act
+        var result = await provider.InvokingAsync(context);
+
+        // Assert - input messages + provided messages merged
+        var messages = result.Messages!.ToList();
+        Assert.Equal(2, messages.Count);
+        Assert.Equal("User input", messages[0].Text);
+        Assert.Equal("Context message", messages[1].Text);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_FiltersInputToExternalOnlyByDefaultAsync()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider(captureFilteredContext: true);
+        var externalMsg = new ChatMessage(ChatRole.User, "External");
+        var chatHistoryMsg = new ChatMessage(ChatRole.User, "History")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src");
+        var contextProviderMsg = new ChatMessage(ChatRole.User, "ContextProvider")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "src");
+        var inputContext = new AIContext { Messages = [externalMsg, chatHistoryMsg, contextProviderMsg] };
+        var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext);
+
+        // Act
+        await provider.InvokingAsync(context);
+
+        // Assert - ProvideAIContextAsync received only External messages
+        Assert.NotNull(provider.LastProvidedContext);
+        var filteredMessages = provider.LastProvidedContext!.AIContext.Messages!.ToList();
+        Assert.Single(filteredMessages);
+        Assert.Equal("External", filteredMessages[0].Text);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_StampsProvidedMessagesWithAIContextProviderSourceAsync()
+    {
+        // Arrange
+        var providedMessages = new[] { new ChatMessage(ChatRole.System, "Provided") };
+        var provider = new TestAIContextProvider(provideContext: new AIContext { Messages = providedMessages });
+        var inputContext = new AIContext { Messages = [] };
+        var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext);
+
+        // Act
+        var result = await provider.InvokingAsync(context);
+
+        // Assert
+        var messages = result.Messages!.ToList();
+        Assert.Single(messages);
+        Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, messages[0].GetAgentRequestMessageSourceType());
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_MergesInstructionsAsync()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider(provideContext: new AIContext { Instructions = "Provided instructions" });
+        var inputContext = new AIContext { Instructions = "Input instructions" };
+        var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext);
+
+        // Act
+        var result = await provider.InvokingAsync(context);
+
+        // Assert - instructions are joined with newline
+        Assert.Equal("Input instructions\nProvided instructions", result.Instructions);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_MergesToolsAsync()
+    {
+        // Arrange
+        var inputTool = AIFunctionFactory.Create(() => "a", "inputTool");
+        var providedTool = AIFunctionFactory.Create(() => "b", "providedTool");
+        var provider = new TestAIContextProvider(provideContext: new AIContext { Tools = [providedTool] });
+        var inputContext = new AIContext { Tools = [inputTool] };
+        var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext);
+
+        // Act
+        var result = await provider.InvokingAsync(context);
+
+        // Assert - both tools present
+        var tools = result.Tools!.ToList();
+        Assert.Equal(2, tools.Count);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_UsesCustomProvideInputFilterAsync()
+    {
+        // Arrange - filter that keeps all messages (not just External)
+        var provider = new TestAIContextProvider(
+            captureFilteredContext: true,
+            provideInputMessageFilter: msgs => msgs);
+        var externalMsg = new ChatMessage(ChatRole.User, "External");
+        var chatHistoryMsg = new ChatMessage(ChatRole.User, "History")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src");
+        var inputContext = new AIContext { Messages = [externalMsg, chatHistoryMsg] };
+        var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext);
+
+        // Act
+        await provider.InvokingAsync(context);
+
+        // Assert - ProvideAIContextAsync received ALL messages (custom filter keeps everything)
+        Assert.NotNull(provider.LastProvidedContext);
+        var filteredMessages = provider.LastProvidedContext!.AIContext.Messages!.ToList();
+        Assert.Equal(2, filteredMessages.Count);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_ReturnsEmptyContextByDefaultAsync()
+    {
+        // Arrange - provider that doesn't override ProvideAIContextAsync
+        var provider = new DefaultAIContextProvider();
+        var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] };
+        var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext);
+
+        // Act
+        var result = await provider.InvokingAsync(context);
+
+        // Assert - only the input messages (no additional provided)
+        var messages = result.Messages!.ToList();
+        Assert.Single(messages);
+        Assert.Equal("Hello", messages[0].Text);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_MergesWithOriginalUnfilteredMessagesAsync()
+    {
+        // Arrange - default filter is External-only, but the MERGED result should include
+        // the original unfiltered input messages plus the provided messages
+        var providedMessages = new[] { new ChatMessage(ChatRole.System, "Provided") };
+        var provider = new TestAIContextProvider(provideContext: new AIContext { Messages = providedMessages });
+        var externalMsg = new ChatMessage(ChatRole.User, "External");
+        var chatHistoryMsg = new ChatMessage(ChatRole.User, "History")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src");
+        var inputContext = new AIContext { Messages = [externalMsg, chatHistoryMsg] };
+        var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext);
+
+        // Act
+        var result = await provider.InvokingAsync(context);
+
+        // Assert - original 2 input messages + 1 provided message
+        var messages = result.Messages!.ToList();
+        Assert.Equal(3, messages.Count);
+        Assert.Equal("External", messages[0].Text);
+        Assert.Equal("History", messages[1].Text);
+        Assert.Equal("Provided", messages[2].Text);
+    }
+
+    #endregion
+
+    #region InvokedCoreAsync Tests
+
+    [Fact]
+    public async Task InvokedCoreAsync_CallsStoreAIContextWithFilteredMessagesAsync()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider();
+        var externalMessage = new ChatMessage(ChatRole.User, "External");
+        var chatHistoryMessage = new ChatMessage(ChatRole.User, "History")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src");
+        var responseMessages = new[] { new ChatMessage(ChatRole.Assistant, "Response") };
+        var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, new[] { externalMessage, chatHistoryMessage }, responseMessages);
+
+        // Act
+        await provider.InvokedAsync(context);
+
+        // Assert - default filter keeps only External messages
+        Assert.NotNull(provider.LastStoredContext);
+        var storedRequest = provider.LastStoredContext!.RequestMessages.ToList();
+        Assert.Single(storedRequest);
+        Assert.Equal("External", storedRequest[0].Text);
+        Assert.Same(responseMessages, provider.LastStoredContext.ResponseMessages);
+    }
+
+    [Fact]
+    public async Task InvokedCoreAsync_SkipsStorageWhenInvokeExceptionIsNotNullAsync()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider();
+        var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "msg")], new InvalidOperationException("Failed"));
+
+        // Act
+        await provider.InvokedAsync(context);
+
+        // Assert - StoreAIContextAsync was NOT called
+        Assert.Null(provider.LastStoredContext);
+    }
+
+    [Fact]
+    public async Task InvokedCoreAsync_UsesCustomStoreInputFilterAsync()
+    {
+        // Arrange - filter that only keeps System messages
+        var provider = new TestAIContextProvider(
+            storeInputMessageFilter: msgs => msgs.Where(m => m.Role == ChatRole.System));
+        var messages = new[]
+        {
+            new ChatMessage(ChatRole.User, "User msg"),
+            new ChatMessage(ChatRole.System, "System msg")
+        };
+        var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, messages, [new ChatMessage(ChatRole.Assistant, "Response")]);
+
+        // Act
+        await provider.InvokedAsync(context);
+
+        // Assert - only System messages were passed to store
+        Assert.NotNull(provider.LastStoredContext);
+        var storedRequest = provider.LastStoredContext!.RequestMessages.ToList();
+        Assert.Single(storedRequest);
+        Assert.Equal("System msg", storedRequest[0].Text);
+    }
+
+    [Fact]
+    public async Task InvokedCoreAsync_DefaultFilterExcludesNonExternalMessagesAsync()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider();
+        var external = new ChatMessage(ChatRole.User, "External");
+        var fromHistory = new ChatMessage(ChatRole.User, "History")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src");
+        var fromContext = new ChatMessage(ChatRole.User, "Context")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "src");
+        var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [external, fromHistory, fromContext], []);
+
+        // Act
+        await provider.InvokedAsync(context);
+
+        // Assert - only External messages kept
+        Assert.NotNull(provider.LastStoredContext);
+        var storedRequest = provider.LastStoredContext!.RequestMessages.ToList();
+        Assert.Single(storedRequest);
+        Assert.Equal("External", storedRequest[0].Text);
+    }
+
+    #endregion
+
     private sealed class TestAIContextProvider : AIContextProvider
     {
-        protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
-            => new(new AIContext());
+        private readonly AIContext? _provideContext;
+        private readonly bool _captureFilteredContext;
+
+        public InvokedContext? LastStoredContext { get; private set; }
+
+        public InvokingContext? LastProvidedContext { get; private set; }
+
+        public TestAIContextProvider(
+            AIContext? provideContext = null,
+            bool captureFilteredContext = false,
+            Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideInputMessageFilter = null,
+            Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputMessageFilter = null)
+            : base(provideInputMessageFilter, storeInputMessageFilter)
+        {
+            this._provideContext = provideContext;
+            this._captureFilteredContext = captureFilteredContext;
+        }
+
+        protected override ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
+        {
+            if (this._captureFilteredContext)
+            {
+                this.LastProvidedContext = context;
+            }
+
+            return new(this._provideContext ?? new AIContext());
+        }
+
+        protected override ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default)
+        {
+            this.LastStoredContext = context;
+            return default;
+        }
     }
+
+    /// <summary>
+    /// A provider that uses only base class defaults (no overrides of ProvideAIContextAsync/StoreAIContextAsync).
+    /// </summary>
+    private sealed class DefaultAIContextProvider : AIContextProvider;
 }

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

@@ -0,0 +1,195 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Moq;
+
+namespace Microsoft.Agents.AI.Abstractions.UnitTests;
+
+/// <summary>
+/// Contains tests for the <see cref="ChatHistoryProvider{TState}"/> class.
+/// </summary>
+public class ChatHistoryProviderTStateTests
+{
+    private static readonly AIAgent s_mockAgent = new Mock<AIAgent>().Object;
+
+    #region GetOrInitializeState Tests
+
+    [Fact]
+    public void GetOrInitializeState_InitializesFromStateInitializerOnFirstCall()
+    {
+        // Arrange
+        var expectedState = new TestState { Value = "initialized" };
+        var provider = new TestChatHistoryProvider(_ => expectedState);
+        var session = new TestAgentSession();
+
+        // Act
+        var state = provider.GetState(session);
+
+        // Assert
+        Assert.Same(expectedState, state);
+    }
+
+    [Fact]
+    public void GetOrInitializeState_ReturnsCachedStateFromStateBagOnSecondCall()
+    {
+        // Arrange
+        var callCount = 0;
+        var provider = new TestChatHistoryProvider(_ =>
+        {
+            callCount++;
+            return new TestState { Value = $"init-{callCount}" };
+        });
+        var session = new TestAgentSession();
+
+        // Act
+        var state1 = provider.GetState(session);
+        var state2 = provider.GetState(session);
+
+        // Assert - initializer called only once; second call reads from StateBag
+        Assert.Equal(1, callCount);
+        Assert.Equal("init-1", state1.Value);
+        Assert.Equal("init-1", state2.Value);
+    }
+
+    [Fact]
+    public void GetOrInitializeState_WorksWhenSessionIsNull()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider(_ => new TestState { Value = "no-session" });
+
+        // Act
+        var state = provider.GetState(null);
+
+        // Assert
+        Assert.Equal("no-session", state.Value);
+    }
+
+    [Fact]
+    public void GetOrInitializeState_ReInitializesWhenSessionIsNull()
+    {
+        // Arrange - without a session, state can't be cached in StateBag
+        var callCount = 0;
+        var provider = new TestChatHistoryProvider(_ =>
+        {
+            callCount++;
+            return new TestState { Value = $"init-{callCount}" };
+        });
+
+        // Act
+        _ = provider.GetState(null);
+        provider.GetState(null);
+
+        // Assert - initializer called each time since there's no session to cache in
+        Assert.Equal(2, callCount);
+    }
+
+    #endregion
+
+    #region SaveState Tests
+
+    [Fact]
+    public void SaveState_SavesToStateBag()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider(_ => new TestState());
+        var session = new TestAgentSession();
+        var state = new TestState { Value = "saved" };
+
+        // Act
+        provider.DoSaveState(session, state);
+        var retrieved = provider.GetState(session);
+
+        // Assert
+        Assert.Equal("saved", retrieved.Value);
+    }
+
+    [Fact]
+    public void SaveState_NoOpWhenSessionIsNull()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider(_ => new TestState { Value = "default" });
+
+        // Act - should not throw
+        provider.DoSaveState(null, new TestState { Value = "saved" });
+
+        // Assert - no exception; can't verify further without a session
+    }
+
+    #endregion
+
+    #region StateKey Tests
+
+    [Fact]
+    public void StateKey_DefaultsToTypeName()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider(_ => new TestState());
+
+        // Act & Assert
+        Assert.Equal(nameof(TestChatHistoryProvider), provider.StateKey);
+    }
+
+    [Fact]
+    public void StateKey_UsesCustomKeyWhenProvided()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider(_ => new TestState(), stateKey: "custom-key");
+
+        // Act & Assert
+        Assert.Equal("custom-key", provider.StateKey);
+    }
+
+    #endregion
+
+    #region Integration Tests
+
+    [Fact]
+    public async Task InvokingCoreAsync_CanUseStateInProvideChatHistoryAsync()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider(_ => new TestState { Value = "state-value" });
+        var session = new TestAgentSession();
+        var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Hi")]);
+
+        // Act
+        var result = (await provider.InvokingAsync(context)).ToList();
+
+        // Assert - the provider uses state to produce history messages
+        Assert.Equal(2, result.Count);
+        Assert.Contains("state-value", result[0].Text);
+    }
+
+    #endregion
+
+    public sealed class TestState
+    {
+        public string Value { get; set; } = string.Empty;
+    }
+
+    private sealed class TestChatHistoryProvider : ChatHistoryProvider<TestState>
+    {
+        public TestChatHistoryProvider(
+            Func<AgentSession?, TestState> stateInitializer,
+            string? stateKey = null)
+            : base(stateInitializer, stateKey, null, null, null)
+        {
+        }
+
+        public TestState GetState(AgentSession? session) => this.GetOrInitializeState(session);
+
+        public void DoSaveState(AgentSession? session, TestState state) => this.SaveState(session, state);
+
+        protected override ValueTask<IEnumerable<ChatMessage>> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default)
+        {
+            var state = this.GetOrInitializeState(context.Session);
+            return new(new[] { new ChatMessage(ChatRole.System, $"History from state: {state.Value}") });
+        }
+    }
+
+    private sealed class TestAgentSession : AgentSession;
+}

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

@@ -274,12 +274,279 @@ public class ChatHistoryProviderTests
 
     #endregion
 
+    #region InvokingAsync / InvokedAsync Null Check Tests
+
+    [Fact]
+    public async Task InvokingAsync_NullContext_ThrowsArgumentNullExceptionAsync()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider();
+
+        // Act & Assert
+        await Assert.ThrowsAsync<ArgumentNullException>(() => provider.InvokingAsync(null!).AsTask());
+    }
+
+    [Fact]
+    public async Task InvokedAsync_NullContext_ThrowsArgumentNullExceptionAsync()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider();
+
+        // Act & Assert
+        await Assert.ThrowsAsync<ArgumentNullException>(() => provider.InvokedAsync(null!).AsTask());
+    }
+
+    #endregion
+
+    #region InvokingCoreAsync Tests
+
+    [Fact]
+    public async Task InvokingCoreAsync_CallsProvideChatHistoryAndReturnsMessagesAsync()
+    {
+        // Arrange
+        var historyMessages = new[] { new ChatMessage(ChatRole.User, "History message") };
+        var provider = new TestChatHistoryProvider(provideMessages: historyMessages);
+        var requestMessages = new[] { new ChatMessage(ChatRole.User, "Request message") };
+        var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, requestMessages);
+
+        // Act
+        var result = (await provider.InvokingAsync(context)).ToList();
+
+        // Assert
+        Assert.Equal(2, result.Count);
+        Assert.Equal("History message", result[0].Text);
+        Assert.Equal("Request message", result[1].Text);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_HistoryAppearsBeforeRequestMessagesAsync()
+    {
+        // Arrange
+        var historyMessages = new[]
+        {
+            new ChatMessage(ChatRole.User, "Hist1"),
+            new ChatMessage(ChatRole.Assistant, "Hist2")
+        };
+        var provider = new TestChatHistoryProvider(provideMessages: historyMessages);
+        var requestMessages = new[] { new ChatMessage(ChatRole.User, "Req1") };
+        var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, requestMessages);
+
+        // Act
+        var result = (await provider.InvokingAsync(context)).ToList();
+
+        // Assert
+        Assert.Equal(3, result.Count);
+        Assert.Equal("Hist1", result[0].Text);
+        Assert.Equal("Hist2", result[1].Text);
+        Assert.Equal("Req1", result[2].Text);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_StampsHistoryMessagesWithChatHistorySourceAsync()
+    {
+        // Arrange
+        var historyMessages = new[] { new ChatMessage(ChatRole.User, "History") };
+        var provider = new TestChatHistoryProvider(provideMessages: historyMessages);
+        var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []);
+
+        // Act
+        var result = (await provider.InvokingAsync(context)).ToList();
+
+        // Assert
+        Assert.Single(result);
+        Assert.Equal(AgentRequestMessageSourceType.ChatHistory, result[0].GetAgentRequestMessageSourceType());
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_NoFilterAppliedWhenProvideOutputFilterIsNullAsync()
+    {
+        // Arrange
+        var historyMessages = new[]
+        {
+            new ChatMessage(ChatRole.User, "User msg"),
+            new ChatMessage(ChatRole.System, "System msg"),
+            new ChatMessage(ChatRole.Assistant, "Assistant msg")
+        };
+        var provider = new TestChatHistoryProvider(provideMessages: historyMessages);
+        var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []);
+
+        // Act
+        var result = (await provider.InvokingAsync(context)).ToList();
+
+        // Assert - all 3 history messages returned (no filter)
+        Assert.Equal(3, result.Count);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_AppliesProvideOutputFilterWhenProvidedAsync()
+    {
+        // Arrange
+        var historyMessages = new[]
+        {
+            new ChatMessage(ChatRole.User, "User msg"),
+            new ChatMessage(ChatRole.System, "System msg"),
+            new ChatMessage(ChatRole.Assistant, "Assistant msg")
+        };
+        var provider = new TestChatHistoryProvider(
+            provideMessages: historyMessages,
+            provideOutputMessageFilter: msgs => msgs.Where(m => m.Role == ChatRole.User));
+        var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []);
+
+        // Act
+        var result = (await provider.InvokingAsync(context)).ToList();
+
+        // Assert - only User messages remain after filter
+        Assert.Single(result);
+        Assert.Equal("User msg", result[0].Text);
+    }
+
+    [Fact]
+    public async Task InvokingCoreAsync_ReturnsEmptyHistoryByDefaultAsync()
+    {
+        // Arrange - provider that doesn't override ProvideChatHistoryAsync (uses base default)
+        var provider = new DefaultChatHistoryProvider();
+        var requestMessages = new[] { new ChatMessage(ChatRole.User, "Hello") };
+        var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, requestMessages);
+
+        // Act
+        var result = (await provider.InvokingAsync(context)).ToList();
+
+        // Assert - only the request message (no history)
+        Assert.Single(result);
+        Assert.Equal("Hello", result[0].Text);
+    }
+
+    #endregion
+
+    #region InvokedCoreAsync Tests
+
+    [Fact]
+    public async Task InvokedCoreAsync_CallsStoreChatHistoryWithFilteredMessagesAsync()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider();
+        var externalMessage = new ChatMessage(ChatRole.User, "External");
+        var chatHistoryMessage = new ChatMessage(ChatRole.User, "From history")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "source");
+        var responseMessages = new[] { new ChatMessage(ChatRole.Assistant, "Response") };
+        var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, new[] { externalMessage, chatHistoryMessage }, responseMessages);
+
+        // Act
+        await provider.InvokedAsync(context);
+
+        // Assert - default filter excludes ChatHistory-sourced messages
+        Assert.NotNull(provider.LastStoredContext);
+        var storedRequest = provider.LastStoredContext!.RequestMessages.ToList();
+        Assert.Single(storedRequest);
+        Assert.Equal("External", storedRequest[0].Text);
+        Assert.Same(responseMessages, provider.LastStoredContext.ResponseMessages);
+    }
+
+    [Fact]
+    public async Task InvokedCoreAsync_SkipsStorageWhenInvokeExceptionIsNotNullAsync()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider();
+        var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "msg")], new InvalidOperationException("Failed"));
+
+        // Act
+        await provider.InvokedAsync(context);
+
+        // Assert - StoreChatHistoryAsync was NOT called
+        Assert.Null(provider.LastStoredContext);
+    }
+
+    [Fact]
+    public async Task InvokedCoreAsync_UsesCustomStoreInputFilterAsync()
+    {
+        // Arrange - filter that only keeps System messages
+        var provider = new TestChatHistoryProvider(
+            storeInputMessageFilter: msgs => msgs.Where(m => m.Role == ChatRole.System));
+        var messages = new[]
+        {
+            new ChatMessage(ChatRole.User, "User msg"),
+            new ChatMessage(ChatRole.System, "System msg")
+        };
+        var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, messages, [new ChatMessage(ChatRole.Assistant, "Response")]);
+
+        // Act
+        await provider.InvokedAsync(context);
+
+        // Assert - only System messages were passed to store
+        Assert.NotNull(provider.LastStoredContext);
+        var storedRequest = provider.LastStoredContext!.RequestMessages.ToList();
+        Assert.Single(storedRequest);
+        Assert.Equal("System msg", storedRequest[0].Text);
+    }
+
+    [Fact]
+    public async Task InvokedCoreAsync_DefaultFilterExcludesChatHistorySourcedMessagesAsync()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider();
+        var external = new ChatMessage(ChatRole.User, "External");
+        var fromHistory = new ChatMessage(ChatRole.User, "History")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src");
+        var fromContext = new ChatMessage(ChatRole.User, "Context")
+            .WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "src");
+        var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [external, fromHistory, fromContext], []);
+
+        // Act
+        await provider.InvokedAsync(context);
+
+        // Assert - External and AIContextProvider messages kept, ChatHistory excluded
+        Assert.NotNull(provider.LastStoredContext);
+        var storedRequest = provider.LastStoredContext!.RequestMessages.ToList();
+        Assert.Equal(2, storedRequest.Count);
+        Assert.Equal("External", storedRequest[0].Text);
+        Assert.Equal("Context", storedRequest[1].Text);
+    }
+
+    [Fact]
+    public async Task InvokedCoreAsync_PassesResponseMessagesToStoreAsync()
+    {
+        // Arrange
+        var provider = new TestChatHistoryProvider();
+        var responseMessages = new[] { new ChatMessage(ChatRole.Assistant, "Resp1"), new ChatMessage(ChatRole.Assistant, "Resp2") };
+        var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "msg")], responseMessages);
+
+        // Act
+        await provider.InvokedAsync(context);
+
+        // Assert
+        Assert.NotNull(provider.LastStoredContext);
+        Assert.Same(responseMessages, provider.LastStoredContext!.ResponseMessages);
+    }
+
+    #endregion
+
     private sealed class TestChatHistoryProvider : ChatHistoryProvider
     {
-        protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
-            => new(new ChatMessage[] { new(ChatRole.User, "Test Message") }.Concat(context.RequestMessages));
+        private readonly IEnumerable<ChatMessage>? _provideMessages;
 
-        protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
-            => default;
+        public InvokedContext? LastStoredContext { get; private set; }
+
+        public TestChatHistoryProvider(
+            IEnumerable<ChatMessage>? provideMessages = null,
+            Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideOutputMessageFilter = null,
+            Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputMessageFilter = null)
+            : base(provideOutputMessageFilter, storeInputMessageFilter)
+        {
+            this._provideMessages = provideMessages;
+        }
+
+        protected override ValueTask<IEnumerable<ChatMessage>> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default)
+            => new(this._provideMessages ?? []);
+
+        protected override ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default)
+        {
+            this.LastStoredContext = context;
+            return default;
+        }
     }
+
+    /// <summary>
+    /// A provider that uses only base class defaults (no overrides of ProvideChatHistoryAsync/StoreChatHistoryAsync).
+    /// </summary>
+    private sealed class DefaultChatHistoryProvider : ChatHistoryProvider;
 }

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

@@ -446,7 +446,7 @@ public class InMemoryChatHistoryProviderTests
         var session = CreateMockSession();
         var provider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions
         {
-            RetrievalOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User)
+            ProvideOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User)
         });
         provider.SetMessages(session,
         [

dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs

@@ -881,12 +881,12 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
         this.SkipIfEmulatorNotAvailable();
         var session = CreateMockSession();
         var conversationId = Guid.NewGuid().ToString();
-        using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId,
-            _ => new CosmosChatHistoryProvider.State(conversationId))
-        {
-            // Custom filter: only store External messages (also exclude AIContextProvider)
-            StorageInputMessageFilter = messages => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External)
-        };
+        using var provider = new CosmosChatHistoryProvider(
+            this._connectionString,
+            s_testDatabaseId,
+            TestContainerId,
+            _ => new CosmosChatHistoryProvider.State(conversationId),
+            storeInputMessageFilter: messages => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External));
 
         var requestMessages = new[]
         {
@@ -919,12 +919,12 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
         this.SkipIfEmulatorNotAvailable();
         var session = CreateMockSession();
         var conversationId = Guid.NewGuid().ToString();
-        using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId,
-            _ => new CosmosChatHistoryProvider.State(conversationId))
-        {
-            // Only return User messages when retrieving
-            RetrievalOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User)
-        };
+        using var provider = new CosmosChatHistoryProvider(
+            this._connectionString,
+            s_testDatabaseId,
+            TestContainerId,
+            _ => new CosmosChatHistoryProvider.State(conversationId),
+            provideOutputMessageFilter: messages => messages.Where(m => m.Role == ChatRole.User));
 
         var requestMessages = new[]
         {
@@ -943,7 +943,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
         var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []);
         var messages = (await provider.InvokingAsync(invokingContext)).ToList();
 
-        // Assert - Only User messages returned (System and Assistant filtered by RetrievalOutputMessageFilter)
+        // Assert - Only User messages returned (System and Assistant filtered by ProvideOutputMessageFilter)
         Assert.Single(messages);
         Assert.Equal("User message", messages[0].Text);
         Assert.Equal(ChatRole.User, messages[0].Role);

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

@@ -115,8 +115,8 @@ public class ChatClientAgentOptionsTests
         const string Description = "Test description";
         var tools = new List<AITool> { AIFunctionFactory.Create(() => "test") };
 
-        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>().Object;
-        var mockAIContextProvider = new Mock<AIContextProvider>().Object;
+        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null).Object;
+        var mockAIContextProvider = new Mock<AIContextProvider>(null, null).Object;
 
         var original = new ChatClientAgentOptions()
         {
@@ -149,8 +149,8 @@ public class ChatClientAgentOptionsTests
     public void Clone_WithoutProvidingChatOptions_ClonesCorrectly()
     {
         // Arrange
-        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>().Object;
-        var mockAIContextProvider = new Mock<AIContextProvider>().Object;
+        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null).Object;
+        var mockAIContextProvider = new Mock<AIContextProvider>(null, null).Object;
 
         var original = new ChatClientAgentOptions
         {

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

@@ -488,7 +488,7 @@ public partial class ChatClientAgentTests
             })
             .ReturnsAsync(new ChatResponse(responseMessages));
 
-        var mockProvider = new Mock<AIContextProvider>();
+        var mockProvider = new Mock<AIContextProvider>(null, null);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -559,7 +559,7 @@ public partial class ChatClientAgentTests
                 It.IsAny<CancellationToken>()))
             .Throws(new InvalidOperationException("downstream failure"));
 
-        var mockProvider = new Mock<AIContextProvider>();
+        var mockProvider = new Mock<AIContextProvider>(null, null);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -617,7 +617,7 @@ public partial class ChatClientAgentTests
             })
             .ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
 
-        var mockProvider = new Mock<AIContextProvider>();
+        var mockProvider = new Mock<AIContextProvider>(null, null);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -677,7 +677,7 @@ public partial class ChatClientAgentTests
             .ReturnsAsync(new ChatResponse(responseMessages));
 
         // Provider 1: adds a system message and a tool
-        var mockProvider1 = new Mock<AIContextProvider>();
+        var mockProvider1 = new Mock<AIContextProvider>(null, null);
         mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
         mockProvider1
             .Protected()
@@ -696,7 +696,7 @@ public partial class ChatClientAgentTests
 
         // Provider 2: adds another system message and verifies it receives accumulated context from provider 1
         AIContext? provider2ReceivedContext = null;
-        var mockProvider2 = new Mock<AIContextProvider>();
+        var mockProvider2 = new Mock<AIContextProvider>(null, null);
         mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
         mockProvider2
             .Protected()
@@ -784,7 +784,7 @@ public partial class ChatClientAgentTests
                 It.IsAny<CancellationToken>()))
             .ThrowsAsync(new InvalidOperationException("downstream failure"));
 
-        var mockProvider1 = new Mock<AIContextProvider>();
+        var mockProvider1 = new Mock<AIContextProvider>(null, null);
         mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
         mockProvider1
             .Protected()
@@ -801,7 +801,7 @@ public partial class ChatClientAgentTests
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
             .Returns(new ValueTask());
 
-        var mockProvider2 = new Mock<AIContextProvider>();
+        var mockProvider2 = new Mock<AIContextProvider>(null, null);
         mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
         mockProvider2
             .Protected()
@@ -869,7 +869,7 @@ public partial class ChatClientAgentTests
             })
             .Returns(ToAsyncEnumerableAsync(responseUpdates));
 
-        var mockProvider1 = new Mock<AIContextProvider>();
+        var mockProvider1 = new Mock<AIContextProvider>(null, null);
         mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
         mockProvider1
             .Protected()
@@ -886,7 +886,7 @@ public partial class ChatClientAgentTests
             .Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
             .Returns(new ValueTask());
 
-        var mockProvider2 = new Mock<AIContextProvider>();
+        var mockProvider2 = new Mock<AIContextProvider>(null, null);
         mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
         mockProvider2
             .Protected()
@@ -1828,7 +1828,7 @@ public partial class ChatClientAgentTests
             })
             .Returns(ToAsyncEnumerableAsync(responseUpdates));
 
-        var mockProvider = new Mock<AIContextProvider>();
+        var mockProvider = new Mock<AIContextProvider>(null, null);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -1907,7 +1907,7 @@ public partial class ChatClientAgentTests
                 It.IsAny<CancellationToken>()))
             .Throws(new InvalidOperationException("downstream failure"));
 
-        var mockProvider = new Mock<AIContextProvider>();
+        var mockProvider = new Mock<AIContextProvider>(null, null);
         mockProvider
             .Protected()
             .Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())

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

@@ -338,7 +338,7 @@ public class ChatClientAgent_BackgroundResponsesTests
         List<ChatMessage> capturedMessages = [];
 
         // Create a mock chat history provider that would normally provide messages
-        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>();
+        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null);
         mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
         mockChatHistoryProvider
             .Protected()
@@ -346,7 +346,7 @@ public class ChatClientAgent_BackgroundResponsesTests
             .ReturnsAsync([new(ChatRole.User, "Message from chat history provider")]);
 
         // Create a mock AI context provider that would normally provide context
-        var mockContextProvider = new Mock<AIContextProvider>();
+        var mockContextProvider = new Mock<AIContextProvider>(null, null);
         mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
         mockContextProvider
             .Protected()
@@ -407,7 +407,7 @@ public class ChatClientAgent_BackgroundResponsesTests
         List<ChatMessage> capturedMessages = [];
 
         // Create a mock chat history provider that would normally provide messages
-        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>();
+        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null);
         mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
         mockChatHistoryProvider
             .Protected()
@@ -415,7 +415,7 @@ public class ChatClientAgent_BackgroundResponsesTests
             .ReturnsAsync([new(ChatRole.User, "Message from chat history provider")]);
 
         // Create a mock AI context provider that would normally provide context
-        var mockContextProvider = new Mock<AIContextProvider>();
+        var mockContextProvider = new Mock<AIContextProvider>(null, null);
         mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
         mockContextProvider
             .Protected()
@@ -638,7 +638,7 @@ public class ChatClientAgent_BackgroundResponsesTests
             .Returns(ToAsyncEnumerableAsync(returnUpdates));
 
         List<ChatMessage> capturedMessagesAddedToProvider = [];
-        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>();
+        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null);
         mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
         mockChatHistoryProvider
             .Protected()
@@ -647,7 +647,7 @@ public class ChatClientAgent_BackgroundResponsesTests
             .Returns(new ValueTask());
 
         AIContextProvider.InvokedContext? capturedInvokedContext = null;
-        var mockContextProvider = new Mock<AIContextProvider>();
+        var mockContextProvider = new Mock<AIContextProvider>(null, null);
         mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
         mockContextProvider
             .Protected()
@@ -702,7 +702,7 @@ public class ChatClientAgent_BackgroundResponsesTests
             .Returns(ToAsyncEnumerableAsync(Array.Empty<ChatResponseUpdate>()));
 
         List<ChatMessage> capturedMessagesAddedToProvider = [];
-        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>();
+        var mockChatHistoryProvider = new Mock<ChatHistoryProvider>(null, null);
         mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
         mockChatHistoryProvider
             .Protected()
@@ -711,7 +711,7 @@ public class ChatClientAgent_BackgroundResponsesTests
             .Returns(new ValueTask());
 
         AIContextProvider.InvokedContext? capturedInvokedContext = null;
-        var mockContextProvider = new Mock<AIContextProvider>();
+        var mockContextProvider = new Mock<AIContextProvider>(null, null);
         mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
         mockContextProvider
             .Protected()

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

@@ -185,7 +185,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
                 It.IsAny<ChatOptions>(),
                 It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
 
-        Mock<ChatHistoryProvider> mockChatHistoryProvider = new();
+        Mock<ChatHistoryProvider> mockChatHistoryProvider = new(null, null);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -240,7 +240,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
                 It.IsAny<ChatOptions>(),
                 It.IsAny<CancellationToken>())).Throws(new InvalidOperationException("Test Error"));
 
-        Mock<ChatHistoryProvider> mockChatHistoryProvider = new();
+        Mock<ChatHistoryProvider> mockChatHistoryProvider = new(null, null);
         mockChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -311,7 +311,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
                 It.IsAny<CancellationToken>())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
 
         // Arrange a chat history provider to override the factory provided one.
-        Mock<ChatHistoryProvider> mockOverrideChatHistoryProvider = new();
+        Mock<ChatHistoryProvider> mockOverrideChatHistoryProvider = new(null, null);
         mockOverrideChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
@@ -324,7 +324,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
 
         // Arrange a chat history provider to provide to the agent at construction time.
         // This one shouldn't be used since it is being overridden.
-        Mock<ChatHistoryProvider> mockAgentOptionsChatHistoryProvider = new();
+        Mock<ChatHistoryProvider> mockAgentOptionsChatHistoryProvider = new(null, null);
         mockAgentOptionsChatHistoryProvider
             .Protected()
             .Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())