chore: import upstream snapshot with attribution

dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs

@@ -0,0 +1,591 @@
+// 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.UnitTests;
+
+/// <summary>
+/// Unit tests for the <see cref="AIContextProviderChatClient"/> class and
+/// the <see cref="AIContextProviderChatClientBuilderExtensions.UseAIContextProviders(ChatClientBuilder, AIContextProvider[])"/> builder extension.
+/// </summary>
+public class AIContextProviderChatClientTests
+{
+    private static readonly AgentSession s_mockSession = new Mock<AgentSession>().Object;
+
+    #region Constructor Tests
+
+    [Fact]
+    public void Constructor_NullInnerClient_ThrowsArgumentNullException()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider("key1");
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => new AIContextProviderChatClient(null!, [provider]));
+    }
+
+    [Fact]
+    public void Constructor_NullProviders_ThrowsArgumentNullException()
+    {
+        // Arrange
+        var innerClient = new Mock<IChatClient>().Object;
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => new AIContextProviderChatClient(innerClient, null!));
+    }
+
+    [Fact]
+    public void Constructor_EmptyProviders_ThrowsArgumentException()
+    {
+        // Arrange
+        var innerClient = new Mock<IChatClient>().Object;
+
+        // Act & Assert
+        Assert.Throws<ArgumentException>(() => new AIContextProviderChatClient(innerClient, []));
+    }
+
+    #endregion
+
+    #region GetResponseAsync Tests
+
+    [Fact]
+    public async Task GetResponseAsync_NoRunContext_ThrowsInvalidOperationExceptionAsync()
+    {
+        // Arrange
+        var innerClient = new Mock<IChatClient>();
+        var provider = new TestAIContextProvider("key1");
+        var chatClient = new AIContextProviderChatClient(innerClient.Object, [provider]);
+
+        // Act & Assert — no AIAgent.CurrentRunContext is set
+        await Assert.ThrowsAsync<InvalidOperationException>(
+            () => chatClient.GetResponseAsync([new ChatMessage(ChatRole.User, "Hello")]));
+    }
+
+    [Fact]
+    public async Task GetResponseAsync_SingleProvider_EnrichesMessagesAsync()
+    {
+        // Arrange
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerClient = CreateMockChatClient(
+            onGetResponse: (messages, _, _) =>
+            {
+                capturedMessages = messages;
+                return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")]));
+            });
+
+        var provider = new TestAIContextProvider("key1", provideMessages: [new ChatMessage(ChatRole.System, "Extra context")]);
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        // Act — run through an agent so CurrentRunContext is set
+        await RunWithAgentContextAsync(chatClient);
+
+        // Assert
+        Assert.NotNull(capturedMessages);
+        var messageList = capturedMessages!.ToList();
+        Assert.Equal(2, messageList.Count);
+        Assert.Equal("Hello", messageList[0].Text);
+        Assert.Contains("Extra context", messageList[1].Text);
+    }
+
+    [Fact]
+    public async Task GetResponseAsync_MultipleProviders_CalledInSequenceAsync()
+    {
+        // Arrange
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerClient = CreateMockChatClient(
+            onGetResponse: (messages, _, _) =>
+            {
+                capturedMessages = messages;
+                return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")]));
+            });
+
+        var provider1 = new TestAIContextProvider("key1", provideMessages: [new ChatMessage(ChatRole.System, "From P1")]);
+        var provider2 = new TestAIContextProvider("key2", provideMessages: [new ChatMessage(ChatRole.System, "From P2")]);
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider1, provider2]);
+
+        // Act
+        await RunWithAgentContextAsync(chatClient);
+
+        // Assert
+        Assert.NotNull(capturedMessages);
+        var messageList = capturedMessages!.ToList();
+        Assert.Equal(3, messageList.Count);
+    }
+
+    [Fact]
+    public async Task GetResponseAsync_Provider_EnrichesToolsAndInstructionsAsync()
+    {
+        // Arrange
+        ChatOptions? capturedOptions = null;
+        var innerClient = CreateMockChatClient(
+            onGetResponse: (_, options, _) =>
+            {
+                capturedOptions = options;
+                return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")]));
+            });
+
+        var provider = new TestAIContextProvider("key1", provideInstructions: "Extra instructions", provideTools: [new TestAITool()]);
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        // Act
+        await RunWithAgentContextAsync(chatClient);
+
+        // Assert
+        Assert.NotNull(capturedOptions);
+        Assert.Equal("Extra instructions", capturedOptions!.Instructions);
+        Assert.Single(capturedOptions.Tools!);
+    }
+
+    [Fact]
+    public async Task GetResponseAsync_OnSuccess_InvokedAsyncCalledAsync()
+    {
+        // Arrange
+        var innerClient = CreateMockChatClient(
+            onGetResponse: (_, _, _) => Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")])));
+
+        var provider = new TestAIContextProvider("key1");
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        // Act
+        await RunWithAgentContextAsync(chatClient);
+
+        // Assert
+        Assert.True(provider.InvokedAsyncCalled);
+        Assert.Null(provider.LastInvokedContext!.InvokeException);
+        Assert.NotNull(provider.LastInvokedContext.ResponseMessages);
+    }
+
+    [Fact]
+    public async Task GetResponseAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync()
+    {
+        // Arrange
+        var expectedException = new InvalidOperationException("Chat failed");
+        var innerClient = CreateMockChatClient(
+            onGetResponse: (_, _, _) => throw expectedException);
+
+        var provider = new TestAIContextProvider("key1");
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        // Act & Assert
+        await Assert.ThrowsAsync<InvalidOperationException>(() => RunWithAgentContextAsync(chatClient));
+
+        Assert.True(provider.InvokedAsyncCalled);
+        Assert.Same(expectedException, provider.LastInvokedContext!.InvokeException);
+    }
+
+    #endregion
+
+    #region GetStreamingResponseAsync Tests
+
+    [Fact]
+    public async Task GetStreamingResponseAsync_SingleProvider_EnrichesAndStreamsAsync()
+    {
+        // Arrange
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerClient = CreateMockStreamingChatClient(
+            onGetStreamingResponse: (messages, _, _) =>
+            {
+                capturedMessages = messages;
+                return ToAsyncEnumerableAsync(
+                    new ChatResponseUpdate(ChatRole.Assistant, "Part1"),
+                    new ChatResponseUpdate(ChatRole.Assistant, "Part2"));
+            });
+
+        var provider = new TestAIContextProvider("key1", provideMessages: [new ChatMessage(ChatRole.System, "Extra context")]);
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        // Act
+        var updates = new List<ChatResponseUpdate>();
+        await RunStreamingWithAgentContextAsync(chatClient, updates);
+
+        // Assert
+        Assert.Equal(2, updates.Count);
+        Assert.NotNull(capturedMessages);
+        Assert.Equal(2, capturedMessages!.ToList().Count);
+    }
+
+    [Fact]
+    public async Task GetStreamingResponseAsync_OnSuccess_InvokedAsyncCalledAsync()
+    {
+        // Arrange
+        var innerClient = CreateMockStreamingChatClient(
+            onGetStreamingResponse: (_, _, _) => ToAsyncEnumerableAsync(
+                new ChatResponseUpdate(ChatRole.Assistant, "Response")));
+
+        var provider = new TestAIContextProvider("key1");
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        // Act
+        await RunStreamingWithAgentContextAsync(chatClient, []);
+
+        // Assert
+        Assert.True(provider.InvokedAsyncCalled);
+        Assert.Null(provider.LastInvokedContext!.InvokeException);
+    }
+
+    [Fact]
+    public async Task GetStreamingResponseAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync()
+    {
+        // Arrange
+        var expectedException = new InvalidOperationException("Stream failed");
+        var innerClient = CreateMockStreamingChatClient(
+            onGetStreamingResponse: (_, _, _) => throw expectedException);
+
+        var provider = new TestAIContextProvider("key1");
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        // Act & Assert
+        await Assert.ThrowsAsync<InvalidOperationException>(
+            () => RunStreamingWithAgentContextAsync(chatClient, []));
+
+        Assert.True(provider.InvokedAsyncCalled);
+        Assert.Same(expectedException, provider.LastInvokedContext!.InvokeException);
+    }
+
+    #endregion
+
+    #region Shared Options Tests
+
+    [Fact]
+    public async Task GetResponseAsync_SharedOptions_ProviderToolsDoNotAccumulateAcrossCallsAsync()
+    {
+        // Arrange: track tool count seen by the inner client on each call
+        var toolCountsSeenByInner = new List<int>();
+
+        var innerClient = CreateMockChatClient(
+            onGetResponse: (_, options, _) =>
+            {
+                toolCountsSeenByInner.Add(options?.Tools?.Count ?? 0);
+                return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")]));
+            });
+
+        var provider = new TestAIContextProvider("key1", provideTools: [new TestAITool()]);
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        var sharedOptions = new ChatOptions
+        {
+            Tools = new List<AITool> { new TestAITool() }
+        };
+
+        // Act: make 3 calls reusing the same ChatOptions
+        for (int i = 0; i < 3; i++)
+        {
+            await RunWithAgentContextAsync(chatClient, sharedOptions);
+        }
+
+        // Assert: each call should see exactly 2 tools (1 baseline + 1 injected)
+        Assert.Equal(3, toolCountsSeenByInner.Count);
+        Assert.All(toolCountsSeenByInner, count => Assert.Equal(2, count));
+    }
+
+    [Fact]
+    public async Task GetResponseAsync_SharedOptions_OriginalToolsNotMutatedAsync()
+    {
+        // Arrange
+        var innerClient = CreateMockChatClient(
+            onGetResponse: (_, _, _) =>
+                Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")])));
+
+        var provider = new TestAIContextProvider("key1", provideTools: [new TestAITool()]);
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        var baselineTool = new TestAITool();
+        var originalTools = new List<AITool> { baselineTool };
+        var sharedOptions = new ChatOptions
+        {
+            Tools = originalTools
+        };
+
+        // Act
+        await RunWithAgentContextAsync(chatClient, sharedOptions);
+
+        // Assert: the original list should still contain only the baseline tool
+        Assert.Single(originalTools);
+        Assert.Same(baselineTool, originalTools[0]);
+        Assert.Same(originalTools, sharedOptions.Tools);
+        Assert.Same(baselineTool, originalTools[0]);
+    }
+
+    [Fact]
+    public async Task GetStreamingResponseAsync_SharedOptions_ProviderToolsDoNotAccumulateAcrossCallsAsync()
+    {
+        // Arrange
+        var toolCountsSeenByInner = new List<int>();
+
+        var innerClient = CreateMockStreamingChatClient(
+            onGetStreamingResponse: (_, options, _) =>
+            {
+                toolCountsSeenByInner.Add(options?.Tools?.Count ?? 0);
+                return ToAsyncEnumerableAsync(
+                    new ChatResponseUpdate(ChatRole.Assistant, "Response"));
+            });
+
+        var provider = new TestAIContextProvider("key1", provideTools: [new TestAITool()]);
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        var sharedOptions = new ChatOptions
+        {
+            Tools = new List<AITool> { new TestAITool() }
+        };
+
+        // Act: make 3 streaming calls reusing the same ChatOptions
+        for (int i = 0; i < 3; i++)
+        {
+            await RunStreamingWithAgentContextAsync(chatClient, [], sharedOptions);
+        }
+
+        // Assert: each call should see exactly 2 tools (1 baseline + 1 injected)
+        Assert.Equal(3, toolCountsSeenByInner.Count);
+        Assert.All(toolCountsSeenByInner, count => Assert.Equal(2, count));
+    }
+
+    [Fact]
+    public async Task GetStreamingResponseAsync_SharedOptions_OriginalToolsNotMutatedAsync()
+    {
+        // Arrange
+        var innerClient = CreateMockStreamingChatClient(
+            onGetStreamingResponse: (_, _, _) => ToAsyncEnumerableAsync(
+                new ChatResponseUpdate(ChatRole.Assistant, "Response")));
+
+        var provider = new TestAIContextProvider("key1", provideTools: [new TestAITool()]);
+        var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
+
+        var baselineTool = new TestAITool();
+        var originalTools = new List<AITool> { baselineTool };
+        var sharedOptions = new ChatOptions
+        {
+            Tools = originalTools
+        };
+
+        // Act
+        await RunStreamingWithAgentContextAsync(chatClient, [], sharedOptions);
+
+        // Assert: the original list should still contain only the baseline tool
+        Assert.Single(originalTools);
+        Assert.Same(baselineTool, originalTools[0]);
+    }
+
+    #endregion
+
+    #region Builder Extension Tests
+
+    [Fact]
+    public void UseExtension_NullBuilder_ThrowsArgumentNullException()
+    {
+        // Arrange
+        var provider = new TestAIContextProvider("key1");
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() =>
+            AIContextProviderChatClientBuilderExtensions.UseAIContextProviders(null!, provider));
+    }
+
+    [Fact]
+    public async Task UseExtension_CreatesWorkingPipelineAsync()
+    {
+        // Arrange
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerClient = CreateMockChatClient(
+            onGetResponse: (messages, _, _) =>
+            {
+                capturedMessages = messages;
+                return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")]));
+            });
+
+        var provider = new TestAIContextProvider("key1", provideMessages: [new ChatMessage(ChatRole.System, "Pipeline context")]);
+
+        var pipeline = new ChatClientBuilder(innerClient)
+            .UseAIContextProviders(provider)
+            .Build();
+
+        // Act — wrap in an agent to set CurrentRunContext
+        var agent = new TestAIAgent
+        {
+            RunAsyncFunc = async (messages, session, options, ct) =>
+            {
+                var response = await pipeline.GetResponseAsync(messages, cancellationToken: ct);
+                return new AgentResponse(response);
+            }
+        };
+
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+
+        // Assert
+        Assert.NotNull(capturedMessages);
+        var messageList = capturedMessages!.ToList();
+        Assert.Equal(2, messageList.Count);
+    }
+
+    #endregion
+
+    #region Helpers
+
+    /// <summary>
+    /// Runs a chat client within an agent context so that AIAgent.CurrentRunContext is set.
+    /// </summary>
+    private static async Task RunWithAgentContextAsync(AIContextProviderChatClient chatClient)
+    {
+        var agent = new TestAIAgent
+        {
+            RunAsyncFunc = async (messages, session, options, ct) =>
+            {
+                var response = await chatClient.GetResponseAsync(messages, cancellationToken: ct);
+                return new AgentResponse(response);
+            }
+        };
+
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+    }
+
+    /// <summary>
+    /// Runs a streaming chat client within an agent context so that AIAgent.CurrentRunContext is set.
+    /// </summary>
+    private static async Task RunStreamingWithAgentContextAsync(AIContextProviderChatClient chatClient, List<ChatResponseUpdate> updates)
+    {
+        var agent = new TestAIAgent
+        {
+            RunAsyncFunc = async (messages, session, options, ct) =>
+            {
+                await foreach (var update in chatClient.GetStreamingResponseAsync(messages, cancellationToken: ct))
+                {
+                    updates.Add(update);
+                }
+
+                return new AgentResponse([new ChatMessage(ChatRole.Assistant, "done")]);
+            }
+        };
+
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+    }
+
+    /// <summary>
+    /// Runs a chat client within an agent context with the specified options.
+    /// </summary>
+    private static async Task RunWithAgentContextAsync(AIContextProviderChatClient chatClient, ChatOptions options)
+    {
+        var agent = new TestAIAgent
+        {
+            RunAsyncFunc = async (messages, session, agentOptions, ct) =>
+            {
+                var response = await chatClient.GetResponseAsync(messages, options, ct);
+                return new AgentResponse(response);
+            }
+        };
+
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+    }
+
+    /// <summary>
+    /// Runs a streaming chat client within an agent context with the specified options.
+    /// </summary>
+    private static async Task RunStreamingWithAgentContextAsync(AIContextProviderChatClient chatClient, List<ChatResponseUpdate> updates, ChatOptions options)
+    {
+        var agent = new TestAIAgent
+        {
+            RunAsyncFunc = async (messages, session, agentOptions, ct) =>
+            {
+                await foreach (var update in chatClient.GetStreamingResponseAsync(messages, options, ct))
+                {
+                    updates.Add(update);
+                }
+
+                return new AgentResponse([new ChatMessage(ChatRole.Assistant, "done")]);
+            }
+        };
+
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+    }
+
+    private static IChatClient CreateMockChatClient(
+        Func<IEnumerable<ChatMessage>, ChatOptions?, CancellationToken, Task<ChatResponse>> onGetResponse)
+    {
+        var mock = new Mock<IChatClient>();
+        mock.Setup(c => c.GetResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions?>(),
+                It.IsAny<CancellationToken>()))
+            .Returns((IEnumerable<ChatMessage> m, ChatOptions? o, CancellationToken ct) => onGetResponse(m, o, ct));
+        return mock.Object;
+    }
+
+    private static IChatClient CreateMockStreamingChatClient(
+        Func<IEnumerable<ChatMessage>, ChatOptions?, CancellationToken, IAsyncEnumerable<ChatResponseUpdate>> onGetStreamingResponse)
+    {
+        var mock = new Mock<IChatClient>();
+        mock.Setup(c => c.GetStreamingResponseAsync(
+                It.IsAny<IEnumerable<ChatMessage>>(),
+                It.IsAny<ChatOptions?>(),
+                It.IsAny<CancellationToken>()))
+            .Returns((IEnumerable<ChatMessage> m, ChatOptions? o, CancellationToken ct) => onGetStreamingResponse(m, o, ct));
+        return mock.Object;
+    }
+
+    private static async IAsyncEnumerable<ChatResponseUpdate> ToAsyncEnumerableAsync(params ChatResponseUpdate[] updates)
+    {
+        foreach (var update in updates)
+        {
+            yield return update;
+        }
+
+        await Task.CompletedTask;
+    }
+
+    /// <summary>
+    /// A test AIContextProvider that provides configurable messages, tools, and instructions.
+    /// </summary>
+    private sealed class TestAIContextProvider : AIContextProvider
+    {
+        private readonly IReadOnlyList<string> _stateKeys;
+        private readonly IEnumerable<ChatMessage> _provideMessages;
+        private readonly string? _provideInstructions;
+        private readonly IEnumerable<AITool>? _provideTools;
+
+        public bool InvokedAsyncCalled { get; private set; }
+
+        public InvokedContext? LastInvokedContext { get; private set; }
+
+        public override IReadOnlyList<string> StateKeys => this._stateKeys;
+
+        public TestAIContextProvider(
+            string stateKey,
+            IEnumerable<ChatMessage>? provideMessages = null,
+            string? provideInstructions = null,
+            IEnumerable<AITool>? provideTools = null)
+        {
+            this._stateKeys = [stateKey];
+            this._provideMessages = provideMessages ?? [];
+            this._provideInstructions = provideInstructions;
+            this._provideTools = provideTools;
+        }
+
+        protected override ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
+        {
+            return new ValueTask<AIContext>(new AIContext
+            {
+                Messages = this._provideMessages,
+                Instructions = this._provideInstructions,
+                Tools = this._provideTools,
+            });
+        }
+
+        protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+        {
+            this.InvokedAsyncCalled = true;
+            this.LastInvokedContext = context;
+            return default;
+        }
+    }
+
+    /// <summary>
+    /// A minimal AITool for testing.
+    /// </summary>
+    private sealed class TestAITool : AITool;
+
+    #endregion
+}

dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs

@@ -0,0 +1,469 @@
+// 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.UnitTests;
+
+/// <summary>
+/// Unit tests for the <see cref="MessageAIContextProviderAgent"/> class and
+/// the <see cref="AIAgentBuilder.UseAIContextProviders(MessageAIContextProvider[])"/> builder extension.
+/// </summary>
+public class MessageAIContextProviderAgentTests
+{
+    private static readonly AgentSession s_mockSession = new Mock<AgentSession>().Object;
+
+    #region Constructor Tests
+
+    [Fact]
+    public void Constructor_NullInnerAgent_ThrowsArgumentNullException()
+    {
+        // Arrange
+        var provider = new TestProvider();
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => new MessageAIContextProviderAgent(null!, [provider]));
+    }
+
+    [Fact]
+    public void Constructor_NullProviders_ThrowsArgumentNullException()
+    {
+        // Arrange
+        var agent = CreateTestAgent();
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => new MessageAIContextProviderAgent(agent, null!));
+    }
+
+    [Fact]
+    public void Constructor_EmptyProviders_ThrowsArgumentOutOfRangeException()
+    {
+        // Arrange
+        var agent = CreateTestAgent();
+
+        // Act & Assert
+        Assert.Throws<ArgumentOutOfRangeException>(() => new MessageAIContextProviderAgent(agent, []));
+    }
+
+    #endregion
+
+    #region RunAsync Tests
+
+    [Fact]
+    public async Task RunAsync_SingleProvider_EnrichesMessagesAndDelegatesToInnerAgentAsync()
+    {
+        // Arrange
+        var contextMessage = new ChatMessage(ChatRole.System, "Extra context");
+        var provider = new TestProvider(provideMessages: [contextMessage]);
+
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerAgent = CreateTestAgent(
+            runFunc: (messages, _, _, _) =>
+            {
+                capturedMessages = messages;
+                return Task.FromResult(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Response")]));
+            });
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider]);
+
+        // Act
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+
+        // Assert - inner agent received enriched messages (input + provider's message)
+        Assert.NotNull(capturedMessages);
+        var messageList = capturedMessages!.ToList();
+        Assert.Equal(2, messageList.Count);
+        Assert.Equal("Hello", messageList[0].Text);
+        Assert.Contains("Extra context", messageList[1].Text);
+    }
+
+    [Fact]
+    public async Task RunAsync_MultipleProviders_CalledInSequenceAsync()
+    {
+        // Arrange
+        var provider1 = new TestProvider(provideMessages: [new ChatMessage(ChatRole.System, "From provider 1")]);
+        var provider2 = new TestProvider(provideMessages: [new ChatMessage(ChatRole.System, "From provider 2")]);
+
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerAgent = CreateTestAgent(
+            runFunc: (messages, _, _, _) =>
+            {
+                capturedMessages = messages;
+                return Task.FromResult(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Response")]));
+            });
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider1, provider2]);
+
+        // Act
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+
+        // Assert - inner agent received messages from both providers in sequence
+        Assert.NotNull(capturedMessages);
+        var messageList = capturedMessages!.ToList();
+        Assert.Equal(3, messageList.Count);
+        Assert.Equal("Hello", messageList[0].Text);
+        Assert.Contains("From provider 1", messageList[1].Text);
+        Assert.Contains("From provider 2", messageList[2].Text);
+    }
+
+    [Fact]
+    public async Task RunAsync_SequentialProviders_EachReceivesPreviousOutputAsync()
+    {
+        // Arrange - provider 2 captures the filtered messages it receives in ProvideMessagesAsync.
+        // The default filter only includes External messages, so provider 1's stamped messages
+        // (marked as AIContextProvider) are filtered out before reaching provider 2's ProvideMessagesAsync.
+        // However, the full unfiltered output from provider 1 is passed to provider 2's InvokingAsync,
+        // and the inner agent receives the full merged output from both providers.
+        IEnumerable<ChatMessage>? provider2ReceivedMessages = null;
+        var provider1 = new TestProvider(provideMessages: [new ChatMessage(ChatRole.System, "From provider 1")]);
+        var provider2 = new TestProvider(
+            provideMessages: [new ChatMessage(ChatRole.System, "From provider 2")],
+            onInvoking: messages => provider2ReceivedMessages = messages.ToList());
+
+        var innerAgent = CreateTestAgent(
+            runFunc: (_, _, _, _) => Task.FromResult(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Response")])));
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider1, provider2]);
+
+        // Act
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+
+        // Assert - provider 2's ProvideMessagesAsync received only External messages (filtered)
+        Assert.NotNull(provider2ReceivedMessages);
+        var received = provider2ReceivedMessages!.ToList();
+        Assert.Single(received);
+        Assert.Equal("Hello", received[0].Text);
+    }
+
+    [Fact]
+    public async Task RunAsync_OnSuccess_InvokedAsyncCalledOnAllProvidersAsync()
+    {
+        // Arrange
+        var provider1 = new TestProvider();
+        var provider2 = new TestProvider();
+        var innerAgent = CreateTestAgent(
+            runFunc: (_, _, _, _) => Task.FromResult(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Response")])));
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider1, provider2]);
+
+        // Act
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+
+        // Assert
+        Assert.True(provider1.InvokedAsyncCalled);
+        Assert.True(provider2.InvokedAsyncCalled);
+        Assert.Null(provider1.LastInvokedContext!.InvokeException);
+        Assert.Null(provider2.LastInvokedContext!.InvokeException);
+    }
+
+    [Fact]
+    public async Task RunAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync()
+    {
+        // Arrange
+        var provider = new TestProvider();
+        var expectedException = new InvalidOperationException("Agent failed");
+        var innerAgent = CreateTestAgent(
+            runFunc: (_, _, _, _) => throw expectedException);
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider]);
+
+        // Act & Assert
+        await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession));
+
+        Assert.True(provider.InvokedAsyncCalled);
+        Assert.Same(expectedException, provider.LastInvokedContext!.InvokeException);
+    }
+
+    [Fact]
+    public async Task RunAsync_OnSuccess_InvokedContextContainsResponseMessagesAsync()
+    {
+        // Arrange
+        var provider = new TestProvider();
+        var responseMessage = new ChatMessage(ChatRole.Assistant, "Response text");
+        var innerAgent = CreateTestAgent(
+            runFunc: (_, _, _, _) => Task.FromResult(new AgentResponse([responseMessage])));
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider]);
+
+        // Act
+        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+
+        // Assert
+        Assert.NotNull(provider.LastInvokedContext?.ResponseMessages);
+        Assert.Contains(provider.LastInvokedContext!.ResponseMessages!, m => m.Text == "Response text");
+    }
+
+    #endregion
+
+    #region RunStreamingAsync Tests
+
+    [Fact]
+    public async Task RunStreamingAsync_SingleProvider_EnrichesMessagesAndStreamsAsync()
+    {
+        // Arrange
+        var contextMessage = new ChatMessage(ChatRole.System, "Extra context");
+        var provider = new TestProvider(provideMessages: [contextMessage]);
+
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerAgent = CreateTestAgent(
+            runStreamingFunc: (messages, _, _, _) =>
+            {
+                capturedMessages = messages;
+                return ToAsyncEnumerableAsync(
+                    new AgentResponseUpdate(ChatRole.Assistant, "Part1"),
+                    new AgentResponseUpdate(ChatRole.Assistant, "Part2"));
+            });
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider]);
+
+        // Act
+        var updates = new List<AgentResponseUpdate>();
+        await foreach (var update in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession))
+        {
+            updates.Add(update);
+        }
+
+        // Assert - streaming updates received
+        Assert.Equal(2, updates.Count);
+        // Assert - inner agent received enriched messages
+        Assert.NotNull(capturedMessages);
+        var messageList = capturedMessages!.ToList();
+        Assert.Equal(2, messageList.Count);
+    }
+
+    [Fact]
+    public async Task RunStreamingAsync_OnSuccess_InvokedAsyncCalledAfterAllUpdatesAsync()
+    {
+        // Arrange
+        var provider = new TestProvider();
+        var innerAgent = CreateTestAgent(
+            runStreamingFunc: (_, _, _, _) => ToAsyncEnumerableAsync(
+                new AgentResponseUpdate(ChatRole.Assistant, "Response")));
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider]);
+
+        // Act - consume all updates
+        await foreach (var _ in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession))
+        {
+        }
+
+        // Assert
+        Assert.True(provider.InvokedAsyncCalled);
+        Assert.Null(provider.LastInvokedContext!.InvokeException);
+    }
+
+    [Fact]
+    public async Task RunStreamingAsync_OnSuccess_InvokedContextContainsAccumulatedResponseAsync()
+    {
+        // Arrange
+        var provider = new TestProvider();
+        var innerAgent = CreateTestAgent(
+            runStreamingFunc: (_, _, _, _) => ToAsyncEnumerableAsync(
+                new AgentResponseUpdate(ChatRole.Assistant, "Hello "),
+                new AgentResponseUpdate(ChatRole.Assistant, "World")));
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider]);
+
+        // Act - consume all updates
+        await foreach (var _ in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession))
+        {
+        }
+
+        // Assert - InvokedAsync received the accumulated response messages
+        Assert.NotNull(provider.LastInvokedContext?.ResponseMessages);
+        var responseMessages = provider.LastInvokedContext!.ResponseMessages!.ToList();
+        Assert.True(responseMessages.Count > 0);
+    }
+
+    [Fact]
+    public async Task RunStreamingAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync()
+    {
+        // Arrange
+        var provider = new TestProvider();
+        var expectedException = new InvalidOperationException("Stream failed");
+        var innerAgent = CreateTestAgent(
+            runStreamingFunc: (_, _, _, _) => throw expectedException);
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider]);
+
+        // Act & Assert
+        await Assert.ThrowsAsync<InvalidOperationException>(async () =>
+        {
+            await foreach (var _ in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession))
+            {
+            }
+        });
+
+        Assert.True(provider.InvokedAsyncCalled);
+        Assert.Same(expectedException, provider.LastInvokedContext!.InvokeException);
+    }
+
+    [Fact]
+    public async Task RunStreamingAsync_MultipleProviders_CalledInSequenceAsync()
+    {
+        // Arrange
+        var provider1 = new TestProvider(provideMessages: [new ChatMessage(ChatRole.System, "From provider 1")]);
+        var provider2 = new TestProvider(provideMessages: [new ChatMessage(ChatRole.System, "From provider 2")]);
+
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerAgent = CreateTestAgent(
+            runStreamingFunc: (messages, _, _, _) =>
+            {
+                capturedMessages = messages;
+                return ToAsyncEnumerableAsync(new AgentResponseUpdate(ChatRole.Assistant, "Response"));
+            });
+
+        var agent = new MessageAIContextProviderAgent(innerAgent, [provider1, provider2]);
+
+        // Act
+        await foreach (var _ in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession))
+        {
+        }
+
+        // Assert
+        Assert.NotNull(capturedMessages);
+        var messageList = capturedMessages!.ToList();
+        Assert.Equal(3, messageList.Count);
+        Assert.Equal("Hello", messageList[0].Text);
+        Assert.Contains("From provider 1", messageList[1].Text);
+        Assert.Contains("From provider 2", messageList[2].Text);
+    }
+
+    #endregion
+
+    #region Builder Extension Tests
+
+    [Fact]
+    public async Task UseExtension_CreatesWorkingPipelineAsync()
+    {
+        // Arrange
+        var contextMessage = new ChatMessage(ChatRole.System, "Pipeline context");
+        var provider = new TestProvider(provideMessages: [contextMessage]);
+
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerAgent = CreateTestAgent(
+            runFunc: (messages, _, _, _) =>
+            {
+                capturedMessages = messages;
+                return Task.FromResult(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Response")]));
+            });
+
+        var pipeline = new AIAgentBuilder(innerAgent)
+            .UseAIContextProviders([provider])
+            .Build();
+
+        // Act
+        await pipeline.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+
+        // Assert
+        Assert.NotNull(capturedMessages);
+        var messageList = capturedMessages!.ToList();
+        Assert.Equal(2, messageList.Count);
+        Assert.Equal("Hello", messageList[0].Text);
+        Assert.Contains("Pipeline context", messageList[1].Text);
+    }
+
+    [Fact]
+    public async Task UseExtension_MultipleProviders_AllAppliedAsync()
+    {
+        // Arrange
+        var provider1 = new TestProvider(provideMessages: [new ChatMessage(ChatRole.System, "P1")]);
+        var provider2 = new TestProvider(provideMessages: [new ChatMessage(ChatRole.System, "P2")]);
+
+        IEnumerable<ChatMessage>? capturedMessages = null;
+        var innerAgent = CreateTestAgent(
+            runFunc: (messages, _, _, _) =>
+            {
+                capturedMessages = messages;
+                return Task.FromResult(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Response")]));
+            });
+
+        var pipeline = new AIAgentBuilder(innerAgent)
+            .UseAIContextProviders([provider1, provider2])
+            .Build();
+
+        // Act
+        await pipeline.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
+
+        // Assert
+        Assert.NotNull(capturedMessages);
+        var messageList = capturedMessages!.ToList();
+        Assert.Equal(3, messageList.Count);
+    }
+
+    #endregion
+
+    #region Helpers
+
+    private static TestAIAgent CreateTestAgent(
+        Func<IEnumerable<ChatMessage>, AgentSession?, AgentRunOptions?, CancellationToken, Task<AgentResponse>>? runFunc = null,
+        Func<IEnumerable<ChatMessage>, AgentSession?, AgentRunOptions?, CancellationToken, IAsyncEnumerable<AgentResponseUpdate>>? runStreamingFunc = null)
+    {
+        var agent = new TestAIAgent();
+        if (runFunc is not null)
+        {
+            agent.RunAsyncFunc = runFunc;
+        }
+
+        if (runStreamingFunc is not null)
+        {
+            agent.RunStreamingAsyncFunc = runStreamingFunc;
+        }
+
+        return agent;
+    }
+
+    private static async IAsyncEnumerable<AgentResponseUpdate> ToAsyncEnumerableAsync(params AgentResponseUpdate[] updates)
+    {
+        foreach (var update in updates)
+        {
+            yield return update;
+        }
+
+        await Task.CompletedTask;
+    }
+
+    /// <summary>
+    /// A test implementation of <see cref="MessageAIContextProvider"/> that records invocation calls.
+    /// </summary>
+    private sealed class TestProvider : MessageAIContextProvider
+    {
+        private readonly IEnumerable<ChatMessage> _provideMessages;
+        private readonly Action<IEnumerable<ChatMessage>>? _onInvoking;
+
+        public bool InvokedAsyncCalled { get; private set; }
+
+        public InvokedContext? LastInvokedContext { get; private set; }
+
+        public TestProvider(
+            IEnumerable<ChatMessage>? provideMessages = null,
+            Action<IEnumerable<ChatMessage>>? onInvoking = null)
+        {
+            this._provideMessages = provideMessages ?? [];
+            this._onInvoking = onInvoking;
+        }
+
+        protected override ValueTask<IEnumerable<ChatMessage>> ProvideMessagesAsync(
+            InvokingContext context,
+            CancellationToken cancellationToken = default)
+        {
+            this._onInvoking?.Invoke(context.RequestMessages);
+            return new ValueTask<IEnumerable<ChatMessage>>(this._provideMessages);
+        }
+
+        protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+        {
+            this.InvokedAsyncCalled = true;
+            this.LastInvokedContext = context;
+            return default;
+        }
+    }
+
+    #endregion
+}