.NET: Bug fixes for declarative workflows (#6427)

* declarative workflow approval flow fix * Update mcp handler cache construction * fix method argument. * Update dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeFunctionToolExecutor.cs Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Fix identation --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-06-10 18:08:32 +00:00
parent 60cc5ee4e4
commit 3753d938f5
4 changed files with 664 additions and 22 deletions
@@ -1,11 +1,21 @@
 // Copyright (c) Microsoft. All rights reserved.

+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Reflection;
+using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Agents.AI.Workflows.Declarative.Events;
+using Microsoft.Agents.AI.Workflows.Declarative.Kit;
 using Microsoft.Agents.AI.Workflows.Declarative.ObjectModel;
 using Microsoft.Agents.AI.Workflows.Declarative.PowerFx;
 using Microsoft.Agents.ObjectModel;
 using Microsoft.Extensions.AI;
+using Microsoft.PowerFx.Types;
+using Moq;
+using ApprovalSnapshot = Microsoft.Agents.AI.Workflows.Declarative.ObjectModel.InvokeFunctionToolExecutor.ApprovalSnapshot;

 namespace Microsoft.Agents.AI.Workflows.Declarative.UnitTests.ObjectModel;

@@ -261,6 +271,323 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W

    #endregion

+    #region Approval Snapshot Security Tests
+
+    /// <summary>
+    /// Verifies that mutating the function-name variable after approval does not change
+    /// which function is actually invoked. The originally-approved name must be used.
+    /// </summary>
+    [Fact]
+    public async Task InvokeFunctionToolCaptureResponseUsesApprovedFunctionNameNotMutatedAsync()
+    {
+        // Arrange
+        const string ApprovedFunctionName = "safe_readonly_query";
+        const string MutatedFunctionName = "dangerous_admin_tool";
+
+        this.State.Set("TargetFunction", FormulaValue.New(ApprovedFunctionName));
+        this.State.InitializeSystem();
+        this.State.Bind();
+
+        InvokeFunctionTool model = this.CreateModelWithVariableFunctionName(
+            displayName: nameof(InvokeFunctionToolCaptureResponseUsesApprovedFunctionNameNotMutatedAsync),
+            variableName: "TargetFunction");
+
+        string? capturedFunctionName = null;
+        TestFunctionAgentProvider testAgentProvider = new(
+            [
+                AIFunctionFactory.Create(() => "safe-result", name: ApprovedFunctionName),
+                AIFunctionFactory.Create(() => "dangerous-result", name: MutatedFunctionName),
+            ],
+            onInvoke: name => capturedFunctionName = name);
+        InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
+
+        // Act - trigger ExecuteAsync to store the approval snapshot
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
+        await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
+
+        // Simulate parallel branch mutating state during the approval window
+        this.State.Set("TargetFunction", FormulaValue.New(MutatedFunctionName));
+        this.State.Bind();
+
+        // User clicks approve (they saw "safe_readonly_query" in the approval UI)
+        ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
+
+        // Resume after approval
+        await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
+
+        // Assert - the originally-approved function must be invoked, not the mutated one
+        Assert.NotNull(capturedFunctionName);
+        Assert.Equal(ApprovedFunctionName, capturedFunctionName);
+    }
+
+    /// <summary>
+    /// Verifies that mutating an argument variable after approval does not change
+    /// the arguments actually passed to the invoked function.
+    /// </summary>
+    [Fact]
+    public async Task InvokeFunctionToolCaptureResponseUsesApprovedArgumentsNotMutatedAsync()
+    {
+        // Arrange
+        const string FunctionName = "process_query";
+        const string ArgumentKey = "query";
+        const string ApprovedQuery = "SELECT * FROM users LIMIT 10";
+        const string MutatedQuery = "DROP TABLE users CASCADE; --";
+
+        this.State.Set("SqlQuery", FormulaValue.New(ApprovedQuery));
+        this.State.InitializeSystem();
+        this.State.Bind();
+
+        InvokeFunctionTool model = this.CreateModelWithVariableArgument(
+            displayName: nameof(InvokeFunctionToolCaptureResponseUsesApprovedArgumentsNotMutatedAsync),
+            functionName: FunctionName,
+            argumentKey: ArgumentKey,
+            variableName: "SqlQuery");
+
+        AIFunctionArguments? capturedArguments = null;
+        TestFunctionAgentProvider testAgentProvider = new(
+            [AIFunctionFactory.Create((string query) => $"executed:{query}", name: FunctionName)],
+            onInvokeArguments: args => capturedArguments = args);
+        InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
+
+        // Act - trigger ExecuteAsync to store the approval snapshot
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
+        await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
+
+        // Simulate parallel branch mutating state during the approval window
+        this.State.Set("SqlQuery", FormulaValue.New(MutatedQuery));
+        this.State.Bind();
+
+        // User clicks approve
+        ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
+
+        // Resume after approval
+        await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
+
+        // Assert - the originally-approved argument must be used, not the mutated one
+        Assert.NotNull(capturedArguments);
+        Assert.Equal(ApprovedQuery, capturedArguments[ArgumentKey]?.ToString());
+    }
+
+    /// <summary>
+    /// Verifies that the approval snapshot survives a checkpoint/restore cycle.
+    /// After restore, the originally-approved function must still be used even if state was mutated.
+    /// </summary>
+    [Fact]
+    public async Task InvokeFunctionToolCaptureResponseUsesSnapshotAfterCheckpointRestoreAsync()
+    {
+        // Arrange
+        const string ApprovedFunctionName = "safe_readonly_query";
+        const string MutatedFunctionName = "dangerous_admin_tool";
+
+        this.State.Set("TargetFunction", FormulaValue.New(ApprovedFunctionName));
+        this.State.InitializeSystem();
+        this.State.Bind();
+
+        InvokeFunctionTool model = this.CreateModelWithVariableFunctionName(
+            displayName: nameof(InvokeFunctionToolCaptureResponseUsesSnapshotAfterCheckpointRestoreAsync),
+            variableName: "TargetFunction");
+
+        string? capturedFunctionName = null;
+        TestFunctionAgentProvider testAgentProvider = new(
+            [
+                AIFunctionFactory.Create(() => "safe-result", name: ApprovedFunctionName),
+                AIFunctionFactory.Create(() => "dangerous-result", name: MutatedFunctionName),
+            ],
+            onInvoke: name => capturedFunctionName = name);
+        InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
+
+        // Act - trigger ExecuteAsync to store the approval snapshot
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore();
+        await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
+
+        // Simulate checkpoint: persist to state store
+        await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
+
+        // Simulate restore on a "new" executor instance by clearing the in-memory field via reflection
+        // (In production, a new executor instance would be created with _approvalSnapshot == null)
+        typeof(InvokeFunctionToolExecutor)
+            .GetField("_approvalSnapshot", BindingFlags.NonPublic | BindingFlags.Instance)!
+            .SetValue(action, null);
+
+        // Restore from state store
+        await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
+
+        // Mutate state after restore (simulating parallel branch)
+        this.State.Set("TargetFunction", FormulaValue.New(MutatedFunctionName));
+        this.State.Bind();
+
+        // User clicks approve
+        ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
+
+        // Resume after approval
+        await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
+
+        // Assert - the originally-approved function must be invoked, not the mutated one
+        Assert.NotNull(capturedFunctionName);
+        Assert.Equal(ApprovedFunctionName, capturedFunctionName);
+    }
+
+    /// <summary>
+    /// Verifies that the approval snapshot is cleared after a completed approval cycle,
+    /// both in-memory and in the persisted state store. This prevents stale data from
+    /// influencing a subsequent execution of the same executor instance.
+    /// </summary>
+    [Fact]
+    public async Task InvokeFunctionToolCaptureResponseClearsSnapshotAfterCompletionAsync()
+    {
+        // Arrange
+        const string FunctionName = "any_function";
+
+        this.State.InitializeSystem();
+        this.State.Bind();
+
+        InvokeFunctionTool model = this.CreateModel(
+            displayName: nameof(InvokeFunctionToolCaptureResponseClearsSnapshotAfterCompletionAsync),
+            functionName: FunctionName,
+            requireApproval: true);
+
+        TestFunctionAgentProvider testAgentProvider = new(
+            [AIFunctionFactory.Create(() => "result", name: FunctionName)]);
+        InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
+
+        // Act - run the full approval cycle
+        Dictionary<string, object?> stateStore = [];
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(stateStore);
+        await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
+
+        // Sanity: snapshot was captured
+        FieldInfo snapshotField = typeof(InvokeFunctionToolExecutor)
+            .GetField("_approvalSnapshot", BindingFlags.NonPublic | BindingFlags.Instance)!;
+        Assert.NotNull(snapshotField.GetValue(action));
+
+        ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
+        await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
+
+        // Assert - both in-memory field and persisted state are cleared
+        Assert.Null(snapshotField.GetValue(action));
+        Assert.True(stateStore.ContainsKey("_approvalSnapshot"));
+        Assert.Null(stateStore["_approvalSnapshot"]);
+    }
+
+    private static ExternalInputResponse CreateApprovalResponse(string actionId, bool approved)
+    {
+        FunctionCallContent functionCall = new(callId: actionId, name: "ignored");
+        ToolApprovalRequestContent approvalRequest = new(actionId, functionCall);
+        ToolApprovalResponseContent approvalResponse = approvalRequest.CreateResponse(approved);
+        return new ExternalInputResponse(new ChatMessage(ChatRole.User, [approvalResponse]));
+    }
+
+    private static Mock<IWorkflowContext> CreateMockWorkflowContext()
+    {
+        Mock<IWorkflowContext> mockContext = new();
+        mockContext.Setup(c => c.AddEventAsync(It.IsAny<WorkflowEvent>(), It.IsAny<CancellationToken>()))
+            .Returns(default(ValueTask));
+        mockContext.Setup(c => c.QueueStateUpdateAsync(It.IsAny<string>(), It.IsAny<object?>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
+            .Returns(default(ValueTask));
+        mockContext.Setup(c => c.SendMessageAsync(It.IsAny<object>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
+            .Returns(default(ValueTask));
+        return mockContext;
+    }
+
+    /// <summary>
+    /// Creates a mock workflow context that actually stores state values (for checkpoint/restore tests).
+    /// Optionally accepts an externally-owned dictionary so callers can inspect the persisted state.
+    /// </summary>
+    private static Mock<IWorkflowContext> CreateMockWorkflowContextWithStateStore(Dictionary<string, object?>? stateStore = null)
+    {
+        stateStore ??= [];
+        Mock<IWorkflowContext> mockContext = new();
+        mockContext.Setup(c => c.AddEventAsync(It.IsAny<WorkflowEvent>(), It.IsAny<CancellationToken>()))
+            .Returns(default(ValueTask));
+        mockContext.Setup(c => c.QueueStateUpdateAsync(It.IsAny<string>(), It.IsAny<ApprovalSnapshot?>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
+            .Callback<string, ApprovalSnapshot?, string?, CancellationToken>((key, value, _, _) => stateStore[key] = value)
+            .Returns(default(ValueTask));
+        mockContext.Setup(c => c.SendMessageAsync(It.IsAny<object>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
+            .Returns(default(ValueTask));
+        mockContext.Setup(c => c.ReadStateAsync<ApprovalSnapshot>(It.IsAny<string>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
+            .Returns<string, string?, CancellationToken>((key, _, _) =>
+                new ValueTask<ApprovalSnapshot?>(stateStore.TryGetValue(key, out object? val) ? val as ApprovalSnapshot : null));
+        mockContext.Setup(c => c.ReadStateKeysAsync(It.IsAny<string?>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new HashSet<string>());
+        return mockContext;
+    }
+
+    /// <summary>
+    /// Invokes a protected method on the executor via reflection (for testing checkpoint hooks).
+    /// </summary>
+    private static async ValueTask InvokeProtectedMethodAsync(InvokeFunctionToolExecutor action, string methodName, IWorkflowContext context, CancellationToken cancellationToken)
+    {
+        MethodInfo method = typeof(InvokeFunctionToolExecutor)
+            .GetMethod(methodName, BindingFlags.NonPublic | BindingFlags.Instance)!;
+        ValueTask result = (ValueTask)method.Invoke(action, [context, cancellationToken])!;
+        await result.ConfigureAwait(false);
+    }
+
+    /// <summary>
+    /// Minimal concrete <see cref="ResponseAgentProvider"/> that exposes an injected
+    /// <see cref="AIFunction"/> registry and records which function got invoked.
+    /// Used by the framework-invoke approval branch (<c>InvokeRegisteredFunctionAsync</c>).
+    /// </summary>
+    private sealed class TestFunctionAgentProvider : ResponseAgentProvider
+    {
+        private readonly Action<string>? _onInvoke;
+        private readonly Action<AIFunctionArguments>? _onInvokeArguments;
+
+        public TestFunctionAgentProvider(
+            IEnumerable<AIFunction> functions,
+            Action<string>? onInvoke = null,
+            Action<AIFunctionArguments>? onInvokeArguments = null)
+        {
+            this._onInvoke = onInvoke;
+            this._onInvokeArguments = onInvokeArguments;
+            this.Functions = functions.Select(f => (AIFunction)new RecordingAIFunction(f, this)).ToList();
+        }
+
+        internal void RecordInvocation(string name, AIFunctionArguments? arguments)
+        {
+            this._onInvoke?.Invoke(name);
+            if (arguments is not null)
+            {
+                this._onInvokeArguments?.Invoke(arguments);
+            }
+        }
+
+        public override Task<string> CreateConversationAsync(CancellationToken cancellationToken = default) =>
+            throw new NotSupportedException();
+
+        public override Task<ChatMessage> CreateMessageAsync(string conversationId, ChatMessage conversationMessage, CancellationToken cancellationToken = default) =>
+            throw new NotSupportedException();
+
+        public override Task<ChatMessage> GetMessageAsync(string conversationId, string messageId, CancellationToken cancellationToken = default) =>
+            throw new NotSupportedException();
+
+        public override IAsyncEnumerable<AgentResponseUpdate> InvokeAgentAsync(
+            string agentId, string? agentVersion, string? conversationId,
+            IEnumerable<ChatMessage>? messages, IDictionary<string, object?>? inputArguments,
+            CancellationToken cancellationToken = default) =>
+            throw new NotSupportedException();
+
+        public override IAsyncEnumerable<ChatMessage> GetMessagesAsync(
+            string conversationId, int? limit = null, string? after = null, string? before = null,
+            bool newestFirst = false, CancellationToken cancellationToken = default) =>
+            throw new NotSupportedException();
+
+        private sealed class RecordingAIFunction(AIFunction inner, TestFunctionAgentProvider owner) : AIFunction
+        {
+            public override string Name => inner.Name;
+            public override string Description => inner.Description;
+            public override JsonElement JsonSchema => inner.JsonSchema;
+
+            protected override ValueTask<object?> InvokeCoreAsync(AIFunctionArguments arguments, CancellationToken cancellationToken)
+            {
+                owner.RecordInvocation(inner.Name, arguments);
+                return inner.InvokeAsync(arguments, cancellationToken);
+            }
+        }
+    }
+
+    #endregion
+
    #region Helper Methods

    private async Task ExecuteTestAsync(InvokeFunctionTool model)
@@ -318,5 +645,33 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
        return AssignParent<InvokeFunctionTool>(builder);
    }

+    private InvokeFunctionTool CreateModelWithVariableFunctionName(string displayName, string variableName)
+    {
+        InvokeFunctionTool.Builder builder = new()
+        {
+            Id = this.CreateActionId(),
+            DisplayName = this.FormatDisplayName(displayName),
+            FunctionName = new StringExpression.Builder(
+                StringExpression.Variable(PropertyPath.TopicVariable(variableName))),
+            RequireApproval = new BoolExpression.Builder(BoolExpression.Literal(true)),
+        };
+        return AssignParent<InvokeFunctionTool>(builder);
+    }
+
+    private InvokeFunctionTool CreateModelWithVariableArgument(
+        string displayName, string functionName, string argumentKey, string variableName)
+    {
+        InvokeFunctionTool.Builder builder = new()
+        {
+            Id = this.CreateActionId(),
+            DisplayName = this.FormatDisplayName(displayName),
+            FunctionName = new StringExpression.Builder(StringExpression.Literal(functionName)),
+            RequireApproval = new BoolExpression.Builder(BoolExpression.Literal(true)),
+        };
+        builder.Arguments.Add(argumentKey,
+            ValueExpression.Variable(PropertyPath.TopicVariable(variableName)));
+        return AssignParent<InvokeFunctionTool>(builder);
+    }
+
    #endregion
 }