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_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
+}