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