tenant.openagents/omega

1use crate::{
2    CompletionSource, Copilot, CopilotEditPrediction,
3    request::{
4        DidShowCompletion, DidShowCompletionParams, DidShowInlineEdit, DidShowInlineEditParams,
5        InlineCompletionItem,
6    },
7};
8use anyhow::Result;
9use edit_prediction_types::{
10    EditPrediction, EditPredictionDelegate, EditPredictionDiscardReason, EditPredictionIconSet,
11    EditPredictionRequestTrigger, interpolate_edits,
12};
13use gpui::{App, Context, Entity, Task, TaskExt};
14use icons::IconName;
15use language::{Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, ToPointUtf16};
16use std::{ops::Range, sync::Arc, time::Duration};
17
18pub const COPILOT_DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(75);
19
20pub struct CopilotEditPredictionDelegate {
21    completion: Option<(CopilotEditPrediction, EditPreview)>,
22    pending_refresh: Option<Task<Result<()>>>,
23    copilot: Entity<Copilot>,
24}
25
26impl CopilotEditPredictionDelegate {
27    pub fn new(copilot: Entity<Copilot>) -> Self {
28        Self {
29            completion: None,
30            pending_refresh: None,
31            copilot,
32        }
33    }
34
35    fn active_completion(&self) -> Option<&(CopilotEditPrediction, EditPreview)> {
36        self.completion.as_ref()
37    }
38}
39
40impl EditPredictionDelegate for CopilotEditPredictionDelegate {
41    fn name() -> &'static str {
42        "copilot"
43    }
44
45    fn display_name() -> &'static str {
46        "Copilot"
47    }
48
49    fn show_predictions_in_menu() -> bool {
50        true
51    }
52
53    fn show_tab_accept_marker() -> bool {
54        true
55    }
56
57    fn icons(&self, _cx: &App) -> EditPredictionIconSet {
58        EditPredictionIconSet::new(IconName::Copilot)
59            .with_disabled(IconName::CopilotDisabled)
60            .with_error(IconName::CopilotError)
61    }
62
63    fn is_refreshing(&self, _cx: &App) -> bool {
64        self.pending_refresh.is_some() && self.completion.is_none()
65    }
66
67    fn is_enabled(
68        &self,
69        _buffer: &Entity<Buffer>,
70        _cursor_position: language::Anchor,
71        cx: &App,
72    ) -> bool {
73        self.copilot.read(cx).status().is_authorized()
74    }
75
76    fn refresh(
77        &mut self,
78        buffer: Entity<Buffer>,
79        cursor_position: language::Anchor,
80        debounce: bool,
81        _trigger: EditPredictionRequestTrigger,
82        cx: &mut Context<Self>,
83    ) {
84        let copilot = self.copilot.clone();
85        self.pending_refresh = Some(cx.spawn(async move |this, cx| {
86            if debounce {
87                cx.background_executor()
88                    .timer(COPILOT_DEBOUNCE_TIMEOUT)
89                    .await;
90            }
91
92            let completions = copilot
93                .update(cx, |copilot, cx| {
94                    copilot.completions(&buffer, cursor_position, cx)
95                })
96                .await?;
97
98            if let Some(mut completion) = completions.into_iter().next()
99                && let Some((trimmed_range, trimmed_text, snapshot)) =
100                    cx.update(|cx| trim_completion(&completion, cx))
101            {
102                let preview = buffer
103                    .update(cx, |this, cx| {
104                        this.preview_edits(
105                            Arc::from([(trimmed_range.clone(), trimmed_text.clone())].as_slice()),
106                            cx,
107                        )
108                    })
109                    .await;
110                this.update(cx, |this, cx| {
111                    this.pending_refresh = None;
112                    completion.range = trimmed_range;
113                    completion.text = trimmed_text.to_string();
114                    completion.snapshot = snapshot;
115                    this.completion = Some((completion, preview));
116
117                    cx.notify();
118                })?;
119            }
120
121            Ok(())
122        }));
123    }
124
125    fn accept(&mut self, cx: &mut Context<Self>) {
126        if let Some((completion, _)) = self.active_completion() {
127            self.copilot
128                .update(cx, |copilot, cx| copilot.accept_completion(completion, cx))
129                .detach_and_log_err(cx);
130        }
131    }
132
133    fn discard(&mut self, _reason: EditPredictionDiscardReason, _: &mut Context<Self>) {
134        self.completion.take();
135    }
136
137    fn suggest(
138        &mut self,
139        buffer: &Entity<Buffer>,
140        _: language::Anchor,
141        cx: &mut Context<Self>,
142    ) -> Option<EditPrediction> {
143        let buffer_id = buffer.entity_id();
144        let buffer = buffer.read(cx);
145        let (completion, edit_preview) = self.active_completion()?;
146
147        if Some(buffer_id) != Some(completion.buffer.entity_id())
148            || !completion.range.start.is_valid(buffer)
149            || !completion.range.end.is_valid(buffer)
150        {
151            return None;
152        }
153        let edits = vec![(
154            completion.range.clone(),
155            Arc::from(completion.text.as_ref()),
156        )];
157        let edits = interpolate_edits(&completion.snapshot, &buffer.snapshot(), &edits)
158            .filter(|edits| !edits.is_empty())?;
159        self.copilot.update(cx, |this, _| {
160            if let Ok(server) = this.server.as_authenticated() {
161                match completion.source {
162                    CompletionSource::NextEditSuggestion => {
163                        if let Some(cmd) = completion.command.as_ref() {
164                            _ = server
165                                .lsp
166                                .notify::<DidShowInlineEdit>(DidShowInlineEditParams {
167                                    item: serde_json::json!({"command": {"arguments": cmd.arguments}}),
168                                });
169                        }
170                    }
171                    CompletionSource::InlineCompletion => {
172                        _ = server.lsp.notify::<DidShowCompletion>(DidShowCompletionParams {
173                            item: InlineCompletionItem {
174                                insert_text: completion.text.clone(),
175                                range: lsp::Range::new(
176                                    language::point_to_lsp(
177                                        completion.range.start.to_point_utf16(&completion.snapshot),
178                                    ),
179                                    language::point_to_lsp(
180                                        completion.range.end.to_point_utf16(&completion.snapshot),
181                                    ),
182                                ),
183                                command: completion.command.clone(),
184                            },
185                        });
186                    }
187                }
188            }
189        });
190        Some(EditPrediction::Local {
191            id: None,
192            edits,
193            cursor_position: None,
194            edit_preview: Some(edit_preview.clone()),
195        })
196    }
197}
198
199fn trim_completion(
200    completion: &CopilotEditPrediction,
201    cx: &mut App,
202) -> Option<(Range<Anchor>, Arc<str>, BufferSnapshot)> {
203    let buffer = completion.buffer.read(cx);
204    let mut completion_range = completion.range.to_offset(buffer);
205    let prefix_len = common_prefix(
206        buffer.chars_for_range(completion_range.clone()),
207        completion.text.chars(),
208    );
209    completion_range.start += prefix_len;
210    let suffix_len = common_prefix(
211        buffer.reversed_chars_for_range(completion_range.clone()),
212        completion.text[prefix_len..].chars().rev(),
213    );
214    completion_range.end = completion_range.end.saturating_sub(suffix_len);
215    let completion_text = &completion.text[prefix_len..completion.text.len() - suffix_len];
216    if completion_text.trim().is_empty() {
217        None
218    } else {
219        let snapshot = buffer.snapshot();
220        let completion_range = snapshot.anchor_after(completion_range.start)
221            ..snapshot.anchor_after(completion_range.end);
222
223        Some((completion_range, Arc::from(completion_text), snapshot))
224    }
225}
226
227fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
228    a.zip(b)
229        .take_while(|(a, b)| a == b)
230        .map(|(a, _)| a.len_utf8())
231        .sum()
232}
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237    use edit_prediction_types::EditPredictionGranularity;
238    use editor::{
239        Editor, MultiBuffer, MultiBufferOffset, PathKey, SelectionEffects,
240        test::{editor_content_with_blocks, editor_lsp_test_context::EditorLspTestContext},
241    };
242    use fs::FakeFs;
243    use futures::StreamExt;
244    use gpui::{AppContext as _, BackgroundExecutor, TestAppContext, UpdateGlobal};
245    use indoc::indoc;
246    use language::{
247        Point,
248        language_settings::{CompletionSettingsContent, LspInsertMode, WordsCompletionMode},
249    };
250    use lsp::Uri;
251    use project::Project;
252    use serde_json::json;
253    use settings::{AllLanguageSettingsContent, SettingsStore};
254    use std::future::Future;
255    use util::{
256        path,
257        test::{TextRangeMarker, marked_text_ranges_by},
258    };
259
260    #[gpui::test(iterations = 10)]
261    async fn test_copilot(executor: BackgroundExecutor, cx: &mut TestAppContext) {
262        // flaky
263        init_test(cx, |settings| {
264            settings.defaults.completions = Some(CompletionSettingsContent {
265                words: Some(WordsCompletionMode::Disabled),
266                words_min_length: Some(0),
267                lsp_insert_mode: Some(LspInsertMode::Insert),
268                ..Default::default()
269            });
270        });
271
272        let (copilot, copilot_lsp) = Copilot::fake(cx);
273        let mut cx = EditorLspTestContext::new_rust(
274            lsp::ServerCapabilities {
275                completion_provider: Some(lsp::CompletionOptions {
276                    trigger_characters: Some(vec![".".to_string(), ":".to_string()]),
277                    ..Default::default()
278                }),
279                ..Default::default()
280            },
281            cx,
282        )
283        .await;
284        let copilot_provider = cx.new(|_| CopilotEditPredictionDelegate::new(copilot));
285        cx.update_editor(|editor, window, cx| {
286            editor.set_edit_prediction_provider(
287                Some(copilot_provider),
288                EditPredictionRequestTrigger::EditorCreated,
289                window,
290                cx,
291            )
292        });
293
294        cx.set_state(indoc! {"
295            oneˇ
296            two
297            three
298        "});
299        cx.simulate_keystroke(".");
300        drop(handle_completion_request(
301            &mut cx,
302            indoc! {"
303                one.|<>
304                two
305                three
306            "},
307            vec!["completion_a", "completion_b"],
308        ));
309        handle_copilot_completion_request(
310            &copilot_lsp,
311            vec![crate::request::NextEditSuggestion {
312                text: "one.copilot1".into(),
313                range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 4)),
314                command: None,
315                text_document: lsp::VersionedTextDocumentIdentifier {
316                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
317                    version: 0,
318                },
319            }],
320        );
321        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
322        cx.update_editor(|editor, window, cx| {
323            assert!(editor.context_menu_visible());
324            assert!(editor.has_active_edit_prediction());
325            // Since we have both, the copilot suggestion is existing but does not show up as ghost text
326            assert_eq!(editor.text(cx), "one.\ntwo\nthree\n");
327            assert_eq!(editor.display_text(cx), "one.\ntwo\nthree\n");
328
329            // Confirming a non-copilot completion inserts it and hides the context menu, without showing
330            // the copilot suggestion afterwards.
331            editor
332                .confirm_completion(&Default::default(), window, cx)
333                .unwrap()
334                .detach();
335            assert!(!editor.context_menu_visible());
336            assert!(!editor.has_active_edit_prediction());
337            assert_eq!(editor.text(cx), "one.completion_a\ntwo\nthree\n");
338            assert_eq!(editor.display_text(cx), "one.completion_a\ntwo\nthree\n");
339        });
340
341        // Reset editor and only return copilot suggestions
342        cx.set_state(indoc! {"
343            oneˇ
344            two
345            three
346        "});
347        cx.simulate_keystroke(".");
348
349        drop(handle_completion_request(
350            &mut cx,
351            indoc! {"
352                one.|<>
353                two
354                three
355            "},
356            vec![],
357        ));
358        handle_copilot_completion_request(
359            &copilot_lsp,
360            vec![crate::request::NextEditSuggestion {
361                text: "one.copilot1".into(),
362                range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 4)),
363                command: None,
364                text_document: lsp::VersionedTextDocumentIdentifier {
365                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
366                    version: 0,
367                },
368            }],
369        );
370        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
371        cx.update_editor(|editor, _, cx| {
372            assert!(!editor.context_menu_visible());
373            assert!(editor.has_active_edit_prediction());
374            // Since only the copilot is available, it's shown inline
375            assert_eq!(editor.display_text(cx), "one.copilot1\ntwo\nthree\n");
376            assert_eq!(editor.text(cx), "one.\ntwo\nthree\n");
377        });
378
379        // Ensure existing edit prediction is interpolated when inserting again.
380        cx.simulate_keystroke("c");
381        executor.run_until_parked();
382        cx.update_editor(|editor, _, cx| {
383            assert!(!editor.context_menu_visible());
384            assert!(editor.has_active_edit_prediction());
385            assert_eq!(editor.display_text(cx), "one.copilot1\ntwo\nthree\n");
386            assert_eq!(editor.text(cx), "one.c\ntwo\nthree\n");
387        });
388
389        // After debouncing, new Copilot completions should be requested.
390        handle_copilot_completion_request(
391            &copilot_lsp,
392            vec![crate::request::NextEditSuggestion {
393                text: "one.copilot2".into(),
394                range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 5)),
395                command: None,
396                text_document: lsp::VersionedTextDocumentIdentifier {
397                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
398                    version: 0,
399                },
400            }],
401        );
402        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
403        cx.update_editor(|editor, window, cx| {
404            assert!(!editor.context_menu_visible());
405            assert!(editor.has_active_edit_prediction());
406            assert_eq!(editor.display_text(cx), "one.copilot2\ntwo\nthree\n");
407            assert_eq!(editor.text(cx), "one.c\ntwo\nthree\n");
408
409            // Canceling should remove the active Copilot suggestion.
410            editor.cancel(&Default::default(), window, cx);
411            assert!(!editor.has_active_edit_prediction());
412            assert_eq!(editor.display_text(cx), "one.c\ntwo\nthree\n");
413            assert_eq!(editor.text(cx), "one.c\ntwo\nthree\n");
414
415            // After canceling, tabbing shouldn't insert the previously shown suggestion.
416            editor.tab(&Default::default(), window, cx);
417            assert!(!editor.has_active_edit_prediction());
418            assert_eq!(editor.display_text(cx), "one.c   \ntwo\nthree\n");
419            assert_eq!(editor.text(cx), "one.c   \ntwo\nthree\n");
420
421            // When undoing the previously active suggestion isn't shown again.
422            editor.undo(&Default::default(), window, cx);
423            assert!(!editor.has_active_edit_prediction());
424            assert_eq!(editor.display_text(cx), "one.c\ntwo\nthree\n");
425            assert_eq!(editor.text(cx), "one.c\ntwo\nthree\n");
426        });
427        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
428        cx.editor(|editor, _, cx| {
429            assert!(editor.has_active_edit_prediction());
430            assert_eq!(editor.display_text(cx), "one.copilot2\ntwo\nthree\n");
431            assert_eq!(editor.text(cx), "one.c\ntwo\nthree\n");
432        });
433
434        // If an edit occurs outside of this editor, the suggestion is still correctly interpolated.
435        cx.update_buffer(|buffer, cx| buffer.edit([(5..5, "o")], None, cx));
436        cx.update_editor(|editor, window, cx| {
437            assert!(editor.has_active_edit_prediction());
438            assert_eq!(editor.display_text(cx), "one.copilot2\ntwo\nthree\n");
439            assert_eq!(editor.text(cx), "one.co\ntwo\nthree\n");
440
441            // AcceptEditPrediction when there is an active suggestion inserts it.
442            editor.accept_edit_prediction(&Default::default(), window, cx);
443            assert!(!editor.has_active_edit_prediction());
444            assert_eq!(editor.display_text(cx), "one.copilot2\ntwo\nthree\n");
445            assert_eq!(editor.text(cx), "one.copilot2\ntwo\nthree\n");
446
447            // When undoing the previously active suggestion is shown again.
448            editor.undo(&Default::default(), window, cx);
449            assert!(editor.has_active_edit_prediction());
450            assert_eq!(editor.display_text(cx), "one.copilot2\ntwo\nthree\n");
451            assert_eq!(editor.text(cx), "one.co\ntwo\nthree\n");
452
453            // Hide suggestion.
454            editor.cancel(&Default::default(), window, cx);
455            assert!(!editor.has_active_edit_prediction());
456            assert_eq!(editor.display_text(cx), "one.co\ntwo\nthree\n");
457            assert_eq!(editor.text(cx), "one.co\ntwo\nthree\n");
458        });
459
460        // If an edit occurs outside of this editor but no suggestion is being shown,
461        // we won't make it visible.
462        cx.update_buffer(|buffer, cx| buffer.edit([(6..6, "p")], None, cx));
463        cx.update_editor(|editor, _, cx| {
464            assert!(!editor.has_active_edit_prediction());
465            assert_eq!(editor.display_text(cx), "one.cop\ntwo\nthree\n");
466            assert_eq!(editor.text(cx), "one.cop\ntwo\nthree\n");
467        });
468    }
469
470    #[gpui::test(iterations = 10)]
471    async fn test_accept_partial_copilot_suggestion(
472        executor: BackgroundExecutor,
473        cx: &mut TestAppContext,
474    ) {
475        // flaky
476        init_test(cx, |settings| {
477            settings.defaults.completions = Some(CompletionSettingsContent {
478                words: Some(WordsCompletionMode::Disabled),
479                words_min_length: Some(0),
480                lsp_insert_mode: Some(LspInsertMode::Insert),
481                ..Default::default()
482            });
483        });
484
485        let (copilot, copilot_lsp) = Copilot::fake(cx);
486        let mut cx = EditorLspTestContext::new_rust(
487            lsp::ServerCapabilities {
488                completion_provider: Some(lsp::CompletionOptions {
489                    trigger_characters: Some(vec![".".to_string(), ":".to_string()]),
490                    ..Default::default()
491                }),
492                ..Default::default()
493            },
494            cx,
495        )
496        .await;
497        let copilot_provider = cx.new(|_| CopilotEditPredictionDelegate::new(copilot));
498        cx.update_editor(|editor, window, cx| {
499            editor.set_edit_prediction_provider(
500                Some(copilot_provider),
501                EditPredictionRequestTrigger::EditorCreated,
502                window,
503                cx,
504            )
505        });
506
507        // Setup the editor with a completion request.
508        cx.set_state(indoc! {"
509            oneˇ
510            two
511            three
512        "});
513        cx.simulate_keystroke(".");
514        drop(handle_completion_request(
515            &mut cx,
516            indoc! {"
517                one.|<>
518                two
519                three
520            "},
521            vec![],
522        ));
523        handle_copilot_completion_request(
524            &copilot_lsp,
525            vec![crate::request::NextEditSuggestion {
526                text: "one.copilot1".into(),
527                range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 4)),
528                command: None,
529                text_document: lsp::VersionedTextDocumentIdentifier {
530                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
531                    version: 0,
532                },
533            }],
534        );
535        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
536        cx.update_editor(|editor, window, cx| {
537            assert!(editor.has_active_edit_prediction());
538
539            // Accepting the first word of the suggestion should only accept the first word and still show the rest.
540            editor.accept_partial_edit_prediction(EditPredictionGranularity::Word, window, cx);
541
542            assert!(editor.has_active_edit_prediction());
543            assert_eq!(editor.text(cx), "one.copilot\ntwo\nthree\n");
544            assert_eq!(editor.display_text(cx), "one.copilot1\ntwo\nthree\n");
545
546            // Accepting next word should accept the non-word and copilot suggestion should be gone
547            editor.accept_partial_edit_prediction(EditPredictionGranularity::Word, window, cx);
548
549            assert!(!editor.has_active_edit_prediction());
550            assert_eq!(editor.text(cx), "one.copilot1\ntwo\nthree\n");
551            assert_eq!(editor.display_text(cx), "one.copilot1\ntwo\nthree\n");
552        });
553
554        // Reset the editor and check non-word and whitespace completion
555        cx.set_state(indoc! {"
556            oneˇ
557            two
558            three
559        "});
560        cx.simulate_keystroke(".");
561        drop(handle_completion_request(
562            &mut cx,
563            indoc! {"
564                one.|<>
565                two
566                three
567            "},
568            vec![],
569        ));
570        handle_copilot_completion_request(
571            &copilot_lsp,
572            vec![crate::request::NextEditSuggestion {
573                text: "one.123. copilot\n 456".into(),
574                range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 4)),
575                command: None,
576                text_document: lsp::VersionedTextDocumentIdentifier {
577                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
578                    version: 0,
579                },
580            }],
581        );
582        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
583        cx.update_editor(|editor, window, cx| {
584            assert!(editor.has_active_edit_prediction());
585
586            // Accepting the first word (non-word) of the suggestion should only accept the first word and still show the rest.
587            editor.accept_partial_edit_prediction(EditPredictionGranularity::Word, window, cx);
588            assert!(editor.has_active_edit_prediction());
589            assert_eq!(editor.text(cx), "one.123. \ntwo\nthree\n");
590            assert_eq!(
591                editor.display_text(cx),
592                "one.123. copilot\n 456\ntwo\nthree\n"
593            );
594
595            // Accepting next word should accept the next word and copilot suggestion should still exist
596            editor.accept_partial_edit_prediction(EditPredictionGranularity::Word, window, cx);
597            assert!(editor.has_active_edit_prediction());
598            assert_eq!(editor.text(cx), "one.123. copilot\ntwo\nthree\n");
599            assert_eq!(
600                editor.display_text(cx),
601                "one.123. copilot\n 456\ntwo\nthree\n"
602            );
603
604            // Accepting the whitespace should accept the non-word/whitespaces with newline and copilot suggestion should be gone
605            editor.accept_partial_edit_prediction(EditPredictionGranularity::Word, window, cx);
606            assert!(!editor.has_active_edit_prediction());
607            assert_eq!(editor.text(cx), "one.123. copilot\n 456\ntwo\nthree\n");
608            assert_eq!(
609                editor.display_text(cx),
610                "one.123. copilot\n 456\ntwo\nthree\n"
611            );
612        });
613    }
614
615    #[gpui::test]
616    async fn test_copilot_completion_invalidation(
617        executor: BackgroundExecutor,
618        cx: &mut TestAppContext,
619    ) {
620        init_test(cx, |_| {});
621
622        let (copilot, copilot_lsp) = Copilot::fake(cx);
623        let mut cx = EditorLspTestContext::new_rust(
624            lsp::ServerCapabilities {
625                completion_provider: Some(lsp::CompletionOptions {
626                    trigger_characters: Some(vec![".".to_string(), ":".to_string()]),
627                    ..Default::default()
628                }),
629                ..Default::default()
630            },
631            cx,
632        )
633        .await;
634        let copilot_provider = cx.new(|_| CopilotEditPredictionDelegate::new(copilot));
635        cx.update_editor(|editor, window, cx| {
636            editor.set_edit_prediction_provider(
637                Some(copilot_provider),
638                EditPredictionRequestTrigger::EditorCreated,
639                window,
640                cx,
641            )
642        });
643
644        cx.set_state(indoc! {"
645            one
646            twˇ
647            three
648        "});
649
650        handle_copilot_completion_request(
651            &copilot_lsp,
652            vec![crate::request::NextEditSuggestion {
653                text: "two.foo()".into(),
654                range: lsp::Range::new(lsp::Position::new(1, 0), lsp::Position::new(1, 2)),
655                command: None,
656                text_document: lsp::VersionedTextDocumentIdentifier {
657                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
658                    version: 0,
659                },
660            }],
661        );
662        cx.update_editor(|editor, window, cx| {
663            editor.show_edit_prediction(&Default::default(), window, cx)
664        });
665        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
666        cx.update_editor(|editor, window, cx| {
667            assert!(editor.has_active_edit_prediction());
668            assert_eq!(editor.display_text(cx), "one\ntwo.foo()\nthree\n");
669            assert_eq!(editor.text(cx), "one\ntw\nthree\n");
670
671            editor.backspace(&Default::default(), window, cx);
672        });
673        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
674        cx.run_until_parked();
675        cx.update_editor(|editor, window, cx| {
676            assert!(editor.has_active_edit_prediction());
677            assert_eq!(editor.display_text(cx), "one\ntwo.foo()\nthree\n");
678            assert_eq!(editor.text(cx), "one\nt\nthree\n");
679
680            editor.backspace(&Default::default(), window, cx);
681        });
682        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
683        cx.run_until_parked();
684        cx.update_editor(|editor, window, cx| {
685            assert!(editor.has_active_edit_prediction());
686            assert_eq!(editor.display_text(cx), "one\ntwo.foo()\nthree\n");
687            assert_eq!(editor.text(cx), "one\n\nthree\n");
688            // Deleting across the original suggestion range invalidates it.
689            editor.backspace(&Default::default(), window, cx);
690            assert!(!editor.has_active_edit_prediction());
691            assert_eq!(editor.display_text(cx), "one\nthree\n");
692            assert_eq!(editor.text(cx), "one\nthree\n");
693
694            // Undoing the deletion restores the suggestion.
695            editor.undo(&Default::default(), window, cx);
696            assert!(editor.has_active_edit_prediction());
697            assert_eq!(editor.display_text(cx), "one\ntwo.foo()\nthree\n");
698            assert_eq!(editor.text(cx), "one\n\nthree\n");
699        });
700    }
701
702    #[gpui::test]
703    async fn test_copilot_multibuffer(executor: BackgroundExecutor, cx: &mut TestAppContext) {
704        init_test(cx, |_| {});
705
706        let (copilot, copilot_lsp) = Copilot::fake(cx);
707
708        let buffer_1 = cx.new(|cx| Buffer::local("a = 1\nb = 2\n", cx));
709        let buffer_2 = cx.new(|cx| Buffer::local("c = 3\nd = 4\n", cx));
710        let multibuffer = cx.new(|cx| {
711            let mut multibuffer = MultiBuffer::new(language::Capability::ReadWrite);
712            multibuffer.set_excerpts_for_path(
713                PathKey::sorted(0),
714                buffer_1.clone(),
715                [Point::new(0, 0)..Point::new(1, 0)],
716                0,
717                cx,
718            );
719            multibuffer.set_excerpts_for_path(
720                PathKey::sorted(1),
721                buffer_2.clone(),
722                [Point::new(0, 0)..Point::new(1, 0)],
723                0,
724                cx,
725            );
726            multibuffer
727        });
728        let (editor, cx) =
729            cx.add_window_view(|window, cx| Editor::for_multibuffer(multibuffer, None, window, cx));
730        editor.update_in(cx, |editor, window, cx| {
731            use gpui::Focusable;
732            window.focus(&editor.focus_handle(cx), cx);
733        });
734        let copilot_provider = cx.new(|_| CopilotEditPredictionDelegate::new(copilot));
735        editor.update_in(cx, |editor, window, cx| {
736            editor.set_edit_prediction_provider(
737                Some(copilot_provider),
738                EditPredictionRequestTrigger::EditorCreated,
739                window,
740                cx,
741            )
742        });
743
744        handle_copilot_completion_request(
745            &copilot_lsp,
746            vec![crate::request::NextEditSuggestion {
747                text: "b = 2 + a".into(),
748                range: lsp::Range::new(lsp::Position::new(1, 0), lsp::Position::new(1, 5)),
749                command: None,
750                text_document: lsp::VersionedTextDocumentIdentifier {
751                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
752                    version: 0,
753                },
754            }],
755        );
756        _ = editor.update_in(cx, |editor, window, cx| {
757            // Ensure copilot suggestions are shown for the first excerpt.
758            editor.change_selections(SelectionEffects::no_scroll(), window, cx, |s| {
759                s.select_ranges([Point::new(1, 5)..Point::new(1, 5)])
760            });
761            editor.show_edit_prediction(&Default::default(), window, cx);
762        });
763        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
764        _ = editor.update_in(cx, |editor, _, _| {
765            assert!(editor.has_active_edit_prediction());
766        });
767        pretty_assertions::assert_eq!(
768            editor_content_with_blocks(&editor, cx),
769            indoc! { "
770                § <no file>
771                § -----
772                a = 1
773                b = 2 + a
774                § <no file>
775                § -----
776                c = 3
777                d = 4"
778            }
779        );
780
781        handle_copilot_completion_request(
782            &copilot_lsp,
783            vec![crate::request::NextEditSuggestion {
784                text: "d = 4 + c".into(),
785                range: lsp::Range::new(lsp::Position::new(1, 0), lsp::Position::new(1, 6)),
786                command: None,
787                text_document: lsp::VersionedTextDocumentIdentifier {
788                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
789                    version: 0,
790                },
791            }],
792        );
793        _ = editor.update_in(cx, |editor, window, cx| {
794            // Move to another excerpt, ensuring the suggestion gets cleared.
795            editor.change_selections(SelectionEffects::no_scroll(), window, cx, |s| {
796                s.select_ranges([Point::new(4, 5)..Point::new(4, 5)])
797            });
798            assert!(!editor.has_active_edit_prediction());
799        });
800        pretty_assertions::assert_eq!(
801            editor_content_with_blocks(&editor, cx),
802            indoc! { "
803                § <no file>
804                § -----
805                a = 1
806                b = 2
807                § <no file>
808                § -----
809                c = 3
810                d = 4"}
811        );
812        editor.update_in(cx, |editor, window, cx| {
813            // Type a character, ensuring we don't even try to interpolate the previous suggestion.
814            editor.handle_input(" ", window, cx);
815            assert!(!editor.has_active_edit_prediction());
816        });
817        pretty_assertions::assert_eq!(
818            editor_content_with_blocks(&editor, cx),
819            indoc! {"
820                § <no file>
821                § -----
822                a = 1
823                b = 2
824                § <no file>
825                § -----
826                c = 3
827                d = 4\x20"
828            },
829        );
830
831        // Ensure the new suggestion is displayed when the debounce timeout expires.
832        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
833        _ = editor.update(cx, |editor, _| {
834            assert!(editor.has_active_edit_prediction());
835        });
836        assert_eq!(
837            editor_content_with_blocks(&editor, cx),
838            indoc! {"
839               § <no file>
840               § -----
841               a = 1
842               b = 2
843               § <no file>
844               § -----
845               c = 3
846               d = 4 + c"}
847        );
848    }
849
850    #[gpui::test]
851    async fn test_copilot_does_not_prevent_completion_triggers(
852        executor: BackgroundExecutor,
853        cx: &mut TestAppContext,
854    ) {
855        init_test(cx, |_| {});
856
857        let (copilot, copilot_lsp) = Copilot::fake(cx);
858        let mut cx = EditorLspTestContext::new_rust(
859            lsp::ServerCapabilities {
860                completion_provider: Some(lsp::CompletionOptions {
861                    trigger_characters: Some(vec![".".to_string(), ":".to_string()]),
862                    ..lsp::CompletionOptions::default()
863                }),
864                ..lsp::ServerCapabilities::default()
865            },
866            cx,
867        )
868        .await;
869        let copilot_provider = cx.new(|_| CopilotEditPredictionDelegate::new(copilot));
870        cx.update_editor(|editor, window, cx| {
871            editor.set_edit_prediction_provider(
872                Some(copilot_provider),
873                EditPredictionRequestTrigger::EditorCreated,
874                window,
875                cx,
876            )
877        });
878
879        cx.set_state(indoc! {"
880                one
881                twˇ
882                three
883            "});
884
885        drop(handle_completion_request(
886            &mut cx,
887            indoc! {"
888                one
889                tw|<>
890                three
891            "},
892            vec!["completion_a", "completion_b"],
893        ));
894        handle_copilot_completion_request(
895            &copilot_lsp,
896            vec![crate::request::NextEditSuggestion {
897                text: "two.foo()".into(),
898                range: lsp::Range::new(lsp::Position::new(1, 0), lsp::Position::new(1, 2)),
899                command: None,
900                text_document: lsp::VersionedTextDocumentIdentifier {
901                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
902                    version: 0,
903                },
904            }],
905        );
906        cx.update_editor(|editor, window, cx| {
907            editor.show_edit_prediction(&Default::default(), window, cx)
908        });
909        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
910        cx.update_editor(|editor, _, cx| {
911            assert!(!editor.context_menu_visible());
912            assert!(editor.has_active_edit_prediction());
913            assert_eq!(editor.display_text(cx), "one\ntwo.foo()\nthree\n");
914            assert_eq!(editor.text(cx), "one\ntw\nthree\n");
915        });
916
917        cx.simulate_keystroke("o");
918        drop(handle_completion_request(
919            &mut cx,
920            indoc! {"
921                one
922                two|<>
923                three
924            "},
925            vec!["completion_a_2", "completion_b_2"],
926        ));
927        handle_copilot_completion_request(
928            &copilot_lsp,
929            vec![crate::request::NextEditSuggestion {
930                text: "two.foo()".into(),
931                range: lsp::Range::new(lsp::Position::new(1, 0), lsp::Position::new(1, 3)),
932                command: None,
933                text_document: lsp::VersionedTextDocumentIdentifier {
934                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
935                    version: 0,
936                },
937            }],
938        );
939        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
940        cx.update_editor(|editor, _, cx| {
941            assert!(!editor.context_menu_visible());
942            assert!(editor.has_active_edit_prediction());
943            assert_eq!(editor.display_text(cx), "one\ntwo.foo()\nthree\n");
944            assert_eq!(editor.text(cx), "one\ntwo\nthree\n");
945        });
946
947        cx.simulate_keystroke(".");
948        drop(handle_completion_request(
949            &mut cx,
950            indoc! {"
951                one
952                two.|<>
953                three
954            "},
955            vec!["something_else()"],
956        ));
957        handle_copilot_completion_request(
958            &copilot_lsp,
959            vec![crate::request::NextEditSuggestion {
960                text: "two.foo()".into(),
961                range: lsp::Range::new(lsp::Position::new(1, 0), lsp::Position::new(1, 4)),
962                command: None,
963                text_document: lsp::VersionedTextDocumentIdentifier {
964                    uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
965                    version: 0,
966                },
967            }],
968        );
969        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
970        cx.update_editor(|editor, _, cx| {
971            assert!(editor.context_menu_visible());
972            assert!(editor.has_active_edit_prediction());
973            assert_eq!(editor.text(cx), "one\ntwo.\nthree\n");
974            assert_eq!(editor.display_text(cx), "one\ntwo.\nthree\n");
975        });
976    }
977
978    #[gpui::test]
979    async fn test_copilot_disabled_globs(executor: BackgroundExecutor, cx: &mut TestAppContext) {
980        init_test(cx, |settings| {
981            settings
982                .edit_predictions
983                .get_or_insert(Default::default())
984                .disabled_globs = Some(vec![".env*".to_string()]);
985        });
986
987        let (copilot, copilot_lsp) = Copilot::fake(cx);
988
989        let fs = FakeFs::new(cx.executor());
990        fs.insert_tree(
991            path!("/test"),
992            json!({
993                ".env": "SECRET=something\n",
994                "README.md": "hello\nworld\nhow\nare\nyou\ntoday"
995            }),
996        )
997        .await;
998        let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
999
1000        let private_buffer = project
1001            .update(cx, |project, cx| {
1002                project.open_local_buffer(path!("/test/.env"), cx)
1003            })
1004            .await
1005            .unwrap();
1006        let public_buffer = project
1007            .update(cx, |project, cx| {
1008                project.open_local_buffer(path!("/test/README.md"), cx)
1009            })
1010            .await
1011            .unwrap();
1012
1013        let multibuffer = cx.new(|cx| {
1014            let mut multibuffer = MultiBuffer::new(language::Capability::ReadWrite);
1015            multibuffer.set_excerpts_for_path(
1016                PathKey::sorted(0),
1017                private_buffer.clone(),
1018                [Point::new(0, 0)..Point::new(1, 0)],
1019                0,
1020                cx,
1021            );
1022            multibuffer.set_excerpts_for_path(
1023                PathKey::sorted(1),
1024                public_buffer.clone(),
1025                [Point::new(0, 0)..Point::new(6, 0)],
1026                0,
1027                cx,
1028            );
1029            multibuffer
1030        });
1031        let editor =
1032            cx.add_window(|window, cx| Editor::for_multibuffer(multibuffer, None, window, cx));
1033        editor
1034            .update(cx, |editor, window, cx| {
1035                use gpui::Focusable;
1036                window.focus(&editor.focus_handle(cx), cx)
1037            })
1038            .unwrap();
1039        let copilot_provider = cx.new(|_| CopilotEditPredictionDelegate::new(copilot));
1040        editor
1041            .update(cx, |editor, window, cx| {
1042                editor.set_edit_prediction_provider(
1043                    Some(copilot_provider),
1044                    EditPredictionRequestTrigger::EditorCreated,
1045                    window,
1046                    cx,
1047                )
1048            })
1049            .unwrap();
1050
1051        let mut copilot_requests = copilot_lsp
1052            .set_request_handler::<crate::request::NextEditSuggestions, _, _>(
1053                move |_params, _cx| async move {
1054                    Ok(crate::request::NextEditSuggestionsResult {
1055                        edits: vec![crate::request::NextEditSuggestion {
1056                            text: "next line".into(),
1057                            range: lsp::Range::new(
1058                                lsp::Position::new(1, 0),
1059                                lsp::Position::new(1, 0),
1060                            ),
1061                            command: None,
1062                            text_document: lsp::VersionedTextDocumentIdentifier {
1063                                uri: Uri::from_file_path(path!("/root/dir/file.rs")).unwrap(),
1064                                version: 0,
1065                            },
1066                        }],
1067                    })
1068                },
1069            );
1070
1071        _ = editor.update(cx, |editor, window, cx| {
1072            editor.change_selections(SelectionEffects::no_scroll(), window, cx, |selections| {
1073                selections.select_ranges([Point::new(0, 0)..Point::new(0, 0)])
1074            });
1075            editor.refresh_edit_prediction(
1076                true,
1077                false,
1078                EditPredictionRequestTrigger::BufferEdit,
1079                window,
1080                cx,
1081            );
1082        });
1083
1084        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
1085        assert!(copilot_requests.try_recv().is_err());
1086
1087        _ = editor.update(cx, |editor, window, cx| {
1088            editor.change_selections(SelectionEffects::no_scroll(), window, cx, |s| {
1089                s.select_ranges([Point::new(5, 0)..Point::new(5, 0)])
1090            });
1091            editor.refresh_edit_prediction(
1092                true,
1093                false,
1094                EditPredictionRequestTrigger::BufferEdit,
1095                window,
1096                cx,
1097            );
1098        });
1099
1100        executor.advance_clock(COPILOT_DEBOUNCE_TIMEOUT);
1101        assert!(copilot_requests.try_recv().is_ok());
1102    }
1103
1104    fn handle_copilot_completion_request(
1105        lsp: &lsp::FakeLanguageServer,
1106        completions: Vec<crate::request::NextEditSuggestion>,
1107    ) {
1108        lsp.set_request_handler::<crate::request::NextEditSuggestions, _, _>(
1109            move |_params, _cx| {
1110                let completions = completions.clone();
1111                async move {
1112                    Ok(crate::request::NextEditSuggestionsResult {
1113                        edits: completions.clone(),
1114                    })
1115                }
1116            },
1117        );
1118    }
1119
1120    fn handle_completion_request(
1121        cx: &mut EditorLspTestContext,
1122        marked_string: &str,
1123        completions: Vec<&'static str>,
1124    ) -> impl Future<Output = ()> {
1125        let complete_from_marker: TextRangeMarker = '|'.into();
1126        let replace_range_marker: TextRangeMarker = ('<', '>').into();
1127        let (_, mut marked_ranges) = marked_text_ranges_by(
1128            marked_string,
1129            vec![complete_from_marker, replace_range_marker.clone()],
1130        );
1131
1132        let range = marked_ranges.remove(&replace_range_marker).unwrap()[0].clone();
1133        let replace_range =
1134            cx.to_lsp_range(MultiBufferOffset(range.start)..MultiBufferOffset(range.end));
1135
1136        let mut request =
1137            cx.set_request_handler::<lsp::request::Completion, _, _>(move |url, params, _| {
1138                let completions = completions.clone();
1139                async move {
1140                    assert_eq!(params.text_document_position.text_document.uri, url.clone());
1141                    Ok(Some(lsp::CompletionResponse::Array(
1142                        completions
1143                            .iter()
1144                            .map(|completion_text| lsp::CompletionItem {
1145                                label: completion_text.to_string(),
1146                                text_edit: Some(lsp::CompletionTextEdit::Edit(lsp::TextEdit {
1147                                    range: replace_range,
1148                                    new_text: completion_text.to_string(),
1149                                })),
1150                                ..Default::default()
1151                            })
1152                            .collect(),
1153                    )))
1154                }
1155            });
1156
1157        async move {
1158            request.next().await;
1159        }
1160    }
1161
1162    fn init_test(cx: &mut TestAppContext, f: fn(&mut AllLanguageSettingsContent)) {
1163        cx.update(|cx| {
1164            let store = SettingsStore::test(cx);
1165            cx.set_global(store);
1166            theme_settings::init(theme::LoadThemes::JustBase, cx);
1167            SettingsStore::update_global(cx, |store: &mut SettingsStore, cx| {
1168                store.update_user_settings(cx, |settings| f(&mut settings.project.all_languages));
1169            });
1170        });
1171    }
1172}
1173