From 0b3243e9c074d6d77f8e6fba5b718c0ef89c9d9c Mon Sep 17 00:00:00 2001
From: Nathan Nguyen <nathannguyen927151@gmail.com>
Date: Tue, 21 Jul 2026 06:00:13 +1000
Subject: [PATCH] Fix headless session hang when cleanup throws (#21781)

* test(headless): reproduce cleanup exception dispatch hang

A dispatcher job that throws during work-item cleanup leaves the dispatch task incomplete and prevents later work from running. The session contract requires every queued dispatch to reach a terminal state.

Add behavioral coverage that posts a throwing cleanup job, expects its exception from Dispatch, and verifies the assembly session can process the next dispatch. The shared NUnit test runs under both PerTest and PerAssembly projects.

* fix(headless): report dispatch cleanup exceptions

An exception during application cleanup escapes the work item before its completion source is settled, faults the private consumer task, and leaves current and future Dispatch calls waiting forever.

Capture cleanup failures as the work item's exception so the consumer can continue. Restore synchronization context, locator scope, and dispatcher state in finally blocks before processing later work.

---------

Co-authored-by: Julien Lebosquain <julien@lebosquain.net>
---
 .../HeadlessUnitTestSession.cs                | 46 +++++++++++++++----
 .../HeadlessUnitTestSessionTests.cs           | 32 +++++++++++++
 2 files changed, 69 insertions(+), 9 deletions(-)
 create mode 100644 tests/Avalonia.Headless.UnitTests/HeadlessUnitTestSessionTests.cs

diff --git a/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs b/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs
index 5613244461..f215d867f0 100644
--- a/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs
+++ b/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs
@@ -145,7 +145,16 @@ public sealed class HeadlessUnitTestSession : IDisposable, IAsyncDisposable
             }
             finally
             {
-                application.Dispose();
+                try
+                {
+                    application.Dispose();
+                }
+                catch (Exception ex)
+                {
+                    // Cleanup runs before the TCS is completed, so its failure must be
+                    // reported by this work item instead of escaping the consumer loop.
+                    caught = ex;
+                }
             }
 
             if (caught != null)
@@ -173,8 +182,14 @@ public sealed class HeadlessUnitTestSession : IDisposable, IAsyncDisposable
 
         return Disposable.Create(() =>
         {
-            Dispatcher.UIThread.RunJobs();
-            SynchronizationContext.SetSynchronizationContext(oldContext);
+            try
+            {
+                Dispatcher.UIThread.RunJobs();
+            }
+            finally
+            {
+                SynchronizationContext.SetSynchronizationContext(oldContext);
+            }
         });
     }
 
@@ -195,12 +210,25 @@ public sealed class HeadlessUnitTestSession : IDisposable, IAsyncDisposable
 
         return Disposable.Create(() =>
         {
-            ((ToolTipService?)AvaloniaLocator.Current.GetService<IToolTipService>())?.Dispose();
-            (AvaloniaLocator.Current.GetService<FontManager>() as IDisposable)?.Dispose();
-            Dispatcher.ResetForUnitTests();
-            scope.Dispose();
-            Dispatcher.ResetBeforeUnitTests();
-            SynchronizationContext.SetSynchronizationContext(oldContext);
+            try
+            {
+                ((ToolTipService?)AvaloniaLocator.Current.GetService<IToolTipService>())?.Dispose();
+                (AvaloniaLocator.Current.GetService<FontManager>() as IDisposable)?.Dispose();
+                Dispatcher.ResetForUnitTests();
+            }
+            finally
+            {
+                // Cleanup jobs can throw, but the ambient state still belongs to this dispatch.
+                try
+                {
+                    scope.Dispose();
+                }
+                finally
+                {
+                    Dispatcher.ResetBeforeUnitTests();
+                    SynchronizationContext.SetSynchronizationContext(oldContext);
+                }
+            }
         });
     }
 
diff --git a/tests/Avalonia.Headless.UnitTests/HeadlessUnitTestSessionTests.cs b/tests/Avalonia.Headless.UnitTests/HeadlessUnitTestSessionTests.cs
new file mode 100644
index 0000000000..7372dd639a
--- /dev/null
+++ b/tests/Avalonia.Headless.UnitTests/HeadlessUnitTestSessionTests.cs
@@ -0,0 +1,32 @@
+#if NUNIT
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Avalonia.Threading;
+
+namespace Avalonia.Headless.UnitTests;
+
+public class HeadlessUnitTestSessionTests
+{
+    [Test]
+    public async Task Dispatch_Should_Report_Cleanup_Exceptions_And_Continue()
+    {
+        var session = HeadlessUnitTestSession.GetOrStartForAssembly(GetType().Assembly);
+
+        const string message = "Thrown by a dispatcher job during cleanup.";
+        var poisonedDispatch = session.Dispatch(
+            () => Dispatcher.UIThread.Post(() => throw new InvalidOperationException(message)),
+            CancellationToken.None);
+
+        var exception = Assert.ThrowsAsync<InvalidOperationException>(async () =>
+            await poisonedDispatch.WaitAsync(TimeSpan.FromSeconds(10)));
+
+        Assert.That(exception!.Message, Is.EqualTo(message));
+
+        var result = await session.Dispatch(() => 42, CancellationToken.None)
+            .WaitAsync(TimeSpan.FromSeconds(10));
+
+        Assert.That(result, Is.EqualTo(42));
+    }
+}
+#endif