completion_provider.rs

  1use std::ops::Range;
  2use std::sync::Arc;
  3use std::sync::atomic::AtomicBool;
  4
  5use acp_thread::MentionUri;
  6use anyhow::Result;
  7use editor::{CompletionProvider, Editor, ExcerptId};
  8use fuzzy::{StringMatch, StringMatchCandidate};
  9use gpui::{App, Entity, Task, WeakEntity};
 10use language::{Buffer, CodeLabel, HighlightId};
 11use lsp::CompletionContext;
 12use project::{
 13    Completion, CompletionIntent, CompletionResponse, Project, ProjectPath, Symbol, WorktreeId,
 14};
 15use prompt_store::PromptStore;
 16use rope::Point;
 17use text::{Anchor, ToPoint as _};
 18use ui::prelude::*;
 19use workspace::Workspace;
 20
 21use agent::thread_store::{TextThreadStore, ThreadStore};
 22
 23use crate::acp::message_editor::MessageEditor;
 24use crate::context_picker::file_context_picker::{FileMatch, search_files};
 25use crate::context_picker::rules_context_picker::{RulesContextEntry, search_rules};
 26use crate::context_picker::symbol_context_picker::SymbolMatch;
 27use crate::context_picker::symbol_context_picker::search_symbols;
 28use crate::context_picker::thread_context_picker::{
 29    ThreadContextEntry, ThreadMatch, search_threads,
 30};
 31use crate::context_picker::{
 32    ContextPickerAction, ContextPickerEntry, ContextPickerMode, RecentEntry,
 33    available_context_picker_entries, recent_context_picker_entries, selection_ranges,
 34};
 35
 36pub(crate) enum Match {
 37    File(FileMatch),
 38    Symbol(SymbolMatch),
 39    Thread(ThreadMatch),
 40    Fetch(SharedString),
 41    Rules(RulesContextEntry),
 42    Entry(EntryMatch),
 43}
 44
 45pub struct EntryMatch {
 46    mat: Option<StringMatch>,
 47    entry: ContextPickerEntry,
 48}
 49
 50impl Match {
 51    pub fn score(&self) -> f64 {
 52        match self {
 53            Match::File(file) => file.mat.score,
 54            Match::Entry(mode) => mode.mat.as_ref().map(|mat| mat.score).unwrap_or(1.),
 55            Match::Thread(_) => 1.,
 56            Match::Symbol(_) => 1.,
 57            Match::Rules(_) => 1.,
 58            Match::Fetch(_) => 1.,
 59        }
 60    }
 61}
 62
 63fn search(
 64    mode: Option<ContextPickerMode>,
 65    query: String,
 66    cancellation_flag: Arc<AtomicBool>,
 67    recent_entries: Vec<RecentEntry>,
 68    prompt_store: Option<Entity<PromptStore>>,
 69    thread_store: WeakEntity<ThreadStore>,
 70    text_thread_context_store: WeakEntity<assistant_context::ContextStore>,
 71    workspace: Entity<Workspace>,
 72    cx: &mut App,
 73) -> Task<Vec<Match>> {
 74    match mode {
 75        Some(ContextPickerMode::File) => {
 76            let search_files_task =
 77                search_files(query.clone(), cancellation_flag.clone(), &workspace, cx);
 78            cx.background_spawn(async move {
 79                search_files_task
 80                    .await
 81                    .into_iter()
 82                    .map(Match::File)
 83                    .collect()
 84            })
 85        }
 86
 87        Some(ContextPickerMode::Symbol) => {
 88            let search_symbols_task =
 89                search_symbols(query.clone(), cancellation_flag.clone(), &workspace, cx);
 90            cx.background_spawn(async move {
 91                search_symbols_task
 92                    .await
 93                    .into_iter()
 94                    .map(Match::Symbol)
 95                    .collect()
 96            })
 97        }
 98
 99        Some(ContextPickerMode::Thread) => {
100            if let Some((thread_store, context_store)) = thread_store
101                .upgrade()
102                .zip(text_thread_context_store.upgrade())
103            {
104                let search_threads_task = search_threads(
105                    query.clone(),
106                    cancellation_flag.clone(),
107                    thread_store,
108                    context_store,
109                    cx,
110                );
111                cx.background_spawn(async move {
112                    search_threads_task
113                        .await
114                        .into_iter()
115                        .map(Match::Thread)
116                        .collect()
117                })
118            } else {
119                Task::ready(Vec::new())
120            }
121        }
122
123        Some(ContextPickerMode::Fetch) => {
124            if !query.is_empty() {
125                Task::ready(vec![Match::Fetch(query.into())])
126            } else {
127                Task::ready(Vec::new())
128            }
129        }
130
131        Some(ContextPickerMode::Rules) => {
132            if let Some(prompt_store) = prompt_store.as_ref() {
133                let search_rules_task =
134                    search_rules(query.clone(), cancellation_flag.clone(), prompt_store, cx);
135                cx.background_spawn(async move {
136                    search_rules_task
137                        .await
138                        .into_iter()
139                        .map(Match::Rules)
140                        .collect::<Vec<_>>()
141                })
142            } else {
143                Task::ready(Vec::new())
144            }
145        }
146
147        None => {
148            if query.is_empty() {
149                let mut matches = recent_entries
150                    .into_iter()
151                    .map(|entry| match entry {
152                        RecentEntry::File {
153                            project_path,
154                            path_prefix,
155                        } => Match::File(FileMatch {
156                            mat: fuzzy::PathMatch {
157                                score: 1.,
158                                positions: Vec::new(),
159                                worktree_id: project_path.worktree_id.to_usize(),
160                                path: project_path.path,
161                                path_prefix,
162                                is_dir: false,
163                                distance_to_relative_ancestor: 0,
164                            },
165                            is_recent: true,
166                        }),
167                        RecentEntry::Thread(thread_context_entry) => Match::Thread(ThreadMatch {
168                            thread: thread_context_entry,
169                            is_recent: true,
170                        }),
171                    })
172                    .collect::<Vec<_>>();
173
174                matches.extend(
175                    available_context_picker_entries(
176                        &prompt_store,
177                        &Some(thread_store.clone()),
178                        &workspace,
179                        cx,
180                    )
181                    .into_iter()
182                    .map(|mode| {
183                        Match::Entry(EntryMatch {
184                            entry: mode,
185                            mat: None,
186                        })
187                    }),
188                );
189
190                Task::ready(matches)
191            } else {
192                let executor = cx.background_executor().clone();
193
194                let search_files_task =
195                    search_files(query.clone(), cancellation_flag.clone(), &workspace, cx);
196
197                let entries = available_context_picker_entries(
198                    &prompt_store,
199                    &Some(thread_store.clone()),
200                    &workspace,
201                    cx,
202                );
203                let entry_candidates = entries
204                    .iter()
205                    .enumerate()
206                    .map(|(ix, entry)| StringMatchCandidate::new(ix, entry.keyword()))
207                    .collect::<Vec<_>>();
208
209                cx.background_spawn(async move {
210                    let mut matches = search_files_task
211                        .await
212                        .into_iter()
213                        .map(Match::File)
214                        .collect::<Vec<_>>();
215
216                    let entry_matches = fuzzy::match_strings(
217                        &entry_candidates,
218                        &query,
219                        false,
220                        true,
221                        100,
222                        &Arc::new(AtomicBool::default()),
223                        executor,
224                    )
225                    .await;
226
227                    matches.extend(entry_matches.into_iter().map(|mat| {
228                        Match::Entry(EntryMatch {
229                            entry: entries[mat.candidate_id],
230                            mat: Some(mat),
231                        })
232                    }));
233
234                    matches.sort_by(|a, b| {
235                        b.score()
236                            .partial_cmp(&a.score())
237                            .unwrap_or(std::cmp::Ordering::Equal)
238                    });
239
240                    matches
241                })
242            }
243        }
244    }
245}
246
247pub struct ContextPickerCompletionProvider {
248    workspace: WeakEntity<Workspace>,
249    thread_store: WeakEntity<ThreadStore>,
250    text_thread_store: WeakEntity<TextThreadStore>,
251    message_editor: WeakEntity<MessageEditor>,
252}
253
254impl ContextPickerCompletionProvider {
255    pub fn new(
256        workspace: WeakEntity<Workspace>,
257        thread_store: WeakEntity<ThreadStore>,
258        text_thread_store: WeakEntity<TextThreadStore>,
259        message_editor: WeakEntity<MessageEditor>,
260    ) -> Self {
261        Self {
262            workspace,
263            thread_store,
264            text_thread_store,
265            message_editor,
266        }
267    }
268
269    fn completion_for_entry(
270        entry: ContextPickerEntry,
271        source_range: Range<Anchor>,
272        message_editor: WeakEntity<MessageEditor>,
273        workspace: &Entity<Workspace>,
274        cx: &mut App,
275    ) -> Option<Completion> {
276        match entry {
277            ContextPickerEntry::Mode(mode) => Some(Completion {
278                replace_range: source_range.clone(),
279                new_text: format!("@{} ", mode.keyword()),
280                label: CodeLabel::plain(mode.label().to_string(), None),
281                icon_path: Some(mode.icon().path().into()),
282                documentation: None,
283                source: project::CompletionSource::Custom,
284                insert_text_mode: None,
285                // This ensures that when a user accepts this completion, the
286                // completion menu will still be shown after "@category " is
287                // inserted
288                confirm: Some(Arc::new(|_, _, _| true)),
289            }),
290            ContextPickerEntry::Action(action) => {
291                let (new_text, on_action) = match action {
292                    ContextPickerAction::AddSelections => {
293                        const PLACEHOLDER: &str = "selection ";
294                        let selections = selection_ranges(workspace, cx)
295                            .into_iter()
296                            .enumerate()
297                            .map(|(ix, (buffer, range))| {
298                                (
299                                    buffer,
300                                    range,
301                                    (PLACEHOLDER.len() * ix)..(PLACEHOLDER.len() * (ix + 1) - 1),
302                                )
303                            })
304                            .collect::<Vec<_>>();
305
306                        let new_text: String = PLACEHOLDER.repeat(selections.len());
307
308                        let callback = Arc::new({
309                            let source_range = source_range.clone();
310                            move |_, window: &mut Window, cx: &mut App| {
311                                let selections = selections.clone();
312                                let message_editor = message_editor.clone();
313                                let source_range = source_range.clone();
314                                window.defer(cx, move |window, cx| {
315                                    message_editor
316                                        .update(cx, |message_editor, cx| {
317                                            message_editor.confirm_mention_for_selection(
318                                                source_range,
319                                                selections,
320                                                window,
321                                                cx,
322                                            )
323                                        })
324                                        .ok();
325                                });
326                                false
327                            }
328                        });
329
330                        (new_text, callback)
331                    }
332                };
333
334                Some(Completion {
335                    replace_range: source_range.clone(),
336                    new_text,
337                    label: CodeLabel::plain(action.label().to_string(), None),
338                    icon_path: Some(action.icon().path().into()),
339                    documentation: None,
340                    source: project::CompletionSource::Custom,
341                    insert_text_mode: None,
342                    // This ensures that when a user accepts this completion, the
343                    // completion menu will still be shown after "@category " is
344                    // inserted
345                    confirm: Some(on_action),
346                })
347            }
348        }
349    }
350
351    fn completion_for_thread(
352        thread_entry: ThreadContextEntry,
353        source_range: Range<Anchor>,
354        recent: bool,
355        editor: WeakEntity<MessageEditor>,
356        cx: &mut App,
357    ) -> Completion {
358        let uri = match &thread_entry {
359            ThreadContextEntry::Thread { id, title } => MentionUri::Thread {
360                id: id.clone(),
361                name: title.to_string(),
362            },
363            ThreadContextEntry::Context { path, title } => MentionUri::TextThread {
364                path: path.to_path_buf(),
365                name: title.to_string(),
366            },
367        };
368
369        let icon_for_completion = if recent {
370            IconName::HistoryRerun.path().into()
371        } else {
372            uri.icon_path(cx)
373        };
374
375        let new_text = format!("{} ", uri.as_link());
376
377        let new_text_len = new_text.len();
378        Completion {
379            replace_range: source_range.clone(),
380            new_text,
381            label: CodeLabel::plain(thread_entry.title().to_string(), None),
382            documentation: None,
383            insert_text_mode: None,
384            source: project::CompletionSource::Custom,
385            icon_path: Some(icon_for_completion.clone()),
386            confirm: Some(confirm_completion_callback(
387                thread_entry.title().clone(),
388                source_range.start,
389                new_text_len - 1,
390                editor,
391                uri,
392            )),
393        }
394    }
395
396    fn completion_for_rules(
397        rule: RulesContextEntry,
398        source_range: Range<Anchor>,
399        editor: WeakEntity<MessageEditor>,
400        cx: &mut App,
401    ) -> Completion {
402        let uri = MentionUri::Rule {
403            id: rule.prompt_id.into(),
404            name: rule.title.to_string(),
405        };
406        let new_text = format!("{} ", uri.as_link());
407        let new_text_len = new_text.len();
408        let icon_path = uri.icon_path(cx);
409        Completion {
410            replace_range: source_range.clone(),
411            new_text,
412            label: CodeLabel::plain(rule.title.to_string(), None),
413            documentation: None,
414            insert_text_mode: None,
415            source: project::CompletionSource::Custom,
416            icon_path: Some(icon_path.clone()),
417            confirm: Some(confirm_completion_callback(
418                rule.title.clone(),
419                source_range.start,
420                new_text_len - 1,
421                editor,
422                uri,
423            )),
424        }
425    }
426
427    pub(crate) fn completion_for_path(
428        project_path: ProjectPath,
429        path_prefix: &str,
430        is_recent: bool,
431        is_directory: bool,
432        source_range: Range<Anchor>,
433        message_editor: WeakEntity<MessageEditor>,
434        project: Entity<Project>,
435        cx: &mut App,
436    ) -> Option<Completion> {
437        let (file_name, directory) =
438            crate::context_picker::file_context_picker::extract_file_name_and_directory(
439                &project_path.path,
440                path_prefix,
441            );
442
443        let label =
444            build_code_label_for_full_path(&file_name, directory.as_ref().map(|s| s.as_ref()), cx);
445
446        let abs_path = project.read(cx).absolute_path(&project_path, cx)?;
447
448        let file_uri = MentionUri::File {
449            abs_path,
450            is_directory,
451        };
452
453        let crease_icon_path = file_uri.icon_path(cx);
454        let completion_icon_path = if is_recent {
455            IconName::HistoryRerun.path().into()
456        } else {
457            crease_icon_path.clone()
458        };
459
460        let new_text = format!("{} ", file_uri.as_link());
461        let new_text_len = new_text.len();
462        Some(Completion {
463            replace_range: source_range.clone(),
464            new_text,
465            label,
466            documentation: None,
467            source: project::CompletionSource::Custom,
468            icon_path: Some(completion_icon_path),
469            insert_text_mode: None,
470            confirm: Some(confirm_completion_callback(
471                file_name,
472                source_range.start,
473                new_text_len - 1,
474                message_editor,
475                file_uri,
476            )),
477        })
478    }
479
480    fn completion_for_symbol(
481        symbol: Symbol,
482        source_range: Range<Anchor>,
483        message_editor: WeakEntity<MessageEditor>,
484        workspace: Entity<Workspace>,
485        cx: &mut App,
486    ) -> Option<Completion> {
487        let project = workspace.read(cx).project().clone();
488
489        let label = CodeLabel::plain(symbol.name.clone(), None);
490
491        let abs_path = project.read(cx).absolute_path(&symbol.path, cx)?;
492        let uri = MentionUri::Symbol {
493            path: abs_path,
494            name: symbol.name.clone(),
495            line_range: symbol.range.start.0.row..symbol.range.end.0.row,
496        };
497        let new_text = format!("{} ", uri.as_link());
498        let new_text_len = new_text.len();
499        let icon_path = uri.icon_path(cx);
500        Some(Completion {
501            replace_range: source_range.clone(),
502            new_text,
503            label,
504            documentation: None,
505            source: project::CompletionSource::Custom,
506            icon_path: Some(icon_path.clone()),
507            insert_text_mode: None,
508            confirm: Some(confirm_completion_callback(
509                symbol.name.clone().into(),
510                source_range.start,
511                new_text_len - 1,
512                message_editor,
513                uri,
514            )),
515        })
516    }
517
518    fn completion_for_fetch(
519        source_range: Range<Anchor>,
520        url_to_fetch: SharedString,
521        message_editor: WeakEntity<MessageEditor>,
522        cx: &mut App,
523    ) -> Option<Completion> {
524        let new_text = format!("@fetch {} ", url_to_fetch.clone());
525        let url_to_fetch = url::Url::parse(url_to_fetch.as_ref())
526            .or_else(|_| url::Url::parse(&format!("https://{url_to_fetch}")))
527            .ok()?;
528        let mention_uri = MentionUri::Fetch {
529            url: url_to_fetch.clone(),
530        };
531        let icon_path = mention_uri.icon_path(cx);
532        Some(Completion {
533            replace_range: source_range.clone(),
534            new_text: new_text.clone(),
535            label: CodeLabel::plain(url_to_fetch.to_string(), None),
536            documentation: None,
537            source: project::CompletionSource::Custom,
538            icon_path: Some(icon_path.clone()),
539            insert_text_mode: None,
540            confirm: Some(confirm_completion_callback(
541                url_to_fetch.to_string().into(),
542                source_range.start,
543                new_text.len() - 1,
544                message_editor,
545                mention_uri,
546            )),
547        })
548    }
549}
550
551fn build_code_label_for_full_path(file_name: &str, directory: Option<&str>, cx: &App) -> CodeLabel {
552    let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
553    let mut label = CodeLabel::default();
554
555    label.push_str(&file_name, None);
556    label.push_str(" ", None);
557
558    if let Some(directory) = directory {
559        label.push_str(&directory, comment_id);
560    }
561
562    label.filter_range = 0..label.text().len();
563
564    label
565}
566
567impl CompletionProvider for ContextPickerCompletionProvider {
568    fn completions(
569        &self,
570        _excerpt_id: ExcerptId,
571        buffer: &Entity<Buffer>,
572        buffer_position: Anchor,
573        _trigger: CompletionContext,
574        _window: &mut Window,
575        cx: &mut Context<Editor>,
576    ) -> Task<Result<Vec<CompletionResponse>>> {
577        let state = buffer.update(cx, |buffer, _cx| {
578            let position = buffer_position.to_point(buffer);
579            let line_start = Point::new(position.row, 0);
580            let offset_to_line = buffer.point_to_offset(line_start);
581            let mut lines = buffer.text_for_range(line_start..position).lines();
582            let line = lines.next()?;
583            MentionCompletion::try_parse(line, offset_to_line)
584        });
585        let Some(state) = state else {
586            return Task::ready(Ok(Vec::new()));
587        };
588
589        let Some(workspace) = self.workspace.upgrade() else {
590            return Task::ready(Ok(Vec::new()));
591        };
592
593        let project = workspace.read(cx).project().clone();
594        let snapshot = buffer.read(cx).snapshot();
595        let source_range = snapshot.anchor_before(state.source_range.start)
596            ..snapshot.anchor_after(state.source_range.end);
597
598        let thread_store = self.thread_store.clone();
599        let text_thread_store = self.text_thread_store.clone();
600        let editor = self.message_editor.clone();
601        let Ok((exclude_paths, exclude_threads)) =
602            self.message_editor.update(cx, |message_editor, _cx| {
603                message_editor.mentioned_path_and_threads()
604            })
605        else {
606            return Task::ready(Ok(Vec::new()));
607        };
608
609        let MentionCompletion { mode, argument, .. } = state;
610        let query = argument.unwrap_or_else(|| "".to_string());
611
612        let recent_entries = recent_context_picker_entries(
613            Some(thread_store.clone()),
614            Some(text_thread_store.clone()),
615            workspace.clone(),
616            &exclude_paths,
617            &exclude_threads,
618            cx,
619        );
620
621        let prompt_store = thread_store
622            .read_with(cx, |thread_store, _cx| thread_store.prompt_store().clone())
623            .ok()
624            .flatten();
625
626        let search_task = search(
627            mode,
628            query,
629            Arc::<AtomicBool>::default(),
630            recent_entries,
631            prompt_store,
632            thread_store.clone(),
633            text_thread_store.clone(),
634            workspace.clone(),
635            cx,
636        );
637
638        cx.spawn(async move |_, cx| {
639            let matches = search_task.await;
640
641            let completions = cx.update(|cx| {
642                matches
643                    .into_iter()
644                    .filter_map(|mat| match mat {
645                        Match::File(FileMatch { mat, is_recent }) => {
646                            let project_path = ProjectPath {
647                                worktree_id: WorktreeId::from_usize(mat.worktree_id),
648                                path: mat.path.clone(),
649                            };
650
651                            Self::completion_for_path(
652                                project_path,
653                                &mat.path_prefix,
654                                is_recent,
655                                mat.is_dir,
656                                source_range.clone(),
657                                editor.clone(),
658                                project.clone(),
659                                cx,
660                            )
661                        }
662
663                        Match::Symbol(SymbolMatch { symbol, .. }) => Self::completion_for_symbol(
664                            symbol,
665                            source_range.clone(),
666                            editor.clone(),
667                            workspace.clone(),
668                            cx,
669                        ),
670
671                        Match::Thread(ThreadMatch {
672                            thread, is_recent, ..
673                        }) => Some(Self::completion_for_thread(
674                            thread,
675                            source_range.clone(),
676                            is_recent,
677                            editor.clone(),
678                            cx,
679                        )),
680
681                        Match::Rules(user_rules) => Some(Self::completion_for_rules(
682                            user_rules,
683                            source_range.clone(),
684                            editor.clone(),
685                            cx,
686                        )),
687
688                        Match::Fetch(url) => Self::completion_for_fetch(
689                            source_range.clone(),
690                            url,
691                            editor.clone(),
692                            cx,
693                        ),
694
695                        Match::Entry(EntryMatch { entry, .. }) => Self::completion_for_entry(
696                            entry,
697                            source_range.clone(),
698                            editor.clone(),
699                            &workspace,
700                            cx,
701                        ),
702                    })
703                    .collect()
704            })?;
705
706            Ok(vec![CompletionResponse {
707                completions,
708                // Since this does its own filtering (see `filter_completions()` returns false),
709                // there is no benefit to computing whether this set of completions is incomplete.
710                is_incomplete: true,
711            }])
712        })
713    }
714
715    fn is_completion_trigger(
716        &self,
717        buffer: &Entity<language::Buffer>,
718        position: language::Anchor,
719        _text: &str,
720        _trigger_in_words: bool,
721        _menu_is_open: bool,
722        cx: &mut Context<Editor>,
723    ) -> bool {
724        let buffer = buffer.read(cx);
725        let position = position.to_point(buffer);
726        let line_start = Point::new(position.row, 0);
727        let offset_to_line = buffer.point_to_offset(line_start);
728        let mut lines = buffer.text_for_range(line_start..position).lines();
729        if let Some(line) = lines.next() {
730            MentionCompletion::try_parse(line, offset_to_line)
731                .map(|completion| {
732                    completion.source_range.start <= offset_to_line + position.column as usize
733                        && completion.source_range.end >= offset_to_line + position.column as usize
734                })
735                .unwrap_or(false)
736        } else {
737            false
738        }
739    }
740
741    fn sort_completions(&self) -> bool {
742        false
743    }
744
745    fn filter_completions(&self) -> bool {
746        false
747    }
748}
749
750fn confirm_completion_callback(
751    crease_text: SharedString,
752    start: Anchor,
753    content_len: usize,
754    message_editor: WeakEntity<MessageEditor>,
755    mention_uri: MentionUri,
756) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
757    Arc::new(move |_, window, cx| {
758        let message_editor = message_editor.clone();
759        let crease_text = crease_text.clone();
760        let mention_uri = mention_uri.clone();
761        window.defer(cx, move |window, cx| {
762            message_editor
763                .clone()
764                .update(cx, |message_editor, cx| {
765                    message_editor.confirm_completion(
766                        crease_text,
767                        start,
768                        content_len,
769                        mention_uri,
770                        window,
771                        cx,
772                    )
773                })
774                .ok();
775        });
776        false
777    })
778}
779
780#[derive(Debug, Default, PartialEq)]
781struct MentionCompletion {
782    source_range: Range<usize>,
783    mode: Option<ContextPickerMode>,
784    argument: Option<String>,
785}
786
787impl MentionCompletion {
788    fn try_parse(line: &str, offset_to_line: usize) -> Option<Self> {
789        let last_mention_start = line.rfind('@')?;
790        if last_mention_start >= line.len() {
791            return Some(Self::default());
792        }
793        if last_mention_start > 0
794            && line
795                .chars()
796                .nth(last_mention_start - 1)
797                .map_or(false, |c| !c.is_whitespace())
798        {
799            return None;
800        }
801
802        let rest_of_line = &line[last_mention_start + 1..];
803
804        let mut mode = None;
805        let mut argument = None;
806
807        let mut parts = rest_of_line.split_whitespace();
808        let mut end = last_mention_start + 1;
809        if let Some(mode_text) = parts.next() {
810            end += mode_text.len();
811
812            if let Some(parsed_mode) = ContextPickerMode::try_from(mode_text).ok() {
813                mode = Some(parsed_mode);
814            } else {
815                argument = Some(mode_text.to_string());
816            }
817            match rest_of_line[mode_text.len()..].find(|c: char| !c.is_whitespace()) {
818                Some(whitespace_count) => {
819                    if let Some(argument_text) = parts.next() {
820                        argument = Some(argument_text.to_string());
821                        end += whitespace_count + argument_text.len();
822                    }
823                }
824                None => {
825                    // Rest of line is entirely whitespace
826                    end += rest_of_line.len() - mode_text.len();
827                }
828            }
829        }
830
831        Some(Self {
832            source_range: last_mention_start + offset_to_line..end + offset_to_line,
833            mode,
834            argument,
835        })
836    }
837}
838
839#[cfg(test)]
840mod tests {
841    use super::*;
842
843    #[test]
844    fn test_mention_completion_parse() {
845        assert_eq!(MentionCompletion::try_parse("Lorem Ipsum", 0), None);
846
847        assert_eq!(
848            MentionCompletion::try_parse("Lorem @", 0),
849            Some(MentionCompletion {
850                source_range: 6..7,
851                mode: None,
852                argument: None,
853            })
854        );
855
856        assert_eq!(
857            MentionCompletion::try_parse("Lorem @file", 0),
858            Some(MentionCompletion {
859                source_range: 6..11,
860                mode: Some(ContextPickerMode::File),
861                argument: None,
862            })
863        );
864
865        assert_eq!(
866            MentionCompletion::try_parse("Lorem @file ", 0),
867            Some(MentionCompletion {
868                source_range: 6..12,
869                mode: Some(ContextPickerMode::File),
870                argument: None,
871            })
872        );
873
874        assert_eq!(
875            MentionCompletion::try_parse("Lorem @file main.rs", 0),
876            Some(MentionCompletion {
877                source_range: 6..19,
878                mode: Some(ContextPickerMode::File),
879                argument: Some("main.rs".to_string()),
880            })
881        );
882
883        assert_eq!(
884            MentionCompletion::try_parse("Lorem @file main.rs ", 0),
885            Some(MentionCompletion {
886                source_range: 6..19,
887                mode: Some(ContextPickerMode::File),
888                argument: Some("main.rs".to_string()),
889            })
890        );
891
892        assert_eq!(
893            MentionCompletion::try_parse("Lorem @file main.rs Ipsum", 0),
894            Some(MentionCompletion {
895                source_range: 6..19,
896                mode: Some(ContextPickerMode::File),
897                argument: Some("main.rs".to_string()),
898            })
899        );
900
901        assert_eq!(
902            MentionCompletion::try_parse("Lorem @main", 0),
903            Some(MentionCompletion {
904                source_range: 6..11,
905                mode: None,
906                argument: Some("main".to_string()),
907            })
908        );
909
910        assert_eq!(MentionCompletion::try_parse("test@", 0), None);
911    }
912}