diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/AutoApprovedFunctionRemovingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/AutoApprovedFunctionRemovingChatClient.cs new file mode 100644 index 0000000000..142a65d418 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/AutoApprovedFunctionRemovingChatClient.cs @@ -0,0 +1,285 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// A delegating chat client that automatically removes for tools +/// that do not actually require approval, storing auto-approved results in the session for transparent +/// re-injection on the next request. +/// +/// +/// +/// has an all-or-nothing behavior for approvals: when any tool +/// in a response is an , it converts all +/// items to — even for tools that do not require approval. This +/// decorator sits above in the pipeline and transparently handles +/// the non-approval-required items so callers only see approval requests for tools that truly need them. +/// +/// +/// On outbound responses, the decorator identifies items for tools +/// that are not wrapped in , removes them from the response, and +/// stores them in the session's . On the next inbound request, the stored +/// items are re-injected as pre-approved so that +/// can process them alongside the caller's human-approved responses. +/// +/// +/// This decorator requires an active with a non-null +/// . An is thrown if no +/// run context or session is available. +/// +/// +internal sealed class AutoApprovedFunctionRemovingChatClient : DelegatingChatClient +{ + /// + /// The key used in to store pending auto-approved function calls + /// between agent runs. + /// + internal const string StateBagKey = "_autoApprovedFunctionCalls"; + + /// + /// Initializes a new instance of the class. + /// + /// The underlying chat client (typically a ). + public AutoApprovedFunctionRemovingChatClient(IChatClient innerClient) + : base(innerClient) + { + } + + /// + public override async Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + var session = GetRequiredSession(); + var autoApprovableNames = this.GetAutoApprovableToolNames(options); + + messages = InjectPendingAutoApprovals(messages, session, autoApprovableNames); + + var response = await base.GetResponseAsync(messages, options, cancellationToken).ConfigureAwait(false); + + RemoveAutoApprovedFromMessages(response.Messages, autoApprovableNames, session); + + return response; + } + + /// + public override async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + var session = GetRequiredSession(); + var autoApprovableNames = this.GetAutoApprovableToolNames(options); + + messages = InjectPendingAutoApprovals(messages, session, autoApprovableNames); + List? autoApproved = null; + + await foreach (var update in base.GetStreamingResponseAsync(messages, options, cancellationToken).ConfigureAwait(false)) + { + if (FilterUpdateContents(update, autoApprovableNames, ref autoApproved)) + { + yield return update; + } + } + + if (autoApproved is { Count: > 0 }) + { + session.StateBag.SetValue(StateBagKey, autoApproved, AgentJsonUtilities.DefaultOptions); + } + } + + /// + /// Gets the current from the ambient run context. + /// + /// No run context or session is available. + private static AgentSession GetRequiredSession() + { + var runContext = AIAgent.CurrentRunContext + ?? throw new InvalidOperationException( + $"{nameof(AutoApprovedFunctionRemovingChatClient)} can only be used within the context of a running AIAgent. " + + "Ensure that the chat client is being invoked as part of an AIAgent.RunAsync or AIAgent.RunStreamingAsync call."); + + return runContext.Session + ?? throw new InvalidOperationException( + $"{nameof(AutoApprovedFunctionRemovingChatClient)} requires a session. " + + "Ensure the agent has a resolved session before invoking the chat client."); + } + + /// + /// Checks the session for stored auto-approvals from a previous turn and injects them as + /// a user message containing items appended to the input messages. + /// + private static IEnumerable InjectPendingAutoApprovals( + IEnumerable messages, + AgentSession session, + HashSet autoApprovableNames) + { + if (!session.StateBag.TryGetValue>( + StateBagKey, + out var pendingRequests, + AgentJsonUtilities.DefaultOptions) + || pendingRequests is not { Count: > 0 }) + { + return messages; + } + + session.StateBag.TryRemoveValue(StateBagKey); + + List approvalResponses = []; + foreach (var request in pendingRequests) + { + if (IsAutoApprovable(request, autoApprovableNames)) + { + approvalResponses.Add(request.CreateResponse(approved: true)); + } + } + + if (approvalResponses.Count == 0) + { + return messages; + } + + var userMessage = new ChatMessage(ChatRole.User, approvalResponses); + return messages.Concat([userMessage]); + } + + /// + /// Builds a set of tool names that do not require approval and can be auto-approved, + /// by checking all available tools from and + /// . + /// + private HashSet GetAutoApprovableToolNames(ChatOptions? options) + { + var ficc = this.GetService(); + + var allTools = (options?.Tools ?? Enumerable.Empty()) + .Concat(ficc?.AdditionalTools ?? Enumerable.Empty()); + + return new HashSet( + allTools + .OfType() + .Where(static f => f.GetService() is null) + .Select(static f => f.Name), + StringComparer.Ordinal); + } + + /// + /// Determines whether a can be auto-approved because + /// the underlying tool is not an . + /// + /// + /// if the approval request is for a known tool that does not require approval + /// and can be auto-approved; otherwise. + /// + private static bool IsAutoApprovable(ToolApprovalRequestContent approval, HashSet autoApprovableNames) + { + if (approval.ToolCall is not FunctionCallContent fcc) + { + // Non-function tool calls cannot be auto-approved. + return false; + } + + // Auto-approve only if the tool is known and explicitly does NOT require approval. + // Unknown tools are not in the set and are treated as approval-required (safe default). + return autoApprovableNames.Contains(fcc.Name); + } + + /// + /// Scans response messages for auto-approvable items, + /// removes them from the messages, and stores them in the session for the next request. + /// + private static void RemoveAutoApprovedFromMessages( + IList messages, + HashSet autoApprovableNames, + AgentSession session) + { + List? autoApproved = null; + + foreach (var message in messages) + { + for (int i = message.Contents.Count - 1; i >= 0; i--) + { + if (message.Contents[i] is ToolApprovalRequestContent approval + && IsAutoApprovable(approval, autoApprovableNames)) + { + (autoApproved ??= []).Add(approval); + message.Contents.RemoveAt(i); + } + } + } + + // Remove messages that are now empty after filtering. + for (int i = messages.Count - 1; i >= 0; i--) + { + if (messages[i].Contents.Count == 0) + { + messages.RemoveAt(i); + } + } + + if (autoApproved is { Count: > 0 }) + { + session.StateBag.SetValue(StateBagKey, autoApproved, AgentJsonUtilities.DefaultOptions); + } + } + + /// + /// Filters auto-approvable items from a streaming update's + /// contents, collecting them for later storage. + /// + /// + /// if the update should be yielded (has remaining content or had no + /// approval content to begin with); if the update is now empty and + /// should be skipped. + /// + private static bool FilterUpdateContents( + ChatResponseUpdate update, + HashSet autoApprovableNames, + ref List? autoApproved) + { + bool hasApprovalContent = false; + List filteredContents = []; + bool removedAny = false; + + for (int i = 0; i < update.Contents.Count; i++) + { + var content = update.Contents[i]; + + if (content is ToolApprovalRequestContent approval) + { + hasApprovalContent = true; + + if (IsAutoApprovable(approval, autoApprovableNames)) + { + (autoApproved ??= []).Add(approval); + removedAny = true; + } + else + { + filteredContents.Add(content); + } + } + else + { + filteredContents.Add(content); + } + } + + if (removedAny) + { + update.Contents = filteredContents; + } + + // Yield the update unless it was purely auto-approvable approval content (now empty). + return update.Contents.Count > 0 || !hasApprovalContent; + } +} diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs index 8df9112446..5df85beca1 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs @@ -141,6 +141,36 @@ public sealed class ChatClientAgentOptions [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] public bool PersistChatHistoryAtEndOfRun { get; set; } + /// + /// Gets or sets a value indicating whether to store automatically approved function calls in the session state + /// for tools that do not require approval when they are returned alongside tools that do. + /// + /// + /// + /// has an all-or-nothing behavior for approvals: when any tool + /// in a response is an , it converts all + /// items to , even for tools that do not require approval. + /// + /// + /// Setting this property to injects an + /// decorator above in the pipeline. This decorator identifies approval + /// requests for non-approval-required tools, removes them from the response, and stores them in the session. + /// On the next request, the stored items are automatically re-injected as approved, so the caller only needs + /// to handle approval requests for tools that truly require human approval. + /// + /// + /// This option has no effect when is . + /// When using a custom chat client stack, you can add an + /// manually via the + /// extension method. + /// + /// + /// + /// Default is . + /// + [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] + public bool StoreAutoApprovedFunctionCalls { get; set; } + /// /// Creates a new instance of with the same values as this instance. /// @@ -158,5 +188,6 @@ public ChatClientAgentOptions Clone() WarnOnChatHistoryProviderConflict = this.WarnOnChatHistoryProviderConflict, ThrowOnChatHistoryProviderConflict = this.ThrowOnChatHistoryProviderConflict, PersistChatHistoryAtEndOfRun = this.PersistChatHistoryAtEndOfRun, + StoreAutoApprovedFunctionCalls = this.StoreAutoApprovedFunctionCalls, }; } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientBuilderExtensions.cs index a1e8b5f8a5..e451f6bfd2 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientBuilderExtensions.cs @@ -126,4 +126,35 @@ public static ChatClientBuilder UseChatHistoryPersisting(this ChatClientBuilder { return builder.Use(innerClient => new ChatHistoryPersistingChatClient(innerClient, markOnly)); } + + /// + /// Adds an to the chat client pipeline. + /// + /// + /// + /// This decorator should be positioned above the in the pipeline + /// so that it can intercept approval requests for tools that do not require approval. When + /// converts all function calls to approval requests (because at + /// least one tool requires approval), this decorator removes the requests for non-approval-required tools, + /// stores them in the session, and automatically re-injects them as approved on the next request. + /// + /// + /// This extension method is intended for use with custom chat client stacks when + /// is . + /// When is (the default), + /// the automatically injects this decorator when + /// is . + /// + /// + /// This decorator only works within the context of a running with + /// an active session, and will throw an exception if used in any other stack. + /// + /// + /// The to add the decorator to. + /// The for chaining. + [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] + public static ChatClientBuilder UseAutoApprovedFunctionRemoval(this ChatClientBuilder builder) + { + return builder.Use(innerClient => new AutoApprovedFunctionRemovingChatClient(innerClient)); + } } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientExtensions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientExtensions.cs index fffac628a6..3d95b61340 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientExtensions.cs @@ -53,6 +53,17 @@ internal static IChatClient WithDefaultAgentMiddleware(this IChatClient chatClie { var chatBuilder = chatClient.AsBuilder(); + // AutoApprovedFunctionRemovingChatClient is registered before FunctionInvokingChatClient so that + // it sits above FICC in the pipeline. ChatClientBuilder.Build applies factories in reverse order, + // making the first Use() call outermost. By adding this decorator first, the resulting pipeline is: + // AutoApprovedFunctionRemovingChatClient → FunctionInvokingChatClient → ChatHistoryPersistingChatClient → leaf IChatClient + // This allows the decorator to intercept FICC's responses and remove approval requests for tools + // that don't actually require approval, storing them for automatic re-injection on the next request. + if (options?.StoreAutoApprovedFunctionCalls is true) + { + chatBuilder.Use(innerClient => new AutoApprovedFunctionRemovingChatClient(innerClient)); + } + if (chatClient.GetService() is null) { chatBuilder.Use((innerClient, services) => diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/AutoApprovedFunctionRemovingChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/AutoApprovedFunctionRemovingChatClientTests.cs new file mode 100644 index 0000000000..588d7a003a --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/AutoApprovedFunctionRemovingChatClientTests.cs @@ -0,0 +1,501 @@ +// 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; + +public class AutoApprovedFunctionRemovingChatClientTests +{ + #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 AutoApprovedFunctionRemovingChatClient(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 AutoApprovedFunctionRemovingChatClient(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(contents[0]); + Assert.Equal(0, session.StateBag.Count); + } + + [Fact] + public async Task GetResponseAsync_MixedApproval_RemovesNonApprovalItemsAsync() + { + // 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 AutoApprovedFunctionRemovingChatClient(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(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 AutoApprovedFunctionRemovingChatClient(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>( + AutoApprovedFunctionRemovingChatClient.StateBagKey, out var stored, AgentJsonUtilities.DefaultOptions)); + Assert.NotNull(stored); + Assert.Single(stored!); + Assert.Equal("req1", stored![0].RequestId); + } + + [Fact] + public async Task GetResponseAsync_AllNonApproval_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 AutoApprovedFunctionRemovingChatClient(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_NextRequest_InjectsStoredAutoApprovalsAsync() + { + // Arrange + var fccNormal = new FunctionCallContent("call1", "normalTool"); + var storedApproval = new ToolApprovalRequestContent("req1", fccNormal); + + var session = new ChatClientAgentSession(); + session.StateBag.SetValue( + AutoApprovedFunctionRemovingChatClient.StateBagKey, + new List { storedApproval }, + AgentJsonUtilities.DefaultOptions); + + IEnumerable? capturedMessages = null; + var innerClient = CreateMockChatClient((messages, _, _) => + { + capturedMessages = messages.ToList(); + return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Done")])); + }); + + var decorator = new AutoApprovedFunctionRemovingChatClient(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().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( + AutoApprovedFunctionRemovingChatClient.StateBagKey, + new List { storedApproval }, + AgentJsonUtilities.DefaultOptions); + + var innerClient = CreateMockChatClient((_, _, _) => + Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Done")]))); + + var decorator = new AutoApprovedFunctionRemovingChatClient(innerClient); + + // Act + await RunWithAgentContextAsync(decorator, session); + + // Assert — the stored data should be cleared + Assert.False(session.StateBag.TryGetValue>( + AutoApprovedFunctionRemovingChatClient.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 AutoApprovedFunctionRemovingChatClient(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(response.Messages[0].Contents[0]); + Assert.Equal(0, session.StateBag.Count); + } + + #endregion + + #region GetStreamingResponseAsync Tests + + [Fact] + public async Task GetStreamingResponseAsync_NoApprovalContent_PassesThroughUnchangedAsync() + { + // Arrange + var innerClient = CreateMockStreamingChatClient((_, _, _) => + ToAsyncEnumerableAsync( + new ChatResponseUpdate(ChatRole.Assistant, "Hello"))); + + var decorator = new AutoApprovedFunctionRemovingChatClient(innerClient); + var session = new ChatClientAgentSession(); + + // Act + var updates = new List(); + 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_FiltersNonApprovalItemsAsync() + { + // 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 AutoApprovedFunctionRemovingChatClient(innerClient); + var session = new ChatClientAgentSession(); + var options = new ChatOptions { Tools = [normalTool, approvalTool] }; + + // Act + var updates = new List(); + 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().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 AutoApprovedFunctionRemovingChatClient(innerClient); + var session = new ChatClientAgentSession(); + var options = new ChatOptions { Tools = [normalTool, approvalTool] }; + + // Act + var updates = new List(); + await RunStreamingWithAgentContextAsync(decorator, session, updates, options); + + // Assert — the auto-approved item should be stored in the session + Assert.True(session.StateBag.TryGetValue>( + AutoApprovedFunctionRemovingChatClient.StateBagKey, out var stored, AgentJsonUtilities.DefaultOptions)); + Assert.NotNull(stored); + Assert.Single(stored!); + Assert.Equal("req1", stored![0].RequestId); + } + + [Fact] + public async Task GetStreamingResponseAsync_AllNonApproval_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 AutoApprovedFunctionRemovingChatClient(innerClient); + var session = new ChatClientAgentSession(); + var options = new ChatOptions { Tools = [normalTool] }; + + // Act + var updates = new List(); + 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 Error Handling Tests + + [Fact] + public async Task GetResponseAsync_NoRunContext_ThrowsInvalidOperationExceptionAsync() + { + // Arrange + var innerClient = CreateMockChatClient((_, _, _) => + Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "response")]))); + + var decorator = new AutoApprovedFunctionRemovingChatClient(innerClient); + + // Act & Assert — calling directly without agent context + await Assert.ThrowsAsync( + () => decorator.GetResponseAsync([new ChatMessage(ChatRole.User, "test")])); + } + + [Fact] + public async Task GetResponseAsync_NoSession_ThrowsInvalidOperationExceptionAsync() + { + // Arrange + var innerClient = CreateMockChatClient((_, _, _) => + Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "response")]))); + + var decorator = new AutoApprovedFunctionRemovingChatClient(innerClient); + + // Act & Assert — run with null session + await Assert.ThrowsAsync( + () => RunWithAgentContextAsync(decorator, session: null!)); + } + + #endregion + + #region Builder Extension Tests + + [Fact] + public void UseAutoApprovedFunctionRemoving_AddsDecoratorToPipelineAsync() + { + // Arrange + var innerClient = new Mock().Object; + + // Act + var pipeline = innerClient.AsBuilder() + .UseAutoApprovedFunctionRemoval() + .Build(); + + // Assert + Assert.NotNull(pipeline.GetService()); + } + + #endregion + + #region Helpers + + private static async Task RunWithAgentContextAsync( + AutoApprovedFunctionRemovingChatClient 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 RunWithAgentContextAsync( + AutoApprovedFunctionRemovingChatClient decorator, + AgentSession session) + => RunWithAgentContextAsync(decorator, session, options: null); + + private static async Task RunStreamingWithAgentContextAsync( + AutoApprovedFunctionRemovingChatClient decorator, + AgentSession session, + List 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, ChatOptions?, CancellationToken, Task> onGetResponse) + { + var mock = new Mock(); + mock.Setup(c => c.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns((IEnumerable m, ChatOptions? o, CancellationToken ct) => onGetResponse(m, o, ct)); + return mock.Object; + } + + private static IChatClient CreateMockStreamingChatClient( + Func, ChatOptions?, CancellationToken, IAsyncEnumerable> onGetStreamingResponse) + { + var mock = new Mock(); + mock.Setup(c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns((IEnumerable m, ChatOptions? o, CancellationToken ct) => onGetStreamingResponse(m, o, ct)); + return mock.Object; + } + + private static async IAsyncEnumerable ToAsyncEnumerableAsync(params ChatResponseUpdate[] updates) + { + foreach (var update in updates) + { + yield return update; + } + + await Task.CompletedTask; + } + + #endregion +}