1use edit_prediction_types::{
2 EditPredictionDelegate, EditPredictionIconSet, PredictedCursorPosition,
3};
4use gpui::{Entity, KeyBinding, Modifiers, prelude::*};
5use indoc::indoc;
6use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
7use std::{ops::Range, sync::Arc};
8use text::{Point, ToOffset};
9use ui::prelude::*;
10
11use crate::{
12 AcceptEditPrediction, EditPrediction, MenuEditPredictionsPolicy, editor_tests::init_test,
13 test::editor_test_context::EditorTestContext,
14};
15
16#[gpui::test]
17async fn test_edit_prediction_insert(cx: &mut gpui::TestAppContext) {
18 init_test(cx, |_| {});
19
20 let mut cx = EditorTestContext::new(cx).await;
21 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
22 assign_editor_completion_provider(provider.clone(), &mut cx);
23 cx.set_state("let absolute_zero_celsius = ˇ;");
24
25 propose_edits(&provider, vec![(28..28, "-273.15")], &mut cx);
26 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
27
28 assert_editor_active_edit_completion(&mut cx, |_, edits| {
29 assert_eq!(edits.len(), 1);
30 assert_eq!(edits[0].1.as_ref(), "-273.15");
31 });
32
33 accept_completion(&mut cx);
34
35 cx.assert_editor_state("let absolute_zero_celsius = -273.15ˇ;")
36}
37
38#[gpui::test]
39async fn test_edit_prediction_cursor_position_inside_insertion(cx: &mut gpui::TestAppContext) {
40 init_test(cx, |_| {});
41
42 let mut cx = EditorTestContext::new(cx).await;
43 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
44 assign_editor_completion_provider(provider.clone(), &mut cx);
45 // Buffer: "fn foo() {}" - we'll insert text and position cursor inside the insertion
46 cx.set_state("fn foo() ˇ{}");
47
48 // Insert "bar()" at offset 9, with cursor at offset 2 within the insertion (after "ba")
49 // This tests the case where cursor is inside newly inserted text
50 propose_edits_with_cursor_position_in_insertion(
51 &provider,
52 vec![(9..9, "bar()")],
53 9, // anchor at the insertion point
54 2, // offset 2 within "bar()" puts cursor after "ba"
55 &mut cx,
56 );
57 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
58
59 assert_editor_active_edit_completion(&mut cx, |_, edits| {
60 assert_eq!(edits.len(), 1);
61 assert_eq!(edits[0].1.as_ref(), "bar()");
62 });
63
64 accept_completion(&mut cx);
65
66 // Cursor should be inside the inserted text at "baˇr()"
67 cx.assert_editor_state("fn foo() baˇr(){}");
68}
69
70#[gpui::test]
71async fn test_edit_prediction_cursor_position_outside_edit(cx: &mut gpui::TestAppContext) {
72 init_test(cx, |_| {});
73
74 let mut cx = EditorTestContext::new(cx).await;
75 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
76 assign_editor_completion_provider(provider.clone(), &mut cx);
77 // Buffer: "let x = ;" with cursor before semicolon - we'll insert "42" and position cursor elsewhere
78 cx.set_state("let x = ˇ;");
79
80 // Insert "42" at offset 8, but set cursor_position to offset 4 (the 'x')
81 // This tests that cursor moves to the predicted position, not the end of the edit
82 propose_edits_with_cursor_position(
83 &provider,
84 vec![(8..8, "42")],
85 Some(4), // cursor at offset 4 (the 'x'), NOT at the edit location
86 &mut cx,
87 );
88 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
89
90 assert_editor_active_edit_completion(&mut cx, |_, edits| {
91 assert_eq!(edits.len(), 1);
92 assert_eq!(edits[0].1.as_ref(), "42");
93 });
94
95 accept_completion(&mut cx);
96
97 // Cursor should be at offset 4 (the 'x'), not at the end of the inserted "42"
98 cx.assert_editor_state("let ˇx = 42;");
99}
100
101#[gpui::test]
102async fn test_edit_prediction_cursor_position_fallback(cx: &mut gpui::TestAppContext) {
103 init_test(cx, |_| {});
104
105 let mut cx = EditorTestContext::new(cx).await;
106 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
107 assign_editor_completion_provider(provider.clone(), &mut cx);
108 cx.set_state("let x = ˇ;");
109
110 // Propose an edit without a cursor position - should fall back to end of edit
111 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
112 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
113
114 accept_completion(&mut cx);
115
116 // Cursor should be at the end of the inserted text (default behavior)
117 cx.assert_editor_state("let x = 42ˇ;")
118}
119
120#[gpui::test]
121async fn test_edit_prediction_modification(cx: &mut gpui::TestAppContext) {
122 init_test(cx, |_| {});
123
124 let mut cx = EditorTestContext::new(cx).await;
125 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
126 assign_editor_completion_provider(provider.clone(), &mut cx);
127 cx.set_state("let pi = ˇ\"foo\";");
128
129 propose_edits(&provider, vec![(9..14, "3.14159")], &mut cx);
130 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
131
132 assert_editor_active_edit_completion(&mut cx, |_, edits| {
133 assert_eq!(edits.len(), 1);
134 assert_eq!(edits[0].1.as_ref(), "3.14159");
135 });
136
137 accept_completion(&mut cx);
138
139 cx.assert_editor_state("let pi = 3.14159ˇ;")
140}
141
142#[gpui::test]
143async fn test_edit_prediction_jump_button(cx: &mut gpui::TestAppContext) {
144 init_test(cx, |_| {});
145
146 let mut cx = EditorTestContext::new(cx).await;
147 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
148 assign_editor_completion_provider(provider.clone(), &mut cx);
149
150 // Cursor is 2+ lines above the proposed edit
151 cx.set_state(indoc! {"
152 line 0
153 line ˇ1
154 line 2
155 line 3
156 line
157 "});
158
159 propose_edits(
160 &provider,
161 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
162 &mut cx,
163 );
164
165 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
166 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
167 assert_eq!(move_target.to_point(&snapshot), Point::new(4, 3));
168 });
169
170 // When accepting, cursor is moved to the proposed location
171 accept_completion(&mut cx);
172 cx.assert_editor_state(indoc! {"
173 line 0
174 line 1
175 line 2
176 line 3
177 linˇe
178 "});
179
180 // Cursor is 2+ lines below the proposed edit
181 cx.set_state(indoc! {"
182 line 0
183 line
184 line 2
185 line 3
186 line ˇ4
187 "});
188
189 propose_edits(
190 &provider,
191 vec![(Point::new(1, 3)..Point::new(1, 3), " 1")],
192 &mut cx,
193 );
194
195 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
196 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
197 assert_eq!(move_target.to_point(&snapshot), Point::new(1, 3));
198 });
199
200 // When accepting, cursor is moved to the proposed location
201 accept_completion(&mut cx);
202 cx.assert_editor_state(indoc! {"
203 line 0
204 linˇe
205 line 2
206 line 3
207 line 4
208 "});
209}
210
211#[gpui::test]
212async fn test_edit_prediction_invalidation_range(cx: &mut gpui::TestAppContext) {
213 init_test(cx, |_| {});
214
215 let mut cx = EditorTestContext::new(cx).await;
216 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
217 assign_editor_completion_provider(provider.clone(), &mut cx);
218
219 // Cursor is 3+ lines above the proposed edit
220 cx.set_state(indoc! {"
221 line 0
222 line ˇ1
223 line 2
224 line 3
225 line 4
226 line
227 "});
228 let edit_location = Point::new(5, 3);
229
230 propose_edits(
231 &provider,
232 vec![(edit_location..edit_location, " 5")],
233 &mut cx,
234 );
235
236 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
237 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
238 assert_eq!(move_target.to_point(&snapshot), edit_location);
239 });
240
241 // If we move *towards* the completion, it stays active
242 cx.set_selections_state(indoc! {"
243 line 0
244 line 1
245 line ˇ2
246 line 3
247 line 4
248 line
249 "});
250 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
251 assert_eq!(move_target.to_point(&snapshot), edit_location);
252 });
253
254 // If we move *away* from the completion, it is discarded
255 cx.set_selections_state(indoc! {"
256 line ˇ0
257 line 1
258 line 2
259 line 3
260 line 4
261 line
262 "});
263 cx.editor(|editor, _, _| {
264 assert!(editor.active_edit_prediction.is_none());
265 });
266
267 // Cursor is 3+ lines below the proposed edit
268 cx.set_state(indoc! {"
269 line
270 line 1
271 line 2
272 line 3
273 line ˇ4
274 line 5
275 "});
276 let edit_location = Point::new(0, 3);
277
278 propose_edits(
279 &provider,
280 vec![(edit_location..edit_location, " 0")],
281 &mut cx,
282 );
283
284 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
285 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
286 assert_eq!(move_target.to_point(&snapshot), edit_location);
287 });
288
289 // If we move *towards* the completion, it stays active
290 cx.set_selections_state(indoc! {"
291 line
292 line 1
293 line 2
294 line ˇ3
295 line 4
296 line 5
297 "});
298 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
299 assert_eq!(move_target.to_point(&snapshot), edit_location);
300 });
301
302 // If we move *away* from the completion, it is discarded
303 cx.set_selections_state(indoc! {"
304 line
305 line 1
306 line 2
307 line 3
308 line 4
309 line ˇ5
310 "});
311 cx.editor(|editor, _, _| {
312 assert!(editor.active_edit_prediction.is_none());
313 });
314}
315
316#[gpui::test]
317async fn test_edit_prediction_jump_disabled_for_non_zed_providers(cx: &mut gpui::TestAppContext) {
318 init_test(cx, |_| {});
319
320 let mut cx = EditorTestContext::new(cx).await;
321 let provider = cx.new(|_| FakeNonZedEditPredictionDelegate::default());
322 assign_editor_completion_provider_non_zed(provider.clone(), &mut cx);
323
324 // Cursor is 2+ lines above the proposed edit
325 cx.set_state(indoc! {"
326 line 0
327 line ˇ1
328 line 2
329 line 3
330 line
331 "});
332
333 propose_edits_non_zed(
334 &provider,
335 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
336 &mut cx,
337 );
338
339 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
340
341 // For non-Zed providers, there should be no move completion (jump functionality disabled)
342 cx.editor(|editor, _, _| {
343 if let Some(completion_state) = &editor.active_edit_prediction {
344 // Should be an Edit prediction, not a Move prediction
345 match &completion_state.completion {
346 EditPrediction::Edit { .. } => {
347 // This is expected for non-Zed providers
348 }
349 EditPrediction::MoveWithin { .. } | EditPrediction::MoveOutside { .. } => {
350 panic!(
351 "Non-Zed providers should not show Move predictions (jump functionality)"
352 );
353 }
354 }
355 }
356 });
357}
358
359#[gpui::test]
360async fn test_edit_prediction_preview_cleanup_on_toggle_off(cx: &mut gpui::TestAppContext) {
361 init_test(cx, |_| {});
362
363 // Bind `ctrl-shift-a` to accept the provided edit prediction. The actual key
364 // binding here doesn't matter, we simply need to confirm that holding the
365 // binding's modifiers triggers the edit prediction preview.
366 cx.update(|cx| cx.bind_keys([KeyBinding::new("ctrl-shift-a", AcceptEditPrediction, None)]));
367
368 let mut cx = EditorTestContext::new(cx).await;
369 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
370 assign_editor_completion_provider(provider.clone(), &mut cx);
371 cx.set_state("let x = ˇ;");
372
373 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
374 cx.update_editor(|editor, window, cx| {
375 editor.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
376 editor.update_visible_edit_prediction(window, cx)
377 });
378
379 cx.editor(|editor, _, _| {
380 assert!(editor.has_active_edit_prediction());
381 });
382
383 // Simulate pressing the modifiers for `AcceptEditPrediction`, namely
384 // `ctrl-shift`, so that we can confirm that the edit prediction preview is
385 // activated.
386 let modifiers = Modifiers::control_shift();
387 cx.simulate_modifiers_change(modifiers);
388 cx.run_until_parked();
389
390 cx.editor(|editor, _, _| {
391 assert!(editor.edit_prediction_preview_is_active());
392 });
393
394 // Disable showing edit predictions without issuing a new modifiers changed
395 // event, to confirm that the edit prediction preview is still active.
396 cx.update_editor(|editor, window, cx| {
397 editor.set_show_edit_predictions(Some(false), window, cx);
398 });
399
400 cx.editor(|editor, _, _| {
401 assert!(!editor.has_active_edit_prediction());
402 assert!(editor.edit_prediction_preview_is_active());
403 });
404
405 // Now release the modifiers
406 // Simulate releasing all modifiers, ensuring that even with edit prediction
407 // disabled, the edit prediction preview is cleaned up.
408 cx.simulate_modifiers_change(Modifiers::none());
409 cx.run_until_parked();
410
411 cx.editor(|editor, _, _| {
412 assert!(!editor.edit_prediction_preview_is_active());
413 });
414}
415
416fn assert_editor_active_edit_completion(
417 cx: &mut EditorTestContext,
418 assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, Arc<str>)>),
419) {
420 cx.editor(|editor, _, cx| {
421 let completion_state = editor
422 .active_edit_prediction
423 .as_ref()
424 .expect("editor has no active completion");
425
426 if let EditPrediction::Edit { edits, .. } = &completion_state.completion {
427 assert(editor.buffer().read(cx).snapshot(cx), edits);
428 } else {
429 panic!("expected edit completion");
430 }
431 })
432}
433
434fn assert_editor_active_move_completion(
435 cx: &mut EditorTestContext,
436 assert: impl FnOnce(MultiBufferSnapshot, Anchor),
437) {
438 cx.editor(|editor, _, cx| {
439 let completion_state = editor
440 .active_edit_prediction
441 .as_ref()
442 .expect("editor has no active completion");
443
444 if let EditPrediction::MoveWithin { target, .. } = &completion_state.completion {
445 assert(editor.buffer().read(cx).snapshot(cx), *target);
446 } else {
447 panic!("expected move completion");
448 }
449 })
450}
451
452fn accept_completion(cx: &mut EditorTestContext) {
453 cx.update_editor(|editor, window, cx| {
454 editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)
455 })
456}
457
458fn propose_edits<T: ToOffset>(
459 provider: &Entity<FakeEditPredictionDelegate>,
460 edits: Vec<(Range<T>, &str)>,
461 cx: &mut EditorTestContext,
462) {
463 propose_edits_with_cursor_position(provider, edits, None, cx);
464}
465
466fn propose_edits_with_cursor_position<T: ToOffset>(
467 provider: &Entity<FakeEditPredictionDelegate>,
468 edits: Vec<(Range<T>, &str)>,
469 cursor_offset: Option<usize>,
470 cx: &mut EditorTestContext,
471) {
472 let snapshot = cx.buffer_snapshot();
473 let cursor_position = cursor_offset
474 .map(|offset| PredictedCursorPosition::at_anchor(snapshot.anchor_after(offset)));
475 let edits = edits.into_iter().map(|(range, text)| {
476 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
477 (range, text.into())
478 });
479
480 cx.update(|_, cx| {
481 provider.update(cx, |provider, _| {
482 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
483 id: None,
484 edits: edits.collect(),
485 cursor_position,
486 edit_preview: None,
487 }))
488 })
489 });
490}
491
492fn propose_edits_with_cursor_position_in_insertion<T: ToOffset>(
493 provider: &Entity<FakeEditPredictionDelegate>,
494 edits: Vec<(Range<T>, &str)>,
495 anchor_offset: usize,
496 offset_within_insertion: usize,
497 cx: &mut EditorTestContext,
498) {
499 let snapshot = cx.buffer_snapshot();
500 // Use anchor_before (left bias) so the anchor stays at the insertion point
501 // rather than moving past the inserted text
502 let cursor_position = Some(PredictedCursorPosition::new(
503 snapshot.anchor_before(anchor_offset),
504 offset_within_insertion,
505 ));
506 let edits = edits.into_iter().map(|(range, text)| {
507 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
508 (range, text.into())
509 });
510
511 cx.update(|_, cx| {
512 provider.update(cx, |provider, _| {
513 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
514 id: None,
515 edits: edits.collect(),
516 cursor_position,
517 edit_preview: None,
518 }))
519 })
520 });
521}
522
523fn assign_editor_completion_provider(
524 provider: Entity<FakeEditPredictionDelegate>,
525 cx: &mut EditorTestContext,
526) {
527 cx.update_editor(|editor, window, cx| {
528 editor.set_edit_prediction_provider(Some(provider), window, cx);
529 })
530}
531
532fn propose_edits_non_zed<T: ToOffset>(
533 provider: &Entity<FakeNonZedEditPredictionDelegate>,
534 edits: Vec<(Range<T>, &str)>,
535 cx: &mut EditorTestContext,
536) {
537 let snapshot = cx.buffer_snapshot();
538 let edits = edits.into_iter().map(|(range, text)| {
539 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
540 (range, text.into())
541 });
542
543 cx.update(|_, cx| {
544 provider.update(cx, |provider, _| {
545 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
546 id: None,
547 edits: edits.collect(),
548 cursor_position: None,
549 edit_preview: None,
550 }))
551 })
552 });
553}
554
555fn assign_editor_completion_provider_non_zed(
556 provider: Entity<FakeNonZedEditPredictionDelegate>,
557 cx: &mut EditorTestContext,
558) {
559 cx.update_editor(|editor, window, cx| {
560 editor.set_edit_prediction_provider(Some(provider), window, cx);
561 })
562}
563
564#[derive(Default, Clone)]
565pub struct FakeEditPredictionDelegate {
566 pub completion: Option<edit_prediction_types::EditPrediction>,
567}
568
569impl FakeEditPredictionDelegate {
570 pub fn set_edit_prediction(
571 &mut self,
572 completion: Option<edit_prediction_types::EditPrediction>,
573 ) {
574 self.completion = completion;
575 }
576}
577
578impl EditPredictionDelegate for FakeEditPredictionDelegate {
579 fn name() -> &'static str {
580 "fake-completion-provider"
581 }
582
583 fn display_name() -> &'static str {
584 "Fake Completion Provider"
585 }
586
587 fn show_predictions_in_menu() -> bool {
588 true
589 }
590
591 fn supports_jump_to_edit() -> bool {
592 true
593 }
594
595 fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
596 EditPredictionIconSet::new(IconName::ZedPredict)
597 }
598
599 fn is_enabled(
600 &self,
601 _buffer: &gpui::Entity<language::Buffer>,
602 _cursor_position: language::Anchor,
603 _cx: &gpui::App,
604 ) -> bool {
605 true
606 }
607
608 fn is_refreshing(&self, _cx: &gpui::App) -> bool {
609 false
610 }
611
612 fn refresh(
613 &mut self,
614 _buffer: gpui::Entity<language::Buffer>,
615 _cursor_position: language::Anchor,
616 _debounce: bool,
617 _cx: &mut gpui::Context<Self>,
618 ) {
619 }
620
621 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
622
623 fn discard(&mut self, _cx: &mut gpui::Context<Self>) {}
624
625 fn suggest<'a>(
626 &mut self,
627 _buffer: &gpui::Entity<language::Buffer>,
628 _cursor_position: language::Anchor,
629 _cx: &mut gpui::Context<Self>,
630 ) -> Option<edit_prediction_types::EditPrediction> {
631 self.completion.clone()
632 }
633}
634
635#[derive(Default, Clone)]
636pub struct FakeNonZedEditPredictionDelegate {
637 pub completion: Option<edit_prediction_types::EditPrediction>,
638}
639
640impl FakeNonZedEditPredictionDelegate {
641 pub fn set_edit_prediction(
642 &mut self,
643 completion: Option<edit_prediction_types::EditPrediction>,
644 ) {
645 self.completion = completion;
646 }
647}
648
649impl EditPredictionDelegate for FakeNonZedEditPredictionDelegate {
650 fn name() -> &'static str {
651 "fake-non-zed-provider"
652 }
653
654 fn display_name() -> &'static str {
655 "Fake Non-Zed Provider"
656 }
657
658 fn show_predictions_in_menu() -> bool {
659 false
660 }
661
662 fn supports_jump_to_edit() -> bool {
663 false
664 }
665
666 fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
667 EditPredictionIconSet::new(IconName::ZedPredict)
668 }
669
670 fn is_enabled(
671 &self,
672 _buffer: &gpui::Entity<language::Buffer>,
673 _cursor_position: language::Anchor,
674 _cx: &gpui::App,
675 ) -> bool {
676 true
677 }
678
679 fn is_refreshing(&self, _cx: &gpui::App) -> bool {
680 false
681 }
682
683 fn refresh(
684 &mut self,
685 _buffer: gpui::Entity<language::Buffer>,
686 _cursor_position: language::Anchor,
687 _debounce: bool,
688 _cx: &mut gpui::Context<Self>,
689 ) {
690 }
691
692 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
693
694 fn discard(&mut self, _cx: &mut gpui::Context<Self>) {}
695
696 fn suggest<'a>(
697 &mut self,
698 _buffer: &gpui::Entity<language::Buffer>,
699 _cursor_position: language::Anchor,
700 _cx: &mut gpui::Context<Self>,
701 ) -> Option<edit_prediction_types::EditPrediction> {
702 self.completion.clone()
703 }
704}