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, 241 insertions, 0 deletions

diff --git a/src/platform/linux.rs b/src/platform/linux.rs
index 7dd2879e..c73584e5 100644
--- a/src/platform/linux.rs
+++ b/src/platform/linux.rs
@@ -26,6 +26,9 @@ pub(crate) use super::unix_common::{
 #[cfg(test)]
 mod config_file_tests;
 
+mod shutdown;
+pub(crate) use shutdown::monitor_host_shutdown;
+
 const WSL_MARKER_ENV_VARS: &[&str] = &["WSL_DISTRO_NAME", "WSL_INTEROP"];
 const PROCESS_DETECTION_ENV_VAR: &str = "HERDR_PROCESS_DETECTION";
 const CHILD_GROUPS_SCAN_LIMIT: usize = 64;
diff --git a/src/platform/linux/shutdown.rs b/src/platform/linux/shutdown.rs
new file mode 100644
index 00000000..b678161a
--- /dev/null
+++ b/src/platform/linux/shutdown.rs
@@ -0,0 +1,208 @@
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures_util::StreamExt;
+
+pub(crate) fn monitor_host_shutdown(
+    requested: Arc<AtomicBool>,
+    wake: impl Fn() + Send + Sync + 'static,
+) -> Option<tokio::task::JoinHandle<()>> {
+    Some(tokio::spawn(async move {
+        let mut retry = Duration::from_secs(1);
+        loop {
+            match watch_shutdown(&requested, &wake).await {
+                Ok(()) => retry = Duration::from_secs(1),
+                Err(err) => {
+                    tracing::debug!(err = %err, retry_seconds = retry.as_secs(), "host shutdown notification unavailable");
+                    tokio::time::sleep(retry).await;
+                    retry = (retry * 2).min(Duration::from_secs(60));
+                }
+            }
+        }
+    }))
+}
+
+async fn watch_shutdown(
+    requested: &AtomicBool,
+    wake: &(impl Fn() + Send + Sync),
+) -> zbus::Result<()> {
+    let connection = zbus::Connection::system().await?;
+    watch_connection(connection, requested, wake).await
+}
+
+async fn watch_connection(
+    connection: zbus::Connection,
+    requested: &AtomicBool,
+    wake: &(impl Fn() + Send + Sync),
+) -> zbus::Result<()> {
+    let manager = zbus::Proxy::new(
+        &connection,
+        "org.freedesktop.login1",
+        "/org/freedesktop/login1",
+        "org.freedesktop.login1.Manager",
+    )
+    .await?;
+    let mut owners = manager.receive_owner_changed().await?;
+    let mut signals = manager.receive_signal("PrepareForShutdown").await?;
+    let inhibitor: zbus::zvariant::OwnedFd = manager
+        .call(
+            "Inhibit",
+            &(
+                "shutdown",
+                "Herdr",
+                "Save terminal workspace layout",
+                "delay",
+            ),
+        )
+        .await?;
+    tracing::debug!("host shutdown notification ready");
+
+    let mut preparing: bool = manager.get_property("PreparingForShutdown").await?;
+    loop {
+        if preparing {
+            tracing::info!("host shutdown requested; preserving session before pane termination");
+            requested.store(true, Ordering::Release);
+            wake();
+            // Keep the delay lock until the server has saved and drops its monitor.
+            // logind caps the delay even if the server gets stuck.
+            std::future::pending::<()>().await;
+        }
+        tokio::select! {
+            biased;
+            _ = owners.next() => break,
+            signal = signals.next() => {
+                let Some(signal) = signal else { break };
+                preparing = signal.body().deserialize()?;
+            }
+        }
+    }
+    drop(inhibitor);
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::{BufRead, Read};
+    use std::os::unix::net::UnixStream;
+    use std::process::{Child, Command, Stdio};
+    use std::sync::Mutex;
+
+    struct PrivateBus(Child);
+
+    impl Drop for PrivateBus {
+        fn drop(&mut self) {
+            let _ = self.0.kill();
+            let _ = self.0.wait();
+        }
+    }
+
+    struct LoginManager {
+        preparing: bool,
+        peer: Arc<Mutex<Option<UnixStream>>>,
+    }
+
+    #[zbus::interface(name = "org.freedesktop.login1.Manager")]
+    impl LoginManager {
+        fn inhibit(
+            &self,
+            what: &str,
+            _who: &str,
+            _why: &str,
+            mode: &str,
+        ) -> zbus::fdo::Result<zbus::zvariant::OwnedFd> {
+            assert_eq!((what, mode), ("shutdown", "delay"));
+            let (lock, peer) = UnixStream::pair().unwrap();
+            peer.set_nonblocking(true).unwrap();
+            *self.peer.lock().unwrap() = Some(peer);
+            Ok(std::os::fd::OwnedFd::from(lock).into())
+        }
+
+        #[zbus(property)]
+        fn preparing_for_shutdown(&self) -> bool {
+            self.preparing
+        }
+    }
+
+    #[tokio::test]
+    #[ignore = "requires dbus-daemon; uses a private bus, never requests host shutdown"]
+    async fn shutdown_warning_holds_inhibitor_until_monitor_is_dropped() {
+        for already_preparing in [false, true] {
+            let mut bus = PrivateBus(
+                Command::new("dbus-daemon")
+                    .args(["--session", "--nofork", "--print-address=1"])
+                    .stdout(Stdio::piped())
+                    .spawn()
+                    .unwrap(),
+            );
+            let mut address = String::new();
+            std::io::BufReader::new(bus.0.stdout.take().unwrap())
+                .read_line(&mut address)
+                .unwrap();
+            let peer = Arc::new(Mutex::new(None));
+            let service = zbus::connection::Builder::address(address.trim())
+                .unwrap()
+                .name("org.freedesktop.login1")
+                .unwrap()
+                .serve_at(
+                    "/org/freedesktop/login1",
+                    LoginManager {
+                        preparing: already_preparing,
+                        peer: peer.clone(),
+                    },
+                )
+                .unwrap()
+                .build()
+                .await
+                .unwrap();
+            let client = zbus::connection::Builder::address(address.trim())
+                .unwrap()
+                .build()
+                .await
+                .unwrap();
+            let requested = Arc::new(AtomicBool::new(false));
+            let wake = Arc::new(tokio::sync::Notify::new());
+            let task = tokio::spawn({
+                let requested = requested.clone();
+                let wake = wake.clone();
+                async move {
+                    watch_connection(client, &requested, &move || wake.notify_one())
+                        .await
+                        .unwrap();
+                }
+            });
+            tokio::time::timeout(Duration::from_secs(5), async {
+                while peer.lock().unwrap().is_none() {
+                    tokio::task::yield_now().await;
+                }
+                if !already_preparing {
+                    assert!(!requested.load(Ordering::Acquire));
+                    service
+                        .emit_signal(
+                            None::<&str>,
+                            "/org/freedesktop/login1",
+                            "org.freedesktop.login1.Manager",
+                            "PrepareForShutdown",
+                            &true,
+                        )
+                        .await
+                        .unwrap();
+                }
+                wake.notified().await;
+            })
+            .await
+            .unwrap();
+            assert!(requested.load(Ordering::Acquire));
+            let mut peer = peer.lock().unwrap().take().unwrap();
+            assert_eq!(
+                peer.read(&mut [0]).unwrap_err().kind(),
+                std::io::ErrorKind::WouldBlock,
+                "shutdown must remain inhibited while the server saves"
+            );
+            task.abort();
+            assert!(task.await.unwrap_err().is_cancelled());
+            assert_eq!(peer.read(&mut [0]).unwrap(), 0);
+        }
+    }
+}
diff --git a/src/platform/mod.rs b/src/platform/mod.rs
index 2ba3e6f3..488c5bee 100644
--- a/src/platform/mod.rs
+++ b/src/platform/mod.rs
@@ -3,6 +3,36 @@
 //! Centralizes OS-dependent behavior behind a clean boundary so core
 //! modules don't scatter `#[cfg]` branches through product logic.
 
+pub(crate) struct HostShutdownMonitor {
+    task: Option<tokio::task::JoinHandle<()>>,
+}
+
+impl HostShutdownMonitor {
+    pub(crate) fn start(
+        requested: std::sync::Arc<std::sync::atomic::AtomicBool>,
+        wake: impl Fn() + Send + Sync + 'static,
+    ) -> Self {
+        let task = monitor_host_shutdown(requested, wake);
+        Self { task }
+    }
+}
+
+impl Drop for HostShutdownMonitor {
+    fn drop(&mut self) {
+        if let Some(task) = self.task.take() {
+            task.abort();
+        }
+    }
+}
+
+#[cfg(not(target_os = "linux"))]
+fn monitor_host_shutdown(
+    _requested: std::sync::Arc<std::sync::atomic::AtomicBool>,
+    _wake: impl Fn() + Send + Sync + 'static,
+) -> Option<tokio::task::JoinHandle<()>> {
+    None
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ForegroundProcess {
     pub pid: u32,