completion_provider.rs

   1use std::cell::RefCell;
   2use std::ops::Range;
   3use std::path::{Path, PathBuf};
   4use std::rc::Rc;
   5use std::sync::Arc;
   6use std::sync::atomic::AtomicBool;
   7
   8use anyhow::Result;
   9use editor::{CompletionProvider, Editor, ExcerptId, ToOffset as _};
  10use file_icons::FileIcons;
  11use fuzzy::{StringMatch, StringMatchCandidate};
  12use gpui::{App, Entity, Task, WeakEntity};
  13use http_client::HttpClientWithUrl;
  14use itertools::Itertools;
  15use language::{Buffer, CodeLabel, HighlightId};
  16use lsp::CompletionContext;
  17use project::{Completion, CompletionIntent, ProjectPath, Symbol, WorktreeId};
  18use prompt_store::PromptStore;
  19use rope::Point;
  20use text::{Anchor, OffsetRangeExt, ToPoint};
  21use ui::prelude::*;
  22use util::ResultExt as _;
  23use workspace::Workspace;
  24
  25use crate::Thread;
  26use crate::context::{AgentContextHandle, AgentContextKey, ContextCreasesAddon, RULES_ICON};
  27use crate::context_store::ContextStore;
  28use crate::thread_store::ThreadStore;
  29
  30use super::fetch_context_picker::fetch_url_content;
  31use super::file_context_picker::{FileMatch, search_files};
  32use super::rules_context_picker::{RulesContextEntry, search_rules};
  33use super::symbol_context_picker::SymbolMatch;
  34use super::symbol_context_picker::search_symbols;
  35use super::thread_context_picker::{ThreadContextEntry, ThreadMatch, search_threads};
  36use super::{
  37    ContextPickerAction, ContextPickerEntry, ContextPickerMode, MentionLink, RecentEntry,
  38    available_context_picker_entries, recent_context_picker_entries, selection_ranges,
  39};
  40
  41pub(crate) enum Match {
  42    File(FileMatch),
  43    Symbol(SymbolMatch),
  44    Thread(ThreadMatch),
  45    Fetch(SharedString),
  46    Rules(RulesContextEntry),
  47    Entry(EntryMatch),
  48}
  49
  50pub struct EntryMatch {
  51    mat: Option<StringMatch>,
  52    entry: ContextPickerEntry,
  53}
  54
  55impl Match {
  56    pub fn score(&self) -> f64 {
  57        match self {
  58            Match::File(file) => file.mat.score,
  59            Match::Entry(mode) => mode.mat.as_ref().map(|mat| mat.score).unwrap_or(1.),
  60            Match::Thread(_) => 1.,
  61            Match::Symbol(_) => 1.,
  62            Match::Fetch(_) => 1.,
  63            Match::Rules(_) => 1.,
  64        }
  65    }
  66}
  67
  68fn search(
  69    mode: Option<ContextPickerMode>,
  70    query: String,
  71    cancellation_flag: Arc<AtomicBool>,
  72    recent_entries: Vec<RecentEntry>,
  73    prompt_store: Option<Entity<PromptStore>>,
  74    thread_store: Option<WeakEntity<ThreadStore>>,
  75    workspace: Entity<Workspace>,
  76    cx: &mut App,
  77) -> Task<Vec<Match>> {
  78    match mode {
  79        Some(ContextPickerMode::File) => {
  80            let search_files_task =
  81                search_files(query.clone(), cancellation_flag.clone(), &workspace, cx);
  82            cx.background_spawn(async move {
  83                search_files_task
  84                    .await
  85                    .into_iter()
  86                    .map(Match::File)
  87                    .collect()
  88            })
  89        }
  90
  91        Some(ContextPickerMode::Symbol) => {
  92            let search_symbols_task =
  93                search_symbols(query.clone(), cancellation_flag.clone(), &workspace, cx);
  94            cx.background_spawn(async move {
  95                search_symbols_task
  96                    .await
  97                    .into_iter()
  98                    .map(Match::Symbol)
  99                    .collect()
 100            })
 101        }
 102
 103        Some(ContextPickerMode::Thread) => {
 104            if let Some(thread_store) = thread_store.as_ref().and_then(|t| t.upgrade()) {
 105                let search_threads_task =
 106                    search_threads(query.clone(), cancellation_flag.clone(), thread_store, cx);
 107                cx.background_spawn(async move {
 108                    search_threads_task
 109                        .await
 110                        .into_iter()
 111                        .map(Match::Thread)
 112                        .collect()
 113                })
 114            } else {
 115                Task::ready(Vec::new())
 116            }
 117        }
 118
 119        Some(ContextPickerMode::Fetch) => {
 120            if !query.is_empty() {
 121                Task::ready(vec![Match::Fetch(query.into())])
 122            } else {
 123                Task::ready(Vec::new())
 124            }
 125        }
 126
 127        Some(ContextPickerMode::Rules) => {
 128            if let Some(prompt_store) = prompt_store.as_ref() {
 129                let search_rules_task =
 130                    search_rules(query.clone(), cancellation_flag.clone(), prompt_store, cx);
 131                cx.background_spawn(async move {
 132                    search_rules_task
 133                        .await
 134                        .into_iter()
 135                        .map(Match::Rules)
 136                        .collect::<Vec<_>>()
 137                })
 138            } else {
 139                Task::ready(Vec::new())
 140            }
 141        }
 142
 143        None => {
 144            if query.is_empty() {
 145                let mut matches = recent_entries
 146                    .into_iter()
 147                    .map(|entry| match entry {
 148                        super::RecentEntry::File {
 149                            project_path,
 150                            path_prefix,
 151                        } => Match::File(FileMatch {
 152                            mat: fuzzy::PathMatch {
 153                                score: 1.,
 154                                positions: Vec::new(),
 155                                worktree_id: project_path.worktree_id.to_usize(),
 156                                path: project_path.path,
 157                                path_prefix,
 158                                is_dir: false,
 159                                distance_to_relative_ancestor: 0,
 160                            },
 161                            is_recent: true,
 162                        }),
 163                        super::RecentEntry::Thread(thread_context_entry) => {
 164                            Match::Thread(ThreadMatch {
 165                                thread: thread_context_entry,
 166                                is_recent: true,
 167                            })
 168                        }
 169                    })
 170                    .collect::<Vec<_>>();
 171
 172                matches.extend(
 173                    available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx)
 174                        .into_iter()
 175                        .map(|mode| {
 176                            Match::Entry(EntryMatch {
 177                                entry: mode,
 178                                mat: None,
 179                            })
 180                        }),
 181                );
 182
 183                Task::ready(matches)
 184            } else {
 185                let executor = cx.background_executor().clone();
 186
 187                let search_files_task =
 188                    search_files(query.clone(), cancellation_flag.clone(), &workspace, cx);
 189
 190                let entries =
 191                    available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx);
 192                let entry_candidates = entries
 193                    .iter()
 194                    .enumerate()
 195                    .map(|(ix, entry)| StringMatchCandidate::new(ix, entry.keyword()))
 196                    .collect::<Vec<_>>();
 197
 198                cx.background_spawn(async move {
 199                    let mut matches = search_files_task
 200                        .await
 201                        .into_iter()
 202                        .map(Match::File)
 203                        .collect::<Vec<_>>();
 204
 205                    let entry_matches = fuzzy::match_strings(
 206                        &entry_candidates,
 207                        &query,
 208                        false,
 209                        100,
 210                        &Arc::new(AtomicBool::default()),
 211                        executor,
 212                    )
 213                    .await;
 214
 215                    matches.extend(entry_matches.into_iter().map(|mat| {
 216                        Match::Entry(EntryMatch {
 217                            entry: entries[mat.candidate_id],
 218                            mat: Some(mat),
 219                        })
 220                    }));
 221
 222                    matches.sort_by(|a, b| {
 223                        b.score()
 224                            .partial_cmp(&a.score())
 225                            .unwrap_or(std::cmp::Ordering::Equal)
 226                    });
 227
 228                    matches
 229                })
 230            }
 231        }
 232    }
 233}
 234
 235pub struct ContextPickerCompletionProvider {
 236    workspace: WeakEntity<Workspace>,
 237    context_store: WeakEntity<ContextStore>,
 238    thread_store: Option<WeakEntity<ThreadStore>>,
 239    editor: WeakEntity<Editor>,
 240    excluded_buffer: Option<WeakEntity<Buffer>>,
 241}
 242
 243impl ContextPickerCompletionProvider {
 244    pub fn new(
 245        workspace: WeakEntity<Workspace>,
 246        context_store: WeakEntity<ContextStore>,
 247        thread_store: Option<WeakEntity<ThreadStore>>,
 248        editor: WeakEntity<Editor>,
 249        exclude_buffer: Option<WeakEntity<Buffer>>,
 250    ) -> Self {
 251        Self {
 252            workspace,
 253            context_store,
 254            thread_store,
 255            editor,
 256            excluded_buffer: exclude_buffer,
 257        }
 258    }
 259
 260    fn completion_for_entry(
 261        entry: ContextPickerEntry,
 262        excerpt_id: ExcerptId,
 263        source_range: Range<Anchor>,
 264        editor: Entity<Editor>,
 265        context_store: Entity<ContextStore>,
 266        workspace: &Entity<Workspace>,
 267        cx: &mut App,
 268    ) -> Option<Completion> {
 269        match entry {
 270            ContextPickerEntry::Mode(mode) => Some(Completion {
 271                replace_range: source_range.clone(),
 272                new_text: format!("@{} ", mode.keyword()),
 273                label: CodeLabel::plain(mode.label().to_string(), None),
 274                icon_path: Some(mode.icon().path().into()),
 275                documentation: None,
 276                source: project::CompletionSource::Custom,
 277                insert_text_mode: None,
 278                // This ensures that when a user accepts this completion, the
 279                // completion menu will still be shown after "@category " is
 280                // inserted
 281                confirm: Some(Arc::new(|_, _, _| true)),
 282            }),
 283            ContextPickerEntry::Action(action) => {
 284                let (new_text, on_action) = match action {
 285                    ContextPickerAction::AddSelections => {
 286                        let selections = selection_ranges(workspace, cx);
 287
 288                        let selection_infos = selections
 289                            .iter()
 290                            .map(|(buffer, range)| {
 291                                let full_path = buffer
 292                                    .read(cx)
 293                                    .file()
 294                                    .map(|file| file.full_path(cx))
 295                                    .unwrap_or_else(|| PathBuf::from("untitled"));
 296                                let file_name = full_path
 297                                    .file_name()
 298                                    .unwrap_or_default()
 299                                    .to_string_lossy()
 300                                    .to_string();
 301                                let line_range = range.to_point(&buffer.read(cx).snapshot());
 302
 303                                let link = MentionLink::for_selection(
 304                                    &file_name,
 305                                    &full_path.to_string_lossy(),
 306                                    line_range.start.row as usize..line_range.end.row as usize,
 307                                );
 308                                (file_name, link, line_range)
 309                            })
 310                            .collect::<Vec<_>>();
 311
 312                        let new_text = selection_infos.iter().map(|(_, link, _)| link).join(" ");
 313
 314                        let callback = Arc::new({
 315                            let context_store = context_store.clone();
 316                            let selections = selections.clone();
 317                            let selection_infos = selection_infos.clone();
 318                            move |_, window: &mut Window, cx: &mut App| {
 319                                context_store.update(cx, |context_store, cx| {
 320                                    for (buffer, range) in &selections {
 321                                        context_store.add_selection(
 322                                            buffer.clone(),
 323                                            range.clone(),
 324                                            cx,
 325                                        );
 326                                    }
 327                                });
 328
 329                                let editor = editor.clone();
 330                                let selection_infos = selection_infos.clone();
 331                                window.defer(cx, move |window, cx| {
 332                                    let mut current_offset = 0;
 333                                    for (file_name, link, line_range) in selection_infos.iter() {
 334                                        let snapshot =
 335                                            editor.read(cx).buffer().read(cx).snapshot(cx);
 336                                        let Some(start) = snapshot
 337                                            .anchor_in_excerpt(excerpt_id, source_range.start)
 338                                        else {
 339                                            return;
 340                                        };
 341
 342                                        let offset = start.to_offset(&snapshot) + current_offset;
 343                                        let text_len = link.len();
 344
 345                                        let range = snapshot.anchor_after(offset)
 346                                            ..snapshot.anchor_after(offset + text_len);
 347
 348                                        let crease = super::crease_for_mention(
 349                                            format!(
 350                                                "{} ({}-{})",
 351                                                file_name,
 352                                                line_range.start.row + 1,
 353                                                line_range.end.row + 1
 354                                            )
 355                                            .into(),
 356                                            IconName::Context.path().into(),
 357                                            range,
 358                                            editor.downgrade(),
 359                                        );
 360
 361                                        editor.update(cx, |editor, cx| {
 362                                            editor.insert_creases(vec![crease.clone()], cx);
 363                                            editor.fold_creases(vec![crease], false, window, cx);
 364                                        });
 365
 366                                        current_offset += text_len + 1;
 367                                    }
 368                                });
 369
 370                                false
 371                            }
 372                        });
 373
 374                        (new_text, callback)
 375                    }
 376                };
 377
 378                Some(Completion {
 379                    replace_range: source_range.clone(),
 380                    new_text,
 381                    label: CodeLabel::plain(action.label().to_string(), None),
 382                    icon_path: Some(action.icon().path().into()),
 383                    documentation: None,
 384                    source: project::CompletionSource::Custom,
 385                    insert_text_mode: None,
 386                    // This ensures that when a user accepts this completion, the
 387                    // completion menu will still be shown after "@category " is
 388                    // inserted
 389                    confirm: Some(on_action),
 390                })
 391            }
 392        }
 393    }
 394
 395    fn completion_for_thread(
 396        thread_entry: ThreadContextEntry,
 397        excerpt_id: ExcerptId,
 398        source_range: Range<Anchor>,
 399        recent: bool,
 400        editor: Entity<Editor>,
 401        context_store: Entity<ContextStore>,
 402        thread_store: Entity<ThreadStore>,
 403    ) -> Completion {
 404        let icon_for_completion = if recent {
 405            IconName::HistoryRerun
 406        } else {
 407            IconName::MessageBubbles
 408        };
 409        let new_text = MentionLink::for_thread(&thread_entry);
 410        let new_text_len = new_text.len();
 411        Completion {
 412            replace_range: source_range.clone(),
 413            new_text,
 414            label: CodeLabel::plain(thread_entry.summary.to_string(), None),
 415            documentation: None,
 416            insert_text_mode: None,
 417            source: project::CompletionSource::Custom,
 418            icon_path: Some(icon_for_completion.path().into()),
 419            confirm: Some(confirm_completion_callback(
 420                IconName::MessageBubbles.path().into(),
 421                thread_entry.summary.clone(),
 422                excerpt_id,
 423                source_range.start,
 424                new_text_len,
 425                editor.clone(),
 426                context_store.clone(),
 427                move |cx| {
 428                    let thread_id = thread_entry.id.clone();
 429                    let context_store = context_store.clone();
 430                    let thread_store = thread_store.clone();
 431                    cx.spawn::<_, Option<_>>(async move |cx| {
 432                        let thread: Entity<Thread> = thread_store
 433                            .update(cx, |thread_store, cx| {
 434                                thread_store.open_thread(&thread_id, cx)
 435                            })
 436                            .ok()?
 437                            .await
 438                            .log_err()?;
 439                        let context = context_store
 440                            .update(cx, |context_store, cx| {
 441                                context_store.add_thread(thread, false, cx)
 442                            })
 443                            .ok()??;
 444                        Some(context)
 445                    })
 446                },
 447            )),
 448        }
 449    }
 450
 451    fn completion_for_rules(
 452        rules: RulesContextEntry,
 453        excerpt_id: ExcerptId,
 454        source_range: Range<Anchor>,
 455        editor: Entity<Editor>,
 456        context_store: Entity<ContextStore>,
 457    ) -> Completion {
 458        let new_text = MentionLink::for_rule(&rules);
 459        let new_text_len = new_text.len();
 460        Completion {
 461            replace_range: source_range.clone(),
 462            new_text,
 463            label: CodeLabel::plain(rules.title.to_string(), None),
 464            documentation: None,
 465            insert_text_mode: None,
 466            source: project::CompletionSource::Custom,
 467            icon_path: Some(RULES_ICON.path().into()),
 468            confirm: Some(confirm_completion_callback(
 469                RULES_ICON.path().into(),
 470                rules.title.clone(),
 471                excerpt_id,
 472                source_range.start,
 473                new_text_len,
 474                editor.clone(),
 475                context_store.clone(),
 476                move |cx| {
 477                    let user_prompt_id = rules.prompt_id;
 478                    let context = context_store.update(cx, |context_store, cx| {
 479                        context_store.add_rules(user_prompt_id, false, cx)
 480                    });
 481                    Task::ready(context)
 482                },
 483            )),
 484        }
 485    }
 486
 487    fn completion_for_fetch(
 488        source_range: Range<Anchor>,
 489        url_to_fetch: SharedString,
 490        excerpt_id: ExcerptId,
 491        editor: Entity<Editor>,
 492        context_store: Entity<ContextStore>,
 493        http_client: Arc<HttpClientWithUrl>,
 494    ) -> Completion {
 495        let new_text = MentionLink::for_fetch(&url_to_fetch);
 496        let new_text_len = new_text.len();
 497        Completion {
 498            replace_range: source_range.clone(),
 499            new_text,
 500            label: CodeLabel::plain(url_to_fetch.to_string(), None),
 501            documentation: None,
 502            source: project::CompletionSource::Custom,
 503            icon_path: Some(IconName::Globe.path().into()),
 504            insert_text_mode: None,
 505            confirm: Some(confirm_completion_callback(
 506                IconName::Globe.path().into(),
 507                url_to_fetch.clone(),
 508                excerpt_id,
 509                source_range.start,
 510                new_text_len,
 511                editor.clone(),
 512                context_store.clone(),
 513                move |cx| {
 514                    let context_store = context_store.clone();
 515                    let http_client = http_client.clone();
 516                    let url_to_fetch = url_to_fetch.clone();
 517                    cx.spawn(async move |cx| {
 518                        if let Some(context) = context_store
 519                            .update(cx, |context_store, _| {
 520                                context_store.get_url_context(url_to_fetch.clone())
 521                            })
 522                            .ok()?
 523                        {
 524                            return Some(context);
 525                        }
 526                        let content = cx
 527                            .background_spawn(fetch_url_content(
 528                                http_client,
 529                                url_to_fetch.to_string(),
 530                            ))
 531                            .await
 532                            .log_err()?;
 533                        context_store
 534                            .update(cx, |context_store, cx| {
 535                                context_store.add_fetched_url(url_to_fetch.to_string(), content, cx)
 536                            })
 537                            .ok()
 538                    })
 539                },
 540            )),
 541        }
 542    }
 543
 544    fn completion_for_path(
 545        project_path: ProjectPath,
 546        path_prefix: &str,
 547        is_recent: bool,
 548        is_directory: bool,
 549        excerpt_id: ExcerptId,
 550        source_range: Range<Anchor>,
 551        editor: Entity<Editor>,
 552        context_store: Entity<ContextStore>,
 553        cx: &App,
 554    ) -> Completion {
 555        let (file_name, directory) = super::file_context_picker::extract_file_name_and_directory(
 556            &project_path.path,
 557            path_prefix,
 558        );
 559
 560        let label =
 561            build_code_label_for_full_path(&file_name, directory.as_ref().map(|s| s.as_ref()), cx);
 562        let full_path = if let Some(directory) = directory {
 563            format!("{}{}", directory, file_name)
 564        } else {
 565            file_name.to_string()
 566        };
 567
 568        let crease_icon_path = if is_directory {
 569            FileIcons::get_folder_icon(false, cx).unwrap_or_else(|| IconName::Folder.path().into())
 570        } else {
 571            FileIcons::get_icon(Path::new(&full_path), cx)
 572                .unwrap_or_else(|| IconName::File.path().into())
 573        };
 574        let completion_icon_path = if is_recent {
 575            IconName::HistoryRerun.path().into()
 576        } else {
 577            crease_icon_path.clone()
 578        };
 579
 580        let new_text = MentionLink::for_file(&file_name, &full_path);
 581        let new_text_len = new_text.len();
 582        Completion {
 583            replace_range: source_range.clone(),
 584            new_text,
 585            label,
 586            documentation: None,
 587            source: project::CompletionSource::Custom,
 588            icon_path: Some(completion_icon_path),
 589            insert_text_mode: None,
 590            confirm: Some(confirm_completion_callback(
 591                crease_icon_path,
 592                file_name,
 593                excerpt_id,
 594                source_range.start,
 595                new_text_len,
 596                editor,
 597                context_store.clone(),
 598                move |cx| {
 599                    if is_directory {
 600                        Task::ready(
 601                            context_store
 602                                .update(cx, |context_store, cx| {
 603                                    context_store.add_directory(&project_path, false, cx)
 604                                })
 605                                .log_err()
 606                                .flatten(),
 607                        )
 608                    } else {
 609                        let result = context_store.update(cx, |context_store, cx| {
 610                            context_store.add_file_from_path(project_path.clone(), false, cx)
 611                        });
 612                        cx.spawn(async move |_| result.await.log_err().flatten())
 613                    }
 614                },
 615            )),
 616        }
 617    }
 618
 619    fn completion_for_symbol(
 620        symbol: Symbol,
 621        excerpt_id: ExcerptId,
 622        source_range: Range<Anchor>,
 623        editor: Entity<Editor>,
 624        context_store: Entity<ContextStore>,
 625        workspace: Entity<Workspace>,
 626        cx: &mut App,
 627    ) -> Option<Completion> {
 628        let path_prefix = workspace
 629            .read(cx)
 630            .project()
 631            .read(cx)
 632            .worktree_for_id(symbol.path.worktree_id, cx)?
 633            .read(cx)
 634            .root_name();
 635
 636        let (file_name, directory) = super::file_context_picker::extract_file_name_and_directory(
 637            &symbol.path.path,
 638            path_prefix,
 639        );
 640        let full_path = if let Some(directory) = directory {
 641            format!("{}{}", directory, file_name)
 642        } else {
 643            file_name.to_string()
 644        };
 645
 646        let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
 647        let mut label = CodeLabel::plain(symbol.name.clone(), None);
 648        label.push_str(" ", None);
 649        label.push_str(&file_name, comment_id);
 650
 651        let new_text = MentionLink::for_symbol(&symbol.name, &full_path);
 652        let new_text_len = new_text.len();
 653        Some(Completion {
 654            replace_range: source_range.clone(),
 655            new_text,
 656            label,
 657            documentation: None,
 658            source: project::CompletionSource::Custom,
 659            icon_path: Some(IconName::Code.path().into()),
 660            insert_text_mode: None,
 661            confirm: Some(confirm_completion_callback(
 662                IconName::Code.path().into(),
 663                symbol.name.clone().into(),
 664                excerpt_id,
 665                source_range.start,
 666                new_text_len,
 667                editor.clone(),
 668                context_store.clone(),
 669                move |cx| {
 670                    let symbol = symbol.clone();
 671                    let context_store = context_store.clone();
 672                    let workspace = workspace.clone();
 673                    let result = super::symbol_context_picker::add_symbol(
 674                        symbol.clone(),
 675                        false,
 676                        workspace.clone(),
 677                        context_store.downgrade(),
 678                        cx,
 679                    );
 680                    cx.spawn(async move |_| result.await.log_err()?.0)
 681                },
 682            )),
 683        })
 684    }
 685}
 686
 687fn build_code_label_for_full_path(file_name: &str, directory: Option<&str>, cx: &App) -> CodeLabel {
 688    let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
 689    let mut label = CodeLabel::default();
 690
 691    label.push_str(&file_name, None);
 692    label.push_str(" ", None);
 693
 694    if let Some(directory) = directory {
 695        label.push_str(&directory, comment_id);
 696    }
 697
 698    label.filter_range = 0..label.text().len();
 699
 700    label
 701}
 702
 703impl CompletionProvider for ContextPickerCompletionProvider {
 704    fn completions(
 705        &self,
 706        excerpt_id: ExcerptId,
 707        buffer: &Entity<Buffer>,
 708        buffer_position: Anchor,
 709        _trigger: CompletionContext,
 710        _window: &mut Window,
 711        cx: &mut Context<Editor>,
 712    ) -> Task<Result<Option<Vec<Completion>>>> {
 713        let state = buffer.update(cx, |buffer, _cx| {
 714            let position = buffer_position.to_point(buffer);
 715            let line_start = Point::new(position.row, 0);
 716            let offset_to_line = buffer.point_to_offset(line_start);
 717            let mut lines = buffer.text_for_range(line_start..position).lines();
 718            let line = lines.next()?;
 719            MentionCompletion::try_parse(line, offset_to_line)
 720        });
 721        let Some(state) = state else {
 722            return Task::ready(Ok(None));
 723        };
 724
 725        let Some((workspace, context_store)) =
 726            self.workspace.upgrade().zip(self.context_store.upgrade())
 727        else {
 728            return Task::ready(Ok(None));
 729        };
 730
 731        let snapshot = buffer.read(cx).snapshot();
 732        let source_range = snapshot.anchor_before(state.source_range.start)
 733            ..snapshot.anchor_before(state.source_range.end);
 734
 735        let thread_store = self.thread_store.clone();
 736        let editor = self.editor.clone();
 737        let http_client = workspace.read(cx).client().http_client();
 738
 739        let MentionCompletion { mode, argument, .. } = state;
 740        let query = argument.unwrap_or_else(|| "".to_string());
 741
 742        let excluded_path = self
 743            .excluded_buffer
 744            .as_ref()
 745            .and_then(WeakEntity::upgrade)
 746            .and_then(|b| b.read(cx).file())
 747            .map(|file| ProjectPath::from_file(file.as_ref(), cx));
 748
 749        let recent_entries = recent_context_picker_entries(
 750            context_store.clone(),
 751            thread_store.clone(),
 752            workspace.clone(),
 753            excluded_path.clone(),
 754            cx,
 755        );
 756
 757        let prompt_store = thread_store.as_ref().and_then(|thread_store| {
 758            thread_store
 759                .read_with(cx, |thread_store, _cx| thread_store.prompt_store().clone())
 760                .ok()
 761                .flatten()
 762        });
 763
 764        let search_task = search(
 765            mode,
 766            query,
 767            Arc::<AtomicBool>::default(),
 768            recent_entries,
 769            prompt_store,
 770            thread_store.clone(),
 771            workspace.clone(),
 772            cx,
 773        );
 774
 775        cx.spawn(async move |_, cx| {
 776            let matches = search_task.await;
 777            let Some(editor) = editor.upgrade() else {
 778                return Ok(None);
 779            };
 780
 781            Ok(Some(cx.update(|cx| {
 782                matches
 783                    .into_iter()
 784                    .filter_map(|mat| match mat {
 785                        Match::File(FileMatch { mat, is_recent }) => {
 786                            let project_path = ProjectPath {
 787                                worktree_id: WorktreeId::from_usize(mat.worktree_id),
 788                                path: mat.path.clone(),
 789                            };
 790
 791                            if excluded_path.as_ref() == Some(&project_path) {
 792                                return None;
 793                            }
 794
 795                            Some(Self::completion_for_path(
 796                                project_path,
 797                                &mat.path_prefix,
 798                                is_recent,
 799                                mat.is_dir,
 800                                excerpt_id,
 801                                source_range.clone(),
 802                                editor.clone(),
 803                                context_store.clone(),
 804                                cx,
 805                            ))
 806                        }
 807
 808                        Match::Symbol(SymbolMatch { symbol, .. }) => Self::completion_for_symbol(
 809                            symbol,
 810                            excerpt_id,
 811                            source_range.clone(),
 812                            editor.clone(),
 813                            context_store.clone(),
 814                            workspace.clone(),
 815                            cx,
 816                        ),
 817
 818                        Match::Thread(ThreadMatch {
 819                            thread, is_recent, ..
 820                        }) => {
 821                            let thread_store = thread_store.as_ref().and_then(|t| t.upgrade())?;
 822                            Some(Self::completion_for_thread(
 823                                thread,
 824                                excerpt_id,
 825                                source_range.clone(),
 826                                is_recent,
 827                                editor.clone(),
 828                                context_store.clone(),
 829                                thread_store,
 830                            ))
 831                        }
 832
 833                        Match::Rules(user_rules) => Some(Self::completion_for_rules(
 834                            user_rules,
 835                            excerpt_id,
 836                            source_range.clone(),
 837                            editor.clone(),
 838                            context_store.clone(),
 839                        )),
 840
 841                        Match::Fetch(url) => Some(Self::completion_for_fetch(
 842                            source_range.clone(),
 843                            url,
 844                            excerpt_id,
 845                            editor.clone(),
 846                            context_store.clone(),
 847                            http_client.clone(),
 848                        )),
 849
 850                        Match::Entry(EntryMatch { entry, .. }) => Self::completion_for_entry(
 851                            entry,
 852                            excerpt_id,
 853                            source_range.clone(),
 854                            editor.clone(),
 855                            context_store.clone(),
 856                            &workspace,
 857                            cx,
 858                        ),
 859                    })
 860                    .collect()
 861            })?))
 862        })
 863    }
 864
 865    fn resolve_completions(
 866        &self,
 867        _buffer: Entity<Buffer>,
 868        _completion_indices: Vec<usize>,
 869        _completions: Rc<RefCell<Box<[Completion]>>>,
 870        _cx: &mut Context<Editor>,
 871    ) -> Task<Result<bool>> {
 872        Task::ready(Ok(true))
 873    }
 874
 875    fn is_completion_trigger(
 876        &self,
 877        buffer: &Entity<language::Buffer>,
 878        position: language::Anchor,
 879        _: &str,
 880        _: bool,
 881        cx: &mut Context<Editor>,
 882    ) -> bool {
 883        let buffer = buffer.read(cx);
 884        let position = position.to_point(buffer);
 885        let line_start = Point::new(position.row, 0);
 886        let offset_to_line = buffer.point_to_offset(line_start);
 887        let mut lines = buffer.text_for_range(line_start..position).lines();
 888        if let Some(line) = lines.next() {
 889            MentionCompletion::try_parse(line, offset_to_line)
 890                .map(|completion| {
 891                    completion.source_range.start <= offset_to_line + position.column as usize
 892                        && completion.source_range.end >= offset_to_line + position.column as usize
 893                })
 894                .unwrap_or(false)
 895        } else {
 896            false
 897        }
 898    }
 899
 900    fn sort_completions(&self) -> bool {
 901        false
 902    }
 903
 904    fn filter_completions(&self) -> bool {
 905        false
 906    }
 907}
 908
 909fn confirm_completion_callback(
 910    crease_icon_path: SharedString,
 911    crease_text: SharedString,
 912    excerpt_id: ExcerptId,
 913    start: Anchor,
 914    content_len: usize,
 915    editor: Entity<Editor>,
 916    context_store: Entity<ContextStore>,
 917    add_context_fn: impl Fn(&mut App) -> Task<Option<AgentContextHandle>> + Send + Sync + 'static,
 918) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
 919    Arc::new(move |_, window, cx| {
 920        let context = add_context_fn(cx);
 921
 922        let crease_text = crease_text.clone();
 923        let crease_icon_path = crease_icon_path.clone();
 924        let editor = editor.clone();
 925        let context_store = context_store.clone();
 926        window.defer(cx, move |window, cx| {
 927            let crease_id = crate::context_picker::insert_crease_for_mention(
 928                excerpt_id,
 929                start,
 930                content_len,
 931                crease_text.clone(),
 932                crease_icon_path,
 933                editor.clone(),
 934                window,
 935                cx,
 936            );
 937            cx.spawn(async move |cx| {
 938                let crease_id = crease_id?;
 939                let context = context.await?;
 940                editor
 941                    .update(cx, |editor, cx| {
 942                        if let Some(addon) = editor.addon_mut::<ContextCreasesAddon>() {
 943                            addon.add_creases(
 944                                &context_store,
 945                                AgentContextKey(context),
 946                                [(crease_id, crease_text)],
 947                                cx,
 948                            );
 949                        }
 950                    })
 951                    .ok()
 952            })
 953            .detach();
 954        });
 955        false
 956    })
 957}
 958
 959#[derive(Debug, Default, PartialEq)]
 960struct MentionCompletion {
 961    source_range: Range<usize>,
 962    mode: Option<ContextPickerMode>,
 963    argument: Option<String>,
 964}
 965
 966impl MentionCompletion {
 967    fn try_parse(line: &str, offset_to_line: usize) -> Option<Self> {
 968        let last_mention_start = line.rfind('@')?;
 969        if last_mention_start >= line.len() {
 970            return Some(Self::default());
 971        }
 972        if last_mention_start > 0
 973            && line
 974                .chars()
 975                .nth(last_mention_start - 1)
 976                .map_or(false, |c| !c.is_whitespace())
 977        {
 978            return None;
 979        }
 980
 981        let rest_of_line = &line[last_mention_start + 1..];
 982
 983        let mut mode = None;
 984        let mut argument = None;
 985
 986        let mut parts = rest_of_line.split_whitespace();
 987        let mut end = last_mention_start + 1;
 988        if let Some(mode_text) = parts.next() {
 989            end += mode_text.len();
 990
 991            if let Some(parsed_mode) = ContextPickerMode::try_from(mode_text).ok() {
 992                mode = Some(parsed_mode);
 993            } else {
 994                argument = Some(mode_text.to_string());
 995            }
 996            match rest_of_line[mode_text.len()..].find(|c: char| !c.is_whitespace()) {
 997                Some(whitespace_count) => {
 998                    if let Some(argument_text) = parts.next() {
 999                        argument = Some(argument_text.to_string());
1000                        end += whitespace_count + argument_text.len();
1001                    }
1002                }
1003                None => {
1004                    // Rest of line is entirely whitespace
1005                    end += rest_of_line.len() - mode_text.len();
1006                }
1007            }
1008        }
1009
1010        Some(Self {
1011            source_range: last_mention_start + offset_to_line..end + offset_to_line,
1012            mode,
1013            argument,
1014        })
1015    }
1016}
1017
1018#[cfg(test)]
1019mod tests {
1020    use super::*;
1021    use editor::AnchorRangeExt;
1022    use gpui::{EventEmitter, FocusHandle, Focusable, TestAppContext, VisualTestContext};
1023    use project::{Project, ProjectPath};
1024    use serde_json::json;
1025    use settings::SettingsStore;
1026    use std::ops::Deref;
1027    use util::{path, separator};
1028    use workspace::{AppState, Item};
1029
1030    #[test]
1031    fn test_mention_completion_parse() {
1032        assert_eq!(MentionCompletion::try_parse("Lorem Ipsum", 0), None);
1033
1034        assert_eq!(
1035            MentionCompletion::try_parse("Lorem @", 0),
1036            Some(MentionCompletion {
1037                source_range: 6..7,
1038                mode: None,
1039                argument: None,
1040            })
1041        );
1042
1043        assert_eq!(
1044            MentionCompletion::try_parse("Lorem @file", 0),
1045            Some(MentionCompletion {
1046                source_range: 6..11,
1047                mode: Some(ContextPickerMode::File),
1048                argument: None,
1049            })
1050        );
1051
1052        assert_eq!(
1053            MentionCompletion::try_parse("Lorem @file ", 0),
1054            Some(MentionCompletion {
1055                source_range: 6..12,
1056                mode: Some(ContextPickerMode::File),
1057                argument: None,
1058            })
1059        );
1060
1061        assert_eq!(
1062            MentionCompletion::try_parse("Lorem @file main.rs", 0),
1063            Some(MentionCompletion {
1064                source_range: 6..19,
1065                mode: Some(ContextPickerMode::File),
1066                argument: Some("main.rs".to_string()),
1067            })
1068        );
1069
1070        assert_eq!(
1071            MentionCompletion::try_parse("Lorem @file main.rs ", 0),
1072            Some(MentionCompletion {
1073                source_range: 6..19,
1074                mode: Some(ContextPickerMode::File),
1075                argument: Some("main.rs".to_string()),
1076            })
1077        );
1078
1079        assert_eq!(
1080            MentionCompletion::try_parse("Lorem @file main.rs Ipsum", 0),
1081            Some(MentionCompletion {
1082                source_range: 6..19,
1083                mode: Some(ContextPickerMode::File),
1084                argument: Some("main.rs".to_string()),
1085            })
1086        );
1087
1088        assert_eq!(
1089            MentionCompletion::try_parse("Lorem @main", 0),
1090            Some(MentionCompletion {
1091                source_range: 6..11,
1092                mode: None,
1093                argument: Some("main".to_string()),
1094            })
1095        );
1096
1097        assert_eq!(MentionCompletion::try_parse("test@", 0), None);
1098    }
1099
1100    struct AtMentionEditor(Entity<Editor>);
1101
1102    impl Item for AtMentionEditor {
1103        type Event = ();
1104
1105        fn include_in_nav_history() -> bool {
1106            false
1107        }
1108
1109        fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
1110            "Test".into()
1111        }
1112    }
1113
1114    impl EventEmitter<()> for AtMentionEditor {}
1115
1116    impl Focusable for AtMentionEditor {
1117        fn focus_handle(&self, cx: &App) -> FocusHandle {
1118            self.0.read(cx).focus_handle(cx).clone()
1119        }
1120    }
1121
1122    impl Render for AtMentionEditor {
1123        fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
1124            self.0.clone().into_any_element()
1125        }
1126    }
1127
1128    #[gpui::test]
1129    async fn test_context_completion_provider(cx: &mut TestAppContext) {
1130        init_test(cx);
1131
1132        let app_state = cx.update(AppState::test);
1133
1134        cx.update(|cx| {
1135            language::init(cx);
1136            editor::init(cx);
1137            workspace::init(app_state.clone(), cx);
1138            Project::init_settings(cx);
1139        });
1140
1141        app_state
1142            .fs
1143            .as_fake()
1144            .insert_tree(
1145                path!("/dir"),
1146                json!({
1147                    "editor": "",
1148                    "a": {
1149                        "one.txt": "",
1150                        "two.txt": "",
1151                        "three.txt": "",
1152                        "four.txt": ""
1153                    },
1154                    "b": {
1155                        "five.txt": "",
1156                        "six.txt": "",
1157                        "seven.txt": "",
1158                        "eight.txt": "",
1159                    }
1160                }),
1161            )
1162            .await;
1163
1164        let project = Project::test(app_state.fs.clone(), [path!("/dir").as_ref()], cx).await;
1165        let window = cx.add_window(|window, cx| Workspace::test_new(project.clone(), window, cx));
1166        let workspace = window.root(cx).unwrap();
1167
1168        let worktree = project.update(cx, |project, cx| {
1169            let mut worktrees = project.worktrees(cx).collect::<Vec<_>>();
1170            assert_eq!(worktrees.len(), 1);
1171            worktrees.pop().unwrap()
1172        });
1173        let worktree_id = worktree.update(cx, |worktree, _| worktree.id());
1174
1175        let mut cx = VisualTestContext::from_window(*window.deref(), cx);
1176
1177        let paths = vec![
1178            separator!("a/one.txt"),
1179            separator!("a/two.txt"),
1180            separator!("a/three.txt"),
1181            separator!("a/four.txt"),
1182            separator!("b/five.txt"),
1183            separator!("b/six.txt"),
1184            separator!("b/seven.txt"),
1185            separator!("b/eight.txt"),
1186        ];
1187
1188        let mut opened_editors = Vec::new();
1189        for path in paths {
1190            let buffer = workspace
1191                .update_in(&mut cx, |workspace, window, cx| {
1192                    workspace.open_path(
1193                        ProjectPath {
1194                            worktree_id,
1195                            path: Path::new(path).into(),
1196                        },
1197                        None,
1198                        false,
1199                        window,
1200                        cx,
1201                    )
1202                })
1203                .await
1204                .unwrap();
1205            opened_editors.push(buffer);
1206        }
1207
1208        let editor = workspace.update_in(&mut cx, |workspace, window, cx| {
1209            let editor = cx.new(|cx| {
1210                Editor::new(
1211                    editor::EditorMode::full(),
1212                    multi_buffer::MultiBuffer::build_simple("", cx),
1213                    None,
1214                    window,
1215                    cx,
1216                )
1217            });
1218            workspace.active_pane().update(cx, |pane, cx| {
1219                pane.add_item(
1220                    Box::new(cx.new(|_| AtMentionEditor(editor.clone()))),
1221                    true,
1222                    true,
1223                    None,
1224                    window,
1225                    cx,
1226                );
1227            });
1228            editor
1229        });
1230
1231        let context_store = cx.new(|_| ContextStore::new(project.downgrade(), None));
1232
1233        let editor_entity = editor.downgrade();
1234        editor.update_in(&mut cx, |editor, window, cx| {
1235            let last_opened_buffer = opened_editors.last().and_then(|editor| {
1236                editor
1237                    .downcast::<Editor>()?
1238                    .read(cx)
1239                    .buffer()
1240                    .read(cx)
1241                    .as_singleton()
1242                    .as_ref()
1243                    .map(Entity::downgrade)
1244            });
1245            window.focus(&editor.focus_handle(cx));
1246            editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
1247                workspace.downgrade(),
1248                context_store.downgrade(),
1249                None,
1250                editor_entity,
1251                last_opened_buffer,
1252            ))));
1253        });
1254
1255        cx.simulate_input("Lorem ");
1256
1257        editor.update(&mut cx, |editor, cx| {
1258            assert_eq!(editor.text(cx), "Lorem ");
1259            assert!(!editor.has_visible_completions_menu());
1260        });
1261
1262        cx.simulate_input("@");
1263
1264        editor.update(&mut cx, |editor, cx| {
1265            assert_eq!(editor.text(cx), "Lorem @");
1266            assert!(editor.has_visible_completions_menu());
1267            assert_eq!(
1268                current_completion_labels(editor),
1269                &[
1270                    "seven.txt dir/b/",
1271                    "six.txt dir/b/",
1272                    "five.txt dir/b/",
1273                    "four.txt dir/a/",
1274                    "Files & Directories",
1275                    "Symbols",
1276                    "Fetch"
1277                ]
1278            );
1279        });
1280
1281        // Select and confirm "File"
1282        editor.update_in(&mut cx, |editor, window, cx| {
1283            assert!(editor.has_visible_completions_menu());
1284            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
1285            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
1286            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
1287            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
1288            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
1289        });
1290
1291        cx.run_until_parked();
1292
1293        editor.update(&mut cx, |editor, cx| {
1294            assert_eq!(editor.text(cx), "Lorem @file ");
1295            assert!(editor.has_visible_completions_menu());
1296        });
1297
1298        cx.simulate_input("one");
1299
1300        editor.update(&mut cx, |editor, cx| {
1301            assert_eq!(editor.text(cx), "Lorem @file one");
1302            assert!(editor.has_visible_completions_menu());
1303            assert_eq!(current_completion_labels(editor), vec!["one.txt dir/a/"]);
1304        });
1305
1306        editor.update_in(&mut cx, |editor, window, cx| {
1307            assert!(editor.has_visible_completions_menu());
1308            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
1309        });
1310
1311        editor.update(&mut cx, |editor, cx| {
1312            assert_eq!(editor.text(cx), "Lorem [@one.txt](@file:dir/a/one.txt)",);
1313            assert!(!editor.has_visible_completions_menu());
1314            assert_eq!(
1315                fold_ranges(editor, cx),
1316                vec![Point::new(0, 6)..Point::new(0, 37)]
1317            );
1318        });
1319
1320        cx.simulate_input(" ");
1321
1322        editor.update(&mut cx, |editor, cx| {
1323            assert_eq!(editor.text(cx), "Lorem [@one.txt](@file:dir/a/one.txt) ",);
1324            assert!(!editor.has_visible_completions_menu());
1325            assert_eq!(
1326                fold_ranges(editor, cx),
1327                vec![Point::new(0, 6)..Point::new(0, 37)]
1328            );
1329        });
1330
1331        cx.simulate_input("Ipsum ");
1332
1333        editor.update(&mut cx, |editor, cx| {
1334            assert_eq!(
1335                editor.text(cx),
1336                "Lorem [@one.txt](@file:dir/a/one.txt) Ipsum ",
1337            );
1338            assert!(!editor.has_visible_completions_menu());
1339            assert_eq!(
1340                fold_ranges(editor, cx),
1341                vec![Point::new(0, 6)..Point::new(0, 37)]
1342            );
1343        });
1344
1345        cx.simulate_input("@file ");
1346
1347        editor.update(&mut cx, |editor, cx| {
1348            assert_eq!(
1349                editor.text(cx),
1350                "Lorem [@one.txt](@file:dir/a/one.txt) Ipsum @file ",
1351            );
1352            assert!(editor.has_visible_completions_menu());
1353            assert_eq!(
1354                fold_ranges(editor, cx),
1355                vec![Point::new(0, 6)..Point::new(0, 37)]
1356            );
1357        });
1358
1359        editor.update_in(&mut cx, |editor, window, cx| {
1360            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
1361        });
1362
1363        cx.run_until_parked();
1364
1365        editor.update(&mut cx, |editor, cx| {
1366            assert_eq!(
1367                editor.text(cx),
1368                "Lorem [@one.txt](@file:dir/a/one.txt) Ipsum [@seven.txt](@file:dir/b/seven.txt)"
1369            );
1370            assert!(!editor.has_visible_completions_menu());
1371            assert_eq!(
1372                fold_ranges(editor, cx),
1373                vec![
1374                    Point::new(0, 6)..Point::new(0, 37),
1375                    Point::new(0, 44)..Point::new(0, 79)
1376                ]
1377            );
1378        });
1379
1380        cx.simulate_input("\n@");
1381
1382        editor.update(&mut cx, |editor, cx| {
1383            assert_eq!(
1384                editor.text(cx),
1385                "Lorem [@one.txt](@file:dir/a/one.txt) Ipsum [@seven.txt](@file:dir/b/seven.txt)\n@"
1386            );
1387            assert!(editor.has_visible_completions_menu());
1388            assert_eq!(
1389                fold_ranges(editor, cx),
1390                vec![
1391                    Point::new(0, 6)..Point::new(0, 37),
1392                    Point::new(0, 44)..Point::new(0, 79)
1393                ]
1394            );
1395        });
1396
1397        editor.update_in(&mut cx, |editor, window, cx| {
1398            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
1399        });
1400
1401        cx.run_until_parked();
1402
1403        editor.update(&mut cx, |editor, cx| {
1404            assert_eq!(
1405                editor.text(cx),
1406                "Lorem [@one.txt](@file:dir/a/one.txt) Ipsum [@seven.txt](@file:dir/b/seven.txt)\n[@six.txt](@file:dir/b/six.txt)"
1407            );
1408            assert!(!editor.has_visible_completions_menu());
1409            assert_eq!(
1410                fold_ranges(editor, cx),
1411                vec![
1412                    Point::new(0, 6)..Point::new(0, 37),
1413                    Point::new(0, 44)..Point::new(0, 79),
1414                    Point::new(1, 0)..Point::new(1, 31)
1415                ]
1416            );
1417        });
1418    }
1419
1420    fn fold_ranges(editor: &Editor, cx: &mut App) -> Vec<Range<Point>> {
1421        let snapshot = editor.buffer().read(cx).snapshot(cx);
1422        editor.display_map.update(cx, |display_map, cx| {
1423            display_map
1424                .snapshot(cx)
1425                .folds_in_range(0..snapshot.len())
1426                .map(|fold| fold.range.to_point(&snapshot))
1427                .collect()
1428        })
1429    }
1430
1431    fn current_completion_labels(editor: &Editor) -> Vec<String> {
1432        let completions = editor.current_completions().expect("Missing completions");
1433        completions
1434            .into_iter()
1435            .map(|completion| completion.label.text.to_string())
1436            .collect::<Vec<_>>()
1437    }
1438
1439    pub(crate) fn init_test(cx: &mut TestAppContext) {
1440        cx.update(|cx| {
1441            let store = SettingsStore::test(cx);
1442            cx.set_global(store);
1443            theme::init(theme::LoadThemes::JustBase, cx);
1444            client::init_settings(cx);
1445            language::init(cx);
1446            Project::init_settings(cx);
1447            workspace::init_settings(cx);
1448            editor::init_settings(cx);
1449        });
1450    }
1451}