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);
     }
 }