1use edit_prediction_types::{
2 EditPredictionDelegate, EditPredictionIconSet, PredictedCursorPosition,
3};
4use gpui::{Entity, KeyBinding, Modifiers, prelude::*};
5use indoc::indoc;
6use language::Buffer;
7use language::EditPredictionsMode;
8use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
9use std::{
10 ops::Range,
11 sync::{
12 Arc,
13 atomic::{self, AtomicUsize},
14 },
15};
16use text::{Point, ToOffset};
17use ui::prelude::*;
18
19use crate::{
20 AcceptEditPrediction, EditPrediction, EditPredictionKeybindAction,
21 EditPredictionKeybindSurface, MenuEditPredictionsPolicy,
22 editor_tests::{init_test, update_test_language_settings},
23 test::editor_test_context::EditorTestContext,
24};
25use rpc::proto::PeerId;
26use workspace::CollaboratorId;
27
28#[gpui::test]
29async fn test_edit_prediction_insert(cx: &mut gpui::TestAppContext) {
30 init_test(cx, |_| {});
31
32 let mut cx = EditorTestContext::new(cx).await;
33 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
34 assign_editor_completion_provider(provider.clone(), &mut cx);
35 cx.set_state("let absolute_zero_celsius = ˇ;");
36
37 propose_edits(&provider, vec![(28..28, "-273.15")], &mut cx);
38 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
39
40 assert_editor_active_edit_completion(&mut cx, |_, edits| {
41 assert_eq!(edits.len(), 1);
42 assert_eq!(edits[0].1.as_ref(), "-273.15");
43 });
44
45 accept_completion(&mut cx);
46
47 cx.assert_editor_state("let absolute_zero_celsius = -273.15ˇ;")
48}
49
50#[gpui::test]
51async fn test_edit_prediction_cursor_position_inside_insertion(cx: &mut gpui::TestAppContext) {
52 init_test(cx, |_| {
53 eprintln!("");
54 });
55
56 let mut cx = EditorTestContext::new(cx).await;
57 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
58
59 assign_editor_completion_provider(provider.clone(), &mut cx);
60 // Buffer: "fn foo() {}" - we'll insert text and position cursor inside the insertion
61 cx.set_state("fn foo() ˇ{}");
62
63 // Insert "bar()" at offset 9, with cursor at offset 2 within the insertion (after "ba")
64 // This tests the case where cursor is inside newly inserted text
65 propose_edits_with_cursor_position_in_insertion(
66 &provider,
67 vec![(9..9, "bar()")],
68 9, // anchor at the insertion point
69 2, // offset 2 within "bar()" puts cursor after "ba"
70 &mut cx,
71 );
72 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
73
74 assert_editor_active_edit_completion(&mut cx, |_, edits| {
75 assert_eq!(edits.len(), 1);
76 assert_eq!(edits[0].1.as_ref(), "bar()");
77 });
78
79 accept_completion(&mut cx);
80
81 // Cursor should be inside the inserted text at "baˇr()"
82 cx.assert_editor_state("fn foo() baˇr(){}");
83}
84
85#[gpui::test]
86async fn test_edit_prediction_cursor_position_outside_edit(cx: &mut gpui::TestAppContext) {
87 init_test(cx, |_| {});
88
89 let mut cx = EditorTestContext::new(cx).await;
90 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
91 assign_editor_completion_provider(provider.clone(), &mut cx);
92 // Buffer: "let x = ;" with cursor before semicolon - we'll insert "42" and position cursor elsewhere
93 cx.set_state("let x = ˇ;");
94
95 // Insert "42" at offset 8, but set cursor_position to offset 4 (the 'x')
96 // This tests that cursor moves to the predicted position, not the end of the edit
97 propose_edits_with_cursor_position(
98 &provider,
99 vec![(8..8, "42")],
100 Some(4), // cursor at offset 4 (the 'x'), NOT at the edit location
101 &mut cx,
102 );
103 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
104
105 assert_editor_active_edit_completion(&mut cx, |_, edits| {
106 assert_eq!(edits.len(), 1);
107 assert_eq!(edits[0].1.as_ref(), "42");
108 });
109
110 accept_completion(&mut cx);
111
112 // Cursor should be at offset 4 (the 'x'), not at the end of the inserted "42"
113 cx.assert_editor_state("let ˇx = 42;");
114}
115
116#[gpui::test]
117async fn test_edit_prediction_cursor_position_fallback(cx: &mut gpui::TestAppContext) {
118 init_test(cx, |_| {});
119
120 let mut cx = EditorTestContext::new(cx).await;
121 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
122 assign_editor_completion_provider(provider.clone(), &mut cx);
123 cx.set_state("let x = ˇ;");
124
125 // Propose an edit without a cursor position - should fall back to end of edit
126 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
127 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
128
129 accept_completion(&mut cx);
130
131 // Cursor should be at the end of the inserted text (default behavior)
132 cx.assert_editor_state("let x = 42ˇ;")
133}
134
135#[gpui::test]
136async fn test_edit_prediction_modification(cx: &mut gpui::TestAppContext) {
137 init_test(cx, |_| {});
138
139 let mut cx = EditorTestContext::new(cx).await;
140 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
141 assign_editor_completion_provider(provider.clone(), &mut cx);
142 cx.set_state("let pi = ˇ\"foo\";");
143
144 propose_edits(&provider, vec![(9..14, "3.14159")], &mut cx);
145 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
146
147 assert_editor_active_edit_completion(&mut cx, |_, edits| {
148 assert_eq!(edits.len(), 1);
149 assert_eq!(edits[0].1.as_ref(), "3.14159");
150 });
151
152 accept_completion(&mut cx);
153
154 cx.assert_editor_state("let pi = 3.14159ˇ;")
155}
156
157#[gpui::test]
158async fn test_edit_prediction_jump_button(cx: &mut gpui::TestAppContext) {
159 init_test(cx, |_| {});
160
161 let mut cx = EditorTestContext::new(cx).await;
162 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
163 assign_editor_completion_provider(provider.clone(), &mut cx);
164
165 // Cursor is 2+ lines above the proposed edit
166 cx.set_state(indoc! {"
167 line 0
168 line ˇ1
169 line 2
170 line 3
171 line
172 "});
173
174 propose_edits(
175 &provider,
176 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
177 &mut cx,
178 );
179
180 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
181 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
182 assert_eq!(move_target.to_point(&snapshot), Point::new(4, 3));
183 });
184
185 // When accepting, cursor is moved to the proposed location
186 accept_completion(&mut cx);
187 cx.assert_editor_state(indoc! {"
188 line 0
189 line 1
190 line 2
191 line 3
192 linˇe
193 "});
194
195 // Cursor is 2+ lines below the proposed edit
196 cx.set_state(indoc! {"
197 line 0
198 line
199 line 2
200 line 3
201 line ˇ4
202 "});
203
204 propose_edits(
205 &provider,
206 vec![(Point::new(1, 3)..Point::new(1, 3), " 1")],
207 &mut cx,
208 );
209
210 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
211 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
212 assert_eq!(move_target.to_point(&snapshot), Point::new(1, 3));
213 });
214
215 // When accepting, cursor is moved to the proposed location
216 accept_completion(&mut cx);
217 cx.assert_editor_state(indoc! {"
218 line 0
219 linˇe
220 line 2
221 line 3
222 line 4
223 "});
224}
225
226#[gpui::test]
227async fn test_edit_prediction_invalidation_range(cx: &mut gpui::TestAppContext) {
228 init_test(cx, |_| {});
229
230 let mut cx = EditorTestContext::new(cx).await;
231 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
232 assign_editor_completion_provider(provider.clone(), &mut cx);
233
234 // Cursor is 3+ lines above the proposed edit
235 cx.set_state(indoc! {"
236 line 0
237 line ˇ1
238 line 2
239 line 3
240 line 4
241 line
242 "});
243 let edit_location = Point::new(5, 3);
244
245 propose_edits(
246 &provider,
247 vec![(edit_location..edit_location, " 5")],
248 &mut cx,
249 );
250
251 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
252 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
253 assert_eq!(move_target.to_point(&snapshot), edit_location);
254 });
255
256 // If we move *towards* the completion, it stays active
257 cx.set_selections_state(indoc! {"
258 line 0
259 line 1
260 line ˇ2
261 line 3
262 line 4
263 line
264 "});
265 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
266 assert_eq!(move_target.to_point(&snapshot), edit_location);
267 });
268
269 // If we move *away* from the completion, it is discarded
270 cx.set_selections_state(indoc! {"
271 line ˇ0
272 line 1
273 line 2
274 line 3
275 line 4
276 line
277 "});
278 cx.editor(|editor, _, _| {
279 assert!(editor.active_edit_prediction.is_none());
280 });
281
282 // Cursor is 3+ lines below the proposed edit
283 cx.set_state(indoc! {"
284 line
285 line 1
286 line 2
287 line 3
288 line ˇ4
289 line 5
290 "});
291 let edit_location = Point::new(0, 3);
292
293 propose_edits(
294 &provider,
295 vec![(edit_location..edit_location, " 0")],
296 &mut cx,
297 );
298
299 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
300 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
301 assert_eq!(move_target.to_point(&snapshot), edit_location);
302 });
303
304 // If we move *towards* the completion, it stays active
305 cx.set_selections_state(indoc! {"
306 line
307 line 1
308 line 2
309 line ˇ3
310 line 4
311 line 5
312 "});
313 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
314 assert_eq!(move_target.to_point(&snapshot), edit_location);
315 });
316
317 // If we move *away* from the completion, it is discarded
318 cx.set_selections_state(indoc! {"
319 line
320 line 1
321 line 2
322 line 3
323 line 4
324 line ˇ5
325 "});
326 cx.editor(|editor, _, _| {
327 assert!(editor.active_edit_prediction.is_none());
328 });
329}
330
331#[gpui::test]
332async fn test_edit_prediction_jump_disabled_for_non_zed_providers(cx: &mut gpui::TestAppContext) {
333 init_test(cx, |_| {});
334
335 let mut cx = EditorTestContext::new(cx).await;
336 let provider = cx.new(|_| FakeNonZedEditPredictionDelegate::default());
337 assign_editor_completion_provider_non_zed(provider.clone(), &mut cx);
338
339 // Cursor is 2+ lines above the proposed edit
340 cx.set_state(indoc! {"
341 line 0
342 line ˇ1
343 line 2
344 line 3
345 line
346 "});
347
348 propose_edits_non_zed(
349 &provider,
350 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
351 &mut cx,
352 );
353
354 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
355
356 // For non-Zed providers, there should be no move completion (jump functionality disabled)
357 cx.editor(|editor, _, _| {
358 if let Some(completion_state) = &editor.active_edit_prediction {
359 // Should be an Edit prediction, not a Move prediction
360 match &completion_state.completion {
361 EditPrediction::Edit { .. } => {
362 // This is expected for non-Zed providers
363 }
364 EditPrediction::MoveWithin { .. } | EditPrediction::MoveOutside { .. } => {
365 panic!(
366 "Non-Zed providers should not show Move predictions (jump functionality)"
367 );
368 }
369 }
370 }
371 });
372}
373
374#[gpui::test]
375async fn test_edit_prediction_refresh_suppressed_while_following(cx: &mut gpui::TestAppContext) {
376 init_test(cx, |_| {});
377
378 let mut cx = EditorTestContext::new(cx).await;
379 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
380 assign_editor_completion_provider(provider.clone(), &mut cx);
381 cx.set_state("let x = ˇ;");
382
383 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
384
385 cx.update_editor(|editor, window, cx| {
386 editor.refresh_edit_prediction(false, false, window, cx);
387 editor.update_visible_edit_prediction(window, cx);
388 });
389
390 assert_eq!(
391 provider.read_with(&cx.cx, |provider, _| {
392 provider.refresh_count.load(atomic::Ordering::SeqCst)
393 }),
394 1
395 );
396 cx.editor(|editor, _, _| {
397 assert!(editor.active_edit_prediction.is_some());
398 });
399
400 cx.update_editor(|editor, window, cx| {
401 editor.leader_id = Some(CollaboratorId::PeerId(PeerId::default()));
402 editor.refresh_edit_prediction(false, false, window, cx);
403 });
404
405 assert_eq!(
406 provider.read_with(&cx.cx, |provider, _| {
407 provider.refresh_count.load(atomic::Ordering::SeqCst)
408 }),
409 1
410 );
411 cx.editor(|editor, _, _| {
412 assert!(editor.active_edit_prediction.is_none());
413 });
414
415 cx.update_editor(|editor, window, cx| {
416 editor.leader_id = None;
417 editor.refresh_edit_prediction(false, false, window, cx);
418 });
419
420 assert_eq!(
421 provider.read_with(&cx.cx, |provider, _| {
422 provider.refresh_count.load(atomic::Ordering::SeqCst)
423 }),
424 2
425 );
426}
427
428#[gpui::test]
429async fn test_edit_prediction_preview_cleanup_on_toggle_off(cx: &mut gpui::TestAppContext) {
430 init_test(cx, |_| {});
431
432 // Bind `ctrl-shift-a` to accept the provided edit prediction. The actual key
433 // binding here doesn't matter, we simply need to confirm that holding the
434 // binding's modifiers triggers the edit prediction preview.
435 cx.update(|cx| cx.bind_keys([KeyBinding::new("ctrl-shift-a", AcceptEditPrediction, None)]));
436
437 let mut cx = EditorTestContext::new(cx).await;
438 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
439 assign_editor_completion_provider(provider.clone(), &mut cx);
440 cx.set_state("let x = ˇ;");
441
442 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
443 cx.update_editor(|editor, window, cx| {
444 editor.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
445 editor.update_visible_edit_prediction(window, cx)
446 });
447
448 cx.editor(|editor, _, _| {
449 assert!(editor.has_active_edit_prediction());
450 });
451
452 // Simulate pressing the modifiers for `AcceptEditPrediction`, namely
453 // `ctrl-shift`, so that we can confirm that the edit prediction preview is
454 // activated.
455 let modifiers = Modifiers::control_shift();
456 cx.simulate_modifiers_change(modifiers);
457 cx.run_until_parked();
458
459 cx.editor(|editor, _, _| {
460 assert!(editor.edit_prediction_preview_is_active());
461 });
462
463 // Disable showing edit predictions without issuing a new modifiers changed
464 // event, to confirm that the edit prediction preview is still active.
465 cx.update_editor(|editor, window, cx| {
466 editor.set_show_edit_predictions(Some(false), window, cx);
467 });
468
469 cx.editor(|editor, _, _| {
470 assert!(!editor.has_active_edit_prediction());
471 assert!(editor.edit_prediction_preview_is_active());
472 });
473
474 // Now release the modifiers
475 // Simulate releasing all modifiers, ensuring that even with edit prediction
476 // disabled, the edit prediction preview is cleaned up.
477 cx.simulate_modifiers_change(Modifiers::none());
478 cx.run_until_parked();
479
480 cx.editor(|editor, _, _| {
481 assert!(!editor.edit_prediction_preview_is_active());
482 });
483}
484
485fn load_default_keymap(cx: &mut gpui::TestAppContext) {
486 cx.update(|cx| {
487 cx.bind_keys(
488 settings::KeymapFile::load_asset_allow_partial_failure(
489 settings::DEFAULT_KEYMAP_PATH,
490 cx,
491 )
492 .expect("failed to load default keymap"),
493 );
494 });
495}
496
497#[gpui::test]
498async fn test_tab_is_preferred_accept_binding_over_alt_tab(cx: &mut gpui::TestAppContext) {
499 init_test(cx, |_| {});
500 load_default_keymap(cx);
501
502 let mut cx = EditorTestContext::new(cx).await;
503 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
504 assign_editor_completion_provider(provider.clone(), &mut cx);
505 cx.set_state("let x = ˇ;");
506
507 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
508 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
509
510 cx.update_editor(|editor, window, cx| {
511 assert!(editor.has_active_edit_prediction());
512 let keybind_display = editor.edit_prediction_keybind_display(
513 EditPredictionKeybindSurface::Inline,
514 window,
515 cx,
516 );
517 let keystroke = keybind_display
518 .accept_keystroke
519 .as_ref()
520 .expect("should have an accept binding");
521 assert!(
522 !keystroke.modifiers().modified(),
523 "preferred accept binding should be unmodified (tab), got modifiers: {:?}",
524 keystroke.modifiers()
525 );
526 assert_eq!(
527 keystroke.key(),
528 "tab",
529 "preferred accept binding should be tab"
530 );
531 });
532}
533
534#[gpui::test]
535async fn test_subtle_in_code_indicator_prefers_preview_binding(cx: &mut gpui::TestAppContext) {
536 init_test(cx, |_| {});
537 load_default_keymap(cx);
538 update_test_language_settings(cx, &|settings| {
539 settings.edit_predictions.get_or_insert_default().mode = Some(EditPredictionsMode::Subtle);
540 });
541
542 let mut cx = EditorTestContext::new(cx).await;
543 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
544 assign_editor_completion_provider(provider.clone(), &mut cx);
545 cx.set_state("let x = ˇ;");
546
547 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
548 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
549
550 cx.update_editor(|editor, window, cx| {
551 assert!(editor.has_active_edit_prediction());
552 assert!(
553 editor.edit_prediction_requires_modifier(),
554 "subtle mode should require a modifier"
555 );
556
557 let inline_keybind_display = editor.edit_prediction_keybind_display(
558 EditPredictionKeybindSurface::Inline,
559 window,
560 cx,
561 );
562 let compact_keybind_display = editor.edit_prediction_keybind_display(
563 EditPredictionKeybindSurface::CursorPopoverCompact,
564 window,
565 cx,
566 );
567
568 let accept_keystroke = inline_keybind_display
569 .accept_keystroke
570 .as_ref()
571 .expect("should have an accept binding");
572 let preview_keystroke = inline_keybind_display
573 .preview_keystroke
574 .as_ref()
575 .expect("should have a preview binding");
576 let in_code_keystroke = inline_keybind_display
577 .displayed_keystroke
578 .as_ref()
579 .expect("should have an in-code binding");
580 let compact_cursor_popover_keystroke = compact_keybind_display
581 .displayed_keystroke
582 .as_ref()
583 .expect("should have a compact cursor popover binding");
584
585 assert_eq!(accept_keystroke.key(), "tab");
586 assert!(
587 !editor.has_visible_completions_menu(),
588 "compact cursor-popover branch should be used without a completions menu"
589 );
590 assert!(
591 preview_keystroke.modifiers().modified(),
592 "preview binding should use modifiers in subtle mode"
593 );
594 assert_eq!(
595 compact_cursor_popover_keystroke.key(),
596 preview_keystroke.key(),
597 "subtle compact cursor popover should prefer the preview binding"
598 );
599 assert_eq!(
600 compact_cursor_popover_keystroke.modifiers(),
601 preview_keystroke.modifiers(),
602 "subtle compact cursor popover should use the preview binding modifiers"
603 );
604 assert_eq!(
605 in_code_keystroke.key(),
606 preview_keystroke.key(),
607 "subtle in-code indicator should prefer the preview binding"
608 );
609 assert_eq!(
610 in_code_keystroke.modifiers(),
611 preview_keystroke.modifiers(),
612 "subtle in-code indicator should use the preview binding modifiers"
613 );
614 });
615}
616
617#[gpui::test]
618async fn test_tab_accepts_edit_prediction_over_completion(cx: &mut gpui::TestAppContext) {
619 init_test(cx, |_| {});
620 load_default_keymap(cx);
621
622 let mut cx = EditorTestContext::new(cx).await;
623 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
624 assign_editor_completion_provider(provider.clone(), &mut cx);
625 cx.set_state("let x = ˇ;");
626
627 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
628 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
629
630 assert_editor_active_edit_completion(&mut cx, |_, edits| {
631 assert_eq!(edits.len(), 1);
632 assert_eq!(edits[0].1.as_ref(), "42");
633 });
634
635 cx.simulate_keystroke("tab");
636 cx.run_until_parked();
637
638 cx.assert_editor_state("let x = 42ˇ;");
639}
640
641#[gpui::test]
642async fn test_single_line_prediction_uses_accept_cursor_popover_action(
643 cx: &mut gpui::TestAppContext,
644) {
645 init_test(cx, |_| {});
646 load_default_keymap(cx);
647
648 let mut cx = EditorTestContext::new(cx).await;
649 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
650 assign_editor_completion_provider(provider.clone(), &mut cx);
651 cx.set_state("let x = ˇ;");
652
653 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
654 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
655
656 cx.update_editor(|editor, window, cx| {
657 assert!(editor.has_active_edit_prediction());
658
659 let keybind_display = editor.edit_prediction_keybind_display(
660 EditPredictionKeybindSurface::CursorPopoverExpanded,
661 window,
662 cx,
663 );
664
665 let accept_keystroke = keybind_display
666 .accept_keystroke
667 .as_ref()
668 .expect("should have an accept binding");
669 let preview_keystroke = keybind_display
670 .preview_keystroke
671 .as_ref()
672 .expect("should have a preview binding");
673
674 assert_eq!(
675 keybind_display.action,
676 EditPredictionKeybindAction::Accept,
677 "single-line prediction should show the accept action"
678 );
679 assert_eq!(accept_keystroke.key(), "tab");
680 assert!(preview_keystroke.modifiers().modified());
681 });
682}
683
684#[gpui::test]
685async fn test_multi_line_prediction_uses_preview_cursor_popover_action(
686 cx: &mut gpui::TestAppContext,
687) {
688 init_test(cx, |_| {});
689 load_default_keymap(cx);
690
691 let mut cx = EditorTestContext::new(cx).await;
692 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
693 assign_editor_completion_provider(provider.clone(), &mut cx);
694 cx.set_state("let x = ˇ;");
695
696 propose_edits(&provider, vec![(8..8, "42\n43")], &mut cx);
697 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
698
699 cx.update_editor(|editor, window, cx| {
700 assert!(editor.has_active_edit_prediction());
701
702 let keybind_display = editor.edit_prediction_keybind_display(
703 EditPredictionKeybindSurface::CursorPopoverExpanded,
704 window,
705 cx,
706 );
707 let preview_keystroke = keybind_display
708 .preview_keystroke
709 .as_ref()
710 .expect("should have a preview binding");
711
712 assert_eq!(
713 keybind_display.action,
714 EditPredictionKeybindAction::Preview,
715 "multi-line prediction should show the preview action"
716 );
717 assert!(preview_keystroke.modifiers().modified());
718 });
719}
720
721#[gpui::test]
722async fn test_single_line_prediction_with_preview_uses_accept_cursor_popover_action(
723 cx: &mut gpui::TestAppContext,
724) {
725 init_test(cx, |_| {});
726 load_default_keymap(cx);
727
728 let mut cx = EditorTestContext::new(cx).await;
729 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
730 assign_editor_completion_provider(provider.clone(), &mut cx);
731 cx.set_state("let x = ˇ;");
732
733 propose_edits_with_preview(&provider, vec![(8..8, "42")], &mut cx).await;
734 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
735
736 cx.update_editor(|editor, window, cx| {
737 assert!(editor.has_active_edit_prediction());
738
739 let keybind_display = editor.edit_prediction_keybind_display(
740 EditPredictionKeybindSurface::CursorPopoverExpanded,
741 window,
742 cx,
743 );
744
745 let accept_keystroke = keybind_display
746 .accept_keystroke
747 .as_ref()
748 .expect("should have an accept binding");
749 let preview_keystroke = keybind_display
750 .preview_keystroke
751 .as_ref()
752 .expect("should have a preview binding");
753
754 assert_eq!(
755 keybind_display.action,
756 EditPredictionKeybindAction::Accept,
757 "single-line prediction should show the accept action even with edit_preview"
758 );
759 assert_eq!(accept_keystroke.key(), "tab");
760 assert!(preview_keystroke.modifiers().modified());
761 });
762}
763
764#[gpui::test]
765async fn test_multi_line_prediction_with_preview_uses_preview_cursor_popover_action(
766 cx: &mut gpui::TestAppContext,
767) {
768 init_test(cx, |_| {});
769 load_default_keymap(cx);
770
771 let mut cx = EditorTestContext::new(cx).await;
772 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
773 assign_editor_completion_provider(provider.clone(), &mut cx);
774 cx.set_state("let x = ˇ;");
775
776 propose_edits_with_preview(&provider, vec![(8..8, "42\n43")], &mut cx).await;
777 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
778 cx.update_editor(|editor, window, cx| {
779 assert!(editor.has_active_edit_prediction());
780
781 let keybind_display = editor.edit_prediction_keybind_display(
782 EditPredictionKeybindSurface::CursorPopoverExpanded,
783 window,
784 cx,
785 );
786 let preview_keystroke = keybind_display
787 .preview_keystroke
788 .as_ref()
789 .expect("should have a preview binding");
790
791 assert_eq!(
792 keybind_display.action,
793 EditPredictionKeybindAction::Preview,
794 "multi-line prediction should show the preview action with edit_preview"
795 );
796 assert!(preview_keystroke.modifiers().modified());
797 });
798}
799
800#[gpui::test]
801async fn test_single_line_deletion_of_newline_uses_accept_cursor_popover_action(
802 cx: &mut gpui::TestAppContext,
803) {
804 init_test(cx, |_| {});
805 load_default_keymap(cx);
806
807 let mut cx = EditorTestContext::new(cx).await;
808 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
809 assign_editor_completion_provider(provider.clone(), &mut cx);
810 cx.set_state(indoc! {"
811 fn main() {
812 let value = 1;
813 ˇprintln!(\"done\");
814 }
815 "});
816
817 propose_edits(
818 &provider,
819 vec![(Point::new(1, 18)..Point::new(2, 17), "")],
820 &mut cx,
821 );
822 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
823
824 cx.update_editor(|editor, window, cx| {
825 assert!(editor.has_active_edit_prediction());
826
827 let keybind_display = editor.edit_prediction_keybind_display(
828 EditPredictionKeybindSurface::CursorPopoverExpanded,
829 window,
830 cx,
831 );
832
833 let accept_keystroke = keybind_display
834 .accept_keystroke
835 .as_ref()
836 .expect("should have an accept binding");
837 let preview_keystroke = keybind_display
838 .preview_keystroke
839 .as_ref()
840 .expect("should have a preview binding");
841
842 assert_eq!(
843 keybind_display.action,
844 EditPredictionKeybindAction::Accept,
845 "deleting one newline plus adjacent text should show the accept action"
846 );
847 assert_eq!(accept_keystroke.key(), "tab");
848 assert!(preview_keystroke.modifiers().modified());
849 });
850}
851
852#[gpui::test]
853async fn test_stale_single_line_prediction_does_not_force_preview_cursor_popover_action(
854 cx: &mut gpui::TestAppContext,
855) {
856 init_test(cx, |_| {});
857 load_default_keymap(cx);
858
859 let mut cx = EditorTestContext::new(cx).await;
860 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
861 assign_editor_completion_provider(provider.clone(), &mut cx);
862 cx.set_state("let x = ˇ;");
863
864 propose_edits(&provider, vec![(8..8, "42\n43")], &mut cx);
865 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
866 cx.update_editor(|editor, _window, cx| {
867 assert!(editor.active_edit_prediction.is_some());
868 assert!(editor.stale_edit_prediction_in_menu.is_none());
869 editor.take_active_edit_prediction(cx);
870 assert!(editor.active_edit_prediction.is_none());
871 assert!(editor.stale_edit_prediction_in_menu.is_some());
872 });
873
874 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
875 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
876
877 cx.update_editor(|editor, window, cx| {
878 assert!(editor.has_active_edit_prediction());
879
880 let keybind_display = editor.edit_prediction_keybind_display(
881 EditPredictionKeybindSurface::CursorPopoverExpanded,
882 window,
883 cx,
884 );
885 let accept_keystroke = keybind_display
886 .accept_keystroke
887 .as_ref()
888 .expect("should have an accept binding");
889
890 assert_eq!(
891 keybind_display.action,
892 EditPredictionKeybindAction::Accept,
893 "single-line active prediction should show the accept action"
894 );
895 assert!(
896 editor.stale_edit_prediction_in_menu.is_none(),
897 "refreshing the visible prediction should clear stale menu state"
898 );
899 assert_eq!(accept_keystroke.key(), "tab");
900 });
901}
902
903fn assert_editor_active_edit_completion(
904 cx: &mut EditorTestContext,
905 assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, Arc<str>)>),
906) {
907 cx.editor(|editor, _, cx| {
908 let completion_state = editor
909 .active_edit_prediction
910 .as_ref()
911 .expect("editor has no active completion");
912
913 if let EditPrediction::Edit { edits, .. } = &completion_state.completion {
914 assert(editor.buffer().read(cx).snapshot(cx), edits);
915 } else {
916 panic!("expected edit completion");
917 }
918 })
919}
920
921fn assert_editor_active_move_completion(
922 cx: &mut EditorTestContext,
923 assert: impl FnOnce(MultiBufferSnapshot, Anchor),
924) {
925 cx.editor(|editor, _, cx| {
926 let completion_state = editor
927 .active_edit_prediction
928 .as_ref()
929 .expect("editor has no active completion");
930
931 if let EditPrediction::MoveWithin { target, .. } = &completion_state.completion {
932 assert(editor.buffer().read(cx).snapshot(cx), *target);
933 } else {
934 panic!("expected move completion");
935 }
936 })
937}
938
939fn accept_completion(cx: &mut EditorTestContext) {
940 cx.update_editor(|editor, window, cx| {
941 editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)
942 })
943}
944
945fn propose_edits<T: ToOffset>(
946 provider: &Entity<FakeEditPredictionDelegate>,
947 edits: Vec<(Range<T>, &str)>,
948 cx: &mut EditorTestContext,
949) {
950 propose_edits_with_cursor_position(provider, edits, None, cx);
951}
952
953async fn propose_edits_with_preview<T: ToOffset + Clone>(
954 provider: &Entity<FakeEditPredictionDelegate>,
955 edits: Vec<(Range<T>, &str)>,
956 cx: &mut EditorTestContext,
957) {
958 let snapshot = cx.buffer_snapshot();
959 let edits = edits
960 .into_iter()
961 .map(|(range, text)| {
962 let anchor_range =
963 snapshot.anchor_after(range.start.clone())..snapshot.anchor_before(range.end);
964 (anchor_range, Arc::<str>::from(text))
965 })
966 .collect::<Vec<_>>();
967
968 let preview_edits = edits
969 .iter()
970 .map(|(range, text)| (range.clone(), text.clone()))
971 .collect::<Arc<[_]>>();
972
973 let edit_preview = cx
974 .buffer(|buffer: &Buffer, app| buffer.preview_edits(preview_edits, app))
975 .await;
976
977 let provider_edits = edits.into_iter().collect();
978
979 cx.update(|_, cx| {
980 provider.update(cx, |provider, _| {
981 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
982 id: None,
983 edits: provider_edits,
984 cursor_position: None,
985 edit_preview: Some(edit_preview),
986 }))
987 })
988 });
989}
990
991fn propose_edits_with_cursor_position<T: ToOffset>(
992 provider: &Entity<FakeEditPredictionDelegate>,
993 edits: Vec<(Range<T>, &str)>,
994 cursor_offset: Option<usize>,
995 cx: &mut EditorTestContext,
996) {
997 let snapshot = cx.buffer_snapshot();
998 let cursor_position = cursor_offset
999 .map(|offset| PredictedCursorPosition::at_anchor(snapshot.anchor_after(offset)));
1000 let edits = edits.into_iter().map(|(range, text)| {
1001 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
1002 (range, text.into())
1003 });
1004
1005 cx.update(|_, cx| {
1006 provider.update(cx, |provider, _| {
1007 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
1008 id: None,
1009 edits: edits.collect(),
1010 cursor_position,
1011 edit_preview: None,
1012 }))
1013 })
1014 });
1015}
1016
1017fn propose_edits_with_cursor_position_in_insertion<T: ToOffset>(
1018 provider: &Entity<FakeEditPredictionDelegate>,
1019 edits: Vec<(Range<T>, &str)>,
1020 anchor_offset: usize,
1021 offset_within_insertion: usize,
1022 cx: &mut EditorTestContext,
1023) {
1024 let snapshot = cx.buffer_snapshot();
1025 // Use anchor_before (left bias) so the anchor stays at the insertion point
1026 // rather than moving past the inserted text
1027 let cursor_position = Some(PredictedCursorPosition::new(
1028 snapshot.anchor_before(anchor_offset),
1029 offset_within_insertion,
1030 ));
1031 let edits = edits.into_iter().map(|(range, text)| {
1032 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
1033 (range, text.into())
1034 });
1035
1036 cx.update(|_, cx| {
1037 provider.update(cx, |provider, _| {
1038 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
1039 id: None,
1040 edits: edits.collect(),
1041 cursor_position,
1042 edit_preview: None,
1043 }))
1044 })
1045 });
1046}
1047
1048fn assign_editor_completion_provider(
1049 provider: Entity<FakeEditPredictionDelegate>,
1050 cx: &mut EditorTestContext,
1051) {
1052 cx.update_editor(|editor, window, cx| {
1053 editor.set_edit_prediction_provider(Some(provider), window, cx);
1054 })
1055}
1056
1057fn propose_edits_non_zed<T: ToOffset>(
1058 provider: &Entity<FakeNonZedEditPredictionDelegate>,
1059 edits: Vec<(Range<T>, &str)>,
1060 cx: &mut EditorTestContext,
1061) {
1062 let snapshot = cx.buffer_snapshot();
1063 let edits = edits.into_iter().map(|(range, text)| {
1064 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
1065 (range, text.into())
1066 });
1067
1068 cx.update(|_, cx| {
1069 provider.update(cx, |provider, _| {
1070 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
1071 id: None,
1072 edits: edits.collect(),
1073 cursor_position: None,
1074 edit_preview: None,
1075 }))
1076 })
1077 });
1078}
1079
1080fn assign_editor_completion_provider_non_zed(
1081 provider: Entity<FakeNonZedEditPredictionDelegate>,
1082 cx: &mut EditorTestContext,
1083) {
1084 cx.update_editor(|editor, window, cx| {
1085 editor.set_edit_prediction_provider(Some(provider), window, cx);
1086 })
1087}
1088
1089#[derive(Default, Clone)]
1090pub struct FakeEditPredictionDelegate {
1091 pub completion: Option<edit_prediction_types::EditPrediction>,
1092 pub refresh_count: Arc<AtomicUsize>,
1093}
1094
1095impl FakeEditPredictionDelegate {
1096 pub fn set_edit_prediction(
1097 &mut self,
1098 completion: Option<edit_prediction_types::EditPrediction>,
1099 ) {
1100 self.completion = completion;
1101 }
1102}
1103
1104impl EditPredictionDelegate for FakeEditPredictionDelegate {
1105 fn name() -> &'static str {
1106 "fake-completion-provider"
1107 }
1108
1109 fn display_name() -> &'static str {
1110 "Fake Completion Provider"
1111 }
1112
1113 fn show_predictions_in_menu() -> bool {
1114 true
1115 }
1116
1117 fn supports_jump_to_edit() -> bool {
1118 true
1119 }
1120
1121 fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
1122 EditPredictionIconSet::new(IconName::ZedPredict)
1123 }
1124
1125 fn is_enabled(
1126 &self,
1127 _buffer: &gpui::Entity<language::Buffer>,
1128 _cursor_position: language::Anchor,
1129 _cx: &gpui::App,
1130 ) -> bool {
1131 true
1132 }
1133
1134 fn is_refreshing(&self, _cx: &gpui::App) -> bool {
1135 false
1136 }
1137
1138 fn refresh(
1139 &mut self,
1140 _buffer: gpui::Entity<language::Buffer>,
1141 _cursor_position: language::Anchor,
1142 _debounce: bool,
1143 _cx: &mut gpui::Context<Self>,
1144 ) {
1145 self.refresh_count.fetch_add(1, atomic::Ordering::SeqCst);
1146 }
1147
1148 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
1149
1150 fn discard(
1151 &mut self,
1152 _reason: edit_prediction_types::EditPredictionDiscardReason,
1153 _cx: &mut gpui::Context<Self>,
1154 ) {
1155 }
1156
1157 fn suggest<'a>(
1158 &mut self,
1159 _buffer: &gpui::Entity<language::Buffer>,
1160 _cursor_position: language::Anchor,
1161 _cx: &mut gpui::Context<Self>,
1162 ) -> Option<edit_prediction_types::EditPrediction> {
1163 self.completion.clone()
1164 }
1165}
1166
1167#[derive(Default, Clone)]
1168pub struct FakeNonZedEditPredictionDelegate {
1169 pub completion: Option<edit_prediction_types::EditPrediction>,
1170}
1171
1172impl FakeNonZedEditPredictionDelegate {
1173 pub fn set_edit_prediction(
1174 &mut self,
1175 completion: Option<edit_prediction_types::EditPrediction>,
1176 ) {
1177 self.completion = completion;
1178 }
1179}
1180
1181impl EditPredictionDelegate for FakeNonZedEditPredictionDelegate {
1182 fn name() -> &'static str {
1183 "fake-non-zed-provider"
1184 }
1185
1186 fn display_name() -> &'static str {
1187 "Fake Non-Zed Provider"
1188 }
1189
1190 fn show_predictions_in_menu() -> bool {
1191 false
1192 }
1193
1194 fn supports_jump_to_edit() -> bool {
1195 false
1196 }
1197
1198 fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
1199 EditPredictionIconSet::new(IconName::ZedPredict)
1200 }
1201
1202 fn is_enabled(
1203 &self,
1204 _buffer: &gpui::Entity<language::Buffer>,
1205 _cursor_position: language::Anchor,
1206 _cx: &gpui::App,
1207 ) -> bool {
1208 true
1209 }
1210
1211 fn is_refreshing(&self, _cx: &gpui::App) -> bool {
1212 false
1213 }
1214
1215 fn refresh(
1216 &mut self,
1217 _buffer: gpui::Entity<language::Buffer>,
1218 _cursor_position: language::Anchor,
1219 _debounce: bool,
1220 _cx: &mut gpui::Context<Self>,
1221 ) {
1222 }
1223
1224 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
1225
1226 fn discard(
1227 &mut self,
1228 _reason: edit_prediction_types::EditPredictionDiscardReason,
1229 _cx: &mut gpui::Context<Self>,
1230 ) {
1231 }
1232
1233 fn suggest<'a>(
1234 &mut self,
1235 _buffer: &gpui::Entity<language::Buffer>,
1236 _cursor_position: language::Anchor,
1237 _cx: &mut gpui::Context<Self>,
1238 ) -> Option<edit_prediction_types::EditPrediction> {
1239 self.completion.clone()
1240 }
1241}