agent: Indiciate files and folders in `list_directory` (#31448)

Oleksiy Syvokon created

Otherwise, the agent confuses directories with files in cases where dirs
are named like files (`TODO`, `task.js`, etc.)

Release Notes:

- N/A

Change summary

crates/assistant_tools/src/list_directory_tool.rs | 285 ++++++++++++++++
1 file changed, 277 insertions(+), 8 deletions(-)

Detailed changes

crates/assistant_tools/src/list_directory_tool.rs 🔗

@@ -125,18 +125,287 @@ impl Tool for ListDirectoryTool {
             return Task::ready(Err(anyhow!("{} is not a directory.", input.path))).into();
         }
 
-        let mut output = String::new();
+        let mut folders = Vec::new();
+        let mut files = Vec::new();
+
         for entry in worktree.child_entries(&project_path.path) {
-            writeln!(
-                output,
-                "{}",
-                Path::new(worktree.root_name()).join(&entry.path).display(),
-            )
-            .unwrap();
+            let full_path = Path::new(worktree.root_name())
+                .join(&entry.path)
+                .display()
+                .to_string();
+            if entry.is_dir() {
+                folders.push(full_path);
+            } else {
+                files.push(full_path);
+            }
+        }
+
+        let mut output = String::new();
+
+        if !folders.is_empty() {
+            writeln!(output, "# Folders:\n{}", folders.join("\n")).unwrap();
+        }
+
+        if !files.is_empty() {
+            writeln!(output, "\n# Files:\n{}", files.join("\n")).unwrap();
         }
+
         if output.is_empty() {
-            return Task::ready(Ok(format!("{} is empty.", input.path).into())).into();
+            writeln!(output, "{} is empty.", input.path).unwrap();
         }
+
         Task::ready(Ok(output.into())).into()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use assistant_tool::Tool;
+    use gpui::{AppContext, TestAppContext};
+    use indoc::indoc;
+    use language_model::fake_provider::FakeLanguageModel;
+    use project::{FakeFs, Project};
+    use serde_json::json;
+    use settings::SettingsStore;
+    use util::path;
+
+    fn platform_paths(path_str: &str) -> String {
+        if cfg!(target_os = "windows") {
+            path_str.replace("/", "\\")
+        } else {
+            path_str.to_string()
+        }
+    }
+
+    fn init_test(cx: &mut TestAppContext) {
+        cx.update(|cx| {
+            let settings_store = SettingsStore::test(cx);
+            cx.set_global(settings_store);
+            language::init(cx);
+            Project::init_settings(cx);
+        });
+    }
+
+    #[gpui::test]
+    async fn test_list_directory_separates_files_and_dirs(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let fs = FakeFs::new(cx.executor());
+        fs.insert_tree(
+            "/project",
+            json!({
+                "src": {
+                    "main.rs": "fn main() {}",
+                    "lib.rs": "pub fn hello() {}",
+                    "models": {
+                        "user.rs": "struct User {}",
+                        "post.rs": "struct Post {}"
+                    },
+                    "utils": {
+                        "helper.rs": "pub fn help() {}"
+                    }
+                },
+                "tests": {
+                    "integration_test.rs": "#[test] fn test() {}"
+                },
+                "README.md": "# Project",
+                "Cargo.toml": "[package]"
+            }),
+        )
+        .await;
+
+        let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
+        let action_log = cx.new(|_| ActionLog::new(project.clone()));
+        let model = Arc::new(FakeLanguageModel::default());
+        let tool = Arc::new(ListDirectoryTool);
+
+        // Test listing root directory
+        let input = json!({
+            "path": "project"
+        });
+
+        let result = cx
+            .update(|cx| {
+                tool.clone().run(
+                    input,
+                    Arc::default(),
+                    project.clone(),
+                    action_log.clone(),
+                    model.clone(),
+                    None,
+                    cx,
+                )
+            })
+            .output
+            .await
+            .unwrap();
+
+        let content = result.content.as_str().unwrap();
+        assert_eq!(
+            content,
+            platform_paths(indoc! {"
+                # Folders:
+                project/src
+                project/tests
+
+                # Files:
+                project/Cargo.toml
+                project/README.md
+            "})
+        );
+
+        // Test listing src directory
+        let input = json!({
+            "path": "project/src"
+        });
+
+        let result = cx
+            .update(|cx| {
+                tool.clone().run(
+                    input,
+                    Arc::default(),
+                    project.clone(),
+                    action_log.clone(),
+                    model.clone(),
+                    None,
+                    cx,
+                )
+            })
+            .output
+            .await
+            .unwrap();
+
+        let content = result.content.as_str().unwrap();
+        assert_eq!(
+            content,
+            platform_paths(indoc! {"
+                # Folders:
+                project/src/models
+                project/src/utils
+
+                # Files:
+                project/src/lib.rs
+                project/src/main.rs
+            "})
+        );
+
+        // Test listing directory with only files
+        let input = json!({
+            "path": "project/tests"
+        });
+
+        let result = cx
+            .update(|cx| {
+                tool.clone().run(
+                    input,
+                    Arc::default(),
+                    project.clone(),
+                    action_log.clone(),
+                    model.clone(),
+                    None,
+                    cx,
+                )
+            })
+            .output
+            .await
+            .unwrap();
+
+        let content = result.content.as_str().unwrap();
+        assert!(!content.contains("# Folders:"));
+        assert!(content.contains("# Files:"));
+        assert!(content.contains(&platform_paths("project/tests/integration_test.rs")));
+    }
+
+    #[gpui::test]
+    async fn test_list_directory_empty_directory(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let fs = FakeFs::new(cx.executor());
+        fs.insert_tree(
+            "/project",
+            json!({
+                "empty_dir": {}
+            }),
+        )
+        .await;
+
+        let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
+        let action_log = cx.new(|_| ActionLog::new(project.clone()));
+        let model = Arc::new(FakeLanguageModel::default());
+        let tool = Arc::new(ListDirectoryTool);
+
+        let input = json!({
+            "path": "project/empty_dir"
+        });
+
+        let result = cx
+            .update(|cx| tool.run(input, Arc::default(), project, action_log, model, None, cx))
+            .output
+            .await
+            .unwrap();
+
+        let content = result.content.as_str().unwrap();
+        assert_eq!(content, "project/empty_dir is empty.\n");
+    }
+
+    #[gpui::test]
+    async fn test_list_directory_error_cases(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let fs = FakeFs::new(cx.executor());
+        fs.insert_tree(
+            "/project",
+            json!({
+                "file.txt": "content"
+            }),
+        )
+        .await;
+
+        let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
+        let action_log = cx.new(|_| ActionLog::new(project.clone()));
+        let model = Arc::new(FakeLanguageModel::default());
+        let tool = Arc::new(ListDirectoryTool);
+
+        // Test non-existent path
+        let input = json!({
+            "path": "project/nonexistent"
+        });
+
+        let result = cx
+            .update(|cx| {
+                tool.clone().run(
+                    input,
+                    Arc::default(),
+                    project.clone(),
+                    action_log.clone(),
+                    model.clone(),
+                    None,
+                    cx,
+                )
+            })
+            .output
+            .await;
+
+        assert!(result.is_err());
+        assert!(result.unwrap_err().to_string().contains("Path not found"));
+
+        // Test trying to list a file instead of directory
+        let input = json!({
+            "path": "project/file.txt"
+        });
+
+        let result = cx
+            .update(|cx| tool.run(input, Arc::default(), project, action_log, model, None, cx))
+            .output
+            .await;
+
+        assert!(result.is_err());
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("is not a directory")
+        );
+    }
+}