fix: Fix Checkpoint Restore when Rehydrating Run (#642)

When checkpointing we did not persist the set of instantiated executors. This means, in turn, when we restore from a checkpoint when using Resume(Stream) rather than restoring a checkpoint in the context of an already existing (Streaming)Run, the executors never got reinstantiated and there were no executors to notify that a state should be loaded.

The fix is to ensure we persist the list and reinstantiate the executors on rehydration.

* Also adds a rehydration restore test

dotnet/src/Microsoft.Agents.Workflows/Execution/RunnerStateData.cs

@@ -5,8 +5,9 @@ using Microsoft.Agents.Workflows.Checkpointing;
 
 namespace Microsoft.Agents.Workflows.Execution;
 
-internal class RunnerStateData(Dictionary<ExecutorIdentity, List<ExportedState>> queuedMessages, List<ExternalRequest> outstandingRequests)
+internal class RunnerStateData(HashSet<string> instantiatedExecutors, Dictionary<ExecutorIdentity, List<ExportedState>> queuedMessages, List<ExternalRequest> outstandingRequests)
 {
+    public HashSet<string> InstantiatedExecutors { get; } = instantiatedExecutors;
     public Dictionary<ExecutorIdentity, List<ExportedState>> QueuedMessages { get; } = queuedMessages;
     public List<ExternalRequest> OutstandingRequests { get; } = outstandingRequests;
 }

dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunner.cs

@@ -228,12 +228,12 @@ internal class InProcessRunner<TInput> : ISuperStepRunner, ICheckpointingRunner
             this._workflowInfoCache = this.Workflow.ToWorkflowInfo();
         }
 
-        RunnerStateData runnerData = await this.RunContext.ExportStateAsync().ConfigureAwait(false);
         Dictionary<EdgeConnection, ExportedState> edgeData = await this.EdgeMap.ExportStateAsync().ConfigureAwait(false);
 
         await prepareTask.ConfigureAwait(false);
         await this.RunContext.StateManager.PublishUpdatesAsync(this.StepTracer).ConfigureAwait(false);
 
+        RunnerStateData runnerData = await this.RunContext.ExportStateAsync().ConfigureAwait(false);
         Dictionary<ScopeKey, ExportedState> stateData = await this.RunContext.StateManager.ExportStateAsync().ConfigureAwait(false);
 
         Checkpoint checkpoint = new(this.StepTracer.StepNumber, this._workflowInfoCache, runnerData, stateData, edgeData);
@@ -261,9 +261,9 @@ internal class InProcessRunner<TInput> : ISuperStepRunner, ICheckpointingRunner
         }
 
         await this.RunContext.StateManager.ImportStateAsync(checkpoint).ConfigureAwait(false);
-        Task executorNotifyTask = this.RunContext.NotifyCheckpointLoadedAsync(cancellation);
-
         await this.RunContext.ImportStateAsync(checkpoint).ConfigureAwait(false);
+
+        Task executorNotifyTask = this.RunContext.NotifyCheckpointLoadedAsync(cancellation);
         ValueTask republishRequestsTask = this.RunContext.RepublishUnservicedRequestsAsync(cancellation);
 
         await this.EdgeMap.ImportStateAsync(checkpoint).ConfigureAwait(false);

dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunnerContext.cs

@@ -136,8 +136,9 @@ internal class InProcessRunnerContext<TExternalInput> : IRunnerContext
         }
 
         Dictionary<ExecutorIdentity, List<ExportedState>> queuedMessages = this._nextStep.ExportMessages();
-
-        RunnerStateData result = new(queuedMessages, this._externalRequests.Values.ToList());
+        RunnerStateData result = new(instantiatedExecutors: [.. this._executors.Keys],
+                                     queuedMessages,
+                                     outstandingRequests: [.. this._externalRequests.Values]);
 
         return new(result);
     }
@@ -154,7 +155,7 @@ internal class InProcessRunnerContext<TExternalInput> : IRunnerContext
         }
     }
 
-    internal ValueTask ImportStateAsync(Checkpoint checkpoint)
+    internal async ValueTask ImportStateAsync(Checkpoint checkpoint)
     {
         if (this.QueuedEvents.Count > 0)
         {
@@ -163,6 +164,11 @@ internal class InProcessRunnerContext<TExternalInput> : IRunnerContext
 
         RunnerStateData importedState = checkpoint.RunnerData;
 
+        Task<Executor>[] executorTasks = importedState.InstantiatedExecutors
+                                                      .Where(id => !this._executors.ContainsKey(id))
+                                                      .Select(id => this.EnsureExecutorAsync(id, tracer: null).AsTask())
+                                                      .ToArray();
+
         this._nextStep = new StepContext();
         this._nextStep.ImportMessages(importedState.QueuedMessages);
 
@@ -176,6 +182,6 @@ internal class InProcessRunnerContext<TExternalInput> : IRunnerContext
             this._externalRequests[request.RequestId] = request;
         }
 
-        return default;
+        await Task.WhenAll(executorTasks).ConfigureAwait(false);
     }
 }

dotnet/tests/Microsoft.Agents.Workflows.UnitTests/Sample/05_Simple_Workflow_Checkpointing.cs

@@ -13,7 +13,7 @@ internal static class Step5EntryPoint
 {
     private static CheckpointManager CheckpointManager { get; } = new();
 
-    public static async ValueTask<string> RunAsync(TextWriter writer, Func<string, int> userGuessCallback)
+    public static async ValueTask<string> RunAsync(TextWriter writer, Func<string, int> userGuessCallback, bool rehydrateToRestore = false)
     {
         Workflow<NumberSignal, string> workflow = Step4EntryPoint.CreateWorkflowInstance(out JudgeExecutor judge);
         Checkpointed<StreamingRun<string>> checkpointed =
@@ -30,7 +30,19 @@ internal static class Step5EntryPoint
         checkpoints.Should().HaveCount(6, "we should have two checkpoints, one for each step");
         judge.Tries.Should().Be(2);
 
-        await checkpointed.RestoreCheckpointAsync(checkpoints[2], CancellationToken.None).ConfigureAwait(false);
+        CheckpointInfo targetCheckpoint = checkpoints[2];
+
+        if (rehydrateToRestore)
+        {
+            checkpointed = await InProcessExecution.ResumeStreamAsync(workflow, targetCheckpoint, CheckpointManager, CancellationToken.None)
+                                                   .ConfigureAwait(false);
+            handle = checkpointed.Run;
+        }
+        else
+        {
+            await checkpointed.RestoreCheckpointAsync(checkpoints[2], CancellationToken.None).ConfigureAwait(false);
+        }
+
         judge.Tries.Should().Be(1);
 
         cancellationSource.Dispose();

dotnet/tests/Microsoft.Agents.Workflows.UnitTests/SampleSmokeTest.cs

@@ -103,6 +103,25 @@ public class SampleSmokeTest
         Assert.Equal("You guessed correctly! You Win!", guessResult);
     }
 
+    [Fact]
+    public async Task Test_RunSample_Step5aAsync()
+    {
+        using StringWriter writer = new();
+
+        VerifyingPlaybackResponder<string, int> responder = new(
+            // Iteration 1
+            ("Guess the number.", 50),
+            ("Your guess was too high. Try again.", 23),
+
+            // Iteration 2
+            ("Your guess was too high. Try again.", 23),
+            ("Your guess was too low. Try again.", 42)
+         );
+
+        string guessResult = await Step5EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext, rehydrateToRestore: true);
+        Assert.Equal("You guessed correctly! You Win!", guessResult);
+    }
+
     [Fact]
     public async Task Test_RunSample_Step6Async()
     {