diff --git a/src/args.rs b/src/args.rs index 9014899..a70752a 100644 --- a/src/args.rs +++ b/src/args.rs @@ -1,24 +1,15 @@ +use crate::tools::{DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, DEFAULT_MAX_RESULTS, OutputMode}; use rmcp::schemars::{self, JsonSchema}; use serde::{Deserialize, Serialize}; -// --------------------------------------------------------------------------- -// Default value functions for serde(default = "…") -// --------------------------------------------------------------------------- - -const fn default_zero() -> usize { - 0 -} -const fn default_max_results() -> usize { - 100 -} -const fn default_max_bytes() -> usize { - 5 * 1024 * 1024 // 5 MiB +fn default_max_results() -> usize { + DEFAULT_MAX_RESULTS } -const fn default_max_lines() -> usize { - 2000 +fn default_max_bytes() -> usize { + DEFAULT_MAX_BYTES } -const fn default_false() -> bool { - false +fn default_max_lines() -> usize { + DEFAULT_MAX_LINES } const fn default_true() -> bool { true @@ -51,14 +42,6 @@ impl StringOrVec { } } -fn default_file_extensions() -> Option { - None -} - -fn default_output_mode() -> String { - "files_with_matches".to_string() -} - // --------------------------------------------------------------------------- // Arg structs — serde fills in defaults automatically // --------------------------------------------------------------------------- @@ -69,38 +52,48 @@ fn default_output_mode() -> String { pub struct GrepArgs { #[schemars(description = "Directory to search in")] pub directory: String, - #[schemars(description = "Regex pattern to search for (Rust regex; no lookaround/backrefs). Use (?i) for case-insensitive.")] + #[schemars( + description = "Regex pattern to search for (Rust regex; no lookaround/backrefs). Use (?i) for case-insensitive." + )] pub pattern: String, - #[serde(default = "default_zero")] + #[serde(default)] #[schemars(description = "Number of lines of context before each match (default 0)")] pub before_context: usize, - #[serde(default = "default_zero")] + #[serde(default)] #[schemars(description = "Number of lines of context after each match (default 0)")] pub after_context: usize, #[serde(default = "default_max_results")] #[schemars(description = "Maximum number of results to return (default 100)")] pub max_results: usize, - #[serde(default = "default_false")] - #[schemars(description = "Case-insensitive search (default false). Equivalent to prefixing pattern with (?i).")] + #[serde(default)] + #[schemars( + description = "Case-insensitive search (default false). Equivalent to prefixing pattern with (?i)." + )] pub case_insensitive: bool, - #[serde(default = "default_false")] + #[serde(default)] #[schemars(description = "Include hidden files and directories (default false)")] pub include_hidden: bool, - #[serde(default = "default_false")] + #[serde(default)] #[schemars(description = "Follow symbolic links (default false)")] pub follow_symlinks: bool, #[serde(default = "default_true")] #[schemars(description = "Respect .gitignore files (default true)")] pub respect_gitignore: bool, - #[serde(default = "default_file_extensions")] - #[schemars(description = "Restrict to files with these extensions. Accepts either a single string (\"sql\") or an array ([\"rs\", \"toml\"]). Empty means all files.")] + #[serde(default)] + #[schemars( + description = "Restrict to files with these extensions. Accepts either a single string (\"sql\") or an array ([\"rs\", \"toml\"]). Empty means all files." + )] pub file_extensions: Option, #[serde(default = "default_max_bytes")] - #[schemars(description = "Hard cap on total response size in bytes (default ~5 MiB). Truncates with a marker.")] + #[schemars( + description = "Hard cap on total response size in bytes (default ~5 MiB). Truncates with a marker." + )] pub max_bytes: usize, - #[serde(default = "default_output_mode")] - #[schemars(description = "Output mode: 'files_with_matches' (default — list file paths only), 'content' (matching lines with line numbers), 'count' (per-file match tallies as path: N).")] - pub output_mode: String, + #[serde(default)] + #[schemars( + description = "Output mode: 'files_with_matches' (default — list file paths only), 'content' (matching lines with line numbers), 'count' (per-file match tallies as path: N)." + )] + pub output_mode: OutputMode, } #[derive(Debug, Serialize, Deserialize, JsonSchema)] @@ -108,19 +101,23 @@ pub struct GrepArgs { pub struct FindArgs { #[schemars(description = "Directory to search in")] pub directory: String, - #[schemars(description = "Regex pattern to match against filenames (Rust regex; no lookaround/backrefs)")] + #[schemars( + description = "Regex pattern to match against filenames (Rust regex; no lookaround/backrefs)" + )] pub pattern: String, #[serde(default = "default_max_results")] #[schemars(description = "Maximum number of results to return (default 100)")] pub max_results: usize, - #[serde(default = "default_false")] + #[serde(default)] #[schemars(description = "Include hidden files and directories (default false)")] pub include_hidden: bool, #[serde(default = "default_true")] #[schemars(description = "Respect .gitignore files (default true)")] pub respect_gitignore: bool, #[serde(default = "default_true")] - #[schemars(description = "Match the basename only (default true). Set false to match the full path.")] + #[schemars( + description = "Match the basename only (default true). Set false to match the full path." + )] pub match_basename: bool, } @@ -129,8 +126,10 @@ pub struct FindArgs { pub struct CatArgs { #[schemars(description = "Path to the file to read")] pub file_path: String, - #[serde(default = "default_zero")] - #[schemars(description = "Line offset to start from (0-based, default 0). Use to paginate long files.")] + #[serde(default)] + #[schemars( + description = "Line offset to start from (0-based, default 0). Use to paginate long files." + )] pub offset: usize, #[serde(default = "default_max_lines")] #[schemars(description = "Maximum number of lines to return (default 2000)")] @@ -144,6 +143,8 @@ pub struct CatArgs { #[serde(rename_all = "snake_case")] pub struct MemoriesArgs { #[serde(default)] - #[schemars(description = "Optional memory file name (relative to memory dir, e.g. \"user_role.md\"). If omitted, returns the index from MEMORY.md or a directory listing.")] + #[schemars( + description = "Optional memory file name (relative to memory dir, e.g. \"user_role.md\"). If omitted, returns the index from MEMORY.md or a directory listing." + )] pub name: Option, } diff --git a/src/cli.rs b/src/cli.rs index ede8567..9c8f3ef 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,6 +1,6 @@ -use std::path::PathBuf; use clap::Parser; use std::net::SocketAddr; +use std::path::PathBuf; /// Command-line arguments for `code-mcp`. /// @@ -8,7 +8,10 @@ use std::net::SocketAddr; /// by `main`; the struct itself is `pub` so it can be referenced from other /// crate modules. #[derive(Debug, Parser)] -#[command(name = "code-mcp", about = "Streamable HTTP MCP server for code search/read tools")] +#[command( + name = "code-mcp", + about = "Streamable HTTP MCP server for code search/read tools" +)] pub struct Args { /// Address to bind, e.g. 0.0.0.0:8080 #[arg(long, default_value = "0.0.0.0:8080")] diff --git a/src/error.rs b/src/error.rs index 3b151ff..760e00d 100644 --- a/src/error.rs +++ b/src/error.rs @@ -71,15 +71,11 @@ impl From for ErrorData { | AppError::GrepRegex(_) | AppError::InvalidRequest(_) | AppError::NotFound(_) - | AppError::OutOfScope(_) => ErrorData::invalid_params( - "invalid_params", - Some(json!({"error": err.to_string()})), - ), + | AppError::OutOfScope(_) => { + ErrorData::invalid_params("invalid_params", Some(json!({"error": err.to_string()}))) + } AppError::Io(_) | AppError::Ignore(_) | AppError::Internal(_) | AppError::Axum(_) => { - ErrorData::internal_error( - "internal_error", - Some(json!({"error": err.to_string()})), - ) + ErrorData::internal_error("internal_error", Some(json!({"error": err.to_string()}))) } } } diff --git a/src/gate.rs b/src/gate.rs index de852c7..2643c9b 100644 --- a/src/gate.rs +++ b/src/gate.rs @@ -6,9 +6,7 @@ use axum::extract::{ConnectInfo, State}; use axum::http::{HeaderMap, Method, Request, StatusCode, header}; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; -use rmcp::transport::streamable_http_server::session::{ - SessionId, local::LocalSessionManager, -}; +use rmcp::transport::streamable_http_server::session::{SessionId, local::LocalSessionManager}; use crate::limiter::PeerLimiter; use crate::reaper::ActivityTracker; @@ -107,7 +105,7 @@ fn peer_ip(headers: &HeaderMap, addr: SocketAddr, trust_xff: bool) -> IpAddr { .map(str::trim) .and_then(|s| s.parse::().ok()) { - tracing::info!(peer = %ip, socket = %addr.ip(), "resolved peer IP from X-Forwarded-For"); + tracing::debug!(peer = %ip, socket = %addr.ip(), "resolved peer IP from X-Forwarded-For"); return ip; } addr.ip() @@ -155,7 +153,10 @@ mod tests { activity: Arc::new(ActivityTracker::new()), }); let app = build_app(ctx); - let res = app.oneshot(req(Method::GET, dummy_addr(), false)).await.unwrap(); + let res = app + .oneshot(req(Method::GET, dummy_addr(), false)) + .await + .unwrap(); assert_eq!(res.status(), StatusCode::OK); } @@ -170,7 +171,10 @@ mod tests { activity: Arc::new(ActivityTracker::new()), }); let app = build_app(ctx); - let res = app.oneshot(req(Method::POST, dummy_addr(), true)).await.unwrap(); + let res = app + .oneshot(req(Method::POST, dummy_addr(), true)) + .await + .unwrap(); assert_eq!(res.status(), StatusCode::OK); } @@ -184,7 +188,10 @@ mod tests { activity: Arc::new(ActivityTracker::new()), }); let app = build_app(ctx); - let res = app.oneshot(req(Method::POST, dummy_addr(), false)).await.unwrap(); + let res = app + .oneshot(req(Method::POST, dummy_addr(), false)) + .await + .unwrap(); assert_eq!(res.status(), StatusCode::SERVICE_UNAVAILABLE); assert_eq!(res.headers().get(header::RETRY_AFTER).unwrap(), "5"); } @@ -201,15 +208,26 @@ mod tests { let app = build_app(ctx); for _ in 0..2 { - let res = app.clone().oneshot(req(Method::POST, dummy_addr(), false)).await.unwrap(); + let res = app + .clone() + .oneshot(req(Method::POST, dummy_addr(), false)) + .await + .unwrap(); assert_eq!(res.status(), StatusCode::OK); } - let res = app.clone().oneshot(req(Method::POST, dummy_addr(), false)).await.unwrap(); + let res = app + .clone() + .oneshot(req(Method::POST, dummy_addr(), false)) + .await + .unwrap(); assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS); assert!(res.headers().get(header::RETRY_AFTER).is_some()); // A different peer is unaffected. - let res = app.oneshot(req(Method::POST, other_addr(), false)).await.unwrap(); + let res = app + .oneshot(req(Method::POST, other_addr(), false)) + .await + .unwrap(); assert_eq!(res.status(), StatusCode::OK); } @@ -217,7 +235,10 @@ mod tests { fn peer_ip_uses_socket_addr_by_default() { let h = HeaderMap::new(); let addr: SocketAddr = "10.0.0.5:1234".parse().unwrap(); - assert_eq!(peer_ip(&h, addr, false), IpAddr::V4(Ipv4Addr::new(10, 0, 0, 5))); + assert_eq!( + peer_ip(&h, addr, false), + IpAddr::V4(Ipv4Addr::new(10, 0, 0, 5)) + ); } #[test] @@ -225,7 +246,10 @@ mod tests { let mut h = HeaderMap::new(); h.insert("x-forwarded-for", "1.2.3.4".parse().unwrap()); let addr: SocketAddr = "10.0.0.5:1234".parse().unwrap(); - assert_eq!(peer_ip(&h, addr, false), IpAddr::V4(Ipv4Addr::new(10, 0, 0, 5))); + assert_eq!( + peer_ip(&h, addr, false), + IpAddr::V4(Ipv4Addr::new(10, 0, 0, 5)) + ); } #[test] @@ -236,7 +260,10 @@ mod tests { let mut h = HeaderMap::new(); h.insert("x-forwarded-for", "1.2.3.4, 5.6.7.8".parse().unwrap()); let addr: SocketAddr = "10.0.0.5:1234".parse().unwrap(); - assert_eq!(peer_ip(&h, addr, true), IpAddr::V4(Ipv4Addr::new(5, 6, 7, 8))); + assert_eq!( + peer_ip(&h, addr, true), + IpAddr::V4(Ipv4Addr::new(5, 6, 7, 8)) + ); } #[test] @@ -244,6 +271,9 @@ mod tests { let mut h = HeaderMap::new(); h.insert("x-forwarded-for", "garbage".parse().unwrap()); let addr: SocketAddr = "10.0.0.5:1234".parse().unwrap(); - assert_eq!(peer_ip(&h, addr, true), IpAddr::V4(Ipv4Addr::new(10, 0, 0, 5))); + assert_eq!( + peer_ip(&h, addr, true), + IpAddr::V4(Ipv4Addr::new(10, 0, 0, 5)) + ); } } diff --git a/src/limiter.rs b/src/limiter.rs index c638c50..e42feff 100644 --- a/src/limiter.rs +++ b/src/limiter.rs @@ -102,7 +102,9 @@ mod tests { let l = PeerLimiter::new(2.0, 0.001, 1024); l.try_consume(ip(1)).unwrap(); l.try_consume(ip(1)).unwrap(); - let err = l.try_consume(ip(1)).expect_err("third call should be rate-limited"); + let err = l + .try_consume(ip(1)) + .expect_err("third call should be rate-limited"); assert!(err > Duration::from_secs(0)); } diff --git a/src/main.rs b/src/main.rs index 8b4fa18..6ea0b85 100644 --- a/src/main.rs +++ b/src/main.rs @@ -39,6 +39,7 @@ //! `invalid_params`. See . mod args; +mod cli; mod error; mod gate; mod limiter; @@ -47,26 +48,23 @@ mod reaper; mod scope; mod server; mod tools; -mod cli; +use clap::Parser; use std::sync::Arc; use std::time::Duration; -use clap::Parser; use std::net::SocketAddr; - use rmcp::transport::streamable_http_server::{ - StreamableHttpServerConfig, StreamableHttpService, - session::local::LocalSessionManager, + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, }; use tokio_util::sync::CancellationToken; +use crate::cli::Args; use crate::gate::{GateCtx, gate}; use crate::limiter::PeerLimiter; use crate::scope::Scope; use crate::server::CodeMcpServer; -use crate::cli::Args; use crate::error::AppError; diff --git a/src/memory.rs b/src/memory.rs index e870597..5a69da0 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -1,4 +1,5 @@ use crate::error::AppError; +use crate::tools::ToolResponse; use std::fmt::Write; use std::path::Path; @@ -7,7 +8,11 @@ use std::path::Path; /// - If `name` is `Some`, reads that specific `.md` file (rejecting path traversal). /// - If `name` is `None`, returns `MEMORY.md` if present, otherwise a listing of /// all `.md` files in the directory. -pub fn load_memory(dir: &Path, name: Option<&str>) -> Result { +/// +/// Returns a [`ToolResponse`] (with no truncation/match metadata, since memory +/// files are read whole) so the caller in `server.rs` can treat this like any +/// other tool entry point. +pub fn load_memory(dir: &Path, name: Option<&str>) -> Result { if !dir.is_dir() { return Err(AppError::NotFound(format!( "memory dir does not exist: {}", @@ -26,13 +31,15 @@ pub fn load_memory(dir: &Path, name: Option<&str>) -> Result { if !path.is_file() { return Err(AppError::NotFound(format!("memory not found: {name}"))); } - return Ok(std::fs::read_to_string(&path)?); + return Ok(ToolResponse::text(std::fs::read_to_string(&path)?)); } // No name: prefer MEMORY.md, otherwise list *.md files. let index = dir.join("MEMORY.md"); - if index.is_file() { - return Ok(std::fs::read_to_string(&index)?); + match std::fs::read_to_string(&index) { + Ok(content) => return Ok(ToolResponse::text(content)), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => return Err(AppError::Io(e)), } let mut listing = String::from("# Memory dir contents\n\n"); @@ -59,7 +66,7 @@ pub fn load_memory(dir: &Path, name: Option<&str>) -> Result { or create a `MEMORY.md` index at the top level.\n", ); } - Ok(listing) + Ok(ToolResponse::text(listing)) } #[cfg(test)] @@ -76,7 +83,7 @@ mod tests { fs::write(td.path().join("MEMORY.md"), "# index\n- foo\n")?; fs::write(td.path().join("foo.md"), "ignored\n")?; - let out = load_memory(td.path(), None)?; + let out = load_memory(td.path(), None)?.content; assert!(out.starts_with("# index"), "got {:?}", out); Ok(()) } @@ -88,7 +95,7 @@ mod tests { fs::write(td.path().join("b.md"), "")?; fs::write(td.path().join("ignore.txt"), "")?; - let out = load_memory(td.path(), None)?; + let out = load_memory(td.path(), None)?.content; assert!(out.contains("- a.md"), "got {}", out); assert!(out.contains("- b.md"), "got {}", out); assert!(!out.contains("ignore.txt"), "got {}", out); @@ -100,7 +107,7 @@ mod tests { let td = TempDir::new()?; fs::write(td.path().join("user_role.md"), "data scientist\n")?; - let out = load_memory(td.path(), Some("user_role.md"))?; + let out = load_memory(td.path(), Some("user_role.md"))?.content; assert_eq!(out, "data scientist\n"); Ok(()) } diff --git a/src/reaper.rs b/src/reaper.rs index f356ff4..81d5ce1 100644 --- a/src/reaper.rs +++ b/src/reaper.rs @@ -58,11 +58,7 @@ pub async fn reap_loop( } } -async fn sweep( - manager: &LocalSessionManager, - tracker: &ActivityTracker, - idle_timeout: Duration, -) { +async fn sweep(manager: &LocalSessionManager, tracker: &ActivityTracker, idle_timeout: Duration) { let now = Instant::now(); let live_ids: Vec = { let s = manager.sessions.read().await; diff --git a/src/scope.rs b/src/scope.rs index 7be4f63..644be76 100644 --- a/src/scope.rs +++ b/src/scope.rs @@ -36,9 +36,9 @@ impl Scope { /// rejected. pub fn check>(&self, input: P) -> Result { let input = input.as_ref(); - let canon = input.canonicalize().map_err(|e| { - AppError::NotFound(format!("{}: {}", input.display(), e)) - })?; + let canon = input + .canonicalize() + .map_err(|e| AppError::NotFound(format!("{}: {}", input.display(), e)))?; if !canon.starts_with(&self.root) { return Err(AppError::OutOfScope(format!( "{} is outside project root {}", diff --git a/src/server.rs b/src/server.rs index 2aefe4d..7e58034 100644 --- a/src/server.rs +++ b/src/server.rs @@ -2,17 +2,15 @@ use std::fmt::Write; use std::path::PathBuf; use rmcp::{ - ServerHandler, - handler::server::wrapper::Parameters, - model::CallToolResult, - tool, tool_handler, tool_router, + ServerHandler, handler::server::wrapper::Parameters, model::CallToolResult, tool, tool_handler, + tool_router, }; use crate::args::{CatArgs, FindArgs, GrepArgs, MemoriesArgs, StringOrVec}; -use crate::error::{ToolResult, join_error, tool_error}; +use crate::error::{AppError, ToolResult, join_error, tool_error}; use crate::memory::load_memory; use crate::scope::Scope; -use crate::tools::{self, OutputMode}; +use crate::tools; /// The MCP server handler. Owns the tool router, the optional memory dir, /// the extra instructions loaded at startup, and the filesystem [`Scope`]. @@ -52,18 +50,11 @@ impl CodeMcpServer { #[tool( description = "Regex search across files (parallel, gitignore-aware). output_mode: 'files_with_matches' (default — list matching file paths), 'content' (matching lines with line numbers), 'count' (per-file match tallies). Other options: case_insensitive, file_extensions, before/after_context, include_hidden, follow_symlinks, max_results, max_bytes." )] - async fn grep( - &self, - Parameters(args): Parameters, - ) -> ToolResult { + async fn grep(&self, Parameters(args): Parameters) -> ToolResult { let directory = match self.scope.check(&args.directory) { Ok(d) => d, Err(e) => return Ok(tool_error(e)), }; - let output_mode = match OutputMode::from_str_lossy(&args.output_mode) { - Ok(m) => m, - Err(e) => return Ok(tool_error(e)), - }; let res = tokio::task::spawn_blocking(move || { let opts = tools::GrepOptions { before_context: args.before_context, @@ -78,9 +69,9 @@ impl CodeMcpServer { .map(StringOrVec::into_vec) .unwrap_or_default(), max_bytes: args.max_bytes, - output_mode, + output_mode: args.output_mode, }; - tools::grep(&directory.to_string_lossy(), &args.pattern, opts) + tools::grep(&directory, &args.pattern, opts) }) .await .map_err(join_error)?; @@ -94,10 +85,7 @@ impl CodeMcpServer { #[tool( description = "Find files by regex (matches basename by default; set match_basename=false to match full path). Options: include_hidden, respect_gitignore, max_results." )] - async fn find( - &self, - Parameters(args): Parameters, - ) -> ToolResult { + async fn find(&self, Parameters(args): Parameters) -> ToolResult { let directory = match self.scope.check(&args.directory) { Ok(d) => d, Err(e) => return Ok(tool_error(e)), @@ -109,7 +97,7 @@ impl CodeMcpServer { respect_gitignore: args.respect_gitignore, match_basename: args.match_basename, }; - tools::find(&directory.to_string_lossy(), &args.pattern, opts) + tools::find(&directory, &args.pattern, opts) }) .await .map_err(join_error)?; @@ -129,12 +117,7 @@ impl CodeMcpServer { Err(e) => return Ok(tool_error(e)), }; let res = tokio::task::spawn_blocking(move || { - tools::cat( - &file_path.to_string_lossy(), - args.offset, - args.max_lines, - args.max_bytes, - ) + tools::cat(&file_path, args.offset, args.max_lines, args.max_bytes) }) .await .map_err(join_error)?; @@ -154,16 +137,9 @@ impl CodeMcpServer { ) -> ToolResult { let dir = match self.memory_dir.clone() { Some(d) => d, - None => { - return Ok(CallToolResult { - content: vec![rmcp::model::Content::text( - "memory dir not configured; start server with --memory-dir ", - )], - structured_content: None, - is_error: Some(true), - meta: None, - }); - } + None => return Ok(tool_error(AppError::InvalidRequest( + "memory dir not configured; start server with --memory-dir ".into(), + ))), }; let res = tokio::task::spawn_blocking(move || load_memory(&dir, args.name.as_deref())) @@ -171,18 +147,7 @@ impl CodeMcpServer { .map_err(join_error)?; match res { - Ok(content) => { - let resp = tools::ToolResponse { - content, - truncated: false, - truncation_reason: None, - match_count: None, - entry_error_count: None, - search_error_count: None, - first_error: None, - }; - Ok(resp.into_call_tool_result()) - } + Ok(r) => Ok(r.into_call_tool_result()), Err(e) => Ok(tool_error(e)), } } @@ -191,9 +156,7 @@ impl CodeMcpServer { #[tool_handler] impl ServerHandler for CodeMcpServer { fn get_info(&self) -> rmcp::model::InitializeResult { - let mut instructions = String::from( - "code-mcp: filesystem search and read tools.\n\n", - ); + let mut instructions = String::from("code-mcp: filesystem search and read tools.\n\n"); let _ = write!( instructions, "All paths are scoped to the project root: {}. \ diff --git a/src/tools/cat.rs b/src/tools/cat.rs index 38e287a..028d4b3 100644 --- a/src/tools/cat.rs +++ b/src/tools/cat.rs @@ -4,7 +4,7 @@ use super::response::ToolResponse; use crate::error::AppError; use std::fs::File; use std::io::{BufRead, BufReader}; -use std::path::PathBuf; +use std::path::Path; /// Read file contents with pagination. /// @@ -17,19 +17,18 @@ use std::path::PathBuf; /// Returns [`AppError::InvalidRequest`] if the target is missing or not a /// regular file. pub fn cat( - file_path: &str, + file_path: &Path, offset: usize, max_lines: usize, max_bytes: usize, ) -> Result { - let path = PathBuf::from(file_path); - if !path.is_file() { + if !file_path.is_file() { return Err(AppError::InvalidRequest( "Target is not a file or does not exist".to_string(), )); } - let file = File::open(&path)?; + let file = File::open(file_path)?; let mut reader = BufReader::new(file); // Skip `offset` lines. @@ -38,16 +37,7 @@ pub fn cat( skip_buf.clear(); let n = reader.read_line(&mut skip_buf)?; if n == 0 { - // EOF before reaching the offset — nothing to return. - return Ok(ToolResponse { - content: String::new(), - truncated: false, - truncation_reason: None, - match_count: None, - entry_error_count: None, - search_error_count: None, - first_error: None, - }); + return Ok(ToolResponse::text(String::new())); } } @@ -102,7 +92,7 @@ pub fn cat( #[cfg(test)] mod tests { use super::*; - use crate::tools::testutil::{path_str, TestResult}; + use crate::tools::testutil::TestResult; use crate::tools::{DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES}; use std::fs; @@ -112,12 +102,16 @@ mod tests { let path = td.path().join("a.txt"); fs::write(&path, "L1\nL2\nL3\nL4\nL5\nL6\nL7\n")?; - let res = cat(path_str(&path)?, 2, 3, DEFAULT_MAX_BYTES)?; - assert!(res.content.starts_with("L3\nL4\nL5\n"), "got {:?}", res.content); + let res = cat(&path, 2, 3, DEFAULT_MAX_BYTES)?; + assert!( + res.content.starts_with("L3\nL4\nL5\n"), + "got {:?}", + res.content + ); assert!(res.truncated, "expected truncated=true"); assert_eq!(res.truncation_reason, Some("line_cap".to_string())); - let res = cat(path_str(&path)?, 4, 3, DEFAULT_MAX_BYTES)?; + let res = cat(&path, 4, 3, DEFAULT_MAX_BYTES)?; assert_eq!(res.content, "L5\nL6\nL7\n", "got {:?}", res.content); assert!(!res.truncated); Ok(()) @@ -130,17 +124,26 @@ mod tests { let body = "abcdefghijklmnopqrstuvwxyz\n".repeat(20); fs::write(&path, &body)?; - let res = cat(path_str(&path)?, 0, DEFAULT_MAX_LINES, 50)?; + let res = cat(&path, 0, DEFAULT_MAX_LINES, 50)?; assert!(res.truncated, "expected truncated=true, got {:?}", res); assert_eq!(res.truncation_reason, Some("byte_cap".to_string())); - assert!(res.content.len() < body.len(), "expected truncation, got len {}", res.content.len()); + assert!( + res.content.len() < body.len(), + "expected truncation, got len {}", + res.content.len() + ); Ok(()) } #[test] fn cat_errors_when_path_is_directory() -> TestResult { let td = tempfile::TempDir::new()?; - match cat(path_str(td.path())?, 0, DEFAULT_MAX_LINES, DEFAULT_MAX_BYTES) { + match cat( + td.path(), + 0, + DEFAULT_MAX_LINES, + DEFAULT_MAX_BYTES, + ) { Err(AppError::InvalidRequest(_)) => Ok(()), Err(other) => Err(format!("expected InvalidRequest, got {:?}", other).into()), Ok(s) => Err(format!("expected error, got Ok({:?})", s).into()), diff --git a/src/tools/common.rs b/src/tools/common.rs index 8ddbb9d..001ac05 100644 --- a/src/tools/common.rs +++ b/src/tools/common.rs @@ -1,11 +1,20 @@ //! Shared helpers for the walker-based tools: parallel walker construction, //! extension filtering, error capture, and byte-capped channel draining. -use super::options::GrepOptions; use ignore::WalkBuilder; use std::path::Path; -use std::sync::mpsc::Receiver; use std::sync::Mutex; +use std::sync::mpsc::Receiver; + +/// Trait for option structs that can configure a parallel directory walker. +/// +/// Both [`super::options::GrepOptions`] and [`super::options::FindOptions`] +/// implement this so [`build_parallel_walker`] is reusable across tools. +pub(crate) trait WalkerConfig { + fn include_hidden(&self) -> bool; + fn respect_gitignore(&self) -> bool; + fn follow_symlinks(&self) -> bool; +} /// Record the first error message into `slot`, ignoring later ones. pub(crate) fn record_first(slot: &Mutex>, msg: String) { @@ -16,17 +25,17 @@ pub(crate) fn record_first(slot: &Mutex>, msg: String) { } } -/// Build a parallel walker from the shared walker options in `GrepOptions`. +/// Build a parallel walker from any option struct implementing [`WalkerConfig`]. pub(crate) fn build_parallel_walker( - directory: &str, - opts: &GrepOptions, + directory: &Path, + opts: &impl WalkerConfig, ) -> ignore::WalkParallel { WalkBuilder::new(directory) - .hidden(!opts.include_hidden) - .git_ignore(opts.respect_gitignore) - .git_global(opts.respect_gitignore) - .git_exclude(opts.respect_gitignore) - .follow_links(opts.follow_symlinks) + .hidden(!opts.include_hidden()) + .git_ignore(opts.respect_gitignore()) + .git_global(opts.respect_gitignore()) + .git_exclude(opts.respect_gitignore()) + .follow_links(opts.follow_symlinks()) .build_parallel() } @@ -41,8 +50,6 @@ pub(crate) fn extension_matches(path: &Path, extensions: &[String]) -> bool { .is_some_and(|e| extensions.iter().any(|w| w == e)) } - - /// Drain string chunks from `rx` into a single output buffer, enforcing the /// authoritative `max_bytes` cap. When the cap is hit the final chunk is cut on /// a UTF-8 character boundary and a trailing newline is ensured; remaining diff --git a/src/tools/find.rs b/src/tools/find.rs index 8b7fb96..47e9095 100644 --- a/src/tools/find.rs +++ b/src/tools/find.rs @@ -1,12 +1,13 @@ //! The `find` tool: locate files by regex over a parallel `ignore` walker. -use super::common::record_first; +use super::common::{build_parallel_walker, record_first}; use super::options::FindOptions; use super::response::ToolResponse; use crate::error::AppError; -use ignore::{WalkBuilder, WalkState}; +use ignore::WalkState; use regex::Regex; use std::mem; +use std::path::Path; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::mpsc::channel; use std::sync::{Arc, Mutex}; @@ -17,23 +18,14 @@ use std::sync::{Arc, Mutex}; /// Uses a parallel `ignore` walker (gitignore-aware) and an `AtomicUsize` /// counter for exact `max_results` capping. Walker entry errors are tallied /// and surfaced in the returned [`ToolResponse`] metadata. -pub fn find( - directory: &str, - pattern: &str, - opts: FindOptions, -) -> Result { +pub fn find(directory: &Path, pattern: &str, opts: FindOptions) -> Result { let re = Regex::new(pattern)?; let max_results = opts.max_results; let count = Arc::new(AtomicUsize::new(0)); let entry_errors = Arc::new(AtomicUsize::new(0)); let first_error: Arc>> = Arc::new(Mutex::new(None)); - let walker = WalkBuilder::new(directory) - .hidden(!opts.include_hidden) - .git_ignore(opts.respect_gitignore) - .git_global(opts.respect_gitignore) - .git_exclude(opts.respect_gitignore) - .build_parallel(); + let walker = build_parallel_walker(directory, &opts); let (tx, rx) = channel::(); let match_basename = opts.match_basename; @@ -63,6 +55,10 @@ pub fn find( } }; + if !entry.file_type().is_some_and(|ft| ft.is_file()) { + return WalkState::Continue; + } + let path = entry.path(); let hay: std::borrow::Cow<'_, str> = if match_basename { match path.file_name() { @@ -99,10 +95,7 @@ pub fn find( } let entry_err_n = entry_errors.load(Ordering::Relaxed); - let first_error = first_error - .lock() - .ok() - .and_then(|g| g.clone()); + let first_error = first_error.lock().ok().and_then(|g| g.clone()); let match_count = count.load(Ordering::Relaxed); @@ -120,7 +113,7 @@ pub fn find( #[cfg(test)] mod tests { use super::*; - use crate::tools::testutil::{path_str, write_file, TestResult}; + use crate::tools::testutil::{TestResult, write_file}; #[test] fn find_match_basename_and_full_path() -> TestResult { @@ -130,7 +123,7 @@ mod tests { write_file(root, "sub/bar.rs", "")?; let basename = find( - path_str(root)?, + root, "^foo", FindOptions { match_basename: true, @@ -138,11 +131,19 @@ mod tests { ..Default::default() }, )?; - assert!(basename.content.contains("foo.rs"), "got {}", basename.content); - assert!(!basename.content.contains("bar.rs"), "got {}", basename.content); + assert!( + basename.content.contains("foo.rs"), + "got {}", + basename.content + ); + assert!( + !basename.content.contains("bar.rs"), + "got {}", + basename.content + ); let fullpath_anchored = find( - path_str(root)?, + root, "^foo", FindOptions { match_basename: false, @@ -158,7 +159,7 @@ mod tests { ); let fullpath_ok = find( - path_str(root)?, + root, r"sub.*foo\.rs$", FindOptions { match_basename: false, @@ -166,7 +167,11 @@ mod tests { ..Default::default() }, )?; - assert!(fullpath_ok.content.contains("foo.rs"), "got {}", fullpath_ok.content); + assert!( + fullpath_ok.content.contains("foo.rs"), + "got {}", + fullpath_ok.content + ); Ok(()) } } diff --git a/src/tools/grep.rs b/src/tools/grep.rs index 26f53ce..b5f7f1a 100644 --- a/src/tools/grep.rs +++ b/src/tools/grep.rs @@ -1,8 +1,6 @@ //! The `grep` tool: regex search across files via parallel directory traversal. -use super::common::{ - build_parallel_walker, drain_capped, extension_matches, record_first, -}; +use super::common::{build_parallel_walker, drain_capped, extension_matches, record_first}; use super::options::{GrepOptions, OutputMode}; use super::response::ToolResponse; use super::sinks::{CountSink, FileMatchSink, MatchSink}; @@ -14,7 +12,7 @@ use std::collections::HashMap; use std::mem; use std::path::Path; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::mpsc::{channel, Sender}; +use std::sync::mpsc::{Sender, channel}; use std::sync::{Arc, Mutex}; // ─── Shared error tracking ────────────────────────────────────────────────── @@ -77,11 +75,7 @@ impl ErrorState { /// Walker entry errors and per-file search errors are tallied and surfaced in /// the returned [`ToolResponse`] metadata rather than aborting the search. #[allow(clippy::needless_pass_by_value)] -pub fn grep( - directory: &str, - pattern: &str, - opts: GrepOptions, -) -> Result { +pub fn grep(directory: &Path, pattern: &str, opts: GrepOptions) -> Result { match opts.output_mode { OutputMode::Content => grep_streamed(directory, pattern, &opts, StreamMode::Content), OutputMode::FilesWithMatches => { @@ -103,7 +97,7 @@ enum StreamMode { /// Unified implementation for `content` and `files_with_matches` modes. /// Both use the mpsc pipeline with early-quit on `max_results`. fn grep_streamed( - directory: &str, + directory: &Path, pattern: &str, opts: &GrepOptions, mode: StreamMode, @@ -159,9 +153,7 @@ fn grep_streamed( max_results, max_bytes, }; - if let Err(err) = - local_searcher.search_path(&local_matcher, path, &mut sink) - { + if let Err(err) = local_searcher.search_path(&local_matcher, path, &mut sink) { errors.record_search_error(path, &err); } flush(&tx, &mut buf); @@ -172,9 +164,7 @@ fn grep_streamed( max_results, matched_this_file: false, }; - if let Err(err) = - local_searcher.search_path(&local_matcher, path, &mut sink) - { + if let Err(err) = local_searcher.search_path(&local_matcher, path, &mut sink) { errors.record_search_error(path, &err); } if sink.matched_this_file { @@ -225,7 +215,7 @@ fn flush(tx: &Sender, buf: &mut String) { /// `count` mode — tally matches per file, output as `path: N` lines. /// Does not use a channel; collects into a shared HashMap instead. fn grep_count( - directory: &str, + directory: &Path, pattern: &str, opts: &GrepOptions, ) -> Result { @@ -236,6 +226,7 @@ fn grep_count( let errors = ErrorState::new(); let extensions = opts.file_extensions.clone(); let file_counts: Arc>> = Arc::new(Mutex::new(HashMap::new())); + let file_count = Arc::new(AtomicUsize::new(0)); let walker = build_parallel_walker(directory, opts); walker.run(|| { @@ -244,8 +235,13 @@ fn grep_count( let local_matcher = matcher.clone(); let extensions = extensions.clone(); let file_counts = Arc::clone(&file_counts); + let file_count = Arc::clone(&file_count); Box::new(move |result| { + if file_count.load(Ordering::Relaxed) >= max_results { + return WalkState::Quit; + } + let entry = match result { Ok(e) => e, Err(err) => { @@ -268,6 +264,7 @@ fn grep_count( } if sink.count > 0 { + file_count.fetch_add(1, Ordering::Relaxed); let key = path.to_string_lossy().into_owned(); let mut map = file_counts .lock() @@ -275,7 +272,11 @@ fn grep_count( *map.entry(key).or_insert(0) += sink.count; } - WalkState::Continue + if file_count.load(Ordering::Relaxed) >= max_results { + WalkState::Quit + } else { + WalkState::Continue + } }) }); @@ -328,10 +329,7 @@ fn build_matcher(pattern: &str, opts: &GrepOptions) -> Result grep_searcher::Searcher { +fn build_searcher(opts: &GrepOptions, mode: &StreamMode) -> grep_searcher::Searcher { let mut builder = SearcherBuilder::new(); builder.binary_detection(BinaryDetection::quit(b'\x00')); @@ -353,7 +351,7 @@ fn build_searcher( #[cfg(test)] mod tests { use super::*; - use crate::tools::testutil::{path_str, write_file, TestResult}; + use crate::tools::testutil::{TestResult, write_file}; use std::fs; #[test] @@ -368,7 +366,7 @@ mod tests { respect_gitignore: false, ..Default::default() }; - let res = grep(path_str(root)?, "needle", opts)?; + let res = grep(root, "needle", opts)?; assert!( res.match_count.unwrap() <= 15, "expected match_count <= 15, got {:?}", @@ -385,7 +383,7 @@ mod tests { write_file(root, "a.txt", "Hello World\n")?; let case_sensitive = grep( - path_str(root)?, + root, "hello", GrepOptions { output_mode: OutputMode::Content, @@ -393,10 +391,15 @@ mod tests { ..Default::default() }, )?; - assert_eq!(case_sensitive.match_count, Some(0), "got {}", case_sensitive.content); + assert_eq!( + case_sensitive.match_count, + Some(0), + "got {}", + case_sensitive.content + ); let case_insensitive = grep( - path_str(root)?, + root, "hello", GrepOptions { case_insensitive: true, @@ -421,7 +424,7 @@ mod tests { write_file(root, "b.txt", "fn target() {}\n")?; let res = grep( - path_str(root)?, + root, "target", GrepOptions { file_extensions: vec!["rs".to_string()], @@ -444,25 +447,37 @@ mod tests { write_file(root, "open.txt", "needle\n")?; let respected = grep( - path_str(root)?, + root, "needle", GrepOptions { respect_gitignore: true, ..Default::default() }, )?; - assert!(!respected.content.contains("secrets.txt"), "got {}", respected.content); - assert!(respected.content.contains("open.txt"), "got {}", respected.content); + assert!( + !respected.content.contains("secrets.txt"), + "got {}", + respected.content + ); + assert!( + respected.content.contains("open.txt"), + "got {}", + respected.content + ); let ignored = grep( - path_str(root)?, + root, "needle", GrepOptions { respect_gitignore: false, ..Default::default() }, )?; - assert!(ignored.content.contains("secrets.txt"), "got {}", ignored.content); + assert!( + ignored.content.contains("secrets.txt"), + "got {}", + ignored.content + ); Ok(()) } @@ -475,7 +490,7 @@ mod tests { write_file(root, "c.rs", "needle here\n")?; let res = grep( - path_str(root)?, + root, "needle", GrepOptions { output_mode: OutputMode::FilesWithMatches, @@ -486,7 +501,11 @@ mod tests { assert!(res.content.contains("a.txt"), "got {}", res.content); assert!(!res.content.contains("b.txt"), "got {}", res.content); assert!(res.content.contains("c.rs"), "got {}", res.content); - assert!(!res.content.contains("1:"), "should not have line numbers: {}", res.content); + assert!( + !res.content.contains("1:"), + "should not have line numbers: {}", + res.content + ); assert_eq!(res.match_count, Some(2), "got {:?}", res.match_count); Ok(()) } @@ -500,7 +519,7 @@ mod tests { } let res = grep( - path_str(root)?, + root, "needle", GrepOptions { output_mode: OutputMode::FilesWithMatches, @@ -526,7 +545,7 @@ mod tests { write_file(root, "c.rs", "needle here\n")?; let res = grep( - path_str(root)?, + root, "needle", GrepOptions { output_mode: OutputMode::Count, diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 787c110..e808856 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -27,7 +27,6 @@ pub use response::ToolResponse; pub(crate) const DEFAULT_MAX_BYTES: usize = 5 * 1024 * 1024; // 5 MiB pub(crate) const DEFAULT_MAX_RESULTS: usize = 100; -#[cfg(test)] pub(crate) const DEFAULT_MAX_LINES: usize = 2000; /// Shared test helpers used across the per-tool test modules. @@ -48,8 +47,4 @@ pub(crate) mod testutil { f.write_all(contents.as_bytes())?; Ok(()) } - - pub(crate) fn path_str(p: &Path) -> Result<&str, Box> { - p.to_str().ok_or_else(|| "non-utf8 path".into()) - } } diff --git a/src/tools/options.rs b/src/tools/options.rs index a79c627..a542010 100644 --- a/src/tools/options.rs +++ b/src/tools/options.rs @@ -1,12 +1,15 @@ //! Configuration types for the `grep` and `find` tools. use super::{DEFAULT_MAX_BYTES, DEFAULT_MAX_RESULTS}; -use crate::error::AppError; +use rmcp::schemars::{self, JsonSchema}; +use serde::{Deserialize, Serialize}; /// Controls what the `grep` tool emits for each match. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize, Serialize, JsonSchema)] +#[serde(rename_all = "snake_case")] pub enum OutputMode { /// Emit the file path on the first match, then skip the rest of that file. + #[default] FilesWithMatches, /// Emit matching lines with line numbers (the original/default behaviour). Content, @@ -14,20 +17,6 @@ pub enum OutputMode { Count, } -impl OutputMode { - /// Parse a string into an `OutputMode`, returning an error for unknown values. - pub fn from_str_lossy(s: &str) -> Result { - match s { - "files_with_matches" => Ok(Self::FilesWithMatches), - "content" => Ok(Self::Content), - "count" => Ok(Self::Count), - other => Err(AppError::InvalidRequest(format!( - "unknown output_mode '{other}'; expected one of: files_with_matches, content, count" - ))), - } - } -} - /// Configuration for the `grep` tool. The boolean fields are independent /// search/walker toggles; grouping them into enums would obscure the /// (flat) JSON contract exposed to MCP clients. @@ -73,6 +62,18 @@ impl Default for GrepOptions { } } +impl super::common::WalkerConfig for GrepOptions { + fn include_hidden(&self) -> bool { + self.include_hidden + } + fn respect_gitignore(&self) -> bool { + self.respect_gitignore + } + fn follow_symlinks(&self) -> bool { + self.follow_symlinks + } +} + /// Configuration for the `find` tool. #[derive(Clone, Copy)] pub struct FindOptions { @@ -98,6 +99,18 @@ impl Default for FindOptions { } } +impl super::common::WalkerConfig for FindOptions { + fn include_hidden(&self) -> bool { + self.include_hidden + } + fn respect_gitignore(&self) -> bool { + self.respect_gitignore + } + fn follow_symlinks(&self) -> bool { + false // find does not expose follow_symlinks + } +} + #[cfg(test)] mod tests { use super::*; @@ -106,12 +119,13 @@ mod tests { #[test] fn grep_output_mode_rejects_unknown() -> TestResult { - match OutputMode::from_str_lossy("bogus") { - Err(AppError::InvalidRequest(msg)) => { - assert!(msg.contains("bogus"), "got: {}", msg); - Ok(()) - } - other => Err(format!("expected InvalidRequest, got {:?}", other).into()), - } + let result: Result = serde_json::from_str(r#""bogus""#); + assert!(result.is_err(), "expected deserialization error for unknown output_mode"); + Ok(()) + } + + #[test] + fn grep_output_mode_default_is_files_with_matches() { + assert_eq!(OutputMode::default(), OutputMode::FilesWithMatches); } } diff --git a/src/tools/response.rs b/src/tools/response.rs index 701b721..65aee4b 100644 --- a/src/tools/response.rs +++ b/src/tools/response.rs @@ -29,6 +29,21 @@ pub struct ToolResponse { } impl ToolResponse { + /// Build a minimal `ToolResponse` carrying just text content, with no + /// truncation or match metadata. Used by tools (like `memories`) whose + /// output is a single opaque string with no associated search metrics. + pub fn text(content: String) -> Self { + Self { + content, + truncated: false, + truncation_reason: None, + match_count: None, + entry_error_count: None, + search_error_count: None, + first_error: None, + } + } + /// Build a `CallToolResult` from this response: text content goes into /// `content`, and the structured metadata goes into `structured_content`. /// @@ -36,7 +51,6 @@ impl ToolResponse { /// prefer structured output still get the actual content. pub fn into_call_tool_result(self) -> CallToolResult { let structured = json!({ - "content": self.content, "truncated": self.truncated, "truncation_reason": self.truncation_reason, "match_count": self.match_count, diff --git a/src/tools/sinks.rs b/src/tools/sinks.rs index ee9d517..4d36f8b 100644 --- a/src/tools/sinks.rs +++ b/src/tools/sinks.rs @@ -23,11 +23,7 @@ pub(crate) struct MatchSink<'a> { impl Sink for MatchSink<'_> { type Error = io::Error; - fn matched( - &mut self, - _searcher: &Searcher, - mat: &SinkMatch<'_>, - ) -> Result { + fn matched(&mut self, _searcher: &Searcher, mat: &SinkMatch<'_>) -> Result { // Increment first; if we are over the cap, undo conceptually by stopping. let prev = self.count.fetch_add(1, Ordering::Relaxed); if prev >= self.max_results { @@ -45,11 +41,7 @@ impl Sink for MatchSink<'_> { Ok(true) } - fn context( - &mut self, - _searcher: &Searcher, - ctx: &SinkContext<'_>, - ) -> Result { + fn context(&mut self, _searcher: &Searcher, ctx: &SinkContext<'_>) -> Result { // Context lines do not count toward `max_results` (the cap is on // matches, not surrounding lines), but we still respect the byte cap. if self.buf.len() >= self.max_bytes { @@ -86,11 +78,7 @@ pub(crate) struct FileMatchSink<'a> { impl Sink for FileMatchSink<'_> { type Error = io::Error; - fn matched( - &mut self, - _searcher: &Searcher, - _mat: &SinkMatch<'_>, - ) -> Result { + fn matched(&mut self, _searcher: &Searcher, _mat: &SinkMatch<'_>) -> Result { if self.matched_this_file { // Already recorded this file; stop searching it. return Ok(false); @@ -104,11 +92,7 @@ impl Sink for FileMatchSink<'_> { Ok(false) } - fn context( - &mut self, - _searcher: &Searcher, - _ctx: &SinkContext<'_>, - ) -> Result { + fn context(&mut self, _searcher: &Searcher, _ctx: &SinkContext<'_>) -> Result { // No context needed for files_with_matches mode. Ok(true) } @@ -123,20 +107,12 @@ pub(crate) struct CountSink { impl Sink for CountSink { type Error = io::Error; - fn matched( - &mut self, - _searcher: &Searcher, - _mat: &SinkMatch<'_>, - ) -> Result { + fn matched(&mut self, _searcher: &Searcher, _mat: &SinkMatch<'_>) -> Result { self.count += 1; Ok(true) } - fn context( - &mut self, - _searcher: &Searcher, - _ctx: &SinkContext<'_>, - ) -> Result { + fn context(&mut self, _searcher: &Searcher, _ctx: &SinkContext<'_>) -> Result { Ok(true) } }