agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Declarative workflow bugfix

Browse Source
This commit is contained in:
Peter Ibekwe
2026-06-15 16:43:10 -07:00
parent 8b0405de1b
commit 61378eee01
5 changed files with 1509 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeFunctionToolExecutorTest.cs
+776 -35
View File
@@ -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.Reflection;
@@ -301,8 +302,9 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
onInvoke: name => capturedFunctionName = name);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
// Act - trigger ExecuteAsync to store the approval snapshot
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
// Act - trigger ExecuteAsync to emit the approval request
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Simulate parallel branch mutating state during the approval window
@@ -310,7 +312,7 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
this.State.Bind();
// User clicks approve (they saw "safe_readonly_query" in the approval UI)
ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
// Resume after approval
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -349,8 +351,9 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
onInvokeArguments: args => capturedArguments = args);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
// Act - trigger ExecuteAsync to store the approval snapshot
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
// Act - trigger ExecuteAsync to emit the approval request
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Simulate parallel branch mutating state during the approval window
@@ -358,7 +361,7 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
this.State.Bind();
// User clicks approve
ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
// Resume after approval
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -396,18 +399,20 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
onInvoke: name => capturedFunctionName = name);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
// Act - trigger ExecuteAsync to store the approval snapshot
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore();
// Act - trigger ExecuteAsync to emit the approval request and capture the snapshot
List<ExternalInputRequest> emittedRequests = [];
Dictionary<string, object?> stateStore = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(stateStore, emittedRequests);
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);
// Simulate restore on a "new" executor instance by clearing the in-memory dictionary via reflection
ConcurrentDictionary<string, ApprovalSnapshot> liveSnapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeFunctionToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
.GetValue(action)!;
liveSnapshots.Clear();
// Restore from state store
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
@@ -417,7 +422,7 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
this.State.Bind();
// User clicks approve
ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
// Resume after approval
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -428,9 +433,7 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
}
/// <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.
/// Verifies that the approval snapshot entry is removed after a completed approval cycle.
/// </summary>
[Fact]
public async Task InvokeFunctionToolCaptureResponseClearsSnapshotAfterCompletionAsync()
@@ -451,33 +454,723 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
// Act - run the full approval cycle
List<ExternalInputRequest> emittedRequests = [];
Dictionary<string, object?> stateStore = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(stateStore);
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(stateStore, emittedRequests);
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));
// Sanity: snapshot dict has exactly one entry
FieldInfo snapshotsField = typeof(InvokeFunctionToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!;
ConcurrentDictionary<string, ApprovalSnapshot> snapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)snapshotsField.GetValue(action)!;
Assert.Single(snapshots);
ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, 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"]);
// Assert - in-memory dict is empty after the matching response is captured
Assert.Empty(snapshots);
}
private static ExternalInputResponse CreateApprovalResponse(string actionId, bool approved)
/// <summary>
/// Each ExecuteAsync invocation must produce a unique per-invocation request id on
/// both the FunctionCallContent.CallId and the ToolApprovalRequestContent.RequestId.
/// </summary>
[Fact]
public async Task InvokeFunctionToolEmitsUniqueRequestIdPerInvocationAsync()
{
FunctionCallContent functionCall = new(callId: actionId, name: "ignored");
ToolApprovalRequestContent approvalRequest = new(actionId, functionCall);
// Arrange
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolEmitsUniqueRequestIdPerInvocationAsync),
functionName: "any_function",
requireApproval: true);
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => "result", name: "any_function")]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
// Act - emit two approval requests from the same executor instance
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Assert - two distinct request ids surfaced
Assert.Equal(2, emittedRequests.Count);
string id1 = emittedRequests[0].AgentResponse.Messages
.SelectMany(m => m.Contents).OfType<ToolApprovalRequestContent>().Single().RequestId;
string id2 = emittedRequests[1].AgentResponse.Messages
.SelectMany(m => m.Contents).OfType<ToolApprovalRequestContent>().Single().RequestId;
Assert.NotEqual(id1, id2);
Assert.NotEqual(action.Id, id1);
Assert.NotEqual(action.Id, id2);
// And the matching inner FunctionCallContent uses the same id
FunctionCallContent fcc1 = emittedRequests[0].AgentResponse.Messages
.SelectMany(m => m.Contents).OfType<FunctionCallContent>().Single();
Assert.Equal(id1, fcc1.CallId);
}
/// <summary>
/// Two concurrent pending approvals on the same executor must each resume with their
/// own approved arguments — out-of-order responses must not swap which invocation gets
/// which set of arguments.
/// </summary>
[Fact]
public async Task InvokeFunctionToolConcurrentPendingApprovalsDoNotSwapAsync()
{
// Arrange
const string FunctionName = "process_query";
const string ArgumentKey = "query";
const string ArgumentsA = "A-args";
const string ArgumentsB = "B-args";
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolConcurrentPendingApprovalsDoNotSwapAsync),
functionName: FunctionName,
requireApproval: true,
argumentKey: ArgumentKey,
argumentValue: ArgumentsA);
InvokeFunctionTool modelB = this.CreateModel(
displayName: nameof(InvokeFunctionToolConcurrentPendingApprovalsDoNotSwapAsync) + "B",
functionName: FunctionName,
requireApproval: true,
argumentKey: ArgumentKey,
argumentValue: ArgumentsB);
List<string?> capturedQueries = [];
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create((string query) => $"executed:{query}", name: FunctionName)],
onInvokeArguments: args => capturedQueries.Add(args[ArgumentKey]?.ToString()));
// Two executor instances simulating concurrent fan-in scenarios with different YAML-evaluated args
InvokeFunctionToolExecutor actionA = new(model, testAgentProvider, this.State);
InvokeFunctionToolExecutor actionB = new(modelB, testAgentProvider, this.State);
List<ExternalInputRequest> emittedA = [];
List<ExternalInputRequest> emittedB = [];
Mock<IWorkflowContext> ctxA = CreateMockWorkflowContext(emittedA);
Mock<IWorkflowContext> ctxB = CreateMockWorkflowContext(emittedB);
// Act - both executors emit approval requests
await actionA.HandleAsync(new ActionExecutorResult(actionA.Id), ctxA.Object, CancellationToken.None);
await actionB.HandleAsync(new ActionExecutorResult(actionB.Id), ctxB.Object, CancellationToken.None);
// Deliver responses out of order: B first, then A
await actionB.CaptureResponseAsync(ctxB.Object, CreateApprovalResponseFor(emittedB, approved: true), CancellationToken.None);
await actionA.CaptureResponseAsync(ctxA.Object, CreateApprovalResponseFor(emittedA, approved: true), CancellationToken.None);
// Assert - each invocation executed with its own approved arguments
Assert.Equal([ArgumentsB, ArgumentsA], capturedQueries);
}
/// <summary>
/// When the approval response references a request id that is not in the snapshot map,
/// the executor must surface a structured error and must not invoke any function.
/// </summary>
[Fact]
public async Task InvokeFunctionToolMissingSnapshotReturnsStructuredErrorAsync()
{
// Arrange
const string FunctionName = "any_function";
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolMissingSnapshotReturnsStructuredErrorAsync),
functionName: FunctionName,
requireApproval: true);
bool functionWasInvoked = false;
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => { functionWasInvoked = true; return "result"; }, name: FunctionName)]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
// Act - deliver an approval response whose RequestId has no matching snapshot
FunctionCallContent fcc = new(callId: "stale-id", name: FunctionName);
ToolApprovalRequestContent staleRequest = new("stale-id", fcc);
ToolApprovalResponseContent staleResponse = staleRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [staleResponse]));
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - the registered function must NOT have been invoked. The
// ToolApprovalResponseContent.RequestId did not match any snapshot in the executor's
// map, so the executor does not attempt to invoke the function at all (no silent
// state re-evaluation).
Assert.False(functionWasInvoked);
}
/// <summary>
/// Two non-approval invocations of the same executor must emit distinct per-invocation
/// CallIds so each response is matched to its originating request.
/// </summary>
[Fact]
public async Task InvokeFunctionToolNonApprovalCallIdsAreDistinctPerInvocationAsync()
{
// Arrange
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolNonApprovalCallIdsAreDistinctPerInvocationAsync),
functionName: "any_function",
requireApproval: false);
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => "result", name: "any_function")]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
// Act - emit two non-approval function-call requests
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Assert - distinct CallIds were stamped on the two emitted FunctionCallContents
Assert.Equal(2, emittedRequests.Count);
FunctionCallContent fcc1 = emittedRequests[0].AgentResponse.Messages
.SelectMany(m => m.Contents).OfType<FunctionCallContent>().Single();
FunctionCallContent fcc2 = emittedRequests[1].AgentResponse.Messages
.SelectMany(m => m.Contents).OfType<FunctionCallContent>().Single();
Assert.NotEqual(fcc1.CallId, fcc2.CallId);
Assert.NotEqual(action.Id, fcc1.CallId);
Assert.NotEqual(action.Id, fcc2.CallId);
}
/// <summary>
/// A snapshot persisted at the legacy <c>"_approvalSnapshot"</c> key must be migrated
/// under <c>this.Id</c> after restore so an approval response carrying
/// <c>RequestId == this.Id</c> resumes with the snapshot's arguments.
/// </summary>
[Fact]
public async Task InvokeFunctionToolLegacySingleSnapshotCheckpointIsMigratedAsync()
{
// Arrange
const string FunctionName = "any_function";
const string ArgumentKey = "query";
const string LegacyApprovedArg = "legacy-approved";
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolLegacySingleSnapshotCheckpointIsMigratedAsync),
functionName: FunctionName,
requireApproval: true);
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);
// Seed the state store with a single ApprovalSnapshot at the legacy key.
Dictionary<string, object?> stateStore = new()
{
["_approvalSnapshot"] = new ApprovalSnapshot(
FunctionName,
new Dictionary<string, object?> { [ArgumentKey] = LegacyApprovedArg }),
};
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(stateStore);
// Act - restore migrates the legacy snapshot under this.Id.
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
ConcurrentDictionary<string, ApprovalSnapshot> snapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeFunctionToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
.GetValue(action)!;
Assert.True(snapshots.ContainsKey(action.Id));
// Deliver an approval response with RequestId == action.Id and resume.
FunctionCallContent fcc = new(callId: action.Id, name: FunctionName);
ToolApprovalRequestContent legacyRequest = new(action.Id, fcc);
ToolApprovalResponseContent legacyResponse = legacyRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [legacyResponse]));
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - the function was invoked with the snapshot arguments.
Assert.NotNull(capturedArguments);
Assert.Equal(LegacyApprovedArg, capturedArguments[ArgumentKey]?.ToString());
}
/// <summary>
/// The legacy <c>"_approvalSnapshot"</c> key is removed from the state store after
/// migration so subsequent checkpoints do not carry stale data.
/// </summary>
[Fact]
public async Task InvokeFunctionToolLegacyKeyIsClearedAfterMigrationAsync()
{
// Arrange
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolLegacyKeyIsClearedAfterMigrationAsync),
functionName: "any_function",
requireApproval: true);
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => "result", name: "any_function")]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
Dictionary<string, object?> stateStore = new()
{
["_approvalSnapshot"] = new ApprovalSnapshot("any_function", new Dictionary<string, object?>()),
};
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(stateStore);
// Act
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
// Assert - legacy key was cleared via QueueStateUpdateAsync<ApprovalSnapshot?>(null).
Assert.False(stateStore.ContainsKey("_approvalSnapshot"));
}
/// <summary>
/// Drives ExecuteAsync → checkpoint → ResetAsync → restore → CaptureResponseAsync on a
/// single pending approval and asserts the originally-approved arguments are used,
/// even though ResetAsync cleared the in-memory dict between checkpoint and restore.
/// </summary>
[Fact]
public async Task InvokeFunctionToolResumeAfterResetUsesPersistedSnapshotAsync()
{
// Arrange
const string FunctionName = "process_query";
const string ArgumentKey = "query";
const string ApprovedQuery = "SELECT * FROM users LIMIT 10";
this.State.Set("SqlQuery", FormulaValue.New(ApprovedQuery));
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModelWithVariableArgument(
displayName: nameof(InvokeFunctionToolResumeAfterResetUsesPersistedSnapshotAsync),
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);
List<ExternalInputRequest> emittedRequests = [];
Dictionary<string, object?> stateStore = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(stateStore, emittedRequests);
ConcurrentDictionary<string, ApprovalSnapshot> liveSnapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeFunctionToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
.GetValue(action)!;
// Act - emit, checkpoint, reset (simulates runner end), restore, then capture.
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
await action.ResetAsync();
Assert.Empty(liveSnapshots);
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
Assert.Single(liveSnapshots);
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - the originally-approved argument was used and the entry was removed.
Assert.NotNull(capturedArguments);
Assert.Equal(ApprovedQuery, capturedArguments[ArgumentKey]?.ToString());
Assert.Empty(liveSnapshots);
}
/// <summary>
/// Two pending invocations (A then B) are interleaved with checkpoint/reset/restore
/// cycles; A's snapshot must survive both reset cycles and route A's response to
/// A's arguments, while B remains pending and is later resolved correctly.
/// </summary>
[Fact]
public async Task InvokeFunctionToolMultiplePendingInvocationsSurviveCheckpointResetRestoreAsync()
{
// Arrange
const string FunctionName = "process_query";
const string ArgumentKey = "query";
const string ArgumentsA = "A-args";
const string ArgumentsB = "B-args";
this.State.Set("SqlQuery", FormulaValue.New(ArgumentsA));
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModelWithVariableArgument(
displayName: nameof(InvokeFunctionToolMultiplePendingInvocationsSurviveCheckpointResetRestoreAsync),
functionName: FunctionName,
argumentKey: ArgumentKey,
variableName: "SqlQuery");
List<string?> capturedQueries = [];
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create((string query) => $"executed:{query}", name: FunctionName)],
onInvokeArguments: args => capturedQueries.Add(args[ArgumentKey]?.ToString()));
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
List<ExternalInputRequest> emittedRequests = [];
Dictionary<string, object?> stateStore = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(stateStore, emittedRequests);
ConcurrentDictionary<string, ApprovalSnapshot> liveSnapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeFunctionToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
.GetValue(action)!;
// Act - invocation A with ArgumentsA, then full checkpoint/reset/restore.
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
await action.ResetAsync();
Assert.Empty(liveSnapshots);
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
Assert.Single(liveSnapshots);
// Mutate the source variable, then invocation B with ArgumentsB.
this.State.Set("SqlQuery", FormulaValue.New(ArgumentsB));
this.State.Bind();
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
await action.ResetAsync();
Assert.Empty(liveSnapshots);
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
Assert.Equal(2, liveSnapshots.Count);
// Capture A's response. State has been mutated to ArgumentsB but the per-invocation
// snapshot must still drive invocation with ArgumentsA.
Assert.Equal(2, emittedRequests.Count);
ExternalInputResponse responseA = CreateApprovalResponseForRequest(emittedRequests[0], approved: true);
await action.CaptureResponseAsync(mockContext.Object, responseA, CancellationToken.None);
Assert.Single(liveSnapshots);
Assert.Equal([ArgumentsA], capturedQueries);
// Another checkpoint/reset/restore cycle - B's snapshot survives.
await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
await action.ResetAsync();
Assert.Empty(liveSnapshots);
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
Assert.Single(liveSnapshots);
// Capture B's response.
ExternalInputResponse responseB = CreateApprovalResponseForRequest(emittedRequests[1], approved: true);
await action.CaptureResponseAsync(mockContext.Object, responseB, CancellationToken.None);
// Assert - both invocations executed with their own approved arguments; nothing pending.
Assert.Equal([ArgumentsA, ArgumentsB], capturedQueries);
Assert.Empty(liveSnapshots);
}
/// <summary>
/// An approval response whose RequestId does not match any pending snapshot must
/// NOT invoke the function and must assign a not-approved error to Output.Result.
/// </summary>
[Fact]
public async Task InvokeFunctionToolUnmatchedApprovalAssignsErrorAsync()
{
// Arrange
const string FunctionName = "any_function";
const string ResultVariable = "Result";
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolUnmatchedApprovalAssignsErrorAsync),
functionName: FunctionName,
requireApproval: true,
outputResultVariable: ResultVariable);
bool functionWasInvoked = false;
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => { functionWasInvoked = true; return "result"; }, name: FunctionName)]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
// Act - deliver an approval response whose RequestId has no matching snapshot.
FunctionCallContent fcc = new(callId: "stale-id", name: FunctionName);
ToolApprovalRequestContent staleRequest = new("stale-id", fcc);
ToolApprovalResponseContent staleResponse = staleRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [staleResponse]));
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - function was NOT invoked AND the error string landed at Output.Result.
Assert.False(functionWasInvoked);
Assert.Contains(mockContext.Invocations, i =>
i.Method.Name == nameof(IWorkflowContext.QueueStateUpdateAsync)
&& i.Arguments.Count >= 2
&& i.Arguments[1] is StringValue sv
&& sv.Value.Contains("not approved by user"));
}
/// <summary>
/// An approval response whose RequestId matches a pending snapshot but is
/// Approved == false must NOT invoke the function, must remove the snapshot, and
/// must assign a not-approved error to Output.Result.
/// </summary>
[Fact]
public async Task InvokeFunctionToolRejectedApprovalAssignsErrorAsync()
{
// Arrange
const string FunctionName = "any_function";
const string ResultVariable = "Result";
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolRejectedApprovalAssignsErrorAsync),
functionName: FunctionName,
requireApproval: true,
outputResultVariable: ResultVariable);
bool functionWasInvoked = false;
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => { functionWasInvoked = true; return "result"; }, name: FunctionName)]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
ConcurrentDictionary<string, ApprovalSnapshot> liveSnapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeFunctionToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
.GetValue(action)!;
// Act - emit the approval request, then deliver a rejection for it.
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
Assert.Single(liveSnapshots);
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: false);
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - function not invoked, snapshot removed, error assigned.
Assert.False(functionWasInvoked);
Assert.Empty(liveSnapshots);
Assert.Contains(mockContext.Invocations, i =>
i.Method.Name == nameof(IWorkflowContext.QueueStateUpdateAsync)
&& i.Arguments.Count >= 2
&& i.Arguments[1] is StringValue sv
&& sv.Value.Contains("not approved by user"));
}
/// <summary>
/// When a response contains multiple <see cref="ToolApprovalResponseContent"/> items —
/// e.g. an unrelated / stale approval followed by the valid one — the executor must
/// select the approval whose RequestId matches a pending snapshot and invoke the
/// function, not silently drop the valid approval because a stale one appeared first.
/// </summary>
[Fact]
public async Task InvokeFunctionToolApprovalMatchPrefersPendingSnapshotAsync()
{
// Arrange
const string FunctionName = "any_function";
const string ResultVariable = "Result";
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolApprovalMatchPrefersPendingSnapshotAsync),
functionName: FunctionName,
requireApproval: true,
outputResultVariable: ResultVariable);
bool functionWasInvoked = false;
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => { functionWasInvoked = true; return "result"; }, name: FunctionName)]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
// Emit one valid approval request from this executor.
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
ExternalInputRequest emitted = Assert.Single(emittedRequests);
ToolApprovalRequestContent validRequest = emitted.AgentResponse.Messages
.SelectMany(m => m.Contents)
.OfType<ToolApprovalRequestContent>()
.Single();
// Build a batched response: a stale (unrelated) approval first, then the valid one.
ToolApprovalRequestContent staleRequest = new("stale-id", new FunctionCallContent("stale-id", FunctionName));
ToolApprovalResponseContent staleResponse = staleRequest.CreateResponse(approved: true);
ToolApprovalResponseContent validResponse = validRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [staleResponse, validResponse]));
// Act
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - the valid approval drove invocation; no not-approved error was assigned.
Assert.True(functionWasInvoked);
Assert.DoesNotContain(mockContext.Invocations, i =>
i.Method.Name == nameof(IWorkflowContext.QueueStateUpdateAsync)
&& i.Arguments.Count >= 2
&& i.Arguments[1] is StringValue sv
&& sv.Value.Contains("not approved by user"));
}
/// <summary>
/// Delivering the same approval response twice must invoke the registered function
/// exactly once; the second delivery surfaces the not-approved error path because the
/// snapshot has already been consumed.
/// </summary>
[Fact]
public async Task InvokeFunctionToolDuplicateApprovalDeliveryInvokesFunctionOnceAsync()
{
// Arrange
const string FunctionName = "any_function";
const string ResultVariable = "Result";
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolDuplicateApprovalDeliveryInvokesFunctionOnceAsync),
functionName: FunctionName,
requireApproval: true,
outputResultVariable: ResultVariable);
int invocationCount = 0;
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => { Interlocked.Increment(ref invocationCount); return "result"; }, name: FunctionName)]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
// Emit one approval request.
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Act - deliver the SAME approval response twice.
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - the registered AIFunction was invoked exactly once.
Assert.Equal(1, invocationCount);
// The second delivery surfaced the not-approved error path.
Assert.Contains(mockContext.Invocations, i =>
i.Method.Name == nameof(IWorkflowContext.QueueStateUpdateAsync)
&& i.Arguments.Count >= 2
&& i.Arguments[1] is StringValue sv
&& sv.Value.Contains("not approved by user"));
}
/// <summary>
/// A non-approval <c>FunctionResultContent</c> whose CallId equals <c>this.Id</c> is
/// consumed and assigned to <c>Output.Result</c> when no pendings are tracked.
/// </summary>
[Fact]
public async Task InvokeFunctionToolLegacyNonApprovalResultIsAcceptedAsync()
{
// Arrange - a fresh executor has no tracked pendings.
const string FunctionName = "any_function";
const string ResultVariable = "Result";
const string HostResult = "host-computed-result";
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolLegacyNonApprovalResultIsAcceptedAsync),
functionName: FunctionName,
requireApproval: false,
outputResultVariable: ResultVariable);
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => "should-not-be-called", name: FunctionName)]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
// Act - deliver a FunctionResultContent with CallId == action.Id.
FunctionResultContent legacyResult = new(action.Id, HostResult);
ExternalInputResponse response = new(new ChatMessage(ChatRole.Tool, [legacyResult]));
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - the host-computed result was assigned to Output.Result and no
// not-approved error was emitted.
Assert.Contains(mockContext.Invocations, i =>
i.Method.Name == nameof(IWorkflowContext.QueueStateUpdateAsync)
&& i.Arguments.Count >= 2
&& i.Arguments[1] is StringValue sv
&& sv.Value == HostResult);
Assert.DoesNotContain(mockContext.Invocations, i =>
i.Method.Name == nameof(IWorkflowContext.QueueStateUpdateAsync)
&& i.Arguments.Count >= 2
&& i.Arguments[1] is StringValue sv
&& sv.Value.Contains("not approved by user"));
}
/// <summary>
/// The legacy non-approval backstop must NOT fire when the executor has a tracked
/// pending invocation; a <c>FunctionResultContent</c> with <c>CallId == this.Id</c>
/// is rejected in that state.
/// </summary>
[Fact]
public async Task InvokeFunctionToolLegacyNonApprovalBackstopGatedOnEmptyStateAsync()
{
// Arrange - emit a non-approval call so a per-invocation CallId is tracked.
const string FunctionName = "any_function";
const string ResultVariable = "Result";
this.State.InitializeSystem();
this.State.Bind();
InvokeFunctionTool model = this.CreateModel(
displayName: nameof(InvokeFunctionToolLegacyNonApprovalBackstopGatedOnEmptyStateAsync),
functionName: FunctionName,
requireApproval: false,
outputResultVariable: ResultVariable);
TestFunctionAgentProvider testAgentProvider = new(
[AIFunctionFactory.Create(() => "result", name: FunctionName)]);
InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Act - deliver a FunctionResultContent with CallId == action.Id (not the emitted GUID).
FunctionResultContent staleLegacyResult = new(action.Id, "should-be-rejected");
ExternalInputResponse response = new(new ChatMessage(ChatRole.Tool, [staleLegacyResult]));
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - Output.Result was NOT assigned with the rejected result.
Assert.DoesNotContain(mockContext.Invocations, i =>
i.Method.Name == nameof(IWorkflowContext.QueueStateUpdateAsync)
&& i.Arguments.Count >= 2
&& i.Arguments[1] is StringValue sv
&& sv.Value == "should-be-rejected");
}
/// <summary>
/// Builds an approval response paired to the inner <c>ToolApprovalRequestContent.RequestId</c>
/// of a specific emitted request. Used when multiple requests are emitted and the
/// caller needs to address one by position.
/// </summary>
private static ExternalInputResponse CreateApprovalResponseForRequest(ExternalInputRequest emitted, bool approved)
{
ToolApprovalRequestContent approvalRequest = emitted.AgentResponse.Messages
.SelectMany(m => m.Contents)
.OfType<ToolApprovalRequestContent>()
.Single();
ToolApprovalResponseContent approvalResponse = approvalRequest.CreateResponse(approved);
return new ExternalInputResponse(new ChatMessage(ChatRole.User, [approvalResponse]));
}
private static Mock<IWorkflowContext> CreateMockWorkflowContext()
/// <summary>
/// Extracts the inner <c>ToolApprovalRequestContent.RequestId</c> from the
/// approval request the executor emitted, and builds a paired response. This mirrors
/// the framework's symmetric content-id rewriting at the envelope boundary.
/// </summary>
private static ExternalInputResponse CreateApprovalResponseFor(IReadOnlyList<ExternalInputRequest> emittedRequests, bool approved)
{
ExternalInputRequest emitted = Assert.Single(emittedRequests);
return CreateApprovalResponseForRequest(emitted, approved);
}
private static Mock<IWorkflowContext> CreateMockWorkflowContext(List<ExternalInputRequest>? emittedRequests = null)
{
Mock<IWorkflowContext> mockContext = new();
mockContext.Setup(c => c.AddEventAsync(It.IsAny<WorkflowEvent>(), It.IsAny<CancellationToken>()))
@@ -485,25 +1178,64 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
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>()))
.Callback<object, string?, CancellationToken>((msg, _, _) =>
{
if (emittedRequests is not null && msg is ExternalInputRequest request)
{
emittedRequests.Add(request);
}
})
.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.
/// Optionally accepts an externally-owned dictionary so callers can inspect the persisted state,
/// and an optional emitted-request list so tests can build matching responses.
/// </summary>
private static Mock<IWorkflowContext> CreateMockWorkflowContextWithStateStore(Dictionary<string, object?>? stateStore = null)
private static Mock<IWorkflowContext> CreateMockWorkflowContextWithStateStore(
Dictionary<string, object?>? stateStore = null,
List<ExternalInputRequest>? emittedRequests = 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<Dictionary<string, ApprovalSnapshot>>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
.Callback<string, Dictionary<string, ApprovalSnapshot>, string?, CancellationToken>((key, value, _, _) => stateStore[key] = value)
.Returns(default(ValueTask));
mockContext.Setup(c => c.QueueStateUpdateAsync(It.IsAny<string>(), It.IsAny<List<string>>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
.Callback<string, List<string>, string?, CancellationToken>((key, value, _, _) => stateStore[key] = value)
.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)
.Callback<string, ApprovalSnapshot?, string?, CancellationToken>((key, value, _, _) =>
{
if (value is null)
{
stateStore.Remove(key);
}
else
{
stateStore[key] = value;
}
})
.Returns(default(ValueTask));
mockContext.Setup(c => c.SendMessageAsync(It.IsAny<object>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
.Callback<object, string?, CancellationToken>((msg, _, _) =>
{
if (emittedRequests is not null && msg is ExternalInputRequest request)
{
emittedRequests.Add(request);
}
})
.Returns(default(ValueTask));
mockContext.Setup(c => c.ReadStateAsync<Dictionary<string, ApprovalSnapshot>>(It.IsAny<string>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
.Returns<string, string?, CancellationToken>((key, _, _) =>
new ValueTask<Dictionary<string, ApprovalSnapshot>?>(stateStore.TryGetValue(key, out object? val) ? val as Dictionary<string, ApprovalSnapshot> : null));
mockContext.Setup(c => c.ReadStateAsync<List<string>>(It.IsAny<string>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
.Returns<string, string?, CancellationToken>((key, _, _) =>
new ValueTask<List<string>?>(stateStore.TryGetValue(key, out object? val) ? val as List<string> : null));
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));
@@ -622,7 +1354,8 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
bool? requireApproval = false,
string? conversationId = null,
string? argumentKey = null,
string? argumentValue = null)
string? argumentValue = null,
string? outputResultVariable = null)
{
InvokeFunctionTool.Builder builder = new()
{
@@ -642,6 +1375,14 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
builder.Arguments.Add(argumentKey, ValueExpression.Literal(new StringDataValue(argumentValue)));
}
if (outputResultVariable is not null)
{
builder.Output = new InvokeToolOutput.Builder
{
Result = new InitializablePropertyPath(PropertyPath.TopicVariable(outputResultVariable), isInitializer: false),
};
}
return AssignParent<InvokeFunctionTool>(builder);
}
dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeMcpToolExecutorTest.cs
+509 -39
View File
@@ -1,5 +1,6 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
@@ -423,15 +424,15 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
MockAgentProvider mockAgentProvider = new();
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
Mock<IWorkflowContext> mockContext = new(MockBehavior.Loose);
// Emit the approval request so the executor records the per-invocation snapshot.
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Build an approved response matching this action's request id.
McpServerToolCallContent toolCall = new(action.Id, TestToolName, TestServerLabel);
ToolApprovalRequestContent approvalRequest = new(action.Id, toolCall);
ToolApprovalResponseContent approvalResponse = approvalRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [approvalResponse]));
// Build the matching approved response from the emitted request.
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
// Act - call CaptureResponseAsync directly so the post-approval branch actually executes.
// Act - call CaptureResponseAsync so the post-approval branch actually executes.
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - headers reach the transport invocation on the approved path.
@@ -887,7 +888,8 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
// Act - trigger ExecuteAsync to store the approval snapshot
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Simulate parallel branch mutating state during the approval window
@@ -895,10 +897,7 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
this.State.Bind();
// User clicks approve (they saw "safe_readonly_query" in the approval UI)
McpServerToolCallContent toolCall = new(action.Id, ApprovedToolName, TestServerUrl);
ToolApprovalRequestContent approvalRequest = new(action.Id, toolCall);
ToolApprovalResponseContent approvalResponse = approvalRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [approvalResponse]));
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
// Resume after approval
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -950,7 +949,8 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
// Act - trigger ExecuteAsync to store the approval snapshot
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Simulate parallel branch mutating state during the approval window
@@ -958,10 +958,7 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
this.State.Bind();
// User clicks approve
McpServerToolCallContent toolCall = new(action.Id, TestToolName, TestServerUrl);
ToolApprovalRequestContent approvalRequest = new(action.Id, toolCall);
ToolApprovalResponseContent approvalResponse = approvalRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [approvalResponse]));
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
// Resume after approval
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -1011,7 +1008,8 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
// Act - trigger ExecuteAsync to store the approval snapshot
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Simulate parallel branch mutating state during the approval window
@@ -1019,10 +1017,7 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
this.State.Bind();
// User clicks approve
McpServerToolCallContent toolCall = new(action.Id, TestToolName, ApprovedServerUrl);
ToolApprovalRequestContent approvalRequest = new(action.Id, toolCall);
ToolApprovalResponseContent approvalResponse = approvalRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [approvalResponse]));
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
// Resume after approval
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -1072,17 +1067,18 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
// Act - trigger ExecuteAsync to store the approval snapshot
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore();
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(emittedRequests);
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(InvokeMcpToolExecutor)
.GetField("_approvalSnapshot", BindingFlags.NonPublic | BindingFlags.Instance)!
.SetValue(action, null);
// Simulate restore on a "new" executor instance by clearing the in-memory dictionary via reflection
ConcurrentDictionary<string, ApprovalSnapshot> liveSnapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeMcpToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
.GetValue(action)!;
liveSnapshots.Clear();
// Restore from state store
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
@@ -1092,10 +1088,7 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
this.State.Bind();
// User clicks approve
McpServerToolCallContent toolCall = new(action.Id, ApprovedToolName, TestServerUrl);
ToolApprovalRequestContent approvalRequest = new(action.Id, toolCall);
ToolApprovalResponseContent approvalResponse = approvalRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [approvalResponse]));
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
// Resume after approval
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -1105,7 +1098,440 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
Assert.Equal(ApprovedToolName, capturedToolName);
}
private static Mock<IWorkflowContext> CreateMockWorkflowContext()
/// <summary>
/// Each ExecuteAsync invocation must produce a unique per-invocation request id on
/// both the McpServerToolCallContent and the wrapping ToolApprovalRequestContent.
/// </summary>
[Fact]
public async Task InvokeMcpToolEmitsUniqueRequestIdPerInvocationAsync()
{
// Arrange
this.State.InitializeSystem();
this.State.Bind();
InvokeMcpTool model = this.CreateModelWithApproval(
displayName: nameof(InvokeMcpToolEmitsUniqueRequestIdPerInvocationAsync),
serverUrl: TestServerUrl,
toolName: TestToolName);
Mock<IMcpToolHandler> mockProvider = new();
MockAgentProvider mockAgentProvider = new();
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
List<ExternalInputRequest> emittedRequests = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
// Act
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
// Assert - two distinct request ids surfaced
Assert.Equal(2, emittedRequests.Count);
string id1 = emittedRequests[0].AgentResponse.Messages
.SelectMany(m => m.Contents).OfType<ToolApprovalRequestContent>().Single().RequestId;
string id2 = emittedRequests[1].AgentResponse.Messages
.SelectMany(m => m.Contents).OfType<ToolApprovalRequestContent>().Single().RequestId;
Assert.NotEqual(id1, id2);
Assert.NotEqual(action.Id, id1);
Assert.NotEqual(action.Id, id2);
}
/// <summary>
/// Two concurrent pending MCP approvals on different executor instances (representing
/// concurrent fan-in or interleaved invocations) must each resume with their own
/// approved parameters when responses are delivered out of order.
/// </summary>
[Fact]
public async Task InvokeMcpToolConcurrentPendingApprovalsDoNotSwapAsync()
{
// Arrange
const string ToolA = "tool_alpha";
const string ToolB = "tool_beta";
this.State.InitializeSystem();
this.State.Bind();
InvokeMcpTool modelA = this.CreateModelWithApproval(
displayName: nameof(InvokeMcpToolConcurrentPendingApprovalsDoNotSwapAsync) + "A",
serverUrl: TestServerUrl,
toolName: ToolA);
InvokeMcpTool modelB = this.CreateModelWithApproval(
displayName: nameof(InvokeMcpToolConcurrentPendingApprovalsDoNotSwapAsync) + "B",
serverUrl: TestServerUrl,
toolName: ToolB);
List<string?> capturedToolNames = [];
Mock<IMcpToolHandler> mockProvider = new();
mockProvider.Setup(p => p.InvokeToolAsync(
It.IsAny<string>(),
It.IsAny<string?>(),
It.IsAny<string>(),
It.IsAny<IDictionary<string, object?>?>(),
It.IsAny<IDictionary<string, string>?>(),
It.IsAny<string?>(),
It.IsAny<CancellationToken>()))
.Callback<string, string?, string, IDictionary<string, object?>?, IDictionary<string, string>?, string?, CancellationToken>(
(_, _, toolName, _, _, _, _) => capturedToolNames.Add(toolName))
.ReturnsAsync(new McpServerToolResultContent("capture-call-id")
{
Outputs = [new TextContent("ok")]
});
MockAgentProvider mockAgentProvider = new();
InvokeMcpToolExecutor actionA = new(modelA, mockProvider.Object, mockAgentProvider.Object, this.State);
InvokeMcpToolExecutor actionB = new(modelB, mockProvider.Object, mockAgentProvider.Object, this.State);
List<ExternalInputRequest> emittedA = [];
List<ExternalInputRequest> emittedB = [];
Mock<IWorkflowContext> ctxA = CreateMockWorkflowContext(emittedA);
Mock<IWorkflowContext> ctxB = CreateMockWorkflowContext(emittedB);
// Act - both executors emit approval requests
await actionA.HandleAsync(new ActionExecutorResult(actionA.Id), ctxA.Object, CancellationToken.None);
await actionB.HandleAsync(new ActionExecutorResult(actionB.Id), ctxB.Object, CancellationToken.None);
// Deliver responses out of order
await actionB.CaptureResponseAsync(ctxB.Object, CreateApprovalResponseFor(emittedB, approved: true), CancellationToken.None);
await actionA.CaptureResponseAsync(ctxA.Object, CreateApprovalResponseFor(emittedA, approved: true), CancellationToken.None);
// Assert - each invocation invoked its own approved tool name
Assert.Equal([ToolB, ToolA], capturedToolNames);
}
/// <summary>
/// When the approval response references a request id that is not in the snapshot map,
/// the executor must NOT invoke the MCP tool.
/// </summary>
[Fact]
public async Task InvokeMcpToolMissingSnapshotAssignsErrorAsync()
{
// Arrange
this.State.InitializeSystem();
this.State.Bind();
InvokeMcpTool model = this.CreateModelWithApproval(
displayName: nameof(InvokeMcpToolMissingSnapshotAssignsErrorAsync),
serverUrl: TestServerUrl,
toolName: TestToolName);
Mock<IMcpToolHandler> mockProvider = new();
MockAgentProvider mockAgentProvider = new();
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
// Act - deliver an approval response whose RequestId has no matching snapshot
McpServerToolCallContent toolCall = new("stale-id", TestToolName, TestServerUrl);
ToolApprovalRequestContent staleRequest = new("stale-id", toolCall);
ToolApprovalResponseContent staleResponse = staleRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [staleResponse]));
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - mcpToolHandler.InvokeToolAsync must NOT have been called
mockProvider.Verify(p => p.InvokeToolAsync(
It.IsAny<string>(),
It.IsAny<string?>(),
It.IsAny<string>(),
It.IsAny<IDictionary<string, object?>?>(),
It.IsAny<IDictionary<string, string>?>(),
It.IsAny<string?>(),
It.IsAny<CancellationToken>()), Times.Never);
}
/// <summary>
/// A snapshot persisted at the legacy <c>"_approvalSnapshot"</c> key must be migrated
/// under <c>this.Id</c> after restore so an approval response carrying
/// <c>RequestId == this.Id</c> resumes with the snapshot's tool name.
/// </summary>
[Fact]
public async Task InvokeMcpToolLegacySingleSnapshotCheckpointIsMigratedAsync()
{
// Arrange
const string LegacyApprovedToolName = "legacy_approved_tool";
this.State.InitializeSystem();
this.State.Bind();
InvokeMcpTool model = this.CreateModelWithApproval(
displayName: nameof(InvokeMcpToolLegacySingleSnapshotCheckpointIsMigratedAsync),
serverUrl: TestServerUrl,
toolName: TestToolName);
string? capturedToolName = null;
Mock<IMcpToolHandler> mockProvider = new();
mockProvider.Setup(p => p.InvokeToolAsync(
It.IsAny<string>(),
It.IsAny<string?>(),
It.IsAny<string>(),
It.IsAny<IDictionary<string, object?>?>(),
It.IsAny<IDictionary<string, string>?>(),
It.IsAny<string?>(),
It.IsAny<CancellationToken>()))
.Callback<string, string?, string, IDictionary<string, object?>?, IDictionary<string, string>?, string?, CancellationToken>(
(_, _, toolName, _, _, _, _) => capturedToolName = toolName)
.ReturnsAsync(new McpServerToolResultContent("capture-call-id")
{
Outputs = [new TextContent("ok")]
});
MockAgentProvider mockAgentProvider = new();
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
// Seed the state store with a single ApprovalSnapshot at the legacy key.
Dictionary<string, object?> stateStore = new()
{
["_approvalSnapshot"] = new ApprovalSnapshot(
TestServerUrl, null, LegacyApprovedToolName, new Dictionary<string, object?>(), null),
};
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStoreSeeded(stateStore);
// Act - restore migrates the legacy snapshot under this.Id.
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
ConcurrentDictionary<string, ApprovalSnapshot> snapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeMcpToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
.GetValue(action)!;
Assert.True(snapshots.ContainsKey(action.Id));
// Deliver an approval response with RequestId == action.Id and resume.
McpServerToolCallContent toolCall = new(action.Id, LegacyApprovedToolName, TestServerUrl);
ToolApprovalRequestContent legacyRequest = new(action.Id, toolCall);
ToolApprovalResponseContent legacyResponse = legacyRequest.CreateResponse(approved: true);
ExternalInputResponse response = new(new ChatMessage(ChatRole.User, [legacyResponse]));
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - the MCP tool was invoked with the snapshot's tool name.
Assert.Equal(LegacyApprovedToolName, capturedToolName);
}
/// <summary>
/// The legacy <c>"_approvalSnapshot"</c> key is removed from the state store after
/// migration so subsequent checkpoints do not carry stale data.
/// </summary>
[Fact]
public async Task InvokeMcpToolLegacyKeyIsClearedAfterMigrationAsync()
{
// Arrange
this.State.InitializeSystem();
this.State.Bind();
InvokeMcpTool model = this.CreateModelWithApproval(
displayName: nameof(InvokeMcpToolLegacyKeyIsClearedAfterMigrationAsync),
serverUrl: TestServerUrl,
toolName: TestToolName);
Mock<IMcpToolHandler> mockProvider = new();
MockAgentProvider mockAgentProvider = new();
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
Dictionary<string, object?> stateStore = new()
{
["_approvalSnapshot"] = new ApprovalSnapshot(
TestServerUrl, null, TestToolName, new Dictionary<string, object?>(), null),
};
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStoreSeeded(stateStore);
// Act
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
// Assert - legacy key was cleared via QueueStateUpdateAsync<ApprovalSnapshot?>(null).
Assert.False(stateStore.ContainsKey("_approvalSnapshot"));
}
/// <summary>
/// Variant of CreateMockWorkflowContextWithStateStore that accepts a pre-seeded state
/// store and supports the read/write operations exercised by the legacy-migration path.
/// </summary>
private static Mock<IWorkflowContext> CreateMockWorkflowContextWithStateStoreSeeded(Dictionary<string, object?> 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, _, _) =>
{
if (value is null)
{
stateStore.Remove(key);
}
else
{
stateStore[key] = value;
}
})
.Returns(default(ValueTask));
mockContext.Setup(c => c.ReadStateAsync<Dictionary<string, ApprovalSnapshot>>(It.IsAny<string>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
.Returns<string, string?, CancellationToken>((key, _, _) =>
new ValueTask<Dictionary<string, ApprovalSnapshot>?>(stateStore.TryGetValue(key, out object? val) ? val as Dictionary<string, ApprovalSnapshot> : null));
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>
/// Drives ExecuteAsync → checkpoint → ResetAsync → restore → CaptureResponseAsync on a
/// single pending approval and asserts the originally-approved tool name is used,
/// even though ResetAsync cleared the in-memory dict between checkpoint and restore.
/// </summary>
[Fact]
public async Task InvokeMcpToolResumeAfterResetUsesPersistedSnapshotAsync()
{
// Arrange
const string ApprovedToolName = "approved_tool";
this.State.Set("TargetTool", FormulaValue.New(ApprovedToolName));
this.State.InitializeSystem();
this.State.Bind();
InvokeMcpTool model = this.CreateModelWithVariableToolName(
displayName: nameof(InvokeMcpToolResumeAfterResetUsesPersistedSnapshotAsync),
serverUrl: TestServerUrl,
variableName: "TargetTool");
string? capturedToolName = null;
Mock<IMcpToolHandler> mockProvider = new();
mockProvider.Setup(p => p.InvokeToolAsync(
It.IsAny<string>(),
It.IsAny<string?>(),
It.IsAny<string>(),
It.IsAny<IDictionary<string, object?>?>(),
It.IsAny<IDictionary<string, string>?>(),
It.IsAny<string?>(),
It.IsAny<CancellationToken>()))
.Callback<string, string?, string, IDictionary<string, object?>?, IDictionary<string, string>?, string?, CancellationToken>(
(_, _, toolName, _, _, _, _) => capturedToolName = toolName)
.ReturnsAsync(new McpServerToolResultContent("capture-call-id")
{
Outputs = [new TextContent("ok")]
});
MockAgentProvider mockAgentProvider = new();
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
List<ExternalInputRequest> emittedRequests = [];
Dictionary<string, object?> stateStore = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(emittedRequests, stateStore);
ConcurrentDictionary<string, ApprovalSnapshot> liveSnapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeMcpToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
.GetValue(action)!;
// Act - emit, checkpoint, reset (simulates runner end), restore, then capture.
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
await action.ResetAsync();
Assert.Empty(liveSnapshots);
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
Assert.Single(liveSnapshots);
ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
// Assert - the originally-approved tool name was used and the entry was removed.
Assert.Equal(ApprovedToolName, capturedToolName);
Assert.Empty(liveSnapshots);
}
/// <summary>
/// Two pending invocations (A then B) are interleaved with checkpoint/reset/restore
/// cycles; A's snapshot must survive both reset cycles and route A's response to
/// A's tool name, while B remains pending and is later resolved correctly.
/// </summary>
[Fact]
public async Task InvokeMcpToolMultiplePendingInvocationsSurviveCheckpointResetRestoreAsync()
{
// Arrange
const string ToolA = "tool_alpha";
const string ToolB = "tool_beta";
this.State.Set("TargetTool", FormulaValue.New(ToolA));
this.State.InitializeSystem();
this.State.Bind();
InvokeMcpTool model = this.CreateModelWithVariableToolName(
displayName: nameof(InvokeMcpToolMultiplePendingInvocationsSurviveCheckpointResetRestoreAsync),
serverUrl: TestServerUrl,
variableName: "TargetTool");
List<string?> capturedToolNames = [];
Mock<IMcpToolHandler> mockProvider = new();
mockProvider.Setup(p => p.InvokeToolAsync(
It.IsAny<string>(),
It.IsAny<string?>(),
It.IsAny<string>(),
It.IsAny<IDictionary<string, object?>?>(),
It.IsAny<IDictionary<string, string>?>(),
It.IsAny<string?>(),
It.IsAny<CancellationToken>()))
.Callback<string, string?, string, IDictionary<string, object?>?, IDictionary<string, string>?, string?, CancellationToken>(
(_, _, toolName, _, _, _, _) => capturedToolNames.Add(toolName))
.ReturnsAsync(new McpServerToolResultContent("capture-call-id")
{
Outputs = [new TextContent("ok")]
});
MockAgentProvider mockAgentProvider = new();
InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
List<ExternalInputRequest> emittedRequests = [];
Dictionary<string, object?> stateStore = [];
Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(emittedRequests, stateStore);
ConcurrentDictionary<string, ApprovalSnapshot> liveSnapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeMcpToolExecutor)
.GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
.GetValue(action)!;
// Act - invocation A with ToolA, then full checkpoint/reset/restore.
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
await action.ResetAsync();
Assert.Empty(liveSnapshots);
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
Assert.Single(liveSnapshots);
// Mutate the source variable, then invocation B with ToolB.
this.State.Set("TargetTool", FormulaValue.New(ToolB));
this.State.Bind();
await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
await action.ResetAsync();
Assert.Empty(liveSnapshots);
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
Assert.Equal(2, liveSnapshots.Count);
// Capture A's response. State has been mutated to ToolB but the per-invocation
// snapshot must still drive invocation with ToolA.
Assert.Equal(2, emittedRequests.Count);
ExternalInputResponse responseA = CreateApprovalResponseForRequest(emittedRequests[0], approved: true);
await action.CaptureResponseAsync(mockContext.Object, responseA, CancellationToken.None);
Assert.Single(liveSnapshots);
Assert.Equal([ToolA], capturedToolNames);
// Another checkpoint/reset/restore cycle - B's snapshot survives.
await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
await action.ResetAsync();
Assert.Empty(liveSnapshots);
await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
Assert.Single(liveSnapshots);
// Capture B's response.
ExternalInputResponse responseB = CreateApprovalResponseForRequest(emittedRequests[1], approved: true);
await action.CaptureResponseAsync(mockContext.Object, responseB, CancellationToken.None);
// Assert - both invocations executed with their own approved tool names; nothing pending.
Assert.Equal([ToolA, ToolB], capturedToolNames);
Assert.Empty(liveSnapshots);
}
private InvokeMcpTool CreateModelWithApproval(string displayName, string serverUrl, string toolName)
{
InvokeMcpTool.Builder builder = new()
{
Id = this.CreateActionId(),
DisplayName = this.FormatDisplayName(displayName),
ServerUrl = new StringExpression.Builder(StringExpression.Literal(serverUrl)),
ToolName = new StringExpression.Builder(StringExpression.Literal(toolName)),
RequireApproval = new BoolExpression.Builder(BoolExpression.Literal(true)),
};
return AssignParent<InvokeMcpTool>(builder);
}
private static Mock<IWorkflowContext> CreateMockWorkflowContext(List<ExternalInputRequest>? emittedRequests = null)
{
Mock<IWorkflowContext> mockContext = new();
mockContext.Setup(c => c.AddEventAsync(It.IsAny<WorkflowEvent>(), It.IsAny<CancellationToken>()))
@@ -1113,32 +1539,76 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
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>()))
.Callback<object, string?, CancellationToken>((msg, _, _) =>
{
if (emittedRequests is not null && msg is ExternalInputRequest request)
{
emittedRequests.Add(request);
}
})
.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 state store so callers can drive multi-step
/// checkpoint/reset/restore sequences against the same persisted state.
/// </summary>
private static Mock<IWorkflowContext> CreateMockWorkflowContextWithStateStore()
private static Mock<IWorkflowContext> CreateMockWorkflowContextWithStateStore(
List<ExternalInputRequest>? emittedRequests = null,
Dictionary<string, object?>? stateStore = null)
{
Dictionary<string, object?> stateStore = new();
stateStore ??= new Dictionary<string, object?>();
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)
mockContext.Setup(c => c.QueueStateUpdateAsync(It.IsAny<string>(), It.IsAny<Dictionary<string, ApprovalSnapshot>>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
.Callback<string, Dictionary<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>()))
.Callback<object, string?, CancellationToken>((msg, _, _) =>
{
if (emittedRequests is not null && msg is ExternalInputRequest request)
{
emittedRequests.Add(request);
}
})
.Returns(default(ValueTask));
mockContext.Setup(c => c.ReadStateAsync<ApprovalSnapshot>(It.IsAny<string>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
mockContext.Setup(c => c.ReadStateAsync<Dictionary<string, 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));
new ValueTask<Dictionary<string, ApprovalSnapshot>?>(stateStore.TryGetValue(key, out object? val) ? val as Dictionary<string, ApprovalSnapshot> : null));
mockContext.Setup(c => c.ReadStateKeysAsync(It.IsAny<string?>(), It.IsAny<CancellationToken>()))
.ReturnsAsync(new HashSet<string>());
return mockContext;
}
/// <summary>
/// Builds an approval response paired to the request id stamped on the emitted
/// <c>MCPToolApprovalRequestContent</c>. Mirrors the framework's symmetric
/// content-id rewriting at the envelope boundary.
/// </summary>
private static ExternalInputResponse CreateApprovalResponseFor(IReadOnlyList<ExternalInputRequest> emittedRequests, bool approved)
{
ExternalInputRequest emitted = Assert.Single(emittedRequests);
return CreateApprovalResponseForRequest(emitted, approved);
}
/// <summary>
/// Builds an approval response paired to the inner <c>ToolApprovalRequestContent.RequestId</c>
/// of a specific emitted request. Used when multiple requests are emitted and the
/// caller needs to address one by position.
/// </summary>
private static ExternalInputResponse CreateApprovalResponseForRequest(ExternalInputRequest emitted, bool approved)
{
ToolApprovalRequestContent approvalRequest = emitted.AgentResponse.Messages
.SelectMany(m => m.Contents)
.OfType<ToolApprovalRequestContent>()
.Single();
ToolApprovalResponseContent approvalResponse = approvalRequest.CreateResponse(approved);
return new ExternalInputResponse(new ChatMessage(ChatRole.User, [approvalResponse]));
}
/// <summary>
/// Invokes a protected method on an executor via reflection (for testing checkpoint hooks).
/// </summary>