Merge branch 'main' into dev/dotnet_workflow/fix_flaky_checkpoint_restore_test

dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs

@@ -186,6 +186,8 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable
 
     public async ValueTask RestoreCheckpointAsync(CheckpointInfo checkpointInfo, CancellationToken cancellationToken = default)
     {
+        LockstepRunEventStream? lockstepEventStream = null;
+
         // Clear buffered events from the channel BEFORE restoring to discard stale events from supersteps
         // that occurred after the checkpoint we're restoring to
         // This must happen BEFORE the restore so that events republished during restore aren't cleared
@@ -193,8 +195,9 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable
         {
             streamingEventStream.ClearBufferedEvents();
         }
-        else if (this._eventStream is LockstepRunEventStream lockstepEventStream)
+        else if (this._eventStream is LockstepRunEventStream lES)
         {
+            lockstepEventStream = lES;
             lockstepEventStream.ClearBufferedEvents();
         }
 
@@ -202,6 +205,8 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable
         // This can re-emit pending requests into the already-active event stream.
         await this._checkpointingHandle.RestoreCheckpointAsync(checkpointInfo, cancellationToken).ConfigureAwait(false);
 
+        lockstepEventStream?.UpdateStatus();
+
         // After restore, signal the run loop to process any restored messages. Initial resume
         // paths handle this separately when they create the event stream after restoring state.
         this.SignalInputToRunLoop();

dotnet/src/Microsoft.Agents.AI.Workflows/Execution/LockstepRunEventStream.cs

@@ -46,6 +46,17 @@ internal sealed class LockstepRunEventStream : IRunEventStream
         Activity.Current = previousActivity;
     }
 
+    internal void UpdateStatus()
+    {
+        switch (this.RunStatus)
+        {
+            case RunStatus.Idle:
+            case RunStatus.PendingRequests:
+                this.RunStatus = this._stepRunner.HasUnservicedRequests ? RunStatus.PendingRequests : RunStatus.Idle;
+                break;
+        }
+    }
+
     public async IAsyncEnumerable<WorkflowEvent> TakeEventStreamAsync(bool blockOnPendingRequest, [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
 #if NET

dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs

@@ -217,7 +217,11 @@ internal sealed class StreamingRunEventStream : IRunEventStream
         // Get the current epoch - we'll only respond to completion signals from this epoch or later
         int currentEpoch = Volatile.Read(ref this._completionEpoch);
 
-        bool expectingFreshWork = this._stepRunner.HasUnprocessedMessages || this._runStatus == RunStatus.Running;
+        bool expectingFreshWork = this._stepRunner.HasUnprocessedMessages ||
+                                  this._stepRunner.HasUnservicedRequests ||
+                                  this._runStatus == RunStatus.Running ||
+                                  this._runStatus == RunStatus.PendingRequests;
+
         int myEpoch = expectingFreshWork ? currentEpoch + 1 : currentEpoch;
 
         // Use custom async enumerable to avoid exceptions on cancellation.