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        }
35    }
36
37    async fn execute(&self, params: Value) -> Result<String> {
38        let pattern = params["pattern"]
39            .as_str()
40            .ok_or_else(|| anyhow::anyhow!("missing 'pattern'"))?;
41        let path = params["path"].as_str().unwrap_or(".");
42        let glob_pattern = params["glob"].as_str();
43
44        let pattern = pattern.to_string();
45        let path = path.to_string();
46        let glob_pattern = glob_pattern.map(|s| s.to_string());
47
48        // Use timeout to prevent hanging on large directories
49        timeout(Duration::from_secs(30), async {
50            tokio::task::spawn_blocking(move || {
51                search_files(&pattern, &path, glob_pattern.as_deref())
52            })
53            .await?
54        })
55        .await
56        .map_err(|_| anyhow::anyhow!("Search timeout (30s) - directory may be too large"))?
57    }
58}
59
60/// Maximum files to search before stopping.
61const MAX_FILES: usize = 500;
62
63fn search_files(pattern: &str, path: &str, glob_pattern: Option<&str>) -> Result<String> {
64    use std::fs;
65    use std::path::Path;
66
67    let regex = regex::Regex::new(pattern)?;
68    let mut results = Vec::new();
69    let root = Path::new(path);
70
71    let entries = collect_files(root, glob_pattern)?;
72
73    for file_path in entries {
74        // Skip very large files (> 1MB)
75        match fs::metadata(&file_path) {
76            Ok(meta) if meta.len() > 1_000_000 => continue,
77            Err(_) => continue,
78            Ok(_) => {}
79        }
80
81        let content = match fs::read_to_string(&file_path) {
82            Ok(c) => c,
83            Err(_) => continue,
84        };
85
86        for (line_num, line) in content.lines().enumerate() {
87            if regex.is_match(line) {
88                results.push(format!(
89                    "{}:{}: {}",
90                    file_path.display(),
91                    line_num + 1,
92                    line.trim()
93                ));
94            }
95        }
96
97        if results.len() > 200 {
98            results.push("... (truncated, too many results)".to_string());
99            break;
100        }
101    }
102
103    if results.is_empty() {
104        Ok("No matches found.".to_string())
105    } else {
106        Ok(results.join("\n"))
107    }
108}
109
110fn collect_files(
111    root: &std::path::Path,
112    glob_pattern: Option<&str>,
113) -> Result<Vec<std::path::PathBuf>> {
114    let mut files = Vec::new();
115
116    if root.is_file() {
117        files.push(root.to_path_buf());
118        return Ok(files);
119    }
120
121    let glob_matcher = glob_pattern.map(glob::Pattern::new).transpose()?;
122
123    let mut stack = vec![root.to_path_buf()];
124
125    while let Some(dir) = stack.pop() {
126        let entries = match std::fs::read_dir(&dir) {
127            Ok(e) => e,
128            Err(_) => continue,
129        };
130
131        for entry in entries.flatten() {
132            let path = entry.path();
133            let name = entry.file_name();
134            let name_str = name.to_string_lossy();
135
136            // Skip hidden dirs and common large directories
137            if name_str.starts_with('.')
138                || name_str == "node_modules"
139                || name_str == "target"
140                || name_str == "dist"
141                || name_str == "build"
142                || name_str == ".git"
143            {
144                continue;
145            }
146
147            // Check glob pattern for files
148            if let Some(ref matcher) = glob_matcher
149                && path.is_file()
150                && let Some(name) = path.file_name().and_then(|n| n.to_str())
151                && !matcher.matches(name)
152            {
153                continue;
154            }
155
156            if path.is_dir() {
157                stack.push(path);
158            } else if path.is_file() {
159                files.push(path);
160                // Limit number of files to search
161                if files.len() >= MAX_FILES {
162                    return Ok(files);
163                }
164            }
165        }
166    }
167
168    Ok(files)
169}