.NET: Release the workflow when a run ends after an executor fails to initialize

dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessExecutionEnvironment.cs

@@ -81,7 +81,16 @@ public async ValueTask<StreamingRun> RunStreamingAsync<TInput>(
         AsyncRunHandle runHandle = await this.BeginRunAsync(workflow, sessionId, [], cancellationToken)
                                              .ConfigureAwait(false);
 
-        return await runHandle.EnqueueAndStreamAsync(input, cancellationToken).ConfigureAwait(false);
+        try
+        {
+            return await runHandle.EnqueueAndStreamAsync(input, cancellationToken).ConfigureAwait(false);
+        }
+        catch
+        {
+            // The caller never receives this run, so end it here to release the workflow.
+            await runHandle.DisposeAsync().ConfigureAwait(false);
+            throw;
+        }
     }
 
     [MemberNotNull(nameof(CheckpointManager))]
@@ -142,11 +151,20 @@ private async ValueTask<AsyncRunHandle> BeginRunHandlingChatProtocolAsync<TInput
         AsyncRunHandle runHandle = await this.BeginRunAsync(workflow, sessionId, descriptor.Accepts, cancellationToken)
                                              .ConfigureAwait(false);
 
-        await runHandle.EnqueueMessageAsync(input, cancellationToken).ConfigureAwait(false);
+        try
+        {
+            await runHandle.EnqueueMessageAsync(input, cancellationToken).ConfigureAwait(false);
 
-        if (descriptor.IsChatProtocol() && input is not TurnToken)
+            if (descriptor.IsChatProtocol() && input is not TurnToken)
+            {
+                await runHandle.EnqueueMessageAsync(new TurnToken(emitEvents: true), cancellationToken).ConfigureAwait(false);
+            }
+        }
+        catch
         {
-            await runHandle.EnqueueMessageAsync(new TurnToken(emitEvents: true), cancellationToken).ConfigureAwait(false);
+            // The caller never receives this run, so end it here to release the workflow.
+            await runHandle.DisposeAsync().ConfigureAwait(false);
+            throw;
         }
 
         return runHandle;

dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessRunnerContext.cs

@@ -531,25 +531,40 @@ public async ValueTask EndRunAsync()
     {
         if (Interlocked.Exchange(ref this._runEnded, 1) == 0)
         {
-            foreach (string executorId in this._executors.Keys)
+            try
             {
-                Task<Executor> executorTask = this._executors[executorId];
-                Executor executor = await executorTask.ConfigureAwait(false);
-
-                if (executor is IAsyncDisposable asyncDisposable)
-                {
-                    await asyncDisposable.DisposeAsync().ConfigureAwait(false);
-                }
-                else if (executor is IDisposable disposable)
+                foreach (string executorId in this._executors.Keys)
                 {
-                    disposable.Dispose();
+                    Task<Executor> executorTask = this._executors[executorId];
+                    Executor executor;
+                    try
+                    {
+                        executor = await executorTask.ConfigureAwait(false);
+                    }
+                    catch (Exception) when (executorTask.IsFaulted || executorTask.IsCanceled)
+                    {
+                        // Creating this executor failed or was cancelled, so there is no instance to dispose. A
+                        // creation failure is reported while the run executes.
+                        continue;
+                    }
+
+                    if (executor is IAsyncDisposable asyncDisposable)
+                    {
+                        await asyncDisposable.DisposeAsync().ConfigureAwait(false);
+                    }
+                    else if (executor is IDisposable disposable)
+                    {
+                        disposable.Dispose();
+                    }
                 }
             }
-
-            if (this._ownsWorkflow)
+            finally
             {
-                await this._workflow.ReleaseOwnershipAsync(this, this._previousOwnership).ConfigureAwait(false);
-                this._ownsWorkflow = false;
+                if (this._ownsWorkflow)
+                {
+                    await this._workflow.ReleaseOwnershipAsync(this, this._previousOwnership).ConfigureAwait(false);
+                    this._ownsWorkflow = false;
+                }
             }
         }
     }

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

@@ -6,6 +6,7 @@
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Agents.AI.Workflows.InProc;
 using Microsoft.Extensions.AI;
 
 namespace Microsoft.Agents.AI.Workflows.UnitTests;
@@ -177,6 +178,211 @@ public async Task RunStreamingAsyncWaitToTakeStreamAsync()
         Assert.NotEmpty(outputEvents);
     }
 
+    /// <summary>
+    /// An executor that fails to initialize must not leave the workflow owned by the failed run, so the same
+    /// workflow can run again once initialization succeeds.
+    /// </summary>
+    [Theory]
+    [InlineData(ExecutionEnvironment.InProcess_Lockstep)]
+    [InlineData(ExecutionEnvironment.InProcess_OffThread)]
+    internal async Task RunAsync_ReleasesWorkflowAfterExecutorInitializationFailsAsync(ExecutionEnvironment environment)
+    {
+        // Arrange
+        InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
+        ForwardMessageExecutor<string> start = new("start");
+        FlakyExecutor<string> target = new("target", FlakyStep.Initialization);
+        Workflow workflow = new WorkflowBuilder(start).AddEdge(start, target).WithOutputFrom(target).Build();
+
+        // Act
+        Run failedRun = await env.RunAsync(workflow, "first");
+        await failedRun.DisposeAsync();
+        await using Run run = await env.RunAsync(workflow, "second");
+
+        // Assert
+        Assert.Contains(failedRun.OutgoingEvents, evt => evt is ExecutorFailedEvent { Data: TimeoutException });
+        Assert.Equal(["second"], run.OutgoingEvents.OfType<WorkflowOutputEvent>().Select(evt => evt.Data));
+    }
+
+    /// <summary>
+    /// Disposing a run while an executor is still initializing cancels the initialization, which must not leave
+    /// the workflow owned by the abandoned run.
+    /// </summary>
+    [Fact]
+    public async Task StreamingRun_ReleasesWorkflowWhenDisposedDuringExecutorInitializationAsync()
+    {
+        // Arrange
+        TokenForwardingExecutor start = new("start");
+        FlakyExecutor<string> target = new("target", FlakyStep.InitializationUntilCancelled);
+        Workflow workflow = new WorkflowBuilder(start).AddEdge(start, target).WithOutputFrom(target).Build();
+
+        // Act
+        StreamingRun abandonedRun = await InProcessExecution.OffThread.RunStreamingAsync(workflow, "first");
+        Task initializing = target.FirstInitializationStarted.Task;
+        Assert.Same(initializing, await Task.WhenAny(initializing, Task.Delay(TimeSpan.FromSeconds(10))));
+        await abandonedRun.DisposeAsync();
+        await using StreamingRun run = await InProcessExecution.OffThread.RunStreamingAsync(workflow, "second");
+
+        // Assert
+        Assert.Equal(["second"], await CollectOutputsAsync(run));
+    }
+
+    /// <summary>
+    /// A run whose executor fails to dispose reports that failure, but must still release the workflow.
+    /// </summary>
+    [Theory]
+    [InlineData(ExecutionEnvironment.InProcess_Lockstep)]
+    [InlineData(ExecutionEnvironment.InProcess_OffThread)]
+    internal async Task RunAsync_ReleasesWorkflowAfterExecutorDisposalFailsAsync(ExecutionEnvironment environment)
+    {
+        // Arrange
+        InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
+        FlakyExecutor<string> start = new("start", FlakyStep.Disposal);
+        Workflow workflow = new WorkflowBuilder(start).WithOutputFrom(start).Build();
+
+        // Act
+        Run failedRun = await env.RunAsync(workflow, "first");
+        await Assert.ThrowsAsync<InvalidOperationException>(() => failedRun.DisposeAsync().AsTask());
+        await using Run run = await env.RunAsync(workflow, "second");
+
+        // Assert
+        Assert.Equal(["second"], run.OutgoingEvents.OfType<WorkflowOutputEvent>().Select(evt => evt.Data));
+    }
+
+    /// <summary>
+    /// When the start executor fails to initialize while the first input is enqueued, the caller never receives
+    /// the run, so the run must be ended before the error propagates.
+    /// </summary>
+    [Theory]
+    [InlineData(ExecutionEnvironment.InProcess_Lockstep)]
+    [InlineData(ExecutionEnvironment.InProcess_OffThread)]
+    internal async Task RunStreamingAsync_ReleasesWorkflowAfterStartExecutorInitializationFailsAsync(ExecutionEnvironment environment)
+    {
+        // Arrange
+        InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
+        FlakyExecutor<string> start = new("start", FlakyStep.Initialization);
+        Workflow workflow = new WorkflowBuilder(start).WithOutputFrom(start).Build();
+
+        // Act
+        await Assert.ThrowsAsync<TimeoutException>(() => env.RunStreamingAsync(workflow, "first").AsTask());
+        await using StreamingRun run = await env.RunStreamingAsync(workflow, "second");
+
+        // Assert
+        Assert.Equal(["second"], await CollectOutputsAsync(run));
+    }
+
+    /// <summary>
+    /// <see cref="InProcessExecutionEnvironment.RunAsync{TInput}(Workflow, TInput, string?, CancellationToken)"/> only
+    /// creates the start executor while enqueueing when the input type is not one of its declared input types; a
+    /// failure there must also end the run before the error propagates.
+    /// </summary>
+    [Theory]
+    [InlineData(ExecutionEnvironment.InProcess_Lockstep)]
+    [InlineData(ExecutionEnvironment.InProcess_OffThread)]
+    internal async Task RunAsync_ReleasesWorkflowAfterStartExecutorInitializationFailsAsync(ExecutionEnvironment environment)
+    {
+        // Arrange
+        InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment();
+        FlakyExecutor<object> start = new("start", FlakyStep.Initialization);
+        Workflow workflow = new WorkflowBuilder(start).WithOutputFrom(start).Build();
+
+        // Act
+        await Assert.ThrowsAsync<TimeoutException>(() => env.RunAsync(workflow, "first").AsTask());
+        await using Run run = await env.RunAsync(workflow, "second");
+
+        // Assert
+        Assert.Equal(["second"], run.OutgoingEvents.OfType<WorkflowOutputEvent>().Select(evt => evt.Data));
+    }
+
+    private static async Task<List<object?>> CollectOutputsAsync(StreamingRun run)
+    {
+        List<object?> outputs = [];
+        await foreach (WorkflowEvent evt in run.WatchStreamAsync())
+        {
+            if (evt is WorkflowOutputEvent output)
+            {
+                outputs.Add(output.Data);
+            }
+        }
+
+        return outputs;
+    }
+
+    /// <summary>
+    /// Forwards every message with the handler's cancellation token, which is also the token the target executor
+    /// is initialized with when this send creates it.
+    /// </summary>
+    private sealed class TokenForwardingExecutor(string id) : Executor(id)
+    {
+        protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBuilder)
+        {
+            protocolBuilder.RouteBuilder.AddHandler<string>(
+                (message, context, cancellationToken) => context.SendMessageAsync(message, cancellationToken: cancellationToken));
+
+            return protocolBuilder.SendsMessage<string>();
+        }
+    }
+
+    private enum FlakyStep
+    {
+        /// <summary>The first initialization throws.</summary>
+        Initialization,
+
+        /// <summary>The first initialization waits until it is cancelled.</summary>
+        InitializationUntilCancelled,
+
+        /// <summary>The first disposal throws.</summary>
+        Disposal,
+    }
+
+    /// <summary>
+    /// Yields every message it receives. The first attempt at the given <see cref="FlakyStep"/> fails.
+    /// </summary>
+    private sealed class FlakyExecutor<TMessage>(string id, FlakyStep flakyStep) : Executor(id), IAsyncDisposable where TMessage : notnull
+    {
+        private int _initializations;
+        private int _disposals;
+
+        public TaskCompletionSource<bool> FirstInitializationStarted { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        protected internal override async ValueTask InitializeAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
+        {
+            await Task.Yield();
+
+            if (Interlocked.Increment(ref this._initializations) > 1)
+            {
+                return;
+            }
+
+            this.FirstInitializationStarted.TrySetResult(true);
+            if (flakyStep == FlakyStep.Initialization)
+            {
+                throw new TimeoutException("Connecting to the backend timed out.");
+            }
+
+            if (flakyStep == FlakyStep.InitializationUntilCancelled)
+            {
+                await Task.Delay(Timeout.Infinite, cancellationToken);
+            }
+        }
+
+        public ValueTask DisposeAsync()
+        {
+            if (flakyStep == FlakyStep.Disposal && Interlocked.Increment(ref this._disposals) == 1)
+            {
+                throw new InvalidOperationException("Closing the backend connection failed.");
+            }
+
+            return default;
+        }
+
+        protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBuilder)
+        {
+            protocolBuilder.RouteBuilder.AddHandler<TMessage>((message, context) => context.YieldOutputAsync(message));
+
+            return protocolBuilder.YieldsOutput<TMessage>();
+        }
+    }
+
     /// <summary>
     /// Simple test agent that echoes back the input message.
     /// </summary>