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 uri = if is_directory {
449            MentionUri::Directory { abs_path }
450        } else {
451            MentionUri::File { abs_path }
452        };
453
454        let crease_icon_path = uri.icon_path(cx);
455        let completion_icon_path = if is_recent {
456            IconName::HistoryRerun.path().into()
457        } else {
458            crease_icon_path.clone()
459        };
460
461        let new_text = format!("{} ", uri.as_link());
462        let new_text_len = new_text.len();
463        Some(Completion {
464            replace_range: source_range.clone(),
465            new_text,
466            label,
467            documentation: None,
468            source: project::CompletionSource::Custom,
469            icon_path: Some(completion_icon_path),
470            insert_text_mode: None,
471            confirm: Some(confirm_completion_callback(
472                file_name,
473                source_range.start,
474                new_text_len - 1,
475                message_editor,
476                uri,
477            )),
478        })
479    }
480
481    fn completion_for_symbol(
482        symbol: Symbol,
483        source_range: Range<Anchor>,
484        message_editor: WeakEntity<MessageEditor>,
485        workspace: Entity<Workspace>,
486        cx: &mut App,
487    ) -> Option<Completion> {
488        let project = workspace.read(cx).project().clone();
489
490        let label = CodeLabel::plain(symbol.name.clone(), None);
491
492        let abs_path = project.read(cx).absolute_path(&symbol.path, cx)?;
493        let uri = MentionUri::Symbol {
494            path: abs_path,
495            name: symbol.name.clone(),
496            line_range: symbol.range.start.0.row..symbol.range.end.0.row,
497        };
498        let new_text = format!("{} ", uri.as_link());
499        let new_text_len = new_text.len();
500        let icon_path = uri.icon_path(cx);
501        Some(Completion {
502            replace_range: source_range.clone(),
503            new_text,
504            label,
505            documentation: None,
506            source: project::CompletionSource::Custom,
507            icon_path: Some(icon_path.clone()),
508            insert_text_mode: None,
509            confirm: Some(confirm_completion_callback(
510                symbol.name.clone().into(),
511                source_range.start,
512                new_text_len - 1,
513                message_editor,
514                uri,
515            )),
516        })
517    }
518
519    fn completion_for_fetch(
520        source_range: Range<Anchor>,
521        url_to_fetch: SharedString,
522        message_editor: WeakEntity<MessageEditor>,
523        cx: &mut App,
524    ) -> Option<Completion> {
525        let new_text = format!("@fetch {} ", url_to_fetch.clone());
526        let url_to_fetch = url::Url::parse(url_to_fetch.as_ref())
527            .or_else(|_| url::Url::parse(&format!("https://{url_to_fetch}")))
528            .ok()?;
529        let mention_uri = MentionUri::Fetch {
530            url: url_to_fetch.clone(),
531        };
532        let icon_path = mention_uri.icon_path(cx);
533        Some(Completion {
534            replace_range: source_range.clone(),
535            new_text: new_text.clone(),
536            label: CodeLabel::plain(url_to_fetch.to_string(), None),
537            documentation: None,
538            source: project::CompletionSource::Custom,
539            icon_path: Some(icon_path.clone()),
540            insert_text_mode: None,
541            confirm: Some(confirm_completion_callback(
542                url_to_fetch.to_string().into(),
543                source_range.start,
544                new_text.len() - 1,
545                message_editor,
546                mention_uri,
547            )),
548        })
549    }
550}
551
552fn build_code_label_for_full_path(file_name: &str, directory: Option<&str>, cx: &App) -> CodeLabel {
553    let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
554    let mut label = CodeLabel::default();
555
556    label.push_str(file_name, None);
557    label.push_str(" ", None);
558
559    if let Some(directory) = directory {
560        label.push_str(directory, comment_id);
561    }
562
563    label.filter_range = 0..label.text().len();
564
565    label
566}
567
568impl CompletionProvider for ContextPickerCompletionProvider {
569    fn completions(
570        &self,
571        _excerpt_id: ExcerptId,
572        buffer: &Entity<Buffer>,
573        buffer_position: Anchor,
574        _trigger: CompletionContext,
575        _window: &mut Window,
576        cx: &mut Context<Editor>,
577    ) -> Task<Result<Vec<CompletionResponse>>> {
578        let state = buffer.update(cx, |buffer, _cx| {
579            let position = buffer_position.to_point(buffer);
580            let line_start = Point::new(position.row, 0);
581            let offset_to_line = buffer.point_to_offset(line_start);
582            let mut lines = buffer.text_for_range(line_start..position).lines();
583            let line = lines.next()?;
584            MentionCompletion::try_parse(line, offset_to_line)
585        });
586        let Some(state) = state else {
587            return Task::ready(Ok(Vec::new()));
588        };
589
590        let Some(workspace) = self.workspace.upgrade() else {
591            return Task::ready(Ok(Vec::new()));
592        };
593
594        let project = workspace.read(cx).project().clone();
595        let snapshot = buffer.read(cx).snapshot();
596        let source_range = snapshot.anchor_before(state.source_range.start)
597            ..snapshot.anchor_after(state.source_range.end);
598
599        let thread_store = self.thread_store.clone();
600        let text_thread_store = self.text_thread_store.clone();
601        let editor = self.message_editor.clone();
602        let Ok((exclude_paths, exclude_threads)) =
603            self.message_editor.update(cx, |message_editor, _cx| {
604                message_editor.mentioned_path_and_threads()
605            })
606        else {
607            return Task::ready(Ok(Vec::new()));
608        };
609
610        let MentionCompletion { mode, argument, .. } = state;
611        let query = argument.unwrap_or_else(|| "".to_string());
612
613        let recent_entries = recent_context_picker_entries(
614            Some(thread_store.clone()),
615            Some(text_thread_store.clone()),
616            workspace.clone(),
617            &exclude_paths,
618            &exclude_threads,
619            cx,
620        );
621
622        let prompt_store = thread_store
623            .read_with(cx, |thread_store, _cx| thread_store.prompt_store().clone())
624            .ok()
625            .flatten();
626
627        let search_task = search(
628            mode,
629            query,
630            Arc::<AtomicBool>::default(),
631            recent_entries,
632            prompt_store,
633            thread_store.clone(),
634            text_thread_store.clone(),
635            workspace.clone(),
636            cx,
637        );
638
639        cx.spawn(async move |_, cx| {
640            let matches = search_task.await;
641
642            let completions = cx.update(|cx| {
643                matches
644                    .into_iter()
645                    .filter_map(|mat| match mat {
646                        Match::File(FileMatch { mat, is_recent }) => {
647                            let project_path = ProjectPath {
648                                worktree_id: WorktreeId::from_usize(mat.worktree_id),
649                                path: mat.path.clone(),
650                            };
651
652                            Self::completion_for_path(
653                                project_path,
654                                &mat.path_prefix,
655                                is_recent,
656                                mat.is_dir,
657                                source_range.clone(),
658                                editor.clone(),
659                                project.clone(),
660                                cx,
661                            )
662                        }
663
664                        Match::Symbol(SymbolMatch { symbol, .. }) => Self::completion_for_symbol(
665                            symbol,
666                            source_range.clone(),
667                            editor.clone(),
668                            workspace.clone(),
669                            cx,
670                        ),
671
672                        Match::Thread(ThreadMatch {
673                            thread, is_recent, ..
674                        }) => Some(Self::completion_for_thread(
675                            thread,
676                            source_range.clone(),
677                            is_recent,
678                            editor.clone(),
679                            cx,
680                        )),
681
682                        Match::Rules(user_rules) => Some(Self::completion_for_rules(
683                            user_rules,
684                            source_range.clone(),
685                            editor.clone(),
686                            cx,
687                        )),
688
689                        Match::Fetch(url) => Self::completion_for_fetch(
690                            source_range.clone(),
691                            url,
692                            editor.clone(),
693                            cx,
694                        ),
695
696                        Match::Entry(EntryMatch { entry, .. }) => Self::completion_for_entry(
697                            entry,
698                            source_range.clone(),
699                            editor.clone(),
700                            &workspace,
701                            cx,
702                        ),
703                    })
704                    .collect()
705            })?;
706
707            Ok(vec![CompletionResponse {
708                completions,
709                // Since this does its own filtering (see `filter_completions()` returns false),
710                // there is no benefit to computing whether this set of completions is incomplete.
711                is_incomplete: true,
712            }])
713        })
714    }
715
716    fn is_completion_trigger(
717        &self,
718        buffer: &Entity<language::Buffer>,
719        position: language::Anchor,
720        _text: &str,
721        _trigger_in_words: bool,
722        _menu_is_open: bool,
723        cx: &mut Context<Editor>,
724    ) -> bool {
725        let buffer = buffer.read(cx);
726        let position = position.to_point(buffer);
727        let line_start = Point::new(position.row, 0);
728        let offset_to_line = buffer.point_to_offset(line_start);
729        let mut lines = buffer.text_for_range(line_start..position).lines();
730        if let Some(line) = lines.next() {
731            MentionCompletion::try_parse(line, offset_to_line)
732                .map(|completion| {
733                    completion.source_range.start <= offset_to_line + position.column as usize
734                        && completion.source_range.end >= offset_to_line + position.column as usize
735                })
736                .unwrap_or(false)
737        } else {
738            false
739        }
740    }
741
742    fn sort_completions(&self) -> bool {
743        false
744    }
745
746    fn filter_completions(&self) -> bool {
747        false
748    }
749}
750
751fn confirm_completion_callback(
752    crease_text: SharedString,
753    start: Anchor,
754    content_len: usize,
755    message_editor: WeakEntity<MessageEditor>,
756    mention_uri: MentionUri,
757) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
758    Arc::new(move |_, window, cx| {
759        let message_editor = message_editor.clone();
760        let crease_text = crease_text.clone();
761        let mention_uri = mention_uri.clone();
762        window.defer(cx, move |window, cx| {
763            message_editor
764                .clone()
765                .update(cx, |message_editor, cx| {
766                    message_editor
767                        .confirm_completion(
768                            crease_text,
769                            start,
770                            content_len,
771                            mention_uri,
772                            window,
773                            cx,
774                        )
775                        .detach();
776                })
777                .ok();
778        });
779        false
780    })
781}
782
783#[derive(Debug, Default, PartialEq)]
784struct MentionCompletion {
785    source_range: Range<usize>,
786    mode: Option<ContextPickerMode>,
787    argument: Option<String>,
788}
789
790impl MentionCompletion {
791    fn try_parse(line: &str, offset_to_line: usize) -> Option<Self> {
792        let last_mention_start = line.rfind('@')?;
793        if last_mention_start >= line.len() {
794            return Some(Self::default());
795        }
796        if last_mention_start > 0
797            && line
798                .chars()
799                .nth(last_mention_start - 1)
800                .is_some_and(|c| !c.is_whitespace())
801        {
802            return None;
803        }
804
805        let rest_of_line = &line[last_mention_start + 1..];
806
807        let mut mode = None;
808        let mut argument = None;
809
810        let mut parts = rest_of_line.split_whitespace();
811        let mut end = last_mention_start + 1;
812        if let Some(mode_text) = parts.next() {
813            end += mode_text.len();
814
815            if let Some(parsed_mode) = ContextPickerMode::try_from(mode_text).ok() {
816                mode = Some(parsed_mode);
817            } else {
818                argument = Some(mode_text.to_string());
819            }
820            match rest_of_line[mode_text.len()..].find(|c: char| !c.is_whitespace()) {
821                Some(whitespace_count) => {
822                    if let Some(argument_text) = parts.next() {
823                        argument = Some(argument_text.to_string());
824                        end += whitespace_count + argument_text.len();
825                    }
826                }
827                None => {
828                    // Rest of line is entirely whitespace
829                    end += rest_of_line.len() - mode_text.len();
830                }
831            }
832        }
833
834        Some(Self {
835            source_range: last_mention_start + offset_to_line..end + offset_to_line,
836            mode,
837            argument,
838        })
839    }
840}
841
842#[cfg(test)]
843mod tests {
844    use super::*;
845
846    #[test]
847    fn test_mention_completion_parse() {
848        assert_eq!(MentionCompletion::try_parse("Lorem Ipsum", 0), None);
849
850        assert_eq!(
851            MentionCompletion::try_parse("Lorem @", 0),
852            Some(MentionCompletion {
853                source_range: 6..7,
854                mode: None,
855                argument: None,
856            })
857        );
858
859        assert_eq!(
860            MentionCompletion::try_parse("Lorem @file", 0),
861            Some(MentionCompletion {
862                source_range: 6..11,
863                mode: Some(ContextPickerMode::File),
864                argument: None,
865            })
866        );
867
868        assert_eq!(
869            MentionCompletion::try_parse("Lorem @file ", 0),
870            Some(MentionCompletion {
871                source_range: 6..12,
872                mode: Some(ContextPickerMode::File),
873                argument: None,
874            })
875        );
876
877        assert_eq!(
878            MentionCompletion::try_parse("Lorem @file main.rs", 0),
879            Some(MentionCompletion {
880                source_range: 6..19,
881                mode: Some(ContextPickerMode::File),
882                argument: Some("main.rs".to_string()),
883            })
884        );
885
886        assert_eq!(
887            MentionCompletion::try_parse("Lorem @file main.rs ", 0),
888            Some(MentionCompletion {
889                source_range: 6..19,
890                mode: Some(ContextPickerMode::File),
891                argument: Some("main.rs".to_string()),
892            })
893        );
894
895        assert_eq!(
896            MentionCompletion::try_parse("Lorem @file main.rs Ipsum", 0),
897            Some(MentionCompletion {
898                source_range: 6..19,
899                mode: Some(ContextPickerMode::File),
900                argument: Some("main.rs".to_string()),
901            })
902        );
903
904        assert_eq!(
905            MentionCompletion::try_parse("Lorem @main", 0),
906            Some(MentionCompletion {
907                source_range: 6..11,
908                mode: None,
909                argument: Some("main".to_string()),
910            })
911        );
912
913        assert_eq!(MentionCompletion::try_parse("test@", 0), None);
914    }
915}