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

refs #3842

2 files changed, 74 insertions, 0 deletions

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);