// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; using System.Linq; using System.Runtime.CompilerServices; using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; using Moq; using Moq.Protected; namespace Microsoft.Agents.AI.UnitTests; /// /// Unit tests for the class. /// public class ToolApprovalAgentTests { #region Constructor /// /// Verify that constructor throws ArgumentNullException when innerAgent is null. /// [Fact] public void Constructor_NullInnerAgent_ThrowsAsync() { // Act & Assert Assert.Throws("innerAgent", () => new ToolApprovalAgent(null!)); } /// /// Verify that constructor creates a valid instance. /// [Fact] public void Constructor_ValidInnerAgent_CreatesInstanceAsync() { // Arrange var innerAgent = new Mock().Object; // Act var agent = new ToolApprovalAgent(innerAgent); // Assert Assert.NotNull(agent); } /// /// Verify that constructor accepts custom options. /// [Fact] public void Constructor_CustomOptions_CreatesInstance() { // Arrange var innerAgent = new Mock().Object; var options = new ToolApprovalAgentOptions { JsonSerializerOptions = new JsonSerializerOptions() }; // Act var agent = new ToolApprovalAgent(innerAgent, options); // Assert Assert.NotNull(agent); } #endregion #region RunAsync - Passthrough /// /// Verify that when there are no approval requests, response passes through unchanged. /// [Fact] public async Task RunAsync_NoApprovalRequests_PassesThroughAsync() { // Arrange var expectedResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Hello")]); var innerAgent = CreateMockAgent(expectedResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Act var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Hi")], new ChatClientAgentSession()); // Assert Assert.Equal("Hello", response.Text); } /// /// Verify that approval requests with no matching rules are surfaced to the caller. /// [Fact] public async Task RunAsync_ApprovalRequestNoRule_SurfacesToCallerAsync() { // Arrange var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var responseMessage = new ChatMessage(ChatRole.Assistant, [approvalRequest]); var expectedResponse = new AgentResponse([responseMessage]); var innerAgent = CreateMockAgent(expectedResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Act var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Hi")], new ChatClientAgentSession()); // Assert var requests = response.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(requests); Assert.Equal("MyTool", ((FunctionCallContent)requests[0].ToolCall).Name); } #endregion #region RunAsync - Deferred Auto-Approve /// /// Verify that when a tool-level rule exists, matching approval requests are /// auto-approved immediately and the inner agent is re-called in the same run. /// [Fact] public async Task RunAsync_ToolLevelRule_DeferredAutoApproveAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); // Call 1: establish rule + inner agent returns approval request → auto-approved → re-call inner agent var approvalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]); var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]); var callCount = 0; List? secondCallMessages = null; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((msgs, _, _, _) => { callCount++; if (callCount == 2) { secondCallMessages = msgs.ToList(); } }) .ReturnsAsync(() => callCount == 1 ? approvalResponse : finalResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: send always-approve → establishes rule, inner returns TARc → auto-approved → re-calls inner var alwaysApproveResponse = approvalRequest.CreateAlwaysApproveToolResponse("User said always"); var response1 = await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApproveResponse])], session); // Assert — inner agent was called twice within the same RunAsync call Assert.Equal(2, callCount); Assert.Equal("Done", response1.Text); // Response should NOT surface the auto-approved approval request to caller var surfacedRequests = response1.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Empty(surfacedRequests); // Verify that the re-call received the injected auto-approval response Assert.NotNull(secondCallMessages); var injectedApprovals = secondCallMessages!.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(injectedApprovals); Assert.True(injectedApprovals[0].Approved); } /// /// Verify that a tool+arguments rule stores pending auto-approvals for matching calls. /// [Fact] public async Task RunAsync_ToolWithArgsRule_DeferredAutoApproveAsync() { // Arrange var session = new ChatClientAgentSession(); var args = new Dictionary { ["path"] = "test.txt" }; var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadFile", args)); var approvalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]); var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "File content")]); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => { callCount++; return callCount == 1 ? approvalResponse : finalResponse; }); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: set up the rule var alwaysApproveResponse = approvalRequest.CreateAlwaysApproveToolWithArgumentsResponse(); await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApproveResponse])], session); // Call 2: pending auto-approval injected var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Continue")], session); // Assert Assert.Equal("File content", response.Text); } /// /// Verify that a tool+arguments rule does NOT auto-approve when arguments differ. /// [Fact] public async Task RunAsync_ToolWithArgsRule_DoesNotAutoApproveDifferentArgsAsync() { // Arrange var session = new ChatClientAgentSession(); // Set up rule with args { path: "test.txt" } var ruleArgs = new Dictionary { ["path"] = "test.txt" }; var ruleRequest = new ToolApprovalRequestContent("req0", new FunctionCallContent("call0", "ReadFile", ruleArgs)); var alwaysApproveResponse = ruleRequest.CreateAlwaysApproveToolWithArgumentsResponse(); // Then the inner agent returns an approval for DIFFERENT args var differentArgs = new Dictionary { ["path"] = "other.txt" }; var newApprovalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadFile", differentArgs)); var approvalResponseMsg = new AgentResponse([new ChatMessage(ChatRole.Assistant, [newApprovalRequest])]); var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(approvalResponseMsg); var agent = new ToolApprovalAgent(innerAgent.Object); var inputMessages = new List { new(ChatRole.User, [alwaysApproveResponse]), }; // Act var response = await agent.RunAsync(inputMessages, session); // Assert — the approval request should surface to the caller (not auto-approved) var requests = response.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(requests); Assert.Equal("ReadFile", ((FunctionCallContent)requests[0].ToolCall).Name); } /// /// Verify that approving a no-argument call with the "always approve with exact arguments" /// option does NOT auto-approve a later same-tool call that supplies arguments. The later /// call must still surface for explicit approval (security regression for #6486). /// [Fact] public async Task RunAsync_EmptyArgsToolWithArgsRule_DoesNotAutoApproveCallWithArgumentsAsync() { // Arrange var session = new ChatClientAgentSession(); // Approve a no-argument SendPayment call with "always approve with exact arguments". var ruleRequest = new ToolApprovalRequestContent("req0", new FunctionCallContent("call0", "SendPayment")); var alwaysApproveResponse = ruleRequest.CreateAlwaysApproveToolWithArgumentsResponse(); // The inner agent then requests SendPayment WITH sensitive arguments. var sensitiveArgs = new Dictionary { ["recipient"] = "attacker@example.test", ["amount"] = 5000, }; var newApprovalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "SendPayment", sensitiveArgs)); var approvalResponseMsg = new AgentResponse([new ChatMessage(ChatRole.Assistant, [newApprovalRequest])]); var innerAgent = CreateMockAgent(approvalResponseMsg); var agent = new ToolApprovalAgent(innerAgent.Object); var inputMessages = new List { new(ChatRole.User, [alwaysApproveResponse]), }; // Act var response = await agent.RunAsync(inputMessages, session); // Assert — the argument-bearing request must surface, not be auto-approved. var requests = response.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(requests); Assert.Equal("SendPayment", ((FunctionCallContent)requests[0].ToolCall).Name); } /// /// Verify that approving a no-argument call with the "always approve with exact arguments" /// option still auto-approves a later same-tool call that also supplies no arguments, /// preserving the intended narrow behavior. /// [Fact] public async Task RunAsync_EmptyArgsToolWithArgsRule_AutoApprovesLaterEmptyArgsCallAsync() { // Arrange var session = new ChatClientAgentSession(); var ruleRequest = new ToolApprovalRequestContent("req0", new FunctionCallContent("call0", "SendPayment")); var alwaysApproveResponse = ruleRequest.CreateAlwaysApproveToolWithArgumentsResponse(); // Inner agent first re-requests SendPayment with no arguments, then returns a final response. var emptyArgsRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "SendPayment")); var approvalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [emptyArgsRequest])]); var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Payment sent")]); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => { callCount++; return callCount == 1 ? approvalResponse : finalResponse; }); var agent = new ToolApprovalAgent(innerAgent.Object); var inputMessages = new List { new(ChatRole.User, [alwaysApproveResponse]), }; // Act var response = await agent.RunAsync(inputMessages, session); // Assert — the no-argument call is auto-approved and the inner agent reaches its final response. Assert.Equal("Payment sent", response.Text); var requests = response.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Empty(requests); } #endregion #region Mixed Auto-Approve /// /// Verify that when some approval requests match rules and others don't, /// matching ones are stored as pending and non-matching are surfaced. /// [Fact] public async Task RunAsync_MixedApprovalRequests_SurfacesNonMatchingStoreMatchingAsync() { // Arrange var session = new ChatClientAgentSession(); // Set up a rule for ToolA only var ruleRequest = new ToolApprovalRequestContent("rule-req", new FunctionCallContent("rule-call", "ToolA")); var alwaysApprove = ruleRequest.CreateAlwaysApproveToolResponse(); // Inner agent returns approval requests for both ToolA and ToolB var approvalA = new ToolApprovalRequestContent("reqA", new FunctionCallContent("callA", "ToolA")); var approvalB = new ToolApprovalRequestContent("reqB", new FunctionCallContent("callB", "ToolB")); var mixedResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalA, approvalB])]); var innerAgent = CreateMockAgent(mixedResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: establish rule + get mixed approval response var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session); // Assert — ToolB request surfaced to caller, ToolA auto-approved is removed from response var surfacedRequests = response.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(surfacedRequests); Assert.Equal("ToolB", ((FunctionCallContent)surfacedRequests[0].ToolCall).Name); // Call 2: verify pending auto-approval for ToolA is injected List? capturedMessages = null; var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Final")]); var callCount = 0; innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((msgs, _, _, _) => { callCount++; capturedMessages = msgs.ToList(); }) .ReturnsAsync(finalResponse); // User manually approves ToolB and sends along var toolBApproval = approvalB.CreateResponse(approved: true); await agent.RunAsync( [new ChatMessage(ChatRole.User, [toolBApproval])], session); // Verify pending auto-approval for ToolA was injected Assert.NotNull(capturedMessages); var allApprovals = capturedMessages!.SelectMany(m => m.Contents).OfType().ToList(); Assert.Equal(2, allApprovals.Count); // ToolA auto-approval + ToolB manual approval } #endregion #region Content Ordering /// /// Verify that content ordering is preserved when unwrapping AlwaysApproveToolApprovalResponseContent. /// [Fact] public async Task RunAsync_UnwrapPreservesContentOrderAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var alwaysApprove = approvalRequest.CreateAlwaysApproveToolResponse(); // Message with mixed content: text before, always-approve in middle, text after var textBefore = new TextContent("Before"); var textAfter = new TextContent("After"); var inputMessage = new ChatMessage(ChatRole.User, [textBefore, alwaysApprove, textAfter]); List? capturedMessages = null; var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "OK")]); var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((msgs, _, _, _) => capturedMessages = msgs.ToList()) .ReturnsAsync(finalResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Act await agent.RunAsync([inputMessage], session); // Assert — content order should be: TextContent("Before"), ToolApprovalResponseContent, TextContent("After") Assert.NotNull(capturedMessages); var contents = capturedMessages![0].Contents; Assert.Equal(3, contents.Count); Assert.IsType(contents[0]); Assert.Equal("Before", ((TextContent)contents[0]).Text); Assert.IsType(contents[1]); Assert.IsType(contents[2]); Assert.Equal("After", ((TextContent)contents[2]).Text); } #endregion #region Unwrapping /// /// Verify that AlwaysApproveToolApprovalResponseContent is unwrapped to the inner ToolApprovalResponseContent /// before being forwarded to the inner agent. /// [Fact] public async Task RunAsync_UnwrapsAlwaysApproveResponse_ForwardsInnerResponseAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var alwaysApprove = approvalRequest.CreateAlwaysApproveToolResponse(); List? capturedMessages = null; var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "OK")]); var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((msgs, _, _, _) => capturedMessages = msgs.ToList()) .ReturnsAsync(finalResponse); var agent = new ToolApprovalAgent(innerAgent.Object); var inputMessages = new List { new(ChatRole.User, [alwaysApprove]), }; // Act await agent.RunAsync(inputMessages, session); // Assert — the forwarded message should contain ToolApprovalResponseContent, not AlwaysApproveToolApprovalResponseContent Assert.NotNull(capturedMessages); var contents = capturedMessages!.SelectMany(m => m.Contents).ToList(); Assert.DoesNotContain(contents, c => c is AlwaysApproveToolApprovalResponseContent); Assert.Contains(contents, c => c is ToolApprovalResponseContent); } #endregion #region Rule Persistence /// /// Verify that rules persist across multiple RunAsync calls on the same session, /// and that auto-approved TARc are immediately handled via re-call within the same run. /// [Fact] public async Task RunAsync_RulesPersistAcrossCallsAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var alwaysApprove = approvalRequest.CreateAlwaysApproveToolResponse(); // Call 1: no approval requests in response, just establish the rule var firstResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "OK")]); // Call 2a: return an approval request that should match stored rule → auto-approved → re-call var secondApproval = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "MyTool")); var secondResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [secondApproval])]); // Call 2b: re-call after auto-approve, final response var thirdResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done after auto-approve")]); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => { callCount++; return callCount switch { 1 => firstResponse, 2 => secondResponse, _ => thirdResponse, }; }); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: establish rule (no approval requests returned) await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session); // Call 2: inner agent returns TARc → matches rule → auto-approved → re-call → final response var response2 = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Do something")], session); // Assert — inner agent called 3 times total (1 for rule setup, 2 for auto-approve loop) Assert.Equal("Done after auto-approve", response2.Text); Assert.Equal(3, callCount); } /// /// Verify that collected approval responses are cleared after injection (during the loop re-call). /// A subsequent RunAsync call should not inject them again. /// [Fact] public async Task RunAsync_CollectedApprovalResponses_ClearedAfterInjectionAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var alwaysApprove = approvalRequest.CreateAlwaysApproveToolResponse(); // Call 1: establish rule var firstResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "OK")]); // Call 2a: return approval → auto-approved → loop var approval = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "MyTool")); var secondResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [approval])]); // Call 2b: re-call after auto-approve (injected) var thirdResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Injected")]); // Call 3: no pending (already cleared by the loop) var fourthResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Clean")]); var callCount = 0; List? lastCallMessages = null; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((msgs, _, _, _) => { callCount++; lastCallMessages = msgs.ToList(); }) .ReturnsAsync(() => callCount switch { 1 => firstResponse, 2 => secondResponse, 3 => thirdResponse, _ => fourthResponse, }); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: establish rule (callCount → 1) await agent.RunAsync([new ChatMessage(ChatRole.User, [alwaysApprove])], session); // Call 2: inner returns TARc → auto-approved → loop → callCount → 2, 3 await agent.RunAsync([new ChatMessage(ChatRole.User, "Trigger approval")], session); // Call 3: no pending (cleared by the loop), callCount → 4 var response3 = await agent.RunAsync([new ChatMessage(ChatRole.User, "No pending")], session); // Assert — last call should only have the user message, no injected approvals Assert.NotNull(lastCallMessages); var approvals = lastCallMessages!.SelectMany(m => m.Contents).OfType().ToList(); Assert.Empty(approvals); Assert.Equal("Clean", response3.Text); } #endregion #region RunStreamingAsync /// /// Verify that streaming passthrough works when there are no approval requests. /// [Fact] public async Task RunStreamingAsync_NoApprovalRequests_PassesThroughAsync() { // Arrange var updates = new[] { new AgentResponseUpdate(ChatRole.Assistant, "Hello") }; var innerAgent = CreateMockStreamingAgent(updates); var agent = new ToolApprovalAgent(innerAgent.Object); // Act var results = new List(); await foreach (var update in agent.RunStreamingAsync( [new ChatMessage(ChatRole.User, "Hi")], new ChatClientAgentSession())) { results.Add(update); } // Assert Assert.Single(results); Assert.Equal("Hello", results[0].Text); } #endregion #region MatchesRule /// /// Verify that a tool-level rule matches regardless of arguments. /// [Fact] public void MatchesRule_ToolLevelRule_MatchesAnyArgs() { // Arrange var rules = new List { new() { ToolName = "MyTool" }, }; var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool", new Dictionary { ["x"] = "1" })); // Act & Assert Assert.True(ToolApprovalAgent.MatchesRule(request, rules, AgentJsonUtilities.DefaultOptions)); } /// /// Verify that a tool-level rule does not match a different tool name. /// [Fact] public void MatchesRule_ToolLevelRule_DoesNotMatchDifferentTool() { // Arrange var rules = new List { new() { ToolName = "ToolA" }, }; var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ToolB")); // Act & Assert Assert.False(ToolApprovalAgent.MatchesRule(request, rules, AgentJsonUtilities.DefaultOptions)); } /// /// Verify that a tool+arguments rule matches with exact same arguments. /// [Fact] public void MatchesRule_ToolWithArgsRule_MatchesExactArgs() { // Arrange — rule arguments must be in serialized form (JSON string representation) var rules = new List { new() { ToolName = "ReadFile", Arguments = new Dictionary { ["path"] = "\"test.txt\"" }, }, }; var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadFile", new Dictionary { ["path"] = "test.txt" })); // Act & Assert Assert.True(ToolApprovalAgent.MatchesRule(request, rules, AgentJsonUtilities.DefaultOptions)); } /// /// Verify that a tool+arguments rule does not match when argument values differ. /// [Fact] public void MatchesRule_ToolWithArgsRule_DoesNotMatchDifferentValues() { // Arrange var rules = new List { new() { ToolName = "ReadFile", Arguments = new Dictionary { ["path"] = "\"test.txt\"" }, }, }; var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadFile", new Dictionary { ["path"] = "other.txt" })); // Act & Assert Assert.False(ToolApprovalAgent.MatchesRule(request, rules, AgentJsonUtilities.DefaultOptions)); } /// /// Verify that a tool+arguments rule does not match when argument count differs. /// [Fact] public void MatchesRule_ToolWithArgsRule_DoesNotMatchDifferentArgCount() { // Arrange var rules = new List { new() { ToolName = "ReadFile", Arguments = new Dictionary { ["path"] = "\"test.txt\"" }, }, }; var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadFile", new Dictionary { ["path"] = "test.txt", ["encoding"] = "utf-8", })); // Act & Assert Assert.False(ToolApprovalAgent.MatchesRule(request, rules, AgentJsonUtilities.DefaultOptions)); } /// /// Verify that a non-FunctionCallContent tool call does not match any rule. /// [Fact] public void MatchesRule_NonFunctionCallContent_ReturnsFalse() { // Arrange var rules = new List { new() { ToolName = "MyTool" }, }; var request = new ToolApprovalRequestContent("req1", new ToolCallContent("call1")); // Act & Assert Assert.False(ToolApprovalAgent.MatchesRule(request, rules, AgentJsonUtilities.DefaultOptions)); } /// /// Verify that matching works with JsonElement argument values. /// [Fact] public void MatchesRule_JsonElementArgs_MatchesCorrectly() { // Arrange var rules = new List { new() { ToolName = "MyTool", Arguments = new Dictionary { ["count"] = "42" }, }, }; var jsonElement = JsonDocument.Parse("42").RootElement; var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool", new Dictionary { ["count"] = jsonElement })); // Act & Assert Assert.True(ToolApprovalAgent.MatchesRule(request, rules, AgentJsonUtilities.DefaultOptions)); } /// /// Verify that an empty-arguments rule (non-null but empty) matches only a call that /// supplies no arguments, and never widens into a tool-level match. /// [Fact] public void MatchesRule_EmptyArgumentsRule_MatchesEmptyArgumentCall_ReturnsTrue() { // Arrange — an exact-arguments approval of a no-argument call is stored as an empty dictionary. var rules = new List { new() { ToolName = "SendPayment", Arguments = new Dictionary(), }, }; var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "SendPayment")); // Act & Assert Assert.True(ToolApprovalAgent.MatchesRule(request, rules, AgentJsonUtilities.DefaultOptions)); } /// /// Verify that an empty-arguments rule does NOT match a later same-tool call that supplies /// arguments. This guards against an exact-arguments approval of a no-argument call being /// widened into a tool-level approval. /// [Fact] public void MatchesRule_EmptyArgumentsRule_DoesNotMatchCallWithArguments_ReturnsFalse() { // Arrange var rules = new List { new() { ToolName = "SendPayment", Arguments = new Dictionary(), }, }; var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "SendPayment", new Dictionary { ["recipient"] = "attacker@example.test", ["amount"] = 5000, })); // Act & Assert Assert.False(ToolApprovalAgent.MatchesRule(request, rules, AgentJsonUtilities.DefaultOptions)); } #endregion #region Extension Methods /// /// Verify that CreateAlwaysApproveToolResponse creates the correct wrapper. /// [Fact] public void CreateAlwaysApproveToolResponse_SetsCorrectFlags() { // Arrange var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); // Act var result = request.CreateAlwaysApproveToolResponse("Test reason"); // Assert Assert.True(result.AlwaysApproveTool); Assert.False(result.AlwaysApproveToolWithArguments); Assert.True(result.InnerResponse.Approved); Assert.Equal("Test reason", result.InnerResponse.Reason); Assert.Equal("req1", result.InnerResponse.RequestId); } /// /// Verify that CreateAlwaysApproveToolWithArgumentsResponse creates the correct wrapper. /// [Fact] public void CreateAlwaysApproveToolWithArgumentsResponse_SetsCorrectFlags() { // Arrange var request = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); // Act var result = request.CreateAlwaysApproveToolWithArgumentsResponse(); // Assert Assert.False(result.AlwaysApproveTool); Assert.True(result.AlwaysApproveToolWithArguments); Assert.True(result.InnerResponse.Approved); } /// /// Verify that extension methods throw on null request. /// [Fact] public void CreateAlwaysApproveToolResponse_NullRequest_Throws() { // Act & Assert Assert.Throws("request", () => ((ToolApprovalRequestContent)null!).CreateAlwaysApproveToolResponse()); } /// /// Verify that extension methods throw on null request. /// [Fact] public void CreateAlwaysApproveToolWithArgumentsResponse_NullRequest_Throws() { // Act & Assert Assert.Throws("request", () => ((ToolApprovalRequestContent)null!).CreateAlwaysApproveToolWithArgumentsResponse()); } #endregion #region Duplicate Rule Prevention /// /// Verify that sending the same always-approve response twice does not create duplicate rules. /// [Fact] public async Task RunAsync_DuplicateAlwaysApprove_DoesNotDuplicateRuleAsync() { // Arrange var session = new ChatClientAgentSession(); var request1 = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var request2 = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "MyTool")); var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "OK")]); var innerAgent = CreateMockAgent(finalResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Act — send two always-approve responses for the same tool await agent.RunAsync( [new ChatMessage(ChatRole.User, [request1.CreateAlwaysApproveToolResponse()])], session); await agent.RunAsync( [new ChatMessage(ChatRole.User, [request2.CreateAlwaysApproveToolResponse()])], session); // Assert — verify the state works correctly (rule still matches on subsequent call) var thirdApproval = new ToolApprovalRequestContent("req3", new FunctionCallContent("call3", "MyTool")); var approvalResponseMsg = new AgentResponse([new ChatMessage(ChatRole.Assistant, [thirdApproval])]); var afterAutoResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Auto-approved")]); var callCount = 0; innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => { callCount++; return callCount == 1 ? approvalResponseMsg : afterAutoResponse; }); // Call 3: triggers approval → stored as pending await agent.RunAsync([new ChatMessage(ChatRole.User, "test")], session); // Call 4: pending injected var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "continue")], session); Assert.Equal("Auto-approved", response.Text); } #endregion #region Auto-Approved Removal /// /// Verify that auto-approved requests are removed from the non-streaming response /// and the inner agent is re-called, returning actual content. /// [Fact] public async Task RunAsync_AutoApprovedRequest_RemovedFromResponseAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var alwaysApprove = approvalRequest.CreateAlwaysApproveToolResponse(); // Inner agent returns a TARc on first call, then real content on second var approvalResponseContent = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "MyTool")); var firstResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalResponseContent])]); var secondResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Tool executed successfully")]); var innerAgent = new Mock(); var callCount = 0; innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => ++callCount == 1 ? firstResponse : secondResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Act: establish rule + inner returns matching approval request → auto-approved → re-call var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session); // Assert — auto-approved request removed, inner agent called twice, final response has text var allRequests = response.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Empty(allRequests); Assert.Equal("Tool executed successfully", response.Messages[0].Text); Assert.Equal(2, callCount); } /// /// Verify that auto-approved requests are removed from streaming updates /// and the inner agent is re-called, yielding actual content. /// [Fact] public async Task RunStreamingAsync_AutoApprovedRequest_RemovedFromUpdatesAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); // First establish a rule via non-streaming var alwaysApprove = approvalRequest.CreateAlwaysApproveToolResponse(); var setupResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "OK")]); var innerAgent = new Mock(); // Setup non-streaming for rule establishment innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(setupResponse); // First streaming call: text + auto-approvable TARc. Second call: just text (tool executed). var streamApproval = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "MyTool")); var textUpdate = new AgentResponseUpdate(ChatRole.Assistant, "Hello"); var approvalUpdate = new AgentResponseUpdate(ChatRole.Assistant, new List { streamApproval }); var finalUpdate = new AgentResponseUpdate(ChatRole.Assistant, "Done"); var streamCallCount = 0; innerAgent .Protected() .Setup>("RunCoreStreamingAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Returns(() => ++streamCallCount == 1 ? ToAsyncEnumerableAsync(new[] { textUpdate, approvalUpdate }) : ToAsyncEnumerableAsync(new[] { finalUpdate })); var agent = new ToolApprovalAgent(innerAgent.Object); // Establish the rule await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session); // Act — stream with auto-approvable request var results = new List(); await foreach (var update in agent.RunStreamingAsync( [new ChatMessage(ChatRole.User, "Do something")], session)) { results.Add(update); } // Assert — text from first call + final text from re-call; no TARc surfaced Assert.Equal(2, results.Count); Assert.Equal("Hello", results[0].Text); Assert.Equal("Done", results[1].Text); Assert.Equal(2, streamCallCount); } /// /// Verify that non-matching approval requests remain in streaming updates. /// [Fact] public async Task RunStreamingAsync_NonMatchingRequest_SurfacedToCallerAsync() { // Arrange var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "UnknownTool")); var approvalUpdate = new AgentResponseUpdate(ChatRole.Assistant, new List { approvalRequest }); var innerAgent = CreateMockStreamingAgent([approvalUpdate]); var agent = new ToolApprovalAgent(innerAgent.Object); // Act var results = new List(); await foreach (var update in agent.RunStreamingAsync( [new ChatMessage(ChatRole.User, "Hi")], new ChatClientAgentSession())) { results.Add(update); } // Assert — non-matching request should be surfaced Assert.Single(results); var requests = results[0].Contents.OfType().ToList(); Assert.Single(requests); Assert.Equal("UnknownTool", ((FunctionCallContent)requests[0].ToolCall).Name); } /// /// Verify that a message with only auto-approved content is removed from the response /// when the inner agent is re-called after auto-approving. /// [Fact] public async Task RunAsync_MessageWithOnlyAutoApprovedContent_RemovedEntirelyAsync() { // Arrange var session = new ChatClientAgentSession(); var ruleRequest = new ToolApprovalRequestContent("rule-req", new FunctionCallContent("rule-call", "MyTool")); var alwaysApprove = ruleRequest.CreateAlwaysApproveToolResponse(); // First call: returns TARc-only message + text message. Second call: returns just text (tools executed). var approvalContent = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var approvalMessage = new ChatMessage(ChatRole.Assistant, [approvalContent]); var textMessage = new ChatMessage(ChatRole.Assistant, "Final text"); var firstResponse = new AgentResponse([approvalMessage, textMessage]); var secondResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Tools executed")]); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => ++callCount == 1 ? firstResponse : secondResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Act var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session); // Assert — inner agent re-called after auto-approve, final response from second call Assert.Equal(2, callCount); Assert.Single(response.Messages); Assert.Equal("Tools executed", response.Messages[0].Text); } /// /// Verify that a streaming update with mixed content (auto-approved + non-auto-approved) only removes the auto-approved part. /// [Fact] public async Task RunStreamingAsync_MixedUpdate_OnlyRemovesAutoApprovedAsync() { // Arrange var session = new ChatClientAgentSession(); var ruleRequest = new ToolApprovalRequestContent("rule-req", new FunctionCallContent("rule-call", "ToolA")); var alwaysApprove = ruleRequest.CreateAlwaysApproveToolResponse(); // Setup non-streaming for rule establishment var setupResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "OK")]); var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(setupResponse); // Streaming update with both auto-approvable (ToolA) and non-auto-approvable (ToolB) requests var autoApprovable = new ToolApprovalRequestContent("reqA", new FunctionCallContent("callA", "ToolA")); var notAutoApprovable = new ToolApprovalRequestContent("reqB", new FunctionCallContent("callB", "ToolB")); var mixedUpdate = new AgentResponseUpdate(ChatRole.Assistant, new List { autoApprovable, notAutoApprovable }); innerAgent .Protected() .Setup>("RunCoreStreamingAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Returns(ToAsyncEnumerableAsync(new[] { mixedUpdate })); var agent = new ToolApprovalAgent(innerAgent.Object); // Establish rule await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session); // Act var results = new List(); await foreach (var update in agent.RunStreamingAsync( [new ChatMessage(ChatRole.User, "Do something")], session)) { results.Add(update); } // Assert — only ToolB request should remain Assert.Single(results); var requests = results[0].Contents.OfType().ToList(); Assert.Single(requests); Assert.Equal("ToolB", ((FunctionCallContent)requests[0].ToolCall).Name); } #endregion #region Queue Behavior - Non-Streaming /// /// Verify that when the inner agent returns multiple unapproved TARc, only the first is returned /// and the rest are queued. /// [Fact] public async Task RunAsync_MultipleTARc_ReturnsOnlyFirstAsync() { // Arrange var tarc1 = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ToolA")); var tarc2 = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "ToolB")); var tarc3 = new ToolApprovalRequestContent("req3", new FunctionCallContent("call3", "ToolC")); var innerResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [tarc1, tarc2, tarc3])]); var innerAgent = CreateMockAgent(innerResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Act var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Do something")], new ChatClientAgentSession()); // Assert — only the first TARc should be returned var requests = response.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(requests); Assert.Equal("ToolA", ((FunctionCallContent)requests[0].ToolCall).Name); } /// /// Verify that after approving the first queued item, the second is returned on the next call /// without calling the inner agent. /// [Fact] public async Task RunAsync_ApproveFirst_ReturnsSecondFromQueueAsync() { // Arrange var session = new ChatClientAgentSession(); var tarc1 = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ToolA")); var tarc2 = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "ToolB")); var tarc3 = new ToolApprovalRequestContent("req3", new FunctionCallContent("call3", "ToolC")); var innerResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [tarc1, tarc2, tarc3])]); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((_, _, _, _) => callCount++) .ReturnsAsync(innerResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: get first TARc var response1 = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Do something")], session); // Call 2: approve first, get second from queue var approval1 = tarc1.CreateResponse(approved: true, reason: "OK"); var response2 = await agent.RunAsync( [new ChatMessage(ChatRole.User, [approval1])], session); // Assert — second TARc returned, inner agent called only once Assert.Equal(1, callCount); var requests = response2.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(requests); Assert.Equal("ToolB", ((FunctionCallContent)requests[0].ToolCall).Name); } /// /// Verify that "always approve" on the first queued item auto-approves matching items in the queue. /// [Fact] public async Task RunAsync_AlwaysApproveFirst_AutoApprovesMatchingInQueueAsync() { // Arrange var session = new ChatClientAgentSession(); // Three TARc: two for ToolA (different args), one for ToolB var tarc1 = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ToolA", arguments: new Dictionary { ["q"] = "query1" })); var tarc2 = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "ToolA", arguments: new Dictionary { ["q"] = "query2" })); var tarc3 = new ToolApprovalRequestContent("req3", new FunctionCallContent("call3", "ToolB")); var innerResponse1 = new AgentResponse([new ChatMessage(ChatRole.Assistant, [tarc1, tarc2, tarc3])]); var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((_, _, _, _) => callCount++) .ReturnsAsync(() => callCount == 1 ? innerResponse1 : finalResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: get first TARc (ToolA query1) var response1 = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Do something")], session); Assert.Single(response1.Messages.SelectMany(m => m.Contents).OfType()); // Call 2: always approve ToolA → should auto-approve ToolA query2 in queue, return ToolB var alwaysApprove = tarc1.CreateAlwaysApproveToolResponse("Always approve ToolA"); var response2 = await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session); var requests2 = response2.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(requests2); Assert.Equal("ToolB", ((FunctionCallContent)requests2[0].ToolCall).Name); } /// /// Verify that once all queued items are resolved, the inner agent is called with all collected responses. /// [Fact] public async Task RunAsync_AllQueuedResolved_CallsInnerAgentWithCollectedResponsesAsync() { // Arrange var session = new ChatClientAgentSession(); var tarc1 = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ToolA")); var tarc2 = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "ToolB")); var innerResponse1 = new AgentResponse([new ChatMessage(ChatRole.Assistant, [tarc1, tarc2])]); var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "All tools executed")]); var callCount = 0; List? thirdCallMessages = null; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((msgs, _, _, _) => { callCount++; if (callCount == 2) { thirdCallMessages = msgs.ToList(); } }) .ReturnsAsync(() => callCount == 1 ? innerResponse1 : finalResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: get first TARc await agent.RunAsync( [new ChatMessage(ChatRole.User, "Do something")], session); // Call 2: approve first, get second from queue var approval1 = tarc1.CreateResponse(approved: true, reason: "OK"); await agent.RunAsync( [new ChatMessage(ChatRole.User, [approval1])], session); // Call 3: approve second → queue empty → inner agent called var approval2 = tarc2.CreateResponse(approved: true, reason: "OK"); var response3 = await agent.RunAsync( [new ChatMessage(ChatRole.User, [approval2])], session); // Assert — inner agent called twice (initial + after queue resolved) Assert.Equal(2, callCount); Assert.Equal("All tools executed", response3.Text); // Verify collected responses were injected Assert.NotNull(thirdCallMessages); var injectedResponses = thirdCallMessages!.SelectMany(m => m.Contents).OfType().ToList(); Assert.Equal(2, injectedResponses.Count); } /// /// Verify that a single TARc (no excess) is returned without queueing. /// [Fact] public async Task RunAsync_SingleTARc_NoQueueingAsync() { // Arrange var tarc = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var innerResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [tarc])]); var innerAgent = CreateMockAgent(innerResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Act var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Hi")], new ChatClientAgentSession()); // Assert — single TARc returned directly var requests = response.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(requests); Assert.Equal("MyTool", ((FunctionCallContent)requests[0].ToolCall).Name); } /// /// Verify that denying a queued item collects the denial and proceeds to the next. /// [Fact] public async Task RunAsync_DenyFirst_ReturnsSecondFromQueueAsync() { // Arrange var session = new ChatClientAgentSession(); var tarc1 = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ToolA")); var tarc2 = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "ToolB")); var innerResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [tarc1, tarc2])]); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((_, _, _, _) => callCount++) .ReturnsAsync(innerResponse); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: get first TARc await agent.RunAsync( [new ChatMessage(ChatRole.User, "Do something")], session); // Call 2: deny first, get second from queue var denial = tarc1.CreateResponse(approved: false, reason: "No"); var response2 = await agent.RunAsync( [new ChatMessage(ChatRole.User, [denial])], session); // Assert — second TARc returned, inner agent called only once Assert.Equal(1, callCount); var requests = response2.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(requests); Assert.Equal("ToolB", ((FunctionCallContent)requests[0].ToolCall).Name); } #endregion #region Queue Behavior - Streaming /// /// Verify that when streaming yields multiple unapproved TARc, only the first is yielded. /// [Fact] public async Task RunStreamingAsync_MultipleTARc_YieldsOnlyFirstAsync() { // Arrange var tarc1 = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ToolA")); var tarc2 = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "ToolB")); var approvalUpdate = new AgentResponseUpdate(ChatRole.Assistant, new List { tarc1, tarc2 }); var textUpdate = new AgentResponseUpdate(ChatRole.Assistant, "Searching..."); var innerAgent = CreateMockStreamingAgent([textUpdate, approvalUpdate]); var agent = new ToolApprovalAgent(innerAgent.Object); // Act var results = new List(); await foreach (var update in agent.RunStreamingAsync( [new ChatMessage(ChatRole.User, "Search")], new ChatClientAgentSession())) { results.Add(update); } // Assert — text update + one TARc Assert.Equal(2, results.Count); Assert.Equal("Searching...", results[0].Text); var requests = results[1].Contents.OfType().ToList(); Assert.Single(requests); Assert.Equal("ToolA", ((FunctionCallContent)requests[0].ToolCall).Name); } /// /// Verify that streaming returns queued items one at a time on subsequent calls. /// [Fact] public async Task RunStreamingAsync_ApproveFirst_ReturnsSecondFromQueueAsync() { // Arrange var session = new ChatClientAgentSession(); var tarc1 = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ToolA")); var tarc2 = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "ToolB")); var approvalUpdate = new AgentResponseUpdate(ChatRole.Assistant, new List { tarc1, tarc2 }); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreStreamingAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((_, _, _, _) => callCount++) .Returns(ToAsyncEnumerableAsync(new[] { approvalUpdate })); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: get first TARc var results1 = new List(); await foreach (var update in agent.RunStreamingAsync( [new ChatMessage(ChatRole.User, "Search")], session)) { results1.Add(update); } Assert.Single(results1); Assert.Equal("ToolA", ((FunctionCallContent)results1[0].Contents.OfType().Single().ToolCall).Name); // Call 2: approve first, get second from queue var approval = tarc1.CreateResponse(approved: true, reason: "OK"); var results2 = new List(); await foreach (var update in agent.RunStreamingAsync( [new ChatMessage(ChatRole.User, [approval])], session)) { results2.Add(update); } // Assert — second TARc returned from queue, inner agent called only once Assert.Equal(1, callCount); Assert.Single(results2); Assert.Equal("ToolB", ((FunctionCallContent)results2[0].Contents.OfType().Single().ToolCall).Name); } /// /// Verify that "always approve" during streaming queue auto-approves matching items. /// [Fact] public async Task RunStreamingAsync_AlwaysApproveFirst_AutoApprovesMatchingInQueueAsync() { // Arrange var session = new ChatClientAgentSession(); // Two ToolA calls + one ToolB call var tarc1 = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ToolA")); var tarc2 = new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "ToolA")); var tarc3 = new ToolApprovalRequestContent("req3", new FunctionCallContent("call3", "ToolB")); var approvalUpdate = new AgentResponseUpdate(ChatRole.Assistant, new List { tarc1, tarc2, tarc3 }); var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreStreamingAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Returns(ToAsyncEnumerableAsync(new[] { approvalUpdate })); var agent = new ToolApprovalAgent(innerAgent.Object); // Call 1: get first TARc (ToolA) var results1 = new List(); await foreach (var update in agent.RunStreamingAsync( [new ChatMessage(ChatRole.User, "Search")], session)) { results1.Add(update); } // Call 2: always approve ToolA → second ToolA auto-approved, returns ToolB var alwaysApprove = tarc1.CreateAlwaysApproveToolResponse(); var results2 = new List(); await foreach (var update in agent.RunStreamingAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session)) { results2.Add(update); } // Assert — ToolB returned (ToolA[2] was auto-approved) Assert.Single(results2); var requests = results2[0].Contents.OfType().ToList(); Assert.Single(requests); Assert.Equal("ToolB", ((FunctionCallContent)requests[0].ToolCall).Name); } #endregion #region Helpers private static Mock CreateMockAgent(AgentResponse response) { var mock = new Mock(); mock .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(response); return mock; } private static Mock CreateMockStreamingAgent(AgentResponseUpdate[] updates) { var mock = new Mock(); mock .Protected() .Setup>("RunCoreStreamingAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Returns(ToAsyncEnumerableAsync(updates)); return mock; } private static async IAsyncEnumerable ToAsyncEnumerableAsync( IEnumerable items, [EnumeratorCancellation] CancellationToken cancellationToken = default) { foreach (var item in items) { cancellationToken.ThrowIfCancellationRequested(); yield return item; await Task.Yield(); } } #endregion #region Auto-Approval Rules (Heuristics) /// /// Verify that approves any tool call regardless of name. /// [Theory] [InlineData("ReadTool")] [InlineData("DangerousTool")] [InlineData("file_access_delete")] public async Task AllToolsAutoApprovalRule_ApprovesAnyToolAsync(string toolName) { // Arrange var functionCall = new FunctionCallContent("call1", toolName); // Act bool approved = await ToolApprovalAgent.AllToolsAutoApprovalRule(functionCall); // Assert Assert.True(approved); } /// /// Verify that wiring auto-approves an approval /// request rather than surfacing it to the caller. /// [Fact] public async Task RunAsync_AllToolsAutoApprovalRule_ApprovesAnyToolAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "DangerousTool")); // Inner agent: first call returns approval request, second returns final response. var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => { callCount++; if (callCount == 1) { return new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]); } return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]); }); var options = new ToolApprovalAgentOptions { AutoApprovalRules = [ToolApprovalAgent.AllToolsAutoApprovalRule] }; var agent = new ToolApprovalAgent(innerAgent.Object, options); // Act var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Hi")], session); // Assert — the approval request was auto-approved, inner agent called twice, nothing surfaced to caller. Assert.Equal(2, callCount); Assert.Equal("Done", response.Text); Assert.Empty(response.Messages.SelectMany(m => m.Contents).OfType()); } /// /// Verify that an auto-approval rule can approve a function call that would otherwise need user approval. /// [Fact] public async Task RunAsync_AutoApprovalRule_ApprovesMatchingToolAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadTool")); // Inner agent: first call returns approval request, second returns final response. var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => { callCount++; if (callCount == 1) { return new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]); } return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]); }); var options = new ToolApprovalAgentOptions { AutoApprovalRules = [fcc => new ValueTask(fcc.Name == "ReadTool")] }; var agent = new ToolApprovalAgent(innerAgent.Object, options); // Act var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Hi")], session); // Assert — the approval request was auto-approved, inner agent called twice Assert.Equal(2, callCount); Assert.Equal("Done", response.Text); } /// /// Verify that when auto-approval rule does not match, request is surfaced to the caller. /// [Fact] public async Task RunAsync_AutoApprovalRule_DoesNotMatchSurfacesToCallerAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "DangerousTool")); var innerAgent = CreateMockAgent(new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])])); var options = new ToolApprovalAgentOptions { AutoApprovalRules = [fcc => new ValueTask(fcc.Name == "ReadTool")] // Only approves ReadTool }; var agent = new ToolApprovalAgent(innerAgent.Object, options); // Act var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, "Hi")], session); // Assert — request surfaced to caller since heuristic doesn't match var requests = response.Messages.SelectMany(m => m.Contents).OfType().ToList(); Assert.Single(requests); Assert.Equal("DangerousTool", ((FunctionCallContent)requests[0].ToolCall).Name); } /// /// Verify that multiple auto-approval rules are evaluated in order; first match wins. /// [Fact] public async Task RunAsync_MultipleAutoApprovalRules_FirstMatchWinsAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "SpecialTool")); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => { callCount++; if (callCount == 1) { return new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]); } return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]); }); var rule1Called = false; var rule2Called = false; var options = new ToolApprovalAgentOptions { AutoApprovalRules = [ fcc => { rule1Called = true; return new ValueTask(fcc.Name == "SpecialTool"); }, fcc => { rule2Called = true; return new ValueTask(true); } // Should not be reached ] }; var agent = new ToolApprovalAgent(innerAgent.Object, options); // Act await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session); // Assert — first rule matched, second was never called Assert.True(rule1Called); Assert.False(rule2Called); } /// /// Verify that standing rules are evaluated before auto-approval rules. /// [Fact] public async Task RunAsync_StandingRuleTakesPrecedenceOverAutoApprovalRuleAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool")); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .ReturnsAsync(() => { callCount++; if (callCount <= 2) { return new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]); } return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]); }); var heuristicCalled = false; var options = new ToolApprovalAgentOptions { AutoApprovalRules = [fcc => { heuristicCalled = true; return new ValueTask(true); }] }; var agent = new ToolApprovalAgent(innerAgent.Object, options); // Call 1: heuristic should be called (no standing rule yet) var response1 = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session); Assert.True(heuristicCalled); Assert.Equal("Done", response1.Text); // Now establish a standing rule by sending AlwaysApprove heuristicCalled = false; callCount = 0; var alwaysApprove = new AlwaysApproveToolApprovalResponseContent( approvalRequest.CreateResponse(approved: true), alwaysApproveTool: true, alwaysApproveToolWithArguments: false); // Call 2: standing rule should match first, heuristic should NOT be called var response2 = await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session); Assert.False(heuristicCalled); Assert.Equal("Done", response2.Text); } /// /// Verify that when a batch contains a mix of heuristic-approved and standing-rule-approved /// requests, the collected approval responses preserve the original request order rather than /// being grouped by approval kind. /// [Fact] public async Task RunAsync_MixedAutoApprovals_PreserveOriginalOrderAsync() { // Arrange var session = new ChatClientAgentSession(); // Batch ordering: first request is approved by a heuristic, second by a standing rule. var heuristicRequest = new ToolApprovalRequestContent("reqA", new FunctionCallContent("callA", "HeuristicTool")); var standingRequest = new ToolApprovalRequestContent("reqB", new FunctionCallContent("callB", "StandingTool")); var batchResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [heuristicRequest, standingRequest])]); var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]); var callCount = 0; List? secondCallMessages = null; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Callback, AgentSession?, AgentRunOptions?, CancellationToken>((msgs, _, _, _) => { callCount++; if (callCount == 2) { secondCallMessages = msgs.ToList(); } }) .ReturnsAsync(() => callCount == 1 ? batchResponse : finalResponse); var options = new ToolApprovalAgentOptions { AutoApprovalRules = [fcc => new ValueTask(fcc.Name == "HeuristicTool")] }; var agent = new ToolApprovalAgent(innerAgent.Object, options); // Establish a standing rule for "StandingTool" via an AlwaysApprove response in the same call. var alwaysApprove = standingRequest.CreateAlwaysApproveToolResponse("User said always"); // Act — both requests auto-approve (heuristic + standing rule), so the inner agent is re-invoked. var response = await agent.RunAsync( [new ChatMessage(ChatRole.User, [alwaysApprove])], session); // Assert — inner agent re-called and final response returned. Assert.Equal(2, callCount); Assert.Equal("Done", response.Text); // The injected approval responses must preserve the original request order: reqA before reqB, // even though reqA was approved by a heuristic and reqB by a standing rule. Assert.NotNull(secondCallMessages); var injected = secondCallMessages! .SelectMany(m => m.Contents) .OfType() .Where(r => r.RequestId is "reqA" or "reqB") .ToList(); Assert.Equal(2, injected.Count); Assert.Equal("reqA", injected[0].RequestId); Assert.Equal("reqB", injected[1].RequestId); Assert.All(injected, r => Assert.True(r.Approved)); } /// /// Verify that auto-approval rules work in the streaming path. /// [Fact] public async Task RunStreamingAsync_AutoApprovalRule_ApprovesMatchingToolAsync() { // Arrange var session = new ChatClientAgentSession(); var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadTool")); var callCount = 0; var innerAgent = new Mock(); innerAgent .Protected() .Setup>("RunCoreStreamingAsync", ItExpr.IsAny>(), ItExpr.IsAny(), ItExpr.IsAny(), ItExpr.IsAny()) .Returns(() => { callCount++; if (callCount == 1) { return ToAsyncEnumerableAsync([new AgentResponseUpdate(ChatRole.Assistant, [approvalRequest])]); } return ToAsyncEnumerableAsync([new AgentResponseUpdate(ChatRole.Assistant, "Done")]); }); var options = new ToolApprovalAgentOptions { AutoApprovalRules = [fcc => new ValueTask(fcc.Name == "ReadTool")] }; var agent = new ToolApprovalAgent(innerAgent.Object, options); // Act var updates = new List(); await foreach (var update in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hi")], session)) { updates.Add(update); } // Assert — the approval request was auto-approved, inner agent streamed twice Assert.Equal(2, callCount); Assert.Single(updates); Assert.Equal("Done", updates[0].Text); } #endregion }