.NET: fix: RunState is sometimes Running when it should be Idle

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

@@ -56,6 +56,9 @@ public async IAsyncEnumerable<WorkflowEvent> TakeEventStreamAsync(bool blockOnPe
         using Activity? activity = s_activitySource.StartActivity(ActivityNames.WorkflowRun);
         activity?.SetTag(Tags.WorkflowId, this._stepRunner.StartExecutorId).SetTag(Tags.RunId, this._stepRunner.RunId);
 
+        bool hadException = false;
+        bool hadCancellation = false;
+
         try
         {
             this.RunStatus = RunStatus.Running;
@@ -77,14 +80,21 @@ public async IAsyncEnumerable<WorkflowEvent> TakeEventStreamAsync(bool blockOnPe
                     }
                     catch (OperationCanceledException)
                     {
+                        hadCancellation = true;
                     }
-                    catch (Exception ex) when (activity is not null)
+                    catch (Exception ex)
                     {
-                        activity.AddEvent(new ActivityEvent(EventNames.WorkflowError, tags: new() {
-                             { Tags.ErrorType, ex.GetType().FullName },
-                             { Tags.BuildErrorMessage, ex.Message },
-                        }));
-                        activity.CaptureException(ex);
+                        hadException = true;
+
+                        if (activity != null)
+                        {
+                            activity.AddEvent(new ActivityEvent(EventNames.WorkflowError, tags: new() {
+                                 { Tags.ErrorType, ex.GetType().FullName },
+                                 { Tags.BuildErrorMessage, ex.Message },
+                            }));
+                            activity.CaptureException(ex);
+                        }
+
                         throw;
                     }
 
@@ -136,7 +146,19 @@ public async IAsyncEnumerable<WorkflowEvent> TakeEventStreamAsync(bool blockOnPe
         }
         finally
         {
-            this.RunStatus = this._stepRunner.HasUnservicedRequests ? RunStatus.PendingRequests : RunStatus.Idle;
+            if (hadException || hadCancellation || linkedSource.Token.IsCancellationRequested)
+            {
+                this.RunStatus = RunStatus.Ended;
+            }
+            else if (this._stepRunner.HasUnservicedRequests)
+            {
+                this.RunStatus = RunStatus.PendingRequests;
+            }
+            else
+            {
+                this.RunStatus = RunStatus.Idle;
+            }
+
             this._stepRunner.OutgoingEvents.EventRaised -= OnWorkflowEventAsync;
         }
 

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

@@ -72,7 +72,6 @@ private async Task RunLoopAsync(CancellationToken cancellationToken)
             // The consumer will call EnqueueMessageAsync which signals the run loop
             await this._inputWaiter.WaitForInputAsync(cancellationToken: linkedSource.Token).ConfigureAwait(false);
 
-            this._runStatus = RunStatus.Running;
             activity?.AddEvent(new ActivityEvent(EventNames.WorkflowStarted));
 
             while (!linkedSource.Token.IsCancellationRequested)
@@ -81,6 +80,9 @@ private async Task RunLoopAsync(CancellationToken cancellationToken)
                 // Events are streamed out in real-time as they happen via the event handler
                 while (this._stepRunner.HasUnprocessedMessages && !linkedSource.Token.IsCancellationRequested)
                 {
+                    this._runStatus = RunStatus.Running; // Ensure we do not inappropriately signal Running status unless
+                                                         // messages are being processed.
+
                     await this._stepRunner.RunSuperStepAsync(linkedSource.Token).ConfigureAwait(false);
                 }
 
@@ -99,9 +101,6 @@ private async Task RunLoopAsync(CancellationToken cancellationToken)
                 // Wait for next input from the consumer
                 // Works for both Idle (no work) and PendingRequests (waiting for responses)
                 await this._inputWaiter.WaitForInputAsync(TimeSpan.FromSeconds(1), linkedSource.Token).ConfigureAwait(false);
-
-                // When signaled, resume running
-                this._runStatus = RunStatus.Running;
             }
         }
         catch (OperationCanceledException)
@@ -122,11 +121,12 @@ private async Task RunLoopAsync(CancellationToken cancellationToken)
         }
         finally
         {
+            // Mark as ended when run loop exits
+            this._runStatus = RunStatus.Ended;
+
             this._stepRunner.OutgoingEvents.EventRaised -= OnEventRaisedAsync;
             this._eventChannel.Writer.Complete();
 
-            // Mark as ended when run loop exits
-            this._runStatus = RunStatus.Ended;
             activity?.AddEvent(new ActivityEvent(EventNames.WorkflowCompleted));
         }
 

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/DelayValueTaskSource.cs

@@ -0,0 +1,29 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Threading.Tasks;
+using System.Threading.Tasks.Sources;
+
+namespace Microsoft.Agents.AI.Workflows.UnitTests;
+
+internal sealed class DelayValueTaskSource<T> : IValueTaskSource<T>
+{
+    private readonly TestValueTaskSource<T> _innerSource = new();
+    private readonly T _value;
+
+    public DelayValueTaskSource(T value)
+    {
+        this._value = value;
+    }
+
+    public ValueTask ReleaseSucceededAsync() => this._innerSource.SetSucceededAsync(this._value);
+    public ValueTask ReleaseFaultedAsync(Exception exception) => this._innerSource.SetFaultedAsync(exception);
+    public ValueTask ReleaseCanceledAsync() => this._innerSource.SetCanceledAsync();
+
+    public T GetResult(short token) => this._innerSource.GetResult(token);
+
+    public ValueTaskSourceStatus GetStatus(short token) => this._innerSource.GetStatus(token);
+
+    public void OnCompleted(Action<object?> continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags)
+        => this._innerSource.OnCompleted(continuation, state, token, flags);
+}