assistant2: List directories in directory context picker (#22300)

This PR updates the directory context picker in Assistant2 to show the
available directories.

Release Notes:

- N/A

crates/assistant2/src/context_picker/directory_context_picker.rs

@@ -1,6 +1,8 @@
-// TODO: Remove this once we've implemented the functionality.
+// TODO: Remove this when we finish the implementation.
 #![allow(unused)]
 
+use std::path::Path;
+use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 
 use fuzzy::PathMatch;
@@ -11,6 +13,7 @@ use ui::{prelude::*, ListItem};
 use util::ResultExt as _;
 use workspace::Workspace;
 
+use crate::context::ContextKind;
 use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
 
@@ -75,6 +78,65 @@ impl DirectoryContextPickerDelegate {
             selected_index: 0,
         }
     }
+
+    fn search(
+        &mut self,
+        query: String,
+        cancellation_flag: Arc<AtomicBool>,
+        workspace: &View<Workspace>,
+        cx: &mut ViewContext<Picker<Self>>,
+    ) -> Task<Vec<PathMatch>> {
+        if query.is_empty() {
+            let workspace = workspace.read(cx);
+            let project = workspace.project().read(cx);
+            let directory_matches = project.worktrees(cx).flat_map(|worktree| {
+                let worktree = worktree.read(cx);
+                let path_prefix: Arc<str> = worktree.root_name().into();
+                worktree.directories(false, 0).map(move |entry| PathMatch {
+                    score: 0.,
+                    positions: Vec::new(),
+                    worktree_id: worktree.id().to_usize(),
+                    path: entry.path.clone(),
+                    path_prefix: path_prefix.clone(),
+                    distance_to_relative_ancestor: 0,
+                    is_dir: true,
+                })
+            });
+
+            Task::ready(directory_matches.collect())
+        } else {
+            let worktrees = workspace.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
+            let candidate_sets = worktrees
+                .into_iter()
+                .map(|worktree| {
+                    let worktree = worktree.read(cx);
+
+                    PathMatchCandidateSet {
+                        snapshot: worktree.snapshot(),
+                        include_ignored: worktree
+                            .root_entry()
+                            .map_or(false, |entry| entry.is_ignored),
+                        include_root_name: true,
+                        candidates: project::Candidates::Directories,
+                    }
+                })
+                .collect::<Vec<_>>();
+
+            let executor = cx.background_executor().clone();
+            cx.foreground_executor().spawn(async move {
+                fuzzy::match_path_sets(
+                    candidate_sets.as_slice(),
+                    query.as_str(),
+                    None,
+                    false,
+                    100,
+                    &cancellation_flag,
+                    executor,
+                )
+                .await
+            })
+        }
+    }
 }
 
 impl PickerDelegate for DirectoryContextPickerDelegate {
@@ -88,7 +150,7 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
         self.selected_index
     }
 
-    fn set_selected_index(&mut self, ix: usize, cx: &mut ViewContext<Picker<Self>>) {
+    fn set_selected_index(&mut self, ix: usize, _cx: &mut ViewContext<Picker<Self>>) {
         self.selected_index = ix;
     }
 
@@ -96,17 +158,65 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
         "Search folders…".into()
     }
 
-    fn update_matches(&mut self, _query: String, _cx: &mut ViewContext<Picker<Self>>) -> Task<()> {
-        // TODO: Implement this once we fix the issues with the file context picker.
-        Task::ready(())
+    fn update_matches(&mut self, query: String, cx: &mut ViewContext<Picker<Self>>) -> Task<()> {
+        let Some(workspace) = self.workspace.upgrade() else {
+            return Task::ready(());
+        };
+
+        let search_task = self.search(query, Arc::<AtomicBool>::default(), &workspace, cx);
+
+        cx.spawn(|this, mut cx| async move {
+            let mut paths = search_task.await;
+            let empty_path = Path::new("");
+            paths.retain(|path_match| path_match.path.as_ref() != empty_path);
+
+            this.update(&mut cx, |this, _cx| {
+                this.delegate.matches = paths;
+            })
+            .log_err();
+        })
     }
 
     fn confirm(&mut self, _secondary: bool, cx: &mut ViewContext<Picker<Self>>) {
-        // TODO: Implement this once we fix the issues with the file context picker.
-        match self.confirm_behavior {
-            ConfirmBehavior::KeepOpen => {}
-            ConfirmBehavior::Close => self.dismissed(cx),
-        }
+        let mat = &self.matches[self.selected_index];
+
+        let workspace = self.workspace.clone();
+        let Some(project) = workspace
+            .upgrade()
+            .map(|workspace| workspace.read(cx).project().clone())
+        else {
+            return;
+        };
+        let path = mat.path.clone();
+        let worktree_id = WorktreeId::from_usize(mat.worktree_id);
+        let confirm_behavior = self.confirm_behavior;
+        cx.spawn(|this, mut cx| async move {
+            this.update(&mut cx, |this, cx| {
+                let mut text = String::new();
+
+                // TODO: Add the files from the selected directory.
+
+                this.delegate
+                    .context_store
+                    .update(cx, |context_store, cx| {
+                        context_store.insert_context(
+                            ContextKind::Directory,
+                            path.to_string_lossy().to_string(),
+                            text,
+                        );
+                    })?;
+
+                match confirm_behavior {
+                    ConfirmBehavior::KeepOpen => {}
+                    ConfirmBehavior::Close => this.delegate.dismissed(cx),
+                }
+
+                anyhow::Ok(())
+            })??;
+
+            anyhow::Ok(())
+        })
+        .detach_and_log_err(cx)
     }
 
     fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
@@ -120,10 +230,18 @@ impl PickerDelegate for DirectoryContextPickerDelegate {
 
     fn render_match(
         &self,
-        _ix: usize,
-        _selected: bool,
+        ix: usize,
+        selected: bool,
         _cx: &mut ViewContext<Picker<Self>>,
     ) -> Option<Self::ListItem> {
-        None
+        let path_match = &self.matches[ix];
+        let directory_name = path_match.path.to_string_lossy().to_string();
+
+        Some(
+            ListItem::new(ix)
+                .inset(true)
+                .toggle_state(selected)
+                .child(h_flex().gap_2().child(Label::new(directory_name))),
+        )
     }
 }