microsoft--agent-framework/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalNotRequiredFunctionBypassingChatClientTests.cs

// 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 Microsoft.Extensions.Logging;
using Moq;

namespace Microsoft.Agents.AI.UnitTests;

public class ApprovalNotRequiredFunctionBypassingChatClientTests
{
    #region GetResponseAsync Tests

    [Fact]
    public async Task GetResponseAsync_NoApprovalContent_PassesThroughUnchangedAsync()
    {
        // Arrange
        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Hello")])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();

        // Act
        var response = await RunWithAgentContextAsync(decorator, session);

        // Assert
        Assert.Single(response.Messages);
        Assert.Equal("Hello", response.Messages[0].Text);
        Assert.Equal(0, session.StateBag.Count);
    }

    [Fact]
    public async Task GetResponseAsync_AllToolsRequireApproval_PassesThroughUnchangedAsync()
    {
        // Arrange
        var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));
        var fcc = new FunctionCallContent("call1", "approvalTool");
        var approval = new ToolApprovalRequestContent("req1", fcc);

        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, [approval])])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();
        var options = new ChatOptions { Tools = [approvalTool] };

        // Act
        var response = await RunWithAgentContextAsync(decorator, session, options);

        // Assert — approval request should remain
        Assert.Single(response.Messages);
        var contents = response.Messages[0].Contents;
        Assert.Single(contents);
        Assert.IsType<ToolApprovalRequestContent>(contents[0]);
        Assert.Equal(0, session.StateBag.Count);
    }

    [Fact]
    public async Task GetResponseAsync_MixedApproval_RemovesApprovalNotRequiredItemsAsync()
    {
        // Arrange
        var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
        var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));

        var fccNormal = new FunctionCallContent("call1", "normalTool");
        var fccApproval = new FunctionCallContent("call2", "approvalTool");
        var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
        var approvalRequired = new ToolApprovalRequestContent("req2", fccApproval);

        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([
                new ChatMessage(ChatRole.Assistant, [approvalNormal, approvalRequired])
            ])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();
        var options = new ChatOptions { Tools = [normalTool, approvalTool] };

        // Act
        var response = await RunWithAgentContextAsync(decorator, session, options);

        // Assert — only the approval-required item remains in the response
        Assert.Single(response.Messages);
        var contents = response.Messages[0].Contents;
        Assert.Single(contents);
        var remainingApproval = Assert.IsType<ToolApprovalRequestContent>(contents[0]);
        Assert.Equal("req2", remainingApproval.RequestId);
    }

    [Fact]
    public async Task GetResponseAsync_MixedApproval_StoresAutoApprovedInSessionAsync()
    {
        // Arrange
        var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
        var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));

        var fccNormal = new FunctionCallContent("call1", "normalTool");
        var fccApproval = new FunctionCallContent("call2", "approvalTool");
        var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
        var approvalRequired = new ToolApprovalRequestContent("req2", fccApproval);

        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([
                new ChatMessage(ChatRole.Assistant, [approvalNormal, approvalRequired])
            ])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();
        var options = new ChatOptions { Tools = [normalTool, approvalTool] };

        // Act
        await RunWithAgentContextAsync(decorator, session, options);

        // Assert — the auto-approved item should be stored in the session
        Assert.True(session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(
            ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey, out var stored, AgentJsonUtilities.DefaultOptions));
        Assert.NotNull(stored);
        Assert.Single(stored!);
        Assert.Equal("req1", stored![0].RequestId);
    }

    [Fact]
    public async Task GetResponseAsync_AllApprovalNotRequired_RemovesAllApprovalsAndRemovesEmptyMessageAsync()
    {
        // Arrange
        var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");

        var fccNormal = new FunctionCallContent("call1", "normalTool");
        var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);

        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([
                new ChatMessage(ChatRole.Assistant, [approvalNormal])
            ])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();
        var options = new ChatOptions { Tools = [normalTool] };

        // Act
        var response = await RunWithAgentContextAsync(decorator, session, options);

        // Assert — the message should be removed since it's now empty
        Assert.Empty(response.Messages);
    }

    [Fact]
    public async Task GetResponseAsync_PreExistingEmptyMessage_IsPreservedAsync()
    {
        // Arrange — the response contains a metadata-only message that was already empty
        // (no content items) before the decorator runs, alongside a normal text message.
        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([
                new ChatMessage(ChatRole.Assistant, contents: []),
                new ChatMessage(ChatRole.Assistant, "Hello")
            ])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();

        // Act
        var response = await RunWithAgentContextAsync(decorator, session);

        // Assert — the decorator must not drop a message it did not empty itself.
        Assert.Equal(2, response.Messages.Count);
        Assert.Empty(response.Messages[0].Contents);
        Assert.Equal("Hello", response.Messages[1].Text);
    }

    [Fact]
    public async Task GetResponseAsync_NextRequest_InjectsStoredAutoApprovalsAsync()
    {
        // Arrange
        var fccNormal = new FunctionCallContent("call1", "normalTool");
        var storedApproval = new ToolApprovalRequestContent("req1", fccNormal);

        var session = new ChatClientAgentSession();
        session.StateBag.SetValue(
            ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey,
            new List<ToolApprovalRequestContent> { storedApproval },
            AgentJsonUtilities.DefaultOptions);

        IEnumerable<ChatMessage>? capturedMessages = null;
        var innerClient = CreateMockChatClient((messages, _, _) =>
        {
            capturedMessages = messages.ToList();
            return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Done")]));
        });

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var options = new ChatOptions { Tools = [AIFunctionFactory.Create(() => "result", "normalTool")] };

        // Act
        await RunWithAgentContextAsync(decorator, session, options);

        // Assert — the inner client should receive injected messages
        Assert.NotNull(capturedMessages);
        var messagesList = capturedMessages!.ToList();

        // Original user message + user message with approved responses.
        Assert.Equal(2, messagesList.Count);
        Assert.Equal(ChatRole.User, messagesList[0].Role);

        // User message with the auto-approved ToolApprovalResponseContent
        Assert.Equal(ChatRole.User, messagesList[1].Role);
        var userContent = messagesList[1].Contents.OfType<ToolApprovalResponseContent>().ToList();
        Assert.Single(userContent);
        Assert.Equal("req1", userContent[0].RequestId);
        Assert.True(userContent[0].Approved);
    }

    [Fact]
    public async Task GetResponseAsync_NextRequest_ClearsStoredAfterInjectionAsync()
    {
        // Arrange
        var fccNormal = new FunctionCallContent("call1", "normalTool");
        var storedApproval = new ToolApprovalRequestContent("req1", fccNormal);

        var session = new ChatClientAgentSession();
        session.StateBag.SetValue(
            ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey,
            new List<ToolApprovalRequestContent> { storedApproval },
            AgentJsonUtilities.DefaultOptions);

        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Done")])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var options = new ChatOptions { Tools = [AIFunctionFactory.Create(() => "result", "normalTool")] };

        // Act
        await RunWithAgentContextAsync(decorator, session, options);

        // Assert — the stored data should be cleared after successful injection
        Assert.False(session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(
            ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey, out _, AgentJsonUtilities.DefaultOptions));
    }

    [Fact]
    public async Task GetResponseAsync_UnknownTool_TreatedAsApprovalRequiredAsync()
    {
        // Arrange — tool is not in ChatOptions.Tools
        var fccUnknown = new FunctionCallContent("call1", "unknownTool");
        var approvalUnknown = new ToolApprovalRequestContent("req1", fccUnknown);

        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([
                new ChatMessage(ChatRole.Assistant, [approvalUnknown])
            ])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();
        var options = new ChatOptions { Tools = [] };

        // Act
        var response = await RunWithAgentContextAsync(decorator, session, options);

        // Assert — unknown tool should NOT be auto-approved
        Assert.Single(response.Messages);
        Assert.Single(response.Messages[0].Contents);
        Assert.IsType<ToolApprovalRequestContent>(response.Messages[0].Contents[0]);
        Assert.Equal(0, session.StateBag.Count);
    }

    [Fact]
    public async Task GetResponseAsync_StoredRequestToolSetChanged_StillInjectsAsApprovedAsync()
    {
        // Arrange — tool was previously non-approval-required but is now wrapped in ApprovalRequiredAIFunction.
        // The LLM still requires a complete set of responses, so we inject unconditionally.
        var fccTool = new FunctionCallContent("call1", "changingTool");
        var storedApproval = new ToolApprovalRequestContent("req1", fccTool);

        var session = new ChatClientAgentSession();
        session.StateBag.SetValue(
            ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey,
            new List<ToolApprovalRequestContent> { storedApproval },
            AgentJsonUtilities.DefaultOptions);

        IEnumerable<ChatMessage>? capturedMessages = null;
        var innerClient = CreateMockChatClient((messages, _, _) =>
        {
            capturedMessages = messages.ToList();
            return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Done")]));
        });

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);

        // The tool is now wrapped in ApprovalRequiredAIFunction — but we still inject unconditionally
        var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "changingTool"));
        var options = new ChatOptions { Tools = [approvalTool] };

        // Act
        await RunWithAgentContextAsync(decorator, session, options);

        // Assert — the stored request should still be injected as approved
        Assert.NotNull(capturedMessages);
        var messagesList = capturedMessages!.ToList();
        Assert.Equal(2, messagesList.Count);
        var userContent = messagesList[1].Contents.OfType<ToolApprovalResponseContent>().ToList();
        Assert.Single(userContent);
        Assert.Equal("req1", userContent[0].RequestId);
        Assert.True(userContent[0].Approved);

        // Session should be cleared
        Assert.False(session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(
            ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey, out _, AgentJsonUtilities.DefaultOptions));
    }

    #endregion

    #region GetStreamingResponseAsync Tests

    [Fact]
    public async Task GetStreamingResponseAsync_NoApprovalContent_PassesThroughUnchangedAsync()
    {
        // Arrange
        var innerClient = CreateMockStreamingChatClient((_, _, _) =>
            ToAsyncEnumerableAsync(
                new ChatResponseUpdate(ChatRole.Assistant, "Hello")));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();

        // Act
        var updates = new List<ChatResponseUpdate>();
        await RunStreamingWithAgentContextAsync(decorator, session, updates);

        // Assert
        Assert.Single(updates);
        Assert.Equal("Hello", updates[0].Text);
        Assert.Equal(0, session.StateBag.Count);
    }

    [Fact]
    public async Task GetStreamingResponseAsync_MixedApproval_FiltersApprovalNotRequiredItemsAsync()
    {
        // Arrange
        var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
        var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));

        var fccNormal = new FunctionCallContent("call1", "normalTool");
        var fccApproval = new FunctionCallContent("call2", "approvalTool");
        var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
        var approvalRequired = new ToolApprovalRequestContent("req2", fccApproval);

        var innerClient = CreateMockStreamingChatClient((_, _, _) =>
            ToAsyncEnumerableAsync(
                new ChatResponseUpdate(ChatRole.Assistant, "text"),
                new ChatResponseUpdate { Contents = [approvalNormal, approvalRequired] }));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();
        var options = new ChatOptions { Tools = [normalTool, approvalTool] };

        // Act
        var updates = new List<ChatResponseUpdate>();
        await RunStreamingWithAgentContextAsync(decorator, session, updates, options);

        // Assert — text update + filtered approval update
        Assert.Equal(2, updates.Count);
        Assert.Equal("text", updates[0].Text);

        // Second update should only have the approval-required item
        var approvalContents = updates[1].Contents.OfType<ToolApprovalRequestContent>().ToList();
        Assert.Single(approvalContents);
        Assert.Equal("req2", approvalContents[0].RequestId);
    }

    [Fact]
    public async Task GetStreamingResponseAsync_MixedApproval_StoresAutoApprovedInSessionAsync()
    {
        // Arrange
        var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
        var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));

        var fccNormal = new FunctionCallContent("call1", "normalTool");
        var fccApproval = new FunctionCallContent("call2", "approvalTool");
        var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
        var approvalRequired = new ToolApprovalRequestContent("req2", fccApproval);

        var innerClient = CreateMockStreamingChatClient((_, _, _) =>
            ToAsyncEnumerableAsync(
                new ChatResponseUpdate { Contents = [approvalNormal, approvalRequired] }));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();
        var options = new ChatOptions { Tools = [normalTool, approvalTool] };

        // Act
        var updates = new List<ChatResponseUpdate>();
        await RunStreamingWithAgentContextAsync(decorator, session, updates, options);

        // Assert — the auto-approved item should be stored in the session
        Assert.True(session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(
            ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey, out var stored, AgentJsonUtilities.DefaultOptions));
        Assert.NotNull(stored);
        Assert.Single(stored!);
        Assert.Equal("req1", stored![0].RequestId);
    }

    [Fact]
    public async Task GetStreamingResponseAsync_AllApprovalNotRequired_SkipsEmptyUpdateAsync()
    {
        // Arrange
        var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");

        var fccNormal = new FunctionCallContent("call1", "normalTool");
        var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);

        var innerClient = CreateMockStreamingChatClient((_, _, _) =>
            ToAsyncEnumerableAsync(
                new ChatResponseUpdate(ChatRole.Assistant, "text"),
                new ChatResponseUpdate { Contents = [approvalNormal] }));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);
        var session = new ChatClientAgentSession();
        var options = new ChatOptions { Tools = [normalTool] };

        // Act
        var updates = new List<ChatResponseUpdate>();
        await RunStreamingWithAgentContextAsync(decorator, session, updates, options);

        // Assert — the approval update should be skipped entirely
        Assert.Single(updates);
        Assert.Equal("text", updates[0].Text);
    }

    #endregion

    #region No-Context Pass-Through Tests

    [Fact]
    public async Task GetResponseAsync_NoRunContext_PassesThroughWithoutBypassingAsync()
    {
        // Arrange
        var fcc = new FunctionCallContent("call1", "normalTool");
        var approval = new ToolApprovalRequestContent("req1", fcc);
        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, [approval])])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);

        // Act — calling directly without agent context; the decorator should no-op and pass through.
        var response = await decorator.GetResponseAsync([new ChatMessage(ChatRole.User, "test")]);

        // Assert — the approval request is surfaced to the caller unchanged (not bypassed).
        var contents = response.Messages.Single().Contents;
        Assert.IsType<ToolApprovalRequestContent>(Assert.Single(contents));
    }

    [Fact]
    public async Task GetResponseAsync_NoSession_PassesThroughWithoutBypassingAsync()
    {
        // Arrange
        var fcc = new FunctionCallContent("call1", "normalTool");
        var approval = new ToolApprovalRequestContent("req1", fcc);
        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, [approval])])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);

        // Act — run with an agent context but a null session; the decorator should no-op and pass through.
        var response = await RunWithAgentContextAsync(decorator, session: null!);

        // Assert — the approval request is surfaced to the caller unchanged (not bypassed).
        var contents = response.Messages.Single().Contents;
        Assert.IsType<ToolApprovalRequestContent>(Assert.Single(contents));
    }

    [Fact]
    public async Task GetStreamingResponseAsync_NoRunContext_PassesThroughWithoutBypassingAsync()
    {
        // Arrange
        var fcc = new FunctionCallContent("call1", "normalTool");
        var approval = new ToolApprovalRequestContent("req1", fcc);
        var innerClient = CreateMockStreamingChatClient((_, _, _) =>
            ToAsyncEnumerableAsync(
                new ChatResponseUpdate(ChatRole.Assistant, [approval])));

        var decorator = new ApprovalNotRequiredFunctionBypassingChatClient(innerClient);

        // Act — calling directly without agent context; the decorator should no-op and pass through.
        var updates = new List<ChatResponseUpdate>();
        await foreach (var update in decorator.GetStreamingResponseAsync([new ChatMessage(ChatRole.User, "test")]))
        {
            updates.Add(update);
        }

        // Assert — the approval request is surfaced to the caller unchanged (not bypassed).
        Assert.Contains(updates.SelectMany(u => u.Contents), c => c is ToolApprovalRequestContent);
    }

    #endregion

    #region Builder Extension Tests

    [Fact]
    public void UseApprovalNotRequiredFunctionBypassing_AddsDecoratorToPipeline()
    {
        // Arrange
        var innerClient = new Mock<IChatClient>().Object;

        // Act
        var pipeline = innerClient.AsBuilder()
            .UseApprovalNotRequiredFunctionBypassing()
            .Build();

        // Assert
        Assert.NotNull(pipeline.GetService<ApprovalNotRequiredFunctionBypassingChatClient>());
    }

    [Fact]
    public async Task UseApprovalNotRequiredFunctionBypassing_ExplicitLoggerFactory_IsUsedForWarningAsync()
    {
        // Arrange
        var loggerMock = new Mock<ILogger>();
        loggerMock.Setup(l => l.IsEnabled(LogLevel.Warning)).Returns(true);
        var loggerFactoryMock = new Mock<ILoggerFactory>();
        loggerFactoryMock.Setup(f => f.CreateLogger(It.IsAny<string>())).Returns(loggerMock.Object);

        var fcc = new FunctionCallContent("call1", "normalTool");
        var approval = new ToolApprovalRequestContent("req1", fcc);
        var innerClient = CreateMockChatClient((_, _, _) =>
            Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, [approval])])));

        var pipeline = innerClient.AsBuilder()
            .UseApprovalNotRequiredFunctionBypassing(loggerFactoryMock.Object)
            .Build();

        // Act — invoked without an agent run context, so the decorator no-ops and logs a warning
        // via the explicitly provided logger factory.
        var response = await pipeline.GetResponseAsync([new ChatMessage(ChatRole.User, "test")]);

        // Assert — the provided factory was used to emit a warning, and the approval request is surfaced.
        loggerFactoryMock.Verify(f => f.CreateLogger(It.IsAny<string>()), Times.Once);
        loggerMock.Verify(
            l => l.Log(
                LogLevel.Warning,
                It.IsAny<EventId>(),
                It.IsAny<It.IsAnyType>(),
                It.IsAny<Exception?>(),
                (Func<It.IsAnyType, Exception?, string>)It.IsAny<object>()),
            Times.Once);
        Assert.IsType<ToolApprovalRequestContent>(Assert.Single(response.Messages.Single().Contents));
    }

    [Fact]
    public void WithDefaultAgentMiddleware_ByDefault_InjectsDecorator()
    {
        // Arrange
        var innerClient = new Mock<IChatClient>().Object;
        var options = new ChatClientAgentOptions();

        // Act
        var pipeline = innerClient.WithDefaultAgentMiddleware(options);

        // Assert
        Assert.NotNull(pipeline.GetService<ApprovalNotRequiredFunctionBypassingChatClient>());
    }

    [Fact]
    public void WithDefaultAgentMiddleware_DisableApprovalNotRequiredFunctionBypassing_DoesNotInjectDecorator()
    {
        // Arrange
        var innerClient = new Mock<IChatClient>().Object;
        var options = new ChatClientAgentOptions { DisableApprovalNotRequiredFunctionBypassing = true };

        // Act
        var pipeline = innerClient.WithDefaultAgentMiddleware(options);

        // Assert
        Assert.Null(pipeline.GetService<ApprovalNotRequiredFunctionBypassingChatClient>());
    }

    #endregion

    #region Helpers

    private static async Task<ChatResponse> RunWithAgentContextAsync(
        ApprovalNotRequiredFunctionBypassingChatClient decorator,
        AgentSession? session,
        ChatOptions? options = null)
    {
        ChatResponse? capturedResponse = null;

        var agent = new TestAIAgent
        {
            RunAsyncFunc = async (messages, agentSession, agentOptions, ct) =>
            {
                capturedResponse = await decorator.GetResponseAsync(messages, options, ct);
                return new AgentResponse(capturedResponse);
            }
        };

        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], session);
        return capturedResponse!;
    }

    private static Task<ChatResponse> RunWithAgentContextAsync(
        ApprovalNotRequiredFunctionBypassingChatClient decorator,
        AgentSession session)
        => RunWithAgentContextAsync(decorator, session, options: null);

    private static async Task RunStreamingWithAgentContextAsync(
        ApprovalNotRequiredFunctionBypassingChatClient decorator,
        AgentSession session,
        List<ChatResponseUpdate> updates,
        ChatOptions? options = null)
    {
        var agent = new TestAIAgent
        {
            RunAsyncFunc = async (messages, agentSession, agentOptions, ct) =>
            {
                await foreach (var update in decorator.GetStreamingResponseAsync(messages, options, ct))
                {
                    updates.Add(update);
                }

                return new AgentResponse([new ChatMessage(ChatRole.Assistant, "done")]);
            }
        };

        await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], session);
    }

    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;
    }

    #endregion
}