Skip to main content

mcp_servers/coding/tools/grep/
mod.rs

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    /// Matching lines with context
23    pub matches: Vec<MatchData>,
24    /// Total number of matches
25    pub total_matches: usize,
26    /// Display metadata for human-friendly rendering
27    #[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    /// Files containing matches
36    pub files: Vec<String>,
37    /// Number of files with matches
38    pub count: usize,
39    /// Display metadata for human-friendly rendering
40    #[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    /// Match counts per file
56    pub counts: Vec<GrepFileCount>,
57    /// Total matches across all files
58    pub total: usize,
59    /// Display metadata for human-friendly rendering
60    #[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    /// The regular expression pattern to search for in file contents
78    pub pattern: String,
79    /// Absolute path to a file or directory to search in. Defaults to the current working directory.
80    pub path: Option<String>,
81    /// Glob pattern to filter files (e.g. "*.js", "*.{ts,tsx}")
82    pub glob: Option<String>,
83    /// File type to search (e.g. "js", "py", "rust"). More efficient than glob for standard file types.
84    #[serde(rename = "type")]
85    pub file_type: Option<String>,
86    /// Output mode: "content" (default, shows matching lines), "filesWithMatches" (file paths only), or "count" (match counts per file)
87    pub output_mode: Option<OutputMode>,
88    /// Case insensitive search
89    pub case_insensitive: Option<bool>,
90    /// Show line numbers in output (content mode only). Defaults to true.
91    pub line_numbers: Option<bool>,
92    /// Number of lines to show before each match (content mode only)
93    pub context_before: Option<u32>,
94    /// Number of lines to show after each match (content mode only)
95    pub context_after: Option<u32>,
96    /// Number of lines to show before and after each match (content mode only). Overrides contextBefore/contextAfter.
97    pub context_around: Option<u32>,
98    /// Limit output to first N entries. In content mode limits match lines, in filesWithMatches mode limits file paths, in count mode limits file entries.
99    pub head_limit: Option<usize>,
100    /// Enable multiline mode where . matches newlines and patterns can span lines
101    pub multiline: Option<bool>,
102}
103
104/// Thread-safe state for parallel grep execution
105struct 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    // Check glob pattern first (more specific)
129    if let Some(glob_set) = glob_pattern {
130        return glob_set.is_match(path);
131    }
132
133    // Check file type filter
134    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
184/// Common search parameters shared across single-file and directory search paths.
185struct SearchConfig<'a> {
186    output_mode: OutputMode,
187    line_numbers: bool,
188    file_type: Option<&'a String>,
189}
190
191/// Collected search results across all searched files.
192struct 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
331/// Filter criteria for file selection during directory walks.
332struct 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}