Preserve focus after workspace transition invalidation

Sources/OmniWM/Core/Controller/LayoutRefreshController.swift

@@ -998,12 +998,14 @@ import QuartzCore
         affectedWorkspaces: Set<WorkspaceDescriptor.ID> = [],
         reason: RefreshReason = .workspaceTransition,
         postLayoutGateWorkspaceIds: Set<WorkspaceDescriptor.ID>? = nil,
+        postLayoutInvalidated: PostLayoutAction? = nil,
         postLayout: PostLayoutAction? = nil
     ) {
         requestImmediateRelayout(
             reason: reason,
             affectedWorkspaceIds: affectedWorkspaces,
             postLayout: postLayout,
+            postLayoutInvalidated: postLayoutInvalidated,
             postLayoutGateWorkspaceIds: postLayoutGateWorkspaceIds
         )
     }

Sources/OmniWM/Core/Controller/WindowActionHandler.swift

@@ -500,10 +500,20 @@ final class WindowActionHandler {
                 )
             )
         }
-        controller.layoutRefreshController
-            .commitWorkspaceTransition(reason: .workspaceTransition) { [weak controller] in
-                controller?.focusWindow(token)
+        let focusTarget: LayoutRefreshController.PostLayoutAction = { [weak controller] in
+            guard let controller,
+                  controller.activeWorkspace()?.id == workspaceId,
+                  controller.workspaceManager.entry(for: token)?.workspaceId == workspaceId
+            else {
+                return
             }
+            controller.focusWindow(token)
+        }
+        controller.layoutRefreshController.commitWorkspaceTransition(
+            reason: .workspaceTransition,
+            postLayoutInvalidated: focusTarget,
+            postLayout: focusTarget
+        )
         return true
     }
 

Tests/OmniWMTests/RuntimeArchitectureTests.swift

@@ -1830,6 +1830,38 @@ final class RuntimeArchitectureTests: XCTestCase {
         XCTAssertTrue(ranInvalidatedAction)
     }
 
+    @MainActor
+    func testWorkspaceTransitionPreservesInvalidatedFocusHandoff() async throws {
+        let controller = Self.controller()
+        let workspaceId = try XCTUnwrap(
+            controller.workspaceManager.workspaceId(for: "1", createIfMissing: true)
+        )
+        controller.layoutRefreshController.layoutState.hasCompletedInitialRefresh = true
+        var ranCurrentAction = false
+        var ranInvalidatedAction = false
+
+        controller.layoutRefreshController.commitWorkspaceTransition(
+            affectedWorkspaces: [workspaceId],
+            postLayoutGateWorkspaceIds: [workspaceId],
+            postLayoutInvalidated: { ranInvalidatedAction = true },
+            postLayout: { ranCurrentAction = true }
+        )
+        controller.workspaceManager.invalidateLayout(for: [workspaceId])
+
+        for _ in 0 ..< 8 {
+            if let task = controller.layoutRefreshController.layoutState.activeRefreshTask {
+                await task.value
+            } else if controller.layoutRefreshController.layoutState.pendingRefresh == nil {
+                break
+            } else {
+                await Task.yield()
+            }
+        }
+
+        XCTAssertFalse(ranCurrentAction)
+        XCTAssertTrue(ranInvalidatedAction)
+    }
+
     @MainActor
     func testLayoutPlanAcceptedSeqIncludesAnimationDirectiveFocusMutation() throws {
         let controller = Self.controller()