diff --git a/crates/assistant2/src/context_picker/directory_context_picker.rs b/crates/assistant2/src/context_picker/directory_context_picker.rs index f67888bca9..a8c76a1554 100644 --- a/crates/assistant2/src/context_picker/directory_context_picker.rs +++ b/crates/assistant2/src/context_picker/directory_context_picker.rs @@ -2,17 +2,16 @@ use std::path::Path; use std::sync::atomic::AtomicBool; use std::sync::Arc; -use anyhow::anyhow; use fuzzy::PathMatch; use gpui::{AppContext, DismissEvent, FocusHandle, FocusableView, Task, View, WeakModel, WeakView}; use picker::{Picker, PickerDelegate}; -use project::{PathMatchCandidateSet, ProjectPath, Worktree, WorktreeId}; +use project::{PathMatchCandidateSet, ProjectPath, WorktreeId}; use ui::{prelude::*, ListItem}; use util::ResultExt as _; use workspace::Workspace; use crate::context_picker::{ConfirmBehavior, ContextPicker}; -use crate::context_store::{push_fenced_codeblock, ContextStore}; +use crate::context_store::ContextStore; pub struct DirectoryContextPicker { picker: View>, @@ -179,107 +178,45 @@ impl PickerDelegate for DirectoryContextPickerDelegate { return; }; - let workspace = self.workspace.clone(); - let Some(project) = workspace - .upgrade() - .map(|workspace| workspace.read(cx).project().clone()) + let project_path = ProjectPath { + worktree_id: WorktreeId::from_usize(mat.worktree_id), + path: mat.path.clone(), + }; + + let Some(task) = self + .context_store + .update(cx, |context_store, cx| { + context_store.add_directory(project_path, cx) + }) + .ok() else { return; }; - let path = mat.path.clone(); - let already_included = self - .context_store - .update(cx, |context_store, _cx| { - if let Some(context_id) = context_store.included_directory(&path) { - context_store.remove_context(&context_id); - true - } else { - false - } - }) - .unwrap_or(true); - if already_included { - return; - } - - let worktree_id = WorktreeId::from_usize(mat.worktree_id); + let workspace = self.workspace.clone(); let confirm_behavior = self.confirm_behavior; 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, &path) - })?; - - let open_buffer_tasks = project.update(&mut cx, |project, cx| { - files - .into_iter() - .map(|file_path| { - project.open_buffer( - ProjectPath { - worktree_id, - path: file_path.clone(), - }, - cx, - ) - }) - .collect::>() - })?; - - let buffers = futures::future::join_all(open_buffer_tasks).await; - - this.update(&mut cx, |this, cx| { - let mut text = String::new(); - - let mut ok_count = 0; - - for buffer in buffers.into_iter().flatten() { - let buffer = buffer.read(cx); - let path = buffer.file().map_or(&path, |file| file.path()); - push_fenced_codeblock(&path, buffer.text(), &mut text); - ok_count += 1; + match task.await { + Ok(()) => { + this.update(&mut cx, |this, cx| match confirm_behavior { + ConfirmBehavior::KeepOpen => {} + ConfirmBehavior::Close => this.delegate.dismissed(cx), + })?; } - - if ok_count == 0 { + Err(err) => { let Some(workspace) = workspace.upgrade() else { return anyhow::Ok(()); }; - workspace.update(cx, |workspace, cx| { - workspace.show_error( - &anyhow::anyhow!( - "Could not read any text files from {}", - path.display() - ), - cx, - ); - }); - - return anyhow::Ok(()); - } - - this.delegate - .context_store - .update(cx, |context_store, _cx| { - context_store.insert_directory(&path, text); + workspace.update(&mut cx, |workspace, cx| { + workspace.show_error(&err, cx); })?; - - match confirm_behavior { - ConfirmBehavior::KeepOpen => {} - ConfirmBehavior::Close => this.delegate.dismissed(cx), } - - anyhow::Ok(()) - })??; + } anyhow::Ok(()) }) - .detach_and_log_err(cx) + .detach_and_log_err(cx); } fn dismissed(&mut self, cx: &mut ViewContext>) { @@ -327,17 +264,3 @@ impl PickerDelegate for DirectoryContextPickerDelegate { ) } } - -fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec> { - let mut files = Vec::new(); - - for entry in worktree.child_entries(path) { - if entry.is_dir() { - files.extend(collect_files_in_path(worktree, &entry.path)); - } else if entry.is_file() { - files.push(entry.path.clone()); - } - } - - files -} diff --git a/crates/assistant2/src/context_store.rs b/crates/assistant2/src/context_store.rs index cdb50ce5b5..03ae83102c 100644 --- a/crates/assistant2/src/context_store.rs +++ b/crates/assistant2/src/context_store.rs @@ -1,11 +1,12 @@ use std::fmt::Write as _; use std::path::{Path, PathBuf}; +use std::sync::Arc; -use anyhow::{anyhow, Result}; +use anyhow::{anyhow, bail, Result}; use collections::{HashMap, HashSet}; use gpui::{ModelContext, SharedString, Task, WeakView}; use language::Buffer; -use project::ProjectPath; +use project::{ProjectPath, Worktree}; use workspace::Workspace; use crate::thread::Thread; @@ -122,6 +123,86 @@ impl ContextStore { }); } + pub fn add_directory( + &mut self, + project_path: ProjectPath, + cx: &mut ModelContext, + ) -> Task> { + 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.included_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_buffer_tasks = project.update(&mut cx, |project, cx| { + files + .into_iter() + .map(|file_path| { + project.open_buffer( + ProjectPath { + worktree_id, + path: file_path.clone(), + }, + cx, + ) + }) + .collect::>() + })?; + + let buffers = futures::future::join_all(open_buffer_tasks).await; + + this.update(&mut cx, |this, cx| { + let mut text = String::new(); + let mut added_files = 0; + + for buffer in buffers.into_iter().flatten() { + let buffer = buffer.read(cx); + let path = buffer.file().map_or(&project_path.path, |file| file.path()); + push_fenced_codeblock(&path, buffer.text(), &mut text); + added_files += 1; + } + + if added_files == 0 { + bail!( + "could not read any text files from {}", + &project_path.path.display() + ); + } + + this.insert_directory(&project_path.path, text); + + anyhow::Ok(()) + })??; + + anyhow::Ok(()) + }) + } + pub fn insert_directory(&mut self, path: &Path, text: impl Into) { let id = self.next_context_id.post_inc(); self.directories.insert(path.to_path_buf(), id); @@ -268,3 +349,17 @@ pub(crate) fn push_fenced_codeblock(path: &Path, content: String, buffer: &mut S buffer.push_str("```\n"); } + +fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec> { + let mut files = Vec::new(); + + for entry in worktree.child_entries(path) { + if entry.is_dir() { + files.extend(collect_files_in_path(worktree, &entry.path)); + } else if entry.is_file() { + files.push(entry.path.clone()); + } + } + + files +}