assistant2: Combine file & directory picker (#26975)

In the process of adding `@mentions` we realized that we do not want to
make a distinction between Files & Directories in the UI, therefore this
PR combines the File & Directory pickers into a unified version



https://github.com/user-attachments/assets/f3bf189c-8b69-4f5f-90ce-0b83b12dbca3

(Ignore the `@mentions`, they are broken also on main)

Release Notes:

- N/A

crates/assistant2/src/context.rs

@@ -43,15 +43,6 @@ pub enum ContextKind {
 }
 
 impl ContextKind {
-    pub fn label(&self) -> &'static str {
-        match self {
-            ContextKind::File => "File",
-            ContextKind::Directory => "Folder",
-            ContextKind::FetchedUrl => "Fetch",
-            ContextKind::Thread => "Thread",
-        }
-    }
-
     pub fn icon(&self) -> IconName {
         match self {
             ContextKind::File => IconName::File,

crates/assistant2/src/context_picker.rs

@@ -1,4 +1,3 @@
-mod directory_context_picker;
 mod fetch_context_picker;
 mod file_context_picker;
 mod thread_context_picker;
@@ -15,8 +14,6 @@ use thread_context_picker::{render_thread_context_entry, ThreadContextEntry};
 use ui::{prelude::*, ContextMenu, ContextMenuEntry, ContextMenuItem};
 use workspace::{notifications::NotifyResultExt, Workspace};
 
-use crate::context::ContextKind;
-use crate::context_picker::directory_context_picker::DirectoryContextPicker;
 use crate::context_picker::fetch_context_picker::FetchContextPicker;
 use crate::context_picker::file_context_picker::FileContextPicker;
 use crate::context_picker::thread_context_picker::ThreadContextPicker;
@@ -30,17 +27,41 @@ pub enum ConfirmBehavior {
     Close,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum ContextPickerMode {
+    File,
+    Fetch,
+    Thread,
+}
+
+impl ContextPickerMode {
+    pub fn label(&self) -> &'static str {
+        match self {
+            Self::File => "File/Directory",
+            Self::Fetch => "Fetch",
+            Self::Thread => "Thread",
+        }
+    }
+
+    pub fn icon(&self) -> IconName {
+        match self {
+            Self::File => IconName::File,
+            Self::Fetch => IconName::Globe,
+            Self::Thread => IconName::MessageCircle,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+enum ContextPickerState {
     Default(Entity<ContextMenu>),
     File(Entity<FileContextPicker>),
-    Directory(Entity<DirectoryContextPicker>),
     Fetch(Entity<FetchContextPicker>),
     Thread(Entity<ThreadContextPicker>),
 }
 
 pub(super) struct ContextPicker {
-    mode: ContextPickerMode,
+    mode: ContextPickerState,
     workspace: WeakEntity<Workspace>,
     editor: WeakEntity<Editor>,
     context_store: WeakEntity<ContextStore>,
@@ -59,7 +80,7 @@ impl ContextPicker {
         cx: &mut Context<Self>,
     ) -> Self {
         ContextPicker {
-            mode: ContextPickerMode::Default(ContextMenu::build(
+            mode: ContextPickerState::Default(ContextMenu::build(
                 window,
                 cx,
                 |menu, _window, _cx| menu,
@@ -73,7 +94,7 @@ impl ContextPicker {
     }
 
     pub fn init(&mut self, window: &mut Window, cx: &mut Context<Self>) {
-        self.mode = ContextPickerMode::Default(self.build_menu(window, cx));
+        self.mode = ContextPickerState::Default(self.build_menu(window, cx));
         cx.notify();
     }
 
@@ -88,13 +109,9 @@ impl ContextPicker {
                 .enumerate()
                 .map(|(ix, entry)| self.recent_menu_item(context_picker.clone(), ix, entry));
 
-            let mut context_kinds = vec![
+            let mut modes = vec![ContextPickerMode::File, ContextPickerMode::Fetch];
-                ContextKind::File,
-                ContextKind::Directory,
-                ContextKind::FetchedUrl,
-            ];
             if self.allow_threads() {
-                context_kinds.push(ContextKind::Thread);
+                modes.push(ContextPickerMode::Thread);
             }
 
             let menu = menu
@@ -112,15 +129,15 @@ impl ContextPicker {
                 })
                 .extend(recent_entries)
                 .when(has_recent, |menu| menu.separator())
-                .extend(context_kinds.into_iter().map(|kind| {
+                .extend(modes.into_iter().map(|mode| {
                     let context_picker = context_picker.clone();
 
-                    ContextMenuEntry::new(kind.label())
+                    ContextMenuEntry::new(mode.label())
-                        .icon(kind.icon())
+                        .icon(mode.icon())
                         .icon_size(IconSize::XSmall)
                         .icon_color(Color::Muted)
                         .handler(move |window, cx| {
-                            context_picker.update(cx, |this, cx| this.select_kind(kind, window, cx))
+                            context_picker.update(cx, |this, cx| this.select_mode(mode, window, cx))
                         })
                 }));
 
@@ -143,12 +160,17 @@ impl ContextPicker {
         self.thread_store.is_some()
     }
 
-    fn select_kind(&mut self, kind: ContextKind, window: &mut Window, cx: &mut Context<Self>) {
+    fn select_mode(
+        &mut self,
+        mode: ContextPickerMode,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
         let context_picker = cx.entity().downgrade();
 
-        match kind {
+        match mode {
-            ContextKind::File => {
+            ContextPickerMode::File => {
-                self.mode = ContextPickerMode::File(cx.new(|cx| {
+                self.mode = ContextPickerState::File(cx.new(|cx| {
                     FileContextPicker::new(
                         context_picker.clone(),
                         self.workspace.clone(),
@@ -160,20 +182,8 @@ impl ContextPicker {
                     )
                 }));
             }
-            ContextKind::Directory => {
+            ContextPickerMode::Fetch => {
-                self.mode = ContextPickerMode::Directory(cx.new(|cx| {
+                self.mode = ContextPickerState::Fetch(cx.new(|cx| {
-                    DirectoryContextPicker::new(
-                        context_picker.clone(),
-                        self.workspace.clone(),
-                        self.context_store.clone(),
-                        self.confirm_behavior,
-                        window,
-                        cx,
-                    )
-                }));
-            }
-            ContextKind::FetchedUrl => {
-                self.mode = ContextPickerMode::Fetch(cx.new(|cx| {
                     FetchContextPicker::new(
                         context_picker.clone(),
                         self.workspace.clone(),
@@ -184,9 +194,9 @@ impl ContextPicker {
                     )
                 }));
             }
-            ContextKind::Thread => {
+            ContextPickerMode::Thread => {
                 if let Some(thread_store) = self.thread_store.as_ref() {
-                    self.mode = ContextPickerMode::Thread(cx.new(|cx| {
+                    self.mode = ContextPickerState::Thread(cx.new(|cx| {
                         ThreadContextPicker::new(
                             thread_store.clone(),
                             context_picker.clone(),
@@ -224,6 +234,7 @@ impl ContextPicker {
                             ElementId::NamedInteger("ctx-recent".into(), ix),
                             &path,
                             &path_prefix,
+                            false,
                             context_store.clone(),
                             cx,
                         )
@@ -392,11 +403,10 @@ impl EventEmitter<DismissEvent> for ContextPicker {}
 impl Focusable for ContextPicker {
     fn focus_handle(&self, cx: &App) -> FocusHandle {
         match &self.mode {
-            ContextPickerMode::Default(menu) => menu.focus_handle(cx),
+            ContextPickerState::Default(menu) => menu.focus_handle(cx),
-            ContextPickerMode::File(file_picker) => file_picker.focus_handle(cx),
+            ContextPickerState::File(file_picker) => file_picker.focus_handle(cx),
-            ContextPickerMode::Directory(directory_picker) => directory_picker.focus_handle(cx),
+            ContextPickerState::Fetch(fetch_picker) => fetch_picker.focus_handle(cx),
-            ContextPickerMode::Fetch(fetch_picker) => fetch_picker.focus_handle(cx),
+            ContextPickerState::Thread(thread_picker) => thread_picker.focus_handle(cx),
-            ContextPickerMode::Thread(thread_picker) => thread_picker.focus_handle(cx),
         }
     }
 }
@@ -407,13 +417,10 @@ impl Render for ContextPicker {
             .w(px(400.))
             .min_w(px(400.))
             .map(|parent| match &self.mode {
-                ContextPickerMode::Default(menu) => parent.child(menu.clone()),
+                ContextPickerState::Default(menu) => parent.child(menu.clone()),
-                ContextPickerMode::File(file_picker) => parent.child(file_picker.clone()),
+                ContextPickerState::File(file_picker) => parent.child(file_picker.clone()),
-                ContextPickerMode::Directory(directory_picker) => {
+                ContextPickerState::Fetch(fetch_picker) => parent.child(fetch_picker.clone()),
-                    parent.child(directory_picker.clone())
+                ContextPickerState::Thread(thread_picker) => parent.child(thread_picker.clone()),
-                }
-                ContextPickerMode::Fetch(fetch_picker) => parent.child(fetch_picker.clone()),
-                ContextPickerMode::Thread(thread_picker) => parent.child(thread_picker.clone()),
             })
     }
 }

crates/assistant2/src/context_picker/directory_context_picker.rs

@@ -1,269 +0,0 @@
-use std::path::Path;
-use std::sync::atomic::AtomicBool;
-use std::sync::Arc;
-
-use fuzzy::PathMatch;
-use gpui::{App, DismissEvent, Entity, FocusHandle, Focusable, Task, WeakEntity};
-use picker::{Picker, PickerDelegate};
-use project::{PathMatchCandidateSet, ProjectPath, WorktreeId};
-use ui::{prelude::*, ListItem};
-use util::ResultExt as _;
-use workspace::{notifications::NotifyResultExt, Workspace};
-
-use crate::context_picker::{ConfirmBehavior, ContextPicker};
-use crate::context_store::ContextStore;
-
-pub struct DirectoryContextPicker {
-    picker: Entity<Picker<DirectoryContextPickerDelegate>>,
-}
-
-impl DirectoryContextPicker {
-    pub fn new(
-        context_picker: WeakEntity<ContextPicker>,
-        workspace: WeakEntity<Workspace>,
-        context_store: WeakEntity<ContextStore>,
-        confirm_behavior: ConfirmBehavior,
-        window: &mut Window,
-        cx: &mut Context<Self>,
-    ) -> Self {
-        let delegate = DirectoryContextPickerDelegate::new(
-            context_picker,
-            workspace,
-            context_store,
-            confirm_behavior,
-        );
-        let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));
-
-        Self { picker }
-    }
-}
-
-impl Focusable for DirectoryContextPicker {
-    fn focus_handle(&self, cx: &App) -> FocusHandle {
-        self.picker.focus_handle(cx)
-    }
-}
-
-impl Render for DirectoryContextPicker {
-    fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
-        self.picker.clone()
-    }
-}
-
-pub struct DirectoryContextPickerDelegate {
-    context_picker: WeakEntity<ContextPicker>,
-    workspace: WeakEntity<Workspace>,
-    context_store: WeakEntity<ContextStore>,
-    confirm_behavior: ConfirmBehavior,
-    matches: Vec<PathMatch>,
-    selected_index: usize,
-}
-
-impl DirectoryContextPickerDelegate {
-    pub fn new(
-        context_picker: WeakEntity<ContextPicker>,
-        workspace: WeakEntity<Workspace>,
-        context_store: WeakEntity<ContextStore>,
-        confirm_behavior: ConfirmBehavior,
-    ) -> Self {
-        Self {
-            context_picker,
-            workspace,
-            context_store,
-            confirm_behavior,
-            matches: Vec::new(),
-            selected_index: 0,
-        }
-    }
-
-    fn search(
-        &mut self,
-        query: String,
-        cancellation_flag: Arc<AtomicBool>,
-        workspace: &Entity<Workspace>,
-        cx: &mut Context<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 {
-    type ListItem = ListItem;
-
-    fn match_count(&self) -> usize {
-        self.matches.len()
-    }
-
-    fn selected_index(&self) -> usize {
-        self.selected_index
-    }
-
-    fn set_selected_index(
-        &mut self,
-        ix: usize,
-        _window: &mut Window,
-        _cx: &mut Context<Picker<Self>>,
-    ) {
-        self.selected_index = ix;
-    }
-
-    fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
-        "Search folders…".into()
-    }
-
-    fn update_matches(
-        &mut self,
-        query: String,
-        _window: &mut Window,
-        cx: &mut Context<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, window: &mut Window, cx: &mut Context<Picker<Self>>) {
-        let Some(mat) = self.matches.get(self.selected_index) else {
-            return;
-        };
-
-        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 confirm_behavior = self.confirm_behavior;
-        cx.spawn_in(window, |this, mut cx| async move {
-            match task.await.notify_async_err(&mut cx) {
-                None => anyhow::Ok(()),
-                Some(()) => this.update_in(&mut cx, |this, window, cx| match confirm_behavior {
-                    ConfirmBehavior::KeepOpen => {}
-                    ConfirmBehavior::Close => this.delegate.dismissed(window, cx),
-                }),
-            }
-        })
-        .detach_and_log_err(cx);
-    }
-
-    fn dismissed(&mut self, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
-        self.context_picker
-            .update(cx, |_, cx| {
-                cx.emit(DismissEvent);
-            })
-            .ok();
-    }
-
-    fn render_match(
-        &self,
-        ix: usize,
-        selected: bool,
-        _window: &mut Window,
-        cx: &mut Context<Picker<Self>>,
-    ) -> Option<Self::ListItem> {
-        let path_match = &self.matches[ix];
-        let directory_name = path_match.path.to_string_lossy().to_string();
-
-        let added = self.context_store.upgrade().map_or(false, |context_store| {
-            context_store
-                .read(cx)
-                .includes_directory(&path_match.path)
-                .is_some()
-        });
-
-        Some(
-            ListItem::new(ix)
-                .inset(true)
-                .toggle_state(selected)
-                .start_slot(
-                    Icon::new(IconName::Folder)
-                        .size(IconSize::XSmall)
-                        .color(Color::Muted),
-                )
-                .child(Label::new(directory_name))
-                .when(added, |el| {
-                    el.end_slot(
-                        h_flex()
-                            .gap_1()
-                            .child(
-                                Icon::new(IconName::Check)
-                                    .size(IconSize::Small)
-                                    .color(Color::Success),
-                            )
-                            .child(Label::new("Added").size(LabelSize::Small)),
-                    )
-                }),
-        )
-    }
-}

crates/assistant2/src/context_picker/file_context_picker.rs

@@ -99,7 +99,6 @@ impl FileContextPickerDelegate {
         query: String,
         cancellation_flag: Arc<AtomicBool>,
         workspace: &Entity<Workspace>,
-
         cx: &mut Context<Picker<Self>>,
     ) -> Task<Vec<PathMatch>> {
         if query.is_empty() {
@@ -124,14 +123,14 @@ impl FileContextPickerDelegate {
             let file_matches = project.worktrees(cx).flat_map(|worktree| {
                 let worktree = worktree.read(cx);
                 let path_prefix: Arc<str> = worktree.root_name().into();
-                worktree.files(false, 0).map(move |entry| PathMatch {
+                worktree.entries(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: false,
+                    is_dir: entry.is_dir(),
                 })
             });
 
@@ -149,7 +148,7 @@ impl FileContextPickerDelegate {
                             .root_entry()
                             .map_or(false, |entry| entry.is_ignored),
                         include_root_name: true,
-                        candidates: project::Candidates::Files,
+                        candidates: project::Candidates::Entries,
                     }
                 })
                 .collect::<Vec<_>>();
@@ -192,7 +191,7 @@ impl PickerDelegate for FileContextPickerDelegate {
     }
 
     fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
-        "Search files…".into()
+        "Search files & directories…".into()
     }
 
     fn update_matches(
@@ -223,13 +222,11 @@ impl PickerDelegate for FileContextPickerDelegate {
             return;
         };
 
-        let Some(file_name) = mat
+        let file_name = mat
             .path
             .file_name()
             .map(|os_str| os_str.to_string_lossy().into_owned())
-        else {
+            .unwrap_or(mat.path_prefix.to_string());
-            return;
-        };
 
         let full_path = mat.path.display().to_string();
 
@@ -238,6 +235,8 @@ impl PickerDelegate for FileContextPickerDelegate {
             path: mat.path.clone(),
         };
 
+        let is_directory = mat.is_dir;
+
         let Some(editor_entity) = self.editor.upgrade() else {
             return;
         };
@@ -288,8 +287,12 @@ impl PickerDelegate for FileContextPickerDelegate {
 
                 editor.insert("\n", window, cx); // Needed to end the fold
 
-                let file_icon = FileIcons::get_icon(&Path::new(&full_path), cx)
+                let file_icon = if is_directory {
-                    .unwrap_or_else(|| SharedString::new(""));
+                    FileIcons::get_folder_icon(false, cx)
+                } else {
+                    FileIcons::get_icon(&Path::new(&full_path), cx)
+                }
+                .unwrap_or_else(|| SharedString::new(""));
 
                 let placeholder = FoldPlaceholder {
                     render: render_fold_icon_button(
@@ -330,7 +333,11 @@ impl PickerDelegate for FileContextPickerDelegate {
         let Some(task) = self
             .context_store
             .update(cx, |context_store, cx| {
-                context_store.add_file_from_path(project_path, cx)
+                if is_directory {
+                    context_store.add_directory(project_path, cx)
+                } else {
+                    context_store.add_file_from_path(project_path, cx)
+                }
             })
             .ok()
         else {
@@ -375,6 +382,7 @@ impl PickerDelegate for FileContextPickerDelegate {
                     ElementId::NamedInteger("file-ctx-picker".into(), ix),
                     &path_match.path,
                     &path_match.path_prefix,
+                    path_match.is_dir,
                     self.context_store.clone(),
                     cx,
                 )),
@@ -386,6 +394,7 @@ pub fn render_file_context_entry(
     id: ElementId,
     path: &Path,
     path_prefix: &Arc<str>,
+    is_directory: bool,
     context_store: WeakEntity<ContextStore>,
     cx: &App,
 ) -> Stateful<Div> {
@@ -409,13 +418,24 @@ pub fn render_file_context_entry(
         (file_name, Some(directory))
     };
 
-    let added = context_store
+    let added = context_store.upgrade().and_then(|context_store| {
-        .upgrade()
+        if is_directory {
-        .and_then(|context_store| context_store.read(cx).will_include_file_path(path, cx));
+            context_store
+                .read(cx)
+                .includes_directory(path)
+                .map(FileInclusion::Direct)
+        } else {
+            context_store.read(cx).will_include_file_path(path, cx)
+        }
+    });
 
-    let file_icon = FileIcons::get_icon(&path, cx)
+    let file_icon = if is_directory {
-        .map(Icon::from_path)
+        FileIcons::get_folder_icon(false, cx)
-        .unwrap_or_else(|| Icon::new(IconName::File));
+    } else {
+        FileIcons::get_icon(&path, cx)
+    }
+    .map(Icon::from_path)
+    .unwrap_or_else(|| Icon::new(IconName::File));
 
     h_flex()
         .id(id)