fix: preserve session layouts across shutdown and restore failures (#4400)

* fix: preserve session layouts across shutdown and restore failures

refs #4320

* test: update subscription fixture for pane restore errors

---------

Co-authored-by: JJ Liebig <jonathan.liebig@gmail.com>

3 files changed, 57 insertions, 1 deletions

diff --git a/src/server/headless.rs b/src/server/headless.rs
index f07c0c37..638559ad 100644
--- a/src/server/headless.rs
+++ b/src/server/headless.rs
@@ -250,6 +250,7 @@ pub struct HeadlessServer {
     pending_handoff_repaint_nudge: bool,
     /// Flag set by Ctrl+C or `server stop` signal.
     should_quit: Arc<AtomicBool>,
+    host_shutdown_requested: Arc<AtomicBool>,
     /// Channel for receiving server events from client connection threads.
     server_event_rx: mpsc::Receiver<ServerEvent>,
     /// Sender for server events (cloned for each client thread).
@@ -373,6 +374,7 @@ impl HeadlessServer {
             headless_size,
             effective_size: headless_size,
             shutting_down: false,
+            host_shutdown_requested: Arc::new(AtomicBool::new(false)),
             handoff_in_progress: false,
             #[cfg(unix)]
             pending_handoff_repaint_nudge: false,
@@ -398,6 +400,13 @@ impl HeadlessServer {
         let should_quit = self.should_quit.clone();
         let quit_notify = self.server_event_tx.clone();
         ctrlc_handler(should_quit, quit_notify);
+        let quit_notify = self.server_event_tx.clone();
+        let _host_shutdown = crate::platform::HostShutdownMonitor::start(
+            self.host_shutdown_requested.clone(),
+            move || {
+                let _ = quit_notify.try_send(ServerEvent::QuitSignal);
+            },
+        );
 
         let mut needs_render = true;
         let mut needs_full_render = true;
@@ -414,6 +423,13 @@ impl HeadlessServer {
                 break;
             }
 
+            // A host shutdown warning precedes process termination. Do not drain pane
+            // deaths here: logind's delay lock stays held until the final session save.
+            if self.host_shutdown_requested.load(Ordering::Acquire) {
+                self.initiate_shutdown();
+                continue;
+            }
+
             // Check if we should start shutting down.
             if self.app.state.should_quit || self.should_quit.load(Ordering::Acquire) {
                 self.drain_internal_events_with_forwarding_up_to(
@@ -638,7 +654,9 @@ impl HeadlessServer {
                 }
             };
 
-            if self.should_quit.load(Ordering::Acquire) {
+            if self.should_quit.load(Ordering::Acquire)
+                || self.host_shutdown_requested.load(Ordering::Acquire)
+            {
                 match event {
                     LoopEvent::Internal(ev) => {
                         self.handle_internal_event_with_forwarding(ev);
diff --git a/src/server/headless/notifications.rs b/src/server/headless/notifications.rs
index 39938b74..722f95c5 100644
--- a/src/server/headless/notifications.rs
+++ b/src/server/headless/notifications.rs
@@ -321,6 +321,9 @@ impl HeadlessServer {
     ///
     /// Returns true if the event changed visual state (requiring a re-render).
     pub(super) fn handle_internal_event_with_forwarding(&mut self, mut ev: AppEvent) -> bool {
+        if self.host_shutdown_requested.load(Ordering::Acquire) {
+            return false;
+        }
         let mut focused_worktree_response = if let AppEvent::WorktreeAddFinished(result) = &mut ev {
             result
                 .api_request
@@ -752,6 +755,9 @@ impl HeadlessServer {
         let mut had_event = false;
         let mut changed = false;
         for _ in 0..limit {
+            if self.host_shutdown_requested.load(Ordering::Acquire) {
+                break;
+            }
             let Ok(ev) = self.app.event_rx.try_recv() else {
                 break;
             };
diff --git a/src/server/headless/tests/mod.rs b/src/server/headless/tests/mod.rs
index 60f89d63..cf6d389b 100644
--- a/src/server/headless/tests/mod.rs
+++ b/src/server/headless/tests/mod.rs
@@ -96,6 +96,7 @@ fn test_headless_server_with_event_hub(event_hub: api::EventHub) -> HeadlessServ
         headless_size,
         effective_size: headless_size,
         shutting_down: false,
+        host_shutdown_requested: Arc::new(AtomicBool::new(false)),
         handoff_in_progress: false,
         #[cfg(unix)]
         pending_handoff_repaint_nudge: false,
@@ -4186,6 +4187,37 @@ fn changed_git_refresh_requests_headless_render() {
 }
 
 #[tokio::test]
+async fn host_shutdown_preserves_panes_from_queued_and_selected_death_events() {
+    let mut server = test_headless_server();
+    let workspace = crate::workspace::Workspace::test_new("host-shutdown");
+    let pane_id = workspace.tabs[0].root_pane;
+    server.app.state.workspaces = vec![workspace];
+    server.app.state.ensure_test_terminals();
+    server.app.state.active = Some(0);
+    let event = || AppEvent::PaneDied {
+        pane_id,
+        exit_reason: crate::platform::ChildExitReason::Exited,
+    };
+    server.app.event_tx.try_send(event()).unwrap();
+    server
+        .host_shutdown_requested
+        .store(true, Ordering::Release);
+    assert_eq!(
+        server.drain_internal_events_with_forwarding_up_to(16),
+        (false, false)
+    );
+    assert!(!server.handle_internal_event_with_forwarding(event()));
+    assert!(server.app.find_pane(pane_id).is_some());
+    assert!(server.app.event_rx.try_recv().is_ok());
+    server
+        .host_shutdown_requested
+        .store(false, Ordering::Release);
+    assert!(server.handle_internal_event_with_forwarding(event()));
+    assert!(server.app.find_pane(pane_id).is_none());
+    shutdown_test_runtimes(&mut server);
+}
+
+#[tokio::test]
 async fn pane_death_reconciles_each_client_view_and_focus() {
     let mut server = test_headless_server();
     let mut workspace = crate::workspace::Workspace::test_new("pane-death-views");