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
22 strip_diff_metadata, strip_diff_path_prefix,22 strip_diff_metadata, strip_diff_path_prefix,
23};23};
24 24
25const INLINE_CURSOR_SENTINEL: &str = "\u{fdd0}";
26
25#[derive(Clone, Debug)]27#[derive(Clone, Debug)]
26pub struct OpenedBuffers(HashMap<String, Entity<Buffer>>);28pub struct OpenedBuffers(HashMap<String, Entity<Buffer>>);
27 29
@@ -54,7 +56,11 @@ pub async fn prediction_edits_for_single_file_diff( · 26 unmodified lines
54 56
55 while let Some(event) = diff.next()? {57 while let Some(event) = diff.next()? {
56 match event {58 match event {
57 DiffEvent::Hunk { path, hunk, status } => {59 DiffEvent::Hunk {
60 path,
61 mut hunk,
62 status,
63 } => {
58 anyhow::ensure!(64 anyhow::ensure!(
59 status == FileStatus::Modified,65 status == FileStatus::Modified,
60 "V4 edit predictions only support modifying existing files"66 "V4 edit predictions only support modifying existing files"
@@ -81,89 +87,37 @@ pub async fn prediction_edits_for_single_file_diff( · 20 unmodified lines
81 }87 }
82 88
83 let (_, _, snapshot) = target_file.as_ref().context("missing target file")?;89 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
85 for (range, text) in resolve_hunk_edits_in_buffer(99 for (range, text) in resolve_hunk_edits_in_buffer(
86 hunk,100 hunk,
87 snapshot,101 snapshot,
88 &[Anchor::min_max_range_for_buffer(snapshot.remote_id())],102 &[Anchor::min_max_range_for_buffer(snapshot.remote_id())],
89 status,103 status,
90 )? {104 )? {
91 let mut remaining = text.as_ref();105 if let Some(marker_offset) = text.find(INLINE_CURSOR_SENTINEL) {
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]);
133 cursor_position.get_or_insert_with(|| {106 cursor_position.get_or_insert_with(|| {
134 PredictedCursorPosition::new(107 PredictedCursorPosition::new(
135 snapshot.anchor_before(range.start.to_offset(snapshot)),108 snapshot.anchor_before(range.start.to_offset(snapshot)),
136 output.len(),109 marker_offset,
137 )110 )
138 });111 });
139 remaining = &remaining[marker_offset + INLINE_CURSOR_MARKER.len()..];112 let text = text.replace(INLINE_CURSOR_SENTINEL, "");
140 }113 if range.start.to_offset(snapshot) != range.end.to_offset(snapshot)
141 114 || !text.is_empty()
142 let marker_prefix_len = (1..=INLINE_CURSOR_MARKER.len().min(remaining.len()))115 {
143 .rev()116 edits.push((range, text.into()));
144 .find(|prefix_len| {117 }
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 ));
155 } else {118 } else {
156 output.push_str(remaining);119 edits.push((range, text));
157 }120 }
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()));
167 }121 }
168 }122 }
169 DiffEvent::FileEnd { renamed_to } => {123 DiffEvent::FileEnd { renamed_to } => {
@@ -698,14 +652,14 @@ mod tests { · 528 unmodified lines
698 }652 }
699 653
700 #[gpui::test]654 #[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(
702 cx: &mut TestAppContext,656 cx: &mut TestAppContext,
703 ) {657 ) {
704 let fs = init_test(cx);658 let fs = init_test(cx);
705 fs.insert_tree(659 fs.insert_tree(
706 path!("/root"),660 path!("/root"),
707 json!({661 json!({
708 "file": "text <|user_cursor\n",662 "file": "before<after\n",
709 }),663 }),
710 )664 )
711 .await;665 .await;
@@ -715,20 +669,25 @@ mod tests { · 3 unmodified lines
715 --- a/file669 --- a/file
716 +++ b/file670 +++ b/file
717 @@ ... @@671 @@ ... @@
718 -text <|user_cursor672 -before<after
719 +text <|user_cursor|>673 +before<|user_cursor|><after
720 "#};674 "#};
721 675
722 let (buffer, _, edits, cursor_position) =676 let (buffer, snapshot, edits, cursor_position) =
723 prediction_edits_for_single_file_diff(diff, &project, &mut cx.to_async())677 prediction_edits_for_single_file_diff(diff, &project, &mut cx.to_async())
724 .await678 .await
725 .unwrap()679 .unwrap()
726 .unwrap();680 .unwrap();
727 681
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 );
729 buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx));688 buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx));
730 buffer.read_with(cx, |buffer, _cx| {689 buffer.read_with(cx, |buffer, _cx| {
731 assert_eq!(buffer.text(), "text <|user_cursor|>\n");690 assert_eq!(buffer.text(), "before<after\n");
732 });691 });
733 }692 }
734 693