use std::path::Path; use std::process::Stdio; use async_trait::async_trait; use serde_json::json; use tokio::process::Command; use tokio::time::timeout; use crate::tools::traits::{Tool, ToolResult}; const MAX_RESULTS: usize = 100; const MAX_OUTPUT_CHARS: usize = 50_000; const TIMEOUT_SECS: u64 = 60; pub struct ContentSearchTool; impl ContentSearchTool { pub fn new() -> Self { Self } fn resolve_dir(&self, dir: Option<&str>) -> String { match dir { Some(d) if !d.is_empty() => d.to_string(), _ => ".".to_string(), } } fn truncate_output(&self, lines: &[String]) -> String { let mut output = String::new(); for (i, line) in lines.iter().enumerate() { if output.len() + line.len() + 1 > MAX_OUTPUT_CHARS { let omitted = lines.len() - i; output.push_str(&format!("\n... ({} matches omitted) ...", omitted)); break; } if !output.is_empty() { output.push('\n'); } output.push_str(line); } output } } impl Default for ContentSearchTool { fn default() -> Self { Self::new() } } #[async_trait] impl Tool for ContentSearchTool { fn name(&self) -> &str { "content_search" } fn description(&self) -> &str { "Search file contents by regex or text pattern. Internally uses ripgrep (rg) for fast searching — if rg is not available, falls back to grep, then pure Rust. Supports context lines, file filtering, and case-sensitivity options." } fn parameters_schema(&self) -> serde_json::Value { json!({ "type": "object", "properties": { "pattern": { "type": "string", "description": "Regex or text pattern to search for in file contents" }, "dir": { "type": "string", "description": "Directory to search in (default: current working directory)" }, "file_pattern": { "type": "string", "description": "Optional glob to restrict which files to search (e.g. '*.rs', '*.{rs,toml}')" }, "case_sensitive": { "type": "boolean", "description": "Whether to match case-sensitively (default: false)" }, "context_lines": { "type": "integer", "description": "Number of context lines to show before and after each match (default: 0)" }, "max_results": { "type": "integer", "description": "Maximum number of matching lines to return (default: 100)" } }, "required": ["pattern"] }) } fn read_only(&self) -> bool { true } async fn execute(&self, args: serde_json::Value) -> anyhow::Result { let pattern = match args.get("pattern").and_then(|v| v.as_str()) { Some(p) if !p.is_empty() => p, _ => { return Ok(ToolResult { success: false, output: String::new(), error: Some("Missing required parameter: pattern".to_string()), }); } }; let dir = self.resolve_dir(args.get("dir").and_then(|v| v.as_str())); let file_pattern = args.get("file_pattern").and_then(|v| v.as_str()); let case_sensitive = args .get("case_sensitive") .and_then(|v| v.as_bool()) .unwrap_or(false); let context_lines = args .get("context_lines") .and_then(|v| v.as_u64()) .unwrap_or(0) as usize; let max_results = args .get("max_results") .and_then(|v| v.as_u64()) .unwrap_or(MAX_RESULTS as u64) as usize; let result = self .run_search( pattern, &dir, file_pattern, case_sensitive, context_lines, max_results, ) .await; match result { Ok(lines) => { let count = lines.len(); let mut output = self.truncate_output(&lines); output.push_str(&format!("\n\n---\n共 {} 条匹配", count)); Ok(ToolResult { success: true, output, error: None, }) } Err(e) => Ok(ToolResult { success: false, output: String::new(), error: Some(e.to_string()), }), } } } impl ContentSearchTool { async fn run_search( &self, pattern: &str, dir: &str, file_pattern: Option<&str>, case_sensitive: bool, context_lines: usize, max_results: usize, ) -> anyhow::Result> { if which::which("rg").is_ok() { match self .search_with_rg( pattern, dir, file_pattern, case_sensitive, context_lines, max_results, ) .await { Ok(lines) => return Ok(lines), Err(e) => tracing::warn!("rg failed: {}, falling back", e), } } if which::which("grep").is_ok() { match self .search_with_grep( pattern, dir, file_pattern, case_sensitive, context_lines, max_results, ) .await { Ok(lines) if !lines.is_empty() => return Ok(lines), Ok(_) => {} Err(e) => tracing::warn!("grep failed: {}, falling back", e), } } tracing::warn!( "No rg/grep available, using built-in content search (much slower). Install ripgrep for better performance." ); self.search_with_rust( pattern, dir, file_pattern, case_sensitive, context_lines, max_results, ) .await } async fn search_with_rg( &self, pattern: &str, dir: &str, file_pattern: Option<&str>, case_sensitive: bool, context_lines: usize, max_results: usize, ) -> anyhow::Result> { let mut cmd = Command::new("rg"); cmd.arg("-n") .arg("--no-heading") .arg("--color") .arg("never") .arg("--max-count") .arg(max_results.to_string()) .arg(pattern) .arg(dir) .stdout(Stdio::piped()) .stderr(Stdio::piped()); if !case_sensitive { cmd.arg("-i"); } if context_lines > 0 { cmd.arg("-C").arg(context_lines.to_string()); } if let Some(fp) = file_pattern { cmd.arg("--glob").arg(fp); } let output = timeout(std::time::Duration::from_secs(TIMEOUT_SECS), cmd.output()) .await .map_err(|_| anyhow::anyhow!("rg timed out after {}s", TIMEOUT_SECS))??; if !output.status.success() && output.status.code() != Some(1) { let stderr = String::from_utf8_lossy(&output.stderr); return Err(anyhow::anyhow!("rg error: {}", stderr.trim())); } let text = String::from_utf8_lossy(&output.stdout); let lines: Vec = text .lines() .take(max_results) .map(|l| l.to_string()) .collect(); Ok(lines) } async fn search_with_grep( &self, pattern: &str, dir: &str, file_pattern: Option<&str>, case_sensitive: bool, context_lines: usize, max_results: usize, ) -> anyhow::Result> { let mut cmd = Command::new("grep"); cmd.arg("-rn") .arg("-E") .arg("--color=never") .arg("--binary-files=without-match") .arg(pattern) .arg(dir) .stdout(Stdio::piped()) .stderr(Stdio::piped()); if !case_sensitive { cmd.arg("-i"); } if context_lines > 0 { cmd.arg("-C").arg(context_lines.to_string()); } if let Some(fp) = file_pattern { cmd.arg("--include").arg(fp); } let output = timeout(std::time::Duration::from_secs(TIMEOUT_SECS), cmd.output()) .await .map_err(|_| anyhow::anyhow!("grep timed out after {}s", TIMEOUT_SECS))??; let text = String::from_utf8_lossy(&output.stdout); let lines: Vec = text .lines() .take(max_results) .map(|l| l.to_string()) .collect(); Ok(lines) } async fn search_with_rust( &self, pattern: &str, dir: &str, file_pattern: Option<&str>, case_sensitive: bool, _context_lines: usize, max_results: usize, ) -> anyhow::Result> { let re = if case_sensitive { regex::Regex::new(pattern) } else { regex::RegexBuilder::new(pattern) .case_insensitive(true) .build() } .map_err(|e| anyhow::anyhow!("Invalid regex pattern '{}': {}", pattern, e))?; let file_re = file_pattern.map(|fp| { let re_str = glob_to_regex(fp); if case_sensitive { regex::Regex::new(&re_str) } else { regex::RegexBuilder::new(&re_str) .case_insensitive(true) .build() } }); let file_re = match file_re { Some(Ok(r)) => Some(r), Some(Err(e)) => return Err(anyhow::anyhow!("Invalid file pattern: {}", e)), None => None, }; let mut results = Vec::new(); grep_dir( Path::new(dir), Path::new(dir), &re, file_re.as_ref(), &mut results, max_results, )?; Ok(results) } } fn glob_to_regex(glob: &str) -> String { let mut regex = String::from("^"); let chars: Vec = glob.chars().collect(); let mut i = 0; while i < chars.len() { match chars[i] { '*' => { if i + 1 < chars.len() && chars[i + 1] == '*' { regex.push_str(".*"); i += 1; } else { regex.push_str("[^/]*"); } } '?' => regex.push_str("[^/]"), '.' | '+' | '(' | ')' | '[' | ']' | '{' | '}' | '^' | '$' | '|' | '\\' => { regex.push('\\'); regex.push(chars[i]); } c => regex.push(c), } i += 1; } regex.push('$'); regex } fn grep_dir( base: &Path, current: &Path, re: ®ex::Regex, file_re: Option<®ex::Regex>, results: &mut Vec, max: usize, ) -> anyhow::Result<()> { if results.len() >= max { return Ok(()); } let entries = match std::fs::read_dir(current) { Ok(e) => e, Err(_) => return Ok(()), }; for entry in entries.flatten() { let path = entry.path(); let rel = match path.strip_prefix(base) { Ok(r) => r, Err(_) => continue, }; if path.is_dir() { if let Some(name) = rel.file_name().and_then(|n| n.to_str()) && name.starts_with('.') && name.len() > 1 { continue; } grep_dir(base, &path, re, file_re, results, max)?; } else if path.is_file() { if let Some(file_re) = file_re && let Some(name) = rel.file_name().and_then(|n| n.to_str()) && !file_re.is_match(name) { continue; } if let Ok(content) = std::fs::read_to_string(&path) { for (line_num, line) in content.lines().enumerate() { if re.is_match(line) { results.push(format!( "{}:{}:{}", rel.to_string_lossy(), line_num + 1, line )); if results.len() >= max { return Ok(()); } } } } } } Ok(()) } #[cfg(test)] mod tests { use super::*; use std::fs; use tempfile::TempDir; #[tokio::test] async fn test_content_search_rust_fallback() { let dir = TempDir::new().unwrap(); fs::write( dir.path().join("main.rs"), "fn main() {\n let x = 42;\n println!(\"hello\");\n}", ) .unwrap(); fs::write( dir.path().join("lib.rs"), "pub fn foo() -> u32 {\n let y = 42;\n y\n}", ) .unwrap(); fs::write(dir.path().join("README.md"), "# Project\nHello world").unwrap(); let tool = ContentSearchTool::new(); let result = tool .execute(json!({ "pattern": "let.*=.*42", "dir": dir.path().to_str().unwrap() })) .await .unwrap(); assert!(result.success); assert!(result.output.contains("main.rs")); assert!(result.output.contains("lib.rs")); assert!(!result.output.contains("README.md")); assert!(result.output.contains("共 2 条匹配")); } #[tokio::test] async fn test_content_search_file_filter() { let dir = TempDir::new().unwrap(); fs::write(dir.path().join("main.rs"), "fn main() {}").unwrap(); fs::write(dir.path().join("config.toml"), "name = \"test\"").unwrap(); let tool = ContentSearchTool::new(); let result = tool .execute(json!({ "pattern": "test", "dir": dir.path().to_str().unwrap(), "file_pattern": "*.toml" })) .await .unwrap(); assert!(result.success); assert!(result.output.contains("config.toml")); assert!(!result.output.contains("main.rs")); } #[tokio::test] async fn test_content_search_max_results() { let dir = TempDir::new().unwrap(); let mut content = String::new(); for i in 0..10 { content.push_str(&format!("match line {}\n", i)); } fs::write(dir.path().join("data.txt"), &content).unwrap(); let tool = ContentSearchTool::new(); let result = tool .execute(json!({ "pattern": "match line", "dir": dir.path().to_str().unwrap(), "max_results": 3 })) .await .unwrap(); assert!(result.success); assert!(result.output.contains("共 3 条匹配")); } }