1use async_trait::async_trait;
2use bamboo_agent_core::{Tool, ToolClass, ToolCtx, ToolError, ToolOutcome, ToolResult};
3use globset::{GlobBuilder, GlobSet};
4use regex::{Regex, RegexBuilder};
5use serde::Deserialize;
6use serde_json::json;
7use std::collections::{BTreeSet, HashMap};
8use std::path::{Path, PathBuf};
9use walkdir::WalkDir;
10
11use super::workspace_state;
12
13const DEFAULT_HEAD_LIMIT: usize = 200;
14const MAX_RESULT_BYTES: usize = 256 * 1024;
15const MAX_MATCHES: usize = 2_000;
16const MAX_SCANNED_FILES: usize = 50_000;
17const MAX_FILE_BYTES: u64 = 2 * 1024 * 1024;
18const SKIP_DIRS: [&str; 8] = [
19 ".git",
20 "node_modules",
21 "target",
22 "dist",
23 "build",
24 ".next",
25 ".cache",
26 "coverage",
27];
28const SEARCH_SCOPE_TOO_BROAD_ERROR: &str =
29 "Search scope too broad. Add path/glob/type or reduce pattern.";
30const MULTILINE_REQUIRES_NARROWED_PATH_ERROR: &str = "Multiline grep requires narrowed path.";
31const RESULT_TOO_LARGE_ERROR: &str = "Result too large; refine query and retry.";
32
33#[derive(Debug, Deserialize, Clone, Copy, Default)]
34#[serde(rename_all = "snake_case")]
35enum OutputMode {
36 Content,
37 #[default]
38 FilesWithMatches,
39 Count,
40}
41
42#[derive(Debug, Deserialize)]
43struct GrepArgs {
44 pattern: String,
45 #[serde(default)]
46 path: Option<String>,
47 #[serde(default)]
48 glob: Option<String>,
49 #[serde(default)]
50 output_mode: Option<OutputMode>,
51 #[serde(rename = "-B", default)]
52 before: Option<usize>,
53 #[serde(rename = "-A", default)]
54 after: Option<usize>,
55 #[serde(rename = "-C", default)]
56 context: Option<usize>,
57 #[serde(rename = "-n", default)]
58 line_numbers: Option<bool>,
59 #[serde(rename = "-i", default)]
60 case_insensitive: Option<bool>,
61 #[serde(default)]
62 r#type: Option<String>,
63 #[serde(default)]
64 head_limit: Option<usize>,
65 #[serde(default)]
66 multiline: Option<bool>,
67}
68
69pub struct GrepTool;
70
71impl GrepTool {
72 pub fn new() -> Self {
73 Self
74 }
75
76 fn extension_map() -> HashMap<&'static str, &'static [&'static str]> {
77 HashMap::from([
78 ("js", &["js", "mjs", "cjs"] as &[_]),
79 ("ts", &["ts", "tsx"]),
80 ("py", &["py"]),
81 ("rust", &["rs"]),
82 ("go", &["go"]),
83 ("java", &["java"]),
84 ("cpp", &["cc", "cpp", "cxx", "hpp", "h"]),
85 ("c", &["c", "h"]),
86 ("json", &["json"]),
87 ("yaml", &["yaml", "yml"]),
88 ("toml", &["toml"]),
89 ("md", &["md", "markdown"]),
90 ])
91 }
92
93 fn collect_files(base: &Path, type_filter: Option<&str>) -> Vec<PathBuf> {
94 let ext_map = Self::extension_map();
95 let allowed_ext = type_filter.and_then(|name| ext_map.get(name).copied());
96
97 let mut files = Vec::new();
98 for entry in WalkDir::new(base)
99 .follow_links(false)
100 .into_iter()
101 .filter_entry(|entry| {
102 !entry.file_type().is_dir() || !Self::should_skip_dir(entry.path())
103 })
104 .filter_map(|entry| entry.ok())
105 {
106 if !entry.file_type().is_file() {
107 continue;
108 }
109 if files.len() >= MAX_SCANNED_FILES {
110 break;
111 }
112 let path = entry.path();
113
114 if let Some(extensions) = allowed_ext {
115 let ext = path
116 .extension()
117 .and_then(|v| v.to_str())
118 .unwrap_or_default();
119 if !extensions.iter().any(|candidate| candidate == &ext) {
120 continue;
121 }
122 }
123
124 files.push(path.to_path_buf());
125 }
126
127 files
128 }
129
130 fn should_skip_dir(path: &Path) -> bool {
131 path.file_name()
132 .and_then(|name| name.to_str())
133 .map(|name| SKIP_DIRS.contains(&name))
134 .unwrap_or(false)
135 }
136
137 fn compile_glob(glob: Option<&str>) -> Result<Option<GlobSet>, ToolError> {
138 let Some(pattern) = glob else {
139 return Ok(None);
140 };
141
142 let mut builder = globset::GlobSetBuilder::new();
143 let glob = GlobBuilder::new(pattern)
144 .literal_separator(false)
145 .build()
146 .map_err(|e| ToolError::InvalidArguments(format!("Invalid glob pattern: {}", e)))?;
147 builder.add(glob);
148 builder
149 .build()
150 .map(Some)
151 .map_err(|e| ToolError::Execution(format!("Failed to compile glob: {}", e)))
152 }
153
154 fn compile_regex(
155 pattern: &str,
156 case_insensitive: bool,
157 multiline: bool,
158 ) -> Result<Regex, ToolError> {
159 let mut builder = RegexBuilder::new(pattern);
160 builder.case_insensitive(case_insensitive);
161 builder.dot_matches_new_line(multiline);
162 builder.multi_line(multiline);
163 builder
164 .build()
165 .map_err(|e| ToolError::InvalidArguments(format!("Invalid regex pattern: {}", e)))
166 }
167
168 fn byte_to_line(line_starts: &[usize], byte: usize) -> usize {
169 match line_starts.binary_search(&byte) {
170 Ok(idx) => idx,
171 Err(idx) => idx.saturating_sub(1),
172 }
173 }
174
175 fn format_content_hits(
176 path: &Path,
177 content: &str,
178 regex: &Regex,
179 multiline: bool,
180 before: usize,
181 after: usize,
182 line_numbers: bool,
183 ) -> Vec<String> {
184 let lines: Vec<&str> = content.lines().collect();
185 if lines.is_empty() {
186 return Vec::new();
187 }
188
189 let mut selected_lines = BTreeSet::new();
190
191 if multiline {
192 let mut line_starts = vec![0usize];
193 for (idx, byte) in content.bytes().enumerate() {
194 if byte == b'\n' {
195 line_starts.push(idx + 1);
196 }
197 }
198
199 for mat in regex.find_iter(content) {
200 let start_line = Self::byte_to_line(&line_starts, mat.start());
201 let end_line = Self::byte_to_line(&line_starts, mat.end().saturating_sub(1));
202 let range_start = start_line.saturating_sub(before);
203 let range_end = (end_line + after).min(lines.len().saturating_sub(1));
204 for line_idx in range_start..=range_end {
205 selected_lines.insert(line_idx);
206 }
207 }
208 } else {
209 for (idx, line) in lines.iter().enumerate() {
210 if regex.is_match(line) {
211 let range_start = idx.saturating_sub(before);
212 let range_end = (idx + after).min(lines.len().saturating_sub(1));
213 for line_idx in range_start..=range_end {
214 selected_lines.insert(line_idx);
215 }
216 }
217 }
218 }
219
220 selected_lines
221 .into_iter()
222 .map(|idx| {
223 let display_path = bamboo_config::paths::path_to_display_string(path);
224 if line_numbers {
225 format!("{}:{}:{}", display_path, idx + 1, lines[idx])
226 } else {
227 format!("{}:{}", display_path, lines[idx])
228 }
229 })
230 .collect()
231 }
232
233 fn resolve_search_root(path: Option<&str>, cwd: &Path) -> PathBuf {
234 match path {
235 Some(path) => {
236 let candidate = PathBuf::from(path);
237 if candidate.is_absolute() {
238 candidate
239 } else {
240 cwd.join(candidate)
241 }
242 }
243 None => cwd.to_path_buf(),
244 }
245 }
246
247 fn validate_scope(
248 args: &GrepArgs,
249 output_mode: OutputMode,
250 multiline: bool,
251 cwd: &Path,
252 ) -> Result<(), ToolError> {
253 if matches!(output_mode, OutputMode::Content)
254 && args.path.is_none()
255 && args.glob.is_none()
256 && args.r#type.is_none()
257 {
258 return Err(ToolError::InvalidArguments(
259 SEARCH_SCOPE_TOO_BROAD_ERROR.to_string(),
260 ));
261 }
262
263 if multiline {
264 let Some(path) = args.path.as_deref() else {
265 return Err(ToolError::InvalidArguments(
266 MULTILINE_REQUIRES_NARROWED_PATH_ERROR.to_string(),
267 ));
268 };
269
270 let resolved = Self::resolve_search_root(Some(path), cwd);
271 if resolved.is_dir() {
272 if let (Ok(resolved_canonical), Ok(cwd_canonical)) =
273 (resolved.canonicalize(), cwd.canonicalize())
274 {
275 if resolved_canonical == cwd_canonical {
276 return Err(ToolError::InvalidArguments(
277 MULTILINE_REQUIRES_NARROWED_PATH_ERROR.to_string(),
278 ));
279 }
280 }
281 }
282 }
283
284 Ok(())
285 }
286}
287
288impl Default for GrepTool {
289 fn default() -> Self {
290 Self::new()
291 }
292}
293
294#[async_trait]
295impl Tool for GrepTool {
296 fn name(&self) -> &str {
297 "Grep"
298 }
299
300 fn description(&self) -> &str {
301 "Search file contents using ripgrep-style regex parameters. Start with files_with_matches or a narrowed path/glob/type before using content or multiline mode."
302 }
303
304 fn classify(&self, _args: &serde_json::Value) -> ToolClass {
305 ToolClass::READONLY_PARALLEL
306 }
307
308 fn parameters_schema(&self) -> serde_json::Value {
309 json!({
310 "type": "object",
311 "properties": {
312 "pattern": { "type": "string", "description": "Regex pattern" },
313 "path": { "type": "string", "description": "File or directory to search. Narrow this for expensive or multiline searches." },
314 "glob": { "type": "string", "description": "Glob file filter used to limit candidate files" },
315 "output_mode": {
316 "type": "string",
317 "enum": ["content", "files_with_matches", "count"],
318 "description": "Output mode. Prefer files_with_matches for broad discovery, then refine with Read or content mode."
319 },
320 "-B": { "type": "number", "description": "Lines before match" },
321 "-A": { "type": "number", "description": "Lines after match" },
322 "-C": { "type": "number", "description": "Lines before and after match" },
323 "-n": { "type": "boolean", "description": "Show line numbers" },
324 "-i": { "type": "boolean", "description": "Case insensitive" },
325 "type": { "type": "string", "description": "File type filter (for example rust, js, ts, py)" },
326 "head_limit": { "type": "number", "description": "Limit output entries. Keep this small for broad queries." },
327 "multiline": { "type": "boolean", "description": "Enable multiline regex. Requires a narrowed path." }
328 },
329 "required": ["pattern"],
330 "additionalProperties": false
331 })
332 }
333
334 async fn invoke(
335 &self,
336 args: serde_json::Value,
337 ctx: ToolCtx,
338 ) -> Result<ToolOutcome, ToolError> {
339 let parsed: GrepArgs = serde_json::from_value(args)
340 .map_err(|e| ToolError::InvalidArguments(format!("Invalid Grep args: {}", e)))?;
341
342 let cwd = workspace_state::workspace_or_process_cwd(ctx.session_id());
343 let root = Self::resolve_search_root(parsed.path.as_deref(), &cwd);
344
345 let output_mode = parsed.output_mode.unwrap_or_default();
346 let context = parsed.context.unwrap_or(0);
347 let before = parsed.before.unwrap_or(context);
348 let after = parsed.after.unwrap_or(context);
349 let line_numbers = parsed.line_numbers.unwrap_or(false);
350 let case_insensitive = parsed.case_insensitive.unwrap_or(false);
351 let multiline = parsed.multiline.unwrap_or(false);
352 let head_limit = parsed.head_limit.unwrap_or(DEFAULT_HEAD_LIMIT);
353
354 Self::validate_scope(&parsed, output_mode, multiline, &cwd)?;
355
356 let regex = Self::compile_regex(&parsed.pattern, case_insensitive, multiline)?;
357 let glob_filter = Self::compile_glob(parsed.glob.as_deref())?;
358
359 let files = if root.is_file() {
360 vec![root.clone()]
361 } else if root.is_dir() {
362 Self::collect_files(&root, parsed.r#type.as_deref())
363 } else {
364 return Err(ToolError::Execution(format!(
365 "Path does not exist: {}",
366 root.display()
367 )));
368 };
369
370 let mut matched_files = Vec::new();
371 let mut count_rows = Vec::new();
372 let mut content_rows = Vec::new();
373 let mut total_matches = 0usize;
374 let mut partial = false;
375
376 for file in files {
377 if let Some(filter) = &glob_filter {
378 let relative = file.strip_prefix(&root).unwrap_or(&file);
379 if !filter.is_match(relative) && !filter.is_match(&file) {
380 continue;
381 }
382 }
383
384 let Ok(metadata) = tokio::fs::metadata(&file).await else {
385 continue;
386 };
387 if metadata.len() > MAX_FILE_BYTES {
388 continue;
389 }
390
391 let Ok(content) = tokio::fs::read_to_string(&file).await else {
392 continue;
393 };
394
395 if content.contains('\0') {
396 continue;
397 }
398
399 let match_count = if multiline {
400 regex.find_iter(&content).count()
401 } else {
402 content.lines().filter(|line| regex.is_match(line)).count()
403 };
404
405 if match_count == 0 {
406 continue;
407 }
408
409 total_matches = total_matches.saturating_add(match_count);
410 if total_matches > MAX_MATCHES {
411 return Err(ToolError::Execution(RESULT_TOO_LARGE_ERROR.to_string()));
412 }
413
414 matched_files.push(bamboo_config::paths::path_to_display_string(&file));
415 count_rows.push(format!(
416 "{}:{}",
417 bamboo_config::paths::path_to_display_string(&file),
418 match_count
419 ));
420
421 if matches!(output_mode, OutputMode::Content) {
422 content_rows.extend(Self::format_content_hits(
423 &file,
424 &content,
425 ®ex,
426 multiline,
427 before,
428 after,
429 line_numbers,
430 ));
431 if content_rows.len() >= head_limit {
432 content_rows.truncate(head_limit);
433 partial = true;
434 break;
435 }
436 }
437
438 if matches!(
439 output_mode,
440 OutputMode::FilesWithMatches | OutputMode::Count
441 ) && matched_files.len() >= head_limit
442 {
443 partial = true;
444 break;
445 }
446 }
447
448 let mut result_lines = match output_mode {
449 OutputMode::FilesWithMatches => matched_files,
450 OutputMode::Count => count_rows,
451 OutputMode::Content => content_rows,
452 };
453
454 if result_lines.len() > head_limit {
455 result_lines.truncate(head_limit);
456 partial = true;
457 }
458 if partial {
459 result_lines
460 .push("[PARTIAL] Output was truncated. Narrow path/pattern and retry.".to_string());
461 }
462
463 let result = result_lines.join("\n");
464 if result.len() > MAX_RESULT_BYTES {
465 return Err(ToolError::Execution(RESULT_TOO_LARGE_ERROR.to_string()));
466 }
467
468 Ok(ToolOutcome::Completed(ToolResult {
469 success: true,
470 result,
471 display_preference: Some("Collapsible".to_string()),
472 images: Vec::new(),
473 }))
474 }
475}
476
477#[cfg(test)]
478mod tests {
479 use super::*;
480 use serde_json::json;
481
482 async fn run(tool: &GrepTool, args: serde_json::Value) -> Result<ToolResult, ToolError> {
483 match tool.invoke(args, ToolCtx::none("t")).await? {
484 ToolOutcome::Completed(r) => Ok(r),
485 _ => panic!("expected Completed"),
486 }
487 }
488
489 fn result_lines(result: &ToolResult) -> Vec<&str> {
490 result
491 .result
492 .lines()
493 .filter(|line| !line.is_empty())
494 .collect()
495 }
496
497 fn non_partial_lines(result: &ToolResult) -> Vec<&str> {
498 result_lines(result)
499 .into_iter()
500 .filter(|line| !line.starts_with("[PARTIAL]"))
501 .collect()
502 }
503
504 #[tokio::test]
505 async fn grep_defaults_to_files_with_matches() {
506 let dir = tempfile::tempdir().unwrap();
507 let file_hit = dir.path().join("match.rs");
508 let file_miss = dir.path().join("miss.txt");
509 tokio::fs::write(&file_hit, "let value = 1;\nneedle\n")
510 .await
511 .unwrap();
512 tokio::fs::write(&file_miss, "nothing to see\n")
513 .await
514 .unwrap();
515
516 let tool = GrepTool::new();
517 let result = run(&tool,json!({
518 "pattern": "needle",
519 "path": dir.path()
520 }))
521 .await
522 .unwrap();
523
524 assert!(result.success);
525 let lines = result_lines(&result);
526 assert_eq!(lines.len(), 1);
527 assert!(lines[0].contains("match.rs"));
528 }
529
530 #[tokio::test]
531 async fn grep_content_mode_supports_context_and_line_numbers() {
532 let dir = tempfile::tempdir().unwrap();
533 let file = dir.path().join("content.txt");
534 tokio::fs::write(&file, "one\ntwo\nneedle\nfour\nfive\n")
535 .await
536 .unwrap();
537
538 let tool = GrepTool::new();
539 let result = run(&tool,json!({
540 "pattern": "needle",
541 "path": file,
542 "output_mode": "content",
543 "-C": 1,
544 "-n": true
545 }))
546 .await
547 .unwrap();
548
549 let output = result.result;
550 assert!(output.contains(":2:two"));
551 assert!(output.contains(":3:needle"));
552 assert!(output.contains(":4:four"));
553 assert!(!output.contains(":1:one"));
554 assert!(!output.contains(":5:five"));
555 }
556
557 #[tokio::test]
558 async fn grep_count_mode_respects_type_filter_and_head_limit() {
559 let dir = tempfile::tempdir().unwrap();
560 let file_rs_a = dir.path().join("a.rs");
561 let file_rs_b = dir.path().join("b.rs");
562 let file_txt = dir.path().join("c.txt");
563 tokio::fs::write(&file_rs_a, "foo\nfoo\n").await.unwrap();
564 tokio::fs::write(&file_rs_b, "foo\n").await.unwrap();
565 tokio::fs::write(&file_txt, "foo\n").await.unwrap();
566
567 let tool = GrepTool::new();
568 let result = run(&tool,json!({
569 "pattern": "foo",
570 "path": dir.path(),
571 "output_mode": "count",
572 "type": "rust",
573 "head_limit": 1
574 }))
575 .await
576 .unwrap();
577
578 let lines = non_partial_lines(&result);
579 assert_eq!(lines.len(), 1);
580 assert!(lines[0].contains(".rs:"));
581 assert!(!lines[0].contains("c.txt"));
582 assert!(result.result.contains("[PARTIAL]"));
583 }
584
585 #[tokio::test]
586 async fn grep_multiline_and_case_insensitive_work_with_glob_filter() {
587 let dir = tempfile::tempdir().unwrap();
588 let file_one = dir.path().join("one.rs");
589 let file_two = dir.path().join("two.rs");
590 tokio::fs::write(&file_one, "Hello\nWORLD\n").await.unwrap();
591 tokio::fs::write(&file_two, "Hello\nplanet\n")
592 .await
593 .unwrap();
594
595 let tool = GrepTool::new();
596 let result = run(&tool,json!({
597 "pattern": "hello\\s+world",
598 "path": dir.path(),
599 "glob": "**/one.rs",
600 "-i": true,
601 "multiline": true
602 }))
603 .await
604 .unwrap();
605
606 let output = result.result;
607 assert!(output.contains("one.rs"));
608 assert!(!output.contains("two.rs"));
609 }
610
611 #[tokio::test]
612 async fn grep_content_mode_requires_scope_hint() {
613 let tool = GrepTool::new();
614 let error = run(&tool,json!({
615 "pattern": "needle",
616 "output_mode": "content"
617 }))
618 .await
619 .expect_err("content mode without scope should fail");
620
621 assert!(matches!(error, ToolError::InvalidArguments(_)));
622 assert!(error.to_string().contains(SEARCH_SCOPE_TOO_BROAD_ERROR));
623 }
624
625 #[tokio::test]
626 async fn grep_multiline_requires_explicit_narrowed_path() {
627 let tool = GrepTool::new();
628 let error = run(&tool,json!({
629 "pattern": "a\\s+b",
630 "multiline": true
631 }))
632 .await
633 .expect_err("multiline without path should fail");
634 assert!(matches!(error, ToolError::InvalidArguments(_)));
635 assert!(error
636 .to_string()
637 .contains(MULTILINE_REQUIRES_NARROWED_PATH_ERROR));
638
639 let cwd = std::env::current_dir().unwrap();
640 let error = run(&tool,json!({
641 "pattern": "a\\s+b",
642 "multiline": true,
643 "path": cwd
644 }))
645 .await
646 .expect_err("multiline at workspace root should fail");
647 assert!(matches!(error, ToolError::InvalidArguments(_)));
648 assert!(error
649 .to_string()
650 .contains(MULTILINE_REQUIRES_NARROWED_PATH_ERROR));
651 }
652
653 #[tokio::test]
654 async fn grep_defaults_head_limit_to_200() {
655 let dir = tempfile::tempdir().unwrap();
656 for idx in 0..260 {
657 let file = dir.path().join(format!("file-{idx}.txt"));
658 tokio::fs::write(&file, "needle\n").await.unwrap();
659 }
660
661 let tool = GrepTool::new();
662 let result = run(&tool,json!({
663 "pattern": "needle",
664 "path": dir.path()
665 }))
666 .await
667 .unwrap();
668
669 let lines = non_partial_lines(&result);
670 assert_eq!(lines.len(), 200);
671 assert!(result.result.contains("[PARTIAL]"));
672 }
673
674 #[tokio::test]
675 async fn grep_rejects_excessive_match_volume() {
676 let dir = tempfile::tempdir().unwrap();
677 let file = dir.path().join("huge.txt");
678 let mut content = String::new();
679 for _ in 0..(MAX_MATCHES + 1) {
680 content.push_str("needle\n");
681 }
682 tokio::fs::write(&file, content).await.unwrap();
683
684 let tool = GrepTool::new();
685 let error = run(&tool,json!({
686 "pattern": "needle",
687 "path": file
688 }))
689 .await
690 .expect_err("should reject oversized results");
691
692 assert!(matches!(error, ToolError::Execution(_)));
693 assert!(error.to_string().contains(RESULT_TOO_LARGE_ERROR));
694 }
695}