From ea789d5e8716f63aa8c5745ae631bb494eb50c42 Mon Sep 17 00:00:00 2001 From: RachelWanggg Date: Wed, 21 Oct 2026 10:10:40 -0700 Subject: [PATCH] .NET: Release the workflow when a run ends after an executor fails to initialize EndRunAsync awaited every executor creation task before releasing ownership, so a factory or InitializeAsync that threw or was cancelled made disposing the run rethrow that error and left the Workflow owned: every later run on it failed with "already owned by another runner". Skip executors that were never created, and release ownership in a finally so a failing executor dispose cannot keep it either. RunStreamingAsync and RunAsync also dropped the run handle when enqueueing the first input threw (the start executor is created there), so the run was never ended. Dispose the handle before rethrowing. --- .../InProc/InProcessExecutionEnvironment.cs | 26 ++- .../InProc/InProcessRunnerContext.cs | 43 ++-- .../InProcessExecutionTests.cs | 206 ++++++++++++++++++ 3 files changed, 257 insertions(+), 18 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessExecutionEnvironment.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessExecutionEnvironment.cs index d08c23c0899..d14b49e4d93 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessExecutionEnvironment.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/InProc/InProcessExecutionEnvironment.cs @@ -81,7 +81,16 @@ public async ValueTask RunStreamingAsync( 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 BeginRunHandlingChatProtocolAsync 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 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; + } } } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessExecutionTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessExecutionTests.cs index 6c7e19fd4b1..79be3fc1114 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessExecutionTests.cs +++ b/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); } + /// + /// 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. + /// + [Theory] + [InlineData(ExecutionEnvironment.InProcess_Lockstep)] + [InlineData(ExecutionEnvironment.InProcess_OffThread)] + internal async Task RunAsync_ReleasesWorkflowAfterExecutorInitializationFailsAsync(ExecutionEnvironment environment) + { + // Arrange + InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment(); + ForwardMessageExecutor start = new("start"); + FlakyExecutor 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().Select(evt => evt.Data)); + } + + /// + /// Disposing a run while an executor is still initializing cancels the initialization, which must not leave + /// the workflow owned by the abandoned run. + /// + [Fact] + public async Task StreamingRun_ReleasesWorkflowWhenDisposedDuringExecutorInitializationAsync() + { + // Arrange + TokenForwardingExecutor start = new("start"); + FlakyExecutor 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)); + } + + /// + /// A run whose executor fails to dispose reports that failure, but must still release the workflow. + /// + [Theory] + [InlineData(ExecutionEnvironment.InProcess_Lockstep)] + [InlineData(ExecutionEnvironment.InProcess_OffThread)] + internal async Task RunAsync_ReleasesWorkflowAfterExecutorDisposalFailsAsync(ExecutionEnvironment environment) + { + // Arrange + InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment(); + FlakyExecutor start = new("start", FlakyStep.Disposal); + Workflow workflow = new WorkflowBuilder(start).WithOutputFrom(start).Build(); + + // Act + Run failedRun = await env.RunAsync(workflow, "first"); + await Assert.ThrowsAsync(() => failedRun.DisposeAsync().AsTask()); + await using Run run = await env.RunAsync(workflow, "second"); + + // Assert + Assert.Equal(["second"], run.OutgoingEvents.OfType().Select(evt => evt.Data)); + } + + /// + /// 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. + /// + [Theory] + [InlineData(ExecutionEnvironment.InProcess_Lockstep)] + [InlineData(ExecutionEnvironment.InProcess_OffThread)] + internal async Task RunStreamingAsync_ReleasesWorkflowAfterStartExecutorInitializationFailsAsync(ExecutionEnvironment environment) + { + // Arrange + InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment(); + FlakyExecutor start = new("start", FlakyStep.Initialization); + Workflow workflow = new WorkflowBuilder(start).WithOutputFrom(start).Build(); + + // Act + await Assert.ThrowsAsync(() => env.RunStreamingAsync(workflow, "first").AsTask()); + await using StreamingRun run = await env.RunStreamingAsync(workflow, "second"); + + // Assert + Assert.Equal(["second"], await CollectOutputsAsync(run)); + } + + /// + /// 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. + /// + [Theory] + [InlineData(ExecutionEnvironment.InProcess_Lockstep)] + [InlineData(ExecutionEnvironment.InProcess_OffThread)] + internal async Task RunAsync_ReleasesWorkflowAfterStartExecutorInitializationFailsAsync(ExecutionEnvironment environment) + { + // Arrange + InProcessExecutionEnvironment env = environment.ToWorkflowExecutionEnvironment(); + FlakyExecutor start = new("start", FlakyStep.Initialization); + Workflow workflow = new WorkflowBuilder(start).WithOutputFrom(start).Build(); + + // Act + await Assert.ThrowsAsync(() => env.RunAsync(workflow, "first").AsTask()); + await using Run run = await env.RunAsync(workflow, "second"); + + // Assert + Assert.Equal(["second"], run.OutgoingEvents.OfType().Select(evt => evt.Data)); + } + + private static async Task> CollectOutputsAsync(StreamingRun run) + { + List outputs = []; + await foreach (WorkflowEvent evt in run.WatchStreamAsync()) + { + if (evt is WorkflowOutputEvent output) + { + outputs.Add(output.Data); + } + } + + return outputs; + } + + /// + /// 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. + /// + private sealed class TokenForwardingExecutor(string id) : Executor(id) + { + protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBuilder) + { + protocolBuilder.RouteBuilder.AddHandler( + (message, context, cancellationToken) => context.SendMessageAsync(message, cancellationToken: cancellationToken)); + + return protocolBuilder.SendsMessage(); + } + } + + private enum FlakyStep + { + /// The first initialization throws. + Initialization, + + /// The first initialization waits until it is cancelled. + InitializationUntilCancelled, + + /// The first disposal throws. + Disposal, + } + + /// + /// Yields every message it receives. The first attempt at the given fails. + /// + private sealed class FlakyExecutor(string id, FlakyStep flakyStep) : Executor(id), IAsyncDisposable where TMessage : notnull + { + private int _initializations; + private int _disposals; + + public TaskCompletionSource 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((message, context) => context.YieldOutputAsync(message)); + + return protocolBuilder.YieldsOutput(); + } + } + /// /// Simple test agent that echoes back the input message. ///