From c0a41d8812f862fec1dc5a6bb700748e1076d959 Mon Sep 17 00:00:00 2001
From: Nicolas Musset <musset.nicolas@gmail.com>
Date: Mon, 27 Apr 2026 13:25:58 +0200
Subject: [PATCH] Fix headless cleanup race: dispose before signalling TCS
 (#20664) (#21223)

* Fix headless cleanup race: dispose before signalling TCS (#20664)

Restructure DispatchCore<TResult> so application.Dispose() runs in a
finally block before tcs.TrySet* is called. Previously, the using-var
scope closed after TrySetResult, leaving s_uiThread set when the next
non-headless test started.

Adds a regression test in IsolationTests that verifies Dispatcher.UIThread
is accessible from the calling thread immediately after a headless dispatch
completes (PerTest isolation only).

* Use Assert.SkipWhen in Dispatch_Cleanup_Should_Complete_Before_Task_Returns

* Restrict Dispatch_Cleanup_Should_Complete_Before_Task_Returns to NUnit

xUnit's AvaloniaFactDiscoverer wraps every [Fact] in session.Dispatch,
so awaiting another session.Dispatch inside the test body deadlocks on
the session's single background thread.

* Remove new test

It isn't reliable enough during CI.

---------

Co-authored-by: Jumar Macato <16554748+jmacato@users.noreply.github.com>
---
 .../HeadlessUnitTestSession.cs                | 45 +++++++++++++------
 1 file changed, 32 insertions(+), 13 deletions(-)

diff --git a/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs b/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs
index cda9c4ae75..0d48dd6919 100644
--- a/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs
+++ b/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs
@@ -96,11 +96,16 @@ public sealed class HeadlessUnitTestSession : IDisposable, IAsyncDisposable
             using var globalCts = token.Register(s => ((CancellationTokenSource)s!).Cancel(), cts, true);
             using var localCts = cancellationToken.Register(s => ((CancellationTokenSource)s!).Cancel(), cts, true);
 
+            var application = _isolated
+                ? EnsureIsolatedApplication()
+                : EnsureSharedApplication();
+
+            bool shouldCancel = false;
+            Exception? caught = null;
+            TResult result = default!;
+
             try
             {
-                using var application = _isolated
-                    ? EnsureIsolatedApplication()
-                    : EnsureSharedApplication();
                 var task = action();
                 if (task.Status != TaskStatus.RanToCompletion)
                 {
@@ -110,22 +115,36 @@ public sealed class HeadlessUnitTestSession : IDisposable, IAsyncDisposable
 
                     if (cts.IsCancellationRequested)
                     {
-                        tcs.TrySetCanceled(cts.Token);
-                        return;
+                        shouldCancel = true;
+                    }
+                    else
+                    {
+                        var frame = new DispatcherFrame();
+                        using var innerCts = cts.Token.Register(() => frame.Continue = false, true);
+                        Dispatcher.UIThread.PushFrame(frame);
+                        result = task.GetAwaiter().GetResult();
                     }
-
-                    var frame = new DispatcherFrame();
-                    using var innerCts = cts.Token.Register(() => frame.Continue = false, true);
-                    Dispatcher.UIThread.PushFrame(frame);
                 }
-
-                var result = task.GetAwaiter().GetResult();
-                tcs.TrySetResult(result);
+                else
+                {
+                    result = task.GetAwaiter().GetResult();
+                }
             }
             catch (Exception ex)
             {
-                tcs.TrySetException(ex);
+                caught = ex;
             }
+            finally
+            {
+                application.Dispose();
+            }
+
+            if (caught != null)
+                tcs.TrySetException(caught);
+            else if (shouldCancel)
+                tcs.TrySetCanceled(cts.Token);
+            else
+                tcs.TrySetResult(result);
         }, executionContext));
         return tcs.Task;
     }