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. ///