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