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, 188 insertions, 17 deletions

diff --git a/src/persist/restore.rs b/src/persist/restore.rs
index 3c5775c9..31ee0a37 100644
--- a/src/persist/restore.rs
+++ b/src/persist/restore.rs
@@ -267,6 +267,14 @@ fn restore_with_imports_and_failures(
     render_notify: Arc<Notify>,
     render_dirty: Arc<RenderSignal>,
 ) -> RestoreFailures<RestoredSession> {
+    let history = history.filter(|history| {
+        let matches = history.layout_fingerprint.is_some()
+            && history.layout_fingerprint == super::snapshot::layout_fingerprint(snapshot);
+        if !matches {
+            tracing::warn!("Ignoring pane history without a matching session layout");
+        }
+        matches
+    });
     let mut workspaces = Vec::new();
     let mut terminals = HashMap::new();
     let mut terminal_runtimes = HashMap::new();
@@ -433,6 +441,34 @@ fn restore_workspace(
     )
 }
 
+fn unavailable_restored_terminal(
+    pane: Option<&super::snapshot::PaneSnapshot>,
+    cwd: PathBuf,
+    reason: String,
+) -> TerminalState {
+    warn!(cwd = %cwd.display(), reason = %reason, "preserving unavailable restored pane");
+    let mut terminal = TerminalState::new(TerminalId::alloc(), cwd);
+    terminal.restore_error = Some(reason);
+    if let Some(pane) = pane {
+        terminal.manual_label = pane.label.clone();
+        terminal.launch_argv = pane.launch_argv.clone();
+        if let Some(session) = restored_terminal_agent_session(pane.agent_session.as_ref(), false) {
+            terminal.set_persisted_agent_session(session);
+        }
+        match (
+            pane.agent_name.as_ref(),
+            pane.managed_agent_kind
+                .as_deref()
+                .and_then(crate::detect::parse_canonical_agent_label),
+        ) {
+            (Some(name), Some(agent)) => terminal.restore_managed_agent(name.clone(), agent),
+            (Some(name), None) => terminal.set_agent_name(name.clone()),
+            _ => {}
+        }
+    }
+    terminal
+}
+
 fn restored_worktree_space_membership(
     space: Option<crate::workspace::WorktreeSpaceMembership>,
 ) -> Option<crate::workspace::WorktreeSpaceMembership> {
@@ -473,22 +509,19 @@ fn restore_tab(
             .map(|p| p.cwd.clone())
             .unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| "/".into()));
 
-        let cwd = if saved_cwd.exists() {
-            saved_cwd
-        } else {
-            warn!(
-                cwd = %saved_cwd.display(),
-                "saved pane cwd does not exist, falling back to HOME"
+        let cwd = saved_cwd;
+        let has_import = old_id.is_some_and(|old_id| imported_panes.contains_key(old_id));
+        if !has_import && !cwd.is_dir() {
+            let terminal = unavailable_restored_terminal(
+                saved_pane,
+                cwd,
+                "Saved directory is unavailable. Restore the directory and restart this session."
+                    .into(),
             );
-            let home = std::env::var("HOME")
-                .map(PathBuf::from)
-                .unwrap_or_else(|_| PathBuf::from("/"));
-            if home.exists() {
-                home
-            } else {
-                PathBuf::from("/")
-            }
-        };
+            panes.insert(*id, PaneState::new(terminal.id.clone()));
+            terminals.push(terminal);
+            continue;
+        }
 
         let saved_label = saved_pane.and_then(|p| p.label.clone());
         let saved_agent_name = saved_pane.and_then(|p| p.agent_name.clone());
@@ -681,8 +714,16 @@ fn restore_tab(
                     tab = ?snap.custom_name,
                     pane_id = id.raw(),
                     err = %e,
-                    "failed to restore pane, skipping"
+                    "failed to restore pane"
                 );
+                if !was_imported {
+                    let terminal = unavailable_restored_terminal(
+                        saved_pane, cwd,
+                        format!("Could not start the saved shell: {e}. Fix the shell configuration and restart this session."),
+                    );
+                    panes.insert(*id, PaneState::new(terminal.id.clone()));
+                    terminals.push(terminal);
+                }
             }
         }
     }
@@ -1168,6 +1209,92 @@ mod tests {
     }
 
     #[tokio::test]
+    async fn failed_cold_restore_preserves_panes_and_saved_directories() {
+        for missing_shell in [false, true] {
+            let mut snapshot: SessionSnapshot = serde_json::from_str(include_str!(
+                "../../tests/fixtures/session/current-herdr-session.json"
+            ))
+            .unwrap();
+            let cwd = std::env::current_dir().unwrap();
+            let missing = cwd.join("__herdr_missing_restore_directory__");
+            assert!(!missing.exists());
+            for workspace in &mut snapshot.workspaces {
+                workspace.identity_cwd = cwd.clone();
+                for tab in &mut workspace.tabs {
+                    for pane in tab.panes.values_mut() {
+                        pane.cwd = cwd.clone();
+                    }
+                }
+            }
+            let failed = snapshot.workspaces[0].tabs[0].panes.get_mut(&1).unwrap();
+            failed.cwd = missing.clone();
+            failed.label = Some("keep my pane".into());
+            failed.agent_session = Some(super::super::snapshot::PaneAgentSessionSnapshot {
+                source: "herdr:opencode".into(),
+                agent: "opencode".into(),
+                kind: crate::agent_resume::AgentSessionRefKind::Id,
+                value: "keep-my-session".into(),
+            });
+            let (events, _rx) = mpsc::channel(32);
+            let (workspaces, terminals, runtimes) = restore(
+                &snapshot,
+                None,
+                24,
+                80,
+                0,
+                if missing_shell {
+                    "__herdr_missing_restore_shell__"
+                } else {
+                    test_restore_shell()
+                },
+                crate::config::ShellModeConfig::NonLogin,
+                false,
+                events,
+                Arc::new(Notify::new()),
+                Arc::new(RenderSignal::new()),
+            );
+            let runtimes = crate::terminal::TerminalRuntimeRegistry::from(runtimes);
+            let captured = crate::persist::capture(&workspaces, &terminals, &runtimes, Some(0), 0);
+            assert_eq!(
+                captured.workspaces.len(),
+                2,
+                "a launch failure must not delete a workspace"
+            );
+            assert_eq!(captured.workspaces[0].tabs.len(), 2);
+            let pane = captured.workspaces[0].tabs[0]
+                .panes
+                .values()
+                .next()
+                .unwrap();
+            assert_eq!(
+                pane.cwd, missing,
+                "fallback cwd must not replace saved intent"
+            );
+            assert_eq!(pane.label.as_deref(), Some("keep my pane"));
+            assert_eq!(
+                pane.agent_session.as_ref().unwrap().value,
+                "keep-my-session"
+            );
+            let root = workspaces[0].tabs[0].root_pane;
+            let terminal_id = workspaces[0].tabs[0].terminal_id(root).unwrap();
+            assert!(
+                runtimes.get(terminal_id).is_none(),
+                "do not open a replacement shell elsewhere"
+            );
+            let healthy = workspaces[1].tabs[0]
+                .terminal_id(workspaces[1].tabs[0].root_pane)
+                .unwrap();
+            assert_eq!(runtimes.get(healthy).is_some(), !missing_shell);
+            assert!(terminals[terminal_id].restore_error.is_some());
+            let mut state = crate::app::AppState::test_new();
+            state.workspaces = workspaces;
+            state.terminals = terminals;
+            state.active = Some(0);
+            state.assert_invariants_for_test();
+        }
+    }
+
+    #[tokio::test]
     async fn restore_carries_persisted_agent_session_metadata() {
         let cwd = std::env::current_dir().unwrap();
         let snapshot = SessionSnapshot {
@@ -1658,6 +1785,48 @@ mod tests {
         let _ = runtime.try_send_bytes(bytes::Bytes::from_static(b"exit\n"));
     }
 
+    #[tokio::test]
+    async fn restore_rejects_history_from_another_layout_or_without_provenance() {
+        for legacy in [false, true] {
+            let (mut snapshot, history) = snapshot_with_saved_pane_history();
+            let mut value = serde_json::to_value(history).unwrap();
+            if legacy {
+                value.as_object_mut().unwrap().remove("layout_fingerprint");
+            } else {
+                snapshot.workspaces[0].tabs[0]
+                    .panes
+                    .get_mut(&0)
+                    .unwrap()
+                    .cwd = std::env::temp_dir();
+            }
+            let history = serde_json::from_value(value).unwrap();
+            let (events, _rx) = mpsc::channel(8);
+            let (_, _, runtimes) = restore(
+                &snapshot,
+                Some(&history),
+                5,
+                80,
+                4096,
+                test_restore_shell(),
+                crate::config::ShellModeConfig::NonLogin,
+                false,
+                events,
+                Arc::new(Notify::new()),
+                Arc::new(RenderSignal::new()),
+            );
+            let runtime = runtimes.values().next().unwrap();
+            assert!(
+                !runtime
+                    .recent_unwrapped_text(10)
+                    .contains("RESTORED_HISTORY"),
+                "screen history must belong to the exact saved layout"
+            );
+            for (_, runtime) in runtimes {
+                runtime.shutdown();
+            }
+        }
+    }
+
     fn snapshot_with_saved_pane_history() -> (SessionSnapshot, SessionHistorySnapshot) {
         let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/"));
         let mut panes = HashMap::new();
@@ -1672,8 +1841,9 @@ mod tests {
                 launch_argv: None,
             },
         );
-        let history = SessionHistorySnapshot {
+        let mut history = SessionHistorySnapshot {
             version: super::super::snapshot::SNAPSHOT_VERSION,
+            layout_fingerprint: None,
             workspaces: vec![WorkspaceHistorySnapshot {
                 tabs: vec![super::super::snapshot::TabHistorySnapshot {
                     panes: HashMap::from([(
@@ -1717,6 +1887,7 @@ mod tests {
             sidebar_section_split: Some(0.5),
             collapsed_space_keys: Default::default(),
         };
+        history.layout_fingerprint = super::super::snapshot::layout_fingerprint(&snapshot);
         (snapshot, history)
     }
 }