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 if path.file_name().and_then(|name| name.to_str()) == Some("worktree")
132 && path
133 .parent()
134 .and_then(Path::file_name)
135 .and_then(|name| name.to_str())
136 == Some(".bamboo")
137 {
138 return true;
139 }
140 path.file_name()
141 .and_then(|name| name.to_str())
142 .map(|name| SKIP_DIRS.contains(&name))
143 .unwrap_or(false)
144 }
145
146 fn compile_glob(glob: Option<&str>) -> Result<Option<GlobSet>, ToolError> {
147 let Some(pattern) = glob else {
148 return Ok(None);
149 };
150
151 let mut builder = globset::GlobSetBuilder::new();
152 let glob = GlobBuilder::new(pattern)
153 .literal_separator(false)
154 .build()
155 .map_err(|e| ToolError::InvalidArguments(format!("Invalid glob pattern: {}", e)))?;
156 builder.add(glob);
157 builder
158 .build()
159 .map(Some)
160 .map_err(|e| ToolError::Execution(format!("Failed to compile glob: {}", e)))
161 }
162
163 fn compile_regex(
164 pattern: &str,
165 case_insensitive: bool,
166 multiline: bool,
167 ) -> Result<Regex, ToolError> {
168 let mut builder = RegexBuilder::new(pattern);
169 builder.case_insensitive(case_insensitive);
170 builder.dot_matches_new_line(multiline);
171 builder.multi_line(multiline);
172 builder
173 .build()
174 .map_err(|e| ToolError::InvalidArguments(format!("Invalid regex pattern: {}", e)))
175 }
176
177 fn byte_to_line(line_starts: &[usize], byte: usize) -> usize {
178 match line_starts.binary_search(&byte) {
179 Ok(idx) => idx,
180 Err(idx) => idx.saturating_sub(1),
181 }
182 }
183
184 fn format_content_hits(
185 path: &Path,
186 content: &str,
187 regex: &Regex,
188 multiline: bool,
189 before: usize,
190 after: usize,
191 line_numbers: bool,
192 ) -> Vec<String> {
193 let lines: Vec<&str> = content.lines().collect();
194 if lines.is_empty() {
195 return Vec::new();
196 }
197
198 let mut selected_lines = BTreeSet::new();
199
200 if multiline {
201 let mut line_starts = vec![0usize];
202 for (idx, byte) in content.bytes().enumerate() {
203 if byte == b'\n' {
204 line_starts.push(idx + 1);
205 }
206 }
207
208 for mat in regex.find_iter(content) {
209 let start_line = Self::byte_to_line(&line_starts, mat.start());
210 let end_line = Self::byte_to_line(&line_starts, mat.end().saturating_sub(1));
211 let range_start = start_line.saturating_sub(before);
212 let range_end = (end_line + 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 } else {
218 for (idx, line) in lines.iter().enumerate() {
219 if regex.is_match(line) {
220 let range_start = idx.saturating_sub(before);
221 let range_end = (idx + after).min(lines.len().saturating_sub(1));
222 for line_idx in range_start..=range_end {
223 selected_lines.insert(line_idx);
224 }
225 }
226 }
227 }
228
229 selected_lines
230 .into_iter()
231 .map(|idx| {
232 let display_path = bamboo_config::paths::path_to_display_string(path);
233 if line_numbers {
234 format!("{}:{}:{}", display_path, idx + 1, lines[idx])
235 } else {
236 format!("{}:{}", display_path, lines[idx])
237 }
238 })
239 .collect()
240 }
241
242 fn resolve_search_root(path: Option<&str>, cwd: &Path) -> PathBuf {
243 match path {
244 Some(path) => {
245 let candidate = PathBuf::from(path);
246 if candidate.is_absolute() {
247 candidate
248 } else {
249 cwd.join(candidate)
250 }
251 }
252 None => cwd.to_path_buf(),
253 }
254 }
255
256 fn validate_scope(
257 args: &GrepArgs,
258 output_mode: OutputMode,
259 multiline: bool,
260 cwd: &Path,
261 ) -> Result<(), ToolError> {
262 if matches!(output_mode, OutputMode::Content)
263 && args.path.is_none()
264 && args.glob.is_none()
265 && args.r#type.is_none()
266 {
267 return Err(ToolError::InvalidArguments(
268 SEARCH_SCOPE_TOO_BROAD_ERROR.to_string(),
269 ));
270 }
271
272 if multiline {
273 let Some(path) = args.path.as_deref() else {
274 return Err(ToolError::InvalidArguments(
275 MULTILINE_REQUIRES_NARROWED_PATH_ERROR.to_string(),
276 ));
277 };
278
279 let resolved = Self::resolve_search_root(Some(path), cwd);
280 if resolved.is_dir() {
281 if let (Ok(resolved_canonical), Ok(cwd_canonical)) =
282 (resolved.canonicalize(), cwd.canonicalize())
283 {
284 if resolved_canonical == cwd_canonical {
285 return Err(ToolError::InvalidArguments(
286 MULTILINE_REQUIRES_NARROWED_PATH_ERROR.to_string(),
287 ));
288 }
289 }
290 }
291 }
292
293 Ok(())
294 }
295}
296
297impl Default for GrepTool {
298 fn default() -> Self {
299 Self::new()
300 }
301}
302
303#[async_trait]
304impl Tool for GrepTool {
305 fn name(&self) -> &str {
306 "Grep"
307 }
308
309 fn description(&self) -> &str {
310 "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."
311 }
312
313 fn classify(&self, _args: &serde_json::Value) -> ToolClass {
314 ToolClass::READONLY_PARALLEL
315 }
316
317 fn parameters_schema(&self) -> serde_json::Value {
318 json!({
319 "type": "object",
320 "properties": {
321 "pattern": { "type": "string", "description": "Regex pattern" },
322 "path": { "type": "string", "description": "File or directory to search. Narrow this for expensive or multiline searches." },
323 "glob": { "type": "string", "description": "Glob file filter used to limit candidate files" },
324 "output_mode": {
325 "type": "string",
326 "enum": ["content", "files_with_matches", "count"],
327 "description": "Output mode. Prefer files_with_matches for broad discovery, then refine with Read or content mode."
328 },
329 "-B": { "type": "number", "description": "Lines before match" },
330 "-A": { "type": "number", "description": "Lines after match" },
331 "-C": { "type": "number", "description": "Lines before and after match" },
332 "-n": { "type": "boolean", "description": "Show line numbers" },
333 "-i": { "type": "boolean", "description": "Case insensitive" },
334 "type": { "type": "string", "description": "File type filter (for example rust, js, ts, py)" },
335 "head_limit": { "type": "number", "description": "Limit output entries. Keep this small for broad queries." },
336 "multiline": { "type": "boolean", "description": "Enable multiline regex. Requires a narrowed path." }
337 },
338 "required": ["pattern"],
339 "additionalProperties": false
340 })
341 }
342
343 async fn invoke(
344 &self,
345 args: serde_json::Value,
346 ctx: ToolCtx,
347 ) -> Result<ToolOutcome, ToolError> {
348 let parsed: GrepArgs = serde_json::from_value(args)
349 .map_err(|e| ToolError::InvalidArguments(format!("Invalid Grep args: {}", e)))?;
350
351 let cwd = workspace_state::workspace_or_process_cwd(ctx.session_id());
352 let root = Self::resolve_search_root(parsed.path.as_deref(), &cwd);
353
354 let output_mode = parsed.output_mode.unwrap_or_default();
355 let context = parsed.context.unwrap_or(0);
356 let before = parsed.before.unwrap_or(context);
357 let after = parsed.after.unwrap_or(context);
358 let line_numbers = parsed.line_numbers.unwrap_or(false);
359 let case_insensitive = parsed.case_insensitive.unwrap_or(false);
360 let multiline = parsed.multiline.unwrap_or(false);
361 let head_limit = parsed.head_limit.unwrap_or(DEFAULT_HEAD_LIMIT);
362
363 Self::validate_scope(&parsed, output_mode, multiline, &cwd)?;
364
365 let regex = Self::compile_regex(&parsed.pattern, case_insensitive, multiline)?;
366 let glob_filter = Self::compile_glob(parsed.glob.as_deref())?;
367
368 let files = if root.is_file() {
369 vec![root.clone()]
370 } else if root.is_dir() {
371 Self::collect_files(&root, parsed.r#type.as_deref())
372 } else {
373 return Err(ToolError::Execution(format!(
374 "Path does not exist: {}",
375 root.display()
376 )));
377 };
378
379 let mut matched_files = Vec::new();
380 let mut count_rows = Vec::new();
381 let mut content_rows = Vec::new();
382 let mut total_matches = 0usize;
383 let mut partial = false;
384
385 for file in files {
386 if let Some(filter) = &glob_filter {
387 let relative = file.strip_prefix(&root).unwrap_or(&file);
388 if !filter.is_match(relative) && !filter.is_match(&file) {
389 continue;
390 }
391 }
392
393 let Ok(metadata) = tokio::fs::metadata(&file).await else {
394 continue;
395 };
396 if metadata.len() > MAX_FILE_BYTES {
397 continue;
398 }
399
400 let Ok(content) = tokio::fs::read_to_string(&file).await else {
401 continue;
402 };
403
404 if content.contains('\0') {
405 continue;
406 }
407
408 let match_count = if multiline {
409 regex.find_iter(&content).count()
410 } else {
411 content.lines().filter(|line| regex.is_match(line)).count()
412 };
413
414 if match_count == 0 {
415 continue;
416 }
417
418 total_matches = total_matches.saturating_add(match_count);
419 if total_matches > MAX_MATCHES {
420 return Err(ToolError::Execution(RESULT_TOO_LARGE_ERROR.to_string()));
421 }
422
423 matched_files.push(bamboo_config::paths::path_to_display_string(&file));
424 count_rows.push(format!(
425 "{}:{}",
426 bamboo_config::paths::path_to_display_string(&file),
427 match_count
428 ));
429
430 if matches!(output_mode, OutputMode::Content) {
431 content_rows.extend(Self::format_content_hits(
432 &file,
433 &content,
434 ®ex,
435 multiline,
436 before,
437 after,
438 line_numbers,
439 ));
440 if content_rows.len() >= head_limit {
441 content_rows.truncate(head_limit);
442 partial = true;
443 break;
444 }
445 }
446
447 if matches!(
448 output_mode,
449 OutputMode::FilesWithMatches | OutputMode::Count
450 ) && matched_files.len() >= head_limit
451 {
452 partial = true;
453 break;
454 }
455 }
456
457 let mut result_lines = match output_mode {
458 OutputMode::FilesWithMatches => matched_files,
459 OutputMode::Count => count_rows,
460 OutputMode::Content => content_rows,
461 };
462
463 if result_lines.len() > head_limit {
464 result_lines.truncate(head_limit);
465 partial = true;
466 }
467 if partial {
468 result_lines
469 .push("[PARTIAL] Output was truncated. Narrow path/pattern and retry.".to_string());
470 }
471
472 let result = result_lines.join("\n");
473 if result.len() > MAX_RESULT_BYTES {
474 return Err(ToolError::Execution(RESULT_TOO_LARGE_ERROR.to_string()));
475 }
476
477 Ok(ToolOutcome::Completed(ToolResult {
478 success: true,
479 result,
480 display_preference: Some("Collapsible".to_string()),
481 images: Vec::new(),
482 }))
483 }
484}
485
486#[cfg(test)]
487mod tests {
488 use super::*;
489 use serde_json::json;
490
491 async fn run(tool: &GrepTool, args: serde_json::Value) -> Result<ToolResult, ToolError> {
492 match tool.invoke(args, ToolCtx::none("t")).await? {
493 ToolOutcome::Completed(r) => Ok(r),
494 _ => panic!("expected Completed"),
495 }
496 }
497
498 fn result_lines(result: &ToolResult) -> Vec<&str> {
499 result
500 .result
501 .lines()
502 .filter(|line| !line.is_empty())
503 .collect()
504 }
505
506 fn non_partial_lines(result: &ToolResult) -> Vec<&str> {
507 result_lines(result)
508 .into_iter()
509 .filter(|line| !line.starts_with("[PARTIAL]"))
510 .collect()
511 }
512
513 #[tokio::test]
514 async fn grep_defaults_to_files_with_matches() {
515 let dir = tempfile::tempdir().unwrap();
516 let file_hit = dir.path().join("match.rs");
517 let file_miss = dir.path().join("miss.txt");
518 tokio::fs::write(&file_hit, "let value = 1;\nneedle\n")
519 .await
520 .unwrap();
521 tokio::fs::write(&file_miss, "nothing to see\n")
522 .await
523 .unwrap();
524
525 let tool = GrepTool::new();
526 let result = run(
527 &tool,
528 json!({
529 "pattern": "needle",
530 "path": dir.path()
531 }),
532 )
533 .await
534 .unwrap();
535
536 assert!(result.success);
537 let lines = result_lines(&result);
538 assert_eq!(lines.len(), 1);
539 assert!(lines[0].contains("match.rs"));
540 }
541
542 #[tokio::test]
543 async fn grep_skips_project_worktree_checkout() {
544 let dir = tempfile::tempdir().unwrap();
545 let kept = dir.path().join("src/kept.txt");
546 let duplicate = dir.path().join(".bamboo/worktree/child/duplicate.txt");
547 tokio::fs::create_dir_all(kept.parent().unwrap())
548 .await
549 .unwrap();
550 tokio::fs::create_dir_all(duplicate.parent().unwrap())
551 .await
552 .unwrap();
553 tokio::fs::write(&kept, "needle").await.unwrap();
554 tokio::fs::write(&duplicate, "needle").await.unwrap();
555
556 let result = run(
557 &GrepTool::new(),
558 json!({"pattern": "needle", "path": dir.path()}),
559 )
560 .await
561 .unwrap();
562 assert!(result.result.contains("kept.txt"));
563 assert!(!result.result.contains("duplicate.txt"));
564 }
565
566 #[tokio::test]
567 async fn grep_content_mode_supports_context_and_line_numbers() {
568 let dir = tempfile::tempdir().unwrap();
569 let file = dir.path().join("content.txt");
570 tokio::fs::write(&file, "one\ntwo\nneedle\nfour\nfive\n")
571 .await
572 .unwrap();
573
574 let tool = GrepTool::new();
575 let result = run(
576 &tool,
577 json!({
578 "pattern": "needle",
579 "path": file,
580 "output_mode": "content",
581 "-C": 1,
582 "-n": true
583 }),
584 )
585 .await
586 .unwrap();
587
588 let output = result.result;
589 assert!(output.contains(":2:two"));
590 assert!(output.contains(":3:needle"));
591 assert!(output.contains(":4:four"));
592 assert!(!output.contains(":1:one"));
593 assert!(!output.contains(":5:five"));
594 }
595
596 #[tokio::test]
597 async fn grep_count_mode_respects_type_filter_and_head_limit() {
598 let dir = tempfile::tempdir().unwrap();
599 let file_rs_a = dir.path().join("a.rs");
600 let file_rs_b = dir.path().join("b.rs");
601 let file_txt = dir.path().join("c.txt");
602 tokio::fs::write(&file_rs_a, "foo\nfoo\n").await.unwrap();
603 tokio::fs::write(&file_rs_b, "foo\n").await.unwrap();
604 tokio::fs::write(&file_txt, "foo\n").await.unwrap();
605
606 let tool = GrepTool::new();
607 let result = run(
608 &tool,
609 json!({
610 "pattern": "foo",
611 "path": dir.path(),
612 "output_mode": "count",
613 "type": "rust",
614 "head_limit": 1
615 }),
616 )
617 .await
618 .unwrap();
619
620 let lines = non_partial_lines(&result);
621 assert_eq!(lines.len(), 1);
622 assert!(lines[0].contains(".rs:"));
623 assert!(!lines[0].contains("c.txt"));
624 assert!(result.result.contains("[PARTIAL]"));
625 }
626
627 #[tokio::test]
628 async fn grep_multiline_and_case_insensitive_work_with_glob_filter() {
629 let dir = tempfile::tempdir().unwrap();
630 let file_one = dir.path().join("one.rs");
631 let file_two = dir.path().join("two.rs");
632 tokio::fs::write(&file_one, "Hello\nWORLD\n").await.unwrap();
633 tokio::fs::write(&file_two, "Hello\nplanet\n")
634 .await
635 .unwrap();
636
637 let tool = GrepTool::new();
638 let result = run(
639 &tool,
640 json!({
641 "pattern": "hello\\s+world",
642 "path": dir.path(),
643 "glob": "**/one.rs",
644 "-i": true,
645 "multiline": true
646 }),
647 )
648 .await
649 .unwrap();
650
651 let output = result.result;
652 assert!(output.contains("one.rs"));
653 assert!(!output.contains("two.rs"));
654 }
655
656 #[tokio::test]
657 async fn grep_content_mode_requires_scope_hint() {
658 let tool = GrepTool::new();
659 let error = run(
660 &tool,
661 json!({
662 "pattern": "needle",
663 "output_mode": "content"
664 }),
665 )
666 .await
667 .expect_err("content mode without scope should fail");
668
669 assert!(matches!(error, ToolError::InvalidArguments(_)));
670 assert!(error.to_string().contains(SEARCH_SCOPE_TOO_BROAD_ERROR));
671 }
672
673 #[tokio::test]
674 async fn grep_multiline_requires_explicit_narrowed_path() {
675 let tool = GrepTool::new();
676 let error = run(
677 &tool,
678 json!({
679 "pattern": "a\\s+b",
680 "multiline": true
681 }),
682 )
683 .await
684 .expect_err("multiline without path should fail");
685 assert!(matches!(error, ToolError::InvalidArguments(_)));
686 assert!(error
687 .to_string()
688 .contains(MULTILINE_REQUIRES_NARROWED_PATH_ERROR));
689
690 let cwd = std::env::current_dir().unwrap();
691 let error = run(
692 &tool,
693 json!({
694 "pattern": "a\\s+b",
695 "multiline": true,
696 "path": cwd
697 }),
698 )
699 .await
700 .expect_err("multiline at workspace root should fail");
701 assert!(matches!(error, ToolError::InvalidArguments(_)));
702 assert!(error
703 .to_string()
704 .contains(MULTILINE_REQUIRES_NARROWED_PATH_ERROR));
705 }
706
707 #[tokio::test]
708 async fn grep_defaults_head_limit_to_200() {
709 let dir = tempfile::tempdir().unwrap();
710 for idx in 0..260 {
711 let file = dir.path().join(format!("file-{idx}.txt"));
712 tokio::fs::write(&file, "needle\n").await.unwrap();
713 }
714
715 let tool = GrepTool::new();
716 let result = run(
717 &tool,
718 json!({
719 "pattern": "needle",
720 "path": dir.path()
721 }),
722 )
723 .await
724 .unwrap();
725
726 let lines = non_partial_lines(&result);
727 assert_eq!(lines.len(), 200);
728 assert!(result.result.contains("[PARTIAL]"));
729 }
730
731 #[tokio::test]
732 async fn grep_rejects_excessive_match_volume() {
733 let dir = tempfile::tempdir().unwrap();
734 let file = dir.path().join("huge.txt");
735 let mut content = String::new();
736 for _ in 0..(MAX_MATCHES + 1) {
737 content.push_str("needle\n");
738 }
739 tokio::fs::write(&file, content).await.unwrap();
740
741 let tool = GrepTool::new();
742 let error = run(
743 &tool,
744 json!({
745 "pattern": "needle",
746 "path": file
747 }),
748 )
749 .await
750 .expect_err("should reject oversized results");
751
752 assert!(matches!(error, ToolError::Execution(_)));
753 assert!(error.to_string().contains(RESULT_TOO_LARGE_ERROR));
754 }
755}