.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>
2026-06-16 21:04:09 +08:00 · 2026-06-05 18:43:07 +01:00
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 1–2: 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 1–2: 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; }
+}