fix: avoid full-screen copies for retained client updates (#3837)

refs #3822

3 files changed, 175 insertions, 67 deletions

diff --git a/src/server/headless/retained_surface.rs b/src/server/headless/retained_surface.rs
index a62a43f6..259714a3 100644
--- a/src/server/headless/retained_surface.rs
+++ b/src/server/headless/retained_surface.rs
@@ -28,7 +28,7 @@ fn patch_intersects_hyperlinks(
         })
 }
 
-fn apply_patch_row(frame: &mut FrameData, row: &protocol::PaneSurfacePatchRow) -> Option<bool> {
+fn patch_row_changed(frame: &FrameData, row: &protocol::PaneSurfacePatchRow) -> Option<bool> {
     if row.y >= frame.height
         || row.x.saturating_add(u16::try_from(row.cells.len()).ok()?) > frame.width
     {
@@ -39,13 +39,11 @@ fn apply_patch_row(frame: &mut FrameData, row: &protocol::PaneSurfacePatchRow) -
     if end > frame.cells.len() {
         return None;
     }
-    let changed = frame.cells[start..end] != row.cells;
-    frame.cells[start..end].clone_from_slice(&row.cells);
-    Some(changed)
+    Some(frame.cells[start..end] != row.cells)
 }
 
-fn apply_rows(
-    frame: &mut FrameData,
+fn changed_rows(
+    frame: &FrameData,
     area: protocol::SurfaceRect,
     patch: &crate::pane::TerminalDirtyPatch,
 ) -> Option<Vec<protocol::PaneSurfacePatchRow>> {
@@ -89,15 +87,12 @@ fn apply_rows(
             offset = end;
         }
     }
-    for row in &rows {
-        apply_patch_row(frame, row)?;
-    }
     Some(rows)
 }
 
 fn retained_scrollbar_patch(
     app: &app::App,
-    frame: &mut FrameData,
+    frame: &FrameData,
     pane: &mut protocol::PaneSurfacePane,
     alternate_screen_active: bool,
     metrics: Option<crate::pane::ScrollMetrics>,
@@ -146,7 +141,7 @@ fn retained_scrollbar_patch(
             y: rect.y.checked_add(u16::try_from(offset).ok()?)?,
             cells: vec![cell],
         };
-        if apply_patch_row(frame, &row)? {
+        if patch_row_changed(frame, &row)? {
             rows.push(row);
         }
     }
@@ -186,9 +181,9 @@ fn retained_cursor(
         })
 }
 
-struct RetainedRecipient {
+struct RetainedRecipient<'a> {
     client_id: u64,
-    surface: protocol::PaneSurfaceFrame,
+    surface: &'a protocol::PaneSurfaceFrame,
 }
 
 struct CollectedPanePatch {
@@ -204,9 +199,11 @@ struct CollectedPanePatch {
 
 struct RetainedRecipientUpdate {
     client_id: u64,
-    surface: protocol::PaneSurfaceFrame,
     patch: protocol::PaneSurfacePatch,
-    graphics_delivery: Option<crate::kitty_graphics::surface::DeliveryCache>,
+    graphics: Option<(
+        protocol::PaneSurfaceFrame,
+        crate::kitty_graphics::surface::DeliveryCache,
+    )>,
 }
 
 impl HeadlessServer {
@@ -283,7 +280,7 @@ impl HeadlessServer {
             }
             recipients.push(RetainedRecipient {
                 client_id: *client_id,
-                surface: surface.clone(),
+                surface,
             });
         }
         if recipients.is_empty() {
@@ -355,7 +352,8 @@ impl HeadlessServer {
         let mut updates = Vec::with_capacity(recipients.len());
         for recipient in recipients {
             let client_id = recipient.client_id;
-            let mut surface = recipient.surface;
+            let surface = recipient.surface;
+            let mut panes = surface.panes.clone();
             let projection_revision = surface.projection_revision;
             let base_surface_revision = surface.surface_revision;
             let mut changed_panes = Vec::with_capacity(collected.len());
@@ -364,8 +362,7 @@ impl HeadlessServer {
             let mut refresh_graphics = !surface.graphics.placements.is_empty()
                 || !surface.graphics.retained_assets.is_empty();
             for collected_pane in &collected {
-                let Some(pane) = surface
-                    .panes
+                let Some(pane) = panes
                     .iter_mut()
                     .find(|pane| pane.pane_id == collected_pane.pane_id)
                 else {
@@ -387,14 +384,14 @@ impl HeadlessServer {
                 refresh_graphics |= collected_pane.graphics_may_have_placements;
                 let previous_pane = pane.clone();
                 let Some(rows) =
-                    apply_rows(&mut surface.frame, pane.inner_rect, &collected_pane.patch)
+                    changed_rows(&surface.frame, pane.inner_rect, &collected_pane.patch)
                 else {
                     fallback!("invalid_patch");
                 };
                 patch_rows.extend(rows);
                 let Some(scrollbar_rows) = retained_scrollbar_patch(
                     &self.app,
-                    &mut surface.frame,
+                    &surface.frame,
                     pane,
                     collected_pane.alternate_screen_active,
                     collected_pane.scroll_metrics,
@@ -417,19 +414,29 @@ impl HeadlessServer {
                 changed_panes.push(pane.clone());
             }
 
-            let cursor = retained_cursor(&self.app, &surface.panes);
+            let cursor = retained_cursor(&self.app, &panes);
             let cursor_changed = cursor != surface.frame.cursor;
-            surface.frame.cursor = cursor.clone();
+            let patch = protocol::PaneSurfacePatch {
+                boot_id: self.client_shell_boot_id.clone(),
+                projection_revision,
+                base_surface_revision,
+                surface_revision: 0,
+                rows: patch_rows,
+                panes: changed_panes,
+                cursor,
+            };
             let mut graphics_changed = false;
-            let graphics_delivery = if refresh_graphics {
+            let graphics = if refresh_graphics {
                 let Some(target) = self.shell_target_for_client(client_id) else {
                     fallback!("graphics_target");
                 };
                 let client = &self.clients[&client_id];
+                let mut next_surface = surface.clone();
+                crate::server::render_stream::apply_pane_surface_patch(&mut next_surface, &patch);
                 let Some((graphics, delivery)) =
                     crate::server::client_shell_graphics::collect_retained(
                         &self.app,
-                        &surface,
+                        &next_surface,
                         target,
                         client.cell_size,
                         &client.shell_graphics_delivery,
@@ -439,28 +446,18 @@ impl HeadlessServer {
                     fallback!("graphics_geometry");
                 };
                 graphics_changed = graphics != surface.graphics;
-                surface.graphics = graphics;
-                Some(delivery)
+                next_surface.graphics = graphics;
+                Some((next_surface, delivery))
             } else {
                 None
             };
-            if patch_rows.is_empty() && !cursor_changed && !metadata_changed && !graphics_changed {
+            if patch.rows.is_empty() && !cursor_changed && !metadata_changed && !graphics_changed {
                 continue;
             }
-            let patch = protocol::PaneSurfacePatch {
-                boot_id: self.client_shell_boot_id.clone(),
-                projection_revision,
-                base_surface_revision,
-                surface_revision: 0,
-                rows: patch_rows,
-                panes: changed_panes,
-                cursor,
-            };
             updates.push(RetainedRecipientUpdate {
                 client_id,
-                surface,
                 patch,
-                graphics_delivery,
+                graphics,
             });
         }
         if updates.is_empty() {
@@ -473,9 +470,8 @@ impl HeadlessServer {
         for update in updates {
             let RetainedRecipientUpdate {
                 client_id,
-                surface,
                 patch,
-                graphics_delivery,
+                graphics,
             } = update;
             let Some(client) = self.clients.get_mut(&client_id) else {
                 continue;
@@ -487,12 +483,13 @@ impl HeadlessServer {
             };
             // The published row patch cannot carry images. Reuse the retained text/layout
             // in a graphics-capable surface message rather than invoking the full renderer.
-            let prepared = if graphics_delivery.is_some() {
-                client.render_state.prepare_pane_surface(surface)
+            let (prepared, graphics_delivery) = if let Some((surface, delivery)) = graphics {
+                (
+                    client.render_state.prepare_pane_surface(surface),
+                    Some(delivery),
+                )
             } else {
-                client
-                    .render_state
-                    .prepare_pane_surface_patch(patch, surface)
+                (client.render_state.prepare_pane_surface_patch(patch), None)
             };
             let Some(prepared) = prepared else {
                 client.defer_full_render();
@@ -573,7 +570,7 @@ mod tests {
 
     #[test]
     fn retained_rows_send_only_changed_cell_spans() {
-        let mut frame = FrameData {
+        let frame = FrameData {
             width: 6,
             height: 2,
             cells: vec![cell(" "); 12],
@@ -585,8 +582,8 @@ mod tests {
             rows: vec![(0, vec![cell(" "), cell("x"), cell("y"), cell(" ")])],
         };
 
-        let rows = apply_rows(
-            &mut frame,
+        let rows = changed_rows(
+            &frame,
             protocol::SurfaceRect {
                 x: 1,
                 y: 1,
@@ -605,13 +602,12 @@ mod tests {
                 cells: vec![cell("x"), cell("y"), cell(" ")],
             }]
         );
-        assert_eq!(frame.cells[8], cell("x"));
-        assert_eq!(frame.cells[9], cell("y"));
+        assert_eq!(frame.cells, vec![cell(" "); 12], "planning must not commit");
     }
 
     #[test]
     fn retained_rows_include_the_cell_after_a_width_transition() {
-        let mut frame = FrameData {
+        let frame = FrameData {
             width: 3,
             height: 1,
             cells: vec![cell("界"), cell("z"), cell("q")],
@@ -623,8 +619,8 @@ mod tests {
             rows: vec![(0, vec![cell("x"), cell("z"), cell("q")])],
         };
 
-        let rows = apply_rows(
-            &mut frame,
+        let rows = changed_rows(
+            &frame,
             protocol::SurfaceRect {
                 x: 0,
                 y: 0,
@@ -647,7 +643,7 @@ mod tests {
 
     #[test]
     fn retained_rows_omit_unchanged_full_dirty_rows() {
-        let mut frame = FrameData {
+        let frame = FrameData {
             width: 4,
             height: 2,
             cells: vec![cell(" "); 8],
@@ -659,8 +655,8 @@ mod tests {
             rows: vec![(0, vec![cell(" "); 4]), (1, vec![cell(" "); 4])],
         };
 
-        let rows = apply_rows(
-            &mut frame,
+        let rows = changed_rows(
+            &frame,
             protocol::SurfaceRect {
                 x: 0,
                 y: 0,
diff --git a/src/server/headless/tests/mod.rs b/src/server/headless/tests/mod.rs
index 63f7264d..41e08164 100644
--- a/src/server/headless/tests/mod.rs
+++ b/src/server/headless/tests/mod.rs
@@ -821,6 +821,13 @@ async fn client_shell_receives_metadata_then_shell_free_pane_surface() {
         other => panic!("expected pane surface, got {other:?}"),
     };
 
+    let baseline = server.clients[&7]
+        .render_state
+        .last_pane_surface()
+        .expect("initial baseline");
+    let cells_ptr = baseline.frame.cells.as_ptr();
+    let untouched_symbol_ptr = baseline.frame.cells.last().unwrap().symbol.as_ptr();
+
     server
         .app
         .state
@@ -846,6 +853,15 @@ async fn client_shell_receives_metadata_then_shell_free_pane_surface() {
         }
         other => panic!("expected pane surface patch, got {other:?}"),
     }
+    let patched = server.clients[&7].render_state.last_pane_surface().unwrap();
+    assert_eq!(
+        (
+            patched.frame.cells.as_ptr(),
+            patched.frame.cells.last().unwrap().symbol.as_ptr()
+        ),
+        (cells_ptr, untouched_symbol_ptr),
+        "a text patch must preserve the frame and unchanged cell storage"
+    );
     server
         .app
         .state
@@ -855,6 +871,7 @@ async fn client_shell_receives_metadata_then_shell_free_pane_surface() {
     assert!(server.render_retained_pane_surface_and_stream(&sources));
     match read_server_message(render_rx.recv().expect("metadata-only pane surface patch")) {
         ServerMessage::PaneSurfacePatch(patch) => {
+            assert!(patch.rows.is_empty(), "mouse modes only change metadata");
             assert_eq!(patch.panes.len(), 1);
             assert!(!patch.panes[0].mouse_reporting);
             assert!(!patch.panes[0].sgr_pixel_mouse);
@@ -864,8 +881,14 @@ async fn client_shell_receives_metadata_then_shell_free_pane_surface() {
     let retained = server.clients[&7]
         .render_state
         .last_pane_surface()
-        .expect("committed retained surface")
-        .clone();
+        .expect("committed retained surface");
+    assert_eq!(retained.frame.cells.as_ptr(), cells_ptr);
+    assert_eq!(
+        retained.frame.cells.last().unwrap().symbol.as_ptr(),
+        untouched_symbol_ptr,
+        "retained updates must not copy unchanged screen cells"
+    );
+    let retained = retained.clone();
     server
         .clients
         .get_mut(&7)
@@ -1135,6 +1158,47 @@ async fn retained_patches_only_reach_shells_viewing_the_dirty_tab() {
 }
 
 #[tokio::test]
+async fn late_retained_fallback_leaves_all_client_baselines_unchanged() {
+    let mut server = test_headless_server();
+    let pane_id = install_shared_view_test_runtime(&mut server);
+    let (_first_control, first_render) = connect_matching_test_shell(&mut server, 7);
+    let (_second_control, second_render) = connect_matching_test_shell(&mut server, 8);
+    server.render_and_stream();
+    let _ = recv_pane_surface(&first_render, "first baseline");
+    let _ = recv_pane_surface(&second_render, "second baseline");
+
+    // Foreground renders last. Its old hyperlink forces a fallback after the first plan.
+    server.foreground_client_id = Some(8);
+    let crate::server::render_stream::ClientRenderState::Semantic { last_surface, .. } =
+        &mut server.clients.get_mut(&8).unwrap().render_state
+    else {
+        panic!("semantic client");
+    };
+    let linked = last_surface.as_mut().unwrap();
+    linked.frame.hyperlinks.push("https://example.com".into());
+    linked.frame.cells[0].hyperlink = Some(0);
+    let before = [7, 8].map(|id| {
+        server.clients[&id]
+            .render_state
+            .last_pane_surface()
+            .unwrap()
+            .clone()
+    });
+
+    write_shared_test_pane(&mut server, pane_id, b"\rNEXT\x1b[?1003h");
+    assert!(!server.render_retained_pane_surface_and_stream(&HashSet::from([pane_id])));
+    assert!(first_render.try_recv().is_err());
+    assert!(second_render.try_recv().is_err());
+    for (id, expected) in [7, 8].into_iter().zip(before) {
+        assert_eq!(
+            server.clients[&id].render_state.last_pane_surface(),
+            Some(&expected)
+        );
+    }
+    shutdown_test_runtimes(&mut server);
+}
+
+#[tokio::test]
 async fn backpressured_shell_does_not_disable_retained_patches_for_responsive_peer() {
     let mut server = test_headless_server();
     let pane_id = install_shared_view_test_runtime(&mut server);
@@ -1156,13 +1220,23 @@ async fn backpressured_shell_does_not_disable_retained_patches_for_responsive_pe
         ServerMessage::PaneSurfacePatch(_)
     ));
 
-    write_shared_test_pane(&mut server, pane_id, b"\rTWO");
+    let slow_baseline = server.clients[&8]
+        .render_state
+        .last_pane_surface()
+        .unwrap()
+        .clone();
+    write_shared_test_pane(&mut server, pane_id, b"\rTWO\x1b[?1003h");
     assert!(server.render_retained_pane_surface_and_stream(&sources));
     assert!(matches!(
         read_server_message(responsive_render.recv().expect("responsive second patch")),
         ServerMessage::PaneSurfacePatch(_)
     ));
     assert_eq!(server.clients[&8].deferred_render(), DeferredRender::Full);
+    assert_eq!(
+        server.clients[&8].render_state.last_pane_surface(),
+        Some(&slow_baseline),
+        "queue-full must not advance cells, metadata, cursor, or revision"
+    );
 
     write_shared_test_pane(&mut server, pane_id, b"\rTHREE");
     assert!(server.render_retained_pane_surface_and_stream(&sources));
diff --git a/src/server/render_stream.rs b/src/server/render_stream.rs
index 41684c26..349454f3 100644
--- a/src/server/render_stream.rs
+++ b/src/server/render_stream.rs
@@ -145,9 +145,8 @@ impl ClientRenderState {
     }
 
     pub(crate) fn prepare_pane_surface_patch(
-        &mut self,
+        &self,
         mut patch: PaneSurfacePatch,
-        mut committed_surface: PaneSurfaceFrame,
     ) -> Option<PreparedRender> {
         let Self::Semantic {
             last_surface,
@@ -165,11 +164,8 @@ impl ClientRenderState {
         }
         let next_revision = surface_revision.saturating_add(1);
         patch.surface_revision = next_revision;
-        committed_surface.surface_revision = next_revision;
-        committed_surface.graphics.assets.clear();
-        Some(PreparedRender::Semantic {
+        Some(PreparedRender::SemanticPatch {
             message: ServerMessage::PaneSurfacePatch(patch),
-            committed_surface: Box::new(committed_surface),
         })
     }
 
@@ -188,6 +184,21 @@ impl ClientRenderState {
                 *last_surface = Some(committed_surface);
             }
             (
+                Self::Semantic {
+                    last_surface,
+                    surface_revision,
+                },
+                PreparedRender::SemanticPatch {
+                    message: ServerMessage::PaneSurfacePatch(patch),
+                },
+            ) => {
+                let surface = last_surface
+                    .as_deref_mut()
+                    .expect("prepared patch baseline");
+                apply_pane_surface_patch(surface, &patch);
+                *surface_revision = patch.surface_revision;
+            }
+            (
                 Self::TerminalAnsi {
                     blit_encoder,
                     seq,
@@ -208,6 +219,28 @@ impl ClientRenderState {
     }
 }
 
+// Planning validates all rows and pane IDs before any send. The server does not yield
+// between planning and commit, so applying the accepted patch cannot fail partway through.
+pub(super) fn apply_pane_surface_patch(surface: &mut PaneSurfaceFrame, patch: &PaneSurfacePatch) {
+    debug_assert_eq!(surface.boot_id, patch.boot_id);
+    debug_assert_eq!(surface.projection_revision, patch.projection_revision);
+    debug_assert_eq!(surface.surface_revision, patch.base_surface_revision);
+    for row in &patch.rows {
+        let start = usize::from(row.y) * usize::from(surface.frame.width) + usize::from(row.x);
+        surface.frame.cells[start..start + row.cells.len()].clone_from_slice(&row.cells);
+    }
+    for updated in &patch.panes {
+        let pane = surface
+            .panes
+            .iter_mut()
+            .find(|pane| pane.pane_id == updated.pane_id)
+            .expect("planned patch pane");
+        pane.clone_from(updated);
+    }
+    surface.frame.cursor.clone_from(&patch.cursor);
+    surface.surface_revision = patch.surface_revision;
+}
+
 fn insert_graphics_before_sync_end(encoded: &mut Vec<u8>, graphics: &[u8]) {
     if graphics.is_empty() {
         return;
@@ -226,6 +259,9 @@ pub(crate) enum PreparedRender {
         message: ServerMessage,
         committed_surface: Box<PaneSurfaceFrame>,
     },
+    SemanticPatch {
+        message: ServerMessage,
+    },
     TerminalAnsi {
         message: ServerMessage,
         frame: FrameData,
@@ -236,7 +272,9 @@ pub(crate) enum PreparedRender {
 impl PreparedRender {
     pub(crate) fn message(&self) -> &ServerMessage {
         match self {
-            Self::Semantic { message, .. } | Self::TerminalAnsi { message, .. } => message,
+            Self::Semantic { message, .. }
+            | Self::SemanticPatch { message }
+            | Self::TerminalAnsi { message, .. } => message,
         }
     }