yapi_core/tools/
grep.rs

1//! `grep`: ripgrep over files, with context lines and pi's output limits.
2
3use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5use std::process::Stdio;
6
7use futures_util::future::BoxFuture;
8use serde_json::{Map, Value, json};
9use tokio::io::AsyncBufReadExt;
10use tokio_util::sync::CancellationToken;
11use yapi_agent::{Tool, UpdateSink};
12use yapi_types::event::ToolResult;
13use yapi_types::message::ToolDeclaration;
14
15use super::external::{ExternalTool, ensure_tool};
16use super::path::{relative, resolve_to_cwd};
17use super::truncate::{DEFAULT_MAX_BYTES, GREP_MAX_LINE_LENGTH, truncate_line};
18use super::{ToolEnv, capped_output, check_abort, js_number, plain_declaration, text_result};
19
20const DEFAULT_LIMIT: f64 = 100.0;
21
22/// The `grep` tool.
23pub struct Grep {
24    env: ToolEnv,
25    declaration: ToolDeclaration,
26}
27
28impl Grep {
29    /// A `grep` tool for `env`.
30    pub fn new(env: ToolEnv) -> Grep {
31        Grep {
32            env,
33            declaration: plain_declaration(
34                "grep",
35                format!(
36                    "Search file contents for a pattern. Returns matching lines with file paths and line numbers. Respects .gitignore. Output is truncated to {} matches or {}KB (whichever is hit first). Long lines are truncated to {GREP_MAX_LINE_LENGTH} chars.",
37                    DEFAULT_LIMIT,
38                    DEFAULT_MAX_BYTES / 1024
39                ),
40                json!({"type":"object","required":["pattern"],"properties":{
41                    "pattern":{"type":"string","description":"Search pattern (regex or literal string)"},
42                    "path":{"type":"string","description":"Directory or file to search (default: current directory)"},
43                    "glob":{"type":"string","description":"Filter files by glob pattern, e.g. '*.ts' or '**/*.spec.ts'"},
44                    "ignoreCase":{"type":"boolean","description":"Case-insensitive search (default: false)"},
45                    "literal":{"type":"boolean","description":"Treat pattern as literal string instead of regex (default: false)"},
46                    "context":{"type":"number","description":"Number of lines to show before and after each match (default: 0)"},
47                    "limit":{"type":"number","description":"Maximum number of matches to return (default: 100)"}}}),
48            ),
49        }
50    }
51}
52
53impl Tool for Grep {
54    fn declaration(&self) -> &ToolDeclaration {
55        &self.declaration
56    }
57
58    fn execute(
59        &self,
60        _call_id: String,
61        args: Value,
62        cancel: CancellationToken,
63        _updates: UpdateSink,
64    ) -> BoxFuture<'_, Result<ToolResult, String>> {
65        Box::pin(async move {
66            check_abort(&cancel)?;
67            self.search(args, cancel).await
68        })
69    }
70}
71
72struct Match {
73    file: String,
74    line: u64,
75    text: Option<String>,
76}
77
78/// Lines of a file with `\r\n` and `\r` as line ends, empty when unreadable.
79fn file_lines(path: &Path) -> Vec<String> {
80    match std::fs::read(path) {
81        Ok(bytes) => String::from_utf8_lossy(&bytes)
82            .replace("\r\n", "\n")
83            .replace('\r', "\n")
84            .split('\n')
85            .map(str::to_owned)
86            .collect(),
87        Err(_) => Vec::new(),
88    }
89}
90
91impl Grep {
92    async fn search(&self, args: Value, cancel: CancellationToken) -> Result<ToolResult, String> {
93        let Some(rg) = ensure_tool(ExternalTool::Rg, &self.env.bin_dir).await else {
94            return Err("ripgrep (rg) is not available and could not be downloaded".into());
95        };
96        let pattern = args["pattern"].as_str().unwrap_or_default();
97        let dir = args["path"]
98            .as_str()
99            .filter(|p| !p.is_empty())
100            .unwrap_or(".");
101        let search = resolve_to_cwd(dir, &self.env.cwd);
102        let is_directory = match std::fs::metadata(&search) {
103            Ok(metadata) => metadata.is_dir(),
104            Err(_) => return Err(format!("Path not found: {}", search.display())),
105        };
106        let context = args["context"].as_f64().filter(|c| *c > 0.0).unwrap_or(0.0) as u64;
107        let limit = args["limit"].as_f64().unwrap_or(DEFAULT_LIMIT).max(1.0);
108        let format_path = |file: &str| -> String {
109            let file = Path::new(file);
110            if is_directory {
111                let relative = relative(&search, file);
112                if !relative.is_empty() && !relative.starts_with("..") {
113                    return relative.replace('\\', "/");
114                }
115            }
116            file.file_name()
117                .map(|name| name.to_string_lossy().into_owned())
118                .unwrap_or_default()
119        };
120
121        let mut command = tokio::process::Command::new(&rg);
122        command.args(["--json", "--line-number", "--color=never", "--hidden"]);
123        if args["ignoreCase"] == true {
124            command.arg("--ignore-case");
125        }
126        if args["literal"] == true {
127            command.arg("--fixed-strings");
128        }
129        if let Some(glob) = args["glob"].as_str().filter(|g| !g.is_empty()) {
130            command.args(["--glob", glob]);
131        }
132        command
133            .arg("--")
134            .arg(pattern)
135            .arg(&search)
136            .stdin(Stdio::null())
137            .stdout(Stdio::piped())
138            .stderr(Stdio::piped())
139            .kill_on_drop(true);
140        let mut child = command
141            .spawn()
142            .map_err(|err| format!("Failed to run ripgrep: {err}"))?;
143        let stdout = child.stdout.take();
144        let stderr = child.stderr.take();
145        let stderr_task = tokio::spawn(async move {
146            let mut text = String::new();
147            if let Some(mut stderr) = stderr {
148                use tokio::io::AsyncReadExt;
149                let mut bytes = Vec::new();
150                let _ = stderr.read_to_end(&mut bytes).await;
151                text = String::from_utf8_lossy(&bytes).into_owned();
152            }
153            text
154        });
155
156        let mut matches: Vec<Match> = Vec::new();
157        let mut limit_reached = false;
158        if let Some(stdout) = stdout {
159            let mut lines = tokio::io::BufReader::new(stdout).lines();
160            loop {
161                let line = tokio::select! {
162                    () = cancel.cancelled() => {
163                        let _ = child.kill().await;
164                        return Err("Operation aborted".into());
165                    }
166                    line = lines.next_line() => line,
167                };
168                let Ok(Some(line)) = line else { break };
169                if line.trim().is_empty() || matches.len() as f64 >= limit {
170                    continue;
171                }
172                let Ok(event) = serde_json::from_str::<Value>(&line) else {
173                    continue;
174                };
175                if event["type"] != "match" {
176                    continue;
177                }
178                let data = &event["data"];
179                if let (Some(file), Some(number)) =
180                    (data["path"]["text"].as_str(), data["line_number"].as_u64())
181                {
182                    matches.push(Match {
183                        file: file.to_owned(),
184                        line: number,
185                        text: data["lines"]["text"].as_str().map(str::to_owned),
186                    });
187                }
188                if matches.len() as f64 >= limit {
189                    limit_reached = true;
190                    let _ = child.start_kill();
191                }
192            }
193        }
194        let status = child.wait().await;
195        let stderr = stderr_task.await.unwrap_or_default();
196        check_abort(&cancel)?;
197        if !limit_reached {
198            let code = status.ok().and_then(|status| status.code());
199            if !matches!(code, Some(0 | 1)) {
200                let message = stderr.trim();
201                return Err(if message.is_empty() {
202                    format!(
203                        "ripgrep exited with code {}",
204                        code.map_or("null".to_owned(), |c| c.to_string())
205                    )
206                } else {
207                    message.to_owned()
208                });
209            }
210        }
211        if matches.is_empty() {
212            return Ok(text_result("No matches found", None));
213        }
214
215        let mut lines_truncated = false;
216        let mut cache: HashMap<String, Vec<String>> = HashMap::new();
217        let mut output: Vec<String> = Vec::new();
218        for found in &matches {
219            let shown = format_path(&found.file);
220            match (&found.text, context) {
221                (Some(text), 0) => {
222                    let text = text.replace("\r\n", "\n").replace('\r', "");
223                    let text = text.strip_suffix('\n').unwrap_or(&text);
224                    let (text, cut) = truncate_line(text, GREP_MAX_LINE_LENGTH);
225                    lines_truncated |= cut;
226                    output.push(format!("{shown}:{}: {text}", found.line));
227                }
228                _ => {
229                    let lines = cache
230                        .entry(found.file.clone())
231                        .or_insert_with(|| file_lines(&PathBuf::from(&found.file)));
232                    if lines.is_empty() {
233                        output.push(format!("{shown}:{}: (unable to read file)", found.line));
234                        continue;
235                    }
236                    let (start, end) = if context > 0 {
237                        (
238                            found.line.saturating_sub(context).max(1),
239                            found.line.saturating_add(context).min(lines.len() as u64),
240                        )
241                    } else {
242                        (found.line, found.line)
243                    };
244                    for current in start..=end {
245                        let text = lines
246                            .get(current as usize - 1)
247                            .map(|line| line.replace('\r', ""))
248                            .unwrap_or_default();
249                        let (text, cut) = truncate_line(&text, GREP_MAX_LINE_LENGTH);
250                        lines_truncated |= cut;
251                        if current == found.line {
252                            output.push(format!("{shown}:{current}: {text}"));
253                        } else {
254                            output.push(format!("{shown}-{current}- {text}"));
255                        }
256                    }
257                }
258            }
259        }
260
261        let mut details = Map::new();
262        let mut notices = Vec::new();
263        if limit_reached {
264            notices.push(format!(
265                "{} matches limit reached. Use limit={} for more, or refine pattern",
266                js_number(limit),
267                js_number(limit * 2.0)
268            ));
269            details.insert("matchLimitReached".into(), json!(limit));
270        }
271        let after = if lines_truncated {
272            vec![format!(
273                "Some lines truncated to {GREP_MAX_LINE_LENGTH} chars. Use read tool to see full lines"
274            )]
275        } else {
276            Vec::new()
277        };
278        let mut result = capped_output(&output.join("\n"), notices, after, details);
279        if lines_truncated {
280            result.details.get_or_insert_with(|| json!({}))["linesTruncated"] = json!(true);
281        }
282        Ok(result)
283    }
284}
yapi_core/tools/grep.rs

yapi_core/tools/
grep.rs
