fix: preserve mouse selection during pane output (#3538)

2 files changed, 157 insertions, 13 deletions

diff --git a/src/client/shell/state.rs b/src/client/shell/state.rs
index 69e77dc5..ffac93da 100644
--- a/src/client/shell/state.rs
+++ b/src/client/shell/state.rs
@@ -4,6 +4,63 @@ pub(super) const MIN_TAB_WIDTH: u16 = 8;
 pub(super) const NEW_TAB_WIDTH: u16 = 3;
 pub(super) const WORKSPACE_HEADER_ROWS: u16 = 2;
 
+fn pane_surface_row<'a>(
+    surface: &'a PaneSurfaceFrame,
+    pane: &crate::protocol::PaneSurfacePane,
+    absolute_row: u32,
+) -> Option<&'a [crate::protocol::CellData]> {
+    let viewport_top = pane
+        .scroll
+        .map(|scroll| {
+            scroll
+                .max_offset_from_bottom
+                .saturating_sub(scroll.offset_from_bottom) as u32
+        })
+        .unwrap_or(0);
+    let viewport_row = u16::try_from(absolute_row.checked_sub(viewport_top)?).ok()?;
+    if viewport_row >= pane.inner_rect.height {
+        return None;
+    }
+    let start = (usize::from(pane.inner_rect.y) + usize::from(viewport_row))
+        * usize::from(surface.frame.width)
+        + usize::from(pane.inner_rect.x);
+    surface
+        .frame
+        .cells
+        .get(start..start + usize::from(pane.inner_rect.width))
+}
+
+fn selection_cells_unchanged(
+    selection: &crate::selection::Selection<String>,
+    previous_surface: &PaneSurfaceFrame,
+    previous_pane: &crate::protocol::PaneSurfacePane,
+    next_surface: &PaneSurfaceFrame,
+    next_pane: &crate::protocol::PaneSurfacePane,
+) -> bool {
+    let ((start_row, start_col), (end_row, end_col)) = selection.ordered_cells();
+    (start_row..=end_row).all(|row| {
+        let first_col = if row == start_row { start_col } else { 0 };
+        let last_col = if row == end_row {
+            end_col
+        } else {
+            previous_pane.inner_rect.width.saturating_sub(1)
+        };
+        pane_surface_row(previous_surface, previous_pane, row)
+            .zip(pane_surface_row(next_surface, next_pane, row))
+            .and_then(|(previous, next)| {
+                previous
+                    .get(usize::from(first_col)..=usize::from(last_col))
+                    .zip(next.get(usize::from(first_col)..=usize::from(last_col)))
+            })
+            .is_some_and(|(previous, next)| {
+                previous
+                    .iter()
+                    .zip(next)
+                    .all(|(previous, next)| previous.symbol == next.symbol)
+            })
+    })
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) enum ClientShellKeybindingSource {
     Local,
@@ -1481,21 +1538,34 @@ impl ClientShellState {
             self.popup_pending_deadline = None;
         }
         let selection_content_changed = self.selection.as_ref().is_some_and(|selection| {
-            let previous_revision = self.pane_surface.as_ref().and_then(|previous| {
-                previous
-                    .panes
-                    .iter()
-                    .find(|pane| pane.pane_id == selection.pane_id)
-                    .map(|pane| pane.content_revision)
-            });
-            let next_revision = surface
+            let Some(previous_surface) = self.pane_surface.as_ref() else {
+                return false;
+            };
+            let previous = previous_surface
+                .panes
+                .iter()
+                .find(|pane| pane.pane_id == selection.pane_id);
+            let next = surface
                 .panes
                 .iter()
-                .find(|pane| pane.pane_id == selection.pane_id)
-                .map(|pane| pane.content_revision);
-            previous_revision.is_some()
-                && next_revision.is_some()
-                && previous_revision != next_revision
+                .find(|pane| pane.pane_id == selection.pane_id);
+            let (Some(previous), Some(next)) = (previous, next) else {
+                return false;
+            };
+            previous.content_revision != next.content_revision
+                && (!selection.is_in_progress()
+                    || !previous.content_revision.is_multiple_of(2)
+                    || !next.content_revision.is_multiple_of(2)
+                    || previous.inner_rect.width != next.inner_rect.width
+                    || previous.inner_rect.height != next.inner_rect.height
+                    || previous.alternate_screen_active != next.alternate_screen_active
+                    || !selection_cells_unchanged(
+                        selection,
+                        previous_surface,
+                        previous,
+                        &surface,
+                        next,
+                    ))
         });
         if selection_content_changed {
             self.selection = None;
diff --git a/src/client/shell/tests/mouse_selection.rs b/src/client/shell/tests/mouse_selection.rs
index 9734db7c..a8f5c532 100644
--- a/src/client/shell/tests/mouse_selection.rs
+++ b/src/client/shell/tests/mouse_selection.rs
@@ -316,6 +316,80 @@ fn client_double_click_selects_and_copies_endpoint_row_word() {
 }
 
 #[test]
+fn pane_content_updates_preserve_active_selection_only_when_selected_cells_stay_stable() {
+    let mut state = ClientShellState::new(ClientShellConfig::from_config(&Config::default()));
+    state.set_snapshot(Box::new(snapshot()));
+    let surface_at = |surface_revision, content_revision, alternate_screen_active| {
+        let mut pane_surface = surface();
+        pane_surface.surface_revision = surface_revision;
+        pane_surface.panes[0].content_revision = content_revision;
+        pane_surface.panes[0].scroll = Some(crate::protocol::PaneSurfaceScrollMetrics {
+            offset_from_bottom: 0,
+            max_offset_from_bottom: 11,
+            viewport_rows: 2,
+        });
+        pane_surface.panes[0].alternate_screen_active = alternate_screen_active;
+        pane_surface
+    };
+    state.set_pane_surface(surface_at(1, 0, true));
+    state.compose(106, 20).expect("composed frame");
+    let pane = state.hits.panes[0].clone();
+    let mouse = |kind, column, row| {
+        RawInputEvent::Mouse(crossterm::event::MouseEvent {
+            kind,
+            column,
+            row,
+            modifiers: KeyModifiers::empty(),
+        })
+    };
+
+    state.handle_raw_events(vec![mouse(
+        MouseEventKind::Down(MouseButton::Left),
+        pane.inner_rect.x,
+        pane.inner_rect.y + 1,
+    )]);
+    let mut updated_surface = surface_at(2, 2, true);
+    updated_surface.frame.cells[0].symbol = "W".into();
+    state.set_pane_surface(updated_surface);
+    state.compose(106, 20).expect("updated frame");
+
+    let drag = state.handle_raw_events(vec![mouse(
+        MouseEventKind::Drag(MouseButton::Left),
+        pane.inner_rect.x + 1,
+        pane.inner_rect.y + 1,
+    )]);
+
+    assert!(drag.repaint);
+    let selection = state.selection.as_ref().expect("visible selection");
+    assert!(selection.is_visible());
+    assert_eq!(selection.ordered_cells(), ((12, 0), (12, 1)));
+
+    state.selection = Some(crate::selection::Selection::absolute_anchor(
+        "pane_1".to_owned(),
+        (12, 0),
+    ));
+    let mut replaced_surface = surface_at(3, 4, true);
+    replaced_surface.frame.cells[4].symbol = "X".into();
+    state.set_pane_surface(replaced_surface);
+    assert!(state.selection.is_none());
+
+    for (surface_revision, content_revision, width, alternate_screen_active) in
+        [(4, 6, 4, false), (5, 8, 3, false), (6, 9, 3, false)]
+    {
+        state.selection = Some(crate::selection::Selection::absolute_anchor(
+            "pane_1".to_owned(),
+            (12, 0),
+        ));
+        let mut changed_surface =
+            surface_at(surface_revision, content_revision, alternate_screen_active);
+        changed_surface.panes[0].inner_rect.width = width;
+        changed_surface.panes[0].alternate_screen_active = alternate_screen_active;
+        state.set_pane_surface(changed_surface);
+        assert!(state.selection.is_none());
+    }
+}
+
+#[test]
 fn pane_mouse_input_keeps_stable_target_and_endpoint_encoding() {
     let mut state = ClientShellState::new(ClientShellConfig::from_config(&Config::default()));
     state.set_snapshot(Box::new(snapshot()));