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>

1 files changed, 157 insertions, 16 deletions

diff --git a/src/persist/snapshot.rs b/src/persist/snapshot.rs
index 07aaf3b9..d7db3e97 100644
--- a/src/persist/snapshot.rs
+++ b/src/persist/snapshot.rs
@@ -33,6 +33,8 @@ pub struct SessionHistorySnapshot {
     /// Format version follows the matching session snapshot version.
     #[serde(default)]
     pub version: u32,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub layout_fingerprint: Option<String>,
     pub workspaces: Vec<WorkspaceHistorySnapshot>,
 }
 
@@ -281,12 +283,20 @@ fn capture_workspace(
     >,
     terminal_runtimes: &TerminalRuntimeRegistry,
 ) -> WorkspaceSnapshot {
+    let tabs: Vec<_> = ws
+        .tabs
+        .iter()
+        .map(|tab| capture_tab(tab, terminals, terminal_runtimes))
+        .collect();
+    let identity_cwd = tabs
+        .first()
+        .and_then(|tab| tab.root_pane.and_then(|id| tab.panes.get(&id)))
+        .map(|pane| pane.cwd.clone())
+        .unwrap_or_else(|| ws.identity_cwd.clone());
     WorkspaceSnapshot {
         id: Some(ws.id.clone()),
         custom_name: ws.custom_name.clone(),
-        identity_cwd: ws
-            .resolved_identity_cwd_from(terminals, terminal_runtimes)
-            .unwrap_or_else(|| ws.identity_cwd.clone()),
+        identity_cwd,
         worktree_space: ws.worktree_space.clone(),
         public_pane_numbers: ws
             .public_pane_numbers
@@ -296,11 +306,7 @@ fn capture_workspace(
         next_public_pane_number: ws.next_public_pane_number,
         public_tab_numbers: ws.tabs.iter().map(|tab| tab.number).collect(),
         next_public_tab_number: ws.next_public_tab_number,
-        tabs: ws
-            .tabs
-            .iter()
-            .map(|tab| capture_tab(tab, terminals, terminal_runtimes))
-            .collect(),
+        tabs,
         active_tab: ws.active_tab,
     }
 }
@@ -315,13 +321,13 @@ fn capture_tab(
 ) -> TabSnapshot {
     let mut panes = HashMap::new();
     for id in tab.panes.keys() {
-        let cwd = tab
-            .cwd_for_pane(*id, terminals, terminal_runtimes)
+        let terminal_id = tab.terminal_id(*id);
+        let terminal = terminal_id.and_then(|id| terminals.get(id));
+        let cwd = terminal_id
+            .and_then(|id| terminal_runtimes.get(id))
+            .and_then(|runtime| runtime.cwd_for_persistence())
+            .or_else(|| terminal.map(|terminal| terminal.cwd.clone()))
             .unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| "/".into()));
-        let terminal = tab
-            .panes
-            .get(id)
-            .and_then(|pane| terminals.get(&pane.attached_terminal_id));
         let label = terminal.and_then(|terminal| terminal.manual_label.clone());
         let (agent_name, managed_agent_kind) = terminal
             .filter(|terminal| !terminal.managed_agent_launch_pending())
@@ -378,13 +384,27 @@ fn capture_tab(
     }
 }
 
+pub(super) fn layout_fingerprint(snapshot: &SessionSnapshot) -> Option<String> {
+    use sha2::{Digest, Sha256};
+
+    let mut value = serde_json::to_value(snapshot).ok()?;
+    // Sets serialize as arrays; normalize their order as well as JSON object keys.
+    let mut collapsed: Vec<_> = snapshot.collapsed_space_keys.iter().collect();
+    collapsed.sort_unstable();
+    value["collapsed_space_keys"] = serde_json::to_value(collapsed).ok()?;
+    let bytes = serde_json::to_vec(&value).ok()?;
+    Some(format!("{:x}", Sha256::digest(bytes)))
+}
+
 /// Capture pane screen history separately from the structural session snapshot.
 pub fn capture_history(
+    snapshot: &SessionSnapshot,
     workspaces: &[Workspace],
     terminal_runtimes: &TerminalRuntimeRegistry,
 ) -> SessionHistorySnapshot {
     SessionHistorySnapshot {
         version: SNAPSHOT_VERSION,
+        layout_fingerprint: layout_fingerprint(snapshot),
         workspaces: workspaces
             .iter()
             .map(|workspace| WorkspaceHistorySnapshot {
@@ -546,7 +566,8 @@ mod tests {
         state: &AppState,
         terminal_runtimes: &TerminalRuntimeRegistry,
     ) -> SessionHistorySnapshot {
-        capture_history(&state.workspaces, terminal_runtimes)
+        let snapshot = capture_from_state_with_runtimes(state, terminal_runtimes);
+        capture_history(&snapshot, &state.workspaces, terminal_runtimes)
     }
 
     fn root_split_ratio(tab: &TabSnapshot) -> Option<f32> {
@@ -594,6 +615,28 @@ mod tests {
     }
 
     #[test]
+    fn layout_fingerprint_survives_json_round_trip() {
+        let mut snapshot = parse_snapshot(include_str!(
+            "../../tests/fixtures/session/current-herdr-session.json"
+        ))
+        .unwrap();
+        snapshot.collapsed_space_keys = ["z", "a", "m"].map(String::from).into();
+        let expected = layout_fingerprint(&snapshot).unwrap();
+        for _ in 0..16 {
+            snapshot = parse_snapshot(&serde_json::to_string(&snapshot).unwrap()).unwrap();
+            assert_eq!(
+                layout_fingerprint(&snapshot).as_deref(),
+                Some(expected.as_str())
+            );
+        }
+        snapshot.workspaces.swap(0, 1);
+        assert_ne!(
+            layout_fingerprint(&snapshot).as_deref(),
+            Some(expected.as_str())
+        );
+    }
+
+    #[test]
     fn round_trip_empty_session() {
         let snap = SessionSnapshot {
             version: SNAPSHOT_VERSION,
@@ -1002,6 +1045,104 @@ mod tests {
         assert_eq!(workspace.next_public_tab_number, 3);
     }
 
+    #[cfg(unix)]
+    #[tokio::test]
+    async fn capture_prefers_live_shell_cwd_and_keeps_it_after_exit() {
+        let old = std::env::current_dir().unwrap();
+        let new = std::env::temp_dir().join(format!(
+            "herdr-persist-cwd-{}-{}",
+            std::process::id(),
+            std::time::SystemTime::now()
+                .duration_since(std::time::UNIX_EPOCH)
+                .unwrap()
+                .as_nanos()
+        ));
+        std::fs::create_dir(&new).unwrap();
+        let new = std::fs::canonicalize(new).unwrap();
+        let mut state = AppState::test_new();
+        state.workspaces = vec![Workspace::test_new("cwd-source")];
+        state.workspaces[0].identity_cwd = old.clone();
+        state.active = Some(0);
+        state.ensure_test_terminals();
+        let pane_id = state.workspaces[0].tabs[0].root_pane;
+        let terminal_id = state.workspaces[0].terminal_id(pane_id).unwrap().clone();
+        let (events, _rx) = tokio::sync::mpsc::channel(32);
+        let runtime = crate::terminal::TerminalRuntime::spawn(
+            pane_id,
+            24,
+            80,
+            old.clone(),
+            0,
+            Default::default(),
+            None,
+            crate::pane::PaneShellConfig::new("/bin/sh", crate::config::ShellModeConfig::NonLogin),
+            &crate::pane::PaneLaunchEnv::default(),
+            events,
+            std::sync::Arc::new(tokio::sync::Notify::new()),
+            std::sync::Arc::new(crate::render_signal::RenderSignal::new()),
+        )
+        .unwrap();
+        let pid = runtime.child_pid().unwrap();
+        runtime
+            .try_send_bytes(bytes::Bytes::from(format!(
+                "cd '{}'; printf '\\033]7;file://{}\\007'; exec sleep 30\n",
+                new.display(),
+                old.display()
+            )))
+            .unwrap();
+        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
+        while (crate::platform::process_cwd(pid).as_ref() != Some(&new)
+            || runtime.cwd().as_ref() != Some(&old))
+            && std::time::Instant::now() < deadline
+        {
+            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+        }
+        assert_eq!(crate::platform::process_cwd(pid), Some(new.clone()));
+        assert_eq!(
+            runtime.cwd(),
+            Some(old.clone()),
+            "existing reported-cwd accessor is unchanged"
+        );
+        let mut runtimes = TerminalRuntimeRegistry::new();
+        runtimes.insert(terminal_id, runtime);
+        let before = capture_from_state_with_runtimes(&state, &runtimes);
+        assert_eq!(
+            before.workspaces[0].tabs[0]
+                .panes
+                .values()
+                .next()
+                .unwrap()
+                .cwd,
+            new
+        );
+        assert_eq!(before.workspaces[0].identity_cwd, new);
+        assert_eq!(runtimes.values().next().unwrap().cwd(), Some(old.clone()));
+        crate::platform::signal_processes(&[pid], crate::platform::Signal::Kill);
+        let exit_deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
+        while crate::platform::process_cwd(pid).is_some()
+            && std::time::Instant::now() < exit_deadline
+        {
+            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+        }
+        assert!(crate::platform::process_cwd(pid).is_none());
+        let after = capture_from_state_with_runtimes(&state, &runtimes);
+        assert_eq!(
+            after.workspaces[0].tabs[0]
+                .panes
+                .values()
+                .next()
+                .unwrap()
+                .cwd,
+            new
+        );
+        assert_eq!(after.workspaces[0].identity_cwd, new);
+        assert_eq!(runtimes.values().next().unwrap().cwd(), Some(old));
+        for (_, runtime) in runtimes.drain() {
+            runtime.shutdown();
+        }
+        std::fs::remove_dir(new).unwrap();
+    }
+
     #[test]
     fn capture_contract_tracks_workspace_identity_and_pane_cwds() {
         let mut state = state_with_workspaces(&["one"]);
@@ -1199,7 +1340,7 @@ mod tests {
     }
 
     #[test]
-    fn restore_falls_back_to_home_when_cwd_missing() {
+    fn snapshot_parsing_preserves_missing_cwd() {
         let mut panes = HashMap::new();
         panes.insert(
             0,