diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs index 16cd61f6e1..8fadf99539 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs +++ b/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(); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/LockstepRunEventStream.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/LockstepRunEventStream.cs index cdd8cc7686..4daf3fde7a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/LockstepRunEventStream.cs +++ b/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 TakeEventStreamAsync(bool blockOnPendingRequest, [EnumeratorCancellation] CancellationToken cancellationToken = default) { #if NET diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs index 4eb1290961..49f146d46a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs @@ -212,7 +212,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.