macos: basic undo close window, not very robust yet

macos/Sources/Features/Splits/SplitTree.swift

@@ -107,6 +107,18 @@ extension SplitTree {
         self.init(root: .leaf(view: view), zoomed: nil)
     }
 
+    /// Checks if the tree contains the specified node.
+    ///
+    /// Note that SplitTree implements Sequence on views so there's already a `contains`
+    /// for views too.
+    ///
+    /// - Parameter node: The node to search for in the tree
+    /// - Returns: True if the node exists in the tree, false otherwise
+    func contains(_ node: Node) -> Bool {
+        guard let root else { return false }
+        return root.path(to: node) != nil
+    }
+
     /// Insert a new view at the given view point by creating a split in the given direction.
     /// This will always reset the zoomed state of the tree.
     func insert(view: ViewType, at: ViewType, direction: NewDirection) throws -> Self {
@@ -1078,3 +1090,29 @@ extension SplitTree.Node: Sequence {
         return leaves().makeIterator()
     }
 }
+
+// MARK: SplitTree Collection
+
+extension SplitTree: Collection {
+    typealias Index = Int
+    typealias Element = ViewType
+
+    var startIndex: Int {
+        return 0
+    }
+
+    var endIndex: Int {
+        return root?.leaves().count ?? 0
+    }
+
+    subscript(position: Int) -> ViewType {
+        precondition(position >= 0 && position < endIndex, "Index out of bounds")
+        let leaves = root?.leaves() ?? []
+        return leaves[position]
+    }
+
+    func index(after i: Int) -> Int {
+        precondition(i < endIndex, "Cannot increment index beyond endIndex")
+        return i + 1
+    }
+}

macos/Sources/Features/Terminal/BaseTerminalController.swift

@@ -76,7 +76,7 @@ class BaseTerminalController: NSWindowController,
     private var focusedSurfaceCancellables: Set<AnyCancellable> = []
 
     /// The time that undo/redo operations that contain running ptys are valid for.
-    private var undoExpiration: Duration {
+    var undoExpiration: Duration {
         .seconds(5)
     }
 
@@ -277,7 +277,11 @@ class BaseTerminalController: NSWindowController,
     }
     
     /// Remove a node from the surface tree and move focus appropriately.
-    private func removeSurfaceAndMoveFocus(_ node: SplitTree<Ghostty.SurfaceView>.Node) {
+    ///
+    /// This also updates the undo manager to support restoring this node.
+    ///
+    /// This does no confirmation and assumes confirmation is already done.
+    private func removeSurfaceNode(_ node: SplitTree<Ghostty.SurfaceView>.Node) {
         let nextTarget = findNextFocusTargetAfterClosing(node: node)
         let oldFocused = focusedSurface
         let focused = node.contains { $0 == focusedSurface }
@@ -311,8 +315,8 @@ class BaseTerminalController: NSWindowController,
                 undoManager.registerUndo(
                     withTarget: target,
                     expiresAfter: target.undoExpiration) { target in
-                    target.closeSurface(
+                    target.closeSurfaceNode(
-                        node.leftmostLeaf(),
+                        node,
                         withConfirmation: node.contains {
                             $0.needsConfirmQuit
                         }
@@ -397,25 +401,26 @@ class BaseTerminalController: NSWindowController,
 
     @objc private func ghosttyDidCloseSurface(_ notification: Notification) {
         guard let target = notification.object as? Ghostty.SurfaceView else { return }
-        closeSurface(
+        guard let node = surfaceTree.root?.node(view: target) else { return }
-            target,
+        closeSurfaceNode(
+            node,
             withConfirmation: (notification.userInfo?["process_alive"] as? Bool) ?? false,
         )
     }
 
-    /// Close a surface view, requesting confirmation if necessary.
+    /// Close a surface node (which may contain splits), requesting confirmation if necessary.
     ///
     /// This will also insert the proper undo stack information in.
-    private func closeSurface(
+    func closeSurfaceNode(
-        _ target: Ghostty.SurfaceView,
+        _ node: SplitTree<Ghostty.SurfaceView>.Node,
         withConfirmation: Bool = true,
     ) {
-        // The target must be within our tree
+        // This node must be part of our tree
-        guard let node = surfaceTree.root?.node(view: target) else { return }
+        guard surfaceTree.contains(node) else { return }
 
         // If the child process is not alive, then we exit immediately
         guard withConfirmation else {
-            removeSurfaceAndMoveFocus(node)
+            removeSurfaceNode(node)
             return
         }
 
@@ -429,7 +434,7 @@ class BaseTerminalController: NSWindowController,
             informativeText: "The terminal still has a running process. If you close the terminal the process will be killed."
         ) { [weak self] in
             if let self {
-                self.removeSurfaceAndMoveFocus(node)
+                self.removeSurfaceNode(node)
             }
         }
     }

macos/Sources/Features/Terminal/TerminalController.swift

@@ -386,6 +386,93 @@ class TerminalController: BaseTerminalController {
         return frame
     }
 
+    /// This is called anytime a node in the surface tree is being removed.
+    override func closeSurfaceNode(
+        _ node: SplitTree<Ghostty.SurfaceView>.Node,
+        withConfirmation: Bool = true
+    ) {
+        // If this isn't the root then we're dealing with a split closure.
+        if surfaceTree.root != node {
+            super.closeSurfaceNode(node, withConfirmation: withConfirmation)
+            return
+        }
+
+        // More than 1 window means we have tabs and we're closing a tab
+        if window?.tabGroup?.windows.count ?? 0 > 1 {
+            closeTab(nil)
+            return
+        }
+
+        // 1 window, closing the window
+        closeWindow(nil)
+    }
+
+    /// Closes the current window (including any other tabs) immediately and without
+    /// confirmation. This will setup proper undo state so the action can be undone.
+    private func closeWindowImmediately(_ sender: Any?) {
+        guard let window = window else { return }
+
+        // Regardless of tabs vs no tabs, what we want to do here is keep
+        // track of the window frame to restore, the surface tree, and the
+        // the focused surface. We want to restore that with undo even
+        // if we end up closing.
+        if let undoManager {
+            // Capture current state for undo
+            let currentFrame = window.frame
+            let currentSurfaceTree = surfaceTree
+            let currentFocusedSurface = focusedSurface
+
+            // Register undo action to restore the window
+            undoManager.setActionName("Close Window")
+            undoManager.registerUndo(
+                withTarget: ghostty,
+                expiresAfter: undoExpiration) { ghostty in
+
+                // Create a new window controller with the saved state
+                let newController = TerminalController(
+                    ghostty,
+                    withSurfaceTree: currentSurfaceTree
+                )
+                
+                // Show the window and restore its frame
+                newController.showWindow(nil)
+                if let newWindow = newController.window {
+                    newWindow.setFrame(currentFrame, display: true)
+                    
+                    // Restore focus to the previously focused surface
+                    if let focusTarget = currentFocusedSurface {
+                        DispatchQueue.main.async {
+                            Ghostty.moveFocus(to: focusTarget, from: nil)
+                        }
+                    }
+                }
+                
+                // Register redo action
+                undoManager.registerUndo(
+                    withTarget: newController,
+                    expiresAfter: newController.undoExpiration) { target in
+                    // For redo, we close the window again
+                    target.closeWindowImmediately(sender)
+                }
+            }
+        }
+
+        guard let tabGroup = window.tabGroup else {
+            // No tabs, no tab group, just perform a normal close.
+            window.close()
+            return
+        }
+
+        // If have one window then we just do a normal close
+        if tabGroup.windows.count == 1 {
+            window.close()
+            return
+        }
+
+
+        tabGroup.windows.forEach { $0.close() }
+    }
+
     //MARK: - NSWindowController
 
     override func windowWillLoad() {
@@ -635,13 +722,13 @@ class TerminalController: BaseTerminalController {
         guard let window = window else { return }
         guard let tabGroup = window.tabGroup else {
             // No tabs, no tab group, just perform a normal close.
-            window.performClose(sender)
+            closeWindowImmediately(sender)
             return
         }
 
         // If have one window then we just do a normal close
         if tabGroup.windows.count == 1 {
-            window.performClose(sender)
+            closeWindowImmediately(sender)
             return
         }
 
@@ -655,7 +742,7 @@ class TerminalController: BaseTerminalController {
 
         // If none need confirmation then we can just close all the windows.
         if !needsConfirm {
-            tabGroup.windows.forEach { $0.close() }
+            closeWindowImmediately(sender)
             return
         }
 
@@ -663,7 +750,7 @@ class TerminalController: BaseTerminalController {
             messageText: "Close Window?",
             informativeText: "All terminal sessions in this window will be terminated."
         ) {
-            tabGroup.windows.forEach { $0.close() }
+            self.closeWindowImmediately(sender)
         }
     }