.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>
2026-06-16 21:04:09 +08:00 · 2025-11-03 08:34:55 -08:00
parent e87eed573b
commit c83011b30d
3 changed files with 121 additions and 11 deletions
@@ -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;
+
+/// <summary>
+/// A custom IAsyncEnumerable implementation that reads from a ChannelReader,
+/// and suppresses OperationCanceledException when the cancellation token is triggered.
+/// </summary>
+internal sealed class NonThrowingChannelReaderAsyncEnumerable<T>(ChannelReader<T> reader) : IAsyncEnumerable<T>
+{
+    private class Enumerator(ChannelReader<T> reader, CancellationToken cancellationToken) : IAsyncEnumerator<T>
+    {
+        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;
+        }
+
+        /// <summary>
+        /// Moves to the next item in the channel.
+        /// </summary>
+        /// <returns>If successful, returns <c>true</c>, otherwise <c>false</c>.</returns>
+        public async ValueTask<bool> 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;
+        }
+    }
+
+    /// <summary>
+    /// Returns an async enumerator that reads items from the channel.
+    /// If cancellation is requested, the enumeration exits silently without throwing.
+    /// </summary>
+    /// <param name="cancellationToken">An optional cancellation token from the caller.</param>
+    /// <returns>An async enumerator over the channel items.</returns>
+    public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+        => new Enumerator(reader, cancellationToken);
+}
@@ -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<WorkflowEvent> 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)
@@ -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<DeclarativeActionInvokedEvent>(), e => e.ActionId == expectedExecutedId);
+        Assert.DoesNotContain(this.WorkflowEvents.OfType<DeclarativeActionCompletedEvent>(), 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<TInput>(string workflowPath, TInput workflowInput) where TInput : notnull
    {
-        using StreamReader yamlReader = File.OpenText(Path.Combine("Workflows", workflowPath));
-        Mock<WorkflowAgentProvider> mockAgentProvider = CreateMockProvider($"{workflowInput}");
-        DeclarativeWorkflowOptions workflowContext = new(mockAgentProvider.Object) { LoggerFactory = this.Output };
-
-        Workflow workflow = DeclarativeWorkflowBuilder.Build<TInput>(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<TInput>(string workflowPath, TInput workflowInput) where TInput : notnull
+    {
+        using StreamReader yamlReader = File.OpenText(Path.Combine("Workflows", workflowPath));
+        Mock<WorkflowAgentProvider> mockAgentProvider = CreateMockProvider($"{workflowInput}");
+        DeclarativeWorkflowOptions workflowContext = new(mockAgentProvider.Object) { LoggerFactory = this.Output };
+        return DeclarativeWorkflowBuilder.Build<TInput>(yamlReader, workflowContext);
+    }
+
    private static Mock<WorkflowAgentProvider> CreateMockProvider(string input)
    {
        Mock<WorkflowAgentProvider> mockAgentProvider = new(MockBehavior.Strict);