fix(core): preserve cancellation across retries

Author: thomhurst

src/TUnit.Core/TestContext.Execution.cs

@@ -167,17 +167,27 @@ internal void InitializeTestCancellation(CancellationToken cancellationToken)
     {
         lock (Lock)
         {
+            // Cancel() targets the whole execution, so retries inherit an accepted request.
+            var preserveCancellationRequest = CurrentRetryAttempt > 0 && _testCancellationRequested;
+
             if (_testCancellationTokenSource is { } previousTestCancellationTokenSource)
             {
                 (_retiredTestCancellationTokenSources ??= []).Add(previousTestCancellationTokenSource);
             }
 
-            _testCancellationTokenSource = cancellationToken.CanBeCanceled
+            var testCancellationTokenSource = cancellationToken.CanBeCanceled
                 ? CancellationTokenSource.CreateLinkedTokenSource(cancellationToken)
                 : new CancellationTokenSource();
+
+            if (preserveCancellationRequest)
+            {
+                testCancellationTokenSource.Cancel();
+            }
+
+            _testCancellationTokenSource = testCancellationTokenSource;
             _acceptingTestCancellation = true;
-            _testCancellationRequested = false;
-            _baseCancellationToken = _testCancellationTokenSource.Token;
+            _testCancellationRequested = preserveCancellationRequest;
+            _baseCancellationToken = testCancellationTokenSource.Token;
             RebuildLinkedCancellationTokenSource();
         }
     }

tests/TUnit.UnitTests/TestContextCancellationTests.cs

@@ -7,17 +7,7 @@ public class TestContextCancellationTests
     [Test]
     public async Task CleanupDefersSourceDisposalWhileAcceptedCancellationIsInProgress()
     {
-        var currentContext = TestContext.Current!;
-        var context = new TestContext(
-            nameof(CleanupDefersSourceDisposalWhileAcceptedCancellationIsInProgress),
-            currentContext.ServiceProvider,
-            currentContext.ClassContext,
-            new TestBuilderContext
-            {
-                TestMetadata = currentContext.TestDetails.MethodMetadata
-            },
-            CancellationToken.None);
-
+        var context = CreateContext();
         context.InitializeTestCancellation(CancellationToken.None);
 
         var cancellationToken = context.TestCancellationToken;
@@ -47,4 +37,54 @@ public async Task CleanupDefersSourceDisposalWhileAcceptedCancellationIsInProgre
             context.RemoveFromRegistry();
         }
     }
+
+    [Test]
+    public async Task RetrySourcePreservesAcceptedCancellation()
+    {
+        var context = CreateContext();
+        context.InitializeTestCancellation(CancellationToken.None);
+
+        var callbackStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        using var releaseCallback = new ManualResetEventSlim();
+        using var registration = context.TestCancellationToken.Register(() =>
+        {
+            callbackStarted.TrySetResult();
+            releaseCallback.Wait(TimeSpan.FromSeconds(10));
+        });
+
+        var cancellationTask = Task.Run(context.Cancel);
+
+        try
+        {
+            await callbackStarted.Task.WaitAsync(TimeSpan.FromSeconds(10));
+
+            context.CurrentRetryAttempt = 1;
+            context.InitializeTestCancellation(CancellationToken.None);
+
+            await Assert.That(context.IsTestCancellationRequested).IsTrue();
+            await Assert.That(context.TestCancellationToken.IsCancellationRequested).IsTrue();
+        }
+        finally
+        {
+            releaseCallback.Set();
+            await cancellationTask;
+            context.DisposeLinkedCancellationTokenSources();
+            context.RemoveFromRegistry();
+        }
+    }
+
+    private static TestContext CreateContext()
+    {
+        var currentContext = TestContext.Current!;
+
+        return new TestContext(
+            nameof(TestContextCancellationTests),
+            currentContext.ServiceProvider,
+            currentContext.ClassContext,
+            new TestBuilderContext
+            {
+                TestMetadata = currentContext.TestDetails.MethodMetadata
+            },
+            CancellationToken.None);
+    }
 }