1use std::ops::Range;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4use std::sync::atomic::AtomicBool;
5
6use acp_thread::MentionUri;
7use anyhow::{Context as _, Result};
8use collections::HashMap;
9use editor::display_map::CreaseId;
10use editor::{CompletionProvider, Editor, ExcerptId};
11use file_icons::FileIcons;
12use futures::future::try_join_all;
13use gpui::{App, Entity, Task, WeakEntity};
14use language::{Buffer, CodeLabel, HighlightId};
15use lsp::CompletionContext;
16use parking_lot::Mutex;
17use project::{Completion, CompletionIntent, CompletionResponse, Project, ProjectPath, WorktreeId};
18use rope::Point;
19use text::{Anchor, ToPoint};
20use ui::prelude::*;
21use workspace::Workspace;
22
23use crate::context_picker::MentionLink;
24use crate::context_picker::file_context_picker::{extract_file_name_and_directory, search_files};
25
26#[derive(Default)]
27pub struct MentionSet {
28 paths_by_crease_id: HashMap<CreaseId, MentionUri>,
29}
30
31impl MentionSet {
32 pub fn insert(&mut self, crease_id: CreaseId, path: PathBuf) {
33 self.paths_by_crease_id
34 .insert(crease_id, MentionUri::File(path));
35 }
36
37 pub fn drain(&mut self) -> impl Iterator<Item = CreaseId> {
38 self.paths_by_crease_id.drain().map(|(id, _)| id)
39 }
40
41 pub fn contents(
42 &self,
43 project: Entity<Project>,
44 cx: &mut App,
45 ) -> Task<Result<HashMap<CreaseId, Mention>>> {
46 let contents = self
47 .paths_by_crease_id
48 .iter()
49 .map(|(crease_id, uri)| match uri {
50 MentionUri::File(path) => {
51 let crease_id = *crease_id;
52 let uri = uri.clone();
53 let path = path.to_path_buf();
54 let buffer_task = project.update(cx, |project, cx| {
55 let path = project
56 .find_project_path(path, cx)
57 .context("Failed to find project path")?;
58 anyhow::Ok(project.open_buffer(path, cx))
59 });
60
61 cx.spawn(async move |cx| {
62 let buffer = buffer_task?.await?;
63 let content = buffer.read_with(cx, |buffer, _cx| buffer.text())?;
64
65 anyhow::Ok((crease_id, Mention { uri, content }))
66 })
67 }
68 _ => {
69 // TODO
70 unimplemented!()
71 }
72 })
73 .collect::<Vec<_>>();
74
75 cx.spawn(async move |_cx| {
76 let contents = try_join_all(contents).await?.into_iter().collect();
77 anyhow::Ok(contents)
78 })
79 }
80}
81
82pub struct Mention {
83 pub uri: MentionUri,
84 pub content: String,
85}
86
87pub struct ContextPickerCompletionProvider {
88 workspace: WeakEntity<Workspace>,
89 editor: WeakEntity<Editor>,
90 mention_set: Arc<Mutex<MentionSet>>,
91}
92
93impl ContextPickerCompletionProvider {
94 pub fn new(
95 mention_set: Arc<Mutex<MentionSet>>,
96 workspace: WeakEntity<Workspace>,
97 editor: WeakEntity<Editor>,
98 ) -> Self {
99 Self {
100 mention_set,
101 workspace,
102 editor,
103 }
104 }
105
106 pub(crate) fn completion_for_path(
107 project_path: ProjectPath,
108 path_prefix: &str,
109 is_recent: bool,
110 is_directory: bool,
111 excerpt_id: ExcerptId,
112 source_range: Range<Anchor>,
113 editor: Entity<Editor>,
114 mention_set: Arc<Mutex<MentionSet>>,
115 project: Entity<Project>,
116 cx: &App,
117 ) -> Completion {
118 let (file_name, directory) =
119 extract_file_name_and_directory(&project_path.path, path_prefix);
120
121 let label =
122 build_code_label_for_full_path(&file_name, directory.as_ref().map(|s| s.as_ref()), cx);
123 let full_path = if let Some(directory) = directory {
124 format!("{}{}", directory, file_name)
125 } else {
126 file_name.to_string()
127 };
128
129 let crease_icon_path = if is_directory {
130 FileIcons::get_folder_icon(false, cx).unwrap_or_else(|| IconName::Folder.path().into())
131 } else {
132 FileIcons::get_icon(Path::new(&full_path), cx)
133 .unwrap_or_else(|| IconName::File.path().into())
134 };
135 let completion_icon_path = if is_recent {
136 IconName::HistoryRerun.path().into()
137 } else {
138 crease_icon_path.clone()
139 };
140
141 let new_text = format!("{} ", MentionLink::for_file(&file_name, &full_path));
142 let new_text_len = new_text.len();
143 Completion {
144 replace_range: source_range.clone(),
145 new_text,
146 label,
147 documentation: None,
148 source: project::CompletionSource::Custom,
149 icon_path: Some(completion_icon_path),
150 insert_text_mode: None,
151 confirm: Some(confirm_completion_callback(
152 crease_icon_path,
153 file_name,
154 project_path,
155 excerpt_id,
156 source_range.start,
157 new_text_len - 1,
158 editor,
159 mention_set,
160 project,
161 )),
162 }
163 }
164}
165
166fn build_code_label_for_full_path(file_name: &str, directory: Option<&str>, cx: &App) -> CodeLabel {
167 let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
168 let mut label = CodeLabel::default();
169
170 label.push_str(&file_name, None);
171 label.push_str(" ", None);
172
173 if let Some(directory) = directory {
174 label.push_str(&directory, comment_id);
175 }
176
177 label.filter_range = 0..label.text().len();
178
179 label
180}
181
182impl CompletionProvider for ContextPickerCompletionProvider {
183 fn completions(
184 &self,
185 excerpt_id: ExcerptId,
186 buffer: &Entity<Buffer>,
187 buffer_position: Anchor,
188 _trigger: CompletionContext,
189 _window: &mut Window,
190 cx: &mut Context<Editor>,
191 ) -> Task<Result<Vec<CompletionResponse>>> {
192 let state = buffer.update(cx, |buffer, _cx| {
193 let position = buffer_position.to_point(buffer);
194 let line_start = Point::new(position.row, 0);
195 let offset_to_line = buffer.point_to_offset(line_start);
196 let mut lines = buffer.text_for_range(line_start..position).lines();
197 let line = lines.next()?;
198 MentionCompletion::try_parse(line, offset_to_line)
199 });
200 let Some(state) = state else {
201 return Task::ready(Ok(Vec::new()));
202 };
203
204 let Some(workspace) = self.workspace.upgrade() else {
205 return Task::ready(Ok(Vec::new()));
206 };
207
208 let project = workspace.read(cx).project().clone();
209 let snapshot = buffer.read(cx).snapshot();
210 let source_range = snapshot.anchor_before(state.source_range.start)
211 ..snapshot.anchor_after(state.source_range.end);
212
213 let editor = self.editor.clone();
214 let mention_set = self.mention_set.clone();
215 let MentionCompletion { argument, .. } = state;
216 let query = argument.unwrap_or_else(|| "".to_string());
217
218 let search_task = search_files(query.clone(), Arc::<AtomicBool>::default(), &workspace, cx);
219
220 cx.spawn(async move |_, cx| {
221 let matches = search_task.await;
222 let Some(editor) = editor.upgrade() else {
223 return Ok(Vec::new());
224 };
225
226 let completions = cx.update(|cx| {
227 matches
228 .into_iter()
229 .map(|mat| {
230 let path_match = &mat.mat;
231 let project_path = ProjectPath {
232 worktree_id: WorktreeId::from_usize(path_match.worktree_id),
233 path: path_match.path.clone(),
234 };
235
236 Self::completion_for_path(
237 project_path,
238 &path_match.path_prefix,
239 mat.is_recent,
240 path_match.is_dir,
241 excerpt_id,
242 source_range.clone(),
243 editor.clone(),
244 mention_set.clone(),
245 project.clone(),
246 cx,
247 )
248 })
249 .collect()
250 })?;
251
252 Ok(vec![CompletionResponse {
253 completions,
254 // Since this does its own filtering (see `filter_completions()` returns false),
255 // there is no benefit to computing whether this set of completions is incomplete.
256 is_incomplete: true,
257 }])
258 })
259 }
260
261 fn is_completion_trigger(
262 &self,
263 buffer: &Entity<language::Buffer>,
264 position: language::Anchor,
265 _text: &str,
266 _trigger_in_words: bool,
267 _menu_is_open: bool,
268 cx: &mut Context<Editor>,
269 ) -> bool {
270 let buffer = buffer.read(cx);
271 let position = position.to_point(buffer);
272 let line_start = Point::new(position.row, 0);
273 let offset_to_line = buffer.point_to_offset(line_start);
274 let mut lines = buffer.text_for_range(line_start..position).lines();
275 if let Some(line) = lines.next() {
276 MentionCompletion::try_parse(line, offset_to_line)
277 .map(|completion| {
278 completion.source_range.start <= offset_to_line + position.column as usize
279 && completion.source_range.end >= offset_to_line + position.column as usize
280 })
281 .unwrap_or(false)
282 } else {
283 false
284 }
285 }
286
287 fn sort_completions(&self) -> bool {
288 false
289 }
290
291 fn filter_completions(&self) -> bool {
292 false
293 }
294}
295
296fn confirm_completion_callback(
297 crease_icon_path: SharedString,
298 crease_text: SharedString,
299 project_path: ProjectPath,
300 excerpt_id: ExcerptId,
301 start: Anchor,
302 content_len: usize,
303 editor: Entity<Editor>,
304 mention_set: Arc<Mutex<MentionSet>>,
305 project: Entity<Project>,
306) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
307 Arc::new(move |_, window, cx| {
308 let crease_text = crease_text.clone();
309 let crease_icon_path = crease_icon_path.clone();
310 let editor = editor.clone();
311 let project_path = project_path.clone();
312 let mention_set = mention_set.clone();
313 let project = project.clone();
314 window.defer(cx, move |window, cx| {
315 let crease_id = crate::context_picker::insert_crease_for_mention(
316 excerpt_id,
317 start,
318 content_len,
319 crease_text.clone(),
320 crease_icon_path,
321 editor.clone(),
322 window,
323 cx,
324 );
325
326 let Some(path) = project.read(cx).absolute_path(&project_path, cx) else {
327 return;
328 };
329
330 if let Some(crease_id) = crease_id {
331 mention_set.lock().insert(crease_id, path);
332 }
333 });
334 false
335 })
336}
337
338#[derive(Debug, Default, PartialEq)]
339struct MentionCompletion {
340 source_range: Range<usize>,
341 argument: Option<String>,
342}
343
344impl MentionCompletion {
345 fn try_parse(line: &str, offset_to_line: usize) -> Option<Self> {
346 let last_mention_start = line.rfind('@')?;
347 if last_mention_start >= line.len() {
348 return Some(Self::default());
349 }
350 if last_mention_start > 0
351 && line
352 .chars()
353 .nth(last_mention_start - 1)
354 .map_or(false, |c| !c.is_whitespace())
355 {
356 return None;
357 }
358
359 let rest_of_line = &line[last_mention_start + 1..];
360 let mut argument = None;
361
362 let mut parts = rest_of_line.split_whitespace();
363 let mut end = last_mention_start + 1;
364 if let Some(argument_text) = parts.next() {
365 end += argument_text.len();
366 argument = Some(argument_text.to_string());
367 }
368
369 Some(Self {
370 source_range: last_mention_start + offset_to_line..end + offset_to_line,
371 argument,
372 })
373 }
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379 use gpui::{EventEmitter, FocusHandle, Focusable, TestAppContext, VisualTestContext};
380 use project::{Project, ProjectPath};
381 use serde_json::json;
382 use settings::SettingsStore;
383 use std::{ops::Deref, rc::Rc};
384 use util::path;
385 use workspace::{AppState, Item};
386
387 #[test]
388 fn test_mention_completion_parse() {
389 assert_eq!(MentionCompletion::try_parse("Lorem Ipsum", 0), None);
390
391 assert_eq!(
392 MentionCompletion::try_parse("Lorem @", 0),
393 Some(MentionCompletion {
394 source_range: 6..7,
395 argument: None,
396 })
397 );
398
399 assert_eq!(
400 MentionCompletion::try_parse("Lorem @main", 0),
401 Some(MentionCompletion {
402 source_range: 6..11,
403 argument: Some("main".to_string()),
404 })
405 );
406
407 assert_eq!(MentionCompletion::try_parse("test@", 0), None);
408 }
409
410 struct AtMentionEditor(Entity<Editor>);
411
412 impl Item for AtMentionEditor {
413 type Event = ();
414
415 fn include_in_nav_history() -> bool {
416 false
417 }
418
419 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
420 "Test".into()
421 }
422 }
423
424 impl EventEmitter<()> for AtMentionEditor {}
425
426 impl Focusable for AtMentionEditor {
427 fn focus_handle(&self, cx: &App) -> FocusHandle {
428 self.0.read(cx).focus_handle(cx).clone()
429 }
430 }
431
432 impl Render for AtMentionEditor {
433 fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
434 self.0.clone().into_any_element()
435 }
436 }
437
438 #[gpui::test]
439 async fn test_context_completion_provider(cx: &mut TestAppContext) {
440 init_test(cx);
441
442 let app_state = cx.update(AppState::test);
443
444 cx.update(|cx| {
445 language::init(cx);
446 editor::init(cx);
447 workspace::init(app_state.clone(), cx);
448 Project::init_settings(cx);
449 });
450
451 app_state
452 .fs
453 .as_fake()
454 .insert_tree(
455 path!("/dir"),
456 json!({
457 "editor": "",
458 "a": {
459 "one.txt": "",
460 "two.txt": "",
461 "three.txt": "",
462 "four.txt": ""
463 },
464 "b": {
465 "five.txt": "",
466 "six.txt": "",
467 "seven.txt": "",
468 "eight.txt": "",
469 }
470 }),
471 )
472 .await;
473
474 let project = Project::test(app_state.fs.clone(), [path!("/dir").as_ref()], cx).await;
475 let window = cx.add_window(|window, cx| Workspace::test_new(project.clone(), window, cx));
476 let workspace = window.root(cx).unwrap();
477
478 let worktree = project.update(cx, |project, cx| {
479 let mut worktrees = project.worktrees(cx).collect::<Vec<_>>();
480 assert_eq!(worktrees.len(), 1);
481 worktrees.pop().unwrap()
482 });
483 let worktree_id = worktree.read_with(cx, |worktree, _| worktree.id());
484
485 let mut cx = VisualTestContext::from_window(*window.deref(), cx);
486
487 let paths = vec![
488 path!("a/one.txt"),
489 path!("a/two.txt"),
490 path!("a/three.txt"),
491 path!("a/four.txt"),
492 path!("b/five.txt"),
493 path!("b/six.txt"),
494 path!("b/seven.txt"),
495 path!("b/eight.txt"),
496 ];
497
498 let mut opened_editors = Vec::new();
499 for path in paths {
500 let buffer = workspace
501 .update_in(&mut cx, |workspace, window, cx| {
502 workspace.open_path(
503 ProjectPath {
504 worktree_id,
505 path: Path::new(path).into(),
506 },
507 None,
508 false,
509 window,
510 cx,
511 )
512 })
513 .await
514 .unwrap();
515 opened_editors.push(buffer);
516 }
517
518 let editor = workspace.update_in(&mut cx, |workspace, window, cx| {
519 let editor = cx.new(|cx| {
520 Editor::new(
521 editor::EditorMode::full(),
522 multi_buffer::MultiBuffer::build_simple("", cx),
523 None,
524 window,
525 cx,
526 )
527 });
528 workspace.active_pane().update(cx, |pane, cx| {
529 pane.add_item(
530 Box::new(cx.new(|_| AtMentionEditor(editor.clone()))),
531 true,
532 true,
533 None,
534 window,
535 cx,
536 );
537 });
538 editor
539 });
540
541 let mention_set = Arc::new(Mutex::new(MentionSet::default()));
542
543 let editor_entity = editor.downgrade();
544 editor.update_in(&mut cx, |editor, window, cx| {
545 window.focus(&editor.focus_handle(cx));
546 editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
547 mention_set.clone(),
548 workspace.downgrade(),
549 editor_entity,
550 ))));
551 });
552
553 cx.simulate_input("Lorem ");
554
555 editor.update(&mut cx, |editor, cx| {
556 assert_eq!(editor.text(cx), "Lorem ");
557 assert!(!editor.has_visible_completions_menu());
558 });
559
560 cx.simulate_input("@");
561
562 editor.update(&mut cx, |editor, cx| {
563 assert_eq!(editor.text(cx), "Lorem @");
564 assert!(editor.has_visible_completions_menu());
565 assert_eq!(
566 current_completion_labels(editor),
567 &[
568 "eight.txt dir/b/",
569 "seven.txt dir/b/",
570 "six.txt dir/b/",
571 "five.txt dir/b/",
572 "four.txt dir/a/",
573 "three.txt dir/a/",
574 "two.txt dir/a/",
575 "one.txt dir/a/",
576 "dir ",
577 "a dir/",
578 "four.txt dir/a/",
579 "one.txt dir/a/",
580 "three.txt dir/a/",
581 "two.txt dir/a/",
582 "b dir/",
583 "eight.txt dir/b/",
584 "five.txt dir/b/",
585 "seven.txt dir/b/",
586 "six.txt dir/b/",
587 "editor dir/"
588 ]
589 );
590 });
591
592 // Select and confirm "File"
593 editor.update_in(&mut cx, |editor, window, cx| {
594 assert!(editor.has_visible_completions_menu());
595 editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
596 editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
597 editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
598 editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
599 editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
600 });
601
602 cx.run_until_parked();
603
604 editor.update(&mut cx, |editor, cx| {
605 assert_eq!(editor.text(cx), "Lorem [@four.txt](@file:dir/a/four.txt) ");
606 });
607 }
608
609 fn current_completion_labels(editor: &Editor) -> Vec<String> {
610 let completions = editor.current_completions().expect("Missing completions");
611 completions
612 .into_iter()
613 .map(|completion| completion.label.text.to_string())
614 .collect::<Vec<_>>()
615 }
616
617 pub(crate) fn init_test(cx: &mut TestAppContext) {
618 cx.update(|cx| {
619 let store = SettingsStore::test(cx);
620 cx.set_global(store);
621 theme::init(theme::LoadThemes::JustBase, cx);
622 client::init_settings(cx);
623 language::init(cx);
624 Project::init_settings(cx);
625 workspace::init_settings(cx);
626 editor::init_settings(cx);
627 });
628 }
629}