Better reactive test cancellation

TUnit.Engine.Tests/CanCancelTests.cs

@@ -0,0 +1,27 @@
+using Shouldly;
+using TUnit.Core.Enums;
+using TUnit.Engine.Tests.Enums;
+
+namespace TUnit.Engine.Tests;
+
+/// <summary>
+/// Tests that verify test cancellation works correctly when graceful cancellation is requested.
+/// Skipped on Windows because CliWrap's graceful cancellation uses GenerateConsoleCtrlEvent,
+/// which doesn't work properly for subprocess control.
+/// </summary>
+[ExcludeOn(OS.Windows)]
+public class CanCancelTests(TestMode testMode) : InvokableTestBase(testMode)
+{
+    [Test, Timeout(30_000)]
+    public async Task Test(CancellationToken ct)
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+        await RunTestsWithFilter(
+            "/*/*/CanCancelTests/*",
+            [
+                result => result.ResultSummary.Outcome.ShouldBe("Failed"),
+                result => TimeSpan.Parse(result.Duration).ShouldBeLessThan(TimeSpan.FromSeconds(30))
+            ],
+            new RunOptions().WithGracefulCancellationToken(cts.Token));
+    }
+}

TUnit.Engine.Tests/InvokableTestBase.cs

@@ -1,4 +1,5 @@
 using System.Runtime.CompilerServices;
+using System.Text;
 using CliWrap;
 using CliWrap.Buffered;
 using TrxTools.TrxParser;
@@ -72,7 +73,7 @@ private async Task RunWithoutAot(string filter,
             .WithWorkingDirectory(testProject.DirectoryName!)
             .WithValidation(CommandResultValidation.None);
 
-        await RunWithFailureLogging(command, trxFilename, assertions, assertionExpression);
+        await RunWithFailureLogging(command, runOptions, trxFilename, assertions, assertionExpression);
     }
 
     private async Task RunWithAot(string filter, List<Action<TestRun>> assertions,
@@ -106,7 +107,7 @@ private async Task RunWithAot(string filter, List<Action<TestRun>> assertions,
             )
             .WithValidation(CommandResultValidation.None);
 
-        await RunWithFailureLogging(command, trxFilename, assertions, assertionExpression);
+        await RunWithFailureLogging(command, runOptions, trxFilename, assertions, assertionExpression);
     }
 
     private async Task RunWithSingleFile(string filter,
@@ -140,7 +141,7 @@ private async Task RunWithSingleFile(string filter,
             )
             .WithValidation(CommandResultValidation.None);
 
-        await RunWithFailureLogging(command, trxFilename, assertions, assertionExpression);
+        await RunWithFailureLogging(command, runOptions, trxFilename, assertions, assertionExpression);
     }
 
     protected static FileInfo? FindFile(Func<FileInfo, bool> predicate)
@@ -153,13 +154,27 @@ private async Task RunWithSingleFile(string filter,
         return FileSystemHelpers.FindFolder(predicate);
     }
 
-    private async Task RunWithFailureLogging(Command command, string trxFilename, List<Action<TestRun>> assertions, string assertionExpression)
+    private async Task<CommandTask<BufferedCommandResult>> RunWithFailureLogging(Command command, RunOptions runOptions,
+        string trxFilename, List<Action<TestRun>> assertions, string assertionExpression)
     {
+        var commandTask = command.ExecuteBufferedAsync
+        (
+            gracefulCancellationToken: runOptions.GracefulCancellationToken,
+            forcefulCancellationToken: runOptions.ForcefulCancellationToken,
+            standardOutputEncoding: runOptions.StandardOutputEncoding,
+            standardErrorEncoding: runOptions.StandardErrorEncoding
+        );
+
         BufferedCommandResult? commandResult = null;
 
         try
         {
-            commandResult = await command.ExecuteBufferedAsync();
+            foreach (var onExecutingDelegate in runOptions.OnExecutingDelegates)
+            {
+                await onExecutingDelegate(commandTask);
+            }
+
+            commandResult = await commandTask;
 
             await TrxAsserter.AssertTrx(testMode, command, commandResult, assertions, trxFilename, assertionExpression: assertionExpression);
         }
@@ -176,16 +191,38 @@ private async Task RunWithFailureLogging(Command command, string trxFilename, Li
                                  Expression: {assertionExpression}
                                  """);
         }
+
+        return commandTask;
     }
 }
 
 public record RunOptions
 {
+    public CancellationToken GracefulCancellationToken { get; set; } = CancellationToken.None;
+    public CancellationToken ForcefulCancellationToken { get; set; } = CancellationToken.None;
+
+    public Encoding StandardOutputEncoding { get; set; } = Encoding.UTF8;
+    public Encoding StandardErrorEncoding { get; set; } = Encoding.UTF8;
+
     public List<string> AdditionalArguments { get; init; } = [];
 
+    public List<Func<CommandTask<BufferedCommandResult>, Task>> OnExecutingDelegates { get; init; } = [];
+
     public RunOptions WithArgument(string argument)
     {
         AdditionalArguments.Add(argument);
         return this;
     }
+
+    public RunOptions WithGracefulCancellationToken(CancellationToken token)
+    {
+        GracefulCancellationToken = token;
+        return this;
+    }
+
+    public RunOptions WithForcefulCancellationToken(CancellationToken token)
+    {
+        ForcefulCancellationToken = token;
+        return this;
+    }
 }

TUnit.Engine/Helpers/TimeoutHelper.cs

@@ -1,5 +1,3 @@
-using System.Diagnostics.CodeAnalysis;
-
 namespace TUnit.Engine.Helpers;
 
 /// <summary>
@@ -11,103 +9,56 @@ internal static class TimeoutHelper
     /// Executes a task with an optional timeout. If the timeout elapses before the task completes,
     /// control is returned to the caller immediately with a TimeoutException.
     /// </summary>
-    /// <param name="taskFactory">Factory function that creates the task to execute</param>
-    /// <param name="timeout">Optional timeout duration. If null, no timeout is applied</param>
-    /// <param name="cancellationToken">Cancellation token for the operation</param>
-    /// <param name="timeoutMessage">Optional custom timeout message. If null, uses default message</param>
-    /// <returns>The completed task</returns>
-    /// <exception cref="TimeoutException">Thrown when the timeout elapses before task completion</exception>
     public static async Task ExecuteWithTimeoutAsync(
         Func<CancellationToken, Task> taskFactory,
         TimeSpan? timeout,
         CancellationToken cancellationToken,
         string? timeoutMessage = null)
     {
-        if (!timeout.HasValue)
-        {
-            await taskFactory(cancellationToken).ConfigureAwait(false);
-            return;
-        }
-
-        using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-        timeoutCts.CancelAfter(timeout.Value);
-
-        var executionTask = taskFactory(timeoutCts.Token);
-
-        // Use a cancellable timeout task to avoid leaving Task.Delay running in the background
-        using var timeoutTaskCts = new CancellationTokenSource();
-        var timeoutTask = Task.Delay(timeout.Value, timeoutTaskCts.Token);
-
-        var completedTask = await Task.WhenAny(executionTask, timeoutTask).ConfigureAwait(false);
-
-        if (completedTask == timeoutTask)
-        {
-            // Timeout occurred - cancel the execution task and wait briefly for cleanup
-            timeoutCts.Cancel();
-
-            // Give the execution task a chance to handle cancellation gracefully
-            try
-            {
-                await executionTask.ConfigureAwait(false);
-            }
-            catch (OperationCanceledException)
-            {
-                // Expected when cancellation is properly handled
-            }
-            catch
+        await ExecuteWithTimeoutAsync(
+            async ct =>
             {
-                // Ignore other exceptions from the cancelled task
-            }
-
-            var message = timeoutMessage ?? $"Operation timed out after {timeout.Value}";
-            throw new TimeoutException(message);
-        }
-
-        // Task completed normally - cancel the timeout task to free resources immediately
-        timeoutTaskCts.Cancel();
-
-        // Await the result to propagate any exceptions
-        await executionTask.ConfigureAwait(false);
+                await taskFactory(ct).ConfigureAwait(false);
+                return true;
+            },
+            timeout,
+            cancellationToken,
+            timeoutMessage).ConfigureAwait(false);
     }
 
     /// <summary>
     /// Executes a task with an optional timeout and returns a result. If the timeout elapses before the task completes,
     /// control is returned to the caller immediately with a TimeoutException.
     /// </summary>
-    /// <typeparam name="T">The type of result returned by the task</typeparam>
-    /// <param name="taskFactory">Factory function that creates the task to execute</param>
-    /// <param name="timeout">Optional timeout duration. If null, no timeout is applied</param>
-    /// <param name="cancellationToken">Cancellation token for the operation</param>
-    /// <param name="timeoutMessage">Optional custom timeout message. If null, uses default message</param>
-    /// <returns>The result of the completed task</returns>
-    /// <exception cref="TimeoutException">Thrown when the timeout elapses before task completion</exception>
     public static async Task<T> ExecuteWithTimeoutAsync<T>(
         Func<CancellationToken, Task<T>> taskFactory,
         TimeSpan? timeout,
         CancellationToken cancellationToken,
         string? timeoutMessage = null)
     {
-        if (!timeout.HasValue)
-        {
-            return await taskFactory(cancellationToken).ConfigureAwait(false);
-        }
+        var executionTask = taskFactory(cancellationToken);
 
-        using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-        timeoutCts.CancelAfter(timeout.Value);
+        // Create a task that completes when cancellation is requested
+        var tcs = new TaskCompletionSource<bool>();
+        using var registration = cancellationToken.Register(() => tcs.TrySetResult(true));
 
-        var executionTask = taskFactory(timeoutCts.Token);
-
-        // Use a cancellable timeout task to avoid leaving Task.Delay running in the background
-        using var timeoutTaskCts = new CancellationTokenSource();
-        var timeoutTask = Task.Delay(timeout.Value, timeoutTaskCts.Token);
+        // Create timeout task if timeout is specified
+        using var timeoutCts = new CancellationTokenSource();
+        var timeoutTask = timeout.HasValue
+            ? Task.Delay(timeout.Value, timeoutCts.Token)
+            : Task.Delay(Timeout.Infinite, timeoutCts.Token);
 
-        var completedTask = await Task.WhenAny(executionTask, timeoutTask).ConfigureAwait(false);
+        var winningTask = await Task.WhenAny(executionTask, tcs.Task, timeoutTask).ConfigureAwait(false);
 
-        if (completedTask == timeoutTask)
+        // Cancellation requested
+        if (winningTask == tcs.Task)
+        {
+            throw new OperationCanceledException(cancellationToken);
+        }
+
+        // Timeout occurred
+        if (winningTask == timeoutTask)
         {
-            // Timeout occurred - cancel the execution task and wait briefly for cleanup
-            timeoutCts.Cancel();
-
             // Give the execution task a chance to handle cancellation gracefully
             try
             {
@@ -121,15 +72,14 @@ public static async Task<T> ExecuteWithTimeoutAsync<T>(
             {
                 // Ignore other exceptions from the cancelled task
             }
-            
-            var message = timeoutMessage ?? $"Operation timed out after {timeout.Value}";
+
+            var message = timeoutMessage ?? $"Operation timed out after {timeout!.Value}";
             throw new TimeoutException(message);
         }
 
-        // Task completed normally - cancel the timeout task to free resources immediately
-        timeoutTaskCts.Cancel();
+        // Task completed normally - cancel the timeout task to free resources
+        timeoutCts.Cancel();
 
-        // Await the result to propagate any exceptions
         return await executionTask.ConfigureAwait(false);
     }
-}
+}

TUnit.Engine/TestSessionCoordinator.cs

@@ -90,7 +90,8 @@ private async Task ExecuteTestsCore(List<AbstractExecutableTest> testList, Cance
         // Combine cancellation tokens
         using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(
             cancellationToken,
-            _serviceProvider.FailFastCancellationSource.Token);
+            _serviceProvider.FailFastCancellationSource.Token,
+            _serviceProvider.CancellationToken.Token);
 
         // Schedule and execute tests (batch approach to preserve ExecutionContext)
         var success = await _testScheduler.ScheduleAndExecuteAsync(testList, linkedCts.Token);

TUnit.TestProject/CanCancelTests.cs

@@ -0,0 +1,10 @@
+namespace TUnit.TestProject;
+
+public class CanCancelTests
+{
+    [Test, Explicit]
+    public async Task CanCancel()
+    {
+        await Task.Delay(TimeSpan.FromMinutes(5));
+    }
+}