From a8c7a30d6f8d9fad63dc13e915e0c6c1810d33af Mon Sep 17 00:00:00 2001 From: Naveen Chatlapalli Date: Thu, 17 Sep 2026 21:16:38 -0500 Subject: [PATCH 1/2] fix: preserve cooperative workflow handler cancellation --- .../Microsoft.Agents.AI.Workflows/Executor.cs | 6 + .../ExecutorCancellationTests.cs | 159 ++++++++++++++++++ 2 files changed, 165 insertions(+) create mode 100644 dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ExecutorCancellationTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs index 9f092e8e887..2d9beb1b647 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs @@ -265,6 +265,12 @@ protected internal virtual ValueTask InitializeAsync(IWorkflowContext context, C .ConfigureAwait(false); ExecutorEvent executionResult; + if (result is { IsSuccess: false } && cancellationToken.IsCancellationRequested && + (result.IsCancelled || result.Exception is OperationCanceledException)) + { + cancellationToken.ThrowIfCancellationRequested(); + } + if (result?.IsSuccess is not false) { executionResult = new ExecutorCompletedEvent(this.Id, result?.Result); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ExecutorCancellationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ExecutorCancellationTests.cs new file mode 100644 index 00000000000..b08e1bb5bb2 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ExecutorCancellationTests.cs @@ -0,0 +1,159 @@ +// Copyright (c) Microsoft. All rights reserved. + +#pragma warning disable CS0618 // Verify cancellation in the supported legacy reflection path too. + +using System; +using System.Collections.Generic; +using System.Reflection; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Agents.AI.Workflows.Reflection; + +namespace Microsoft.Agents.AI.Workflows.UnitTests; + +public class ExecutorCancellationTests +{ + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task RuntimeCancellationDoesNotEmitFailureAsync(bool useReflection) + { + // Arrange + using CancellationTokenSource source = new(); + async ValueTask CancelAsync(string message, IWorkflowContext context, CancellationToken cancellationToken) + { + source.Cancel(); + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + } + + Executor executor = useReflection + ? new ReflectingHandler(CancelAsync) + : new FunctionExecutor("cancel", CancelAsync); + TestWorkflowContext context = new(executor.Id); + + // Act + OperationCanceledException exception = await Assert.ThrowsAnyAsync( + () => executor.ExecuteCoreAsync("input", new(typeof(string)), context, source.Token).AsTask()); + + // Assert + Assert.Equal(source.Token, exception.CancellationToken); + Assert.DoesNotContain(context.EmittedEvents, evt => evt is ExecutorFailedEvent or ExecutorCompletedEvent); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task OrdinaryFailureRemainsFailureWhenRuntimeIsCancelledAsync(bool useReflection) + { + // Arrange + using CancellationTokenSource source = new(); + InvalidOperationException expected = new("handler failed"); + async ValueTask FailAsync(string message, IWorkflowContext context, CancellationToken cancellationToken) + { + await Task.Yield(); + source.Cancel(); + throw expected; + } + + Executor executor = useReflection + ? new ReflectingHandler(FailAsync) + : new FunctionExecutor("fail", FailAsync); + TestWorkflowContext context = new(executor.Id); + + // Act + TargetInvocationException exception = await Assert.ThrowsAsync( + () => executor.ExecuteCoreAsync("input", new(typeof(string)), context, source.Token).AsTask()); + + // Assert + Assert.Same(expected, exception.InnerException); + Assert.Contains(context.EmittedEvents, evt => evt is ExecutorFailedEvent); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task CancellationWithoutRuntimeCancellationRetainsFailureBehaviorAsync(bool useReflection) + { + // Arrange + async ValueTask CancelAsync(string message, IWorkflowContext context, CancellationToken cancellationToken) + { + await Task.Yield(); + throw new OperationCanceledException(); + } + + Executor executor = useReflection + ? new ReflectingHandler(CancelAsync) + : new FunctionExecutor("cancel", CancelAsync); + TestWorkflowContext context = new(executor.Id); + + // Act + TargetInvocationException exception = await Assert.ThrowsAsync( + () => executor.ExecuteCoreAsync("input", new(typeof(string)), context).AsTask()); + + // Assert: preserve the existing distinction between the two routing paths. + if (useReflection) + { + Assert.Null(exception.InnerException); + } + else + { + Assert.IsType(exception.InnerException); + } + + Assert.Contains(context.EmittedEvents, evt => evt is ExecutorFailedEvent); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task RuntimeCancellationDoesNotSurfaceAsWorkflowErrorAsync(bool offThread) + { + // Arrange + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(30)); + TaskCompletionSource started = new(TaskCreationOptions.RunContinuationsAsynchronously); + Task deadline = Task.Delay(Timeout.InfiniteTimeSpan, timeout.Token); + FunctionExecutor executor = new("cancel", async (message, context, cancellationToken) => + { + started.SetResult(true); + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + }); + Workflow workflow = new WorkflowBuilder(executor).Build(); + var environment = offThread ? InProcessExecution.OffThread : InProcessExecution.Lockstep; + List events = []; + + // Act + await using StreamingRun run = await environment.RunStreamingAsync(workflow, "input"); + async Task ReadEventsAsync() + { + try + { + await foreach (WorkflowEvent evt in run.WatchStreamAsync(timeout.Token)) + { + events.Add(evt); + } + } + catch (OperationCanceledException) when (!timeout.IsCancellationRequested) + { + // The stream may terminate through cancellation rather than normal completion. + } + } + + Task reading = ReadEventsAsync(); + Assert.Same(started.Task, await Task.WhenAny(started.Task, deadline)); + await run.CancelRunAsync(); + Assert.Same(reading, await Task.WhenAny(reading, deadline)); + await reading; + + // Assert + Assert.False(timeout.IsCancellationRequested); + Assert.DoesNotContain(events, evt => evt is ExecutorFailedEvent or WorkflowErrorEvent); + timeout.Cancel(); + } + + private sealed class ReflectingHandler(Func handler) + : ReflectingExecutor("reflecting"), IMessageHandler + { + public ValueTask HandleAsync(string message, IWorkflowContext context, CancellationToken cancellationToken = default) + => handler(message, context, cancellationToken); + } +} From d03a31b375c5b2c483d0414c6f0ea1f941fcf56b Mon Sep 17 00:00:00 2001 From: Naveen Chatlapalli Date: Thu, 17 Sep 2026 23:26:47 -0500 Subject: [PATCH 2/2] fix: preserve cancellation identity across workflow routing --- .../Execution/CallResult.cs | 8 +++- .../Microsoft.Agents.AI.Workflows/Executor.cs | 6 ++- .../Reflection/MessageHandlerInfo.cs | 4 +- .../ExecutorCancellationTests.cs | 40 +++++++++++++++++++ 4 files changed, 53 insertions(+), 5 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/CallResult.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/CallResult.cs index 952b48cd6ce..b7b19dde02d 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/CallResult.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/CallResult.cs @@ -32,6 +32,11 @@ internal sealed class CallResult /// public bool IsCancelled { get; init; } + /// + /// Gets the original cancellation exception, including the token observed by the handler. + /// + public OperationCanceledException? CancellationException { get; init; } + /// /// Indicates whether the call was successful. A call is considered successful if it returned /// without throwing an exception. @@ -64,7 +69,8 @@ private CallResult(bool isVoid = false, bool isCancelled = false) /// A boolean specifying whether the call was void (was not expected to return /// a value). /// A indicating the result of the call. - public static CallResult Cancelled(bool wasVoid) => new(wasVoid, isCancelled: true); + /// The original cancellation exception from the handler. + public static CallResult Cancelled(bool wasVoid, OperationCanceledException exception) => new(wasVoid, isCancelled: true) { CancellationException = exception }; /// /// Create a indicating that an exception was raised during the call. diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs index 2d9beb1b647..c0669da20a8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Executor.cs @@ -250,6 +250,8 @@ protected internal virtual ValueTask InitializeAsync(IWorkflowContext context, C /// The default is . /// A ValueTask representing the asynchronous operation, wrapping the output from the executor. /// No handler found for the message type. + /// The handler observes cancellation of the supplied + /// and throws a cancellation exception carrying that token. /// An exception is generated while handling the message. public ValueTask ExecuteCoreAsync(object message, TypeId messageType, IWorkflowContext context, CancellationToken cancellationToken = default) => this.ExecuteCoreAsync(message, messageType, context, WorkflowTelemetryContext.Disabled, cancellationToken); @@ -265,8 +267,8 @@ protected internal virtual ValueTask InitializeAsync(IWorkflowContext context, C .ConfigureAwait(false); ExecutorEvent executionResult; - if (result is { IsSuccess: false } && cancellationToken.IsCancellationRequested && - (result.IsCancelled || result.Exception is OperationCanceledException)) + OperationCanceledException? cancellation = result?.CancellationException ?? result?.Exception as OperationCanceledException; + if (cancellationToken.IsCancellationRequested && cancellation?.CancellationToken == cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/MessageHandlerInfo.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/MessageHandlerInfo.cs index f655c27cd45..4c53d8c8df8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/MessageHandlerInfo.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Reflection/MessageHandlerInfo.cs @@ -117,10 +117,10 @@ async ValueTask InvokeHandlerAsync(object message, IWorkflowContext return CallResult.ReturnResult(result); } - catch (OperationCanceledException) + catch (OperationCanceledException exception) { // If the operation was canceled, return a canceled CallResult. - return CallResult.Cancelled(wasVoid: expectingVoid); + return CallResult.Cancelled(wasVoid: expectingVoid, exception); } catch (Exception ex) { diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ExecutorCancellationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ExecutorCancellationTests.cs index b08e1bb5bb2..475caae2f67 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ExecutorCancellationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/ExecutorCancellationTests.cs @@ -150,6 +150,46 @@ async Task ReadEventsAsync() timeout.Cancel(); } + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task ForeignCancellationRemainsFailureWhenRuntimeIsCancelledAsync(bool useReflection, bool useDefaultToken) + { + // Arrange + using CancellationTokenSource runtime = new(); + using CancellationTokenSource foreign = new(); + foreign.Cancel(); + OperationCanceledException expected = new(useDefaultToken ? CancellationToken.None : foreign.Token); + async ValueTask CancelAsync(string message, IWorkflowContext context, CancellationToken cancellationToken) + { + await Task.Yield(); + runtime.Cancel(); + throw expected; + } + + Executor executor = useReflection + ? new ReflectingHandler(CancelAsync) + : new FunctionExecutor("foreign", CancelAsync); + TestWorkflowContext context = new(executor.Id); + + // Act + TargetInvocationException exception = await Assert.ThrowsAsync( + () => executor.ExecuteCoreAsync("input", new(typeof(string)), context, runtime.Token).AsTask()); + + // Assert: preserve each routing path's existing failure shape. + if (useReflection) + { + Assert.Null(exception.InnerException); + } + else + { + Assert.Same(expected, exception.InnerException); + } + + Assert.Contains(context.EmittedEvents, evt => evt is ExecutorFailedEvent); + } private sealed class ReflectingHandler(Func handler) : ReflectingExecutor("reflecting"), IMessageHandler {