934 lines · 30.7 KB · Rust
RawDownload
mod completion_provider;mod fetch_context_picker;mod file_context_picker;mod rules_context_picker;mod symbol_context_picker;mod thread_context_picker; use std::ops::Range;use std::path::PathBuf;use std::sync::Arc; use anyhow::{Result, anyhow};use editor::display_map::{Crease, FoldId};use editor::{Anchor, AnchorRangeExt as _, Editor, ExcerptId, FoldPlaceholder, ToOffset};use file_context_picker::render_file_context_entry;use gpui::{    App, DismissEvent, Empty, Entity, EventEmitter, FocusHandle, Focusable, Subscription, Task,    WeakEntity,};use language::Buffer;use multi_buffer::MultiBufferRow;use project::{Entry, ProjectPath};use prompt_store::UserPromptId;use rules_context_picker::RulesContextEntry;use symbol_context_picker::SymbolContextPicker;use thread_context_picker::{ThreadContextEntry, render_thread_context_entry};use ui::{    ButtonLike, ContextMenu, ContextMenuEntry, ContextMenuItem, Disclosure, TintColor, prelude::*,};use uuid::Uuid;use workspace::{Workspace, notifications::NotifyResultExt}; use crate::AssistantPanel;use crate::context::RULES_ICON;pub use crate::context_picker::completion_provider::ContextPickerCompletionProvider;use crate::context_picker::fetch_context_picker::FetchContextPicker;use crate::context_picker::file_context_picker::FileContextPicker;use crate::context_picker::rules_context_picker::RulesContextPicker;use crate::context_picker::thread_context_picker::ThreadContextPicker;use crate::context_store::ContextStore;use crate::thread::ThreadId;use crate::thread_store::ThreadStore; #[derive(Debug, Clone, Copy, PartialEq, Eq)]enum ContextPickerEntry {    Mode(ContextPickerMode),    Action(ContextPickerAction),} impl ContextPickerEntry {    pub fn keyword(&self) -> &'static str {        match self {            Self::Mode(mode) => mode.keyword(),            Self::Action(action) => action.keyword(),        }    }     pub fn label(&self) -> &'static str {        match self {            Self::Mode(mode) => mode.label(),            Self::Action(action) => action.label(),        }    }     pub fn icon(&self) -> IconName {        match self {            Self::Mode(mode) => mode.icon(),            Self::Action(action) => action.icon(),        }    }} #[derive(Debug, Clone, Copy, PartialEq, Eq)]enum ContextPickerMode {    File,    Symbol,    Fetch,    Thread,    Rules,} #[derive(Debug, Clone, Copy, PartialEq, Eq)]enum ContextPickerAction {    AddSelections,} impl ContextPickerAction {    pub fn keyword(&self) -> &'static str {        match self {            Self::AddSelections => "selection",        }    }     pub fn label(&self) -> &'static str {        match self {            Self::AddSelections => "Selection",        }    }     pub fn icon(&self) -> IconName {        match self {            Self::AddSelections => IconName::Context,        }    }} impl TryFrom<&str> for ContextPickerMode {    type Error = String;     fn try_from(value: &str) -> Result<Self, Self::Error> {        match value {            "file" => Ok(Self::File),            "symbol" => Ok(Self::Symbol),            "fetch" => Ok(Self::Fetch),            "thread" => Ok(Self::Thread),            "rules" => Ok(Self::Rules),            _ => Err(format!("Invalid context picker mode: {}", value)),        }    }} impl ContextPickerMode {    pub fn keyword(&self) -> &'static str {        match self {            Self::File => "file",            Self::Symbol => "symbol",            Self::Fetch => "fetch",            Self::Thread => "thread",            Self::Rules => "rules",        }    }     pub fn label(&self) -> &'static str {        match self {            Self::File => "Files & Directories",            Self::Symbol => "Symbols",            Self::Fetch => "Fetch",            Self::Thread => "Threads",            Self::Rules => "Rules",        }    }     pub fn icon(&self) -> IconName {        match self {            Self::File => IconName::File,            Self::Symbol => IconName::Code,            Self::Fetch => IconName::Globe,            Self::Thread => IconName::MessageBubbles,            Self::Rules => RULES_ICON,        }    }} #[derive(Debug, Clone)]enum ContextPickerState {    Default(Entity<ContextMenu>),    File(Entity<FileContextPicker>),    Symbol(Entity<SymbolContextPicker>),    Fetch(Entity<FetchContextPicker>),    Thread(Entity<ThreadContextPicker>),    Rules(Entity<RulesContextPicker>),} pub(super) struct ContextPicker {    mode: ContextPickerState,    workspace: WeakEntity<Workspace>,    context_store: WeakEntity<ContextStore>,    thread_store: Option<WeakEntity<ThreadStore>>,    _subscriptions: Vec<Subscription>,} impl ContextPicker {    pub fn new(        workspace: WeakEntity<Workspace>,        thread_store: Option<WeakEntity<ThreadStore>>,        context_store: WeakEntity<ContextStore>,        window: &mut Window,        cx: &mut Context<Self>,    ) -> Self {        let subscriptions = context_store            .upgrade()            .map(|context_store| {                cx.observe(&context_store, |this, _, cx| this.notify_current_picker(cx))            })            .into_iter()            .chain(                thread_store                    .as_ref()                    .and_then(|thread_store| thread_store.upgrade())                    .map(|thread_store| {                        cx.observe(&thread_store, |this, _, cx| this.notify_current_picker(cx))                    }),            )            .collect::<Vec<Subscription>>();         ContextPicker {            mode: ContextPickerState::Default(ContextMenu::build(                window,                cx,                |menu, _window, _cx| menu,            )),            workspace,            context_store,            thread_store,            _subscriptions: subscriptions,        }    }     pub fn init(&mut self, window: &mut Window, cx: &mut Context<Self>) {        self.mode = ContextPickerState::Default(self.build_menu(window, cx));        cx.notify();    }     fn build_menu(&mut self, window: &mut Window, cx: &mut Context<Self>) -> Entity<ContextMenu> {        let context_picker = cx.entity().clone();         let menu = ContextMenu::build(window, cx, move |menu, _window, cx| {            let recent = self.recent_entries(cx);            let has_recent = !recent.is_empty();            let recent_entries = recent                .into_iter()                .enumerate()                .map(|(ix, entry)| self.recent_menu_item(context_picker.clone(), ix, entry));             let entries = self                .workspace                .upgrade()                .map(|workspace| {                    available_context_picker_entries(&self.thread_store, &workspace, cx)                })                .unwrap_or_default();             menu.when(has_recent, |menu| {                menu.custom_row(|_, _| {                    div()                        .mb_1()                        .child(                            Label::new("Recent")                                .color(Color::Muted)                                .size(LabelSize::Small),                        )                        .into_any_element()                })            })            .extend(recent_entries)            .when(has_recent, |menu| menu.separator())            .extend(entries.into_iter().map(|entry| {                let context_picker = context_picker.clone();                 ContextMenuEntry::new(entry.label())                    .icon(entry.icon())                    .icon_size(IconSize::XSmall)                    .icon_color(Color::Muted)                    .handler(move |window, cx| {                        context_picker.update(cx, |this, cx| this.select_entry(entry, window, cx))                    })            }))            .keep_open_on_confirm()        });         cx.subscribe(&menu, move |_, _, _: &DismissEvent, cx| {            cx.emit(DismissEvent);        })        .detach();         menu    }     /// Whether threads are allowed as context.    pub fn allow_threads(&self) -> bool {        self.thread_store.is_some()    }     fn select_entry(        &mut self,        entry: ContextPickerEntry,        window: &mut Window,        cx: &mut Context<Self>,    ) {        let context_picker = cx.entity().downgrade();         match entry {            ContextPickerEntry::Mode(mode) => match mode {                ContextPickerMode::File => {                    self.mode = ContextPickerState::File(cx.new(|cx| {                        FileContextPicker::new(                            context_picker.clone(),                            self.workspace.clone(),                            self.context_store.clone(),                            window,                            cx,                        )                    }));                }                ContextPickerMode::Symbol => {                    self.mode = ContextPickerState::Symbol(cx.new(|cx| {                        SymbolContextPicker::new(                            context_picker.clone(),                            self.workspace.clone(),                            self.context_store.clone(),                            window,                            cx,                        )                    }));                }                ContextPickerMode::Rules => {                    if let Some(thread_store) = self.thread_store.as_ref() {                        self.mode = ContextPickerState::Rules(cx.new(|cx| {                            RulesContextPicker::new(                                thread_store.clone(),                                context_picker.clone(),                                self.context_store.clone(),                                window,                                cx,                            )                        }));                    }                }                ContextPickerMode::Fetch => {                    self.mode = ContextPickerState::Fetch(cx.new(|cx| {                        FetchContextPicker::new(                            context_picker.clone(),                            self.workspace.clone(),                            self.context_store.clone(),                            window,                            cx,                        )                    }));                }                ContextPickerMode::Thread => {                    if let Some(thread_store) = self.thread_store.as_ref() {                        self.mode = ContextPickerState::Thread(cx.new(|cx| {                            ThreadContextPicker::new(                                thread_store.clone(),                                context_picker.clone(),                                self.context_store.clone(),                                window,                                cx,                            )                        }));                    }                }            },            ContextPickerEntry::Action(action) => match action {                ContextPickerAction::AddSelections => {                    if let Some((context_store, workspace)) =                        self.context_store.upgrade().zip(self.workspace.upgrade())                    {                        add_selections_as_context(&context_store, &workspace, cx);                    }                     cx.emit(DismissEvent);                }            },        }         cx.notify();        cx.focus_self(window);    }     fn recent_menu_item(        &self,        context_picker: Entity<ContextPicker>,        ix: usize,        entry: RecentEntry,    ) -> ContextMenuItem {        match entry {            RecentEntry::File {                project_path,                path_prefix,            } => {                let context_store = self.context_store.clone();                let worktree_id = project_path.worktree_id;                let path = project_path.path.clone();                 ContextMenuItem::custom_entry(                    move |_window, cx| {                        render_file_context_entry(                            ElementId::NamedInteger("ctx-recent".into(), ix),                            worktree_id,                            &path,                            &path_prefix,                            false,                            context_store.clone(),                            cx,                        )                        .into_any()                    },                    move |window, cx| {                        context_picker.update(cx, |this, cx| {                            this.add_recent_file(project_path.clone(), window, cx);                        })                    },                )            }            RecentEntry::Thread(thread) => {                let context_store = self.context_store.clone();                let view_thread = thread.clone();                 ContextMenuItem::custom_entry(                    move |_window, cx| {                        render_thread_context_entry(&view_thread, context_store.clone(), cx)                            .into_any()                    },                    move |_window, cx| {                        context_picker.update(cx, |this, cx| {                            this.add_recent_thread(thread.clone(), cx)                                .detach_and_log_err(cx);                        })                    },                )            }        }    }     fn add_recent_file(        &self,        project_path: ProjectPath,        window: &mut Window,        cx: &mut Context<Self>,    ) {        let Some(context_store) = self.context_store.upgrade() else {            return;        };         let task = context_store.update(cx, |context_store, cx| {            context_store.add_file_from_path(project_path.clone(), true, cx)        });         cx.spawn_in(window, async move |_, cx| task.await.notify_async_err(cx))            .detach();         cx.notify();    }     fn add_recent_thread(        &self,        thread: ThreadContextEntry,        cx: &mut Context<Self>,    ) -> Task<Result<()>> {        let Some(context_store) = self.context_store.upgrade() else {            return Task::ready(Err(anyhow!("context store not available")));        };         let Some(thread_store) = self            .thread_store            .as_ref()            .and_then(|thread_store| thread_store.upgrade())        else {            return Task::ready(Err(anyhow!("thread store not available")));        };         let open_thread_task = thread_store.update(cx, |this, cx| this.open_thread(&thread.id, cx));        cx.spawn(async move |this, cx| {            let thread = open_thread_task.await?;            context_store.update(cx, |context_store, cx| {                context_store.add_thread(thread, true, cx);            })?;             this.update(cx, |_this, cx| cx.notify())        })    }     fn recent_entries(&self, cx: &mut App) -> Vec<RecentEntry> {        let Some(workspace) = self.workspace.upgrade() else {            return vec![];        };         let Some(context_store) = self.context_store.upgrade() else {            return vec![];        };         recent_context_picker_entries(context_store, self.thread_store.clone(), workspace, cx)    }     fn notify_current_picker(&mut self, cx: &mut Context<Self>) {        match &self.mode {            ContextPickerState::Default(entity) => entity.update(cx, |_, cx| cx.notify()),            ContextPickerState::File(entity) => entity.update(cx, |_, cx| cx.notify()),            ContextPickerState::Symbol(entity) => entity.update(cx, |_, cx| cx.notify()),            ContextPickerState::Fetch(entity) => entity.update(cx, |_, cx| cx.notify()),            ContextPickerState::Thread(entity) => entity.update(cx, |_, cx| cx.notify()),            ContextPickerState::Rules(entity) => entity.update(cx, |_, cx| cx.notify()),        }    }} impl EventEmitter<DismissEvent> for ContextPicker {} impl Focusable for ContextPicker {    fn focus_handle(&self, cx: &App) -> FocusHandle {        match &self.mode {            ContextPickerState::Default(menu) => menu.focus_handle(cx),            ContextPickerState::File(file_picker) => file_picker.focus_handle(cx),            ContextPickerState::Symbol(symbol_picker) => symbol_picker.focus_handle(cx),            ContextPickerState::Fetch(fetch_picker) => fetch_picker.focus_handle(cx),            ContextPickerState::Thread(thread_picker) => thread_picker.focus_handle(cx),            ContextPickerState::Rules(user_rules_picker) => user_rules_picker.focus_handle(cx),        }    }

Lines 1–500 of 934.