[BREAKING] .NET: Workflow Off-Thread Execution Mode (#1233)

* Updates to async run loop.

* fix: Workflow Onwership can be release by nonowner

* fix: Incorrect handling of blockOnPending in StreamingRun

Depending on whether we are running in streaming on non-streaming mode, we may be using the StreamingRun in different ways. Unfortunately, the only place we can really know what is the actual state of execution is in the RunEventStream implementations.

This resulted in blocking where blocking was unneeded and occasionally not-blocking when blocking was needed.

The fix is to move the logic of handling this blocking into RunEventStream implementations.

* fix: Fix cleanup on error and end run

This ensures we clean up the background resources correctly.

* fix: Ensure we let the run loop proceed when shutting down

* fix: Add timeout for Input Waiting

* fix: Make the samples properly clean up `Run`s and `StreamingRun`s

* fix: Simplify Declarative Workflow Run disposal pattern

* Also fixes missing .Disposal() in Integration tests

---------

Co-authored-by: Ben Thomas <ben.thomas@microsoft.com>

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

@@ -11,13 +11,12 @@ using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Agents.AI.Workflows.Execution;
 
-internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, IInputCoordinator
+internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable
 {
-    private readonly AsyncCoordinator _waitForResponseCoordinator = new();
     private readonly ISuperStepRunner _stepRunner;
     private readonly ICheckpointingHandle _checkpointingHandle;
 
-    private readonly LockstepRunEventStream _eventStream;
+    private readonly IRunEventStream _eventStream;
     private readonly CancellationTokenSource _endRunSource = new();
     private int _isDisposed;
     private int _isEventStreamTaken;
@@ -29,17 +28,26 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, I
 
         this._eventStream = mode switch
         {
-            //ExecutionMode.OffThread => Not supported yet
+            ExecutionMode.OffThread => new StreamingRunEventStream(stepRunner),
+            ExecutionMode.Subworkflow => new StreamingRunEventStream(stepRunner, disableRunLoop: true),
             ExecutionMode.Lockstep => new LockstepRunEventStream(stepRunner),
             _ => throw new ArgumentOutOfRangeException(nameof(mode), $"Unknown execution mode {mode}")
         };
+
         this._eventStream.Start();
+
+        // If there are already unprocessed messages (e.g., from a checkpoint restore that happened
+        // before this handle was created), signal the run loop to start processing them
+        if (stepRunner.HasUnprocessedMessages)
+        {
+            this.SignalInputToRunLoop();
+        }
     }
 
-    public ValueTask<bool> WaitForNextInputAsync(CancellationToken cancellation = default)
-        => this._waitForResponseCoordinator.WaitForCoordinationAsync(cancellation);
+    //private readonly AsyncCoordinator _waitForResponseCoordinator = new();
 
-    public void ReleaseResponseWaiter() => this._waitForResponseCoordinator.MarkCoordinationPoint();
+    //public ValueTask<bool> WaitForNextInputAsync(CancellationToken cancellation = default)
+    //    => this._waitForResponseCoordinator.WaitForCoordinationAsync(cancellation);
 
     public string RunId => this._stepRunner.RunId;
 
@@ -48,28 +56,39 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, I
     public ValueTask<RunStatus> GetStatusAsync(CancellationToken cancellation = default)
         => this._eventStream.GetStatusAsync(cancellation);
 
-    public async IAsyncEnumerable<WorkflowEvent> TakeEventStreamAsync(bool breakOnHalt, [EnumeratorCancellation] CancellationToken cancellation = default)
+    public async IAsyncEnumerable<WorkflowEvent> TakeEventStreamAsync(bool blockOnPendingRequest, [EnumeratorCancellation] CancellationToken cancellation = default)
     {
-        // Create a linked cancellation token that combines the provided token with the end-run token
-        using var linkedSource = CancellationTokenSource.CreateLinkedTokenSource(cancellation, this._endRunSource.Token);
-
-        // Only one enumerator of this is allowed at a time
+        //Debug.Assert(breakOnHalt);
+        // Enforce single active enumerator (this runs when enumeration begins)
         if (Interlocked.CompareExchange(ref this._isEventStreamTaken, 1, 0) != 0)
         {
             throw new InvalidOperationException("The event stream has already been taken. Only one enumerator is allowed at a time.");
         }
 
+        CancellationTokenSource? linked = null;
         try
         {
-            await foreach (WorkflowEvent @event in this._eventStream.TakeEventStreamAsync(linkedSource.Token)
-                                                                    .ConfigureAwait(false))
+            linked = CancellationTokenSource.CreateLinkedTokenSource(cancellation, this._endRunSource.Token);
+            var token = linked.Token;
+
+            // Build the inner stream before the loop so synchronous exceptions still release the gate
+            var inner = this._eventStream.TakeEventStreamAsync(blockOnPendingRequest, token);
+
+            await foreach (var ev in inner.WithCancellation(token).ConfigureAwait(false))
             {
-                yield return @event;
+                // Filter out the RequestHaltEvent, since it is an internal signalling event.
+                if (ev is RequestHaltEvent)
+                {
+                    yield break;
+                }
+
+                yield return ev;
             }
         }
         finally
         {
-            Volatile.Write(ref this._isEventStreamTaken, 0);
+            linked?.Dispose();
+            Interlocked.Exchange(ref this._isEventStreamTaken, 0);
         }
     }
 
@@ -80,7 +99,7 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, I
     {
         if (message is ExternalResponse response)
         {
-            // EnqueueResponseAsync marks the coordination point itself
+            // EnqueueResponseAsync handles signaling
             await this.EnqueueResponseAsync(response, cancellation)
                       .ConfigureAwait(false);
 
@@ -90,7 +109,8 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, I
         bool result = await this._stepRunner.EnqueueMessageAsync(message, cancellation)
                                             .ConfigureAwait(false);
 
-        this._waitForResponseCoordinator.MarkCoordinationPoint();
+        // Signal the run loop that new input is available
+        this.SignalInputToRunLoop();
 
         return result;
     }
@@ -104,7 +124,7 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, I
 
         if (declaredType != null && typeof(ExternalResponse).IsAssignableFrom(declaredType))
         {
-            // EnqueueResponseAsync marks the coordination point itself
+            // EnqueueResponseAsync handles signaling
             await this.EnqueueResponseAsync((ExternalResponse)message, cancellation)
                       .ConfigureAwait(false);
 
@@ -112,7 +132,7 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, I
         }
         else if (declaredType == null && message is ExternalResponse response)
         {
-            // EnqueueResponseAsync marks the coordination point itself
+            // EnqueueResponseAsync handles signaling
             await this.EnqueueResponseAsync(response, cancellation)
                       .ConfigureAwait(false);
 
@@ -122,7 +142,8 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, I
         bool result = await this._stepRunner.EnqueueMessageUntypedAsync(message, declaredType ?? message.GetType(), cancellation)
                                             .ConfigureAwait(false);
 
-        this._waitForResponseCoordinator.MarkCoordinationPoint();
+        // Signal the run loop that new input is available
+        this.SignalInputToRunLoop();
 
         return result;
     }
@@ -131,13 +152,13 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, I
     {
         await this._stepRunner.EnqueueResponseAsync(response, cancellation).ConfigureAwait(false);
 
-        this._waitForResponseCoordinator.MarkCoordinationPoint();
+        // Signal the run loop that new input is available
+        this.SignalInputToRunLoop();
     }
 
-    public ValueTask RequestEndRunAsync()
+    private void SignalInputToRunLoop()
     {
-        this._endRunSource.Cancel();
-        return this._stepRunner.RequestEndRunAsync();
+        this._eventStream.SignalInput();
     }
 
     public async ValueTask DisposeAsync()
@@ -145,13 +166,32 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable, I
         if (Interlocked.Exchange(ref this._isDisposed, 1) == 0)
         {
             this._endRunSource.Cancel();
-            await this.RequestEndRunAsync().ConfigureAwait(false);
+
+            await this._eventStream.StopAsync().ConfigureAwait(false);
+            await this._stepRunner.RequestEndRunAsync().ConfigureAwait(false);
+
             this._endRunSource.Dispose();
 
             await this._eventStream.DisposeAsync().ConfigureAwait(false);
         }
     }
 
-    public ValueTask RestoreCheckpointAsync(CheckpointInfo checkpointInfo, CancellationToken cancellationToken = default)
-        => this._checkpointingHandle.RestoreCheckpointAsync(checkpointInfo, cancellationToken);
+    public async ValueTask RestoreCheckpointAsync(CheckpointInfo checkpointInfo, CancellationToken cancellationToken = default)
+    {
+        // 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
+        if (this._eventStream is StreamingRunEventStream streamingEventStream)
+        {
+            streamingEventStream.ClearBufferedEvents();
+        }
+
+        // Restore the workflow state - this will republish unserviced requests as new events
+        await this._checkpointingHandle.RestoreCheckpointAsync(checkpointInfo, cancellationToken).ConfigureAwait(false);
+
+        // After restore, signal the run loop to process any restored messages
+        // This is necessary because ClearBufferedEvents() doesn't signal, and the restored
+        // queued messages won't automatically wake up the run loop
+        this.SignalInputToRunLoop();
+    }
 }