1use edit_prediction::EditPredictionProvider;
2use gpui::{Entity, prelude::*};
3use indoc::indoc;
4use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
5use project::Project;
6use std::ops::Range;
7use text::{Point, ToOffset};
8
9use crate::{
10 EditPrediction,
11 editor_tests::{init_test, update_test_language_settings},
12 test::editor_test_context::EditorTestContext,
13};
14
15#[gpui::test]
16async fn test_edit_prediction_insert(cx: &mut gpui::TestAppContext) {
17 init_test(cx, |_| {});
18
19 let mut cx = EditorTestContext::new(cx).await;
20 let provider = cx.new(|_| FakeEditPredictionProvider::default());
21 assign_editor_completion_provider(provider.clone(), &mut cx);
22 cx.set_state("let absolute_zero_celsius = ˇ;");
23
24 propose_edits(&provider, vec![(28..28, "-273.15")], &mut cx);
25 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
26
27 assert_editor_active_edit_completion(&mut cx, |_, edits| {
28 assert_eq!(edits.len(), 1);
29 assert_eq!(edits[0].1.as_str(), "-273.15");
30 });
31
32 accept_completion(&mut cx);
33
34 cx.assert_editor_state("let absolute_zero_celsius = -273.15ˇ;")
35}
36
37#[gpui::test]
38async fn test_edit_prediction_modification(cx: &mut gpui::TestAppContext) {
39 init_test(cx, |_| {});
40
41 let mut cx = EditorTestContext::new(cx).await;
42 let provider = cx.new(|_| FakeEditPredictionProvider::default());
43 assign_editor_completion_provider(provider.clone(), &mut cx);
44 cx.set_state("let pi = ˇ\"foo\";");
45
46 propose_edits(&provider, vec![(9..14, "3.14159")], &mut cx);
47 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
48
49 assert_editor_active_edit_completion(&mut cx, |_, edits| {
50 assert_eq!(edits.len(), 1);
51 assert_eq!(edits[0].1.as_str(), "3.14159");
52 });
53
54 accept_completion(&mut cx);
55
56 cx.assert_editor_state("let pi = 3.14159ˇ;")
57}
58
59#[gpui::test]
60async fn test_edit_prediction_jump_button(cx: &mut gpui::TestAppContext) {
61 init_test(cx, |_| {});
62
63 let mut cx = EditorTestContext::new(cx).await;
64 let provider = cx.new(|_| FakeEditPredictionProvider::default());
65 assign_editor_completion_provider(provider.clone(), &mut cx);
66
67 // Cursor is 2+ lines above the proposed edit
68 cx.set_state(indoc! {"
69 line 0
70 line ˇ1
71 line 2
72 line 3
73 line
74 "});
75
76 propose_edits(
77 &provider,
78 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
79 &mut cx,
80 );
81
82 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
83 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
84 assert_eq!(move_target.to_point(&snapshot), Point::new(4, 3));
85 });
86
87 // When accepting, cursor is moved to the proposed location
88 accept_completion(&mut cx);
89 cx.assert_editor_state(indoc! {"
90 line 0
91 line 1
92 line 2
93 line 3
94 linˇe
95 "});
96
97 // Cursor is 2+ lines below the proposed edit
98 cx.set_state(indoc! {"
99 line 0
100 line
101 line 2
102 line 3
103 line ˇ4
104 "});
105
106 propose_edits(
107 &provider,
108 vec![(Point::new(1, 3)..Point::new(1, 3), " 1")],
109 &mut cx,
110 );
111
112 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
113 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
114 assert_eq!(move_target.to_point(&snapshot), Point::new(1, 3));
115 });
116
117 // When accepting, cursor is moved to the proposed location
118 accept_completion(&mut cx);
119 cx.assert_editor_state(indoc! {"
120 line 0
121 linˇe
122 line 2
123 line 3
124 line 4
125 "});
126}
127
128#[gpui::test]
129async fn test_edit_prediction_invalidation_range(cx: &mut gpui::TestAppContext) {
130 init_test(cx, |_| {});
131
132 let mut cx = EditorTestContext::new(cx).await;
133 let provider = cx.new(|_| FakeEditPredictionProvider::default());
134 assign_editor_completion_provider(provider.clone(), &mut cx);
135
136 // Cursor is 3+ lines above the proposed edit
137 cx.set_state(indoc! {"
138 line 0
139 line ˇ1
140 line 2
141 line 3
142 line 4
143 line
144 "});
145 let edit_location = Point::new(5, 3);
146
147 propose_edits(
148 &provider,
149 vec![(edit_location..edit_location, " 5")],
150 &mut cx,
151 );
152
153 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
154 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
155 assert_eq!(move_target.to_point(&snapshot), edit_location);
156 });
157
158 // If we move *towards* the completion, it stays active
159 cx.set_selections_state(indoc! {"
160 line 0
161 line 1
162 line ˇ2
163 line 3
164 line 4
165 line
166 "});
167 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
168 assert_eq!(move_target.to_point(&snapshot), edit_location);
169 });
170
171 // If we move *away* from the completion, it is discarded
172 cx.set_selections_state(indoc! {"
173 line ˇ0
174 line 1
175 line 2
176 line 3
177 line 4
178 line
179 "});
180 cx.editor(|editor, _, _| {
181 assert!(editor.active_edit_prediction.is_none());
182 });
183
184 // Cursor is 3+ lines below the proposed edit
185 cx.set_state(indoc! {"
186 line
187 line 1
188 line 2
189 line 3
190 line ˇ4
191 line 5
192 "});
193 let edit_location = Point::new(0, 3);
194
195 propose_edits(
196 &provider,
197 vec![(edit_location..edit_location, " 0")],
198 &mut cx,
199 );
200
201 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
202 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
203 assert_eq!(move_target.to_point(&snapshot), edit_location);
204 });
205
206 // If we move *towards* the completion, it stays active
207 cx.set_selections_state(indoc! {"
208 line
209 line 1
210 line 2
211 line ˇ3
212 line 4
213 line 5
214 "});
215 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
216 assert_eq!(move_target.to_point(&snapshot), edit_location);
217 });
218
219 // If we move *away* from the completion, it is discarded
220 cx.set_selections_state(indoc! {"
221 line
222 line 1
223 line 2
224 line 3
225 line 4
226 line ˇ5
227 "});
228 cx.editor(|editor, _, _| {
229 assert!(editor.active_edit_prediction.is_none());
230 });
231}
232
233#[gpui::test]
234async fn test_edit_prediction_jump_disabled_for_non_zed_providers(cx: &mut gpui::TestAppContext) {
235 init_test(cx, |_| {});
236
237 let mut cx = EditorTestContext::new(cx).await;
238 let provider = cx.new(|_| FakeNonZedEditPredictionProvider::default());
239 assign_editor_completion_provider_non_zed(provider.clone(), &mut cx);
240
241 // Cursor is 2+ lines above the proposed edit
242 cx.set_state(indoc! {"
243 line 0
244 line ˇ1
245 line 2
246 line 3
247 line
248 "});
249
250 propose_edits_non_zed(
251 &provider,
252 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
253 &mut cx,
254 );
255
256 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
257
258 // For non-Zed providers, there should be no move completion (jump functionality disabled)
259 cx.editor(|editor, _, _| {
260 if let Some(completion_state) = &editor.active_edit_prediction {
261 // Should be an Edit prediction, not a Move prediction
262 match &completion_state.completion {
263 EditPrediction::Edit { .. } => {
264 // This is expected for non-Zed providers
265 }
266 EditPrediction::Move { .. } => {
267 panic!(
268 "Non-Zed providers should not show Move predictions (jump functionality)"
269 );
270 }
271 }
272 }
273 });
274}
275
276#[gpui::test]
277async fn test_edit_predictions_disabled_in_scope(cx: &mut gpui::TestAppContext) {
278 init_test(cx, |_| {});
279
280 update_test_language_settings(cx, |settings| {
281 settings.defaults.edit_predictions_disabled_in = Some(vec!["string".to_string()]);
282 });
283
284 let mut cx = EditorTestContext::new(cx).await;
285 let provider = cx.new(|_| FakeEditPredictionProvider::default());
286 assign_editor_completion_provider(provider.clone(), &mut cx);
287
288 let language = languages::language("javascript", tree_sitter_typescript::LANGUAGE_TSX.into());
289 cx.update_buffer(|buffer, cx| buffer.set_language(Some(language), cx));
290
291 // Test enabled outside of string
292 cx.set_state("const x = \"hello world\"; ˇ");
293 propose_edits(&provider, vec![(24..24, "// comment")], &mut cx);
294 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
295 cx.editor(|editor, _, _| {
296 assert!(
297 editor.active_edit_prediction.is_some(),
298 "Edit predictions should work outside of disabled scopes"
299 );
300 });
301
302 // Test disabled inside of string
303 cx.set_state("const x = \"hello ˇworld\";");
304 propose_edits(&provider, vec![(17..17, "beautiful ")], &mut cx);
305 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
306 cx.editor(|editor, _, _| {
307 assert!(
308 editor.active_edit_prediction.is_none(),
309 "Edit predictions should be disabled in string scopes when configured in edit_predictions_disabled_in"
310 );
311 });
312
313 // Test even when disabled in scope, user requested prediction should show
314 cx.update_editor(|editor, window, cx| {
315 editor.show_edit_prediction(&crate::ShowEditPrediction, window, cx)
316 });
317 cx.editor(|editor, _, _| {
318 assert!(
319 editor.active_edit_prediction.is_some(),
320 "User-requested ShowEditPrediction should display despite disabled scope"
321 );
322 });
323}
324
325#[gpui::test]
326async fn test_user_requested_predictions(cx: &mut gpui::TestAppContext) {
327 init_test(cx, |_| {});
328
329 update_test_language_settings(cx, |settings| {
330 settings.defaults.show_edit_predictions = Some(false);
331 });
332
333 let mut cx = EditorTestContext::new(cx).await;
334 let provider = cx.new(|_| FakeEditPredictionProvider::default());
335 assign_editor_completion_provider(provider.clone(), &mut cx);
336
337 cx.set_state("let value = ˇ;");
338 propose_edits(&provider, vec![(12..12, "42")], &mut cx);
339
340 // Test with global setting disabled, non user requested update should not show
341 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(false, window, cx));
342 cx.editor(|editor, _, _| {
343 assert!(editor.active_edit_prediction.is_none());
344 });
345
346 // Test ShowEditPrediction shows prediction even `show_edit_predictions` as `false`
347 cx.update_editor(|editor, window, cx| {
348 editor.show_edit_prediction(&crate::ShowEditPrediction, window, cx)
349 });
350 assert_editor_active_edit_completion(&mut cx, |_, edits| {
351 assert_eq!(edits.len(), 1);
352 assert_eq!(edits[0].1.as_str(), "42");
353 });
354
355 // clear
356 cx.update_editor(|editor, _, cx| editor.take_active_edit_prediction(cx));
357 cx.editor(|editor, _, _| {
358 assert!(editor.active_edit_prediction.is_none());
359 });
360
361 // Test PreviousEditPrediction shows prediction even `show_edit_predictions` as `false`
362 cx.update_editor(|editor, window, cx| {
363 editor.previous_edit_prediction(&crate::PreviousEditPrediction, window, cx)
364 });
365 assert_editor_active_edit_completion(&mut cx, |_, edits| {
366 assert_eq!(edits.len(), 1);
367 assert_eq!(edits[0].1.as_str(), "42");
368 });
369
370 // clear
371 cx.update_editor(|editor, _, cx| editor.take_active_edit_prediction(cx));
372 cx.editor(|editor, _, _| {
373 assert!(editor.active_edit_prediction.is_none());
374 });
375
376 // Test NextEditPrediction shows prediction even `show_edit_predictions` as `false`
377 cx.update_editor(|editor, window, cx| {
378 editor.next_edit_prediction(&crate::NextEditPrediction, window, cx)
379 });
380 assert_editor_active_edit_completion(&mut cx, |_, edits| {
381 assert_eq!(edits.len(), 1);
382 assert_eq!(edits[0].1.as_str(), "42");
383 });
384}
385
386fn assert_editor_active_edit_completion(
387 cx: &mut EditorTestContext,
388 assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, String)>),
389) {
390 cx.editor(|editor, _, cx| {
391 let completion_state = editor
392 .active_edit_prediction
393 .as_ref()
394 .expect("editor has no active completion");
395
396 if let EditPrediction::Edit { edits, .. } = &completion_state.completion {
397 assert(editor.buffer().read(cx).snapshot(cx), edits);
398 } else {
399 panic!("expected edit completion");
400 }
401 })
402}
403
404fn assert_editor_active_move_completion(
405 cx: &mut EditorTestContext,
406 assert: impl FnOnce(MultiBufferSnapshot, Anchor),
407) {
408 cx.editor(|editor, _, cx| {
409 let completion_state = editor
410 .active_edit_prediction
411 .as_ref()
412 .expect("editor has no active completion");
413
414 if let EditPrediction::Move { target, .. } = &completion_state.completion {
415 assert(editor.buffer().read(cx).snapshot(cx), *target);
416 } else {
417 panic!("expected move completion");
418 }
419 })
420}
421
422fn accept_completion(cx: &mut EditorTestContext) {
423 cx.update_editor(|editor, window, cx| {
424 editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)
425 })
426}
427
428fn propose_edits<T: ToOffset>(
429 provider: &Entity<FakeEditPredictionProvider>,
430 edits: Vec<(Range<T>, &str)>,
431 cx: &mut EditorTestContext,
432) {
433 let snapshot = cx.buffer_snapshot();
434 let edits = edits.into_iter().map(|(range, text)| {
435 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
436 (range, text.into())
437 });
438
439 cx.update(|_, cx| {
440 provider.update(cx, |provider, _| {
441 provider.set_edit_prediction(Some(edit_prediction::EditPrediction {
442 id: None,
443 edits: edits.collect(),
444 edit_preview: None,
445 }))
446 })
447 });
448}
449
450fn assign_editor_completion_provider(
451 provider: Entity<FakeEditPredictionProvider>,
452 cx: &mut EditorTestContext,
453) {
454 cx.update_editor(|editor, window, cx| {
455 editor.set_edit_prediction_provider(Some(provider), window, cx);
456 })
457}
458
459fn propose_edits_non_zed<T: ToOffset>(
460 provider: &Entity<FakeNonZedEditPredictionProvider>,
461 edits: Vec<(Range<T>, &str)>,
462 cx: &mut EditorTestContext,
463) {
464 let snapshot = cx.buffer_snapshot();
465 let edits = edits.into_iter().map(|(range, text)| {
466 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
467 (range, text.into())
468 });
469
470 cx.update(|_, cx| {
471 provider.update(cx, |provider, _| {
472 provider.set_edit_prediction(Some(edit_prediction::EditPrediction {
473 id: None,
474 edits: edits.collect(),
475 edit_preview: None,
476 }))
477 })
478 });
479}
480
481fn assign_editor_completion_provider_non_zed(
482 provider: Entity<FakeNonZedEditPredictionProvider>,
483 cx: &mut EditorTestContext,
484) {
485 cx.update_editor(|editor, window, cx| {
486 editor.set_edit_prediction_provider(Some(provider), window, cx);
487 })
488}
489
490#[derive(Default, Clone)]
491pub struct FakeEditPredictionProvider {
492 pub completion: Option<edit_prediction::EditPrediction>,
493}
494
495impl FakeEditPredictionProvider {
496 pub fn set_edit_prediction(&mut self, completion: Option<edit_prediction::EditPrediction>) {
497 self.completion = completion;
498 }
499}
500
501impl EditPredictionProvider for FakeEditPredictionProvider {
502 fn name() -> &'static str {
503 "fake-completion-provider"
504 }
505
506 fn display_name() -> &'static str {
507 "Fake Completion Provider"
508 }
509
510 fn show_completions_in_menu() -> bool {
511 false
512 }
513
514 fn supports_jump_to_edit() -> bool {
515 true
516 }
517
518 fn is_enabled(
519 &self,
520 _buffer: &gpui::Entity<language::Buffer>,
521 _cursor_position: language::Anchor,
522 _cx: &gpui::App,
523 ) -> bool {
524 true
525 }
526
527 fn is_refreshing(&self) -> bool {
528 false
529 }
530
531 fn refresh(
532 &mut self,
533 _project: Option<Entity<Project>>,
534 _buffer: gpui::Entity<language::Buffer>,
535 _cursor_position: language::Anchor,
536 _debounce: bool,
537 _cx: &mut gpui::Context<Self>,
538 ) {
539 }
540
541 fn cycle(
542 &mut self,
543 _buffer: gpui::Entity<language::Buffer>,
544 _cursor_position: language::Anchor,
545 _direction: edit_prediction::Direction,
546 _cx: &mut gpui::Context<Self>,
547 ) {
548 }
549
550 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
551
552 fn discard(&mut self, _cx: &mut gpui::Context<Self>) {}
553
554 fn suggest<'a>(
555 &mut self,
556 _buffer: &gpui::Entity<language::Buffer>,
557 _cursor_position: language::Anchor,
558 _cx: &mut gpui::Context<Self>,
559 ) -> Option<edit_prediction::EditPrediction> {
560 self.completion.clone()
561 }
562}
563
564#[derive(Default, Clone)]
565pub struct FakeNonZedEditPredictionProvider {
566 pub completion: Option<edit_prediction::EditPrediction>,
567}
568
569impl FakeNonZedEditPredictionProvider {
570 pub fn set_edit_prediction(&mut self, completion: Option<edit_prediction::EditPrediction>) {
571 self.completion = completion;
572 }
573}
574
575impl EditPredictionProvider for FakeNonZedEditPredictionProvider {
576 fn name() -> &'static str {
577 "fake-non-zed-provider"
578 }
579
580 fn display_name() -> &'static str {
581 "Fake Non-Zed Provider"
582 }
583
584 fn show_completions_in_menu() -> bool {
585 false
586 }
587
588 fn supports_jump_to_edit() -> bool {
589 false
590 }
591
592 fn is_enabled(
593 &self,
594 _buffer: &gpui::Entity<language::Buffer>,
595 _cursor_position: language::Anchor,
596 _cx: &gpui::App,
597 ) -> bool {
598 true
599 }
600
601 fn is_refreshing(&self) -> bool {
602 false
603 }
604
605 fn refresh(
606 &mut self,
607 _project: Option<Entity<Project>>,
608 _buffer: gpui::Entity<language::Buffer>,
609 _cursor_position: language::Anchor,
610 _debounce: bool,
611 _cx: &mut gpui::Context<Self>,
612 ) {
613 }
614
615 fn cycle(
616 &mut self,
617 _buffer: gpui::Entity<language::Buffer>,
618 _cursor_position: language::Anchor,
619 _direction: edit_prediction::Direction,
620 _cx: &mut gpui::Context<Self>,
621 ) {
622 }
623
624 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
625
626 fn discard(&mut self, _cx: &mut gpui::Context<Self>) {}
627
628 fn suggest<'a>(
629 &mut self,
630 _buffer: &gpui::Entity<language::Buffer>,
631 _cursor_position: language::Anchor,
632 _cx: &mut gpui::Context<Self>,
633 ) -> Option<edit_prediction::EditPrediction> {
634 self.completion.clone()
635 }
636}