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>

2 files changed, 264 insertions, 14 deletions

diff --git a/tests/client_mode.rs b/tests/client_mode.rs
index 0b94b32e..9bcb4b2c 100644
--- a/tests/client_mode.rs
+++ b/tests/client_mode.rs
@@ -1796,10 +1796,7 @@ fn client_receives_pane_surface_after_pane_output() {
 }
 
 #[test]
-fn pane_spawn_cwd_fallback_in_server() {
-    // Pane spawn failure cwd fallback in server context.
-    // This test verifies that the server can start even with invalid
-    // session data pointing to non-existent directories.
+fn unavailable_restored_pane_keeps_saved_cwd_in_server() {
     let _lock = test_lock();
     let base = unique_test_dir();
     let config_home = base.join("config");
@@ -1854,12 +1851,11 @@ fn pane_spawn_cwd_fallback_in_server() {
     assert_eq!(pane["result"]["pane"]["workspace_id"], workspace_id);
     let cwd = pane["result"]["pane"]["cwd"]
         .as_str()
-        .expect("restored pane should report fallback cwd");
-    assert_ne!(cwd, missing_cwd);
-    assert!(
-        std::path::Path::new(cwd).exists(),
-        "fallback cwd should exist: {cwd}"
-    );
+        .expect("restored pane should retain saved cwd");
+    assert_eq!(cwd, missing_cwd);
+    assert!(pane["result"]["pane"]["restore_error"]
+        .as_str()
+        .is_some_and(|error| error.contains("directory")));
 
     let client_shell = spawn_client_shell_process(&config_home, &runtime_dir, &api_socket);
     let output = spawn_pty_drain(
@@ -1872,14 +1868,47 @@ fn pane_spawn_cwd_fallback_in_server() {
     );
     assert!(
         wait_until(Duration::from_secs(8), Duration::from_millis(20), || {
-            read_output(&output).contains("missing-cwd")
+            let screen = read_output(&output);
+            screen.contains("missing-cwd") && screen.contains("unavailable")
         }),
-        "client shell should render the restored session; output: {:?}",
+        "client shell should render the unavailable pane; output: {:?}",
         read_output(&output)
     );
-
+    drop(client_shell);
+    let stopped = send_json_request(
+        &api_socket,
+        r#"{"id":"stop","method":"server.stop","params":{}}"#,
+    );
+    assert!(stopped.get("error").is_none(), "{stopped}");
+    let mut spawned = spawned;
+    assert!(wait_until(
+        Duration::from_secs(10),
+        Duration::from_millis(20),
+        || { spawned.child.try_wait().unwrap().is_some() }
+    ));
     drop(spawned);
-    cleanup_spawned_herdr(client_shell, base);
+
+    fs::create_dir(missing_cwd).unwrap();
+    let restarted = spawn_server(&config_home, &runtime_dir, &api_socket, &client_socket);
+    wait_for_socket(&api_socket, Duration::from_secs(10));
+    let recovered = send_json_request(
+        &api_socket,
+        &format!(r#"{{"id":"recovered","method":"pane.get","params":{{"pane_id":"{pane_id}"}}}}"#),
+    );
+    assert_eq!(
+        std::fs::canonicalize(recovered["result"]["pane"]["cwd"].as_str().unwrap()).unwrap(),
+        std::fs::canonicalize(missing_cwd).unwrap()
+    );
+    assert!(recovered["result"]["pane"]["restore_error"].is_null());
+    let sent = send_json_request(
+        &api_socket,
+        &serde_json::json!({"id": "type", "method": "pane.send_text", "params": {
+            "pane_id": pane_id, "text": "printf 'RESTORE_RETRY_OK\\n'\n"
+        }})
+        .to_string(),
+    );
+    assert!(sent.get("error").is_none(), "{sent}");
+    cleanup_spawned_herdr(restarted, base);
 }
 
 #[test]
diff --git a/tests/host_shutdown.rs b/tests/host_shutdown.rs
new file mode 100644
index 00000000..f988a52a
--- /dev/null
+++ b/tests/host_shutdown.rs
@@ -0,0 +1,221 @@
+#![cfg(target_os = "linux")]
+
+use std::io::{BufRead, Read, Write};
+use std::os::unix::net::UnixStream;
+use std::path::Path;
+use std::process::{Child, Command, Stdio};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+struct ChildGuard(Child);
+
+impl Drop for ChildGuard {
+    fn drop(&mut self) {
+        let _ = self.0.kill();
+        let _ = self.0.wait();
+    }
+}
+
+struct LoginManager {
+    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 {
+        false
+    }
+}
+
+fn private_bus(address: &str) -> ChildGuard {
+    let mut child = ChildGuard(
+        Command::new("dbus-daemon")
+            .args([
+                "--session",
+                "--nofork",
+                "--print-address=1",
+                "--address",
+                address,
+            ])
+            .stdout(Stdio::piped())
+            .spawn()
+            .unwrap(),
+    );
+    let mut ready = String::new();
+    std::io::BufReader::new(child.0.stdout.take().unwrap())
+        .read_line(&mut ready)
+        .unwrap();
+    child
+}
+
+async fn login_service(address: &str, peer: Arc<Mutex<Option<UnixStream>>>) -> zbus::Connection {
+    zbus::connection::Builder::address(address)
+        .unwrap()
+        .name("org.freedesktop.login1")
+        .unwrap()
+        .serve_at("/org/freedesktop/login1", LoginManager { peer })
+        .unwrap()
+        .build()
+        .await
+        .unwrap()
+}
+
+async fn wait_for_inhibitor(peer: &Mutex<Option<UnixStream>>) {
+    tokio::time::timeout(Duration::from_secs(10), async {
+        while peer.lock().unwrap().is_none() {
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+    })
+    .await
+    .unwrap();
+}
+
+fn api(socket: &Path, method: &str, params: serde_json::Value) -> serde_json::Value {
+    let mut stream = UnixStream::connect(socket).unwrap();
+    stream
+        .set_read_timeout(Some(Duration::from_secs(5)))
+        .unwrap();
+    writeln!(
+        stream,
+        "{}",
+        serde_json::json!({"id":"test", "method":method, "params":params})
+    )
+    .unwrap();
+    let mut line = String::new();
+    std::io::BufReader::new(stream)
+        .read_line(&mut line)
+        .unwrap();
+    let response: serde_json::Value = serde_json::from_str(&line).unwrap();
+    assert!(response.get("error").is_none(), "{response}");
+    response
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[ignore = "requires dbus-daemon; exercises a private bus and disposable named server, not host shutdown"]
+async fn host_shutdown_saves_layout_before_releasing_delay_lock() {
+    let base = std::path::PathBuf::from(format!(
+        "/var/tmp/hhs-{}-{}",
+        std::process::id(),
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos()
+    ));
+    std::fs::create_dir_all(&base).unwrap();
+    let address = format!("unix:path={}", base.join("bus").display());
+    let mut bus = private_bus(&address);
+    let peer = Arc::new(Mutex::new(None));
+    let mut service = login_service(&address, peer.clone()).await;
+    let socket = base.join("herdr-dev/sessions/shutdown/herdr.sock");
+    let config = base.join("config.toml");
+    std::fs::write(&config, "onboarding = false\n[experimental]\nallow_nested = true\n[terminal]\ndefault_shell = \"/bin/sh\"\n").unwrap();
+    let mut server = ChildGuard(
+        Command::new(env!("CARGO_BIN_EXE_herdr"))
+            .args(["--session", "shutdown", "server"])
+            .env("XDG_CONFIG_HOME", &base)
+            .env("XDG_STATE_HOME", &base)
+            .env("XDG_RUNTIME_DIR", &base)
+            .env("HERDR_CONFIG_PATH", &config)
+            .env_remove("HERDR_SOCKET_PATH")
+            .env("DBUS_SYSTEM_BUS_ADDRESS", address.trim())
+            .env_remove("HERDR_CLIENT_SOCKET_PATH")
+            .env_remove("HERDR_SESSION")
+            .env_remove("HERDR_WORKSPACE_ID")
+            .env_remove("HERDR_TAB_ID")
+            .env_remove("HERDR_PANE_ID")
+            .stdin(Stdio::null())
+            .stdout(Stdio::null())
+            .spawn()
+            .unwrap(),
+    );
+    tokio::time::timeout(Duration::from_secs(10), async {
+        while !socket.exists() || peer.lock().unwrap().is_none() {
+            assert!(
+                server.0.try_wait().unwrap().is_none(),
+                "server exited during startup"
+            );
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+    })
+    .await
+    .unwrap();
+    // Neither a login1 owner change nor a bus restart may leave a stale inhibitor.
+    peer.lock().unwrap().take();
+    service
+        .release_name("org.freedesktop.login1")
+        .await
+        .unwrap();
+    drop(service);
+    service = login_service(&address, peer.clone()).await;
+    wait_for_inhibitor(&peer).await;
+    peer.lock().unwrap().take();
+    drop(service);
+    drop(bus);
+    if base.join("bus").exists() {
+        std::fs::remove_file(base.join("bus")).unwrap();
+    }
+    bus = private_bus(&address);
+    service = login_service(&address, peer.clone()).await;
+    wait_for_inhibitor(&peer).await;
+
+    for label in ["one", "two", "three"] {
+        api(
+            &socket,
+            "workspace.create",
+            serde_json::json!({"cwd":base,"label":label,"focus":true}),
+        );
+    }
+    service
+        .emit_signal(
+            None::<&str>,
+            "/org/freedesktop/login1",
+            "org.freedesktop.login1.Manager",
+            "PrepareForShutdown",
+            &true,
+        )
+        .await
+        .unwrap();
+    let mut peer = peer.lock().unwrap().take().unwrap();
+    tokio::time::timeout(Duration::from_secs(10), async {
+        loop {
+            match peer.read(&mut [0]) {
+                Ok(0) => break,
+                Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {}
+                other => panic!("unexpected inhibitor state: {other:?}"),
+            }
+            tokio::time::sleep(Duration::from_millis(5)).await;
+        }
+    })
+    .await
+    .unwrap();
+    let saved = base.join("herdr-dev/sessions/shutdown/session.json");
+    let layout: serde_json::Value = serde_json::from_slice(&std::fs::read(saved).unwrap()).unwrap();
+    assert_eq!(layout["workspaces"].as_array().unwrap().len(), 3);
+    assert_eq!(layout["workspaces"][2]["custom_name"], "three");
+    tokio::time::timeout(Duration::from_secs(5), async {
+        while server.0.try_wait().unwrap().is_none() {
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+    })
+    .await
+    .unwrap();
+    drop(server);
+    drop(service);
+    drop(bus);
+    std::fs::remove_dir_all(base).unwrap();
+}