fix: preserve live handoff terminal state

12 files changed, 129 insertions, 4 deletions

diff --git a/src/app/actions.rs b/src/app/actions.rs
index 3dd5562f..c2329531 100644
--- a/src/app/actions.rs
+++ b/src/app/actions.rs
@@ -2177,6 +2177,7 @@ impl AppState {
 
         let pane_terminal_id = self.terminal_id_for_pane(ws_idx, pane_id);
         let workspace_terminal_ids = self.terminal_ids_for_workspace(ws_idx);
+        self.pane_id_aliases.retain(|_, alias| *alias != pane_id);
         let should_close_workspace = {
             let ws = &mut self.workspaces[ws_idx];
             ws.remove_pane(pane_id)
diff --git a/src/app/agents.rs b/src/app/agents.rs
index b9534ce2..8868fb9c 100644
--- a/src/app/agents.rs
+++ b/src/app/agents.rs
@@ -332,6 +332,8 @@ impl App {
         self.state.terminals.insert(terminal.id.clone(), terminal);
         self.state.workspaces.push(ws);
         let ws_idx = self.state.workspaces.len() - 1;
+        self.state
+            .remove_alias_shadowed_by_new_pane(self.state.workspaces[ws_idx].tabs[0].root_pane);
         if focus || self.state.active.is_none() {
             self.state.switch_workspace(ws_idx);
             self.state.mode = Mode::Terminal;
@@ -380,6 +382,8 @@ impl App {
         self.terminal_runtimes
             .insert(result.1.terminal.id.clone(), result.1.runtime);
         self.state
+            .remove_alias_shadowed_by_new_pane(result.1.pane_id);
+        self.state
             .terminals
             .insert(result.1.terminal.id.clone(), result.1.terminal);
         if focus {
diff --git a/src/app/api.rs b/src/app/api.rs
index a69dccec..562bc449 100644
--- a/src/app/api.rs
+++ b/src/app/api.rs
@@ -12,6 +12,12 @@ mod worktrees;
 use super::{api_helpers::pane_agent_status, App, Mode, OverlayPaneState, ToastKind};
 use crate::events::AppEvent;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum RuntimeExitAction {
+    RespawnShell,
+    ClosePane,
+}
+
 impl App {
     pub(crate) fn handle_internal_event(&mut self, ev: AppEvent) {
         if let AppEvent::ClipboardWrite { content } = ev {
@@ -43,6 +49,17 @@ impl App {
             return;
         }
 
+        if let AppEvent::PaneDied { pane_id } = &ev {
+            if self.runtime_exit_action(*pane_id) == RuntimeExitAction::RespawnShell
+                && self.respawn_shell_for_launch_pane(*pane_id)
+            {
+                self.overlay_panes.remove(pane_id);
+                self.render_dirty.store(true, Ordering::Release);
+                self.render_notify.notify_one();
+                return;
+            }
+        }
+
         let overlay_state = if let AppEvent::PaneDied { pane_id } = &ev {
             self.overlay_panes.remove(pane_id)
         } else {
@@ -210,6 +227,71 @@ impl App {
         }
     }
 
+    fn runtime_exit_action(&self, pane_id: crate::layout::PaneId) -> RuntimeExitAction {
+        let Some((_, pane_state)) = self.find_pane(pane_id) else {
+            return RuntimeExitAction::ClosePane;
+        };
+        let Some(terminal) = self.state.terminals.get(&pane_state.attached_terminal_id) else {
+            return RuntimeExitAction::ClosePane;
+        };
+
+        if terminal.respawn_shell_on_exit {
+            RuntimeExitAction::RespawnShell
+        } else {
+            RuntimeExitAction::ClosePane
+        }
+    }
+
+    fn respawn_shell_for_launch_pane(&mut self, pane_id: crate::layout::PaneId) -> bool {
+        let Some((ws_idx, pane_state)) = self.find_pane(pane_id) else {
+            return false;
+        };
+        let terminal_id = pane_state.attached_terminal_id.clone();
+        let Some(terminal) = self.state.terminals.get(&terminal_id) else {
+            return false;
+        };
+
+        let cwd = terminal.cwd.clone();
+        let (rows, cols) = self
+            .terminal_runtimes
+            .get(&terminal_id)
+            .map(|runtime| runtime.current_size())
+            .unwrap_or_else(|| self.state.estimate_pane_size());
+        let runtime = match crate::terminal::TerminalRuntime::spawn(
+            pane_id,
+            rows,
+            cols,
+            cwd,
+            self.state.pane_scrollback_limit_bytes,
+            self.state.host_terminal_theme,
+            &self.state.default_shell,
+            self.event_tx.clone(),
+            self.render_notify.clone(),
+            self.render_dirty.clone(),
+        ) {
+            Ok(runtime) => runtime,
+            Err(err) => {
+                tracing::warn!(
+                    pane = pane_id.raw(),
+                    terminal = %terminal_id,
+                    err = %err,
+                    "failed to respawn shell after launch command exited"
+                );
+                return false;
+            }
+        };
+
+        self.terminal_runtimes.insert(terminal_id.clone(), runtime);
+        if let Some(terminal) = self.state.terminals.get_mut(&terminal_id) {
+            terminal.launch_argv = None;
+            terminal.respawn_shell_on_exit = false;
+            terminal.clear_agent_name();
+        }
+        self.state.focus_pane_in_workspace(ws_idx, pane_id);
+        self.schedule_session_save();
+        true
+    }
+
     pub(crate) fn emit_pane_state_update(&self, update: &crate::app::actions::PaneStateUpdate) {
         let Some(pane_id) = self.public_pane_id(update.ws_idx, update.pane_id) else {
             return;
diff --git a/src/app/api/panes.rs b/src/app/api/panes.rs
index a81a6080..4620329f 100644
--- a/src/app/api/panes.rs
+++ b/src/app/api/panes.rs
@@ -66,6 +66,8 @@ impl App {
         self.terminal_runtimes
             .insert(new_pane.terminal.id.clone(), new_pane.runtime);
         self.state
+            .remove_alias_shadowed_by_new_pane(new_pane.pane_id);
+        self.state
             .terminals
             .insert(new_pane.terminal.id.clone(), new_pane.terminal);
         self.schedule_session_save();
diff --git a/src/app/api/tabs.rs b/src/app/api/tabs.rs
index 0344fcea..f605d143 100644
--- a/src/app/api/tabs.rs
+++ b/src/app/api/tabs.rs
@@ -93,6 +93,9 @@ impl App {
             Ok((tab_idx, terminal, runtime)) => {
                 self.terminal_runtimes.insert(terminal.id.clone(), runtime);
                 self.state.terminals.insert(terminal.id.clone(), terminal);
+                self.state.remove_alias_shadowed_by_new_pane(
+                    self.state.workspaces[ws_idx].tabs[tab_idx].root_pane,
+                );
                 if let Some(label) = label {
                     let workspace_id = self.state.workspaces[ws_idx].id.clone();
                     let tab_id = self
diff --git a/src/app/creation.rs b/src/app/creation.rs
index a034fdc4..aa727dd5 100644
--- a/src/app/creation.rs
+++ b/src/app/creation.rs
@@ -112,8 +112,10 @@ impl App {
             self.state.host_terminal_theme,
             &self.state.default_shell,
         )?;
+        let root_pane = ws.tabs[idx].root_pane;
         self.terminal_runtimes.insert(terminal.id.clone(), runtime);
         self.state.terminals.insert(terminal.id.clone(), terminal);
+        self.state.remove_alias_shadowed_by_new_pane(root_pane);
         if focus {
             self.state.switch_workspace_tab(ws_idx, idx);
             self.state.mode = Mode::Terminal;
@@ -149,6 +151,8 @@ impl App {
         self.state.terminals.insert(terminal.id.clone(), terminal);
         self.state.workspaces.push(ws);
         let idx = self.state.workspaces.len() - 1;
+        self.state
+            .remove_alias_shadowed_by_new_pane(self.state.workspaces[idx].tabs[0].root_pane);
         let workspace_id = self.state.workspaces[idx].id.clone();
         let root_pane = self.state.workspaces[idx].tabs[0].root_pane.raw();
         crate::logging::workspace_created(&workspace_id, root_pane);
diff --git a/src/app/ids.rs b/src/app/ids.rs
index 490e6051..a18efd2d 100644
--- a/src/app/ids.rs
+++ b/src/app/ids.rs
@@ -1,7 +1,7 @@
 use super::App;
 
 impl App {
-    pub(super) fn find_pane(
+    pub(crate) fn find_pane(
         &self,
         pane_id: crate::layout::PaneId,
     ) -> Option<(usize, &crate::pane::PaneState)> {
@@ -57,15 +57,27 @@ impl App {
         Some((ws_idx, tab_idx))
     }
 
+    fn resolve_raw_pane_id(&self, raw: u32) -> Option<crate::layout::PaneId> {
+        if let Some(alias) = self.state.pane_id_aliases.get(&raw).copied() {
+            return self.find_pane(alias).map(|_| alias);
+        }
+        let pane_id = crate::layout::PaneId::from_raw(raw);
+        if self.find_pane(pane_id).is_some() {
+            return Some(pane_id);
+        }
+        None
+    }
+
     pub(super) fn parse_pane_id(&self, id: &str) -> Option<(usize, crate::layout::PaneId)> {
         if let Some(rest) = id.strip_prefix("p_") {
             if let Some((ws_raw, pane_raw)) = rest.rsplit_once('_') {
                 let ws_idx = self.parse_workspace_id(ws_raw)?;
-                let pane_id = crate::layout::PaneId::from_raw(pane_raw.parse::<u32>().ok()?);
+                let pane_id = self.resolve_raw_pane_id(pane_raw.parse::<u32>().ok()?)?;
+                self.state.workspaces.get(ws_idx)?.pane_state(pane_id)?;
                 return Some((ws_idx, pane_id));
             }
 
-            let pane_id = crate::layout::PaneId::from_raw(rest.parse::<u32>().ok()?);
+            let pane_id = self.resolve_raw_pane_id(rest.parse::<u32>().ok()?)?;
             return self.find_pane(pane_id).map(|(ws_idx, _)| (ws_idx, pane_id));
         }
 
diff --git a/src/app/input/mod.rs b/src/app/input/mod.rs
index 8ba8e4bd..d9092d40 100644
--- a/src/app/input/mod.rs
+++ b/src/app/input/mod.rs
@@ -419,6 +419,7 @@ impl AppState {
             ) {
                 let new_id = new_pane.pane_id;
                 terminal_runtimes.insert(new_pane.terminal.id.clone(), new_pane.runtime);
+                self.remove_alias_shadowed_by_new_pane(new_id);
                 self.terminals
                     .insert(new_pane.terminal.id.clone(), new_pane.terminal);
                 self.record_pane_focus_change(previous_focus, ws_idx, new_id);
diff --git a/src/app/input/mouse.rs b/src/app/input/mouse.rs
index 6db1a4a9..a9badfa9 100644
--- a/src/app/input/mouse.rs
+++ b/src/app/input/mouse.rs
@@ -2302,6 +2302,7 @@ mod tests {
             mouse_protocol_mode: crate::input::MouseProtocolMode::ButtonMotion,
             mouse_protocol_encoding: crate::input::MouseProtocolEncoding::Sgr,
             mouse_alternate_scroll: true,
+            modify_other_keys: false,
         };
 
         assert_eq!(wheel_routing(input_state), WheelRouting::MouseReport);
@@ -2626,6 +2627,7 @@ mod tests {
             mouse_protocol_mode: crate::input::MouseProtocolMode::None,
             mouse_protocol_encoding: crate::input::MouseProtocolEncoding::Default,
             mouse_alternate_scroll: true,
+            modify_other_keys: false,
         };
 
         assert_eq!(wheel_routing(input_state), WheelRouting::AlternateScroll);
@@ -2641,6 +2643,7 @@ mod tests {
             mouse_protocol_mode: crate::input::MouseProtocolMode::None,
             mouse_protocol_encoding: crate::input::MouseProtocolEncoding::Default,
             mouse_alternate_scroll: true,
+            modify_other_keys: false,
         };
 
         assert_eq!(wheel_routing(input_state), WheelRouting::HostScroll);
diff --git a/src/app/input/navigate.rs b/src/app/input/navigate.rs
index 3b8435df..032ba8ae 100644
--- a/src/app/input/navigate.rs
+++ b/src/app/input/navigate.rs
@@ -348,6 +348,7 @@ impl App {
                 temp_files,
             },
         );
+        self.state.remove_alias_shadowed_by_new_pane(new_pane_id);
         self.state.mode = Mode::Terminal;
         Ok(())
     }
diff --git a/src/app/mod.rs b/src/app/mod.rs
index 12fb704e..fbbacd61 100644
--- a/src/app/mod.rs
+++ b/src/app/mod.rs
@@ -393,6 +393,7 @@ impl App {
         let mut state = AppState {
             terminals: std::collections::HashMap::new(),
             direct_attach_resize_locks: std::collections::HashSet::new(),
+            pane_id_aliases: std::collections::HashMap::new(),
             workspaces,
             active,
             previous_pane_focus: None,
@@ -594,7 +595,10 @@ impl App {
         api_rx: tokio::sync::mpsc::UnboundedReceiver<crate::api::ApiRequestMessage>,
         event_hub: crate::api::EventHub,
         snapshot: &crate::persist::SessionSnapshot,
-        imports: &mut std::collections::HashMap<u32, crate::persist::ImportedPaneRuntime>,
+        imports: &mut std::collections::HashMap<
+            u32,
+            crate::handoff_runtime::ImportedHandoffRuntime,
+        >,
     ) -> io::Result<Self> {
         let mut app = Self::new(config, true, config_diagnostic, api_rx, event_hub);
         let (workspaces, terminals, runtimes) = crate::persist::restore_handoff(
@@ -606,9 +610,11 @@ impl App {
             app.render_notify.clone(),
             app.render_dirty.clone(),
         )?;
+        let pane_id_aliases = crate::persist::handoff_pane_aliases(snapshot, &workspaces);
 
         app.no_session = false;
         app.state.detach_exits = false;
+        app.state.pane_id_aliases = pane_id_aliases;
         app.state.workspaces = workspaces;
         app.state.terminals = terminals;
         app.terminal_runtimes = runtimes.into();
diff --git a/src/app/state.rs b/src/app/state.rs
index 7c47973c..37edc9df 100644
--- a/src/app/state.rs
+++ b/src/app/state.rs
@@ -1046,6 +1046,7 @@ pub struct AppState {
         std::collections::HashMap<crate::terminal::TerminalId, crate::terminal::TerminalState>,
     /// Terminal ids whose size is currently owned by a direct attach client.
     pub direct_attach_resize_locks: std::collections::HashSet<crate::terminal::TerminalId>,
+    pub(crate) pane_id_aliases: std::collections::HashMap<u32, PaneId>,
     pub workspaces: Vec<Workspace>,
     pub active: Option<usize>,
     pub(crate) previous_pane_focus: Option<PaneFocusTarget>,
@@ -1182,6 +1183,10 @@ impl AppState {
         self.session_dirty = true;
     }
 
+    pub(crate) fn remove_alias_shadowed_by_new_pane(&mut self, pane_id: PaneId) {
+        self.pane_id_aliases.remove(&pane_id.raw());
+    }
+
     pub fn sound_enabled(&self) -> bool {
         self.sound.enabled
     }
@@ -1349,6 +1354,7 @@ impl AppState {
         Self {
             terminals: std::collections::HashMap::new(),
             direct_attach_resize_locks: std::collections::HashSet::new(),
+            pane_id_aliases: std::collections::HashMap::new(),
             workspaces: Vec::new(),
             active: None,
             previous_pane_focus: None,