use std::path::{Path, PathBuf};use std::sync::Arc; use anyhow::{anyhow, bail, Result};use collections::{BTreeMap, HashMap, HashSet};use futures::{self, future, Future, FutureExt};use gpui::{App, AppContext as _, AsyncApp, Context, Entity, SharedString, Task, WeakEntity};use language::Buffer;use project::{ProjectPath, Worktree};use rope::Rope;use text::BufferId;use workspace::Workspace; use crate::context::{    AssistantContext, ContextBuffer, ContextId, ContextSnapshot, DirectoryContext,    FetchedUrlContext, FileContext, ThreadContext,};use crate::context_strip::SuggestedContext;use crate::thread::{Thread, ThreadId}; pub struct ContextStore {    workspace: WeakEntity<Workspace>,    context: Vec<AssistantContext>,    // TODO: If an EntityId is used for all context types (like BufferId), can remove ContextId.    next_context_id: ContextId,    files: BTreeMap<BufferId, ContextId>,    directories: HashMap<PathBuf, ContextId>,    threads: HashMap<ThreadId, ContextId>,    fetched_urls: HashMap<String, ContextId>,} impl ContextStore {    pub fn new(workspace: WeakEntity<Workspace>) -> Self {        Self {            workspace,            context: Vec::new(),            next_context_id: ContextId(0),            files: BTreeMap::default(),            directories: HashMap::default(),            threads: HashMap::default(),            fetched_urls: HashMap::default(),        }    }     pub fn snapshot<'a>(&'a self, cx: &'a App) -> impl Iterator<Item = ContextSnapshot> + 'a {        self.context()            .iter()            .flat_map(|context| context.snapshot(cx))    }     pub fn context(&self) -> &Vec<AssistantContext> {        &self.context    }     pub fn clear(&mut self) {        self.context.clear();        self.files.clear();        self.directories.clear();        self.threads.clear();        self.fetched_urls.clear();    }     pub fn add_file_from_path(        &mut self,        project_path: ProjectPath,        cx: &mut Context<Self>,    ) -> Task<Result<()>> {        let workspace = self.workspace.clone();         let Some(project) = workspace            .upgrade()            .map(|workspace| workspace.read(cx).project().clone())        else {            return Task::ready(Err(anyhow!("failed to read project")));        };         cx.spawn(|this, mut cx| async move {            let open_buffer_task = project.update(&mut cx, |project, cx| {                project.open_buffer(project_path.clone(), cx)            })?;             let buffer_entity = open_buffer_task.await?;            let buffer_id = this.update(&mut cx, |_, cx| buffer_entity.read(cx).remote_id())?;             let already_included = this.update(&mut cx, |this, _cx| {                match this.will_include_buffer(buffer_id, &project_path.path) {                    Some(FileInclusion::Direct(context_id)) => {                        this.remove_context(context_id);                        true                    }                    Some(FileInclusion::InDirectory(_)) => true,                    None => false,                }            })?;             if already_included {                return anyhow::Ok(());            }             let (buffer_info, text_task) = this.update(&mut cx, |_, cx| {                let buffer = buffer_entity.read(cx);                collect_buffer_info_and_text(                    project_path.path.clone(),                    buffer_entity,                    buffer,                    cx.to_async(),                )            })?;             let text = text_task.await;             this.update(&mut cx, |this, _cx| {                this.insert_file(make_context_buffer(buffer_info, text));            })?;             anyhow::Ok(())        })    }     pub fn add_file_from_buffer(        &mut self,        buffer_entity: Entity<Buffer>,        cx: &mut Context<Self>,    ) -> Task<Result<()>> {        cx.spawn(|this, mut cx| async move {            let (buffer_info, text_task) = this.update(&mut cx, |_, cx| {                let buffer = buffer_entity.read(cx);                let Some(file) = buffer.file() else {                    return Err(anyhow!("Buffer has no path."));                };                Ok(collect_buffer_info_and_text(                    file.path().clone(),                    buffer_entity,                    buffer,                    cx.to_async(),                ))            })??;             let text = text_task.await;             this.update(&mut cx, |this, _cx| {                this.insert_file(make_context_buffer(buffer_info, text))            })?;             anyhow::Ok(())        })    }     fn insert_file(&mut self, context_buffer: ContextBuffer) {        let id = self.next_context_id.post_inc();        self.files.insert(context_buffer.id, id);        self.context            .push(AssistantContext::File(FileContext { id, context_buffer }));    }     pub fn add_directory(        &mut self,        project_path: ProjectPath,        cx: &mut Context<Self>,    ) -> Task<Result<()>> {        let workspace = self.workspace.clone();        let Some(project) = workspace            .upgrade()            .map(|workspace| workspace.read(cx).project().clone())        else {            return Task::ready(Err(anyhow!("failed to read project")));        };         let already_included = if let Some(context_id) = self.includes_directory(&project_path.path)        {            self.remove_context(context_id);            true        } else {            false        };        if already_included {            return Task::ready(Ok(()));        }         let worktree_id = project_path.worktree_id;        cx.spawn(|this, mut cx| async move {            let worktree = project.update(&mut cx, |project, cx| {                project                    .worktree_for_id(worktree_id, cx)                    .ok_or_else(|| anyhow!("no worktree found for {worktree_id:?}"))            })??;             let files = worktree.update(&mut cx, |worktree, _cx| {                collect_files_in_path(worktree, &project_path.path)            })?;             let open_buffers_task = project.update(&mut cx, |project, cx| {                let tasks = files.iter().map(|file_path| {                    project.open_buffer(                        ProjectPath {                            worktree_id,                            path: file_path.clone(),                        },                        cx,                    )                });                future::join_all(tasks)            })?;             let buffers = open_buffers_task.await;             let mut buffer_infos = Vec::new();            let mut text_tasks = Vec::new();            this.update(&mut cx, |_, cx| {                for (path, buffer_entity) in files.into_iter().zip(buffers) {                    let buffer_entity = buffer_entity?;                    let buffer = buffer_entity.read(cx);                    let (buffer_info, text_task) =                        collect_buffer_info_and_text(path, buffer_entity, buffer, cx.to_async());                    buffer_infos.push(buffer_info);                    text_tasks.push(text_task);                }                anyhow::Ok(())            })??;             let buffer_texts = future::join_all(text_tasks).await;            let context_buffers = buffer_infos                .into_iter()                .zip(buffer_texts)                .map(|(info, text)| make_context_buffer(info, text))                .collect::<Vec<_>>();             if context_buffers.is_empty() {                bail!("No text files found in {}", &project_path.path.display());            }             this.update(&mut cx, |this, _| {                this.insert_directory(&project_path.path, context_buffers);            })?;             anyhow::Ok(())        })    }     fn insert_directory(&mut self, path: &Path, context_buffers: Vec<ContextBuffer>) {        let id = self.next_context_id.post_inc();        self.directories.insert(path.to_path_buf(), id);         self.context            .push(AssistantContext::Directory(DirectoryContext::new(                id,                path,                context_buffers,            )));    }     pub fn add_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {        if let Some(context_id) = self.includes_thread(&thread.read(cx).id()) {            self.remove_context(context_id);        } else {            self.insert_thread(thread, cx);        }    }     fn insert_thread(&mut self, thread: Entity<Thread>, cx: &App) {        let id = self.next_context_id.post_inc();        let text = thread.read(cx).text().into();         self.threads.insert(thread.read(cx).id().clone(), id);        self.context            .push(AssistantContext::Thread(ThreadContext { id, thread, text }));    }     pub fn add_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {        if self.includes_url(&url).is_none() {            self.insert_fetched_url(url, text);        }    }     fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {        let id = self.next_context_id.post_inc();         self.fetched_urls.insert(url.clone(), id);        self.context            .push(AssistantContext::FetchedUrl(FetchedUrlContext {                id,                url: url.into(),                text: text.into(),            }));    }     pub fn accept_suggested_context(        &mut self,        suggested: &SuggestedContext,        cx: &mut Context<ContextStore>,    ) -> Task<Result<()>> {        match suggested {            SuggestedContext::File {                buffer,                icon_path: _,                name: _,            } => {                if let Some(buffer) = buffer.upgrade() {                    return self.add_file_from_buffer(buffer, cx);                };            }            SuggestedContext::Thread { thread, name: _ } => {                if let Some(thread) = thread.upgrade() {                    self.insert_thread(thread, cx);                };            }        }        Task::ready(Ok(()))    }     pub fn remove_context(&mut self, id: ContextId) {        let Some(ix) = self.context.iter().position(|context| context.id() == id) else {            return;        };         match self.context.remove(ix) {            AssistantContext::File(_) => {                self.files.retain(|_, context_id| *context_id != id);            }            AssistantContext::Directory(_) => {                self.directories.retain(|_, context_id| *context_id != id);            }            AssistantContext::FetchedUrl(_) => {                self.fetched_urls.retain(|_, context_id| *context_id != id);            }            AssistantContext::Thread(_) => {                self.threads.retain(|_, context_id| *context_id != id);            }        }    }     /// Returns whether the buffer is already included directly in the context, or if it will be    /// included in the context via a directory. Directory inclusion is based on paths rather than    /// buffer IDs as the directory will be re-scanned.    pub fn will_include_buffer(&self, buffer_id: BufferId, path: &Path) -> Option<FileInclusion> {        if let Some(context_id) = self.files.get(&buffer_id) {            return Some(FileInclusion::Direct(*context_id));        }         self.will_include_file_path_via_directory(path)    }     /// Returns whether this file path is already included directly in the context, or if it will be    /// included in the context via a directory.    pub fn will_include_file_path(&self, path: &Path, cx: &App) -> Option<FileInclusion> {        if !self.files.is_empty() {            let found_file_context = self.context.iter().find(|context| match &context {                AssistantContext::File(file_context) => {                    let buffer = file_context.context_buffer.buffer.read(cx);                    if let Some(file_path) = buffer_path_log_err(buffer) {                        *file_path == *path                    } else {                        false                    }                }                _ => false,            });            if let Some(context) = found_file_context {                return Some(FileInclusion::Direct(context.id()));            }        }         self.will_include_file_path_via_directory(path)    }     fn will_include_file_path_via_directory(&self, path: &Path) -> Option<FileInclusion> {        if self.directories.is_empty() {            return None;        }         let mut buf = path.to_path_buf();         while buf.pop() {            if let Some(_) = self.directories.get(&buf) {                return Some(FileInclusion::InDirectory(buf));            }        }         None    }     pub fn includes_directory(&self, path: &Path) -> Option<ContextId> {        self.directories.get(path).copied()    }     pub fn includes_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {        self.threads.get(thread_id).copied()    }     pub fn includes_url(&self, url: &str) -> Option<ContextId> {        self.fetched_urls.get(url).copied()    }     /// Replaces the context that matches the ID of the new context, if any match.    fn replace_context(&mut self, new_context: AssistantContext) {        let id = new_context.id();        for context in self.context.iter_mut() {            if context.id() == id {                *context = new_context;                break;            }        }    }     pub fn file_paths(&self, cx: &App) -> HashSet<PathBuf> {        self.context            .iter()            .filter_map(|context| match context {                AssistantContext::File(file) => {                    let buffer = file.context_buffer.buffer.read(cx);                    buffer_path_log_err(buffer).map(|p| p.to_path_buf())                }                AssistantContext::Directory(_)                | AssistantContext::FetchedUrl(_)                | AssistantContext::Thread(_) => None,            })            .collect()    }     pub fn thread_ids(&self) -> HashSet<ThreadId> {        self.threads.keys().cloned().collect()    }} pub enum FileInclusion {    Direct(ContextId),    InDirectory(PathBuf),} // ContextBuffer without text.struct BufferInfo {    buffer_entity: Entity<Buffer>,    id: BufferId,    version: clock::Global,} fn make_context_buffer(info: BufferInfo, text: SharedString) -> ContextBuffer {    ContextBuffer {        id: info.id,        buffer: info.buffer_entity,        version: info.version,        text,    }} fn collect_buffer_info_and_text(    path: Arc<Path>,    buffer_entity: Entity<Buffer>,    buffer: &Buffer,    cx: AsyncApp,) -> (BufferInfo, Task<SharedString>) {    let buffer_info = BufferInfo {        id: buffer.remote_id(),        buffer_entity,        version: buffer.version(),    };    // Important to collect version at the same time as content so that staleness logic is correct.    let content = buffer.as_rope().clone();    let text_task = cx.background_spawn(async move { to_fenced_codeblock(&path, content) });    (buffer_info, text_task)} pub fn buffer_path_log_err(buffer: &Buffer) -> Option<Arc<Path>> {    if let Some(file) = buffer.file() {        Some(file.path().clone())    } else {        log::error!("Buffer that had a path unexpectedly no longer has a path.");        None    }} fn to_fenced_codeblock(path: &Path, content: Rope) -> SharedString {    let path_extension = path.extension().and_then(|ext| ext.to_str());    let path_string = path.to_string_lossy();    let capacity = 3        + path_extension.map_or(0, |extension| extension.len() + 1)        + path_string.len()        + 1        + content.len()        + 5;    let mut buffer = String::with_capacity(capacity);     buffer.push_str("```");     if let Some(extension) = path_extension {        buffer.push_str(extension);        buffer.push(' ');    }    buffer.push_str(&path_string);     buffer.push('\n');    for chunk in content.chunks() {        buffer.push_str(&chunk);    }     if !buffer.ends_with('\n') {        buffer.push('\n');    }     buffer.push_str("```\n");

Lines 1–500 of 657. Load more