1506 lines · 54.3 KB · Rust
RawDownload
use std::cell::RefCell;use std::ops::Range;use std::path::{Path, PathBuf};use std::rc::Rc;use std::sync::Arc;use std::sync::atomic::AtomicBool; use anyhow::Result;use editor::{CompletionProvider, Editor, ExcerptId, ToOffset as _};use file_icons::FileIcons;use fuzzy::{StringMatch, StringMatchCandidate};use gpui::{App, Entity, Task, WeakEntity};use http_client::HttpClientWithUrl;use itertools::Itertools;use language::{Buffer, CodeLabel, HighlightId};use lsp::CompletionContext;use project::{Completion, CompletionIntent, CompletionResponse, ProjectPath, Symbol, WorktreeId};use prompt_store::PromptStore;use rope::Point;use text::{Anchor, OffsetRangeExt, ToPoint};use ui::prelude::*;use util::ResultExt as _;use workspace::Workspace; use crate::Thread;use crate::context::{AgentContextHandle, AgentContextKey, ContextCreasesAddon, RULES_ICON};use crate::context_store::ContextStore;use crate::thread_store::{TextThreadStore, ThreadStore}; use super::fetch_context_picker::fetch_url_content;use super::file_context_picker::{FileMatch, search_files};use super::rules_context_picker::{RulesContextEntry, search_rules};use super::symbol_context_picker::SymbolMatch;use super::symbol_context_picker::search_symbols;use super::thread_context_picker::{ThreadContextEntry, ThreadMatch, search_threads};use super::{    ContextPickerAction, ContextPickerEntry, ContextPickerMode, MentionLink, RecentEntry,    available_context_picker_entries, recent_context_picker_entries, selection_ranges,}; pub(crate) enum Match {    File(FileMatch),    Symbol(SymbolMatch),    Thread(ThreadMatch),    Fetch(SharedString),    Rules(RulesContextEntry),    Entry(EntryMatch),} pub struct EntryMatch {    mat: Option<StringMatch>,    entry: ContextPickerEntry,} impl Match {    pub fn score(&self) -> f64 {        match self {            Match::File(file) => file.mat.score,            Match::Entry(mode) => mode.mat.as_ref().map(|mat| mat.score).unwrap_or(1.),            Match::Thread(_) => 1.,            Match::Symbol(_) => 1.,            Match::Fetch(_) => 1.,            Match::Rules(_) => 1.,        }    }} fn search(    mode: Option<ContextPickerMode>,    query: String,    cancellation_flag: Arc<AtomicBool>,    recent_entries: Vec<RecentEntry>,    prompt_store: Option<Entity<PromptStore>>,    thread_store: Option<WeakEntity<ThreadStore>>,    text_thread_context_store: Option<WeakEntity<assistant_context_editor::ContextStore>>,    workspace: Entity<Workspace>,    cx: &mut App,) -> Task<Vec<Match>> {    match mode {        Some(ContextPickerMode::File) => {            let search_files_task =                search_files(query.clone(), cancellation_flag.clone(), &workspace, cx);            cx.background_spawn(async move {                search_files_task                    .await                    .into_iter()                    .map(Match::File)                    .collect()            })        }         Some(ContextPickerMode::Symbol) => {            let search_symbols_task =                search_symbols(query.clone(), cancellation_flag.clone(), &workspace, cx);            cx.background_spawn(async move {                search_symbols_task                    .await                    .into_iter()                    .map(Match::Symbol)                    .collect()            })        }         Some(ContextPickerMode::Thread) => {            if let Some((thread_store, context_store)) = thread_store                .as_ref()                .and_then(|t| t.upgrade())                .zip(text_thread_context_store.as_ref().and_then(|t| t.upgrade()))            {                let search_threads_task = search_threads(                    query.clone(),                    cancellation_flag.clone(),                    thread_store,                    context_store,                    cx,                );                cx.background_spawn(async move {                    search_threads_task                        .await                        .into_iter()                        .map(Match::Thread)                        .collect()                })            } else {                Task::ready(Vec::new())            }        }         Some(ContextPickerMode::Fetch) => {            if !query.is_empty() {                Task::ready(vec![Match::Fetch(query.into())])            } else {                Task::ready(Vec::new())            }        }         Some(ContextPickerMode::Rules) => {            if let Some(prompt_store) = prompt_store.as_ref() {                let search_rules_task =                    search_rules(query.clone(), cancellation_flag.clone(), prompt_store, cx);                cx.background_spawn(async move {                    search_rules_task                        .await                        .into_iter()                        .map(Match::Rules)                        .collect::<Vec<_>>()                })            } else {                Task::ready(Vec::new())            }        }         None => {            if query.is_empty() {                let mut matches = recent_entries                    .into_iter()                    .map(|entry| match entry {                        super::RecentEntry::File {                            project_path,                            path_prefix,                        } => Match::File(FileMatch {                            mat: fuzzy::PathMatch {                                score: 1.,                                positions: Vec::new(),                                worktree_id: project_path.worktree_id.to_usize(),                                path: project_path.path,                                path_prefix,                                is_dir: false,                                distance_to_relative_ancestor: 0,                            },                            is_recent: true,                        }),                        super::RecentEntry::Thread(thread_context_entry) => {                            Match::Thread(ThreadMatch {                                thread: thread_context_entry,                                is_recent: true,                            })                        }                    })                    .collect::<Vec<_>>();                 matches.extend(                    available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx)                        .into_iter()                        .map(|mode| {                            Match::Entry(EntryMatch {                                entry: mode,                                mat: None,                            })                        }),                );                 Task::ready(matches)            } else {                let executor = cx.background_executor().clone();                 let search_files_task =                    search_files(query.clone(), cancellation_flag.clone(), &workspace, cx);                 let entries =                    available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx);                let entry_candidates = entries                    .iter()                    .enumerate()                    .map(|(ix, entry)| StringMatchCandidate::new(ix, entry.keyword()))                    .collect::<Vec<_>>();                 cx.background_spawn(async move {                    let mut matches = search_files_task                        .await                        .into_iter()                        .map(Match::File)                        .collect::<Vec<_>>();                     let entry_matches = fuzzy::match_strings(                        &entry_candidates,                        &query,                        false,                        100,                        &Arc::new(AtomicBool::default()),                        executor,                    )                    .await;                     matches.extend(entry_matches.into_iter().map(|mat| {                        Match::Entry(EntryMatch {                            entry: entries[mat.candidate_id],                            mat: Some(mat),                        })                    }));                     matches.sort_by(|a, b| {                        b.score()                            .partial_cmp(&a.score())                            .unwrap_or(std::cmp::Ordering::Equal)                    });                     matches                })            }        }    }} pub struct ContextPickerCompletionProvider {    workspace: WeakEntity<Workspace>,    context_store: WeakEntity<ContextStore>,    thread_store: Option<WeakEntity<ThreadStore>>,    text_thread_store: Option<WeakEntity<TextThreadStore>>,    editor: WeakEntity<Editor>,    excluded_buffer: Option<WeakEntity<Buffer>>,} impl ContextPickerCompletionProvider {    pub fn new(        workspace: WeakEntity<Workspace>,        context_store: WeakEntity<ContextStore>,        thread_store: Option<WeakEntity<ThreadStore>>,        text_thread_store: Option<WeakEntity<TextThreadStore>>,        editor: WeakEntity<Editor>,        exclude_buffer: Option<WeakEntity<Buffer>>,    ) -> Self {        Self {            workspace,            context_store,            thread_store,            text_thread_store,            editor,            excluded_buffer: exclude_buffer,        }    }     fn completion_for_entry(        entry: ContextPickerEntry,        excerpt_id: ExcerptId,        source_range: Range<Anchor>,        editor: Entity<Editor>,        context_store: Entity<ContextStore>,        workspace: &Entity<Workspace>,        cx: &mut App,    ) -> Option<Completion> {        match entry {            ContextPickerEntry::Mode(mode) => Some(Completion {                replace_range: source_range.clone(),                new_text: format!("@{} ", mode.keyword()),                label: CodeLabel::plain(mode.label().to_string(), None),                icon_path: Some(mode.icon().path().into()),                documentation: None,                source: project::CompletionSource::Custom,                insert_text_mode: None,                // This ensures that when a user accepts this completion, the                // completion menu will still be shown after "@category " is                // inserted                confirm: Some(Arc::new(|_, _, _| true)),            }),            ContextPickerEntry::Action(action) => {                let (new_text, on_action) = match action {                    ContextPickerAction::AddSelections => {                        let selections = selection_ranges(workspace, cx);                         let selection_infos = selections                            .iter()                            .map(|(buffer, range)| {                                let full_path = buffer                                    .read(cx)                                    .file()                                    .map(|file| file.full_path(cx))                                    .unwrap_or_else(|| PathBuf::from("untitled"));                                let file_name = full_path                                    .file_name()                                    .unwrap_or_default()                                    .to_string_lossy()                                    .to_string();                                let line_range = range.to_point(&buffer.read(cx).snapshot());                                 let link = MentionLink::for_selection(                                    &file_name,                                    &full_path.to_string_lossy(),                                    line_range.start.row as usize..line_range.end.row as usize,                                );                                (file_name, link, line_range)                            })                            .collect::<Vec<_>>();                         let new_text = format!(                            "{} ",                            selection_infos.iter().map(|(_, link, _)| link).join(" ")                        );                         let callback = Arc::new({                            let context_store = context_store.clone();                            let selections = selections.clone();                            let selection_infos = selection_infos.clone();                            move |_, window: &mut Window, cx: &mut App| {                                context_store.update(cx, |context_store, cx| {                                    for (buffer, range) in &selections {                                        context_store.add_selection(                                            buffer.clone(),                                            range.clone(),                                            cx,                                        );                                    }                                });                                 let editor = editor.clone();                                let selection_infos = selection_infos.clone();                                window.defer(cx, move |window, cx| {                                    let mut current_offset = 0;                                    for (file_name, link, line_range) in selection_infos.iter() {                                        let snapshot =                                            editor.read(cx).buffer().read(cx).snapshot(cx);                                        let Some(start) = snapshot                                            .anchor_in_excerpt(excerpt_id, source_range.start)                                        else {                                            return;                                        };                                         let offset = start.to_offset(&snapshot) + current_offset;                                        let text_len = link.len();                                         let range = snapshot.anchor_after(offset)                                            ..snapshot.anchor_after(offset + text_len);                                         let crease = super::crease_for_mention(                                            format!(                                                "{} ({}-{})",                                                file_name,                                                line_range.start.row + 1,                                                line_range.end.row + 1                                            )                                            .into(),                                            IconName::Context.path().into(),                                            range,                                            editor.downgrade(),                                        );                                         editor.update(cx, |editor, cx| {                                            editor.insert_creases(vec![crease.clone()], cx);                                            editor.fold_creases(vec![crease], false, window, cx);                                        });                                         current_offset += text_len + 1;                                    }                                });                                 false                            }                        });                         (new_text, callback)                    }                };                 Some(Completion {                    replace_range: source_range.clone(),                    new_text,                    label: CodeLabel::plain(action.label().to_string(), None),                    icon_path: Some(action.icon().path().into()),                    documentation: None,                    source: project::CompletionSource::Custom,                    insert_text_mode: None,                    // This ensures that when a user accepts this completion, the                    // completion menu will still be shown after "@category " is                    // inserted                    confirm: Some(on_action),                })            }        }    }     fn completion_for_thread(        thread_entry: ThreadContextEntry,        excerpt_id: ExcerptId,        source_range: Range<Anchor>,        recent: bool,        editor: Entity<Editor>,        context_store: Entity<ContextStore>,        thread_store: Entity<ThreadStore>,        text_thread_store: Entity<TextThreadStore>,    ) -> Completion {        let icon_for_completion = if recent {            IconName::HistoryRerun        } else {            IconName::MessageBubbles        };        let new_text = format!("{} ", MentionLink::for_thread(&thread_entry));        let new_text_len = new_text.len();        Completion {            replace_range: source_range.clone(),            new_text,            label: CodeLabel::plain(thread_entry.title().to_string(), None),            documentation: None,            insert_text_mode: None,            source: project::CompletionSource::Custom,            icon_path: Some(icon_for_completion.path().into()),            confirm: Some(confirm_completion_callback(                IconName::MessageBubbles.path().into(),                thread_entry.title().clone(),                excerpt_id,                source_range.start,                new_text_len - 1,                editor.clone(),                context_store.clone(),                move |window, cx| match &thread_entry {                    ThreadContextEntry::Thread { id, .. } => {                        let thread_id = id.clone();                        let context_store = context_store.clone();                        let thread_store = thread_store.clone();                        window.spawn::<_, Option<_>>(cx, async move |cx| {                            let thread: Entity<Thread> = thread_store                                .update_in(cx, |thread_store, window, cx| {                                    thread_store.open_thread(&thread_id, window, cx)                                })                                .ok()?                                .await                                .log_err()?;                            let context = context_store                                .update(cx, |context_store, cx| {                                    context_store.add_thread(thread, false, cx)                                })                                .ok()??;                            Some(context)                        })                    }                    ThreadContextEntry::Context { path, .. } => {                        let path = path.clone();                        let context_store = context_store.clone();                        let text_thread_store = text_thread_store.clone();                        cx.spawn::<_, Option<_>>(async move |cx| {                            let thread = text_thread_store                                .update(cx, |store, cx| store.open_local_context(path, cx))                                .ok()?                                .await                                .log_err()?;                            let context = context_store                                .update(cx, |context_store, cx| {                                    context_store.add_text_thread(thread, false, cx)                                })                                .ok()??;                            Some(context)                        })                    }                },            )),        }    }     fn completion_for_rules(        rules: RulesContextEntry,        excerpt_id: ExcerptId,        source_range: Range<Anchor>,        editor: Entity<Editor>,        context_store: Entity<ContextStore>,    ) -> Completion {        let new_text = format!("{} ", MentionLink::for_rule(&rules));        let new_text_len = new_text.len();        Completion {            replace_range: source_range.clone(),            new_text,            label: CodeLabel::plain(rules.title.to_string(), None),

Lines 1–500 of 1506.