Declarative workflow bugfix

2026-06-16 21:04:09 +08:00 · 2026-06-15 16:43:10 -07:00
parent 8b0405de1b
commit 61378eee01
5 changed files with 1509 additions and 166 deletions
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.

 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text.Json;
@@ -28,12 +29,23 @@ internal sealed class InvokeFunctionToolExecutor(
    WorkflowFormulaState state) :
    DeclarativeActionExecutor<InvokeFunctionTool>(model, state)
 {
-    private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshot);
+    private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshots);
+    private const string PendingCallIdsStateKey = nameof(_pendingNonApprovalCallIds);
+    private const string LegacyApprovalSnapshotStateKey = "_approvalSnapshot";

    /// <summary>
-    /// Snapshot of evaluated parameters at approval-request time.
+    /// Snapshots of evaluated parameters captured at approval-request time, keyed by
+    /// per-invocation request id. Each pending approval lives here until the matching
+    /// response is captured.
    /// </summary>
-    private ApprovalSnapshot? _approvalSnapshot;
+    private readonly ConcurrentDictionary<string, ApprovalSnapshot> _approvalSnapshots = new(StringComparer.Ordinal);
+
+    /// <summary>
+    /// Per-invocation call ids for in-flight non-approval requests; used to match the
+    /// returning <see cref="FunctionResultContent"/> on the response path. Used as a set;
+    /// the byte value is ignored.
+    /// </summary>
+    private readonly ConcurrentDictionary<string, byte> _pendingNonApprovalCallIds = new(StringComparer.Ordinal);

    /// <summary>
    /// Step identifiers for the function tool invocation workflow.
@@ -65,9 +77,12 @@ internal sealed class InvokeFunctionToolExecutor(
        bool requireApproval = this.GetRequireApproval();
        Dictionary<string, object?>? arguments = this.GetArguments();

+        // Per-invocation request id stamped on the outbound content.
+        string requestId = Guid.NewGuid().ToString("N");
+
        // Create the function call content to send to the caller
        FunctionCallContent functionCall = new(
-            callId: this.Id,
+            callId: requestId,
            name: functionName,
            arguments: arguments);

@@ -77,11 +92,15 @@ internal sealed class InvokeFunctionToolExecutor(
        // If approval is required, add user input request content
        if (requireApproval)
        {
-            // Snapshot the evaluated parameters.
-            // If state mutates during the approval window, the approved values are used on resume.
-            this._approvalSnapshot = new ApprovalSnapshot(functionName, arguments);
+            // Capture the evaluated parameters keyed by request id; the matching response
+            // resumes from this snapshot.
+            this._approvalSnapshots[requestId] = new ApprovalSnapshot(functionName, arguments);

-            requestMessage.Contents.Add(new ToolApprovalRequestContent(this.Id, functionCall));
+            requestMessage.Contents.Add(new ToolApprovalRequestContent(requestId, functionCall));
+        }
+        else
+        {
+            this._pendingNonApprovalCallIds.TryAdd(requestId, 0);
        }

        AgentResponse agentResponse = new([requestMessage]);
@@ -108,13 +127,24 @@ internal sealed class InvokeFunctionToolExecutor(
        bool autoSend = this.GetAutoSendValue();
        string? conversationId = this.GetConversationId();

-        // Extract function results from the response
-        IEnumerable<FunctionResultContent> functionResults = response.Messages
+        // Match the inbound result by its per-invocation call id.
+        FunctionResultContent? matchingResult = response.Messages
            .SelectMany(m => m.Contents)
-            .OfType<FunctionResultContent>();
+            .OfType<FunctionResultContent>()
+            .FirstOrDefault(r => this.IsKnownPendingId(r.CallId));

-        FunctionResultContent? matchingResult = functionResults
-            .FirstOrDefault(r => r.CallId == this.Id);
+        // Legacy non-approval backstop: when no pendings are tracked, accept a result
+        // whose CallId equals this.Id. The runtime has already routed the response to
+        // this executor's port and the framework does not invoke a function here.
+        if (matchingResult is null
+            && this._pendingNonApprovalCallIds.IsEmpty
+            && this._approvalSnapshots.IsEmpty)
+        {
+            matchingResult = response.Messages
+                .SelectMany(m => m.Contents)
+                .OfType<FunctionResultContent>()
+                .FirstOrDefault(r => string.Equals(r.CallId, this.Id, StringComparison.Ordinal));
+        }

        // When the caller approved an approval-required function call but didn't execute it
        // locally (the hosted Foundry scenario, where mcp_approval_response is converted to a
@@ -123,14 +153,43 @@ internal sealed class InvokeFunctionToolExecutor(
        // SendActivity/PropertyPath consumers like {Local.Result}).
        if (matchingResult is null)
        {
-            ToolApprovalResponseContent? approval = response.Messages
+            List<ToolApprovalResponseContent> approvals = response.Messages
                .SelectMany(m => m.Contents)
                .OfType<ToolApprovalResponseContent>()
-                .FirstOrDefault(r => r.RequestId == this.Id);
+                .ToList();

-            if (approval is { Approved: true })
+            // Prefer an approval matching a pending snapshot; otherwise take the first
+            // present approval.
+            ToolApprovalResponseContent? approval =
+                approvals.FirstOrDefault(r => this._approvalSnapshots.ContainsKey(r.RequestId))
+                ?? approvals.FirstOrDefault();
+
+            if (approval is not null)
            {
-                matchingResult = await this.InvokeRegisteredFunctionAsync(cancellationToken).ConfigureAwait(false);
+                if (!this._approvalSnapshots.ContainsKey(approval.RequestId))
+                {
+                    this.Logger.LogWarning(
+                        "Approval response '{RequestId}' did not match any pending invocation on '{ActionId}'.",
+                        approval.RequestId, this.Id);
+                    await this.AssignErrorAsync(context, "Function invocation was not approved by user.").ConfigureAwait(false);
+                }
+                else if (!approval.Approved)
+                {
+                    this._approvalSnapshots.TryRemove(approval.RequestId, out _);
+                    await this.AssignErrorAsync(context, "Function invocation was not approved by user.").ConfigureAwait(false);
+                }
+                else if (this._approvalSnapshots.TryRemove(approval.RequestId, out ApprovalSnapshot? snapshot))
+                {
+                    matchingResult = await this.InvokeRegisteredFunctionAsync(approval.RequestId, snapshot, cancellationToken).ConfigureAwait(false);
+                }
+                else
+                {
+                    // Snapshot was consumed by a concurrent delivery; surface the not-approved error.
+                    this.Logger.LogWarning(
+                        "Approval response '{RequestId}' had no remaining pending snapshot on '{ActionId}'.",
+                        approval.RequestId, this.Id);
+                    await this.AssignErrorAsync(context, "Function invocation was not approved by user.").ConfigureAwait(false);
+                }
            }
        }

@@ -145,6 +204,10 @@ internal sealed class InvokeFunctionToolExecutor(
                AgentResponse resultResponse = new([new ChatMessage(ChatRole.Tool, [matchingResult])]);
                await context.AddEventAsync(new AgentResponseEvent(this.Id, resultResponse), cancellationToken).ConfigureAwait(false);
            }
+
+            // Drop the per-invocation entry now that the response has been processed.
+            this._pendingNonApprovalCallIds.TryRemove(matchingResult.CallId, out _);
+            this._approvalSnapshots.TryRemove(matchingResult.CallId, out _);
        }

        // Store messages if output path is configured
@@ -167,31 +230,76 @@ internal sealed class InvokeFunctionToolExecutor(

        // Completes the action after processing the function result.
        await context.RaiseCompletionEventAsync(this.Model, cancellationToken).ConfigureAwait(false);
+    }

-        // Clear the approval snapshot after the action completes so a subsequent
-        // execution of the same executor instance doesn't reuse stale data.
-        this._approvalSnapshot = null;
-        await context.QueueStateUpdateAsync<ApprovalSnapshot?>(ApprovalSnapshotStateKey, null, null, cancellationToken).ConfigureAwait(false);
+    private bool IsKnownPendingId(string callId) =>
+        this._pendingNonApprovalCallIds.ContainsKey(callId) || this._approvalSnapshots.ContainsKey(callId);
+
+    /// <inheritdoc/>
+    public override ValueTask ResetAsync()
+    {
+        this._approvalSnapshots.Clear();
+        this._pendingNonApprovalCallIds.Clear();
+        return default;
    }

    /// <inheritdoc/>
    /// <remarks>
-    /// Persists the approval snapshot to workflow state so it survives checkpoint/restore cycles.
+    /// Persists pending approval snapshots and non-approval call ids so they survive
+    /// checkpoint/restore cycles.
    /// </remarks>
    protected override async ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
    {
-        await context.QueueStateUpdateAsync(ApprovalSnapshotStateKey, this._approvalSnapshot, null, cancellationToken).ConfigureAwait(false);
+        Dictionary<string, ApprovalSnapshot> snapshotCopy = this._approvalSnapshots.ToDictionary(kvp => kvp.Key, kvp => kvp.Value, StringComparer.Ordinal);
+        await context.QueueStateUpdateAsync(ApprovalSnapshotStateKey, snapshotCopy, null, cancellationToken).ConfigureAwait(false);
+
+        List<string> pendingCopy = [.. this._pendingNonApprovalCallIds.Keys];
+        await context.QueueStateUpdateAsync(PendingCallIdsStateKey, pendingCopy, null, cancellationToken).ConfigureAwait(false);
+
        await base.OnCheckpointingAsync(context, cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc/>
    /// <remarks>
-    /// Restores the approval snapshot from workflow state after a checkpoint restore.
+    /// Restores pending approval snapshots and non-approval call ids from workflow state
+    /// after a checkpoint restore.
    /// </remarks>
    protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
    {
        await base.OnCheckpointRestoredAsync(context, cancellationToken).ConfigureAwait(false);
-        this._approvalSnapshot = await context.ReadStateAsync<ApprovalSnapshot>(ApprovalSnapshotStateKey, null, cancellationToken).ConfigureAwait(false);
+
+        this._approvalSnapshots.Clear();
+        Dictionary<string, ApprovalSnapshot>? snapshots = await context.ReadStateAsync<Dictionary<string, ApprovalSnapshot>>(
+            ApprovalSnapshotStateKey, null, cancellationToken).ConfigureAwait(false);
+        if (snapshots is not null)
+        {
+            foreach (KeyValuePair<string, ApprovalSnapshot> entry in snapshots)
+            {
+                this._approvalSnapshots[entry.Key] = entry.Value;
+            }
+        }
+
+        this._pendingNonApprovalCallIds.Clear();
+        List<string>? pending = await context.ReadStateAsync<List<string>>(
+            PendingCallIdsStateKey, null, cancellationToken).ConfigureAwait(false);
+        if (pending is not null)
+        {
+            foreach (string id in pending)
+            {
+                this._pendingNonApprovalCallIds.TryAdd(id, 0);
+            }
+        }
+
+        // Migrate a single ApprovalSnapshot at the legacy key under this.Id so the
+        // legacy approval response matches the per-invocation map; clear the legacy key.
+        ApprovalSnapshot? legacy = await context.ReadStateAsync<ApprovalSnapshot>(
+            LegacyApprovalSnapshotStateKey, null, cancellationToken).ConfigureAwait(false);
+        if (legacy is not null)
+        {
+            this._approvalSnapshots.TryAdd(this.Id, legacy);
+            await context.QueueStateUpdateAsync<ApprovalSnapshot?>(
+                LegacyApprovalSnapshotStateKey, null, null, cancellationToken).ConfigureAwait(false);
+        }
    }

    /// <summary>
@@ -280,6 +388,14 @@ internal sealed class InvokeFunctionToolExecutor(
        await this.AssignAsync(this.Model.Output.Result?.Path, resultValue.ToFormula(), context).ConfigureAwait(false);
    }

+    private async ValueTask AssignErrorAsync(IWorkflowContext context, string errorMessage)
+    {
+        if (this.Model.Output?.Result is not null)
+        {
+            await this.AssignAsync(this.Model.Output.Result?.Path, $"Error: {errorMessage}".ToFormula(), context).ConfigureAwait(false);
+        }
+    }
+
    private string GetFunctionName() =>
        this.Evaluator.GetValue(
            Throw.IfNull(
@@ -297,32 +413,19 @@ internal sealed class InvokeFunctionToolExecutor(
        return conversationIdValue.Length == 0 ? null : conversationIdValue;
    }

-    private async ValueTask<FunctionResultContent?> InvokeRegisteredFunctionAsync(CancellationToken cancellationToken)
+    private async ValueTask<FunctionResultContent?> InvokeRegisteredFunctionAsync(string callId, ApprovalSnapshot snapshot, CancellationToken cancellationToken)
    {
-        string functionName;
-        Dictionary<string, object?>? arguments;
-
-        if (this._approvalSnapshot is { } snapshot)
-        {
-            // Use the snapshot captured at approval-request time so we invoke exactly what
-            // the user approved, even if Power Fx state has mutated during the approval window.
-            functionName = snapshot.FunctionName;
-            arguments = snapshot.Arguments;
-        }
-        else
-        {
-            // Fallback for checkpoints created before approval snapshots were introduced.
-            this.Logger.LogWarning("Approval snapshot missing for '{ActionId}'; falling back to expression re-evaluation.", this.Id);
-            functionName = this.GetFunctionName();
-            arguments = this.GetArguments();
-        }
+        // Use the snapshot captured at approval-request time so we invoke exactly what
+        // the user approved, even if Power Fx state has mutated during the approval window.
+        string functionName = snapshot.FunctionName;
+        Dictionary<string, object?>? arguments = snapshot.Arguments;

        AIFunction? function = agentProvider.Functions?.FirstOrDefault(
            f => string.Equals(f.Name, functionName, StringComparison.Ordinal));

        if (function is null)
        {
-            return new FunctionResultContent(this.Id, result: null)
+            return new FunctionResultContent(callId, result: null)
            {
                Exception = new InvalidOperationException(
                    $"Function '{functionName}' is not registered with the agent provider."),
@@ -338,7 +441,7 @@ internal sealed class InvokeFunctionToolExecutor(
        }
        catch (Exception ex) when (ex is not OperationCanceledException)
        {
-            return new FunctionResultContent(this.Id, result: null) { Exception = ex };
+            return new FunctionResultContent(callId, result: null) { Exception = ex };
        }

        // Match FunctionInvokingChatClient's serialization: pass strings through as-is and
@@ -352,7 +455,7 @@ internal sealed class InvokeFunctionToolExecutor(
            _ => JsonSerializer.Serialize(result, AIJsonUtilities.DefaultOptions.GetTypeInfo(result.GetType())),
        };

-        return new FunctionResultContent(this.Id, serialized);
+        return new FunctionResultContent(callId, serialized);
    }

    private bool GetRequireApproval()
@@ -396,9 +499,8 @@ internal sealed class InvokeFunctionToolExecutor(
    }

    /// <summary>
-    /// Stores the evaluated parameters at approval-request time so that
-    /// <see cref="CaptureResponseAsync"/> uses the values the user reviewed,
-    /// even if <see cref="WorkflowFormulaState"/> mutates during the approval window.
+    /// Captured invocation parameters used by <see cref="CaptureResponseAsync"/> on
+    /// resume so the approved values are invoked regardless of subsequent state changes.
    /// </summary>
    internal sealed record ApprovalSnapshot(
        string FunctionName,