diff --git a/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs b/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs
index 312064064f..b44d530442 100644
--- a/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs
+++ b/src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs
@@ -64,13 +64,14 @@ public sealed class HeadlessUnitTestSession : IDisposable
         _queue.Add(() =>
         {
             var cts = new CancellationTokenSource();
-            using var globalCts = token.Register(s => ((CancellationTokenSource)s!).Cancel(), cts);
-            using var localCts = cancellationToken.Register(s => ((CancellationTokenSource)s!).Cancel(), cts);
+            using var globalCts = token.Register(s => ((CancellationTokenSource)s!).Cancel(), cts, true);
+            using var localCts = cancellationToken.Register(s => ((CancellationTokenSource)s!).Cancel(), cts, true);
 
             try
             {
                 var task = action();
-                task.ContinueWith((_, s) => ((CancellationTokenSource)s!).Cancel(), cts);
+                task.ContinueWith((_, s) => ((CancellationTokenSource)s!).Cancel(), cts,
+                    TaskScheduler.FromCurrentSynchronizationContext());
 
                 if (cts.IsCancellationRequested)
                 {
@@ -78,7 +79,7 @@ public sealed class HeadlessUnitTestSession : IDisposable
                 }
 
                 var frame = new DispatcherFrame();
-                using var innerCts = cts.Token.Register(() => frame.Continue = false);
+                using var innerCts = cts.Token.Register(() => frame.Continue = false, true);
                 Dispatcher.UIThread.PushFrame(frame);
 
                 var result = task.GetAwaiter().GetResult();
diff --git a/src/Headless/Avalonia.Headless/HeadlessWindowExtensions.cs b/src/Headless/Avalonia.Headless/HeadlessWindowExtensions.cs
index 8fbc5ec6ef..a3bda25b19 100644
--- a/src/Headless/Avalonia.Headless/HeadlessWindowExtensions.cs
+++ b/src/Headless/Avalonia.Headless/HeadlessWindowExtensions.cs
@@ -44,53 +44,54 @@ public static class HeadlessWindowExtensions
     /// Simulates keyboard press on the headless window/toplevel.
     /// </summary>
     public static void KeyPress(this TopLevel topLevel, Key key, RawInputModifiers modifiers) =>
-        RunJobsAndGetImpl(topLevel).KeyPress(key, modifiers);
+        RunJobsOnImpl(topLevel, w => w.KeyPress(key, modifiers));
 
     /// <summary>
     /// Simulates keyboard release on the headless window/toplevel.
     /// </summary>
     public static void KeyRelease(this TopLevel topLevel, Key key, RawInputModifiers modifiers) =>
-        RunJobsAndGetImpl(topLevel).KeyRelease(key, modifiers);
+        RunJobsOnImpl(topLevel, w => w.KeyRelease(key, modifiers));
 
     /// <summary>
     /// Simulates mouse down on the headless window/toplevel.
     /// </summary>
     public static void MouseDown(this TopLevel topLevel, Point point, MouseButton button,
         RawInputModifiers modifiers = RawInputModifiers.None) =>
-        RunJobsAndGetImpl(topLevel).MouseDown(point, button, modifiers);
+        RunJobsOnImpl(topLevel, w => w.MouseDown(point, button, modifiers));
 
     /// <summary>
     /// Simulates mouse move on the headless window/toplevel.
     /// </summary>
     public static void MouseMove(this TopLevel topLevel, Point point,
         RawInputModifiers modifiers = RawInputModifiers.None) =>
-        RunJobsAndGetImpl(topLevel).MouseMove(point, modifiers);
+        RunJobsOnImpl(topLevel, w => w.MouseMove(point, modifiers));
 
     /// <summary>
     /// Simulates mouse up on the headless window/toplevel.
     /// </summary>
     public static void MouseUp(this TopLevel topLevel, Point point, MouseButton button,
         RawInputModifiers modifiers = RawInputModifiers.None) =>
-        RunJobsAndGetImpl(topLevel).MouseUp(point, button, modifiers);
+        RunJobsOnImpl(topLevel, w => w.MouseUp(point, button, modifiers));
 
     /// <summary>
     /// Simulates mouse wheel on the headless window/toplevel.
     /// </summary>
     public static void MouseWheel(this TopLevel topLevel, Point point, Vector delta,
         RawInputModifiers modifiers = RawInputModifiers.None) =>
-        RunJobsAndGetImpl(topLevel).MouseWheel(point, delta, modifiers);
+        RunJobsOnImpl(topLevel, w => w.MouseWheel(point, delta, modifiers));
 
     /// <summary>
     /// Simulates drag'n'drop target on the headless window/toplevel.
     /// </summary>
     public static void DragDrop(this TopLevel topLevel, Point point, RawDragEventType type, IDataObject data,
         DragDropEffects effects, RawInputModifiers modifiers = RawInputModifiers.None) =>
-        RunJobsAndGetImpl(topLevel).DragDrop(point, type, data, effects, modifiers);
+        RunJobsOnImpl(topLevel, w => w.DragDrop(point, type, data, effects, modifiers));
 
-    private static IHeadlessWindow RunJobsAndGetImpl(this TopLevel topLevel)
+    private static void RunJobsOnImpl(this TopLevel topLevel, Action<IHeadlessWindow> action)
     {
         Dispatcher.UIThread.RunJobs();
-        return GetImpl(topLevel);
+        action(GetImpl(topLevel));
+        Dispatcher.UIThread.RunJobs();
     }
 
     private static IHeadlessWindow GetImpl(this TopLevel topLevel)
diff --git a/tests/Avalonia.Headless.UnitTests/InputTests.cs b/tests/Avalonia.Headless.UnitTests/InputTests.cs
index 5a7d3faae9..5e3b6e762f 100644
--- a/tests/Avalonia.Headless.UnitTests/InputTests.cs
+++ b/tests/Avalonia.Headless.UnitTests/InputTests.cs
@@ -46,7 +46,7 @@ public class InputTests
 
         _window.MouseDown(new Point(50, 50), MouseButton.Left);
         _window.MouseUp(new Point(50, 50), MouseButton.Left);
-        
+
         Assert.True(buttonClicked);
     }