use std::ops::Range;use std::path::Path;use std::sync::Arc; use anyhow::{Context as _, Result, anyhow};use collections::{BTreeMap, HashMap, HashSet};use futures::future::join_all;use futures::{self, Future, FutureExt, future};use gpui::{App, AppContext as _, Context, Entity, SharedString, Task, WeakEntity};use language::{Buffer, File};use project::{Project, ProjectItem, ProjectPath, Worktree};use rope::Rope;use text::{Anchor, BufferId, OffsetRangeExt};use util::{ResultExt as _, maybe}; use crate::ThreadStore;use crate::context::{    AssistantContext, ContextBuffer, ContextId, ContextSymbol, ContextSymbolId, DirectoryContext,    FetchedUrlContext, FileContext, SymbolContext, ThreadContext,};use crate::context_strip::SuggestedContext;use crate::thread::{Thread, ThreadId}; pub struct ContextStore {    project: WeakEntity<Project>,    context: Vec<AssistantContext>,    thread_store: Option<WeakEntity<ThreadStore>>,    // 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<ProjectPath, ContextId>,    symbols: HashMap<ContextSymbolId, ContextId>,    symbol_buffers: HashMap<ContextSymbolId, Entity<Buffer>>,    symbols_by_path: HashMap<ProjectPath, Vec<ContextSymbolId>>,    threads: HashMap<ThreadId, ContextId>,    thread_summary_tasks: Vec<Task<()>>,    fetched_urls: HashMap<String, ContextId>,} impl ContextStore {    pub fn new(        project: WeakEntity<Project>,        thread_store: Option<WeakEntity<ThreadStore>>,    ) -> Self {        Self {            project,            thread_store,            context: Vec::new(),            next_context_id: ContextId(0),            files: BTreeMap::default(),            directories: HashMap::default(),            symbols: HashMap::default(),            symbol_buffers: HashMap::default(),            symbols_by_path: HashMap::default(),            threads: HashMap::default(),            thread_summary_tasks: Vec::new(),            fetched_urls: HashMap::default(),        }    }     pub fn context(&self) -> &Vec<AssistantContext> {        &self.context    }     pub fn context_for_id(&self, id: ContextId) -> Option<&AssistantContext> {        self.context().iter().find(|context| context.id() == id)    }     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,        remove_if_exists: bool,        cx: &mut Context<Self>,    ) -> Task<Result<()>> {        let Some(project) = self.project.upgrade() else {            return Task::ready(Err(anyhow!("failed to read project")));        };         cx.spawn(async move |this, cx| {            let open_buffer_task = project.update(cx, |project, cx| {                project.open_buffer(project_path.clone(), cx)            })?;             let buffer = open_buffer_task.await?;            let buffer_id = this.update(cx, |_, cx| buffer.read(cx).remote_id())?;             let already_included = this.update(cx, |this, cx| {                match this.will_include_buffer(buffer_id, &project_path) {                    Some(FileInclusion::Direct(context_id)) => {                        if remove_if_exists {                            this.remove_context(context_id, cx);                        }                        true                    }                    Some(FileInclusion::InDirectory(_)) => true,                    None => false,                }            })?;             if already_included {                return anyhow::Ok(());            }             let (buffer_info, text_task) =                this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, None, cx))??;             let text = text_task.await;             this.update(cx, |this, cx| {                this.insert_file(make_context_buffer(buffer_info, text), cx);            })?;             anyhow::Ok(())        })    }     pub fn add_file_from_buffer(        &mut self,        buffer: Entity<Buffer>,        cx: &mut Context<Self>,    ) -> Task<Result<()>> {        cx.spawn(async move |this, cx| {            let (buffer_info, text_task) =                this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, None, cx))??;             let text = text_task.await;             this.update(cx, |this, cx| {                this.insert_file(make_context_buffer(buffer_info, text), cx)            })?;             anyhow::Ok(())        })    }     fn insert_file(&mut self, context_buffer: ContextBuffer, cx: &mut Context<Self>) {        let id = self.next_context_id.post_inc();        self.files.insert(context_buffer.id, id);        self.context            .push(AssistantContext::File(FileContext { id, context_buffer }));        cx.notify();    }     pub fn add_directory(        &mut self,        project_path: ProjectPath,        remove_if_exists: bool,        cx: &mut Context<Self>,    ) -> Task<Result<()>> {        let Some(project) = self.project.upgrade() else {            return Task::ready(Err(anyhow!("failed to read project")));        };         let already_included = match self.includes_directory(&project_path) {            Some(FileInclusion::Direct(context_id)) => {                if remove_if_exists {                    self.remove_context(context_id, cx);                }                true            }            Some(FileInclusion::InDirectory(_)) => true,            None => false,        };        if already_included {            return Task::ready(Ok(()));        }         let worktree_id = project_path.worktree_id;        cx.spawn(async move |this, cx| {            let worktree = project.update(cx, |project, cx| {                project                    .worktree_for_id(worktree_id, cx)                    .ok_or_else(|| anyhow!("no worktree found for {worktree_id:?}"))            })??;             let files = worktree.update(cx, |worktree, _cx| {                collect_files_in_path(worktree, &project_path.path)            })?;             let open_buffers_task = project.update(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(cx, |_, cx| {                // Skip all binary files and other non-UTF8 files                for buffer in buffers.into_iter().flatten() {                    if let Some((buffer_info, text_task)) =                        collect_buffer_info_and_text(buffer, None, cx).log_err()                    {                        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() {                let full_path = cx.update(|cx| worktree.read(cx).full_path(&project_path.path))?;                return Err(anyhow!("No text files found in {}", &full_path.display()));            }             this.update(cx, |this, cx| {                this.insert_directory(worktree, project_path, context_buffers, cx);            })?;             anyhow::Ok(())        })    }     fn insert_directory(        &mut self,        worktree: Entity<Worktree>,        project_path: ProjectPath,        context_buffers: Vec<ContextBuffer>,        cx: &mut Context<Self>,    ) {        let id = self.next_context_id.post_inc();        let path = project_path.path.clone();        self.directories.insert(project_path, id);         self.context            .push(AssistantContext::Directory(DirectoryContext {                id,                worktree,                path,                context_buffers,            }));        cx.notify();    }     pub fn add_symbol(        &mut self,        buffer: Entity<Buffer>,        symbol_name: SharedString,        symbol_range: Range<Anchor>,        symbol_enclosing_range: Range<Anchor>,        remove_if_exists: bool,        cx: &mut Context<Self>,    ) -> Task<Result<bool>> {        let buffer_ref = buffer.read(cx);        let Some(project_path) = buffer_ref.project_path(cx) else {            return Task::ready(Err(anyhow!("buffer has no path")));        };         if let Some(symbols_for_path) = self.symbols_by_path.get(&project_path) {            let mut matching_symbol_id = None;            for symbol in symbols_for_path {                if &symbol.name == &symbol_name {                    let snapshot = buffer_ref.snapshot();                    if symbol.range.to_offset(&snapshot) == symbol_range.to_offset(&snapshot) {                        matching_symbol_id = self.symbols.get(symbol).cloned();                        break;                    }                }            }             if let Some(id) = matching_symbol_id {                if remove_if_exists {                    self.remove_context(id, cx);                }                return Task::ready(Ok(false));            }        }         let (buffer_info, collect_content_task) =            match collect_buffer_info_and_text(buffer, Some(symbol_enclosing_range.clone()), cx) {                Ok((buffer_info, collect_context_task)) => (buffer_info, collect_context_task),                Err(err) => return Task::ready(Err(err)),            };         cx.spawn(async move |this, cx| {            let content = collect_content_task.await;             this.update(cx, |this, cx| {                this.insert_symbol(                    make_context_symbol(                        buffer_info,                        project_path,                        symbol_name,                        symbol_range,                        symbol_enclosing_range,                        content,                    ),                    cx,                )            })?;            anyhow::Ok(true)        })    }     fn insert_symbol(&mut self, context_symbol: ContextSymbol, cx: &mut Context<Self>) {        let id = self.next_context_id.post_inc();        self.symbols.insert(context_symbol.id.clone(), id);        self.symbols_by_path            .entry(context_symbol.id.path.clone())            .or_insert_with(Vec::new)            .push(context_symbol.id.clone());        self.symbol_buffers            .insert(context_symbol.id.clone(), context_symbol.buffer.clone());        self.context.push(AssistantContext::Symbol(SymbolContext {            id,            context_symbol,        }));        cx.notify();    }     pub fn add_thread(        &mut self,        thread: Entity<Thread>,        remove_if_exists: bool,        cx: &mut Context<Self>,    ) {        if let Some(context_id) = self.includes_thread(&thread.read(cx).id()) {            if remove_if_exists {                self.remove_context(context_id, cx);            }        } else {            self.insert_thread(thread, cx);        }    }     pub fn wait_for_summaries(&mut self, cx: &App) -> Task<()> {        let tasks = std::mem::take(&mut self.thread_summary_tasks);         cx.spawn(async move |_cx| {            join_all(tasks).await;        })    }     fn insert_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {        if let Some(summary_task) =            thread.update(cx, |thread, cx| thread.generate_detailed_summary(cx))        {            let thread = thread.clone();            let thread_store = self.thread_store.clone();             self.thread_summary_tasks.push(cx.spawn(async move |_, cx| {                summary_task.await;                 if let Some(thread_store) = thread_store {                    // Save thread so its summary can be reused later                    let save_task = thread_store                        .update(cx, |thread_store, cx| thread_store.save_thread(&thread, cx));                     if let Some(save_task) = save_task.ok() {                        save_task.await.log_err();                    }                }            }));        }         let id = self.next_context_id.post_inc();         let text = thread.read(cx).latest_detailed_summary_or_text();         self.threads.insert(thread.read(cx).id().clone(), id);        self.context            .push(AssistantContext::Thread(ThreadContext { id, thread, text }));        cx.notify();    }     pub fn add_fetched_url(        &mut self,        url: String,        text: impl Into<SharedString>,        cx: &mut Context<ContextStore>,    ) {        if self.includes_url(&url).is_none() {            self.insert_fetched_url(url, text, cx);        }    }     fn insert_fetched_url(        &mut self,        url: String,        text: impl Into<SharedString>,        cx: &mut Context<ContextStore>,    ) {        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(),            }));        cx.notify();    }     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, cx: &mut Context<Self>) {        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::Symbol(symbol) => {                if let Some(symbols_in_path) =                    self.symbols_by_path.get_mut(&symbol.context_symbol.id.path)                {                    symbols_in_path.retain(|s| {                        self.symbols                            .get(s)                            .map_or(false, |context_id| *context_id != id)                    });                }                self.symbol_buffers.remove(&symbol.context_symbol.id);                self.symbols.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);            }        }         cx.notify();    }     /// 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,        project_path: &ProjectPath,    ) -> 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(project_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,        project_path: &ProjectPath,        cx: &App,    ) -> Option<FileInclusion> {

Lines 1–500 of 947. Load more