From c83011b30d91cadac8d688be7579261dbe604dc3 Mon Sep 17 00:00:00 2001 From: Peter Ibekwe <109177538+peibekwe@users.noreply.github.com> Date: Mon, 3 Nov 2025 08:34:55 -0800 Subject: [PATCH] .NET: Peibekwe/workflows cancellation token fix (#1740) * Propagate cancellation token down the stack * Added unit tests to cover workflow cancellation scenarios * Updated tests based on feedback to simplify assert. * Create custom AsyncEnumrable to gracefully handle cancellation for Channel reader. Tailor cancellation tests to declarative scenarios. * Update comment and naming for readability. * Fixing minor stylistic recommendation. --------- Co-authored-by: Chris <66376200+crickman@users.noreply.github.com> --- ...NonThrowingChannelReaderAsyncEnumerable.cs | 61 ++++++++++++++++++ .../Execution/StreamingRunEventStream.cs | 8 +-- .../DeclarativeWorkflowTest.cs | 63 +++++++++++++++++-- 3 files changed, 121 insertions(+), 11 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows/Execution/NonThrowingChannelReaderAsyncEnumerable.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/NonThrowingChannelReaderAsyncEnumerable.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/NonThrowingChannelReaderAsyncEnumerable.cs new file mode 100644 index 0000000000..aaae42f2f1 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/NonThrowingChannelReaderAsyncEnumerable.cs @@ -0,0 +1,61 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Channels; +using System.Threading.Tasks; + +namespace Microsoft.Agents.AI.Workflows.Execution; + +/// +/// A custom IAsyncEnumerable implementation that reads from a ChannelReader, +/// and suppresses OperationCanceledException when the cancellation token is triggered. +/// +internal sealed class NonThrowingChannelReaderAsyncEnumerable(ChannelReader reader) : IAsyncEnumerable +{ + private class Enumerator(ChannelReader reader, CancellationToken cancellationToken) : IAsyncEnumerator + { + private T? _current; + public T Current => this._current ?? throw new InvalidOperationException("Enumeration not started."); + + public ValueTask DisposeAsync() + { + // no-op - the reader should not be disposed. + return default; + } + + /// + /// Moves to the next item in the channel. + /// + /// If successful, returns true, otherwise false. + public async ValueTask MoveNextAsync() + { + try + { + bool hasData = await reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false); + if (hasData) + { + this._current = await reader.ReadAsync(cancellationToken).ConfigureAwait(false); + return true; + } + } + catch (OperationCanceledException) + { + // Swallow cancellation exceptions to prevent throwing from the enumerator + // Enables clean cancellation and aligns with the expected behavior of IAsyncEnumerable. + } + + return false; + } + } + + /// + /// Returns an async enumerator that reads items from the channel. + /// If cancellation is requested, the enumeration exits silently without throwing. + /// + /// An optional cancellation token from the caller. + /// An async enumerator over the channel items. + public IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) + => new Enumerator(reader, cancellationToken); +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs index 718ebcd11c..e6afbb5440 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StreamingRunEventStream.cs @@ -139,11 +139,9 @@ internal sealed class StreamingRunEventStream : IRunEventStream // Get the current epoch - we'll only respond to completion signals from this epoch or later int myEpoch = Volatile.Read(ref this._completionEpoch) + 1; - // Simply read from channel - all coordination is handled by Channel infrastructure - // Note: When cancellation is requested, ReadAllAsync may throw OperationCanceledException - // or may complete the enumeration. We check IsCancellationRequested explicitly at superstep - // boundaries to ensure clean cancellation. - await foreach (WorkflowEvent evt in this._eventChannel.Reader.ReadAllAsync(cancellationToken).ConfigureAwait(false)) + // Use custom async enumerable to avoid exceptions on cancellation. + NonThrowingChannelReaderAsyncEnumerable eventStream = new(this._eventChannel.Reader); + await foreach (WorkflowEvent evt in eventStream.WithCancellation(cancellationToken).ConfigureAwait(false)) { // Filter out internal signals used for run loop coordination if (evt is InternalHaltSignal completionSignal) diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/DeclarativeWorkflowTest.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/DeclarativeWorkflowTest.cs index ae650352c5..6f2282902f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/DeclarativeWorkflowTest.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/DeclarativeWorkflowTest.cs @@ -225,6 +225,54 @@ public sealed class DeclarativeWorkflowTest(ITestOutputHelper output) : Workflow Assert.True(visitor.HasUnsupportedActions); } + [Theory] + [InlineData("CaseInsensitive.yaml", "end_when_match")] + [InlineData("ClearAllVariables.yaml", "clear_all")] + [InlineData("Condition.yaml", "setVariable_test")] + [InlineData("ConditionElse.yaml", "setVariable_test")] + [InlineData("EndConversation.yaml", "end_all")] + [InlineData("EndDialog.yaml", "end_all")] + [InlineData("EditTable.yaml", "edit_var")] + [InlineData("EditTableV2.yaml", "edit_var")] + [InlineData("Goto.yaml", "goto_end")] + [InlineData("LoopBreak.yaml", "break_loop_now")] + [InlineData("LoopContinue.yaml", "foreach_loop")] + [InlineData("LoopEach.yaml", "foreach_loop")] + [InlineData("MixedScopes.yaml", "activity_input")] + [InlineData("ParseValue.yaml", "parse_var")] + [InlineData("ParseValueList.yaml", "parse_var")] + [InlineData("ResetVariable.yaml", "clear_var")] + [InlineData("SendActivity.yaml", "activity_input")] + [InlineData("SetVariable.yaml", "set_var")] + [InlineData("SetTextVariable.yaml", "set_text")] + public async Task CancelRunAsync(string workflowPath, string expectedExecutedId) + { + // Arrange + const string WorkflowInput = "Test input message"; + Workflow workflow = this.CreateWorkflow(workflowPath, WorkflowInput); + await using StreamingRun run = await InProcessExecution.StreamAsync(workflow: workflow, input: WorkflowInput); + + // Act + await foreach (WorkflowEvent workflowEvent in run.WatchStreamAsync()) + { + this.WorkflowEvents.Add(workflowEvent); + + if (workflowEvent is DeclarativeActionInvokedEvent actionInvokedEvent && actionInvokedEvent.ActionId == expectedExecutedId) + { + // Cancel run after the specified declarative action is invoked. + await run.CancelRunAsync(); + } + } + RunStatus currentRunStatus = await run.GetStatusAsync(); + this.WorkflowEventCounts = this.WorkflowEvents.GroupBy(e => e.GetType()).ToDictionary(e => e.Key, e => e.Count()); + + // Assert + Assert.Equal(expected: RunStatus.Ended, actual: currentRunStatus); + Assert.NotEmpty(this.WorkflowEventCounts); + Assert.Contains(this.WorkflowEvents.OfType(), e => e.ActionId == expectedExecutedId); + Assert.DoesNotContain(this.WorkflowEvents.OfType(), e => e.ActionId == expectedExecutedId); + } + private void AssertExecutionCount(int expectedCount) { Assert.Equal(expectedCount + 2, this.WorkflowEventCounts[typeof(ExecutorInvokedEvent)]); @@ -256,12 +304,7 @@ public sealed class DeclarativeWorkflowTest(ITestOutputHelper output) : Workflow private async Task RunWorkflowAsync(string workflowPath, TInput workflowInput) where TInput : notnull { - using StreamReader yamlReader = File.OpenText(Path.Combine("Workflows", workflowPath)); - Mock mockAgentProvider = CreateMockProvider($"{workflowInput}"); - DeclarativeWorkflowOptions workflowContext = new(mockAgentProvider.Object) { LoggerFactory = this.Output }; - - Workflow workflow = DeclarativeWorkflowBuilder.Build(yamlReader, workflowContext); - + Workflow workflow = this.CreateWorkflow(workflowPath, workflowInput); await using StreamingRun run = await InProcessExecution.StreamAsync(workflow, workflowInput); await foreach (WorkflowEvent workflowEvent in run.WatchStreamAsync()) @@ -303,6 +346,14 @@ public sealed class DeclarativeWorkflowTest(ITestOutputHelper output) : Workflow this.WorkflowEventCounts = this.WorkflowEvents.GroupBy(e => e.GetType()).ToDictionary(e => e.Key, e => e.Count()); } + private Workflow CreateWorkflow(string workflowPath, TInput workflowInput) where TInput : notnull + { + using StreamReader yamlReader = File.OpenText(Path.Combine("Workflows", workflowPath)); + Mock mockAgentProvider = CreateMockProvider($"{workflowInput}"); + DeclarativeWorkflowOptions workflowContext = new(mockAgentProvider.Object) { LoggerFactory = this.Output }; + return DeclarativeWorkflowBuilder.Build(yamlReader, workflowContext); + } + private static Mock CreateMockProvider(string input) { Mock mockAgentProvider = new(MockBehavior.Strict);