Skip to main content

matrixcode_core/tools/
search.rs

1use anyhow::Result;
2use async_trait::async_trait;
3use serde_json::{Value, json};
4use tokio::time::{Duration, timeout};
5
6use super::{Tool, ToolDefinition};
7
8pub struct SearchTool;
9
10#[async_trait]
11impl Tool for SearchTool {
12    fn definition(&self) -> ToolDefinition {
13        ToolDefinition {
14            name: "search".to_string(),
15            description: "在文件中搜索模式,类似 grep 功能".to_string(),
16            parameters: json!({
17                "type": "object",
18                "properties": {
19                    "pattern": {
20                        "type": "string",
21                        "description": "要搜索的正则表达式模式"
22                    },
23                    "path": {
24                        "type": "string",
25                        "description": "搜索的目录或文件路径(默认 '.')"
26                    },
27                    "glob": {
28                        "type": "string",
29                        "description": "文件过滤的 glob 模式(如 '*.rs')"
30                    }
31                },
32                "required": ["pattern"]
33            }),
34            ..Default::default()
35        }
36    }
37
38    async fn execute(&self, params: Value) -> Result<String> {
39        let pattern = params["pattern"]
40            .as_str()
41            .ok_or_else(|| anyhow::anyhow!("missing 'pattern'"))?;
42        let path = params["path"].as_str().unwrap_or(".");
43        let glob_pattern = params["glob"].as_str();
44
45        let pattern = pattern.to_string();
46        let path = path.to_string();
47        let glob_pattern = glob_pattern.map(|s| s.to_string());
48
49        // Use timeout to prevent hanging on large directories
50        timeout(Duration::from_secs(30), async {
51            tokio::task::spawn_blocking(move || {
52                search_files(&pattern, &path, glob_pattern.as_deref())
53            })
54            .await?
55        })
56        .await
57        .map_err(|_| anyhow::anyhow!("Search timeout (30s) - directory may be too large"))?
58    }
59}
60
61/// Maximum files to search before stopping.
62const MAX_FILES: usize = 500;
63
64fn search_files(pattern: &str, path: &str, glob_pattern: Option<&str>) -> Result<String> {
65    use std::fs;
66    use std::path::Path;
67
68    let regex = regex::Regex::new(pattern)?;
69    let mut results = Vec::new();
70    let root = Path::new(path);
71
72    let entries = collect_files(root, glob_pattern)?;
73
74    for file_path in entries {
75        // Skip very large files (> 1MB)
76        match fs::metadata(&file_path) {
77            Ok(meta) if meta.len() > 1_000_000 => continue,
78            Err(_) => continue,
79            Ok(_) => {}
80        }
81
82        let content = match fs::read_to_string(&file_path) {
83            Ok(c) => c,
84            Err(_) => continue,
85        };
86
87        for (line_num, line) in content.lines().enumerate() {
88            if regex.is_match(line) {
89                results.push(format!(
90                    "{}:{}: {}",
91                    file_path.display(),
92                    line_num + 1,
93                    line.trim()
94                ));
95            }
96        }
97
98        if results.len() > 200 {
99            results.push("... (truncated, too many results)".to_string());
100            break;
101        }
102    }
103
104    if results.is_empty() {
105        Ok("No matches found.".to_string())
106    } else {
107        Ok(results.join("\n"))
108    }
109}
110
111fn collect_files(
112    root: &std::path::Path,
113    glob_pattern: Option<&str>,
114) -> Result<Vec<std::path::PathBuf>> {
115    let mut files = Vec::new();
116
117    if root.is_file() {
118        files.push(root.to_path_buf());
119        return Ok(files);
120    }
121
122    let glob_matcher = glob_pattern.map(glob::Pattern::new).transpose()?;
123
124    let mut stack = vec![root.to_path_buf()];
125
126    while let Some(dir) = stack.pop() {
127        let entries = match std::fs::read_dir(&dir) {
128            Ok(e) => e,
129            Err(_) => continue,
130        };
131
132        for entry in entries.flatten() {
133            let path = entry.path();
134            let name = entry.file_name();
135            let name_str = name.to_string_lossy();
136
137            // Skip hidden dirs and common large directories
138            if name_str.starts_with('.')
139                || name_str == "node_modules"
140                || name_str == "target"
141                || name_str == "dist"
142                || name_str == "build"
143                || name_str == ".git"
144            {
145                continue;
146            }
147
148            // Check glob pattern for files
149            if let Some(ref matcher) = glob_matcher
150                && path.is_file()
151                && let Some(name) = path.file_name().and_then(|n| n.to_str())
152                && !matcher.matches(name)
153            {
154                continue;
155            }
156
157            if path.is_dir() {
158                stack.push(path);
159            } else if path.is_file() {
160                files.push(path);
161                // Limit number of files to search
162                if files.len() >= MAX_FILES {
163                    return Ok(files);
164                }
165            }
166        }
167    }
168
169    Ok(files)
170}