Fix native tab window replacement

Sources/AppBundle/model/KnownBundleId.swift

@@ -9,6 +9,7 @@ enum KnownBundleId: String, Equatable {
     case codex = "com.openai.codex"
     case emacs = "org.gnu.Emacs"
     case finder = "com.apple.finder"
+    case fork = "com.DanPristupov.Fork"
     case ghostty = "com.mitchellh.ghostty"
     case gimp = "org.gimp.gimp-2.10"
     case iphonesimulator = "com.apple.iphonesimulator"
@@ -43,4 +44,8 @@ enum KnownBundleId: String, Equatable {
     var isVscode: Bool {
         self == .vscode || self == .vscodium
     }
+
+    var exposesInactiveNativeTabsAsWindows: Bool {
+        self == .fork || self == .ghostty
+    }
 }

Sources/AppBundle/tree/MacApp.swift

@@ -118,13 +118,44 @@ final class MacApp: AbstractApp {
 
     // todo merge together with detectNewWindows
     func getFocusedWindow(_ cm: CancellationMode) async throws -> Window? {
-        let windowId = try await thread?.runInLoop(cm) { [nsApp, axApp, windows] job in
-            try axApp.threadGuarded.get(Ax.focusedWindowAttr)
-                .flatMap { try windows.threadGuarded.getOrRegisterAxWindow(windowId: $0.windowId, $0.ax.cast, nsApp, job) }?
-                .windowId
+        let previousFocusedWindowId = lastNativeFocusedWindowId
+        let focused = try await thread?.runInLoop(cm) { [nsApp, appId, axApp, windows] job -> (UInt32, [UInt32])? in
+            guard let axFocusedWindow = axApp.threadGuarded.get(Ax.focusedWindowAttr),
+                  let focusedWindow = try windows.threadGuarded.getOrRegisterAxWindow(
+                      windowId: axFocusedWindow.windowId,
+                      axFocusedWindow.ax.cast,
+                      nsApp,
+                      job,
+                  )
+            else {
+                return nil
+            }
+            guard appId?.exposesInactiveNativeTabsAsWindows == true,
+                  !isLeftMouseButtonDown,
+                  let axWindows = axApp.threadGuarded.get(Ax.windowsAttr)
+            else {
+                return (focusedWindow.windowId, [])
+            }
+
+            // Inactive native tabs keep valid window IDs, so AXWindows is the source of truth for these apps.
+            var activeWindowIds = Set(axWindows.map { $0.0 })
+            activeWindowIds.insert(focusedWindow.windowId)
+            let staleWindowIds = windows.threadGuarded.keys.filter { !activeWindowIds.contains($0) }
+            for staleWindowId in staleWindowIds {
+                windows.threadGuarded.removeValue(forKey: staleWindowId)
+            }
+            return (focusedWindow.windowId, staleWindowIds)
+        }
+        guard let (windowId, staleWindowIds) = focused else { return nil }
+        for staleWindowId in staleWindowIds {
+            setFrameJobs.removeValue(forKey: staleWindowId)?.cancel()
         }
-        guard let windowId else { return nil }
-        return try await MacWindow.getOrRegister(windowId: windowId, macApp: self)
+        return try await MacWindow.getOrRegister(
+            windowId: windowId,
+            macApp: self,
+            replacingNativeTabWindowIds: staleWindowIds,
+            lastFocusedWindowId: previousFocusedWindowId,
+        )
     }
 
     @MainActor func nativeFocus(_ windowId: UInt32) {
@@ -303,19 +334,27 @@ final class MacApp: AbstractApp {
             return []
         }
         guard let thread else { return [] }
-        let (alive, dead) = try await thread.runInLoop(.cancellable) { [nsApp, windows, axApp] (job) -> ([UInt32], [UInt32]) in
+        let (alive, dead) = try await thread.runInLoop(.cancellable) { [nsApp, appId, windows, axApp] (job) -> ([UInt32], [UInt32]) in
             var alive: [UInt32: AxWindow] = windows.threadGuarded
             var dead = [UInt32: AxWindow]()
+            let axWindows = axApp.threadGuarded.get(Ax.windowsAttr)
+            let activeNativeTabWindowIds: Set<UInt32>? =
+                appId?.exposesInactiveNativeTabsAsWindows == true && !isLeftMouseButtonDown
+                    ? axWindows.map { Set($0.map { $0.0 }) }
+                    : nil
             // Second line of defence against lock screen. See the first line of defence: closedWindowsCache
             // Second and third lines of defence are technically needed only to avoid potential flickering
             if frontmostAppBundleId != lockScreenAppBundleId {
                 (alive, dead) = try alive.partition {
                     try job.checkCancellation()
-                    return $0.value.ax.containingWindowId() != nil
+                    if $0.value.ax.containingWindowId() == nil {
+                        return false
+                    }
+                    return activeNativeTabWindowIds?.contains($0.key) != false
                 }
             }
 
-            for (id, window) in axApp.threadGuarded.get(Ax.windowsAttr) ?? [] {
+            for (id, window) in axWindows ?? [] {
                 try job.checkCancellation()
                 try alive.getOrRegisterAxWindow(windowId: id, window, nsApp, job)
             }

Sources/AppBundle/tree/MacWindow.swift

@@ -16,9 +16,27 @@ final class MacWindow: Window {
 
     @MainActor
     @discardableResult
-    static func getOrRegister(windowId: UInt32, macApp: MacApp) async throws -> MacWindow {
-        if let existing = allWindowsMap[windowId] { return existing }
+    static func getOrRegister(
+        windowId: UInt32,
+        macApp: MacApp,
+        replacingNativeTabWindowIds: [UInt32] = [],
+        lastFocusedWindowId: UInt32? = nil,
+    ) async throws -> MacWindow {
+        if let existing = allWindowsMap[windowId] {
+            discardNativeTabWindows(replacingNativeTabWindowIds, from: macApp)
+            return existing
+        }
         let rect = try await macApp.getAxRect(windowId, .cancellable)
+        if let replacement = registerNativeTabReplacement(
+            windowId: windowId,
+            macApp: macApp,
+            rect: rect,
+            staleWindowIds: replacingNativeTabWindowIds,
+            lastFocusedWindowId: lastFocusedWindowId,
+        ) {
+            try await debugWindowsIfRecording(replacement, .cancellable)
+            return replacement
+        }
         let data = try await unbindAndGetBindingDataForNewWindow(
             windowId,
             macApp,
@@ -30,9 +48,23 @@ final class MacWindow: Window {
         )
 
         // atomic synchronous section
-        if let existing = allWindowsMap[windowId] { return existing }
+        if let existing = allWindowsMap[windowId] {
+            discardNativeTabWindows(replacingNativeTabWindowIds, from: macApp)
+            return existing
+        }
+        if let replacement = registerNativeTabReplacement(
+            windowId: windowId,
+            macApp: macApp,
+            rect: rect,
+            staleWindowIds: replacingNativeTabWindowIds,
+            lastFocusedWindowId: lastFocusedWindowId,
+        ) {
+            try await debugWindowsIfRecording(replacement, .cancellable)
+            return replacement
+        }
         let window = MacWindow(windowId, macApp, lastFloatingSize: rect?.size, parent: data.parent, adaptiveWeight: data.adaptiveWeight, index: data.index)
         allWindowsMap[windowId] = window
+        discardNativeTabWindows(replacingNativeTabWindowIds, from: macApp)
 
         try await debugWindowsIfRecording(window, .cancellable)
         if try await !restoreClosedWindowsCacheIfNeeded(newlyDetectedWindow: window) {
@@ -41,6 +73,55 @@ final class MacWindow: Window {
         return window
     }
 
+    @MainActor
+    private static func registerNativeTabReplacement(
+        windowId: UInt32,
+        macApp: MacApp,
+        rect: Rect?,
+        staleWindowIds: [UInt32],
+        lastFocusedWindowId: UInt32?,
+    ) -> MacWindow? {
+        let staleWindows = staleWindowIds.compactMap { allWindowsMap[$0] }.filter { $0.macApp === macApp }
+        guard let (oldAbstractWindow, bindingData) = takeNativeTabReplacementBinding(
+            from: staleWindows,
+            lastFocusedWindowId: lastFocusedWindowId,
+        ) else {
+            return nil
+        }
+        let oldWindow = oldAbstractWindow as! MacWindow
+        _ = allWindowsMap.removeValue(forKey: oldWindow.windowId)
+
+        let window = MacWindow(
+            windowId,
+            macApp,
+            lastFloatingSize: oldWindow.lastFloatingSize ?? rect?.size,
+            parent: bindingData.parent,
+            adaptiveWeight: bindingData.adaptiveWeight,
+            index: bindingData.index,
+        )
+        window.isFullscreen = oldWindow.isFullscreen
+        window.noOuterGapsInFullscreen = oldWindow.noOuterGapsInFullscreen
+        window.layoutReason = oldWindow.layoutReason
+        window.lastAppliedLayoutVirtualRect = oldWindow.lastAppliedLayoutVirtualRect
+        window.lastAppliedLayoutPhysicalRect = oldWindow.lastAppliedLayoutPhysicalRect
+        window.prevUnhiddenProportionalPositionInsideWorkspaceRect = oldWindow.prevUnhiddenProportionalPositionInsideWorkspaceRect
+        allWindowsMap[windowId] = window
+        discardNativeTabWindows(staleWindowIds, from: macApp)
+        // A replacement tab is the same logical window, so don't run on-window-detected again.
+        return window
+    }
+
+    @MainActor
+    private static func discardNativeTabWindows(_ windowIds: [UInt32], from macApp: MacApp) {
+        for windowId in windowIds {
+            guard let window = allWindowsMap[windowId], window.macApp === macApp else { continue }
+            allWindowsMap.removeValue(forKey: windowId)
+            if window.isBound {
+                window.unbindFromParent()
+            }
+        }
+    }
+
     // var description: String {
     //     let description = [
     //         ("title", title),
@@ -199,6 +280,19 @@ final class MacWindow: Window {
     }
 }
 
+@MainActor
+func takeNativeTabReplacementBinding(
+    from staleWindows: [Window],
+    lastFocusedWindowId: UInt32?,
+) -> (window: Window, bindingData: BindingData)? {
+    let staleWindows = staleWindows.filter(\.isBound)
+    let lastFocusedWindow = lastFocusedWindowId.flatMap { id in
+        staleWindows.first { $0.windowId == id }
+    }
+    guard let replacement = lastFocusedWindow ?? staleWindows.singleOrNil() else { return nil }
+    return (replacement, replacement.unbindFromParent())
+}
+
 extension Window {
     @MainActor
     func relayoutWindow(on workspace: Workspace, _ cm: CancellationMode, forceTile: Bool = false) async throws {

Sources/AppBundleTests/tree/NativeTabWindowReplacementTest.swift

@@ -0,0 +1,66 @@
+@testable import AppBundle
+import XCTest
+
+@MainActor
+final class NativeTabWindowReplacementTest: XCTestCase {
+    override func setUp() async throws { setUpWorkspacesForTests() }
+
+    func testReplacementPreservesNestedTreeSlotAndWeight() throws {
+        let root = Workspace.get(byName: name).rootTilingContainer
+        let container = TilingContainer.newHTiles(parent: root, adaptiveWeight: 1)
+        TestWindow.new(id: 1, parent: container)
+        let oldWindow = TestWindow.new(id: 2, parent: container, adaptiveWeight: 3)
+        TestWindow.new(id: 3, parent: container)
+
+        let replacement = takeNativeTabReplacementBinding(
+            from: [oldWindow],
+            lastFocusedWindowId: nil,
+        )
+        let bindingData = try XCTUnwrap(replacement).bindingData
+        let newWindow = TestWindow.new(id: 4, parent: root)
+        newWindow.bind(
+            to: bindingData.parent,
+            adaptiveWeight: bindingData.adaptiveWeight,
+            index: bindingData.index,
+        )
+
+        assertEquals(container.layoutDescription, .h_tiles([.window(1), .window(4), .window(3)]))
+        assertEquals(newWindow.getWeight(.h), 3)
+    }
+
+    func testReplacementPrefersLastFocusedWindow() {
+        let root = Workspace.get(byName: name).rootTilingContainer
+        let first = TestWindow.new(id: 1, parent: root)
+        let lastFocused = TestWindow.new(id: 2, parent: root)
+
+        let replacement = takeNativeTabReplacementBinding(
+            from: [first, lastFocused],
+            lastFocusedWindowId: lastFocused.windowId,
+        )
+
+        assertEquals(replacement?.window.windowId, lastFocused.windowId)
+        XCTAssertTrue(first.isBound)
+        XCTAssertFalse(lastFocused.isBound)
+    }
+
+    func testReplacementDoesNotGuessWhenStaleWindowsAreAmbiguous() {
+        let root = Workspace.get(byName: name).rootTilingContainer
+        let first = TestWindow.new(id: 1, parent: root)
+        let second = TestWindow.new(id: 2, parent: root)
+
+        let replacement = takeNativeTabReplacementBinding(
+            from: [first, second],
+            lastFocusedWindowId: nil,
+        )
+
+        XCTAssertNil(replacement)
+        XCTAssertTrue(first.isBound)
+        XCTAssertTrue(second.isBound)
+    }
+
+    func testAffectedBundleIds() {
+        XCTAssertTrue(KnownBundleId.ghostty.exposesInactiveNativeTabsAsWindows)
+        XCTAssertTrue(KnownBundleId.fork.exposesInactiveNativeTabsAsWindows)
+        XCTAssertFalse(KnownBundleId.iterm2.exposesInactiveNativeTabsAsWindows)
+    }
+}