1293 lines · 41.7 KB · Rust
RawDownload
use std::fmt::{self, Display, Formatter, Write as _};use std::hash::{Hash, Hasher};use std::path::PathBuf;use std::{ops::Range, path::Path, sync::Arc}; use assistant_context_editor::AssistantContext;use assistant_tool::outline;use collections::{HashMap, HashSet};use editor::display_map::CreaseId;use editor::{Addon, Editor};use futures::future;use futures::{FutureExt, future::Shared};use gpui::{App, AppContext as _, Entity, SharedString, Subscription, Task};use language::{Buffer, ParseStatus};use language_model::{LanguageModelImage, LanguageModelRequestMessage, MessageContent};use project::{Project, ProjectEntryId, ProjectPath, Worktree};use prompt_store::{PromptStore, UserPromptId};use ref_cast::RefCast;use rope::Point;use text::{Anchor, OffsetRangeExt as _};use ui::{Context, ElementId, IconName};use util::markdown::MarkdownCodeBlock;use util::{ResultExt as _, post_inc}; use crate::context_store::{ContextStore, ContextStoreEvent};use crate::thread::Thread; pub const RULES_ICON: IconName = IconName::Context; pub enum ContextKind {    File,    Directory,    Symbol,    Selection,    FetchedUrl,    Thread,    TextThread,    Rules,    Image,} impl ContextKind {    pub fn icon(&self) -> IconName {        match self {            ContextKind::File => IconName::File,            ContextKind::Directory => IconName::Folder,            ContextKind::Symbol => IconName::Code,            ContextKind::Selection => IconName::Context,            ContextKind::FetchedUrl => IconName::Globe,            ContextKind::Thread => IconName::MessageBubbles,            ContextKind::TextThread => IconName::MessageBubbles,            ContextKind::Rules => RULES_ICON,            ContextKind::Image => IconName::Image,        }    }} /// Handle for context that can be attached to a user message.////// This uses IDs that are stable enough for tracking renames and identifying when context has/// already been added to the thread. To use this in a set, wrap it in `AgentContextKey` to opt in/// to `PartialEq` and `Hash` impls that use the subset of the fields used for this stable identity.#[derive(Debug, Clone)]pub enum AgentContextHandle {    File(FileContextHandle),    Directory(DirectoryContextHandle),    Symbol(SymbolContextHandle),    Selection(SelectionContextHandle),    FetchedUrl(FetchedUrlContext),    Thread(ThreadContextHandle),    TextThread(TextThreadContextHandle),    Rules(RulesContextHandle),    Image(ImageContext),} impl AgentContextHandle {    pub fn id(&self) -> ContextId {        match self {            Self::File(context) => context.context_id,            Self::Directory(context) => context.context_id,            Self::Symbol(context) => context.context_id,            Self::Selection(context) => context.context_id,            Self::FetchedUrl(context) => context.context_id,            Self::Thread(context) => context.context_id,            Self::TextThread(context) => context.context_id,            Self::Rules(context) => context.context_id,            Self::Image(context) => context.context_id,        }    }     pub fn element_id(&self, name: SharedString) -> ElementId {        ElementId::NamedInteger(name, self.id().0)    }} /// Loaded context that can be attached to a user message. This can be thought of as a/// snapshot of the context along with an `AgentContextHandle`.#[derive(Debug, Clone)]pub enum AgentContext {    File(FileContext),    Directory(DirectoryContext),    Symbol(SymbolContext),    Selection(SelectionContext),    FetchedUrl(FetchedUrlContext),    Thread(ThreadContext),    TextThread(TextThreadContext),    Rules(RulesContext),    Image(ImageContext),} impl AgentContext {    pub fn handle(&self) -> AgentContextHandle {        match self {            AgentContext::File(context) => AgentContextHandle::File(context.handle.clone()),            AgentContext::Directory(context) => {                AgentContextHandle::Directory(context.handle.clone())            }            AgentContext::Symbol(context) => AgentContextHandle::Symbol(context.handle.clone()),            AgentContext::Selection(context) => {                AgentContextHandle::Selection(context.handle.clone())            }            AgentContext::FetchedUrl(context) => AgentContextHandle::FetchedUrl(context.clone()),            AgentContext::Thread(context) => AgentContextHandle::Thread(context.handle.clone()),            AgentContext::TextThread(context) => {                AgentContextHandle::TextThread(context.handle.clone())            }            AgentContext::Rules(context) => AgentContextHandle::Rules(context.handle.clone()),            AgentContext::Image(context) => AgentContextHandle::Image(context.clone()),        }    }} /// ID created at time of context add, for use in ElementId. This is not the stable identity of a/// context, instead that's handled by the `PartialEq` and `Hash` impls of `AgentContextKey`.#[derive(Debug, Copy, Clone)]pub struct ContextId(u64); impl ContextId {    pub fn zero() -> Self {        ContextId(0)    }     fn for_lookup() -> Self {        ContextId(u64::MAX)    }     pub fn post_inc(&mut self) -> Self {        Self(post_inc(&mut self.0))    }} /// File context provides the entire contents of a file.////// This holds an `Entity<Buffer>` so that file path renames affect its display and so that it can/// be opened even if the file has been deleted. An alternative might be to use `ProjectEntryId`,/// but then when deleted there is no path info or ability to open.#[derive(Debug, Clone)]pub struct FileContextHandle {    pub buffer: Entity<Buffer>,    pub context_id: ContextId,} #[derive(Debug, Clone)]pub struct FileContext {    pub handle: FileContextHandle,    pub full_path: Arc<Path>,    pub text: SharedString,    pub is_outline: bool,} impl FileContextHandle {    pub fn eq_for_key(&self, other: &Self) -> bool {        self.buffer == other.buffer    }     pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {        self.buffer.hash(state)    }     pub fn project_path(&self, cx: &App) -> Option<ProjectPath> {        let file = self.buffer.read(cx).file()?;        Some(ProjectPath {            worktree_id: file.worktree_id(cx),            path: file.path().clone(),        })    }     fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {        let buffer_ref = self.buffer.read(cx);        let Some(file) = buffer_ref.file() else {            log::error!("file context missing path");            return Task::ready(None);        };        let full_path: Arc<Path> = file.full_path(cx).into();        let rope = buffer_ref.as_rope().clone();        let buffer = self.buffer.clone();         cx.spawn(async move |cx| {            // For large files, use outline instead of full content            if rope.len() > outline::AUTO_OUTLINE_SIZE {                // Wait until the buffer has been fully parsed, so we can read its outline                if let Ok(mut parse_status) =                    buffer.read_with(cx, |buffer, _| buffer.parse_status())                {                    while *parse_status.borrow() != ParseStatus::Idle {                        parse_status.changed().await.log_err();                    }                     if let Ok(snapshot) = buffer.read_with(cx, |buffer, _| buffer.snapshot()) {                        if let Some(outline) = snapshot.outline(None) {                            let items = outline                                .items                                .into_iter()                                .map(|item| item.to_point(&snapshot));                             if let Ok(outline_text) =                                outline::render_outline(items, None, 0, usize::MAX).await                            {                                let context = AgentContext::File(FileContext {                                    handle: self,                                    full_path,                                    text: outline_text.into(),                                    is_outline: true,                                });                                return Some((context, vec![buffer]));                            }                        }                    }                }            }             // Fallback to full content if we couldn't build an outline            // (or didn't need to because the file was small enough)            let context = AgentContext::File(FileContext {                handle: self,                full_path,                text: rope.to_string().into(),                is_outline: false,            });            Some((context, vec![buffer]))        })    }} impl Display for FileContext {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        write!(            f,            "{}",            MarkdownCodeBlock {                tag: &codeblock_tag(&self.full_path, None),                text: &self.text,            }        )    }} /// Directory contents provides the entire contents of text files in a directory.////// This has a `ProjectEntryId` so that it follows renames.#[derive(Debug, Clone)]pub struct DirectoryContextHandle {    pub entry_id: ProjectEntryId,    pub context_id: ContextId,} #[derive(Debug, Clone)]pub struct DirectoryContext {    pub handle: DirectoryContextHandle,    pub full_path: Arc<Path>,    pub descendants: Vec<DirectoryContextDescendant>,} #[derive(Debug, Clone)]pub struct DirectoryContextDescendant {    /// Path within the directory.    pub rel_path: Arc<Path>,    pub fenced_codeblock: SharedString,} impl DirectoryContextHandle {    pub fn eq_for_key(&self, other: &Self) -> bool {        self.entry_id == other.entry_id    }     pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {        self.entry_id.hash(state)    }     fn load(        self,        project: Entity<Project>,        cx: &mut App,    ) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {        let Some(worktree) = project.read(cx).worktree_for_entry(self.entry_id, cx) else {            return Task::ready(None);        };        let worktree_ref = worktree.read(cx);        let Some(entry) = worktree_ref.entry_for_id(self.entry_id) else {            return Task::ready(None);        };        if entry.is_file() {            log::error!("DirectoryContext unexpectedly refers to a file.");            return Task::ready(None);        }         let directory_path = entry.path.clone();        let directory_full_path = worktree_ref.full_path(&directory_path).into();         let file_paths = collect_files_in_path(worktree_ref, &directory_path);        let descendants_future = future::join_all(file_paths.into_iter().map(|path| {            let worktree_ref = worktree.read(cx);            let worktree_id = worktree_ref.id();            let full_path = worktree_ref.full_path(&path);             let rel_path = path                .strip_prefix(&directory_path)                .log_err()                .map_or_else(|| path.clone(), |rel_path| rel_path.into());             let open_task = project.update(cx, |project, cx| {                project.buffer_store().update(cx, |buffer_store, cx| {                    let project_path = ProjectPath { worktree_id, path };                    buffer_store.open_buffer(project_path, cx)                })            });             // TODO: report load errors instead of just logging            let rope_task = cx.spawn(async move |cx| {                let buffer = open_task.await.log_err()?;                let rope = buffer                    .read_with(cx, |buffer, _cx| buffer.as_rope().clone())                    .log_err()?;                Some((rope, buffer))            });             cx.background_spawn(async move {                let (rope, buffer) = rope_task.await?;                let fenced_codeblock = MarkdownCodeBlock {                    tag: &codeblock_tag(&full_path, None),                    text: &rope.to_string(),                }                .to_string()                .into();                let descendant = DirectoryContextDescendant {                    rel_path,                    fenced_codeblock,                };                Some((descendant, buffer))            })        }));         cx.background_spawn(async move {            let (descendants, buffers) = descendants_future.await.into_iter().flatten().unzip();            let context = AgentContext::Directory(DirectoryContext {                handle: self,                full_path: directory_full_path,                descendants,            });            Some((context, buffers))        })    }} impl Display for DirectoryContext {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        let mut is_first = true;        for descendant in &self.descendants {            if !is_first {                write!(f, "\n")?;            } else {                is_first = false;            }            write!(f, "{}", descendant.fenced_codeblock)?;        }        Ok(())    }} #[derive(Debug, Clone)]pub struct SymbolContextHandle {    pub buffer: Entity<Buffer>,    pub symbol: SharedString,    pub range: Range<Anchor>,    /// The range that fully contains the symbol. e.g. for function symbol, this will include not    /// only the signature, but also the body. Not used by `PartialEq` or `Hash` for    /// `AgentContextKey`.    pub enclosing_range: Range<Anchor>,    pub context_id: ContextId,} #[derive(Debug, Clone)]pub struct SymbolContext {    pub handle: SymbolContextHandle,    pub full_path: Arc<Path>,    pub line_range: Range<Point>,    pub text: SharedString,} impl SymbolContextHandle {    pub fn eq_for_key(&self, other: &Self) -> bool {        self.buffer == other.buffer && self.symbol == other.symbol && self.range == other.range    }     pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {        self.buffer.hash(state);        self.symbol.hash(state);        self.range.hash(state);    }     pub fn full_path(&self, cx: &App) -> Option<PathBuf> {        Some(self.buffer.read(cx).file()?.full_path(cx))    }     pub fn enclosing_line_range(&self, cx: &App) -> Range<Point> {        self.enclosing_range            .to_point(&self.buffer.read(cx).snapshot())    }     pub fn text(&self, cx: &App) -> SharedString {        self.buffer            .read(cx)            .text_for_range(self.enclosing_range.clone())            .collect::<String>()            .into()    }     fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {        let buffer_ref = self.buffer.read(cx);        let Some(file) = buffer_ref.file() else {            log::error!("symbol context's file has no path");            return Task::ready(None);        };        let full_path = file.full_path(cx).into();        let line_range = self.enclosing_range.to_point(&buffer_ref.snapshot());        let text = self.text(cx);        let buffer = self.buffer.clone();        let context = AgentContext::Symbol(SymbolContext {            handle: self,            full_path,            line_range,            text,        });        Task::ready(Some((context, vec![buffer])))    }} impl Display for SymbolContext {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        let code_block = MarkdownCodeBlock {            tag: &codeblock_tag(&self.full_path, Some(self.line_range.clone())),            text: &self.text,        };        write!(f, "{code_block}",)    }} #[derive(Debug, Clone)]pub struct SelectionContextHandle {    pub buffer: Entity<Buffer>,    pub range: Range<Anchor>,    pub context_id: ContextId,} #[derive(Debug, Clone)]pub struct SelectionContext {    pub handle: SelectionContextHandle,    pub full_path: Arc<Path>,    pub line_range: Range<Point>,    pub text: SharedString,} impl SelectionContextHandle {    pub fn eq_for_key(&self, other: &Self) -> bool {        self.buffer == other.buffer && self.range == other.range    }     pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {        self.buffer.hash(state);        self.range.hash(state);    }     pub fn full_path(&self, cx: &App) -> Option<PathBuf> {        Some(self.buffer.read(cx).file()?.full_path(cx))    }     pub fn line_range(&self, cx: &App) -> Range<Point> {        self.range.to_point(&self.buffer.read(cx).snapshot())    }     pub fn text(&self, cx: &App) -> SharedString {        self.buffer            .read(cx)            .text_for_range(self.range.clone())            .collect::<String>()            .into()    }     fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {        let Some(full_path) = self.full_path(cx) else {

Lines 1–500 of 1293.