fix: restore saved machine screens after reconnect (#4126)

refs #3842

7 files changed, 286 insertions, 28 deletions

diff --git a/src/client/shell/endpoints.rs b/src/client/shell/endpoints.rs
index d3823d54..c27358e5 100644
--- a/src/client/shell/endpoints.rs
+++ b/src/client/shell/endpoints.rs
@@ -208,6 +208,7 @@ impl ClientShellState {
         let Some(snapshot) = endpoint.snapshot.clone() else {
             return false;
         };
+        let generation = endpoint.snapshot_generation;
         let switching_endpoint = endpoint_id != &self.active_endpoint_id;
         let agent_scroll = self.agent_scroll;
         if switching_endpoint {
@@ -215,7 +216,7 @@ impl ClientShellState {
             self.pane_surface = None;
             self.pending_pane_surface = None;
         }
-        self.apply_active_snapshot(snapshot);
+        self.apply_active_snapshot(snapshot, generation);
         if switching_endpoint {
             // The aggregate agent list belongs to the client, not one endpoint.
             self.agent_scroll = agent_scroll;
@@ -653,16 +654,21 @@ impl ClientShellState {
     }
 
     fn apply_cached_endpoint_snapshot(&mut self, endpoint_id: &ClientEndpointId) {
-        let Some(snapshot) = self
+        let Some((snapshot, generation)) = self
             .endpoints
             .iter()
             .find(|endpoint| &endpoint.endpoint_id == endpoint_id)
-            .and_then(|endpoint| endpoint.snapshot.clone())
+            .and_then(|endpoint| {
+                endpoint
+                    .snapshot
+                    .clone()
+                    .map(|snapshot| (snapshot, endpoint.snapshot_generation))
+            })
         else {
             return;
         };
         if endpoint_id == &self.active_endpoint_id {
-            self.apply_active_snapshot(snapshot);
+            self.apply_active_snapshot(snapshot, generation);
         }
     }
 }
diff --git a/src/client/shell/state.rs b/src/client/shell/state.rs
index 8f2aec43..04e87169 100644
--- a/src/client/shell/state.rs
+++ b/src/client/shell/state.rs
@@ -895,6 +895,8 @@ pub(super) struct ClientCopyModeState {
 pub(crate) struct ClientShellState {
     pub(super) config: ClientShellConfig,
     pub(super) snapshot: Option<Box<ClientShellSnapshot>>,
+    pub(super) active_snapshot_generation: Option<u64>,
+    pub(super) pane_surface_generation: Option<u64>,
     pub(super) pane_surface: Option<PaneSurfaceFrame>,
     /// A future projection surface waits here until its matching snapshot arrives. The visible
     /// pane surface always remains an exact snapshot pair.
@@ -1054,6 +1056,8 @@ impl ClientShellState {
         Self {
             config,
             snapshot: None,
+            active_snapshot_generation: None,
+            pane_surface_generation: None,
             pane_surface: None,
             pending_pane_surface: None,
             graphics: crate::kitty_graphics::surface::ClientState::default(),
@@ -1309,7 +1313,11 @@ impl ClientShellState {
         self.dismissed_product_announcement = None;
     }
 
-    pub(super) fn apply_active_snapshot(&mut self, mut snapshot: Box<ClientShellSnapshot>) {
+    pub(super) fn apply_active_snapshot(
+        &mut self,
+        mut snapshot: Box<ClientShellSnapshot>,
+        generation: Option<u64>,
+    ) {
         snapshot
             .commands
             .retain(|command| command.action != crate::protocol::ClientShellCommandAction::Unknown);
@@ -1320,13 +1328,21 @@ impl ClientShellState {
         };
         let endpoint_boot_changed =
             self.snapshot.is_some() && self.graphics.scope() != graphics_scope;
+        let generation_changed = self.active_snapshot_generation != generation;
         if !endpoint_boot_changed
+            && !generation_changed
             && self.snapshot.as_ref().is_some_and(|current| {
                 current.boot_id == snapshot.boot_id && snapshot.revision < current.revision
             })
         {
             return;
         }
+        // Screen revisions restart per connection. Keep the displayed surface for selection
+        // content comparisons, but retire speculative frames from the old connection.
+        if generation_changed {
+            self.pending_pane_surface = None;
+        }
+        self.active_snapshot_generation = generation;
         self.graphics.set_scope(&graphics_scope);
         let command_bindings_changed = self.snapshot.as_ref().is_none_or(|current| {
             current.commands.len() != snapshot.commands.len()
@@ -1618,7 +1634,8 @@ impl ClientShellState {
             return;
         }
         if self.pane_surface.as_ref().is_some_and(|current| {
-            current.boot_id == surface.boot_id
+            self.pane_surface_generation == self.active_snapshot_generation
+                && current.boot_id == surface.boot_id
                 && (surface.projection_revision < current.projection_revision
                     || (surface.projection_revision == current.projection_revision
                         && surface.surface_revision < current.surface_revision))
@@ -1646,7 +1663,8 @@ impl ClientShellState {
             || surface.projection_revision < snapshot.revision
             || (!retain_future && surface.projection_revision != snapshot.revision)
             || self.pane_surface.as_ref().is_some_and(|current| {
-                current.boot_id == surface.boot_id
+                self.pane_surface_generation == self.active_snapshot_generation
+                    && current.boot_id == surface.boot_id
                     && (surface.projection_revision < current.projection_revision
                         || (surface.projection_revision == current.projection_revision
                             && surface.surface_revision < current.surface_revision))
@@ -1817,6 +1835,7 @@ impl ClientShellState {
         self.graphics
             .set_scene(std::mem::take(&mut surface.graphics));
         self.pane_surface = Some(surface);
+        self.pane_surface_generation = self.active_snapshot_generation;
         self.invalidate_link_hover();
         self.resume_mobile_switcher_if_ready();
         self.reconcile_input_source();
diff --git a/src/client/shell/surface_patch.rs b/src/client/shell/surface_patch.rs
index 87f947fa..6c829acf 100644
--- a/src/client/shell/surface_patch.rs
+++ b/src/client/shell/surface_patch.rs
@@ -113,7 +113,8 @@ impl ClientShellState {
         let Some(current) = self.pane_surface.as_ref() else {
             return ClientPaneSurfacePatchOutcome::Rejected;
         };
-        if patch.boot_id != current.boot_id
+        if self.pane_surface_generation != self.active_snapshot_generation
+            || patch.boot_id != current.boot_id
             || patch.projection_revision != current.projection_revision
             || patch.base_surface_revision != current.surface_revision
             || patch.surface_revision != current.surface_revision.saturating_add(1)
diff --git a/src/client/shell/tests/endpoints.rs b/src/client/shell/tests/endpoints.rs
index 1edba8e6..4ac8183c 100644
--- a/src/client/shell/tests/endpoints.rs
+++ b/src/client/shell/tests/endpoints.rs
@@ -1527,6 +1527,53 @@ fn new_connection_generation_accepts_a_lower_same_boot_projection_revision() {
 }
 
 #[test]
+fn reconnect_same_endpoint_accepts_new_generation_surface_revision() {
+    for previous_revision in [9, 1] {
+        let (mut state, endpoint_id) = state_with_remote();
+        let mut previous = snapshot();
+        previous.boot_id = "shared-server-boot".into();
+        previous.revision = previous_revision;
+        state.cache_endpoint_snapshot_inactive_for_generation(&endpoint_id, 4, Box::new(previous));
+        assert!(state.activate_endpoint_projection(&endpoint_id));
+        let mut previous_surface = surface();
+        previous_surface.boot_id = "shared-server-boot".into();
+        previous_surface.projection_revision = previous_revision;
+        previous_surface.surface_revision = 9;
+        state.set_pane_surface(previous_surface.clone());
+        previous_surface.projection_revision += 1;
+        state.set_pane_surface(previous_surface);
+        assert!(state.pending_pane_surface.is_some());
+        state.agent_scroll = 7;
+
+        state.mark_endpoint_disconnected(&endpoint_id);
+        let mut reconnected = snapshot();
+        reconnected.boot_id = "shared-server-boot".into();
+        reconnected.revision = 1;
+        state.cache_endpoint_snapshot_inactive_for_generation(
+            &endpoint_id,
+            5,
+            Box::new(reconnected),
+        );
+        assert_eq!(state.snapshot.as_ref().unwrap().revision, previous_revision);
+        assert_eq!(state.pane_surface.as_ref().unwrap().surface_revision, 9);
+
+        state.set_endpoint_status(&endpoint_id, ClientEndpointStatus::Online);
+        assert!(state.activate_endpoint_projection(&endpoint_id));
+        let mut reconnected_surface = surface();
+        reconnected_surface.boot_id = "shared-server-boot".into();
+        reconnected_surface.projection_revision = 1;
+        reconnected_surface.surface_revision = 1;
+        state.set_pane_surface(reconnected_surface);
+
+        assert_eq!(state.snapshot.as_ref().unwrap().revision, 1);
+        assert_eq!(state.pane_surface.as_ref().unwrap().projection_revision, 1);
+        assert_eq!(state.pane_surface.as_ref().unwrap().surface_revision, 1);
+        assert!(state.pending_pane_surface.is_none());
+        assert_eq!(state.agent_scroll, 7);
+    }
+}
+
+#[test]
 fn reconnect_snapshot_waits_for_coherent_activation_before_replacing_projection() {
     let (mut state, endpoint_id) = state_with_remote();
     assert!(state.activate_endpoint_projection(&endpoint_id));
diff --git a/src/client/shell/tests/mouse_selection.rs b/src/client/shell/tests/mouse_selection.rs
index 0a54203e..ccb9b86d 100644
--- a/src/client/shell/tests/mouse_selection.rs
+++ b/src/client/shell/tests/mouse_selection.rs
@@ -579,6 +579,33 @@ fn double_click_drag_invalidates_cached_boundaries_outside_selected_cells() {
 }
 
 #[test]
+fn reconnect_word_selection_tracks_content_changes() {
+    for content_changed in [false, true] {
+        let mut state = word_drag_state(true);
+        let initial = start_word_drag(&mut state);
+        word_row_reply(&mut state, &initial, "alpha bravo charlie");
+        let mut next_surface = state.pane_surface.as_ref().unwrap().clone();
+        if content_changed {
+            next_surface.panes[0].content_revision += 2;
+            next_surface.frame.cells[14].symbol = " ".into();
+        }
+        let endpoint_id = state.active_endpoint_id.clone();
+        let snapshot = state.snapshot.as_ref().unwrap().clone();
+        state.mark_endpoint_disconnected(&endpoint_id);
+        state.cache_endpoint_snapshot_inactive_for_generation(&endpoint_id, 1, snapshot);
+        state.set_endpoint_status(
+            &endpoint_id,
+            crate::client::endpoint::ClientEndpointStatus::Online,
+        );
+        assert!(state.activate_endpoint_projection(&endpoint_id));
+        state.set_pane_surface(next_surface);
+
+        assert_eq!(state.selection.is_some(), !content_changed);
+        assert_eq!(state.word_selection_gesture.is_some(), !content_changed);
+    }
+}
+
+#[test]
 fn double_click_release_ignores_reply_after_focus_or_content_changes() {
     for focus_changed in [false, true] {
         let mut state = word_drag_state(true);
diff --git a/tests/client_mode.rs b/tests/client_mode.rs
index d3205e5a..9a6bee7d 100644
--- a/tests/client_mode.rs
+++ b/tests/client_mode.rs
@@ -3,6 +3,8 @@
 #![cfg(unix)]
 
 pub mod support;
+#[path = "support/terminal_screen.rs"]
+mod terminal_screen;
 
 use std::fs;
 use std::io::{BufRead, BufReader, Read, Write};
@@ -656,20 +658,28 @@ fn output_has_mouse_teardown(output: &str) -> bool {
 /// keeps the blocking `Box<dyn Read>` (which has no timeout) off the test's
 /// main thread, so a client that never exits fails the deadline instead of
 /// hanging the whole test forever.
-type SharedOutput = std::sync::Arc<Mutex<String>>;
+#[derive(Default)]
+struct PtyOutput {
+    bytes: Vec<u8>,
+    // Keep legacy raw-string watermarks stable even across split UTF-8 reads.
+    text: String,
+}
+
+type SharedOutput = std::sync::Arc<Mutex<PtyOutput>>;
 
 fn spawn_pty_drain(mut reader: Box<dyn Read + Send>) -> SharedOutput {
-    let output: SharedOutput = std::sync::Arc::new(Mutex::new(String::new()));
+    let output: SharedOutput = std::sync::Arc::new(Mutex::new(PtyOutput::default()));
     let thread_output = output.clone();
     thread::spawn(move || {
         let mut buf = [0u8; 4096];
         loop {
             match reader.read(&mut buf) {
                 Ok(0) => break,
-                Ok(n) => thread_output
-                    .lock()
-                    .unwrap_or_else(|p| p.into_inner())
-                    .push_str(&String::from_utf8_lossy(&buf[..n])),
+                Ok(n) => {
+                    let mut captured = thread_output.lock().unwrap_or_else(|p| p.into_inner());
+                    captured.bytes.extend_from_slice(&buf[..n]);
+                    captured.text.push_str(&String::from_utf8_lossy(&buf[..n]));
+                }
                 Err(_) => break,
             }
         }
@@ -677,15 +687,33 @@ fn spawn_pty_drain(mut reader: Box<dyn Read + Send>) -> SharedOutput {
     output
 }
 
+#[test]
+fn screen_capture_preserves_split_utf8_bytes() {
+    let reader = std::io::Cursor::new(b"caf\xc3").chain(std::io::Cursor::new(b"\xa9"));
+    let output = spawn_pty_drain(Box::new(reader));
+    assert!(wait_until(
+        Duration::from_secs(2),
+        Duration::from_millis(10),
+        || {
+            let bytes = output.lock().unwrap().bytes.clone();
+            bytes == "café".as_bytes() && terminal_screen::text(&bytes, 80, 24).contains("café")
+        }
+    ));
+}
+
 fn read_output(output: &SharedOutput) -> String {
-    output.lock().unwrap_or_else(|p| p.into_inner()).clone()
+    output
+        .lock()
+        .unwrap_or_else(|p| p.into_inner())
+        .text
+        .clone()
 }
 
 /// Current captured byte length, used as a watermark so a test can search only
 /// the output emitted *after* a trigger. The teardown markers also appear in
 /// normal attach-phase output, so matching the whole buffer is meaningless.
 fn output_len(output: &SharedOutput) -> usize {
-    output.lock().unwrap_or_else(|p| p.into_inner()).len()
+    output.lock().unwrap_or_else(|p| p.into_inner()).text.len()
 }
 
 /// Spawns a server + real thin client under a PTY and waits until the client
@@ -890,9 +918,10 @@ fn federated_client_starts_without_local_and_survives_its_restart() {
     let quote =
         |path: &std::path::Path| format!("'{}'", path.display().to_string().replace('\'', "'\\''"));
     let ssh_commands = base.join("ssh-commands");
+    let bridge_pid = base.join("bridge-pid");
     fs::write(bin.join("ssh"), format!(
-        "#!/bin/sh\nexport HOME={} XDG_CONFIG_HOME={} XDG_RUNTIME_DIR={} HERDR_SOCKET_PATH={}\nunset HERDR_CLIENT_SOCKET_PATH HERDR_SESSION\nfor arg do last=\"$arg\"; done\nprintf '%s\\n' \"$last\" >> {}\nexec /bin/sh -c \"$last\"\n",
-        quote(&base.join("home")), quote(&remote_config), quote(&remote_runtime), quote(&remote_api), quote(&ssh_commands),
+        "#!/bin/sh\nexport HOME={} XDG_CONFIG_HOME={} XDG_RUNTIME_DIR={} HERDR_SOCKET_PATH={}\nunset HERDR_CLIENT_SOCKET_PATH HERDR_SESSION\nfor arg do last=\"$arg\"; done\nprintf '%s\\n' \"$last\" >> {}\ncase \"$last\" in *remote-client-bridge*) printf '%s\\n' \"$$\" > {};; esac\nexec /bin/sh -c \"$last\"\n",
+        quote(&base.join("home")), quote(&remote_config), quote(&remote_runtime), quote(&remote_api), quote(&ssh_commands), quote(&bridge_pid),
     )).unwrap();
     fs::set_permissions(bin.join("ssh"), fs::Permissions::from_mode(0o700)).unwrap();
     let path = format!(
@@ -908,9 +937,17 @@ fn federated_client_starts_without_local_and_survives_its_restart() {
         &[("PATH", &path)],
     );
     let output = spawn_pty_drain(client._master.as_ref().unwrap().try_clone_reader().unwrap());
+    let screen_text = || {
+        let bytes = output
+            .lock()
+            .unwrap_or_else(|p| p.into_inner())
+            .bytes
+            .clone();
+        terminal_screen::text(&bytes, 80, 24)
+    };
     assert!(
         wait_until(Duration::from_secs(12), Duration::from_millis(20), || {
-            read_output(&output).contains("REMOTE_INITIAL_FRAME")
+            screen_text().contains("REMOTE_INITIAL_FRAME")
         }),
         "remote must be usable before Local exists: {}",
         read_output(&output)
@@ -922,6 +959,39 @@ fn federated_client_starts_without_local_and_survives_its_restart() {
         "saved machine discovery must opt into the advertised bridge idle timeout"
     );
 
+    let mut input = client._master.as_ref().unwrap().take_writer().unwrap();
+    send_pane_shell_command(&remote_api, remote_pane, "reconnect_survivor=ALIVE");
+    for cycle in 1..=3 {
+        let pid: libc::pid_t = fs::read_to_string(&bridge_pid)
+            .unwrap()
+            .trim()
+            .parse()
+            .unwrap();
+        assert_eq!(unsafe { libc::kill(pid, libc::SIGTERM) }, 0);
+        let marker = format!("REMOTE_RECONNECTED_{cycle}");
+        send_pane_shell_command(&remote_api, remote_pane, &format!("printf '{marker}\\n'"));
+        assert!(
+            wait_until(Duration::from_secs(15), Duration::from_millis(20), || {
+                screen_text().contains(&marker)
+            }),
+            "remote reconnect {cycle} must restore the visible screen without switching machines"
+        );
+        assert!(
+            wait_until(Duration::from_secs(8), Duration::from_millis(100), || {
+                if screen_text().contains(&format!("REMOTE_ALIVE_INPUT_{cycle}")) {
+                    return true;
+                }
+                write!(
+                    input,
+                    "printf 'REMOTE_%s_INPUT_{cycle}\\n' \"$reconnect_survivor\"\r"
+                )
+                .unwrap();
+                false
+            }),
+            "remote reconnect {cycle} must restore visible input and preserve the shell"
+        );
+    }
+
     let mut local = spawn_server(&config_home, &runtime_dir, &api_socket, &client_socket);
     wait_for_socket(&api_socket, Duration::from_secs(10));
     let created = send_json_request(
@@ -936,22 +1006,21 @@ fn federated_client_starts_without_local_and_survives_its_restart() {
     assert!(wait_until(
         Duration::from_secs(10),
         Duration::from_millis(20),
-        || read_output(&output).contains("local-online")
+        || screen_text().contains("local-online")
     ));
 
     local.child.kill().unwrap();
     local.close_master();
     drop(local);
-    let watermark = output_len(&output);
-    let mut input = client._master.as_ref().unwrap().take_writer().unwrap();
     input
         .write_all(b"printf 'REMOTE_%s\\n' SURVIVED\r")
         .unwrap();
     assert!(
         wait_until(Duration::from_secs(8), Duration::from_millis(20), || {
-            read_output(&output)[watermark..].contains("REMOTE_SURVIVED")
+            screen_text().contains("REMOTE_SURVIVED")
         }),
-        "Local loss must not interrupt remote input or output"
+        "Local loss must not interrupt remote input or output: {}",
+        screen_text()
     );
     assert!(client.child.try_wait().unwrap().is_none());
 
@@ -968,19 +1037,19 @@ fn federated_client_starts_without_local_and_survives_its_restart() {
     assert_eq!(created["result"]["type"], "workspace_created");
     assert!(
         wait_until(Duration::from_secs(12), Duration::from_millis(20), || {
-            read_output(&output).contains("local-returned")
+            screen_text().contains("local-returned")
         }),
         "Local must reconnect with fresh metadata"
     );
-    let watermark = output_len(&output);
     input
         .write_all(b"printf 'REMOTE_%s\\n' STILL_SELECTED\r")
         .unwrap();
     assert!(
         wait_until(Duration::from_secs(8), Duration::from_millis(20), || {
-            read_output(&output)[watermark..].contains("REMOTE_STILL_SELECTED")
+            screen_text().contains("REMOTE_STILL_SELECTED")
         }),
-        "Local recovery must not steal selection"
+        "Local recovery must not steal selection: {}",
+        screen_text()
     );
     let watermark = output_len(&output);
     remote_server.child.kill().unwrap();
diff --git a/tests/support/terminal_screen.rs b/tests/support/terminal_screen.rs
new file mode 100644
index 00000000..45d7731b
--- /dev/null
+++ b/tests/support/terminal_screen.rs
@@ -0,0 +1,89 @@
+use std::{mem::size_of, ptr};
+
+// Reuse the generated C API; this test observer needs only terminal parsing and formatting.
+#[allow(
+    dead_code,
+    non_camel_case_types,
+    non_snake_case,
+    non_upper_case_globals
+)]
+#[path = "../../src/ghostty/bindings.rs"]
+mod ffi;
+
+struct Screen {
+    terminal: ffi::GhosttyTerminal,
+    formatter: ffi::GhosttyFormatter,
+}
+
+impl Drop for Screen {
+    fn drop(&mut self) {
+        unsafe {
+            ffi::ghostty_formatter_free(self.formatter);
+            ffi::ghostty_terminal_free(self.terminal);
+        }
+    }
+}
+
+pub fn text(output: &[u8], cols: u16, rows: u16) -> String {
+    let mut screen = Screen {
+        terminal: ptr::null_mut(),
+        formatter: ptr::null_mut(),
+    };
+    let options = ffi::GhosttyFormatterTerminalOptions {
+        size: size_of::<ffi::GhosttyFormatterTerminalOptions>(),
+        emit: ffi::GhosttyFormatterFormat_GHOSTTY_FORMATTER_FORMAT_PLAIN,
+        trim: true,
+        extra: ffi::GhosttyFormatterTerminalExtra {
+            size: size_of::<ffi::GhosttyFormatterTerminalExtra>(),
+            screen: ffi::GhosttyFormatterScreenExtra {
+                size: size_of::<ffi::GhosttyFormatterScreenExtra>(),
+                ..Default::default()
+            },
+            ..Default::default()
+        },
+        ..Default::default()
+    };
+    unsafe {
+        assert_eq!(
+            ffi::ghostty_terminal_new(ptr::null(), &mut screen.terminal, cols, rows),
+            ffi::GhosttyResult_GHOSTTY_SUCCESS
+        );
+        ffi::ghostty_terminal_vt_write(screen.terminal, output.as_ptr(), output.len());
+        assert_eq!(
+            ffi::ghostty_formatter_terminal_new(
+                ptr::null(),
+                &mut screen.formatter,
+                screen.terminal,
+                options,
+            ),
+            ffi::GhosttyResult_GHOSTTY_SUCCESS
+        );
+        let mut len = 0;
+        let result =
+            ffi::ghostty_formatter_format_buf(screen.formatter, ptr::null_mut(), 0, &mut len);
+        assert!(matches!(
+            result,
+            ffi::GhosttyResult_GHOSTTY_SUCCESS | ffi::GhosttyResult_GHOSTTY_OUT_OF_SPACE
+        ));
+        let mut bytes = vec![0; len];
+        assert_eq!(
+            ffi::ghostty_formatter_format_buf(
+                screen.formatter,
+                bytes.as_mut_ptr(),
+                bytes.len(),
+                &mut len
+            ),
+            ffi::GhosttyResult_GHOSTTY_SUCCESS
+        );
+        String::from_utf8_lossy(&bytes[..len]).into_owned()
+    }
+}
+
+#[test]
+fn screen_text_reconstructs_partial_redraws() {
+    let output = b"REMOTE_SURVIVED\x1b[1;8HSTILL_SELECTED";
+    assert!(!output
+        .windows(b"REMOTE_STILL_SELECTED".len())
+        .any(|part| part == b"REMOTE_STILL_SELECTED"));
+    assert!(text(output, 80, 24).contains("REMOTE_STILL_SELECTED"));
+}