1pub mod common;
2
3use crate::coding::error::GrepError;
4use crate::coding::tools::glob_filter::build_glob_set;
5use aether_lspd::extensions_for_alias as extensions_for_type;
6use common::{CountSink, HasMatchSink, MatchCollectorSink, MatchData, OutputMode};
7use grep::{
8 regex::{RegexMatcher, RegexMatcherBuilder},
9 searcher::{BinaryDetection, Searcher, SearcherBuilder},
10};
11use ignore::{WalkBuilder, WalkState};
12use mcp_utils::display_meta::{ToolDisplayMeta, ToolResultMeta, basename};
13use schemars::JsonSchema;
14use serde::{Deserialize, Serialize};
15use std::path::Path;
16use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
17use std::sync::{Arc, Mutex};
18
19#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
20#[serde(rename_all = "camelCase")]
21pub struct GrepContentOutput {
22 pub matches: Vec<MatchData>,
24 pub total_matches: usize,
26 #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
28 #[schemars(skip)]
29 pub meta: Option<ToolResultMeta>,
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
33#[serde(rename_all = "camelCase")]
34pub struct GrepFilesOutput {
35 pub files: Vec<String>,
37 pub count: usize,
39 #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
41 #[schemars(skip)]
42 pub meta: Option<ToolResultMeta>,
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
46#[serde(rename_all = "camelCase")]
47pub struct GrepFileCount {
48 pub file: String,
49 pub count: usize,
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
53#[serde(rename_all = "camelCase")]
54pub struct GrepCountOutput {
55 pub counts: Vec<GrepFileCount>,
57 pub total: usize,
59 #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
61 #[schemars(skip)]
62 pub meta: Option<ToolResultMeta>,
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
66#[serde(tag = "type", rename_all = "camelCase")]
67#[schemars(extend("type" = "object"))]
68pub enum GrepOutput {
69 Content(GrepContentOutput),
70 Files(GrepFilesOutput),
71 Count(GrepCountOutput),
72}
73
74#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
75#[serde(rename_all = "camelCase")]
76pub struct GrepInput {
77 pub pattern: String,
79 pub path: Option<String>,
81 pub glob: Option<String>,
83 #[serde(rename = "type")]
85 pub file_type: Option<String>,
86 pub output_mode: Option<OutputMode>,
88 pub case_insensitive: Option<bool>,
90 pub line_numbers: Option<bool>,
92 pub context_before: Option<u32>,
94 pub context_after: Option<u32>,
96 pub context_around: Option<u32>,
98 pub head_limit: Option<usize>,
100 pub multiline: Option<bool>,
102}
103
104struct ParallelGrepState {
106 matches: Mutex<Vec<MatchData>>,
107 files_with_matches: Mutex<Vec<String>>,
108 file_counts: Mutex<Vec<GrepFileCount>>,
109 total_items: AtomicUsize,
110 max_items: usize,
111 limit_reached: AtomicBool,
112}
113
114impl ParallelGrepState {
115 fn new(max_items: usize) -> Self {
116 Self {
117 matches: Mutex::new(Vec::new()),
118 files_with_matches: Mutex::new(Vec::new()),
119 file_counts: Mutex::new(Vec::new()),
120 total_items: AtomicUsize::new(0),
121 max_items,
122 limit_reached: AtomicBool::new(false),
123 }
124 }
125}
126
127fn should_include_file(path: &Path, file_type: Option<&String>, glob_pattern: Option<&globset::GlobSet>) -> bool {
128 if let Some(glob_set) = glob_pattern {
130 return glob_set.is_match(path);
131 }
132
133 if let Some(ftype) = file_type {
135 if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
136 let extensions = extensions_for_type(ftype);
137 return extensions.contains(&ext);
138 }
139 return false;
140 }
141
142 true
143}
144
145pub async fn perform_grep(mut args: GrepInput) -> Result<GrepOutput, GrepError> {
146 if args.path.as_deref().is_some_and(|p| p.trim().is_empty()) {
147 args.path = None;
148 }
149
150 let glob_set = build_glob_set(args.glob.as_deref())?;
151
152 let matcher = build_matcher(&args.pattern, args.case_insensitive, args.multiline)?;
153
154 let output_mode = args.output_mode.unwrap_or(OutputMode::Content);
155 let line_numbers = args.line_numbers.unwrap_or(true);
156 let searcher_builder = build_searcher(&args);
157
158 let search_path = args.path.as_deref().unwrap_or(".");
159 let path_obj = Path::new(search_path);
160
161 if !path_obj.exists() {
162 return Err(GrepError::PathNotFound(search_path.to_string()));
163 }
164
165 let config = SearchConfig { output_mode, line_numbers, file_type: args.file_type.as_ref() };
166
167 let results = if path_obj.is_file() {
168 search_single_file(
169 path_obj,
170 search_path,
171 args.head_limit,
172 &matcher,
173 glob_set.as_ref(),
174 &config,
175 &searcher_builder,
176 )?
177 } else {
178 search_directory(search_path, &args, glob_set, matcher, &config, &searcher_builder)
179 };
180
181 Ok(build_grep_output(results, output_mode, &args.pattern, search_path))
182}
183
184struct SearchConfig<'a> {
186 output_mode: OutputMode,
187 line_numbers: bool,
188 file_type: Option<&'a String>,
189}
190
191struct SearchResults {
193 matches: Vec<MatchData>,
194 files_with_matches: Vec<String>,
195 file_counts: Vec<GrepFileCount>,
196}
197
198impl SearchResults {
199 fn empty() -> Self {
200 Self { matches: Vec::new(), files_with_matches: Vec::new(), file_counts: Vec::new() }
201 }
202}
203
204fn build_matcher(
205 pattern: &str,
206 case_insensitive: Option<bool>,
207 multiline: Option<bool>,
208) -> Result<RegexMatcher, GrepError> {
209 let mut matcher_builder = RegexMatcherBuilder::new();
210 matcher_builder.case_insensitive(case_insensitive.unwrap_or(false));
211
212 if multiline.unwrap_or(false) {
213 matcher_builder.multi_line(true).dot_matches_new_line(true);
214 }
215
216 matcher_builder.build(pattern).map_err(|e| GrepError::InvalidRegex(e.to_string()))
217}
218
219fn build_searcher(args: &GrepInput) -> SearcherBuilder {
220 let (context_before, context_after) = if let Some(around) = args.context_around {
221 (around, around)
222 } else {
223 (args.context_before.unwrap_or(0), args.context_after.unwrap_or(0))
224 };
225
226 let mut builder = SearcherBuilder::new();
227 builder
228 .binary_detection(BinaryDetection::quit(b'\x00'))
229 .line_number(args.line_numbers.unwrap_or(true))
230 .before_context(context_before as usize)
231 .after_context(context_after as usize);
232
233 if args.multiline.unwrap_or(false) {
234 builder.multi_line(true);
235 }
236
237 builder
238}
239
240fn search_single_file(
241 path_obj: &Path,
242 search_path: &str,
243 head_limit: Option<usize>,
244 matcher: &RegexMatcher,
245 glob_set: Option<&globset::GlobSet>,
246 config: &SearchConfig<'_>,
247 searcher_builder: &SearcherBuilder,
248) -> Result<SearchResults, GrepError> {
249 if !should_include_file(path_obj, config.file_type, glob_set) {
250 return Ok(SearchResults::empty());
251 }
252
253 let mut searcher = searcher_builder.build();
254 let mut results = SearchResults::empty();
255
256 match config.output_mode {
257 OutputMode::Content => {
258 let mut sink = MatchCollectorSink::with_max_results(path_obj, config.line_numbers, head_limit);
259 searcher.search_path(matcher, path_obj, &mut sink).map_err(|e| GrepError::SearchFailed(e.to_string()))?;
260 results.matches = sink.matches;
261 }
262 OutputMode::FilesWithMatches => {
263 let mut sink = HasMatchSink::new();
264 searcher.search_path(matcher, path_obj, &mut sink).map_err(|e| GrepError::SearchFailed(e.to_string()))?;
265 if sink.has_match {
266 results.files_with_matches.push(search_path.to_string());
267 }
268 }
269 OutputMode::Count => {
270 let mut sink = CountSink::new();
271 searcher.search_path(matcher, path_obj, &mut sink).map_err(|e| GrepError::SearchFailed(e.to_string()))?;
272 if sink.count > 0 {
273 results.file_counts.push(GrepFileCount { file: search_path.to_string(), count: sink.count });
274 }
275 }
276 }
277
278 Ok(results)
279}
280
281fn search_directory(
282 search_path: &str,
283 args: &GrepInput,
284 glob_set: Option<globset::GlobSet>,
285 matcher: RegexMatcher,
286 config: &SearchConfig<'_>,
287 searcher_builder: &SearcherBuilder,
288) -> SearchResults {
289 let walker = WalkBuilder::new(search_path).hidden(false).git_ignore(true).build_parallel();
290
291 let max_items = args.head_limit.unwrap_or(usize::MAX);
292 let state = Arc::new(ParallelGrepState::new(max_items));
293 let matcher = Arc::new(matcher);
294 let glob_set = Arc::new(glob_set);
295 let file_type = args.file_type.clone();
296 let output_mode = config.output_mode;
297 let line_numbers = config.line_numbers;
298
299 walker.run(|| {
300 let state = state.clone();
301 let matcher = matcher.clone();
302 let glob_set = glob_set.clone();
303 let file_type = file_type.clone();
304 let mut thread_searcher = searcher_builder.build();
305
306 Box::new(move |result| {
307 let filter = FileFilter { glob_set: glob_set.as_ref().as_ref(), file_type: file_type.as_ref() };
308 search_directory_entry(result, &state, &matcher, &filter, &mut thread_searcher, output_mode, line_numbers)
309 })
310 });
311
312 let mut results = SearchResults {
313 matches: state.matches.lock().unwrap().clone(),
314 files_with_matches: state.files_with_matches.lock().unwrap().clone(),
315 file_counts: state.file_counts.lock().unwrap().clone(),
316 };
317
318 results.matches.sort_by(|a, b| a.file.cmp(&b.file));
319 results.files_with_matches.sort();
320 results.file_counts.sort_by(|a, b| a.file.cmp(&b.file));
321
322 if let Some(limit) = args.head_limit {
323 results.matches.truncate(limit);
324 results.files_with_matches.truncate(limit);
325 results.file_counts.truncate(limit);
326 }
327
328 results
329}
330
331struct FileFilter<'a> {
333 glob_set: Option<&'a globset::GlobSet>,
334 file_type: Option<&'a String>,
335}
336
337fn search_directory_entry(
338 result: std::result::Result<ignore::DirEntry, ignore::Error>,
339 state: &ParallelGrepState,
340 matcher: &RegexMatcher,
341 filter: &FileFilter<'_>,
342 searcher: &mut Searcher,
343 output_mode: OutputMode,
344 line_numbers: bool,
345) -> WalkState {
346 if state.limit_reached.load(Ordering::Relaxed) {
347 return WalkState::Quit;
348 }
349
350 let Ok(entry) = result else {
351 return WalkState::Continue;
352 };
353
354 if !entry.file_type().is_some_and(|ft| ft.is_file()) {
355 return WalkState::Continue;
356 }
357
358 if !should_include_file(entry.path(), filter.file_type, filter.glob_set) {
359 return WalkState::Continue;
360 }
361
362 match output_mode {
363 OutputMode::Content => {
364 let mut sink = MatchCollectorSink::with_max_results(entry.path(), line_numbers, None);
365 if searcher.search_path(matcher, entry.path(), &mut sink).is_ok() && !sink.matches.is_empty() {
366 let new_count = state.total_items.fetch_add(sink.matches.len(), Ordering::SeqCst) + sink.matches.len();
367 if new_count >= state.max_items {
368 state.limit_reached.store(true, Ordering::Release);
369 }
370 if let Ok(mut matches) = state.matches.lock() {
371 matches.extend(sink.matches);
372 }
373 }
374 }
375 OutputMode::FilesWithMatches => {
376 let mut sink = HasMatchSink::new();
377 if searcher.search_path(matcher, entry.path(), &mut sink).is_ok() && sink.has_match {
378 let new_count = state.total_items.fetch_add(1, Ordering::SeqCst) + 1;
379 if new_count >= state.max_items {
380 state.limit_reached.store(true, Ordering::Release);
381 }
382 if let Ok(mut files) = state.files_with_matches.lock() {
383 files.push(entry.path().to_string_lossy().to_string());
384 }
385 }
386 }
387 OutputMode::Count => {
388 let mut sink = CountSink::new();
389 if searcher.search_path(matcher, entry.path(), &mut sink).is_ok() && sink.count > 0 {
390 let new_count = state.total_items.fetch_add(1, Ordering::SeqCst) + 1;
391 if new_count >= state.max_items {
392 state.limit_reached.store(true, Ordering::Release);
393 }
394 if let Ok(mut counts) = state.file_counts.lock() {
395 counts.push(GrepFileCount { file: entry.path().to_string_lossy().to_string(), count: sink.count });
396 }
397 }
398 }
399 }
400
401 WalkState::Continue
402}
403
404fn build_grep_output(results: SearchResults, output_mode: OutputMode, pattern: &str, search_path: &str) -> GrepOutput {
405 let match_count = match output_mode {
406 OutputMode::Content => results.matches.len(),
407 OutputMode::FilesWithMatches => results.files_with_matches.len(),
408 OutputMode::Count => results.file_counts.iter().map(|fc| fc.count).sum(),
409 };
410
411 let display_meta =
412 ToolDisplayMeta::new("Grep", format!("'{}' in {} ({} matches)", pattern, basename(search_path), match_count));
413 let meta = Some(display_meta.into());
414
415 match output_mode {
416 OutputMode::Content => {
417 GrepOutput::Content(GrepContentOutput { matches: results.matches, total_matches: match_count, meta })
418 }
419 OutputMode::FilesWithMatches => {
420 GrepOutput::Files(GrepFilesOutput { files: results.files_with_matches, count: match_count, meta })
421 }
422 OutputMode::Count => {
423 GrepOutput::Count(GrepCountOutput { counts: results.file_counts, total: match_count, meta })
424 }
425 }
426}
427
428#[cfg(test)]
429mod tests {
430 use super::*;
431 use std::fs;
432 use tempfile::TempDir;
433
434 fn input(pattern: &str) -> GrepInput {
435 GrepInput {
436 pattern: pattern.to_string(),
437 path: None,
438 glob: None,
439 file_type: None,
440 output_mode: None,
441 case_insensitive: None,
442 line_numbers: None,
443 context_before: None,
444 context_after: None,
445 context_around: None,
446 head_limit: None,
447 multiline: None,
448 }
449 }
450
451 fn create_test_dir() -> TempDir {
452 let dir = TempDir::new().expect("Failed to create temp dir");
453 let p = dir.path();
454 fs::write(p.join("test.rs"), "fn main() {\n println!(\"Hello, world!\");\n let x = 42;\n}").unwrap();
455 fs::write(p.join("script.py"), "def hello():\n print(\"Hello, world!\")\n x = 42\n").unwrap();
456 fs::write(p.join("app.js"), "function hello() {\n console.log(\"Hello, world!\");\n const x = 42;\n}")
457 .unwrap();
458 dir
459 }
460
461 fn unwrap_content(output: GrepOutput) -> GrepContentOutput {
462 match output {
463 GrepOutput::Content(c) => c,
464 other => panic!("Expected Content, got {other:?}"),
465 }
466 }
467
468 fn unwrap_files(output: GrepOutput) -> GrepFilesOutput {
469 match output {
470 GrepOutput::Files(f) => f,
471 other => panic!("Expected Files, got {other:?}"),
472 }
473 }
474
475 fn unwrap_count(output: GrepOutput) -> GrepCountOutput {
476 match output {
477 GrepOutput::Count(c) => c,
478 other => panic!("Expected Count, got {other:?}"),
479 }
480 }
481
482 async fn grep(args: GrepInput) -> GrepOutput {
483 perform_grep(args).await.expect("grep failed")
484 }
485
486 #[tokio::test]
487 async fn test_file_type_and_glob_filtering() {
488 let temp_dir = create_test_dir();
489 let path = temp_dir.path().to_str().unwrap().to_string();
490
491 let cases: Vec<(Option<String>, Option<String>, &str)> =
492 vec![(Some("rust".into()), None, "test.rs"), (None, Some("*.py".into()), "script.py")];
493 for (file_type, glob, expected_file) in cases {
494 let mut args = input("hello");
495 args.path = Some(path.clone());
496 args.file_type = file_type;
497 args.glob = glob;
498 args.output_mode = Some(OutputMode::Content);
499 args.case_insensitive = Some(true);
500 args.line_numbers = Some(true);
501
502 let content = unwrap_content(grep(args).await);
503 assert!(content.total_matches > 0, "no matches for {expected_file}");
504 assert!(
505 content.matches.iter().all(|m| m.file.contains(expected_file)),
506 "expected all matches in {expected_file}"
507 );
508 }
509 }
510
511 #[tokio::test]
512 async fn test_files_with_matches_output() {
513 let temp_dir = create_test_dir();
514 let mut args = input("hello");
515 args.path = Some(temp_dir.path().to_str().unwrap().to_string());
516 args.output_mode = Some(OutputMode::FilesWithMatches);
517 args.case_insensitive = Some(true);
518
519 let files = unwrap_files(grep(args).await);
520 assert!(files.count >= 2);
521 assert!(files.files.iter().any(|f| f.contains(".py")));
522 assert!(files.files.iter().any(|f| f.contains(".js")));
523 }
524
525 #[tokio::test]
526 async fn test_count_output() {
527 let temp_dir = create_test_dir();
528 let mut args = input("hello");
529 args.path = Some(temp_dir.path().to_str().unwrap().to_string());
530 args.output_mode = Some(OutputMode::Count);
531 args.case_insensitive = Some(true);
532
533 let count = unwrap_count(grep(args).await);
534 assert!(count.counts.len() >= 2);
535 assert!(count.total >= 2);
536 assert!(count.counts.iter().all(|fc| fc.count > 0));
537 }
538
539 #[tokio::test]
540 async fn test_head_limit() {
541 let temp_dir = create_test_dir();
542 let mut args = input("hello");
543 args.path = Some(temp_dir.path().to_str().unwrap().to_string());
544 args.output_mode = Some(OutputMode::Content);
545 args.case_insensitive = Some(true);
546 args.line_numbers = Some(true);
547 args.head_limit = Some(1);
548
549 let content = unwrap_content(grep(args).await);
550 assert!(content.total_matches <= 1);
551 }
552
553 #[tokio::test]
554 async fn test_multiline_mode() {
555 let temp_dir = TempDir::new().expect("Failed to create temp dir");
556 let file = temp_dir.path().join("multiline.txt");
557 fs::write(&file, "start\nmiddle content\nend").unwrap();
558
559 let mut args = input(r"start.*end");
560 args.path = Some(file.to_str().unwrap().to_string());
561 args.output_mode = Some(OutputMode::Content);
562 args.line_numbers = Some(true);
563 args.multiline = Some(true);
564
565 let content = unwrap_content(grep(args).await);
566 assert!(content.total_matches > 0);
567 }
568
569 #[tokio::test]
570 async fn test_context_lines() {
571 let temp_dir = TempDir::new().expect("Failed to create temp dir");
572 let file = temp_dir.path().join("context.txt");
573 fs::write(&file, "line 1\nline 2\ntarget\nline 4\nline 5").unwrap();
574
575 let mut args = input("target");
576 args.path = Some(file.to_str().unwrap().to_string());
577 args.output_mode = Some(OutputMode::Content);
578 args.line_numbers = Some(true);
579 args.context_before = Some(1);
580 args.context_after = Some(1);
581
582 let content = unwrap_content(grep(args).await);
583 assert!(content.total_matches > 0);
584 let m = &content.matches[0];
585 let before = m.before_context.as_ref().expect("missing before_context");
586 let after = m.after_context.as_ref().expect("missing after_context");
587 assert_eq!(before.len(), 1);
588 assert_eq!(after.len(), 1);
589 }
590
591 #[test]
592 fn grep_input_accepts_camel_case_fields() {
593 let args: GrepInput = serde_json::from_value(serde_json::json!({
594 "pattern": "hello",
595 "path": "/tmp",
596 "type": "rust",
597 "outputMode": "filesWithMatches",
598 "caseInsensitive": true,
599 "lineNumbers": true,
600 "contextBefore": 1,
601 "contextAfter": 2,
602 "contextAround": 3,
603 "headLimit": 10,
604 "multiline": false
605 }))
606 .unwrap();
607
608 assert_eq!(args.file_type, Some("rust".to_string()));
609 assert!(matches!(args.output_mode, Some(OutputMode::FilesWithMatches)));
610 assert_eq!(args.case_insensitive, Some(true));
611 assert_eq!(args.line_numbers, Some(true));
612 assert_eq!(args.context_before, Some(1));
613 assert_eq!(args.context_after, Some(2));
614 assert_eq!(args.context_around, Some(3));
615 assert_eq!(args.head_limit, Some(10));
616 assert_eq!(args.multiline, Some(false));
617 }
618
619 #[test]
620 fn output_mode_files_with_matches_snake_case_alias() {
621 let args: GrepInput = serde_json::from_value(serde_json::json!({
622 "pattern": "hello",
623 "outputMode": "files_with_matches"
624 }))
625 .unwrap();
626 assert!(matches!(args.output_mode, Some(OutputMode::FilesWithMatches)));
627 }
628
629 #[tokio::test]
630 async fn empty_path_treated_as_cwd() {
631 let temp_dir = create_test_dir();
632 let _guard = std::env::set_current_dir(temp_dir.path());
633
634 let mut args = input("hello");
635 args.path = Some(String::new());
636 args.output_mode = Some(OutputMode::Content);
637 args.case_insensitive = Some(true);
638
639 assert!(perform_grep(args).await.is_ok());
640 }
641
642 #[tokio::test]
643 async fn nonexistent_path_returns_error() {
644 let mut args = input("hello");
645 args.path = Some("/no/such/path/exists".to_string());
646
647 let err = perform_grep(args).await.unwrap_err();
648 assert!(matches!(err, GrepError::PathNotFound(_)), "Expected PathNotFound, got: {err}");
649 }
650}