Merge branch 'main' into local-branch-python-add-reset-to-workflow

dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs

@@ -62,7 +62,7 @@ internal abstract class DeclarativeActionExecutor : Executor<ActionExecutorResul
     protected virtual bool EmitResultEvent => true;
 
     /// <inheritdoc/>
-    public virtual ValueTask ResetAsync()
+    public ValueTask ResetAsync()
     {
         return default;
     }

dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeFunctionToolExecutor.cs

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text.Json;
@@ -29,23 +28,12 @@ internal sealed class InvokeFunctionToolExecutor(
     WorkflowFormulaState state) :
     DeclarativeActionExecutor<InvokeFunctionTool>(model, state)
 {
-    private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshots);
-    private const string PendingCallIdsStateKey = nameof(_pendingNonApprovalCallIds);
-    private const string LegacyApprovalSnapshotStateKey = "_approvalSnapshot";
+    private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshot);
 
     /// <summary>
-    /// 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.
+    /// Snapshot of evaluated parameters at approval-request time.
     /// </summary>
-    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);
+    private ApprovalSnapshot? _approvalSnapshot;
 
     /// <summary>
     /// Step identifiers for the function tool invocation workflow.
@@ -77,12 +65,9 @@ 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: requestId,
+            callId: this.Id,
             name: functionName,
             arguments: arguments);
 
@@ -92,15 +77,11 @@ internal sealed class InvokeFunctionToolExecutor(
         // If approval is required, add user input request content
         if (requireApproval)
         {
-            // Capture the evaluated parameters keyed by request id; the matching response
-            // resumes from this snapshot.
-            this._approvalSnapshots[requestId] = new ApprovalSnapshot(functionName, arguments);
+            // Snapshot the evaluated parameters.
+            // If state mutates during the approval window, the approved values are used on resume.
+            this._approvalSnapshot = new ApprovalSnapshot(functionName, arguments);
 
-            requestMessage.Contents.Add(new ToolApprovalRequestContent(requestId, functionCall));
-        }
-        else
-        {
-            this._pendingNonApprovalCallIds.TryAdd(requestId, 0);
+            requestMessage.Contents.Add(new ToolApprovalRequestContent(this.Id, functionCall));
         }
 
         AgentResponse agentResponse = new([requestMessage]);
@@ -127,24 +108,13 @@ internal sealed class InvokeFunctionToolExecutor(
         bool autoSend = this.GetAutoSendValue();
         string? conversationId = this.GetConversationId();
 
-        // Match the inbound result by its per-invocation call id.
-        FunctionResultContent? matchingResult = response.Messages
+        // Extract function results from the response
+        IEnumerable<FunctionResultContent> functionResults = response.Messages
             .SelectMany(m => m.Contents)
-            .OfType<FunctionResultContent>()
-            .FirstOrDefault(r => this.IsKnownPendingId(r.CallId));
+            .OfType<FunctionResultContent>();
 
-        // 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));
-        }
+        FunctionResultContent? matchingResult = functionResults
+            .FirstOrDefault(r => r.CallId == this.Id);
 
         // 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
@@ -153,42 +123,14 @@ internal sealed class InvokeFunctionToolExecutor(
         // SendActivity/PropertyPath consumers like {Local.Result}).
         if (matchingResult is null)
         {
-            List<ToolApprovalResponseContent> approvals = response.Messages
+            ToolApprovalResponseContent? approval = response.Messages
                 .SelectMany(m => m.Contents)
                 .OfType<ToolApprovalResponseContent>()
-                .ToList();
+                .FirstOrDefault(r => r.RequestId == this.Id);
 
-            // 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)
+            if (approval is { Approved: true })
             {
-                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, "No pending approval matched the response.").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
-                {
-                    this.Logger.LogWarning(
-                        "Approval response '{RequestId}' had no remaining pending snapshot on '{ActionId}'.",
-                        approval.RequestId, this.Id);
-                    await this.AssignErrorAsync(context, "No pending approval matched the response.").ConfigureAwait(false);
-                }
+                matchingResult = await this.InvokeRegisteredFunctionAsync(cancellationToken).ConfigureAwait(false);
             }
         }
 
@@ -203,10 +145,6 @@ 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
@@ -229,76 +167,31 @@ internal sealed class InvokeFunctionToolExecutor(
 
         // Completes the action after processing the function result.
         await context.RaiseCompletionEventAsync(this.Model, 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;
+        // 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);
     }
 
     /// <inheritdoc/>
     /// <remarks>
-    /// Persists pending approval snapshots and non-approval call ids so they survive
-    /// checkpoint/restore cycles.
+    /// Persists the approval snapshot to workflow state so it survives checkpoint/restore cycles.
     /// </remarks>
     protected override async ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
     {
-        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 context.QueueStateUpdateAsync(ApprovalSnapshotStateKey, this._approvalSnapshot, null, cancellationToken).ConfigureAwait(false);
         await base.OnCheckpointingAsync(context, cancellationToken).ConfigureAwait(false);
     }
 
     /// <inheritdoc/>
     /// <remarks>
-    /// Restores pending approval snapshots and non-approval call ids from workflow state
-    /// after a checkpoint restore.
+    /// Restores the approval snapshot 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._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);
-        }
+        this._approvalSnapshot = await context.ReadStateAsync<ApprovalSnapshot>(ApprovalSnapshotStateKey, null, cancellationToken).ConfigureAwait(false);
     }
 
     /// <summary>
@@ -387,14 +280,6 @@ 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(
@@ -412,19 +297,32 @@ internal sealed class InvokeFunctionToolExecutor(
         return conversationIdValue.Length == 0 ? null : conversationIdValue;
     }
 
-    private async ValueTask<FunctionResultContent?> InvokeRegisteredFunctionAsync(string callId, ApprovalSnapshot snapshot, CancellationToken cancellationToken)
+    private async ValueTask<FunctionResultContent?> InvokeRegisteredFunctionAsync(CancellationToken cancellationToken)
     {
-        // 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;
+        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();
+        }
 
         AIFunction? function = agentProvider.Functions?.FirstOrDefault(
             f => string.Equals(f.Name, functionName, StringComparison.Ordinal));
 
         if (function is null)
         {
-            return new FunctionResultContent(callId, result: null)
+            return new FunctionResultContent(this.Id, result: null)
             {
                 Exception = new InvalidOperationException(
                     $"Function '{functionName}' is not registered with the agent provider."),
@@ -440,7 +338,7 @@ internal sealed class InvokeFunctionToolExecutor(
         }
         catch (Exception ex) when (ex is not OperationCanceledException)
         {
-            return new FunctionResultContent(callId, result: null) { Exception = ex };
+            return new FunctionResultContent(this.Id, result: null) { Exception = ex };
         }
 
         // Match FunctionInvokingChatClient's serialization: pass strings through as-is and
@@ -454,7 +352,7 @@ internal sealed class InvokeFunctionToolExecutor(
             _ => JsonSerializer.Serialize(result, AIJsonUtilities.DefaultOptions.GetTypeInfo(result.GetType())),
         };
 
-        return new FunctionResultContent(callId, serialized);
+        return new FunctionResultContent(this.Id, serialized);
     }
 
     private bool GetRequireApproval()
@@ -498,8 +396,9 @@ internal sealed class InvokeFunctionToolExecutor(
     }
 
     /// <summary>
-    /// Captured invocation parameters used by <see cref="CaptureResponseAsync"/> on
-    /// resume so the approved values are invoked regardless of subsequent state changes.
+    /// 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.
     /// </summary>
     internal sealed record ApprovalSnapshot(
         string FunctionName,

dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs

@@ -1,7 +1,5 @@
 // Copyright (c) Microsoft. All rights reserved.
 
-using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text.Json;
@@ -29,15 +27,13 @@ internal sealed class InvokeMcpToolExecutor(
     WorkflowFormulaState state) :
     DeclarativeActionExecutor<InvokeMcpTool>(model, state)
 {
-    private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshots);
-    private const string LegacyApprovalSnapshotStateKey = "_approvalSnapshot";
+    private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshot);
 
     /// <summary>
-    /// 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.
+    /// Snapshot of evaluated parameters at approval-request time.
+    /// Used to prevent TOCTOU attacks where state mutates during the approval window.
     /// </summary>
-    private readonly ConcurrentDictionary<string, ApprovalSnapshot> _approvalSnapshots = new(StringComparer.Ordinal);
+    private ApprovalSnapshot? _approvalSnapshot;
 
     /// <summary>
     /// Step identifiers for the MCP tool invocation workflow.
@@ -87,22 +83,19 @@ internal sealed class InvokeMcpToolExecutor(
 
         if (requireApproval)
         {
-            // Per-invocation request id stamped on the outbound content.
-            string requestId = Guid.NewGuid().ToString("N");
-
-            // Capture the evaluated parameters keyed by request id; the matching response
-            // resumes from this snapshot.
-            this._approvalSnapshots[requestId] = new ApprovalSnapshot(serverUrl, serverLabel, toolName, arguments, connectionName);
+            // Snapshot the evaluated parameters to prevent TOCTOU attacks.
+            // If state mutates during the approval window, the approved values are used on resume.
+            this._approvalSnapshot = new ApprovalSnapshot(serverUrl, serverLabel, toolName, arguments, connectionName);
 
             // Create tool call content for approval request.
             // Transport headers (e.g. Authorization) are intentionally excluded from the
             // approval event: they must not cross into the externally-surfaced approval request.
-            McpServerToolCallContent toolCall = new(requestId, toolName, serverLabel ?? serverUrl)
+            McpServerToolCallContent toolCall = new(this.Id, toolName, serverLabel ?? serverUrl)
             {
                 Arguments = arguments
             };
 
-            ToolApprovalRequestContent approvalRequest = new(requestId, toolCall);
+            ToolApprovalRequestContent approvalRequest = new(this.Id, toolCall);
 
             ChatMessage requestMessage = new(ChatRole.Assistant, [approvalRequest]);
             AgentResponse agentResponse = new([requestMessage]);
@@ -147,38 +140,31 @@ internal sealed class InvokeMcpToolExecutor(
         ToolApprovalResponseContent? approvalResponse = response.Messages
             .SelectMany(m => m.Contents)
             .OfType<ToolApprovalResponseContent>()
-            .FirstOrDefault(r => this._approvalSnapshots.ContainsKey(r.RequestId));
+            .FirstOrDefault(r => r.RequestId == this.Id);
 
-        if (approvalResponse is null)
+        if (approvalResponse?.Approved != true)
         {
-            await this.AssignErrorAsync(context, "No pending approval matched the response.").ConfigureAwait(false);
-            return;
-        }
-
-        if (!approvalResponse.Approved)
-        {
-            this._approvalSnapshots.TryRemove(approvalResponse.RequestId, out _);
+            // Tool call was rejected
             await this.AssignErrorAsync(context, "MCP tool invocation was not approved by user.").ConfigureAwait(false);
             return;
         }
 
-        // Source invocation fields from the snapshot captured at approval-request time.
-        // Headers are re-evaluated (they may contain auth secrets not persisted to state).
-        if (!this._approvalSnapshots.TryRemove(approvalResponse.RequestId, out ApprovalSnapshot? snapshot))
-        {
-            await this.AssignErrorAsync(context, "No pending approval matched the response.").ConfigureAwait(false);
-            return;
-        }
-
+        // Approved - use the snapshot from approval-request time to prevent TOCTOU attacks.
+        // Headers are re-evaluated (they may contain auth secrets that should not be persisted).
+        string serverUrl = this._approvalSnapshot?.ServerUrl ?? this.GetServerUrl();
+        string? serverLabel = this._approvalSnapshot?.ServerLabel ?? this.GetServerLabel();
+        string toolName = this._approvalSnapshot?.ToolName ?? this.GetToolName();
+        Dictionary<string, object?>? arguments = this._approvalSnapshot?.Arguments ?? this.GetArguments();
         Dictionary<string, string>? headers = this.GetHeaders();
+        string? connectionName = this._approvalSnapshot?.ConnectionName ?? this.GetConnectionName();
 
         McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync(
-            snapshot.ServerUrl,
-            snapshot.ServerLabel,
-            snapshot.ToolName,
-            snapshot.Arguments,
+            serverUrl,
+            serverLabel,
+            toolName,
+            arguments,
             headers,
-            snapshot.ConnectionName,
+            connectionName,
             cancellationToken).ConfigureAwait(false);
 
         await this.ProcessResultAsync(context, resultContent, cancellationToken).ConfigureAwait(false);
@@ -189,57 +175,31 @@ internal sealed class InvokeMcpToolExecutor(
     /// </summary>
     public async ValueTask CompleteAsync(IWorkflowContext context, ActionExecutorResult message, CancellationToken cancellationToken)
     {
+        // Clear the approval snapshot after successful completion.
+        this._approvalSnapshot = null;
+        await ClearSnapshotStateAsync(context, cancellationToken).ConfigureAwait(false);
+
         await context.RaiseCompletionEventAsync(this.Model, cancellationToken).ConfigureAwait(false);
     }
 
-    /// <inheritdoc/>
-    public override ValueTask ResetAsync()
-    {
-        this._approvalSnapshots.Clear();
-        return default;
-    }
-
     /// <inheritdoc/>
     /// <remarks>
-    /// Persists pending approval snapshots to workflow state so they survive
-    /// checkpoint/restore cycles.
+    /// Persists the approval snapshot to workflow state so it survives checkpoint/restore cycles.
     /// </remarks>
     protected override async ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
     {
-        Dictionary<string, ApprovalSnapshot> snapshotCopy = this._approvalSnapshots.ToDictionary(kvp => kvp.Key, kvp => kvp.Value, StringComparer.Ordinal);
-        await context.QueueStateUpdateAsync(ApprovalSnapshotStateKey, snapshotCopy, null, cancellationToken).ConfigureAwait(false);
+        await context.QueueStateUpdateAsync(ApprovalSnapshotStateKey, this._approvalSnapshot, null, cancellationToken).ConfigureAwait(false);
         await base.OnCheckpointingAsync(context, cancellationToken).ConfigureAwait(false);
     }
 
     /// <inheritdoc/>
     /// <remarks>
-    /// Restores pending approval snapshots from workflow state after a checkpoint restore.
+    /// Restores the approval snapshot 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._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;
-            }
-        }
-
-        // 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);
-        }
+        this._approvalSnapshot = await context.ReadStateAsync<ApprovalSnapshot>(ApprovalSnapshotStateKey, null, cancellationToken).ConfigureAwait(false);
     }
 
     private async ValueTask ProcessResultAsync(IWorkflowContext context, McpServerToolResultContent resultContent, CancellationToken cancellationToken)
@@ -444,8 +404,17 @@ internal sealed class InvokeMcpToolExecutor(
     }
 
     /// <summary>
-    /// Captured invocation parameters used by <see cref="CaptureResponseAsync"/> on
-    /// resume so the approved values are invoked regardless of subsequent state changes.
+    /// Clears the persisted approval snapshot state after a successful tool invocation.
+    /// </summary>
+    private static async ValueTask ClearSnapshotStateAsync(IWorkflowContext context, CancellationToken cancellationToken)
+    {
+        await context.QueueStateUpdateAsync<ApprovalSnapshot?>(ApprovalSnapshotStateKey, null, null, cancellationToken).ConfigureAwait(false);
+    }
+
+    /// <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.
     /// </summary>
     internal sealed record ApprovalSnapshot(
         string ServerUrl,

dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/ObjectModel/InvokeFunctionToolExecutorTest.cs

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Reflection;
@@ -302,9 +301,8 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
             onInvoke: name => capturedFunctionName = name);
         InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
 
-        // Act - trigger ExecuteAsync to emit the approval request
-        List<ExternalInputRequest> emittedRequests = [];
-        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
+        // Act - trigger ExecuteAsync to store the approval snapshot
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
         await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
 
         // Simulate parallel branch mutating state during the approval window
@@ -312,7 +310,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 = CreateApprovalResponseFor(emittedRequests, approved: true);
+        ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
 
         // Resume after approval
         await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -351,9 +349,8 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
             onInvokeArguments: args => capturedArguments = args);
         InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
 
-        // Act - trigger ExecuteAsync to emit the approval request
-        List<ExternalInputRequest> emittedRequests = [];
-        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
+        // Act - trigger ExecuteAsync to store the approval snapshot
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
         await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
 
         // Simulate parallel branch mutating state during the approval window
@@ -361,7 +358,7 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
         this.State.Bind();
 
         // User clicks approve
-        ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
+        ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
 
         // Resume after approval
         await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -399,20 +396,18 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
             onInvoke: name => capturedFunctionName = name);
         InvokeFunctionToolExecutor action = new(model, testAgentProvider, this.State);
 
-        // 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);
+        // Act - trigger ExecuteAsync to store the approval snapshot
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore();
         await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
 
         // Simulate checkpoint: persist to state store
         await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
 
-        // Simulate restore on a "new" executor instance by clearing the in-memory dictionary via reflection
-        ConcurrentDictionary<string, ApprovalSnapshot> liveSnapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeFunctionToolExecutor)
-            .GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
-            .GetValue(action)!;
-        liveSnapshots.Clear();
+        // Simulate restore on a "new" executor instance by clearing the in-memory field via reflection
+        // (In production, a new executor instance would be created with _approvalSnapshot == null)
+        typeof(InvokeFunctionToolExecutor)
+            .GetField("_approvalSnapshot", BindingFlags.NonPublic | BindingFlags.Instance)!
+            .SetValue(action, null);
 
         // Restore from state store
         await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
@@ -422,7 +417,7 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
         this.State.Bind();
 
         // User clicks approve
-        ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
+        ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
 
         // Resume after approval
         await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -433,7 +428,9 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
     }
 
     /// <summary>
-    /// Verifies that the approval snapshot entry is removed after a completed approval cycle.
+    /// Verifies that the approval snapshot is cleared after a completed approval cycle,
+    /// both in-memory and in the persisted state store. This prevents stale data from
+    /// influencing a subsequent execution of the same executor instance.
     /// </summary>
     [Fact]
     public async Task InvokeFunctionToolCaptureResponseClearsSnapshotAfterCompletionAsync()
@@ -454,723 +451,33 @@ 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, emittedRequests);
-        await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
-
-        // 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 = CreateApprovalResponseFor(emittedRequests, approved: true);
-        await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
-
-        // Assert - in-memory dict is empty after the matching response is captured
-        Assert.Empty(snapshots);
-    }
-
-    /// <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()
-    {
-        // 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("No pending approval"));
-    }
-
-    /// <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 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.StartsWith("Error:", StringComparison.Ordinal));
-    }
-
-    /// <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);
+        // Sanity: snapshot was captured
+        FieldInfo snapshotField = typeof(InvokeFunctionToolExecutor)
+            .GetField("_approvalSnapshot", BindingFlags.NonPublic | BindingFlags.Instance)!;
+        Assert.NotNull(snapshotField.GetValue(action));
+
+        ExternalInputResponse response = CreateApprovalResponse(action.Id, approved: true);
         await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
 
-        // Assert - the registered AIFunction was invoked exactly once.
-        Assert.Equal(1, invocationCount);
-        // The second delivery surfaced the no-pending-approval error.
-        Assert.Contains(mockContext.Invocations, i =>
-            i.Method.Name == nameof(IWorkflowContext.QueueStateUpdateAsync)
-            && i.Arguments.Count >= 2
-            && i.Arguments[1] is StringValue sv
-            && sv.Value.Contains("No pending approval"));
+        // Assert - both in-memory field and persisted state are cleared
+        Assert.Null(snapshotField.GetValue(action));
+        Assert.True(stateStore.ContainsKey("_approvalSnapshot"));
+        Assert.Null(stateStore["_approvalSnapshot"]);
     }
 
-    /// <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()
+    private static ExternalInputResponse CreateApprovalResponse(string actionId, bool approved)
     {
-        // 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
-        // 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.StartsWith("Error:", StringComparison.Ordinal));
-    }
-
-    /// <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();
+        FunctionCallContent functionCall = new(callId: actionId, name: "ignored");
+        ToolApprovalRequestContent approvalRequest = new(actionId, functionCall);
         ToolApprovalResponseContent approvalResponse = approvalRequest.CreateResponse(approved);
         return new ExternalInputResponse(new ChatMessage(ChatRole.User, [approvalResponse]));
     }
 
-    /// <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)
+    private static Mock<IWorkflowContext> CreateMockWorkflowContext()
     {
         Mock<IWorkflowContext> mockContext = new();
         mockContext.Setup(c => c.AddEventAsync(It.IsAny<WorkflowEvent>(), It.IsAny<CancellationToken>()))
@@ -1178,64 +485,25 @@ 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,
-    /// and an optional emitted-request list so tests can build matching responses.
+    /// Optionally accepts an externally-owned dictionary so callers can inspect the persisted state.
     /// </summary>
-    private static Mock<IWorkflowContext> CreateMockWorkflowContextWithStateStore(
-        Dictionary<string, object?>? stateStore = null,
-        List<ExternalInputRequest>? emittedRequests = null)
+    private static Mock<IWorkflowContext> CreateMockWorkflowContextWithStateStore(Dictionary<string, object?>? stateStore = null)
     {
         stateStore ??= [];
         Mock<IWorkflowContext> mockContext = new();
         mockContext.Setup(c => c.AddEventAsync(It.IsAny<WorkflowEvent>(), It.IsAny<CancellationToken>()))
             .Returns(default(ValueTask));
-        mockContext.Setup(c => c.QueueStateUpdateAsync(It.IsAny<string>(), It.IsAny<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, _, _) =>
-            {
-                if (value is null)
-                {
-                    stateStore.Remove(key);
-                }
-                else
-                {
-                    stateStore[key] = value;
-                }
-            })
+            .Callback<string, ApprovalSnapshot?, string?, CancellationToken>((key, value, _, _) => stateStore[key] = value)
             .Returns(default(ValueTask));
         mockContext.Setup(c => c.SendMessageAsync(It.IsAny<object>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
-            .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));
@@ -1354,8 +622,7 @@ public sealed class InvokeFunctionToolExecutorTest(ITestOutputHelper output) : W
         bool? requireApproval = false,
         string? conversationId = null,
         string? argumentKey = null,
-        string? argumentValue = null,
-        string? outputResultVariable = null)
+        string? argumentValue = null)
     {
         InvokeFunctionTool.Builder builder = new()
         {
@@ -1375,14 +642,6 @@ 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

@@ -1,6 +1,5 @@
 // Copyright (c) Microsoft. All rights reserved.
 
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Reflection;
@@ -424,15 +423,15 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
         MockAgentProvider mockAgentProvider = new();
         InvokeMcpToolExecutor action = new(model, mockProvider.Object, mockAgentProvider.Object, this.State);
 
-        // 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);
+        Mock<IWorkflowContext> mockContext = new(MockBehavior.Loose);
 
-        // Build the matching approved response from the emitted request.
-        ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
+        // 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]));
 
-        // Act - call CaptureResponseAsync so the post-approval branch actually executes.
+        // Act - call CaptureResponseAsync directly 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.
@@ -888,8 +887,7 @@ 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
-        List<ExternalInputRequest> emittedRequests = [];
-        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
         await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
 
         // Simulate parallel branch mutating state during the approval window
@@ -897,7 +895,10 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
         this.State.Bind();
 
         // User clicks approve (they saw "safe_readonly_query" in the approval UI)
-        ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
+        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]));
 
         // Resume after approval
         await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -949,8 +950,7 @@ 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
-        List<ExternalInputRequest> emittedRequests = [];
-        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
         await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
 
         // Simulate parallel branch mutating state during the approval window
@@ -958,7 +958,10 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
         this.State.Bind();
 
         // User clicks approve
-        ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
+        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]));
 
         // Resume after approval
         await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -1008,8 +1011,7 @@ 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
-        List<ExternalInputRequest> emittedRequests = [];
-        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext(emittedRequests);
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContext();
         await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
 
         // Simulate parallel branch mutating state during the approval window
@@ -1017,7 +1019,10 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
         this.State.Bind();
 
         // User clicks approve
-        ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
+        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]));
 
         // Resume after approval
         await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -1067,18 +1072,17 @@ 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
-        List<ExternalInputRequest> emittedRequests = [];
-        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore(emittedRequests);
+        Mock<IWorkflowContext> mockContext = CreateMockWorkflowContextWithStateStore();
         await action.HandleAsync(new ActionExecutorResult(action.Id), mockContext.Object, CancellationToken.None);
 
         // Simulate checkpoint: persist to state store
         await InvokeProtectedMethodAsync(action, "OnCheckpointingAsync", mockContext.Object, CancellationToken.None);
 
-        // Simulate restore on a "new" executor instance by clearing the in-memory dictionary via reflection
-        ConcurrentDictionary<string, ApprovalSnapshot> liveSnapshots = (ConcurrentDictionary<string, ApprovalSnapshot>)typeof(InvokeMcpToolExecutor)
-            .GetField("_approvalSnapshots", BindingFlags.NonPublic | BindingFlags.Instance)!
-            .GetValue(action)!;
-        liveSnapshots.Clear();
+        // 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);
 
         // Restore from state store
         await InvokeProtectedMethodAsync(action, "OnCheckpointRestoredAsync", mockContext.Object, CancellationToken.None);
@@ -1088,7 +1092,10 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
         this.State.Bind();
 
         // User clicks approve
-        ExternalInputResponse response = CreateApprovalResponseFor(emittedRequests, approved: true);
+        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]));
 
         // Resume after approval
         await action.CaptureResponseAsync(mockContext.Object, response, CancellationToken.None);
@@ -1098,440 +1105,7 @@ public sealed class InvokeMcpToolExecutorTest(ITestOutputHelper output) : Workfl
         Assert.Equal(ApprovedToolName, capturedToolName);
     }
 
-    /// <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)
+    private static Mock<IWorkflowContext> CreateMockWorkflowContext()
     {
         Mock<IWorkflowContext> mockContext = new();
         mockContext.Setup(c => c.AddEventAsync(It.IsAny<WorkflowEvent>(), It.IsAny<CancellationToken>()))
@@ -1539,75 +1113,32 @@ 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(
-        List<ExternalInputRequest>? emittedRequests = null,
-        Dictionary<string, object?>? stateStore = null)
+    private static Mock<IWorkflowContext> CreateMockWorkflowContextWithStateStore()
     {
-        stateStore ??= new Dictionary<string, object?>();
+        Dictionary<string, object?> stateStore = new();
         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)
+        mockContext.Setup(c => c.QueueStateUpdateAsync(It.IsAny<string>(), It.IsAny<ApprovalSnapshot?>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
+            .Callback<string, ApprovalSnapshot?, string?, CancellationToken>((key, value, _, _) => stateStore[key] = value)
             .Returns(default(ValueTask));
         mockContext.Setup(c => c.SendMessageAsync(It.IsAny<object>(), It.IsAny<string?>(), It.IsAny<CancellationToken>()))
-            .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>()))
+        mockContext.Setup(c => c.ReadStateAsync<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));
+                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>
-    /// Builds an approval response paired to the request id stamped on the emitted
-    /// <c>ToolApprovalRequestContent</c>.
-    /// </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>

python/packages/core/agent_framework/__init__.py

@@ -258,6 +258,7 @@ from ._workflows._agent_executor import (
 )
 from ._workflows._agent_utils import resolve_agent_id
 from ._workflows._checkpoint import (
+    CheckpointID,
     CheckpointStorage,
     FileCheckpointStorage,
     InMemoryCheckpointStorage,
@@ -301,7 +302,6 @@ from ._workflows._functional import (
     workflow,
 )
 from ._workflows._request_info_mixin import response_handler
-from ._workflows._runner import Runner
 from ._workflows._runner_context import (
     InProcRunnerContext,
     RunnerContext,
@@ -397,6 +397,7 @@ __all__ = [
     "ChatResponse",
     "ChatResponseUpdate",
     "CheckResult",
+    "CheckpointID",
     "CheckpointStorage",
     "ClassSkill",
     "CompactionProvider",
@@ -489,7 +490,6 @@ __all__ = [
     "RoleLiteral",
     "RubricScore",
     "RunContext",
-    "Runner",
     "RunnerContext",
     "SamplingApprovalCallback",
     "SecretString",

python/packages/core/agent_framework/_workflows/_runner.py

@@ -10,7 +10,6 @@ from typing import Any
 from ..exceptions import (
     WorkflowCheckpointException,
     WorkflowConvergenceException,
-    WorkflowRunnerException,
 )
 from ._checkpoint import CheckpointID, CheckpointStorage, WorkflowCheckpoint
 from ._const import EXECUTOR_STATE_KEY
@@ -63,99 +62,105 @@ class Runner:
         self._iteration = 0
         self._max_iterations = max_iterations
         self._state = state
-        self._running = False
-        self._resumed_from_checkpoint = False  # Track whether we resumed
+
+        # Checkpointing related attributes
+        self._resumed_from_checkpoint = False
+        self.previous_checkpoint_id: CheckpointID | None = None
 
     @property
     def context(self) -> RunnerContext:
-        """Get the workflow context."""
+        """Get the runner context for message, event, and checkpoint handling."""
         return self._ctx
 
+    @property
+    def state(self) -> State:
+        """Get the shared state for the workflow."""
+        return self._state
+
     def reset_iteration_count(self) -> None:
-        """Reset the iteration count to zero."""
+        """Reset the iteration count to zero.
+
+        This is useful when the workflow resumes from a new set of messages.
+
+        Note:
+            When a workflow is resumed from a response (for a request_info_event)
+            or a checkpoint, the iteration count is normally NOT reset.
+        """
         self._iteration = 0
 
     async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
         """Run the workflow until no more messages are sent."""
-        if self._running:
-            raise WorkflowRunnerException("Runner is already running.")
+        # Emit any events already produced prior to entering loop
+        if await self._ctx.has_events():
+            logger.info("Yielding pre-loop events")
+            for event in await self._ctx.drain_events():
+                yield event
 
-        self._running = True
-        previous_checkpoint_id: CheckpointID | None = None
-        try:
-            # Emit any events already produced prior to entering loop
-            if await self._ctx.has_events():
-                logger.info("Yielding pre-loop events")
-                for event in await self._ctx.drain_events():
-                    yield event
+        # Create a checkpoint before a run starts. Checkpoints are usually considered to be created at the
+        # end of an iteration, we can think of this checkpoint as being created at the end of "superstep 0"
+        # which captures the states after which the start executor has run. Note that we execute the start
+        # executor outside of the main iteration loop.
+        if await self._ctx.has_messages() and not self._resumed_from_checkpoint:
+            await self.create_checkpoint_if_enabled()
 
-            # Create the first checkpoint. Checkpoints are usually considered to be created at the end of an iteration,
-            # we can think of the first checkpoint as being created at the end of a "superstep 0" which captures the
-            # states after which the start executor has run.  Note that we execute the start executor outside of the
-            # main iteration loop.
-            if await self._ctx.has_messages() and not self._resumed_from_checkpoint:
-                previous_checkpoint_id = await self._create_checkpoint_if_enabled(previous_checkpoint_id)
+        while self._iteration < self._max_iterations:
+            logger.info(f"Starting superstep {self._iteration + 1}")
+            yield WorkflowEvent.superstep_started(iteration=self._iteration + 1)
 
-            while self._iteration < self._max_iterations:
-                logger.info(f"Starting superstep {self._iteration + 1}")
-                yield WorkflowEvent.superstep_started(iteration=self._iteration + 1)
-
-                # Run iteration concurrently with live event streaming: we poll
-                # for new events while the iteration coroutine progresses.
-                iteration_task = asyncio.create_task(self._run_iteration())
-                try:
-                    while not iteration_task.done():
-                        try:
-                            # Wait briefly for any new event; timeout allows progress checks
-                            event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
-                            yield event
-                        except asyncio.TimeoutError:
-                            # Periodically continue to let iteration advance
-                            continue
-                except asyncio.CancelledError:
-                    # Propagate cancellation to the iteration task to avoid orphaned work
-                    iteration_task.cancel()
-                    with contextlib.suppress(asyncio.CancelledError):
-                        await iteration_task
-                    raise
-
-                # Propagate errors from iteration, but first surface any pending events
-                try:
+            # Run iteration concurrently with live event streaming: we poll
+            # for new events while the iteration coroutine progresses.
+            iteration_task = asyncio.create_task(self._run_iteration())
+            try:
+                while not iteration_task.done():
+                    try:
+                        # Wait briefly for any new event; timeout allows progress checks
+                        event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
+                        yield event
+                    except asyncio.TimeoutError:
+                        # Periodically continue to let iteration advance
+                        continue
+            except asyncio.CancelledError:
+                # Propagate cancellation to the iteration task to avoid orphaned work
+                iteration_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
                     await iteration_task
-                except Exception:
-                    # Make sure failure-related events (like ExecutorFailedEvent) are surfaced
-                    if await self._ctx.has_events():
-                        for event in await self._ctx.drain_events():
-                            yield event
-                    raise
-                self._iteration += 1
+                raise
 
-                # Drain any straggler events emitted at tail end
+            # Propagate errors from iteration, but first surface any pending events
+            try:
+                await iteration_task
+            except Exception:
+                # Make sure failure-related events (like ExecutorFailedEvent) are surfaced
                 if await self._ctx.has_events():
                     for event in await self._ctx.drain_events():
                         yield event
+                raise
+            self._iteration += 1
 
-                logger.info(f"Completed superstep {self._iteration}")
+            # Drain any straggler events emitted at tail end
+            if await self._ctx.has_events():
+                for event in await self._ctx.drain_events():
+                    yield event
 
-                # Commit pending state changes at superstep boundary
-                self._state.commit()
+            logger.info(f"Completed superstep {self._iteration}")
 
-                # Create checkpoint after each superstep iteration
-                previous_checkpoint_id = await self._create_checkpoint_if_enabled(previous_checkpoint_id)
+            # Commit pending state changes at superstep boundary
+            self._state.commit()
 
-                yield WorkflowEvent.superstep_completed(iteration=self._iteration)
+            # Create checkpoint after each superstep iteration
+            await self.create_checkpoint_if_enabled()
 
-                # Check for convergence: no more messages to process
-                if not await self._ctx.has_messages():
-                    break
+            yield WorkflowEvent.superstep_completed(iteration=self._iteration)
 
-            if self._iteration >= self._max_iterations and await self._ctx.has_messages():
-                raise WorkflowConvergenceException(f"Runner did not converge after {self._max_iterations} iterations.")
+            # Check for convergence: no more messages to process
+            if not await self._ctx.has_messages():
+                break
 
-            logger.info(f"Workflow completed after {self._iteration} supersteps")
-            self._resumed_from_checkpoint = False  # Reset resume flag for next run
-        finally:
-            self._running = False
+        logger.info(f"Workflow completed after {self._iteration} supersteps")
+        self._resumed_from_checkpoint = False  # Reset resume flag for next run
+
+        if self._iteration >= self._max_iterations and await self._ctx.has_messages():
+            raise WorkflowConvergenceException(f"Runner did not converge after {self._max_iterations} iterations.")
 
     async def _run_iteration(self) -> None:
         """Run a single iteration of the workflow.
@@ -209,10 +214,10 @@ class Runner:
         ]
         await asyncio.gather(*tasks)
 
-    async def _create_checkpoint_if_enabled(self, previous_checkpoint_id: CheckpointID | None) -> CheckpointID | None:
+    async def create_checkpoint_if_enabled(self) -> None:
         """Create a checkpoint if checkpointing is enabled and attach a label and metadata."""
         if not self._ctx.has_checkpointing():
-            return None
+            return
 
         try:
             # Save executor states into the shared state before creating the checkpoint,
@@ -227,22 +232,33 @@ class Runner:
                 self._workflow_name,
                 self._graph_signature_hash,
                 self._state,
-                previous_checkpoint_id,
+                self.previous_checkpoint_id,
                 self._iteration,
             )
 
-            logger.info(f"Created checkpoint: {checkpoint_id}")
-            return checkpoint_id
+            logger.info(
+                "Created checkpoint: %s with parent checkpoint at iteration %d: %s",
+                checkpoint_id,
+                self._iteration,
+                self.previous_checkpoint_id,
+            )
+            self.previous_checkpoint_id = checkpoint_id
         except Exception as e:
-            logger.warning(f"Failed to create checkpoint: {e}")
-            return None
+            logger.warning(
+                "Failed to create checkpoint at iteration %d: %s. "
+                "Note that this does not fail the workflow run. "
+                "The next successfully-created checkpoint will be parented to the last successful checkpoint: %s",
+                self._iteration,
+                e,
+                self.previous_checkpoint_id,
+            )
 
     async def restore_from_checkpoint(
         self,
         checkpoint_id: CheckpointID,
         checkpoint_storage: CheckpointStorage | None = None,
     ) -> None:
-        """Restore workflow state from a checkpoint.
+        """Restore the runner from a checkpoint.
 
         Args:
             checkpoint_id: The ID of the checkpoint to restore from
@@ -290,7 +306,7 @@ class Runner:
             # Apply the checkpoint to the context
             await self._ctx.apply_checkpoint(checkpoint)
             # Mark the runner as resumed
-            self._mark_resumed(checkpoint.iteration_count)
+            self._mark_resumed(checkpoint)
 
             logger.info(f"Successfully restored workflow from checkpoint: {checkpoint_id}")
         except WorkflowCheckpointException:
@@ -356,13 +372,14 @@ class Runner:
 
         return parsed
 
-    def _mark_resumed(self, iteration: int) -> None:
+    def _mark_resumed(self, checkpoint: WorkflowCheckpoint) -> None:
         """Mark the runner as having resumed from a checkpoint.
 
         Optionally set the current iteration and max iterations.
         """
         self._resumed_from_checkpoint = True
-        self._iteration = iteration
+        self._iteration = checkpoint.iteration_count
+        self.previous_checkpoint_id = checkpoint.checkpoint_id
 
     async def _set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:
         """Store executor state in state under a reserved key.

python/packages/core/agent_framework/_workflows/_runner_context.py

@@ -403,12 +403,14 @@ class InProcRunnerContext:
     def reset_for_new_run(self) -> None:
         """Reset the context for a new workflow run.
 
-        This clears messages, events, and resets streaming flag.
-        Runtime checkpoint storage is NOT cleared here as it's managed at the workflow level.
+        Clears messages, the pending event queue, the pending request_info
+        correlation map, and the streaming flag. Runtime checkpoint storage is
+        NOT cleared here as it's managed at the workflow level.
         """
         self._messages.clear()
         # Clear any pending events (best-effort) by recreating the queue
         self._event_queue = asyncio.Queue()
+        self._pending_request_info_events.clear()
         self._streaming = False  # Reset streaming flag
 
     async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:

python/packages/core/agent_framework/_workflows/_workflow.py

@@ -11,14 +11,16 @@ import logging
 import types
 import uuid
 import warnings
+import weakref
 from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Literal, overload
 
 from .._sessions import ContextProvider
 from .._types import ResponseStream
+from ..exceptions import WorkflowCheckpointException, WorkflowException
 from ..observability import OtelAttr, capture_exception, create_workflow_span
-from ._checkpoint import CheckpointStorage
+from ._checkpoint import CheckpointID, CheckpointStorage
 from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY
 from ._edge import (
     EdgeGroup,
@@ -346,25 +348,29 @@ class Workflow(DictConvertible):
         # Store non-serializable runtime objects as private attributes
         self._runner_context = runner_context
         self._runner_context.set_yield_output_classifier(self._output_designation.classify)
-        self._state = State()
         self._runner: Runner = Runner(
             self.edge_groups,
             self.executors,
-            self._state,
+            State(),
             runner_context,
             self.name,
             self.graph_signature_hash,
             max_iterations=max_iterations,
         )
 
-        # Flag to prevent concurrent workflow executions
-        self._is_running = False
-
         # Current run-level status of this workflow instance. Updated in lockstep with
         # the status events emitted from `_run_workflow_with_tracing`. Defaults to IDLE
         # for a freshly built workflow that has not yet been run.
         self._status: WorkflowRunState = WorkflowRunState.IDLE
 
+        # Weak reference to the in-flight run's ``ResponseStream``. Used as the single
+        # concurrency lock: if the previous stream is still alive, ``run()`` rejects a
+        # new run synchronously (before any await). When the stream is fully consumed
+        # ``_run_core``'s finally clears this; if the caller drops the stream without
+        # ever iterating, the weakref dereferences to ``None`` once Python collects it,
+        # so a subsequent ``run()`` is allowed.
+        self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None
+
     @property
     def status(self) -> WorkflowRunState:
         """Return the current run-level status of this workflow instance.
@@ -376,16 +382,6 @@ class Workflow(DictConvertible):
         """
         return self._status
 
-    def _ensure_not_running(self) -> None:
-        """Ensure the workflow is not already running."""
-        if self._is_running:
-            raise RuntimeError("Workflow is already running. Concurrent executions are not allowed.")
-        self._is_running = True
-
-    def _reset_running_flag(self) -> None:
-        """Reset the running flag."""
-        self._is_running = False
-
     def to_dict(self) -> dict[str, Any]:
         """Serialize the workflow definition into a JSON-ready dictionary."""
         data: dict[str, Any] = {
@@ -478,6 +474,50 @@ class Workflow(DictConvertible):
         """Get the list of executors in the workflow."""
         return list(self.executors.values())
 
+    async def create_checkpoint(self, checkpoint_storage: CheckpointStorage | None) -> CheckpointID:
+        """Create a checkpoint of the current workflow state in the provided storage.
+
+        Args:
+            checkpoint_storage: The CheckpointStorage instance where the checkpoint will be stored.
+                If None, will use the workflow's default checkpoint storage if configured, or raise
+                if checkpointing is not enabled.
+
+        Notes:
+            - Checkpoints can only be created when the workflow is idle (not actively running).
+            - Checkpoints are automatically created at the end of each superstep if a checkpoint storage is configured.
+              Use this method only when necessary, for example to capture the initial state of the workflow prior to the
+              first run.
+            - Creating a checkpoint manually will alter the checkpoint lineage. The new checkpoint will become the
+              parent of the next checkpoint created automatically (if checkpointing is enabled by providing a storage).
+        """
+        if self._is_run_active():
+            raise WorkflowException(
+                "Cannot create checkpoint while a workflow run is active. "
+                "Checkpointing is only allowed between runs when the workflow is idle."
+            )
+
+        if checkpoint_storage is None and not self._runner.context.has_checkpointing():
+            raise WorkflowCheckpointException(
+                "Checkpoint storage must be provided to create a checkpoint when checkpointing is not enabled."
+            )
+        if checkpoint_storage is not None:
+            self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage)
+
+        # Capture the runner's checkpoint id before attempting to save. The runner
+        # log-and-swallows storage save errors and only updates
+        # ``previous_checkpoint_id`` on success, so a failed save would otherwise
+        # leave the prior id in place and we'd return it as if a fresh checkpoint
+        # had been created.
+        previous_id_before = self._runner.previous_checkpoint_id
+        try:
+            await self._runner.create_checkpoint_if_enabled()
+            new_id = self._runner.previous_checkpoint_id
+            if new_id is None or new_id == previous_id_before:
+                raise WorkflowCheckpointException("Failed to create checkpoint.")
+            return new_id
+        finally:
+            self._runner.context.clear_runtime_checkpoint_storage()
+
     async def _run_workflow_with_tracing(
         self,
         initial_executor_fn: Callable[[], Awaitable[None]] | None = None,
@@ -535,13 +575,12 @@ class Workflow(DictConvertible):
                 yield in_progress  # noqa: RUF070
 
                 # Per-run reset for fresh-message runs only. We deliberately
-                # do NOT clear shared workflow state (`_state.clear()`) or the
-                # runner context's in-flight messages (`reset_for_new_run()`)
-                # here - state and pending work persist across `run()` calls
-                # so that a `WorkflowAgent` can deliver multi-turn input on
-                # the same instance and have prior turns' context survive.
-                # Iteration counting and per-run kwargs ARE per-run though,
-                # so they're reset here.
+                # do NOT clear shared workflow state or the runner context's
+                # in-flight messages here - state and pending work persist
+                # across `run()` calls so that a `WorkflowAgent` can deliver
+                # multi-turn input on the same instance and have prior turns'
+                # context survive. Iteration counting and per-run kwargs ARE
+                # per-run though, so they're reset here.
                 if not is_continuation:
                     self._runner.reset_iteration_count()
 
@@ -564,14 +603,13 @@ class Workflow(DictConvertible):
                         combined_kwargs["client_kwargs"] = self._resolve_invocation_kwargs(
                             client_kwargs, "client_kwargs"
                         )
-                    self._state.set(WORKFLOW_RUN_KWARGS_KEY, combined_kwargs)
+                    self._runner.state.set(WORKFLOW_RUN_KWARGS_KEY, combined_kwargs)
                 elif not is_continuation:
-                    self._state.set(WORKFLOW_RUN_KWARGS_KEY, {})
-                self._state.commit()  # Commit immediately so kwargs are available
+                    self._runner.state.set(WORKFLOW_RUN_KWARGS_KEY, {})
+                self._runner.state.commit()  # Commit immediately so kwargs are available
 
-                # Set streaming mode (always set explicitly per run since
-                # reset_for_new_run() no longer runs to clear it).
-                self._runner_context.set_streaming(streaming)
+                # Explicitly set streaming mode per run
+                self._runner.context.set_streaming(streaming)
 
                 # Execute initial setup if provided
                 if initial_executor_fn:
@@ -665,7 +703,7 @@ class Workflow(DictConvertible):
             await executor.execute(
                 message,
                 [self.__class__.__name__],
-                self._state,
+                self._runner.state,
                 self._runner.context,
                 trace_contexts=None,
                 source_span_ids=None,
@@ -745,9 +783,22 @@ class Workflow(DictConvertible):
         Raises:
             ValueError: If parameter combination is invalid.
         """
-        # Validate parameters and set running flag eagerly (before any async work)
+        # Validate parameters first so misuse fails before we touch any run state.
         self._validate_run_params(message, responses, checkpoint_id)
-        self._ensure_not_running()
+
+        # Concurrency check: reject a second run synchronously - before constructing
+        # the ResponseStream or yielding control to the event loop - so a concurrent
+        # ``run`` call can't slip past the guard while the first call is suspended
+        # inside its async generator. The ``ResponseStream`` returned below is the
+        # lock: as long as the caller holds a reference to it, ``self._active_run()``
+        # resolves to a live object and a new ``run`` is rejected. When the stream is
+        # fully consumed, ``_run_core``'s finally clears the attribute. When the
+        # caller drops the stream without iterating, garbage collection invalidates
+        # the weakref, so a subsequent ``run`` is permitted.
+        if self._is_run_active():
+            raise WorkflowException(
+                "Workflow is already running; concurrent runs are not allowed on the same instance."
+            )
 
         response_stream = ResponseStream[WorkflowEvent, WorkflowRunResult](
             self._run_core(
@@ -760,10 +811,8 @@ class Workflow(DictConvertible):
                 client_kwargs=client_kwargs,
             ),
             finalizer=functools.partial(self._finalize_events, include_status_events=include_status_events),
-            cleanup_hooks=[
-                functools.partial(self._run_cleanup, checkpoint_storage),
-            ],
         )
+        self._active_run = weakref.ref(response_stream)
 
         if stream:
             return response_stream
@@ -789,51 +838,67 @@ class Workflow(DictConvertible):
         if checkpoint_storage is not None:
             self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage)
 
-        # Async validation: a fresh-message run is only allowed when the
-        # runner context has fully drained from any prior run. If it still
-        # has in-flight executor messages, the prior run didn't complete -
-        # the caller must either resume from a checkpoint or wait for the
-        # prior run to drain. (Pending request_info events are intentionally
-        # NOT blocked here: a follow-up run with message=... is the normal
-        # way to deliver a response to those pending requests, e.g. via
-        # WorkflowAgent._process_pending_requests.)
-        # NOTE: _validate_run_params already enforces that ``message`` is
-        # mutually exclusive with both ``checkpoint_id`` and ``responses``,
-        # so we don't need to re-check those here.
-        if message is not None and await self._runner.context.has_messages():
-            raise RuntimeError(
-                "Cannot start a new run with 'message' while in-flight executor "
-                "messages remain from a prior run. Resume from a checkpoint "
-                "(checkpoint_id=...) or wait for the prior run to complete. "
-                "Workflows that need to recover from a mid-run failure must use "
-                "checkpointing; there is no in-process recovery path."
-            )
+        # Capture the weakref instance ``run()`` installed for *this* run. We
+        # compare by object identity in the finally so a stale finalizer (e.g.
+        # the caller dropped this stream after partial iteration, then started
+        # a new run before async-gen finalization throws ``GeneratorExit`` into
+        # us) does not clobber a successor run's freshly installed weakref.
+        # ``run()`` runs synchronously and assigns ``self._active_run`` before
+        # this generator's body is first iterated, so by the time we read it
+        # here it already points at our own ``ResponseStream``.
+        my_active_run = self._active_run
 
-        initial_executor_fn = self._resolve_execution_mode(message, responses, checkpoint_id, checkpoint_storage)
+        try:
+            # Async validation: a fresh-message run is only allowed when the
+            # runner context has fully drained from any prior run. If it still
+            # has in-flight executor messages, the prior run didn't complete -
+            # the caller must either resume from a checkpoint or wait for the
+            # prior run to drain. (Pending request_info events are intentionally
+            # NOT blocked here: a follow-up run with message=... is the normal
+            # way to deliver a response to those pending requests, e.g. via
+            # WorkflowAgent._process_pending_requests.)
+            # NOTE: _validate_run_params already enforces that ``message`` is
+            # mutually exclusive with both ``checkpoint_id`` and ``responses``,
+            # so we don't need to re-check those here.
+            if message is not None and await self._runner.context.has_messages():
+                raise RuntimeError(
+                    "Cannot start a new run with 'message' while in-flight executor "
+                    "messages remain from a prior run. Resume from a checkpoint "
+                    "(checkpoint_id=...) or wait for the prior run to complete. "
+                    "Workflows that need to recover from a mid-run failure must use "
+                    "checkpointing; there is no in-process recovery path."
+                )
 
-        async for event in self._run_workflow_with_tracing(
-            initial_executor_fn=initial_executor_fn,
-            is_continuation=(message is None),
-            streaming=streaming,
-            function_invocation_kwargs=function_invocation_kwargs,
-            client_kwargs=client_kwargs,
-        ):
-            if event.type == "request_info" and event.request_id in (responses or {}):
-                # Don't yield request_info events for which we have responses to send -
-                # these are considered "handled". This prevents the caller from seeing
-                # events for requests they are already responding to.
-                # This usually happens when responses are provided with a checkpoint
-                # (restore then send), because the request_info events are stored in the
-                # checkpoint and would be emitted on restoration by the runner regardless
-                # of if a response is provided or not.
-                continue
-            yield event
+            initial_executor_fn = self._resolve_execution_mode(message, responses, checkpoint_id, checkpoint_storage)
 
-    async def _run_cleanup(self, checkpoint_storage: CheckpointStorage | None) -> None:
-        """Cleanup hook called after stream consumption."""
-        if checkpoint_storage is not None:
-            self._runner.context.clear_runtime_checkpoint_storage()
-        self._reset_running_flag()
+            async for event in self._run_workflow_with_tracing(
+                initial_executor_fn=initial_executor_fn,
+                is_continuation=(message is None),
+                streaming=streaming,
+                function_invocation_kwargs=function_invocation_kwargs,
+                client_kwargs=client_kwargs,
+            ):
+                if event.type == "request_info" and event.request_id in (responses or {}):
+                    # Don't yield request_info events for which we have responses to send -
+                    # these are considered "handled". This prevents the caller from seeing
+                    # events for requests they are already responding to.
+                    # This usually happens when responses are provided with a checkpoint
+                    # (restore then send), because the request_info events are stored in the
+                    # checkpoint and would be emitted on restoration by the runner regardless
+                    # of if a response is provided or not.
+                    continue
+                yield event
+        finally:
+            # Clear the active-run weakref so a subsequent ``run()`` is allowed,
+            # but only if the slot still holds *our* weakref. If the caller
+            # dropped this stream after partial iteration and a new ``run()``
+            # already installed its own weakref before our async-gen finalizer
+            # ran, ``self._active_run`` now points at the successor; clearing
+            # it would silently break the successor's concurrency guard.
+            if self._active_run is my_active_run:
+                self._active_run = None
+            if checkpoint_storage is not None:
+                self._runner.context.clear_runtime_checkpoint_storage()
 
     @staticmethod
     def _finalize_events(
@@ -935,7 +1000,7 @@ class Workflow(DictConvertible):
 
     async def _send_responses_internal(self, responses: Mapping[str, Any]) -> None:
         """Internal method to validate and send responses to the executors."""
-        pending_requests = await self._runner_context.get_pending_request_info_events()
+        pending_requests = await self._runner.context.get_pending_request_info_events()
         if not pending_requests:
             raise RuntimeError("No pending requests found in workflow context.")
 
@@ -955,7 +1020,7 @@ class Workflow(DictConvertible):
             coerced_responses[request_id] = response
 
         await asyncio.gather(*[
-            self._runner_context.send_request_info_response(request_id, response)
+            self._runner.context.send_request_info_response(request_id, response)
             for request_id, response in coerced_responses.items()
         ])
 
@@ -1151,3 +1216,12 @@ class Workflow(DictConvertible):
             context_providers=context_providers,
             **kwargs,
         )
+
+    def _is_run_active(self) -> bool:
+        """Check if a workflow run is currently active.
+
+        Returns:
+            True if a run is active, False otherwise.
+        """
+        existing_stream = self._active_run() if self._active_run is not None else None
+        return existing_stream is not None

python/packages/core/tests/workflow/test_checkpoint.py

@@ -336,6 +336,97 @@ async def test_workflow_checkpoint_chaining_via_previous_checkpoint_id():
         )
 
 
+async def test_workflow_checkpoint_ancestry_preserved_after_resume():
+    """Resuming from a checkpoint must preserve ancestry: future checkpoints chain back to the resumed one."""
+    from typing_extensions import Never
+
+    from agent_framework import WorkflowBuilder, WorkflowContext, handler
+    from agent_framework._workflows._executor import Executor
+
+    class StartExecutor(Executor):
+        @handler
+        async def run(self, message: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message(message, target_id="middle")
+
+    class MiddleExecutor(Executor):
+        @handler
+        async def process(self, message: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message(message + "-processed", target_id="finish")
+
+    class FinishExecutor(Executor):
+        @handler
+        async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None:
+            await ctx.yield_output(message + "-done")
+
+    storage = InMemoryCheckpointStorage()
+
+    def _build_workflow() -> Any:
+        start = StartExecutor(id="start")
+        middle = MiddleExecutor(id="middle")
+        finish = FinishExecutor(id="finish")
+        return (
+            WorkflowBuilder(
+                name="resume-ancestry-test",
+                max_iterations=10,
+                start_executor=start,
+                checkpoint_storage=storage,
+            )
+            .add_edge(start, middle)
+            .add_edge(middle, finish)
+            .build()
+        )
+
+    # First run: produce an initial chain of checkpoints
+    workflow = _build_workflow()
+    workflow_name = workflow.name
+    _ = [event async for event in workflow.run("hello", stream=True)]
+
+    initial_checkpoints = sorted(await storage.list_checkpoints(workflow_name=workflow_name), key=lambda c: c.timestamp)
+    assert len(initial_checkpoints) >= 3, (
+        f"Need at least 3 initial checkpoints to pick a middle one, got {len(initial_checkpoints)}"
+    )
+    initial_ids = {cp.checkpoint_id for cp in initial_checkpoints}
+
+    # Pick an intermediate checkpoint to resume from (not the first, not the last)
+    resume_from = initial_checkpoints[len(initial_checkpoints) // 2]
+
+    # Resume on a fresh workflow instance (same graph signature) and run to completion
+    resumed_workflow = _build_workflow()
+    assert resumed_workflow.name == workflow_name
+    _ = [event async for event in resumed_workflow.run(checkpoint_id=resume_from.checkpoint_id, stream=True)]
+
+    # Inspect new checkpoints created after resuming
+    all_checkpoints = sorted(await storage.list_checkpoints(workflow_name=workflow_name), key=lambda c: c.timestamp)
+    new_checkpoints = [cp for cp in all_checkpoints if cp.checkpoint_id not in initial_ids]
+    assert new_checkpoints, "Resuming from an intermediate checkpoint should produce new checkpoints"
+
+    # The very first checkpoint created after resuming must chain back to the resumed checkpoint
+    assert new_checkpoints[0].previous_checkpoint_id == resume_from.checkpoint_id, (
+        "First post-resume checkpoint must chain to the checkpoint that was resumed from; "
+        f"got previous_checkpoint_id={new_checkpoints[0].previous_checkpoint_id!r}, "
+        f"expected {resume_from.checkpoint_id!r}"
+    )
+
+    # Subsequent post-resume checkpoints must continue chaining
+    for i in range(1, len(new_checkpoints)):
+        assert new_checkpoints[i].previous_checkpoint_id == new_checkpoints[i - 1].checkpoint_id, (
+            f"Post-resume checkpoint {i} should chain to checkpoint {i - 1}"
+        )
+
+    # Walking the chain backwards from the most recent checkpoint must reach the original root
+    # without breaks (i.e. the full ancestry across the resume boundary is intact).
+    checkpoints_by_id = {cp.checkpoint_id: cp for cp in all_checkpoints}
+    chain: list[str] = []
+    cursor: str | None = new_checkpoints[-1].checkpoint_id
+    while cursor is not None:
+        chain.append(cursor)
+        cursor = checkpoints_by_id[cursor].previous_checkpoint_id
+    # Chain must include the resumed-from checkpoint and terminate at the original root
+    assert resume_from.checkpoint_id in chain
+    assert chain[-1] == initial_checkpoints[0].checkpoint_id
+    assert checkpoints_by_id[chain[-1]].previous_checkpoint_id is None
+
+
 async def test_memory_checkpoint_storage_roundtrip_json_native_types():
     """Test that JSON-native types (str, int, float, bool, None) roundtrip correctly."""
     storage = InMemoryCheckpointStorage()

python/packages/core/tests/workflow/test_runner.py

@@ -17,7 +17,6 @@ from agent_framework import (
     WorkflowContext,
     WorkflowConvergenceException,
     WorkflowEvent,
-    WorkflowRunnerException,
     WorkflowRunState,
     handler,
 )
@@ -305,40 +304,62 @@ async def test_fanout_edge_runner_delivers_to_multiple_targets_concurrently() ->
     assert probe_target.call_count == 1
 
 
-async def test_runner_already_running():
-    """Test that running the runner while it is already running raises an error."""
+async def test_runner_run_until_convergence_runs_sequentially():
+    """run_until_convergence can be invoked back-to-back on the same Runner.
+
+    The Runner itself does not enforce concurrency; that responsibility lives on
+    :class:`Workflow`. This test simply confirms the Runner is reusable across
+    sequential runs.
+    """
+    runner = _make_runner()
+    async for _ in runner.run_until_convergence():
+        pass
+    async for _ in runner.run_until_convergence():
+        pass
+
+
+def _make_runner() -> Runner:
+    """Build a minimal runner for runner-level tests."""
+    return Runner(
+        [],
+        {},
+        State(),
+        InProcRunnerContext(),
+        "test_name",
+        graph_signature_hash="test_hash",
+    )
+
+
+async def test_runner_accepts_new_run_after_previous_failure():
+    """A failed run must not leave the Runner unable to start a new run.
+
+    After the first run raises, ``run_until_convergence()`` must be callable
+    again and not surface any lifecycle-related rejection.
+    """
     executor_a = MockExecutor(id="executor_a")
     executor_b = MockExecutor(id="executor_b")
-
-    # Create a loop
     edges = [
         SingleEdgeGroup(executor_a.id, executor_b.id),
         SingleEdgeGroup(executor_b.id, executor_a.id),
     ]
-
-    executors: dict[str, Executor] = {
-        executor_a.id: executor_a,
-        executor_b.id: executor_b,
-    }
+    executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
     state = State()
     ctx = InProcRunnerContext()
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash", max_iterations=2)
 
-    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
+    await executor_a.execute(MockMessage(data=0), ["START"], state, ctx)
 
-    await executor_a.execute(
-        MockMessage(data=0),
-        ["START"],  # source_executor_ids
-        state,  # state
-        ctx,  # runner_context
-    )
+    with pytest.raises(WorkflowConvergenceException):
+        async for _ in runner.run_until_convergence():
+            pass
 
-    with pytest.raises(WorkflowRunnerException, match="Runner is already running."):
-
-        async def _run():
-            async for _ in runner.run_until_convergence():
-                pass
-
-        await asyncio.gather(_run(), _run())
+    # A second run on the same Runner must not be blocked by stale lifecycle
+    # state from the failed run.
+    try:
+        async for _ in runner.run_until_convergence():
+            pass
+    except Exception as exc:
+        assert "Runner is already running" not in str(exc), "Runner stayed locked after a failed run"
 
 
 async def test_runner_emits_runner_completion_for_agent_response_without_targets():
@@ -862,7 +883,13 @@ async def test_runner_checkpoint_with_resumed_flag():
     state = State()
 
     runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
-    runner._mark_resumed(5)  # pyright: ignore[reportPrivateUsage]
+    resumed_checkpoint = WorkflowCheckpoint(
+        checkpoint_id="resumed-cp",
+        workflow_name="test_name",
+        graph_signature_hash="test_hash",
+        iteration_count=5,
+    )
+    runner._mark_resumed(resumed_checkpoint)  # pyright: ignore[reportPrivateUsage]
 
     # Add a message to trigger the checkpoint creation path
     await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id="START"))
@@ -882,6 +909,86 @@ async def test_runner_checkpoint_with_resumed_flag():
     assert runner._resumed_from_checkpoint is False  # pyright: ignore[reportPrivateUsage]
 
 
+async def test_runner_mark_resumed_sets_previous_checkpoint_id():
+    """_mark_resumed must populate _previous_checkpoint_id so future checkpoints chain back to the resume point."""
+    runner = Runner(
+        [],
+        {},
+        State(),
+        InProcRunnerContext(),
+        "test_name",
+        graph_signature_hash="test_hash",
+    )
+
+    # Pre-condition: nothing to chain back to
+    assert runner.previous_checkpoint_id is None
+
+    resumed_checkpoint = WorkflowCheckpoint(
+        checkpoint_id="resumed-cp-id",
+        workflow_name="test_name",
+        graph_signature_hash="test_hash",
+        iteration_count=3,
+    )
+    runner._mark_resumed(resumed_checkpoint)  # pyright: ignore[reportPrivateUsage]
+
+    assert runner._resumed_from_checkpoint is True  # pyright: ignore[reportPrivateUsage]
+    assert runner._iteration == 3  # pyright: ignore[reportPrivateUsage]
+    assert runner.previous_checkpoint_id == "resumed-cp-id"
+
+
+async def test_runner_post_resume_checkpoint_chains_to_resumed_checkpoint():
+    """After resuming, the next checkpoint created must reference the resumed checkpoint as its parent."""
+    storage = InMemoryCheckpointStorage()
+    ctx = CheckpointingContext(storage)
+    executor_a = MockExecutor(id="executor_a")
+    executor_b = MockExecutor(id="executor_b")
+
+    edges = [
+        SingleEdgeGroup(executor_a.id, executor_b.id),
+        SingleEdgeGroup(executor_b.id, executor_a.id),
+    ]
+
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
+    state = State()
+
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    # Simulate having resumed from a prior checkpoint
+    resumed_checkpoint = WorkflowCheckpoint(
+        checkpoint_id="parent-checkpoint-id",
+        workflow_name="test_name",
+        graph_signature_hash="test_hash",
+        iteration_count=1,
+    )
+    runner._mark_resumed(resumed_checkpoint)  # pyright: ignore[reportPrivateUsage]
+
+    # Seed a message so the runner has work to do (and creates checkpoints at superstep boundaries)
+    await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id=executor_a.id))
+
+    async for _ in runner.run_until_convergence():
+        pass
+
+    # Find the first checkpoint created after the resume point (across all workflows tracked by storage)
+    new_checkpoints = sorted(
+        await storage.list_checkpoints(workflow_name="test_name"),
+        key=lambda c: c.timestamp,
+    )
+    assert new_checkpoints, "Resuming and running should produce at least one new checkpoint"
+
+    # The first new checkpoint must chain to the resumed-from checkpoint, not to None
+    assert new_checkpoints[0].previous_checkpoint_id == "parent-checkpoint-id", (
+        "First post-resume checkpoint must chain to the resumed checkpoint id; "
+        f"got {new_checkpoints[0].previous_checkpoint_id!r}"
+    )
+
+    # Subsequent post-resume checkpoints continue the chain
+    for i in range(1, len(new_checkpoints)):
+        assert new_checkpoints[i].previous_checkpoint_id == new_checkpoints[i - 1].checkpoint_id
+
+
 class ExecutorThatFailsWithEvents(Executor):
     """An executor that emits events and then raises an exception after receiving messages."""
 

python/packages/core/tests/workflow/test_runner_context.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for `InProcRunnerContext`."""
+
+import pytest
+
+from agent_framework import (
+    InProcRunnerContext,
+    WorkflowEvent,
+    WorkflowMessage,
+)
+
+
+def _make_request_info_event(request_id: str, source_executor_id: str = "executor") -> WorkflowEvent[str]:
+    return WorkflowEvent.request_info(
+        request_id=request_id,
+        source_executor_id=source_executor_id,
+        request_data="please respond",
+        response_type=str,
+    )
+
+
+class TestInProcRunnerContextResetForNewRun:
+    """Verify `reset_for_new_run` clears per-run state, including pending request_info events."""
+
+    async def test_reset_clears_pending_request_info_events(self) -> None:
+        ctx = InProcRunnerContext()
+
+        await ctx.add_request_info_event(_make_request_info_event("req-1"))
+        await ctx.add_request_info_event(_make_request_info_event("req-2"))
+
+        assert set((await ctx.get_pending_request_info_events()).keys()) == {"req-1", "req-2"}
+
+        ctx.reset_for_new_run()
+
+        assert await ctx.get_pending_request_info_events() == {}
+
+    async def test_reset_clears_pending_request_info_events_when_already_empty(self) -> None:
+        ctx = InProcRunnerContext()
+
+        assert await ctx.get_pending_request_info_events() == {}
+
+        ctx.reset_for_new_run()
+
+        assert await ctx.get_pending_request_info_events() == {}
+
+    async def test_reset_after_pending_event_blocks_response_correlation(self) -> None:
+        """After `reset_for_new_run`, prior request ids must no longer correlate to a response."""
+        ctx = InProcRunnerContext()
+        await ctx.add_request_info_event(_make_request_info_event("req-1"))
+
+        ctx.reset_for_new_run()
+
+        with pytest.raises(ValueError, match="No pending request found for request_id: req-1"):
+            await ctx.send_request_info_response("req-1", "answer")
+
+    async def test_reset_clears_messages_events_and_streaming_flag(self) -> None:
+        """Sanity-check the other state `reset_for_new_run` is documented to clear."""
+        ctx = InProcRunnerContext()
+        await ctx.send_message(WorkflowMessage(data="hello", source_id="executor"))
+        await ctx.add_event(WorkflowEvent("status", data="running"))
+        ctx.set_streaming(True)
+
+        assert await ctx.has_messages() is True
+        assert await ctx.has_events() is True
+        assert ctx.is_streaming() is True
+
+        ctx.reset_for_new_run()
+
+        assert await ctx.has_messages() is False
+        assert await ctx.has_events() is False
+        assert ctx.is_streaming() is False

python/packages/core/tests/workflow/test_workflow.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+import gc
 import tempfile
 from collections.abc import AsyncIterable, Awaitable, Sequence
 from dataclasses import dataclass, field
@@ -19,6 +20,7 @@ from agent_framework import (
     Content,
     Executor,
     FileCheckpointStorage,
+    InMemoryCheckpointStorage,
     Message,
     ResponseStream,
     WorkflowBuilder,
@@ -26,6 +28,7 @@ from agent_framework import (
     WorkflowContext,
     WorkflowConvergenceException,
     WorkflowEvent,
+    WorkflowException,
     WorkflowMessage,
     WorkflowRunState,
     handler,
@@ -759,8 +762,7 @@ async def test_workflow_concurrent_execution_prevention():
 
     # Try to start a second concurrent execution - this should fail
     with pytest.raises(
-        RuntimeError,
-        match="Workflow is already running. Concurrent executions are not allowed.",
+        WorkflowException, match="Workflow is already running; concurrent runs are not allowed on the same instance."
     ):
         await workflow.run(NumberMessage(data=0))
 
@@ -795,8 +797,7 @@ async def test_workflow_concurrent_execution_prevention_streaming():
 
     # Try to start a second concurrent execution - this should fail
     with pytest.raises(
-        RuntimeError,
-        match="Workflow is already running. Concurrent executions are not allowed.",
+        WorkflowException, match="Workflow is already running; concurrent runs are not allowed on the same instance."
     ):
         await workflow.run(NumberMessage(data=0))
 
@@ -828,14 +829,12 @@ async def test_workflow_concurrent_execution_prevention_mixed_methods():
 
     # Try different execution methods - all should fail
     with pytest.raises(
-        RuntimeError,
-        match="Workflow is already running. Concurrent executions are not allowed.",
+        WorkflowException, match="Workflow is already running; concurrent runs are not allowed on the same instance."
     ):
         await workflow.run(NumberMessage(data=0))
 
     with pytest.raises(
-        RuntimeError,
-        match="Workflow is already running. Concurrent executions are not allowed.",
+        WorkflowException, match="Workflow is already running; concurrent runs are not allowed on the same instance."
     ):
         async for _ in workflow.run(NumberMessage(data=0), stream=True):
             break
@@ -848,6 +847,154 @@ async def test_workflow_concurrent_execution_prevention_mixed_methods():
     assert result.get_final_state() == WorkflowRunState.IDLE
 
 
+async def test_workflow_sequential_runs_after_completion() -> None:
+    """A completed run must release the runner so the next ``run`` succeeds.
+
+    This is the happy-path counterpart to the concurrent-run guard tests:
+    those tests verify that a *concurrent* run is rejected, but they do not
+    verify that the lock is actually released afterwards. This test
+    exercises that release path explicitly across the three call shapes
+    (non-streaming, streaming-iterated, streaming-via-get_final_response)
+    and across multiple consecutive turns to catch lock leaks.
+    """
+    executor = IncrementExecutor(id="seq_executor", limit=3, increment=1)
+    workflow = WorkflowBuilder(start_executor=executor).build()
+
+    # Non-streaming -> non-streaming
+    r1 = await workflow.run(NumberMessage(data=0))
+    assert r1.get_final_state() == WorkflowRunState.IDLE
+
+    r2 = await workflow.run(NumberMessage(data=0))
+    assert r2.get_final_state() == WorkflowRunState.IDLE
+
+    # Non-streaming -> streaming-iterated
+    stream_events: list[WorkflowEvent] = []
+    async for event in workflow.run(NumberMessage(data=0), stream=True):
+        stream_events.append(event)
+    assert any(e.type == "status" and e.state == WorkflowRunState.IDLE for e in stream_events)
+
+    # Streaming -> streaming via get_final_response (no manual iteration)
+    r3 = await workflow.run(NumberMessage(data=0), stream=True).get_final_response()
+    assert r3.get_final_state() == WorkflowRunState.IDLE
+
+    # Streaming -> non-streaming (back to the start)
+    r4 = await workflow.run(NumberMessage(data=0))
+    assert r4.get_final_state() == WorkflowRunState.IDLE
+
+
+async def test_workflow_unconsumed_stream_releases_run_lock() -> None:
+    """An unconsumed stream must not leak the run lock.
+
+    ``Workflow.run`` reserves the runner *synchronously* so that concurrent
+    callers are rejected immediately. The reservation is normally released
+    by ``_run_core``'s ``finally`` once the stream is iterated. If the
+    caller never iterates the stream, a GC-time finalizer must release the
+    reservation instead - otherwise every subsequent ``Workflow.run`` call
+    on this instance would fail with the concurrent-run error.
+    """
+    executor = IncrementExecutor(id="unconsumed_stream_exec", limit=3, increment=1)
+    workflow = WorkflowBuilder(start_executor=executor).build()
+
+    # Build a stream and immediately drop it without iterating.
+    stream = workflow.run(NumberMessage(data=0), stream=True)
+    assert stream is not None  # silence unused-variable warnings; stream is GC'd below
+    del stream
+    gc.collect()
+    # Yield to the event loop so any scheduled finalizer work can run.
+    await asyncio.sleep(0)
+
+    # The runner should be back to IDLE; a fresh run must succeed.
+    result = await workflow.run(NumberMessage(data=0))
+    assert result.get_final_state() == WorkflowRunState.IDLE
+
+
+async def test_workflow_unawaited_run_coroutine_releases_run_lock() -> None:
+    """An un-awaited non-streaming ``run()`` coroutine must also not leak the lock.
+
+    ``Workflow.run`` (non-streaming) returns a coroutine produced by
+    ``ResponseStream.get_final_response``. The underlying ResponseStream is
+    held alive by that coroutine, so dropping the coroutine without
+    awaiting it must still release the reservation via the same GC-time
+    fallback used for unconsumed streams.
+    """
+    executor = IncrementExecutor(id="unawaited_run_exec", limit=3, increment=1)
+    workflow = WorkflowBuilder(start_executor=executor).build()
+
+    coro = workflow.run(NumberMessage(data=0))
+    # Closing suppresses the "coroutine was never awaited" warning. We cast to
+    # ``Any`` because the typed return is ``Awaitable[...]``; in practice it is
+    # a coroutine that exposes ``close``.
+    cast(Any, coro).close()
+    del coro
+    gc.collect()
+    await asyncio.sleep(0)
+
+    result = await workflow.run(NumberMessage(data=0))
+    assert result.get_final_state() == WorkflowRunState.IDLE
+
+
+async def test_workflow_partial_stream_does_not_clobber_successor_active_run() -> None:
+    """A stale ``_run_core`` finalizer must not clear a successor's run lock.
+
+    Repro for the GC-finalizer race the user reported:
+
+    1. Start stream A and consume one event so its body is suspended at a
+       ``yield``. Its ``finally`` is now armed and will run when the
+       generator is closed.
+    2. Drop stream A and ``gc.collect``. The ``_active_run`` weakref's
+       referent is gone, so a subsequent ``run()`` will pass the
+       concurrency guard - but stream A's async-gen finalizer hasn't
+       actually executed yet (``aclose`` is scheduled on the loop).
+    3. Synchronously start stream B; ``run()`` installs a fresh weakref
+       in ``_active_run``.
+    4. Yield to the loop so stream A's stale ``finally`` runs. Without
+       the identity check it unconditionally writes
+       ``self._active_run = None``, silently disabling the concurrency
+       guard for stream B.
+    """
+    executor = IncrementExecutor(id="stale_finalizer_exec", limit=100, increment=1)
+    workflow = WorkflowBuilder(start_executor=executor).build()
+
+    # Step 1: drive stream A's body until it's suspended at its first yield.
+    stream_a = workflow.run(NumberMessage(data=0), stream=True)
+    aiter_a = stream_a.__aiter__()
+    await aiter_a.__anext__()
+
+    # Step 2: drop stream A; GC invalidates the weakref and schedules
+    # async-gen close, but does not run the close inline.
+    del stream_a
+    del aiter_a
+    gc.collect()
+
+    # Step 3: synchronously start stream B *before* yielding to the loop,
+    # so the stale ``aclose`` for stream A hasn't fired yet.
+    stream_b = workflow.run(NumberMessage(data=0), stream=True)
+    ref_b = workflow._active_run  # type: ignore[attr-defined]
+    assert ref_b is not None and ref_b() is stream_b
+
+    # Step 4: yield enough times for stream A's scheduled aclose to drive
+    # its body through ``GeneratorExit`` and into its ``finally``.
+    for _ in range(5):
+        await asyncio.sleep(0)
+
+    # With the fix, stream B's reservation is still in place. Without it,
+    # ``_active_run`` was clobbered to ``None`` and a concurrent run would
+    # be (incorrectly) accepted.
+    assert workflow._active_run is ref_b  # type: ignore[attr-defined]
+    with pytest.raises(
+        WorkflowException,
+        match="Workflow is already running; concurrent runs are not allowed on the same instance.",
+    ):
+        await workflow.run(NumberMessage(data=0))
+
+    # Tear down stream B without iterating it (its body never started, so
+    # closing it is a no-op for workflow state).
+    del stream_b
+    del ref_b
+    gc.collect()
+    await asyncio.sleep(0)
+
+
 class _StreamingTestAgent(BaseAgent):
     """Test agent that supports both streaming and non-streaming modes."""
 
@@ -1269,3 +1416,143 @@ async def test_output_executors_filtering_with_run_responses_streaming() -> None
 
 
 # endregion
+
+
+# region Workflow.create_checkpoint
+
+
+class TestWorkflowCreateCheckpoint:
+    """Tests for :meth:`Workflow.create_checkpoint`."""
+
+    async def test_returns_checkpoint_id_with_runtime_storage(self, simple_executor: Executor) -> None:
+        """Calling `create_checkpoint` with a runtime storage persists a checkpoint and returns its id."""
+        storage = InMemoryCheckpointStorage()
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        checkpoint_id = await workflow.create_checkpoint(storage)
+
+        assert checkpoint_id
+        loaded = await storage.load(checkpoint_id)
+        assert loaded is not None
+        assert loaded.checkpoint_id == checkpoint_id
+        assert loaded.workflow_name == workflow.name
+        assert loaded.graph_signature_hash == workflow.graph_signature_hash
+
+    async def test_uses_buildtime_storage_when_none_provided(self, simple_executor: Executor) -> None:
+        """When called with `None`, the build-time storage is used."""
+        storage = InMemoryCheckpointStorage()
+        workflow = (
+            WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
+            .add_edge(simple_executor, simple_executor)
+            .build()
+        )
+
+        checkpoint_id = await workflow.create_checkpoint(None)
+
+        loaded = await storage.load(checkpoint_id)
+        assert loaded is not None
+        assert loaded.checkpoint_id == checkpoint_id
+
+    async def test_raises_when_no_storage_available(self, simple_executor: Executor) -> None:
+        """Without build-time or runtime storage, `create_checkpoint(None)` raises."""
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        with pytest.raises(WorkflowCheckpointException, match="Checkpoint storage must be provided"):
+            await workflow.create_checkpoint(None)
+
+    async def test_raises_while_run_active(self, simple_executor: Executor) -> None:
+        """`create_checkpoint` must reject while a workflow run is still active."""
+        storage = InMemoryCheckpointStorage()
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        # Hold a live reference to a streaming run without iterating it so that
+        # ``_is_run_active`` remains True (the active-run weakref still resolves).
+        active_stream = workflow.run(WorkflowMessage(data="hi", source_id="test"), stream=True)
+        try:
+            with pytest.raises(WorkflowException, match="Cannot create checkpoint while a workflow run is active"):
+                await workflow.create_checkpoint(storage)
+        finally:
+            # Drain the stream so the run completes cleanly and the active-run
+            # weakref is cleared; otherwise pytest's asyncio teardown can leak
+            # the unconsumed generator.
+            async for _ in active_stream:
+                pass
+
+    async def test_clears_runtime_storage_after_call(self, simple_executor: Executor) -> None:
+        """The runtime storage override must not leak past the call."""
+        storage = InMemoryCheckpointStorage()
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        await workflow.create_checkpoint(storage)
+
+        assert workflow._runner.context.has_checkpointing() is False
+        assert workflow._runner.context._runtime_checkpoint_storage is None  # type: ignore[attr-defined]
+
+    async def test_clears_runtime_storage_after_failure(self, simple_executor: Executor) -> None:
+        """The runtime storage override must be cleared even if checkpoint creation fails."""
+        from unittest.mock import AsyncMock
+
+        storage = InMemoryCheckpointStorage()
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        # The runner logs-and-swallows storage save errors, so a failed save
+        # surfaces as the "Failed to create checkpoint." path when
+        # ``previous_checkpoint_id`` remains ``None``. Either way, the
+        # ``finally`` cleanup must still clear the runtime override.
+        storage.save = AsyncMock(side_effect=RuntimeError("boom"))  # type: ignore[method-assign]
+
+        with pytest.raises(WorkflowCheckpointException, match="Failed to create checkpoint"):
+            await workflow.create_checkpoint(storage)
+
+        assert workflow._runner.context._runtime_checkpoint_storage is None  # type: ignore[attr-defined]
+
+    async def test_alters_lineage_for_next_checkpoint(self, simple_executor: Executor) -> None:
+        """A manually created checkpoint becomes the parent of the next checkpoint."""
+        storage = InMemoryCheckpointStorage()
+        workflow = (
+            WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
+            .add_edge(simple_executor, simple_executor)
+            .build()
+        )
+
+        first_id = await workflow.create_checkpoint(None)
+        second_id = await workflow.create_checkpoint(None)
+
+        assert first_id != second_id
+        second = await storage.load(second_id)
+        assert second is not None
+        assert second.previous_checkpoint_id == first_id
+
+    async def test_raises_when_save_fails_after_prior_success(self, simple_executor: Executor) -> None:
+        """A failed save after an earlier successful checkpoint must not return the stale id.
+
+        The runner log-and-swallows storage save errors and only updates
+        ``previous_checkpoint_id`` on success. Without an explicit transition check,
+        ``create_checkpoint`` would silently return the previously stored id as if a
+        new checkpoint had been created.
+        """
+        from unittest.mock import AsyncMock
+
+        storage = InMemoryCheckpointStorage()
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        # First call succeeds and seeds ``previous_checkpoint_id``.
+        first_id = await workflow.create_checkpoint(storage)
+        assert first_id
+
+        # Second call fails to save, so the runner leaves ``previous_checkpoint_id``
+        # pointing at ``first_id``. The method must detect that the id did not
+        # transition and raise instead of returning the stale value.
+        original_save = storage.save
+        storage.save = AsyncMock(side_effect=RuntimeError("boom"))  # type: ignore[method-assign]
+        try:
+            with pytest.raises(WorkflowCheckpointException, match="Failed to create checkpoint"):
+                await workflow.create_checkpoint(storage)
+        finally:
+            storage.save = original_save  # type: ignore[method-assign]
+
+        # The runner's bookkeeping is unchanged after the failed call.
+        assert workflow._runner.previous_checkpoint_id == first_id  # type: ignore[attr-defined]
+
+
+# endregion

python/packages/declarative/tests/test_http_request_executor.py

@@ -90,7 +90,7 @@ async def _run(yaml_def: dict[str, Any], handler: HttpRequestHandler) -> Any:
 
 def _state(workflow: Any, events: Any) -> dict[str, Any]:
     """Read declarative state out of the workflow after run completes."""
-    return workflow._state.get(DECLARATIVE_STATE_KEY) or {}
+    return workflow._runner.state.get(DECLARATIVE_STATE_KEY) or {}
 
 
 # Helper used by parametrised path tests
@@ -151,7 +151,7 @@ class TestSuccessPath:
         workflow = factory.create_workflow_from_definition(_yaml(_action(method="GET", response="Local.Result")))
         await workflow.run({})
 
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] == {"key": "value", "number": 42}
         assert handler.last_info is not None
         assert handler.last_info.method == "GET"
@@ -164,7 +164,7 @@ class TestSuccessPath:
         workflow = factory.create_workflow_from_definition(_yaml(_action(response="Local.Result")))
         await workflow.run({})
 
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] == "not-json content"
 
     @pytest.mark.asyncio
@@ -174,7 +174,7 @@ class TestSuccessPath:
         workflow = factory.create_workflow_from_definition(_yaml(_action(response="Local.Result")))
         await workflow.run({})
 
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] is None
 
     @pytest.mark.asyncio
@@ -184,7 +184,7 @@ class TestSuccessPath:
         workflow = factory.create_workflow_from_definition(_yaml(_action(response={"path": "Local.Result"})))
         await workflow.run({})
 
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] == {"x": 1}
 
     @pytest.mark.asyncio
@@ -517,7 +517,7 @@ class TestResponseHeaders:
         factory = WorkflowFactory(http_request_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(response_headers="Local.H")))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         h = decl["Local"]["H"]
         assert h["Content-Type"] == "application/json"
         assert h["Set-Cookie"] == "a=1,b=2"
@@ -528,7 +528,7 @@ class TestResponseHeaders:
         factory = WorkflowFactory(http_request_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(response_headers="Local.H")))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["H"] is None
 
     @pytest.mark.asyncio
@@ -538,7 +538,7 @@ class TestResponseHeaders:
         workflow = factory.create_workflow_from_definition(_yaml(_action(response_headers="Local.H")))
         with pytest.raises(DeclarativeActionError):
             await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["H"] == {"X-Trace": "abc"}
 
 
@@ -559,7 +559,7 @@ class TestConversationAppend:
             )
         )
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         conv = decl["System"]["conversations"].get("conv-test-1")
         assert conv is not None
         assert len(conv["messages"]) == 1
@@ -570,7 +570,7 @@ class TestConversationAppend:
         factory = WorkflowFactory(http_request_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(response="Local.Result", conversation_id="")))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         # Auto-init creates an entry for the System.ConversationId conversation,
         # but it should NOT have HTTP-appended messages from us.
         for _cid, conv in decl["System"]["conversations"].items():
@@ -582,7 +582,7 @@ class TestConversationAppend:
         factory = WorkflowFactory(http_request_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(conversation_id="conv-test-1")))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         # No conversation entry should have been created either.
         assert "conv-test-1" not in decl["System"]["conversations"]
 

python/packages/declarative/tests/test_http_request_yaml_integration.py

@@ -73,7 +73,7 @@ async def test_http_request_yaml_roundtrip() -> None:
     workflow = factory.create_workflow_from_yaml_path(FIXTURE_PATH)
     await workflow.run({})
 
-    decl: dict[str, Any] = workflow._state.get(DECLARATIVE_STATE_KEY) or {}
+    decl: dict[str, Any] = workflow._runner.state.get(DECLARATIVE_STATE_KEY) or {}
     local = decl.get("Local") or {}
 
     assert local.get("RepoOwner") == "dotnet"

python/packages/declarative/tests/test_invoke_mcp_tool_executor.py

@@ -244,7 +244,7 @@ class TestOutput:
         factory = WorkflowFactory(mcp_tool_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] == [{"k": "v", "n": 1}]
 
     @pytest.mark.asyncio
@@ -253,7 +253,7 @@ class TestOutput:
         factory = WorkflowFactory(mcp_tool_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] == ["plain text not json"]
 
     @pytest.mark.asyncio
@@ -262,7 +262,7 @@ class TestOutput:
         factory = WorkflowFactory(mcp_tool_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(output={"messages": "Local.Messages"})))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         msg = decl["Local"]["Messages"]
         # Single Tool-role message containing both contents (parity with .NET).
         assert isinstance(msg, Message)
@@ -276,7 +276,7 @@ class TestOutput:
         factory = WorkflowFactory(mcp_tool_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] == ["https://example.com/file.txt"]
 
     @pytest.mark.asyncio
@@ -285,7 +285,7 @@ class TestOutput:
         factory = WorkflowFactory(mcp_tool_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": {"path": "Local.Result"}})))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] == ["ok"]
 
 
@@ -306,7 +306,7 @@ class TestConversation:
             )
         )
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         conv = decl["System"]["conversations"]["conv-42"]
         msgs = conv["messages"] if isinstance(conv, dict) else conv.messages
         assert len(msgs) == 1
@@ -328,7 +328,7 @@ class TestConversation:
             )
         )
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         # Empty conversation id must not produce a `""` entry under System.conversations.
         conversations = decl.get("System", {}).get("conversations", {})
         assert "" not in conversations
@@ -529,7 +529,7 @@ class TestErrorHandling:
         factory = WorkflowFactory(mcp_tool_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] == "Error: server down"
 
     @pytest.mark.asyncio
@@ -538,7 +538,7 @@ class TestErrorHandling:
         factory = WorkflowFactory(mcp_tool_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert decl["Local"]["Result"] == "Error: invalid arguments"
 
     @pytest.mark.asyncio
@@ -547,7 +547,7 @@ class TestErrorHandling:
         factory = WorkflowFactory(mcp_tool_handler=handler)
         workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
         await workflow.run({})
-        decl = workflow._state.get(DECLARATIVE_STATE_KEY)
+        decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         result = decl["Local"]["Result"]
         assert isinstance(result, str)
         assert result.startswith("Error:")

python/packages/declarative/tests/test_workflow_factory.py

@@ -289,11 +289,11 @@ actions:
         # Stamp a marker into the declarative state between turns. The
         # continuation branch must preserve it; a state-clearing run would
         # wipe ``DECLARATIVE_STATE_KEY`` and force re-initialization.
-        state_data = workflow._state.get(DECLARATIVE_STATE_KEY)
+        state_data = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert isinstance(state_data, dict), "Expected declarative state to be initialized after turn 1"
         state_data["Local"] = {"persisted_marker": "kept-from-turn-1"}
-        workflow._state.set(DECLARATIVE_STATE_KEY, state_data)
-        workflow._state.commit()
+        workflow._runner.state.set(DECLARATIVE_STATE_KEY, state_data)
+        workflow._runner.state.commit()
 
         second = await agent.run("turn-2-msg")
         assert second.text == "turn-2-msg", (
@@ -303,7 +303,7 @@ actions:
         # The continuation branch in ``_ensure_state_initialized`` must:
         # 1. preserve the cross-turn marker we stamped above
         # 2. refresh Inputs.input and System.LastMessage* to the new turn
-        post_state = workflow._state.get(DECLARATIVE_STATE_KEY)
+        post_state = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
         assert isinstance(post_state, dict), "declarative state vanished between turns"
         local = post_state.get("Local", {})
         assert local.get("persisted_marker") == "kept-from-turn-1", (

python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py

@@ -17,6 +17,7 @@ from typing import Protocol, cast
 
 from agent_framework import (
     ChatOptions,
+    CheckpointID,
     Content,
     ContextProvider,
     FileCheckpointStorage,
@@ -343,6 +344,7 @@ class ResponsesHostServer(ResponsesAgentServerHost):
 
     # TODO(@taochen): Allow a different checkpoint storage that stores checkpoints externally
     CHECKPOINT_STORAGE_PATH = "/.checkpoints"
+    INITIAL_CHECKPOINT_STORAGE_NAME = "initial"
     FUNCTION_APPROVAL_STORAGE_PATH = "/.function_approvals/approval_requests.json"
 
     def __init__(
@@ -386,7 +388,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
             )
 
         self._is_workflow_agent = False
-        self._checkpoint_storage_path = None
         if isinstance(agent, WorkflowAgent):
             if agent.workflow._runner_context.has_checkpointing():  # pyright: ignore[reportPrivateUsage]
                 raise RuntimeError(
@@ -399,6 +400,12 @@ class ResponsesHostServer(ResponsesAgentServerHost):
                 else os.path.join(os.getcwd(), self.CHECKPOINT_STORAGE_PATH.lstrip("/"))
             )
             self._is_workflow_agent = True
+            # The initial checkpoint storage that stores the workflow's initial state. We will use this checkpoint
+            # to restore the workflow when no conversation_id or previous_response_id is supplied in a request.
+            self._initial_checkpoint_storage = _checkpoint_storage_for_context(
+                self._checkpoint_storage_path, self.INITIAL_CHECKPOINT_STORAGE_NAME
+            )
+            self._initial_checkpoint_id: CheckpointID | None = None
 
         self._agent = agent
         self._approval_storage = (
@@ -580,8 +587,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
 
         # The following should never happen due to the checks above.
         # This is for type safety and defensive programming.
-        if self._checkpoint_storage_path is None:
-            raise RuntimeError("Checkpoint storage path is not configured for workflow agent.")
         if not isinstance(self._agent, WorkflowAgent):
             raise RuntimeError("Agent is not a workflow agent.")
 
@@ -590,6 +595,12 @@ class ResponsesHostServer(ResponsesAgentServerHost):
         # any future async resources owned by the workflow are entered here.
         await self._ensure_agent_ready()
 
+        # Create a checkpoint to store the initial state of the workflow, if it doesn't already exist.
+        # This allows us to restore to a clean slate when no conversation_id or previous_response_id
+        # is supplied in a request.
+        if self._initial_checkpoint_id is None:
+            self._initial_checkpoint_id = await self._agent.workflow.create_checkpoint(self._initial_checkpoint_storage)
+
         # Determine the latest checkpoint (if any) so we can resume the
         # workflow's prior state for this turn. The directory is keyed by
         # the inbound context id (conversation_id when set, otherwise
@@ -599,14 +610,40 @@ class ResponsesHostServer(ResponsesAgentServerHost):
         # the only place that state lives is the workflow checkpoint, so
         # on every turn we restore the latest checkpoint and feed the new
         # input back into the start executor as a continuation rather than
-        # a fresh run.
-        latest_checkpoint_id: str | None = None
-        restore_storage: FileCheckpointStorage | None = None
+        # a fresh run. If no conversation_id or previous_response_id is
+        # supplied, the workflow will be restored to the initial checkpoint
+        # to avoid context bleed between requests.
+        latest_checkpoint_id: str = self._initial_checkpoint_id
+        restore_storage: FileCheckpointStorage = self._initial_checkpoint_storage
         if context_id is not None:
-            restore_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, context_id)
-            latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name)
+            context_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, context_id)
+            latest_checkpoint = await context_storage.get_latest(workflow_name=self._agent.workflow.name)
             if latest_checkpoint is not None:
+                # Only switch the restore storage when a checkpoint was actually
+                # found under the per-context directory. Otherwise the initial
+                # checkpoint id would not resolve in `context_storage` and the
+                # restore call below would fail.
                 latest_checkpoint_id = latest_checkpoint.checkpoint_id
+                restore_storage = context_storage
+
+        # Restore the workflow to the latest checkpoint and run it with the
+        # new input. Events (including request info events) will not be emitted
+        # during restoration (in streaming) or after restoration (in non-streaming)
+        # since we assume the client had already seen those events and we don't want
+        # to emit duplicates.
+        if is_streaming_request:
+            async for _ in self._agent.run(
+                stream=True,
+                checkpoint_id=latest_checkpoint_id,
+                checkpoint_storage=restore_storage,
+            ):
+                pass
+        else:
+            await self._agent.run(
+                stream=False,
+                checkpoint_id=latest_checkpoint_id,
+                checkpoint_storage=restore_storage,
+            )
 
         # Storage that will receive checkpoints written during this turn.
         # When the caller chains with previous_response_id, the next turn
@@ -619,37 +656,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
         write_context_id = context.conversation_id or context.response_id
         write_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, write_context_id)
 
-        # Multi-turn pattern: when we have a prior checkpoint, restore it
-        # first (drive the workflow back to idle with prior state intact),
-        # then make a separate call that delivers the new user input. This
-        # depends on Workflow.run preserving shared state across calls. The
-        # restore-only call may yield events from any pending in-flight
-        # work in the checkpoint; we consume those internally here so they
-        # don't surface to the response stream as duplicates.
-        #
-        # If the restored checkpoint had pending request_info events, the
-        # restore-only call replays them through
-        # ``WorkflowAgent._convert_workflow_event_to_agent_response_updates``
-        # and populates ``self._agent.pending_requests``. That is the correct
-        # state: those requests are genuinely outstanding, and the next
-        # ``run(input_messages, ...)`` call may contain ``function_call_output``
-        # items (carried as FunctionResult/FunctionApprovalResponse content)
-        # that fulfill them via :meth:`WorkflowAgent._process_pending_requests`.
-        if latest_checkpoint_id is not None:
-            if is_streaming_request:
-                async for _ in self._agent.run(
-                    stream=True,
-                    checkpoint_id=latest_checkpoint_id,
-                    checkpoint_storage=restore_storage,
-                ):
-                    pass
-            else:
-                await self._agent.run(
-                    stream=False,
-                    checkpoint_id=latest_checkpoint_id,
-                    checkpoint_storage=restore_storage,
-                )
-
         # Now run the agent with the latest input
         response_event_stream = ResponseEventStream(response_id=context.response_id, model=request.model)
 

python/packages/foundry_hosting/tests/test_responses.py

@@ -3032,6 +3032,7 @@ class TestCheckpointContextPathValidation:
         agent.workflow = MagicMock()
         agent.workflow.name = "wf"
         agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
+        agent.workflow.create_checkpoint = AsyncMock(return_value="cp_initial")
         agent.run = AsyncMock(
             side_effect=[
                 AgentResponse(messages=[]),
@@ -3062,6 +3063,181 @@ class TestCheckpointContextPathValidation:
         assert new_turn_messages[0].text == "next turn"
         assert new_turn_call.kwargs["checkpoint_storage"].storage_path == (root / response_id).resolve()
 
+    async def test_handle_inner_workflow_restores_initial_checkpoint_when_no_context_id(self, tmp_path: Any) -> None:
+        """When neither previous_response_id nor conversation_id is supplied, the workflow
+        must be restored from the initial checkpoint to avoid context bleed between requests.
+        """
+        from agent_framework import WorkflowAgent
+        from azure.ai.agentserver.responses import ResponseContext
+        from azure.ai.agentserver.responses.models import CreateResponse, ItemMessage
+
+        response_id = "resp_current"
+        root = tmp_path / "root"
+        root.mkdir()
+
+        agent = MagicMock(spec=WorkflowAgent)
+        agent.id = "wf-agent"
+        agent.name = "wf"
+        agent.description = ""
+        agent.context_providers = []
+        agent.workflow = MagicMock()
+        agent.workflow.name = "wf"
+        agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
+        agent.workflow.create_checkpoint = AsyncMock(return_value="cp_initial")
+        agent.run = AsyncMock(
+            side_effect=[
+                AgentResponse(messages=[]),
+                AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ok")])]),
+            ]
+        )
+        server = ResponsesHostServer(agent, store=InMemoryResponseProvider())
+        server._checkpoint_storage_path = str(root)  # pyright: ignore[reportPrivateUsage]
+
+        # No previous_response_id and no conversation_id.
+        request = CreateResponse(model="m", input="hi")
+        context = ResponseContext(response_id=response_id, mode_flags=MagicMock())
+        input_item = ItemMessage({"type": "message", "role": "user", "content": "fresh turn"})
+
+        with patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[input_item])):
+            async for _ in server._handle_inner_workflow(request, context):  # pyright: ignore[reportPrivateUsage]
+                pass
+
+        # The initial checkpoint must have been created exactly once, against the
+        # initial checkpoint storage owned by the server.
+        assert agent.workflow.create_checkpoint.await_count == 1
+        (initial_storage_arg,) = agent.workflow.create_checkpoint.await_args.args
+        assert initial_storage_arg is server._initial_checkpoint_storage  # pyright: ignore[reportPrivateUsage]
+
+        # First run() call is the restoration: no positional input, restored from
+        # the initial checkpoint id, using the initial checkpoint storage (NOT a
+        # per-context directory).
+        assert agent.run.call_count == 2
+        restore_call = agent.run.call_args_list[0]
+        assert restore_call.args == ()
+        assert restore_call.kwargs["checkpoint_id"] == "cp_initial"
+        assert restore_call.kwargs["checkpoint_storage"] is server._initial_checkpoint_storage  # pyright: ignore[reportPrivateUsage]
+
+        # Second run() call delivers the new input; checkpoints land under response_id
+        # (the write-sink directory keyed by the current response id).
+        new_turn_call = agent.run.call_args_list[1]
+        new_turn_messages = new_turn_call.args[0]
+        assert len(new_turn_messages) == 1
+        assert new_turn_messages[0].text == "fresh turn"
+        assert new_turn_call.kwargs["checkpoint_storage"].storage_path == (root / response_id).resolve()
+
+    async def test_handle_inner_workflow_creates_initial_checkpoint_once_across_requests(self, tmp_path: Any) -> None:
+        """The initial checkpoint must be created exactly once and reused across
+        subsequent requests, regardless of whether the requests carry a context id.
+        """
+        from agent_framework import WorkflowAgent
+        from azure.ai.agentserver.responses import ResponseContext
+        from azure.ai.agentserver.responses.models import CreateResponse, ItemMessage
+
+        root = tmp_path / "root"
+        root.mkdir()
+
+        agent = MagicMock(spec=WorkflowAgent)
+        agent.id = "wf-agent"
+        agent.name = "wf"
+        agent.description = ""
+        agent.context_providers = []
+        agent.workflow = MagicMock()
+        agent.workflow.name = "wf"
+        agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
+        agent.workflow.create_checkpoint = AsyncMock(return_value="cp_initial")
+        # Four run() calls total: restore + new turn for each of the two requests.
+        agent.run = AsyncMock(return_value=AgentResponse(messages=[]))
+
+        server = ResponsesHostServer(agent, store=InMemoryResponseProvider())
+        server._checkpoint_storage_path = str(root)  # pyright: ignore[reportPrivateUsage]
+
+        request1 = CreateResponse(model="m", input="hi")
+        context1 = ResponseContext(response_id="resp_first", mode_flags=MagicMock())
+        request2 = CreateResponse(model="m", input="hi again")
+        context2 = ResponseContext(response_id="resp_second", mode_flags=MagicMock())
+        input_item = ItemMessage({"type": "message", "role": "user", "content": "turn"})
+
+        with patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[input_item])):
+            async for _ in server._handle_inner_workflow(request1, context1):  # pyright: ignore[reportPrivateUsage]
+                pass
+            async for _ in server._handle_inner_workflow(request2, context2):  # pyright: ignore[reportPrivateUsage]
+                pass
+
+        # Initial checkpoint creation must not be repeated on the second request.
+        assert agent.workflow.create_checkpoint.await_count == 1
+
+        # Both requests' restoration calls must use the same initial checkpoint id
+        # and the same initial checkpoint storage instance.
+        restore_call_1 = agent.run.call_args_list[0]
+        restore_call_2 = agent.run.call_args_list[2]
+        assert restore_call_1.kwargs["checkpoint_id"] == "cp_initial"
+        assert restore_call_2.kwargs["checkpoint_id"] == "cp_initial"
+        assert (
+            restore_call_1.kwargs["checkpoint_storage"]
+            is restore_call_2.kwargs["checkpoint_storage"]
+            is server._initial_checkpoint_storage  # pyright: ignore[reportPrivateUsage]
+        )
+
+    async def test_handle_inner_workflow_falls_back_to_initial_storage_when_context_dir_is_empty(
+        self, tmp_path: Any
+    ) -> None:
+        """When ``previous_response_id`` is supplied but its checkpoint directory has no
+        checkpoints, the restoration must fall back to BOTH the initial checkpoint id
+        and the initial checkpoint storage. Otherwise the initial id would be looked up
+        inside the per-context storage where it does not exist, and the restore would
+        fail.
+        """
+        from agent_framework import WorkflowAgent
+        from azure.ai.agentserver.responses import ResponseContext
+        from azure.ai.agentserver.responses.models import CreateResponse, ItemMessage
+
+        previous_response_id = "resp_previous"
+        response_id = "resp_current"
+        root = tmp_path / "root"
+        root.mkdir()
+        # The per-context storage exists but contains no checkpoints.
+        (root / previous_response_id).mkdir()
+
+        agent = MagicMock(spec=WorkflowAgent)
+        agent.id = "wf-agent"
+        agent.name = "wf"
+        agent.description = ""
+        agent.context_providers = []
+        agent.workflow = MagicMock()
+        agent.workflow.name = "wf"
+        agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
+        agent.workflow.create_checkpoint = AsyncMock(return_value="cp_initial")
+        agent.run = AsyncMock(
+            side_effect=[
+                AgentResponse(messages=[]),
+                AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ok")])]),
+            ]
+        )
+        server = ResponsesHostServer(agent, store=InMemoryResponseProvider())
+        server._checkpoint_storage_path = str(root)  # pyright: ignore[reportPrivateUsage]
+
+        request = CreateResponse(model="m", input="hi", previous_response_id=previous_response_id)
+        context = ResponseContext(
+            response_id=response_id, previous_response_id=previous_response_id, mode_flags=MagicMock()
+        )
+        input_item = ItemMessage({"type": "message", "role": "user", "content": "next turn"})
+
+        with patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[input_item])):
+            async for _ in server._handle_inner_workflow(request, context):  # pyright: ignore[reportPrivateUsage]
+                pass
+
+        # The restoration call must use the initial id AND the initial storage,
+        # not the empty per-context storage. Mismatching the two would attempt
+        # to load ``cp_initial`` from a directory that doesn't contain it.
+        assert agent.run.call_count == 2
+        restore_call = agent.run.call_args_list[0]
+        assert restore_call.kwargs["checkpoint_id"] == "cp_initial"
+        assert restore_call.kwargs["checkpoint_storage"] is server._initial_checkpoint_storage  # pyright: ignore[reportPrivateUsage]
+
+        # The new turn still writes checkpoints under the current response id.
+        new_turn_call = agent.run.call_args_list[1]
+        assert new_turn_call.kwargs["checkpoint_storage"].storage_path == (root / response_id).resolve()
+
     @pytest.mark.parametrize(
         "bad_id",
         [
@@ -3155,6 +3331,8 @@ class TestCheckpointContextPathValidation:
         agent.workflow = MagicMock()
         agent.workflow.name = "wf"
         agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
+        agent.workflow.create_checkpoint = AsyncMock(return_value="cp_initial")
+        agent.run = AsyncMock(return_value=AgentResponse(messages=[]))
 
         # Constructor inspects WorkflowAgent.workflow internals; bypass setup
         # by feeding a configured mock through a normal init.