From 63d8a43f9d4722578a86a98d9e4f3763922fedff Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Tue, 17 Dec 2024 16:53:56 -0500
Subject: [PATCH] assistant2: Fix file context picker (#22163)

This PR fixes an issue with the Assistant2 file context picker where
files weren't being attached properly.

We needed to retrieve the files from the worktree without the worktree
root name in the file path.

Release Notes:

- N/A

---------

Co-authored-by: Max <max@zed.dev>
---
 .../src/context_picker/file_context_picker.rs | 101 ++++++++++--------
 1 file changed, 54 insertions(+), 47 deletions(-)

diff --git a/crates/assistant2/src/context_picker/file_context_picker.rs b/crates/assistant2/src/context_picker/file_context_picker.rs
index be34283826e1f97a2e945b8e03841c41bb6df7c7..9ab5decefaba2205fb29c4bfd4fa3b0cb0ad45bd 100644
--- a/crates/assistant2/src/context_picker/file_context_picker.rs
+++ b/crates/assistant2/src/context_picker/file_context_picker.rs
@@ -1,6 +1,6 @@
 use std::fmt::Write as _;
 use std::ops::RangeInclusive;
-use std::path::{Path, PathBuf};
+use std::path::Path;
 use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 
@@ -79,44 +79,37 @@ impl FileContextPickerDelegate {
         if query.is_empty() {
             let workspace = workspace.read(cx);
             let project = workspace.project().read(cx);
-            let entries = workspace.recent_navigation_history(Some(10), cx);
-
-            let entries = entries
+            let recent_matches = workspace
+                .recent_navigation_history(Some(10), cx)
                 .into_iter()
-                .map(|entries| entries.0)
-                .chain(project.worktrees(cx).flat_map(|worktree| {
-                    let worktree = worktree.read(cx);
-                    let id = worktree.id();
-                    worktree
-                        .child_entries(Path::new(""))
-                        .filter(|entry| entry.kind.is_file())
-                        .map(move |entry| project::ProjectPath {
-                            worktree_id: id,
-                            path: entry.path.clone(),
-                        })
-                }))
-                .collect::<Vec<_>>();
-
-            let path_prefix: Arc<str> = Arc::default();
-            Task::ready(
-                entries
-                    .into_iter()
-                    .filter_map(|entry| {
-                        let worktree = project.worktree_for_id(entry.worktree_id, cx)?;
-                        let mut full_path = PathBuf::from(worktree.read(cx).root_name());
-                        full_path.push(&entry.path);
-                        Some(PathMatch {
-                            score: 0.,
-                            positions: Vec::new(),
-                            worktree_id: entry.worktree_id.to_usize(),
-                            path: full_path.into(),
-                            path_prefix: path_prefix.clone(),
-                            distance_to_relative_ancestor: 0,
-                            is_dir: false,
-                        })
+                .filter_map(|(project_path, _)| {
+                    let worktree = project.worktree_for_id(project_path.worktree_id, cx)?;
+                    Some(PathMatch {
+                        score: 0.,
+                        positions: Vec::new(),
+                        worktree_id: project_path.worktree_id.to_usize(),
+                        path: project_path.path,
+                        path_prefix: worktree.read(cx).root_name().into(),
+                        distance_to_relative_ancestor: 0,
+                        is_dir: false,
                     })
-                    .collect(),
-            )
+                });
+
+            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(true, 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,
+                })
+            });
+
+            Task::ready(recent_matches.chain(file_matches).collect())
         } else {
             let worktrees = workspace.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
             let candidate_sets = worktrees
@@ -253,16 +246,30 @@ impl PickerDelegate for FileContextPickerDelegate {
         _cx: &mut ViewContext<Picker<Self>>,
     ) -> Option<Self::ListItem> {
         let path_match = &self.matches[ix];
-        let file_name = path_match
-            .path
-            .file_name()
-            .unwrap_or_default()
-            .to_string_lossy()
-            .to_string();
-        let directory = path_match
-            .path
-            .parent()
-            .map(|directory| format!("{}/", directory.to_string_lossy()));
+
+        let (file_name, directory) = if path_match.path.as_ref() == Path::new("") {
+            (SharedString::from(path_match.path_prefix.clone()), None)
+        } else {
+            let file_name = path_match
+                .path
+                .file_name()
+                .unwrap_or_default()
+                .to_string_lossy()
+                .to_string()
+                .into();
+
+            let mut directory = format!("{}/", path_match.path_prefix);
+            if let Some(parent) = path_match
+                .path
+                .parent()
+                .filter(|parent| parent != &Path::new(""))
+            {
+                directory.push_str(&parent.to_string_lossy());
+                directory.push('/');
+            }
+
+            (file_name, Some(directory))
+        };
 
         Some(
             ListItem::new(ix).inset(true).toggle_state(selected).child(