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).

Author: Kryptos-FR

src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs

@@ -96,11 +96,16 @@ internal Task<TResult> DispatchCore<TResult>(Func<Task<TResult>> action, bool ca
             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 @@ internal Task<TResult> DispatchCore<TResult>(Func<Task<TResult>> action, bool ca
 
                     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;
     }

tests/Avalonia.Headless.UnitTests/IsolationTests.cs

@@ -1,5 +1,8 @@
 using System;
 using System.Reflection;
+using System.Threading;
+using System.Threading.Tasks;
+using Avalonia.Headless;
 using Avalonia.Threading;
 
 namespace Avalonia.Headless.UnitTests;
@@ -64,4 +67,36 @@ public void Application_Instance_Should_Match_Isolation_Level(int runIndex)
         s_previousAppRef = new WeakReference<Application>(currentApp);
         s_previousDispatcherRef = new WeakReference<Dispatcher>(currentDispatcher);
     }
+
+#if NUNIT
+    [Test]
+#elif XUNIT
+    [Fact]
+#endif
+    public async Task Dispatch_Cleanup_Should_Complete_Before_Task_Returns()
+    {
+        // Regression test for https://github.com/AvaloniaUI/Avalonia/issues/20664.
+        // EnsureIsolatedApplication().Dispose() must complete (resetting s_uiThread to null)
+        // before the dispatch task resolves. If Dispose ran after tcs.TrySetResult, s_uiThread
+        // could still point to the headless dispatcher here, making CheckAccess() return false
+        // on this non-headless thread.
+        //
+        // Only applies to PerTest isolation: PerAssembly uses EnsureSharedApplication which
+        // intentionally keeps s_uiThread set for the lifetime of the assembly, so CheckAccess()
+        // from a non-headless thread would always be false there and the race does not apply.
+        var isolationLevel =
+            GetType().Assembly.GetCustomAttribute<AvaloniaTestIsolationAttribute>()?.IsolationLevel
+            ?? AvaloniaTestIsolationLevel.PerTest;
+
+        if (isolationLevel != AvaloniaTestIsolationLevel.PerTest)
+            return;
+
+        // Uses the shared assembly session (not StartNew) so no competing thread calls
+        // ResetGlobalState() concurrently with other tests in the suite.
+        var session = HeadlessUnitTestSession.GetOrStartForAssembly(GetType().Assembly);
+
+        await session.Dispatch(() => { }, CancellationToken.None);
+
+        AssertHelper.True(Dispatcher.UIThread.CheckAccess());
+    }
 }