.NET: [BREAKING] Add auto-approval rules (heuristics) to ToolApprovalAgent (#6335)

* Add support for approving tools via heuristic rules

* Address PR comments

* Address PR comments

* Apply suggestion from @SergeyMenshykh

Co-authored-by: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com>

---------

Co-authored-by: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com>
This commit is contained in:
westey
2026-06-05 18:43:07 +01:00
committed by GitHub
Unverified
parent ab8ba8fc61
commit 6bd2cfec03
8 changed files with 497 additions and 48 deletions
@@ -138,7 +138,7 @@ public sealed class HarnessAgent : DelegatingAIAgent
if (options?.DisableToolApproval is not true)
{
builder.UseToolApproval();
builder.UseToolApproval(options?.ToolApprovalAgentOptions);
}
if (options?.DisableOpenTelemetry is not true)
@@ -101,6 +101,15 @@ public sealed class HarnessAgentOptions
/// </remarks>
public bool DisableToolApproval { get; set; }
/// <summary>
/// Gets or sets the options for the <see cref="ToolApprovalAgent"/> middleware.
/// </summary>
/// <remarks>
/// When <see langword="null"/>, the <see cref="ToolApprovalAgent"/> uses default settings.
/// This property has no effect when <see cref="DisableToolApproval"/> is <see langword="true"/>.
/// </remarks>
public ToolApprovalAgentOptions? ToolApprovalAgentOptions { get; set; }
/// <summary>
/// Gets or sets a value indicating whether the <see cref="FileMemoryProvider"/> is disabled.
/// </summary>
@@ -51,20 +51,22 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
{
private readonly ProviderSessionState<ToolApprovalState> _sessionState;
private readonly JsonSerializerOptions _jsonSerializerOptions;
private readonly Func<FunctionCallContent, ValueTask<bool>>[]? _autoApprovalRules;
/// <summary>
/// Initializes a new instance of the <see cref="ToolApprovalAgent"/> class.
/// </summary>
/// <param name="innerAgent">The underlying agent to delegate to.</param>
/// <param name="jsonSerializerOptions">
/// Optional <see cref="JsonSerializerOptions"/> used for serializing argument values when storing rules
/// and for persisting state. When <see langword="null"/>, <see cref="AgentJsonUtilities.DefaultOptions"/> is used.
/// <param name="options">
/// Optional <see cref="ToolApprovalAgentOptions"/> for configuring serialization and auto-approval rules.
/// When <see langword="null"/>, default settings are used.
/// </param>
/// <exception cref="ArgumentNullException"><paramref name="innerAgent"/> is <see langword="null"/>.</exception>
public ToolApprovalAgent(AIAgent innerAgent, JsonSerializerOptions? jsonSerializerOptions = null)
public ToolApprovalAgent(AIAgent innerAgent, ToolApprovalAgentOptions? options = null)
: base(innerAgent)
{
this._jsonSerializerOptions = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
this._jsonSerializerOptions = options?.JsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
this._autoApprovalRules = options?.AutoApprovalRules?.ToArray();
this._sessionState = new ProviderSessionState<ToolApprovalState>(
_ => new ToolApprovalState(),
"toolApprovalState",
@@ -79,7 +81,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
CancellationToken cancellationToken = default)
{
// Steps 12: Unwrap AlwaysApprove wrappers, process any queued approval requests.
var (state, callerMessages, nextQueuedItem) = this.PrepareInboundMessages(messages, session);
var (state, callerMessages, nextQueuedItem) = await this.PrepareInboundMessagesAsync(messages, session).ConfigureAwait(false);
if (nextQueuedItem is not null)
{
@@ -98,7 +100,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
var response = await this.InnerAgent.RunAsync(processedMessages, session, options, cancellationToken).ConfigureAwait(false);
// Classify approval requests: auto-approve matching, queue excess, keep first unapproved.
bool allAutoApproved = this.ProcessAndQueueOutboundApprovalRequests(response.Messages, state, session);
bool allAutoApproved = await this.ProcessAndQueueOutboundApprovalRequestsAsync(response.Messages, state, session).ConfigureAwait(false);
if (!allAutoApproved)
{
@@ -119,7 +121,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
[EnumeratorCancellation] CancellationToken cancellationToken = default)
{
// Steps 12: Unwrap AlwaysApprove wrappers, process any queued approval requests.
var (state, callerMessages, nextQueuedItem) = this.PrepareInboundMessages(messages, session);
var (state, callerMessages, nextQueuedItem) = await this.PrepareInboundMessagesAsync(messages, session).ConfigureAwait(false);
if (nextQueuedItem is not null)
{
@@ -197,7 +199,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
yield break;
}
// 4. Classify the collected approval requests against standing rules.
// 4. Classify the collected approval requests against standing rules and auto-approval rules.
List<ToolApprovalRequestContent> unapproved = [];
foreach (var tarc in streamedApprovalRequests)
{
@@ -206,6 +208,11 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by standing rule"));
}
else if (await this.MatchesAutoApprovalRuleAsync(tarc).ConfigureAwait(false))
{
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by auto-approval rule"));
}
else
{
unapproved.Add(tarc);
@@ -291,9 +298,9 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
}
/// <summary>
/// Re-evaluates queued approval requests against current rules and auto-approves any that now match.
/// Re-evaluates queued approval requests against current rules and auto-approval rules, and auto-approves any that now match.
/// </summary>
private void DrainAutoApprovableFromQueue(ToolApprovalState state)
private async ValueTask DrainAutoApprovableFromQueueAsync(ToolApprovalState state)
{
for (int i = state.QueuedApprovalRequests.Count - 1; i >= 0; i--)
{
@@ -303,6 +310,12 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
state.QueuedApprovalRequests[i].CreateResponse(approved: true, reason: "Auto-approved by standing rule"));
state.QueuedApprovalRequests.RemoveAt(i);
}
else if (await this.MatchesAutoApprovalRuleAsync(state.QueuedApprovalRequests[i]).ConfigureAwait(false))
{
state.CollectedApprovalResponses.Add(
state.QueuedApprovalRequests[i].CreateResponse(approved: true, reason: "Auto-approved by auto-approval rule"));
state.QueuedApprovalRequests.RemoveAt(i);
}
}
}
@@ -318,8 +331,8 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
/// A tuple of (state, processed caller messages, next queued item or <see langword="null"/> if the queue is resolved).
/// When the returned item is non-null, the caller should return/yield it without calling the inner agent.
/// </returns>
private (ToolApprovalState State, List<ChatMessage> CallerMessages, ToolApprovalRequestContent? NextQueuedItem)
PrepareInboundMessages(IEnumerable<ChatMessage> messages, AgentSession? session)
private async ValueTask<(ToolApprovalState State, List<ChatMessage> CallerMessages, ToolApprovalRequestContent? NextQueuedItem)>
PrepareInboundMessagesAsync(IEnumerable<ChatMessage> messages, AgentSession? session)
{
var state = this._sessionState.GetOrInitializeState(session);
@@ -337,7 +350,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
// Re-evaluate remaining queued items — the caller may have added new rules
// (e.g., "always approve this tool") that resolve additional items.
this.DrainAutoApprovableFromQueue(state);
await this.DrainAutoApprovableFromQueueAsync(state).ConfigureAwait(false);
if (state.QueuedApprovalRequests.Count > 0)
{
@@ -386,15 +399,18 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
/// <see langword="true"/> if all TARc items were auto-approved (caller should re-invoke the inner agent);
/// <see langword="false"/> otherwise.
/// </returns>
private bool ProcessAndQueueOutboundApprovalRequests(
private async ValueTask<bool> ProcessAndQueueOutboundApprovalRequestsAsync(
IList<ChatMessage> responseMessages,
ToolApprovalState state,
AgentSession? session)
{
// Pass 1: Scan all response messages and classify each approval request as
// auto-approved (matches a standing rule) or unapproved (needs caller decision).
var autoApproved = new List<ToolApprovalRequestContent>();
// Pass 1: Scan all response messages and classify each approval request.
// Auto-approved requests (matching a standing rule or auto-approval rule) have their
// responses collected immediately, preserving the original request order, and are
// marked for removal. Unapproved requests are collected for the caller to decide.
var toRemove = new HashSet<ToolApprovalRequestContent>();
var unapproved = new List<ToolApprovalRequestContent>();
int autoApprovedCount = 0;
foreach (var message in responseMessages)
{
@@ -404,7 +420,17 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
{
if (MatchesRule(tarc, state.Rules, this._jsonSerializerOptions))
{
autoApproved.Add(tarc);
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by standing rule"));
toRemove.Add(tarc);
autoApprovedCount++;
}
else if (await this.MatchesAutoApprovalRuleAsync(tarc).ConfigureAwait(false))
{
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by auto-approval rule"));
toRemove.Add(tarc);
autoApprovedCount++;
}
else
{
@@ -415,18 +441,12 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
}
// Nothing to process: no auto-approved items and at most one unapproved (no queueing needed).
if (autoApproved.Count == 0 && unapproved.Count <= 1)
// No responses were collected above in this case, so state is unmodified and safe to leave.
if (autoApprovedCount == 0 && unapproved.Count <= 1)
{
return false;
}
// Store auto-approved responses for later injection into the inner agent.
foreach (var tarc in autoApproved)
{
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by standing rule"));
}
// If every approval request was auto-approved, strip them all and signal the caller
// to re-invoke the inner agent immediately with the collected responses.
if (unapproved.Count == 0)
@@ -439,14 +459,10 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
// Pass 2: Keep only the first unapproved request in the response (for the caller to decide).
// Queue the remaining unapproved requests for subsequent one-at-a-time delivery.
// Remove all auto-approved and queued items from the response messages.
var toRemove = new HashSet<ToolApprovalRequestContent>(autoApproved);
if (unapproved.Count > 1)
for (int i = 1; i < unapproved.Count; i++)
{
for (int i = 1; i < unapproved.Count; i++)
{
toRemove.Add(unapproved[i]);
state.QueuedApprovalRequests.Add(unapproved[i]);
}
toRemove.Add(unapproved[i]);
state.QueuedApprovalRequests.Add(unapproved[i]);
}
// Walk messages in reverse and strip marked items.
@@ -663,8 +679,36 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
}
/// <summary>
/// Compares stored rule arguments against actual function call arguments for an exact match.
/// Checks whether a <see cref="ToolApprovalRequestContent"/> is approved by any of the configured
/// auto-approval rules (heuristic functions).
/// </summary>
/// <returns>
/// <see langword="true"/> if any auto-approval rule returns <see langword="true"/> for the function call;
/// <see langword="false"/> if no rules are configured, the request is not a function call, or no rule approves it.
/// </returns>
private async ValueTask<bool> MatchesAutoApprovalRuleAsync(ToolApprovalRequestContent request)
{
if (this._autoApprovalRules is not { Length: > 0 })
{
return false;
}
if (request.ToolCall is not FunctionCallContent functionCall)
{
return false;
}
foreach (var rule in this._autoApprovalRules)
{
if (await rule(functionCall).ConfigureAwait(false))
{
return true;
}
}
return false;
}
private static bool ArgumentsMatch(IDictionary<string, string> ruleArguments, IDictionary<string, object?>? callArguments, JsonSerializerOptions jsonSerializerOptions)
{
if (callArguments is null)
@@ -1,7 +1,6 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Diagnostics.CodeAnalysis;
using System.Text.Json;
using Microsoft.Shared.DiagnosticIds;
using Microsoft.Shared.Diagnostics;
@@ -17,9 +16,9 @@ public static class ToolApprovalAgentBuilderExtensions
/// Adds tool approval middleware to the agent pipeline, enabling "don't ask again" approval behavior.
/// </summary>
/// <param name="builder">The <see cref="AIAgentBuilder"/> to which tool approval support will be added.</param>
/// <param name="jsonSerializerOptions">
/// Optional <see cref="JsonSerializerOptions"/> used for serializing argument values when storing rules
/// and for persisting state. When <see langword="null"/>, <see cref="AgentJsonUtilities.DefaultOptions"/> is used.
/// <param name="options">
/// Optional <see cref="ToolApprovalAgentOptions"/> for configuring serialization and auto-approval rules.
/// When <see langword="null"/>, default settings are used.
/// </param>
/// <returns>The <see cref="AIAgentBuilder"/> with tool approval middleware added, enabling method chaining.</returns>
/// <exception cref="System.ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
@@ -32,6 +31,6 @@ public static class ToolApprovalAgentBuilderExtensions
/// </remarks>
public static AIAgentBuilder UseToolApproval(
this AIAgentBuilder builder,
JsonSerializerOptions? jsonSerializerOptions = null)
=> Throw.IfNull(builder).Use(innerAgent => new ToolApprovalAgent(innerAgent, jsonSerializerOptions));
ToolApprovalAgentOptions? options = null)
=> Throw.IfNull(builder).Use(innerAgent => new ToolApprovalAgent(innerAgent, options));
}
@@ -0,0 +1,45 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Text.Json;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Microsoft.Shared.DiagnosticIds;
namespace Microsoft.Agents.AI;
/// <summary>
/// Options for configuring the <see cref="ToolApprovalAgent"/> middleware.
/// </summary>
[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
public class ToolApprovalAgentOptions
{
/// <summary>
/// Gets or sets the <see cref="System.Text.Json.JsonSerializerOptions"/> used for serializing argument values
/// when storing rules and for persisting state.
/// </summary>
/// <remarks>
/// When <see langword="null"/>, <see cref="AgentJsonUtilities.DefaultOptions"/> is used.
/// </remarks>
public JsonSerializerOptions? JsonSerializerOptions { get; set; }
/// <summary>
/// Gets or sets a collection of heuristic functions that can automatically approve function calls
/// that would otherwise require user approval.
/// </summary>
/// <remarks>
/// <para>
/// Each function receives a <see cref="FunctionCallContent"/> representing the tool call that requires approval
/// and returns a <see cref="ValueTask{Boolean}"/> that resolves to <see langword="true"/> to auto-approve
/// the call, or <see langword="false"/> to continue evaluating the next rule.
/// </para>
/// <para>
/// Auto-approval rules are evaluated after standing rules (derived from prior user approvals) but before
/// prompting the user. Rules are evaluated in order; the first rule returning <see langword="true"/>
/// causes the function call to be auto-approved.
/// </para>
/// </remarks>
public IEnumerable<Func<FunctionCallContent, ValueTask<bool>>>? AutoApprovalRules { get; set; }
}
@@ -644,6 +644,51 @@ public class HarnessAgentTests
Assert.Null(agent.GetService<ToolApprovalAgent>());
}
/// <summary>
/// Verify that ToolApprovalAgentOptions auto-approval rules are passed through and actually used.
/// </summary>
[Fact]
public async Task ToolApproval_AutoApprovalRulesAreAppliedAsync()
{
// Arrange — inner client returns an approval request on first call, then final response on second.
var callCount = 0;
var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadTool"));
var mockClient = new Mock<IChatClient>();
mockClient
.Setup(c => c.GetResponseAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<ChatOptions>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync(() =>
{
callCount++;
if (callCount == 1)
{
return new ChatResponse(new ChatMessage(ChatRole.Assistant, [approvalRequest]));
}
return new ChatResponse(new ChatMessage(ChatRole.Assistant, "Done"));
});
var options = CreateAllDisabledOptions();
options.DisableToolApproval = false;
options.ToolApprovalAgentOptions = new ToolApprovalAgentOptions
{
AutoApprovalRules = [fcc => new ValueTask<bool>(fcc.Name == "ReadTool")]
};
var agent = new HarnessAgent(mockClient.Object, TestMaxContextWindowTokens, TestMaxOutputTokens, options);
var session = await agent.CreateSessionAsync();
// Act
var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session);
// Assert — the auto-approval rule approved the request, so we get "Done" (not an approval request)
Assert.Equal(2, callCount);
Assert.Equal("Done", response.Text);
}
#endregion
#region Feature: OpenTelemetry
@@ -59,15 +59,15 @@ public class ToolApprovalAgentBuilderExtensionsTests
/// Verify that UseToolApproval with custom JsonSerializerOptions works correctly.
/// </summary>
[Fact]
public void UseToolApproval_WithCustomJsonSerializerOptions_ReturnsToolApprovalAgent()
public void UseToolApproval_WithCustomOptions_ReturnsToolApprovalAgent()
{
// Arrange
var mockAgent = new Mock<AIAgent>();
var builder = new AIAgentBuilder(mockAgent.Object);
var options = new JsonSerializerOptions();
var options = new ToolApprovalAgentOptions { JsonSerializerOptions = new JsonSerializerOptions() };
// Act
var result = builder.UseToolApproval(jsonSerializerOptions: options).Build();
var result = builder.UseToolApproval(options: options).Build();
// Assert
Assert.IsType<ToolApprovalAgent>(result);
@@ -47,14 +47,14 @@ public class ToolApprovalAgentTests
}
/// <summary>
/// Verify that constructor accepts custom JsonSerializerOptions.
/// Verify that constructor accepts custom options.
/// </summary>
[Fact]
public void Constructor_CustomJsonSerializerOptions_CreatesInstanceAsync()
public void Constructor_CustomOptions_CreatesInstance()
{
// Arrange
var innerAgent = new Mock<AIAgent>().Object;
var options = new JsonSerializerOptions();
var options = new ToolApprovalAgentOptions { JsonSerializerOptions = new JsonSerializerOptions() };
// Act
var agent = new ToolApprovalAgent(innerAgent, options);
@@ -1535,4 +1535,311 @@ public class ToolApprovalAgentTests
}
#endregion
#region Auto-Approval Rules (Heuristics)
/// <summary>
/// Verify that an auto-approval rule can approve a function call that would otherwise need user approval.
/// </summary>
[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<AIAgent>();
innerAgent
.Protected()
.Setup<Task<AgentResponse>>("RunCoreAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.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<bool>(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);
}
/// <summary>
/// Verify that when auto-approval rule does not match, request is surfaced to the caller.
/// </summary>
[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<bool>(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<ToolApprovalRequestContent>().ToList();
Assert.Single(requests);
Assert.Equal("DangerousTool", ((FunctionCallContent)requests[0].ToolCall).Name);
}
/// <summary>
/// Verify that multiple auto-approval rules are evaluated in order; first match wins.
/// </summary>
[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<AIAgent>();
innerAgent
.Protected()
.Setup<Task<AgentResponse>>("RunCoreAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.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<bool>(fcc.Name == "SpecialTool"); },
fcc => { rule2Called = true; return new ValueTask<bool>(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);
}
/// <summary>
/// Verify that standing rules are evaluated before auto-approval rules.
/// </summary>
[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<AIAgent>();
innerAgent
.Protected()
.Setup<Task<AgentResponse>>("RunCoreAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.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<bool>(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);
}
/// <summary>
/// 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.
/// </summary>
[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<ChatMessage>? secondCallMessages = null;
var innerAgent = new Mock<AIAgent>();
innerAgent
.Protected()
.Setup<Task<AgentResponse>>("RunCoreAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.Callback<IEnumerable<ChatMessage>, AgentSession?, AgentRunOptions?, CancellationToken>((msgs, _, _, _) =>
{
callCount++;
if (callCount == 2)
{
secondCallMessages = msgs.ToList();
}
})
.ReturnsAsync(() => callCount == 1 ? batchResponse : finalResponse);
var options = new ToolApprovalAgentOptions
{
AutoApprovalRules = [fcc => new ValueTask<bool>(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<ToolApprovalResponseContent>()
.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));
}
/// <summary>
/// Verify that auto-approval rules work in the streaming path.
/// </summary>
[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<AIAgent>();
innerAgent
.Protected()
.Setup<IAsyncEnumerable<AgentResponseUpdate>>("RunCoreStreamingAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.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<bool>(fcc.Name == "ReadTool")]
};
var agent = new ToolApprovalAgent(innerAgent.Object, options);
// Act
var updates = new List<AgentResponseUpdate>();
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
}