fix: directory tree tool result (#26)

* fix: directory tree result

* chore: fix typo

* chore: clippy warnings

Author: hashemix

rust-toolchain.toml

@@ -0,0 +1,3 @@
+[toolchain]
+channel = "1.88.0"
+components = ["rustfmt", "clippy"]

src/fs_service.rs

@@ -6,6 +6,7 @@ use grep::{
     regex::RegexMatcherBuilder,
     searcher::{sinks::UTF8, BinaryDetection, Searcher},
 };
+use serde_json::{json, Value};
 
 use std::{
     env,
@@ -67,7 +68,7 @@ impl FileSystemService {
             .map_while(|dir| {
                 let expand_result = expand_home(dir.into());
                 if !expand_result.is_dir() {
-                    panic!("{}", format!("Error: {} is not a directory", dir));
+                    panic!("{}", format!("Error: {dir} is not a directory"));
                 }
                 Some(expand_result)
             })
@@ -179,7 +180,7 @@ impl FileSystemService {
         if target_path.exists() {
             return Err(std::io::Error::new(
                 std::io::ErrorKind::AlreadyExists,
-                format!("'{}' already exists!", target_zip_file),
+                format!("'{target_zip_file}' already exists!"),
             )
             .into());
         }
@@ -269,7 +270,7 @@ impl FileSystemService {
         if target_path.exists() {
             return Err(std::io::Error::new(
                 std::io::ErrorKind::AlreadyExists,
-                format!("'{}' already exists!", target_zip_file),
+                format!("'{target_zip_file}' already exists!"),
             )
             .into());
         }
@@ -326,7 +327,7 @@ impl FileSystemService {
         if target_dir_path.exists() {
             return Err(std::io::Error::new(
                 std::io::ErrorKind::AlreadyExists,
-                format!("'{}' directory already exists!", target_dir),
+                format!("'{target_dir}' directory already exists!"),
             )
             .into());
         }
@@ -473,7 +474,7 @@ impl FileSystemService {
                     let glob_pattern = if pattern.contains('*') {
                         pattern.clone()
                     } else {
-                        format!("*{}*", pattern)
+                        format!("*{pattern}*")
                     };
 
                     Pattern::new(&glob_pattern)
@@ -501,6 +502,85 @@ impl FileSystemService {
         Ok(result)
     }
 
+    /// Generates a JSON representation of a directory tree starting at the given path.
+    ///
+    /// This function recursively builds a JSON array object representing the directory structure,
+    /// where each entry includes a `name` (file or directory name), `type` ("file" or "directory"),
+    /// and for directories, a `children` array containing their contents. Files do not have a
+    /// `children` field.
+    ///
+    /// The function supports optional constraints to limit the tree size:
+    /// - `max_depth`: Limits the depth of directory traversal.
+    /// - `max_files`: Limits the total number of entries (files and directories).
+    ///
+    /// # IMPORTANT NOTE
+    ///
+    /// use max_depth or max_files could lead to partial or skewed representations of actual directory tree
+    pub fn directory_tree<P: AsRef<Path>>(
+        &self,
+        root_path: P,
+        max_depth: Option<usize>,
+        max_files: Option<usize>,
+        current_count: &mut usize,
+    ) -> ServiceResult<Value> {
+        let valid_path = self.validate_path(root_path.as_ref())?;
+
+        let metadata = fs::metadata(&valid_path)?;
+        if !metadata.is_dir() {
+            return Err(ServiceError::FromString(
+                "Root path must be a directory".into(),
+            ));
+        }
+
+        let mut children = Vec::new();
+
+        if max_depth != Some(0) {
+            for entry in WalkDir::new(valid_path)
+                .min_depth(1)
+                .max_depth(1)
+                .follow_links(true)
+                .into_iter()
+                .filter_map(|e| e.ok())
+            {
+                let child_path = entry.path();
+                let metadata = fs::metadata(child_path)?;
+
+                let entry_name = child_path
+                    .file_name()
+                    .ok_or(ServiceError::FromString("Invalid path".to_string()))?
+                    .to_string_lossy()
+                    .into_owned();
+
+                // Increment the count for this entry
+                *current_count += 1;
+
+                // Check if we've exceeded max_files (if set)
+                if let Some(max) = max_files {
+                    if *current_count > max {
+                        continue; // Skip this entry but continue processing others
+                    }
+                }
+
+                let mut json_entry = json!({
+                    "name": entry_name,
+                    "type": if metadata.is_dir() { "directory" } else { "file" }
+                });
+
+                if metadata.is_dir() {
+                    let next_depth = max_depth.map(|d| d - 1);
+                    let child_children =
+                        self.directory_tree(child_path, next_depth, max_files, current_count)?;
+                    json_entry
+                        .as_object_mut()
+                        .unwrap()
+                        .insert("children".to_string(), child_children);
+                }
+                children.push(json_entry);
+            }
+        }
+        Ok(Value::Array(children))
+    }
+
     pub fn create_unified_diff(
         &self,
         original_content: &str,
@@ -519,8 +599,8 @@ impl FileSystemService {
         let patch = diff
             .unified_diff()
             .header(
-                format!("{}\toriginal", file_name).as_str(),
-                format!("{}\tmodified", file_name).as_str(),
+                format!("{file_name}\toriginal").as_str(),
+                format!("{file_name}\tmodified").as_str(),
             )
             .context_radius(4)
             .to_string();

src/fs_service/utils.rs

@@ -82,7 +82,7 @@ pub fn format_bytes(bytes: u64) -> String {
             return format!("{:.2} {}", bytes as f64 / threshold as f64, unit);
         }
     }
-    format!("{} bytes", bytes)
+    format!("{bytes} bytes")
 }
 
 pub async fn write_zip_entry(

src/handler.rs

@@ -77,7 +77,7 @@ impl ServerHandler for MyServerHandler {
     ) -> std::result::Result<InitializeResult, RpcError> {
         runtime
             .set_client_details(initialize_request.params.clone())
-            .map_err(|err| RpcError::internal_error().with_message(format!("{}", err)))?;
+            .map_err(|err| RpcError::internal_error().with_message(format!("{err}")))?;
 
         let mut server_info = runtime.server_info().to_owned();
         // Provide compatibility for clients using older MCP protocol versions.

src/tools/directory_tree.rs

@@ -1,16 +1,17 @@
-use std::path::Path;
-
 use rust_mcp_sdk::macros::{mcp_tool, JsonSchema};
 use rust_mcp_sdk::schema::{schema_utils::CallToolError, CallToolResult};
 use serde_json::json;
 
+use crate::error::ServiceError;
 use crate::fs_service::FileSystemService;
 
 #[mcp_tool(
     name = "directory_tree",
     description = concat!("Get a recursive tree view of files and directories as a JSON structure. ",
     "Each entry includes 'name', 'type' (file/directory), and 'children' for directories. ",
     "Files have no children array, while directories always have a children array (which may be empty). ",
+    "If the 'max_depth' parameter is provided, the traversal will be limited to the specified depth. ",
+    "As a result, the returned directory structure may be incomplete or provide a skewed representation of the full directory tree, since deeper-level files and subdirectories beyond the specified depth will be excluded. ",
     "The output is formatted with 2-space indentation for readability. Only works within allowed directories."),
     destructive_hint = false,
     idempotent_hint = false,
@@ -21,28 +22,31 @@ use crate::fs_service::FileSystemService;
 pub struct DirectoryTreeTool {
     /// The root path of the directory tree to generate.
     pub path: String,
+    /// Limits the depth of directory traversal
+    pub max_depth: Option<u64>,
 }
 impl DirectoryTreeTool {
     pub async fn run_tool(
         params: Self,
         context: &FileSystemService,
     ) -> std::result::Result<CallToolResult, CallToolError> {
+        let mut entry_counter: usize = 0;
         let entries = context
-            .list_directory(Path::new(&params.path))
-            .await
+            .directory_tree(
+                params.path,
+                params.max_depth.map(|v| v as usize),
+                None,
+                &mut entry_counter,
+            )
             .map_err(CallToolError::new)?;
 
-        let json_tree: Vec<serde_json::Value> = entries
-            .iter()
-            .map(|entry| {
-                json!({
-                    "name": entry.file_name().to_str().unwrap_or_default(),
-                    "type": if entry.path().is_dir(){"directory"}else{"file"}
-                })
-            })
-            .collect();
-        let json_str =
-            serde_json::to_string_pretty(&json!(json_tree)).map_err(CallToolError::new)?;
+        if entry_counter == 0 {
+            return Err(CallToolError::new(ServiceError::FromString(
+                "Could not find any entries".to_string(),
+            )));
+        }
+
+        let json_str = serde_json::to_string_pretty(&json!(entries)).map_err(CallToolError::new)?;
         Ok(CallToolResult::text_content(json_str, None))
     }
 }

src/tools/read_multiple_files.rs

@@ -40,8 +40,8 @@ impl ReadMultipleFilesTool {
                         .map_err(CallToolError::new);
 
                     content.map_or_else(
-                        |err| format!("{}: Error - {}", path, err),
-                        |value| format!("{}:\n{}\n", path, value),
+                        |err| format!("{path}: Error - {err}"),
+                        |value| format!("{path}:\n{value}\n"),
                     )
                 }
             })