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