History

edit_prediction: Prevent v4 cursor marker leaks (#61957)

# Objective Prevent V4 cursor markers from appearing in edit prediction previews or accepted edits. ## Solution Replace the cursor marker with an internal sentinel before sub-line diffing, so matching source text cannot reconstruct the marker. ## Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content adheres to Zed's UI standards ([UX/UI](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) and [icon](https://github.com/zed-industries/zed/blob/main/crates/icons/README.md) guidelines) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable --- Release Notes: - N/A

Ben Kunkle committed 3 months ago

Commit 4f047ac1 parent b2090001 file · +37 −78Browse files

Unified · Split

crates/edit_prediction/src/udiff.rs+37−78
@@ -22,6 +22,8 @@ pub use zeta_prompt::udiff::{ · 21 unmodified lines
2222 strip_diff_metadata, strip_diff_path_prefix,
2323 };
2424
25+const INLINE_CURSOR_SENTINEL: &str = "\u{fdd0}";
26+
2527 #[derive(Clone, Debug)]
2628 pub struct OpenedBuffers(HashMap<String, Entity<Buffer>>);
2729
@@ -54,7 +56,11 @@ pub async fn prediction_edits_for_single_file_diff( · 26 unmodified lines
5456
5557 while let Some(event) = diff.next()? {
5658 match event {
57− DiffEvent::Hunk { path, hunk, status } => {
59+ DiffEvent::Hunk {
60+ path,
61+ mut hunk,
62+ status,
63+ } => {
5864 anyhow::ensure!(
5965 status == FileStatus::Modified,
6066 "V4 edit predictions only support modifying existing files"
@@ -81,89 +87,37 @@ pub async fn prediction_edits_for_single_file_diff( · 20 unmodified lines
8187 }
8288
8389 let (_, _, snapshot) = target_file.as_ref().context("missing target file")?;
84− let mut pending_marker: Option<(Range<Anchor>, String, usize)> = None;
90+ for edit in &mut hunk.edits {
91+ while let Some(marker_offset) = edit.text.find(INLINE_CURSOR_MARKER) {
92+ edit.text.replace_range(
93+ marker_offset..marker_offset + INLINE_CURSOR_MARKER.len(),
94+ INLINE_CURSOR_SENTINEL,
95+ );
96+ }
97+ }
98+
8599 for (range, text) in resolve_hunk_edits_in_buffer(
86100 hunk,
87101 snapshot,
88102 &[Anchor::min_max_range_for_buffer(snapshot.remote_id())],
89103 status,
90104 )? {
91− let mut remaining = text.as_ref();
92− let mut output = String::new();
93−
94− if let Some((pending_range, mut pending_text, pending_offset)) =
95− pending_marker.take()
96− {
97− let matched_len = INLINE_CURSOR_MARKER[pending_text.len()..]
98− .bytes()
99− .zip(remaining.bytes())
100− .take_while(|(left, right)| left == right)
101− .count();
102−
103− if matched_len == 0 {
104− edits.push((pending_range, pending_text.into()));
105− } else {
106− let marker_len = pending_text.len() + matched_len;
107− if marker_len == INLINE_CURSOR_MARKER.len() {
108− cursor_position.get_or_insert_with(|| {
109− PredictedCursorPosition::new(
110− snapshot
111− .anchor_before(pending_range.start.to_offset(snapshot)),
112− pending_offset,
113− )
114− });
115− remaining = &remaining[matched_len..];
116− } else if matched_len == remaining.len() {
117− pending_text.push_str(
118− &INLINE_CURSOR_MARKER[pending_text.len()..marker_len],
119− );
120− pending_marker =
121− Some((pending_range, pending_text, pending_offset));
122− continue;
123− } else {
124− pending_text.push_str(&remaining[..matched_len]);
125− edits.push((pending_range, pending_text.into()));
126− remaining = &remaining[matched_len..];
127− }
128− }
129− }
130−
131− while let Some(marker_offset) = remaining.find(INLINE_CURSOR_MARKER) {
132− output.push_str(&remaining[..marker_offset]);
105+ if let Some(marker_offset) = text.find(INLINE_CURSOR_SENTINEL) {
133106 cursor_position.get_or_insert_with(|| {
134107 PredictedCursorPosition::new(
135108 snapshot.anchor_before(range.start.to_offset(snapshot)),
136− output.len(),
109+ marker_offset,
137110 )
138111 });
139− remaining = &remaining[marker_offset + INLINE_CURSOR_MARKER.len()..];
140− }
141−
142− let marker_prefix_len = (1..=INLINE_CURSOR_MARKER.len().min(remaining.len()))
143− .rev()
144− .find(|prefix_len| {
145− remaining.ends_with(&INLINE_CURSOR_MARKER[..*prefix_len])
146− });
147− if let Some(marker_prefix_len) = marker_prefix_len {
148− let marker_start = remaining.len() - marker_prefix_len;
149− output.push_str(&remaining[..marker_start]);
150− pending_marker = Some((
151− range.clone(),
152− remaining[marker_start..].to_string(),
153− output.len(),
154− ));
112+ let text = text.replace(INLINE_CURSOR_SENTINEL, "");
113+ if range.start.to_offset(snapshot) != range.end.to_offset(snapshot)
114+ || !text.is_empty()
115+ {
116+ edits.push((range, text.into()));
117+ }
155118 } else {
156− output.push_str(remaining);
119+ edits.push((range, text));
157120 }
158−
159− if range.start.to_offset(snapshot) != range.end.to_offset(snapshot)
160− || !output.is_empty()
161− {
162− edits.push((range, output.into()));
163− }
164− }
165− if let Some((range, text, _)) = pending_marker {
166− edits.push((range, text.into()));
167121 }
168122 }
169123 DiffEvent::FileEnd { renamed_to } => {
@@ -698,14 +652,14 @@ mod tests { · 528 unmodified lines
698652 }
699653
700654 #[gpui::test]
701− async fn test_prediction_edits_for_single_file_diff_does_not_treat_completed_literal_marker_as_cursor(
655+ async fn test_prediction_edits_for_single_file_diff_strips_cursor_marker_that_overlaps_source(
702656 cx: &mut TestAppContext,
703657 ) {
704658 let fs = init_test(cx);
705659 fs.insert_tree(
706660 path!("/root"),
707661 json!({
708− "file": "text <|user_cursor\n",
662+ "file": "before<after\n",
709663 }),
710664 )
711665 .await;
@@ -715,20 +669,25 @@ mod tests { · 3 unmodified lines
715669 --- a/file
716670 +++ b/file
717671 @@ ... @@
718− -text <|user_cursor
719− +text <|user_cursor|>
672+ -before<after
673+ +before<|user_cursor|><after
720674 "#};
721675
722− let (buffer, _, edits, cursor_position) =
676+ let (buffer, snapshot, edits, cursor_position) =
723677 prediction_edits_for_single_file_diff(diff, &project, &mut cx.to_async())
724678 .await
725679 .unwrap()
726680 .unwrap();
727681
728− assert!(cursor_position.is_none());
682+ assert!(edits.is_empty());
683+ let cursor_position = cursor_position.unwrap();
684+ assert_eq!(
685+ cursor_position.anchor.to_offset(&snapshot) + cursor_position.offset,
686+ "before".len()
687+ );
729688 buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx));
730689 buffer.read_with(cx, |buffer, _cx| {
731− assert_eq!(buffer.text(), "text <|user_cursor|>\n");
690+ assert_eq!(buffer.text(), "before<after\n");
732691 });
733692 }
734693