assistant2: Push logic for adding directory context down into the ContextStore (#22852)

This PR takes the logic for adding file context out of the
`DirectoryContextPicker` and pushes it down into the `ContextStore`.

Release Notes:

- N/A

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<Picker<DirectoryContextPickerDelegate>>,
@@ -179,107 +178,45 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
             return;
         };
 
-        let workspace = self.workspace.clone();
+        let project_path = ProjectPath {
-        let Some(project) = workspace
+            worktree_id: WorktreeId::from_usize(mat.worktree_id),
-            .upgrade()
+            path: mat.path.clone(),
-            .map(|workspace| workspace.read(cx).project().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
+        let workspace = self.workspace.clone();
-            .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 confirm_behavior = self.confirm_behavior;
         cx.spawn(|this, mut cx| async move {
-            let worktree = project.update(&mut cx, |project, cx| {
+            match task.await {
-                project
+                Ok(()) => {
-                    .worktree_for_id(worktree_id, cx)
+                    this.update(&mut cx, |this, cx| match confirm_behavior {
-                    .ok_or_else(|| anyhow!("no worktree found for {worktree_id:?}"))
+                        ConfirmBehavior::KeepOpen => {}
-            })??;
+                        ConfirmBehavior::Close => this.delegate.dismissed(cx),
-
+                    })?;
-            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::<Vec<_>>()
-            })?;
-
-            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;
                 }
-
+                Err(err) => {
-                if ok_count == 0 {
                     let Some(workspace) = workspace.upgrade() else {
                         return anyhow::Ok(());
                     };
 
-                    workspace.update(cx, |workspace, cx| {
+                    workspace.update(&mut cx, |workspace, cx| {
-                        workspace.show_error(
+                        workspace.show_error(&err, cx);
-                            &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);
                     })?;
-
-                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<Picker<Self>>) {
@@ -327,17 +264,3 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
         )
     }
 }
-
-fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
-    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
-}

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<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.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::<Vec<_>>()
+            })?;
+
+            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<SharedString>) {
         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<Arc<Path>> {
+    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
+}