mcp_servers/coding/tools/ast_grep/
mod.rs

1use crate::coding::error::AstGrepError;
2use crate::coding::tools::glob_filter::build_glob_set;
3use ast_grep_core::Doc;
4use ast_grep_core::matcher::{NodeMatch, Pattern};
5use ast_grep_language::{Language, LanguageExt, SupportLang};
6use globset::GlobSet;
7use ignore::{WalkBuilder, WalkState};
8use mcp_utils::display_meta::{ToolDisplayMeta, ToolResultMeta, basename};
9use regex::Regex;
10use schemars::JsonSchema;
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13use std::fs::read_to_string;
14use std::path::Path;
15use std::sync::{Arc, Mutex};
16use tokio::task::spawn_blocking;
17
18#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
19#[serde(rename_all = "camelCase")]
20pub struct AstGrepInput {
21    /// ast-grep pattern code.
22    pub pattern: String,
23    /// Language alias understood by ast-grep, e.g. "rs", "rust", "ts", "tsx", "py", "js".
24    pub language: String,
25    /// File or directory to search. Defaults to the coding server workspace root.
26    pub path: Option<String>,
27    /// Optional glob filter, e.g. "**/*.rs" or "*.{ts,tsx}".
28    pub glob: Option<String>,
29    /// Lines before each match.
30    pub context_before: Option<u32>,
31    /// Lines after each match.
32    pub context_after: Option<u32>,
33    /// Lines before and after each match. Overrides contextBefore/contextAfter.
34    pub context_around: Option<u32>,
35    /// Maximum number of matches to return.
36    pub head_limit: Option<usize>,
37    /// Optional regex constraints on metavariable captures.
38    /// Key is the metavariable name without `$`; value is a regex the captured text must match.
39    /// Example: `{"CRATE": "^crossterm"}` with pattern `use $CRATE;`.
40    pub constraints: Option<HashMap<String, String>>,
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
44#[serde(rename_all = "camelCase")]
45pub struct AstGrepOutput {
46    pub matches: Vec<AstGrepMatch>,
47    pub count: usize,
48    pub truncated: bool,
49    pub language: String,
50    pub search_path: String,
51    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
52    #[schemars(skip)]
53    pub meta: Option<ToolResultMeta>,
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
57#[serde(rename_all = "camelCase")]
58pub struct AstGrepMatch {
59    pub file: String,
60    pub range: AstGrepRange,
61    pub text: String,
62    pub captures: Vec<AstGrepCapture>,
63    pub before_context: Option<Vec<String>>,
64    pub after_context: Option<Vec<String>>,
65}
66
67#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
68#[serde(rename_all = "camelCase")]
69pub struct AstGrepRange {
70    /// 1-based line number.
71    pub start_line: usize,
72    /// 1-based character column.
73    pub start_column: usize,
74    /// 1-based line number.
75    pub end_line: usize,
76    /// 1-based character column.
77    pub end_column: usize,
78    /// 0-based byte offset.
79    pub start_byte: usize,
80    /// 0-based exclusive byte offset.
81    pub end_byte: usize,
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
85#[serde(rename_all = "camelCase")]
86pub struct AstGrepCapture {
87    pub name: String,
88    pub text: String,
89}
90
91pub async fn perform_ast_grep(mut args: AstGrepInput) -> Result<AstGrepOutput, AstGrepError> {
92    spawn_blocking(move || {
93        if args.path.as_deref().is_some_and(|p| p.trim().is_empty()) {
94            args.path = None;
95        }
96
97        let lang: SupportLang =
98            args.language.parse().map_err(|_| AstGrepError::UnsupportedLanguage(args.language.clone()))?;
99
100        let pattern = Pattern::try_new(&args.pattern, lang).map_err(|e| AstGrepError::InvalidPattern(e.to_string()))?;
101
102        let constraint_regexes = compile_constraints(args.constraints.as_ref())?;
103        let glob_set = build_glob_set(args.glob.as_deref())?;
104        let search_path = args.path.as_deref().unwrap_or(".");
105        let path = Path::new(search_path);
106        if !path.exists() {
107            return Err(AstGrepError::PathNotFound(search_path.to_string()));
108        }
109
110        let (context_before, context_after) = context_counts(&args);
111        let limit = args.head_limit.unwrap_or(usize::MAX);
112        let mut matches = if path.is_file() {
113            if glob_matches(path, path, glob_set.as_ref()) {
114                search_file(path, lang, &pattern, &constraint_regexes, context_before, context_after)?
115            } else {
116                Vec::new()
117            }
118        } else {
119            search_directory(path, lang, pattern, glob_set, &constraint_regexes, context_before, context_after)
120        };
121
122        matches.sort_by(|a, b| {
123            a.file
124                .cmp(&b.file)
125                .then(a.range.start_byte.cmp(&b.range.start_byte))
126                .then(a.range.end_byte.cmp(&b.range.end_byte))
127        });
128        let truncated = matches.len() > limit;
129        matches.truncate(limit);
130
131        let count = matches.len();
132        let meta = ToolDisplayMeta::new(
133            "AST grep",
134            format!("'{}' in {} ({count} matches)", args.pattern, basename(search_path)),
135        );
136
137        Ok(AstGrepOutput {
138            matches,
139            count,
140            truncated,
141            language: lang.to_string(),
142            search_path: search_path.to_string(),
143            meta: Some(meta.into()),
144        })
145    })
146    .await
147    .map_err(|error| AstGrepError::SearchFailed(error.to_string()))?
148}
149
150fn search_directory(
151    search_path: &Path,
152    lang: SupportLang,
153    pattern: Pattern,
154    glob: Option<GlobSet>,
155    constraints: &HashMap<String, Regex>,
156    context_before: usize,
157    context_after: usize,
158) -> Vec<AstGrepMatch> {
159    let matches = Arc::new(Mutex::new(Vec::new()));
160    let pattern = Arc::new(pattern);
161    let glob = Arc::new(glob);
162    let constraints = Arc::new(constraints.clone());
163    let search_root = Arc::new(search_path.to_path_buf());
164    let mut walker = WalkBuilder::new(search_path);
165    walker.follow_links(false);
166
167    walker.build_parallel().run(|| {
168        let matches = matches.clone();
169        let pattern = pattern.clone();
170        let glob = glob.clone();
171        let constraints = constraints.clone();
172        let search_root = search_root.clone();
173
174        Box::new(move |result| {
175            let Ok(entry) = result else {
176                return WalkState::Continue;
177            };
178
179            if !entry.file_type().is_some_and(|file_type| file_type.is_file()) {
180                return WalkState::Continue;
181            }
182
183            if !file_matches_filters(entry.path(), search_root.as_ref().as_path(), lang, glob.as_ref().as_ref()) {
184                return WalkState::Continue;
185            }
186
187            let Ok(file_matches) =
188                search_file(entry.path(), lang, &pattern, &constraints, context_before, context_after)
189            else {
190                return WalkState::Continue;
191            };
192
193            if !file_matches.is_empty()
194                && let Ok(mut matches) = matches.lock()
195            {
196                matches.extend(file_matches);
197            }
198
199            WalkState::Continue
200        })
201    });
202
203    matches.lock().map(|matches| matches.clone()).unwrap_or_default()
204}
205
206fn file_matches_filters(path: &Path, search_root: &Path, lang: SupportLang, glob: Option<&GlobSet>) -> bool {
207    <SupportLang as Language>::from_path(path) == Some(lang) && glob_matches(path, search_root, glob)
208}
209
210fn glob_matches(path: &Path, search_root: &Path, glob: Option<&GlobSet>) -> bool {
211    let Some(glob) = glob else {
212        return true;
213    };
214
215    glob.is_match(path) || path.strip_prefix(search_root).is_ok_and(|relative| glob.is_match(relative))
216}
217
218fn search_file(
219    path: &Path,
220    lang: SupportLang,
221    pattern: &Pattern,
222    constraints: &HashMap<String, Regex>,
223    context_before: usize,
224    context_after: usize,
225) -> Result<Vec<AstGrepMatch>, AstGrepError> {
226    let source = read_to_string(path)
227        .map_err(|error| AstGrepError::ReadFailed { path: path.display().to_string(), reason: error.to_string() })?;
228
229    let lines: Vec<&str> = source.lines().collect();
230    let root = lang.ast_grep(&source);
231    let mut matches = Vec::new();
232    for node_match in root.root().find_all(pattern.clone()) {
233        if !constraints_match(&node_match, constraints) {
234            continue;
235        }
236        matches.push(build_match(path, &lines, &node_match, context_before, context_after));
237    }
238    Ok(matches)
239}
240
241fn build_match<T: Doc>(
242    path: &Path,
243    lines: &[&str],
244    node_match: &NodeMatch<'_, T>,
245    context_before: usize,
246    context_after: usize,
247) -> AstGrepMatch {
248    let start = node_match.start_pos();
249    let end = node_match.end_pos();
250    let range = node_match.range();
251    let start_line = start.line() + 1;
252    let end_line = end.line() + 1;
253    let (before_context, after_context) = context_for_match(lines, start_line, end_line, context_before, context_after);
254
255    AstGrepMatch {
256        file: path.to_string_lossy().to_string(),
257        range: AstGrepRange {
258            start_line,
259            start_column: start.column(node_match.get_node()) + 1,
260            end_line,
261            end_column: end.column(node_match.get_node()) + 1,
262            start_byte: range.start,
263            end_byte: range.end,
264        },
265        text: node_match.text().to_string(),
266        captures: captures(node_match),
267        before_context,
268        after_context,
269    }
270}
271
272fn captures<T: Doc>(node_match: &NodeMatch<'_, T>) -> Vec<AstGrepCapture> {
273    let env: HashMap<String, String> = HashMap::from(node_match.get_env().clone());
274    let mut captures: Vec<AstGrepCapture> = env.into_iter().map(|(name, text)| AstGrepCapture { name, text }).collect();
275    captures.sort_by(|a, b| a.name.cmp(&b.name));
276    captures
277}
278
279fn context_for_match(
280    lines: &[&str],
281    start_line: usize,
282    end_line: usize,
283    before: usize,
284    after: usize,
285) -> (Option<Vec<String>>, Option<Vec<String>>) {
286    let before_context = if before == 0 {
287        None
288    } else {
289        let start = start_line.saturating_sub(before).max(1);
290        let context: Vec<String> =
291            lines[start - 1..start_line.saturating_sub(1)].iter().map(|line| (*line).to_string()).collect();
292        Some(context)
293    };
294    let after_context = if after == 0 {
295        None
296    } else if end_line >= lines.len() {
297        Some(Vec::new())
298    } else {
299        let end = (end_line + after).min(lines.len());
300        let context: Vec<String> = lines[end_line..end].iter().map(|line| (*line).to_string()).collect();
301        Some(context)
302    };
303    (before_context, after_context)
304}
305
306fn compile_constraints(constraints: Option<&HashMap<String, String>>) -> Result<HashMap<String, Regex>, AstGrepError> {
307    let Some(constraints) = constraints else { return Ok(HashMap::new()) };
308    constraints
309        .iter()
310        .map(|(name, pattern)| {
311            Regex::new(pattern)
312                .map(|regex| (name.clone(), regex))
313                .map_err(|error| AstGrepError::InvalidConstraintRegex { name: name.clone(), reason: error.to_string() })
314        })
315        .collect()
316}
317
318fn constraints_match<T: Doc>(node_match: &NodeMatch<'_, T>, constraints: &HashMap<String, Regex>) -> bool {
319    let env = node_match.get_env();
320    constraints.iter().all(|(name, regex)| env.get_match(name).is_some_and(|node| regex.is_match(&node.text())))
321}
322
323fn context_counts(args: &AstGrepInput) -> (usize, usize) {
324    if let Some(around) = args.context_around {
325        (around as usize, around as usize)
326    } else {
327        (args.context_before.unwrap_or(0) as usize, args.context_after.unwrap_or(0) as usize)
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334    use std::fs;
335    use tempfile::TempDir;
336
337    fn input(pattern: &str, path: &Path) -> AstGrepInput {
338        AstGrepInput {
339            pattern: pattern.to_string(),
340            language: "rs".to_string(),
341            path: Some(path.to_string_lossy().to_string()),
342            glob: None,
343            context_before: None,
344            context_after: None,
345            context_around: None,
346            head_limit: None,
347            constraints: None,
348        }
349    }
350
351    #[tokio::test]
352    async fn finds_rust_function_pattern() {
353        let temp = TempDir::new().unwrap();
354        fs::write(temp.path().join("lib.rs"), "fn foo() {}\nfn bar() { let x = 1; }\n").unwrap();
355
356        let result = perform_ast_grep(input("fn $NAME() { $$$BODY }", temp.path())).await.unwrap();
357
358        assert_eq!(result.count, 2);
359        assert!(result.matches.iter().any(|m| m.text == "fn foo() {}"));
360        assert!(result.matches.iter().any(|m| m.text == "fn bar() { let x = 1; }"));
361    }
362
363    #[tokio::test]
364    async fn returns_metavariable_captures() {
365        let temp = TempDir::new().unwrap();
366        fs::write(temp.path().join("lib.rs"), "fn foo() {}\n").unwrap();
367
368        let result = perform_ast_grep(input("fn $NAME() {}", temp.path())).await.unwrap();
369
370        assert_eq!(result.matches[0].captures[0].name, "NAME");
371        assert_eq!(result.matches[0].captures[0].text, "foo");
372    }
373
374    #[tokio::test]
375    async fn filters_directory_by_language() {
376        let temp = TempDir::new().unwrap();
377        fs::write(temp.path().join("lib.rs"), "fn foo() {}\n").unwrap();
378        fs::write(temp.path().join("test.py"), "fn fake() {}\n").unwrap();
379
380        let result = perform_ast_grep(input("fn $NAME() {}", temp.path())).await.unwrap();
381
382        assert_eq!(result.count, 1);
383        assert!(result.matches[0].file.ends_with("lib.rs"));
384    }
385
386    #[tokio::test]
387    async fn applies_glob_filter() {
388        let temp = TempDir::new().unwrap();
389        fs::create_dir(temp.path().join("src")).unwrap();
390        fs::create_dir(temp.path().join("tests")).unwrap();
391        fs::write(temp.path().join("src/lib.rs"), "fn src_fn() {}\n").unwrap();
392        fs::write(temp.path().join("tests/lib.rs"), "fn test_fn() {}\n").unwrap();
393        let mut args = input("fn $NAME() {}", temp.path());
394        args.glob = Some("src/**/*.rs".to_string());
395        let result = perform_ast_grep(args).await.unwrap();
396        assert_eq!(result.count, 1);
397        assert!(result.matches[0].file.ends_with("src/lib.rs"));
398    }
399
400    #[tokio::test]
401    async fn invalid_language_returns_error() {
402        let temp = TempDir::new().unwrap();
403        let mut args = input("fn $NAME() {}", temp.path());
404        args.language = "not-a-language".to_string();
405        let result = perform_ast_grep(args).await;
406        assert!(matches!(result, Err(AstGrepError::UnsupportedLanguage(_))));
407    }
408
409    #[tokio::test]
410    async fn invalid_pattern_returns_error() {
411        let temp = TempDir::new().unwrap();
412        let result = perform_ast_grep(input("", temp.path())).await;
413
414        assert!(matches!(result, Err(AstGrepError::InvalidPattern(_))));
415    }
416
417    #[tokio::test]
418    async fn head_limit_truncates_results() {
419        let temp = TempDir::new().unwrap();
420        fs::write(temp.path().join("lib.rs"), "fn a() {}\nfn b() {}\n").unwrap();
421        let mut args = input("fn $NAME() {}", temp.path());
422        args.head_limit = Some(1);
423        let result = perform_ast_grep(args).await.unwrap();
424        assert_eq!(result.count, 1);
425        assert!(result.truncated);
426    }
427
428    #[tokio::test]
429    async fn head_limit_zero_returns_no_matches_but_truncated() {
430        let temp = TempDir::new().unwrap();
431        fs::write(temp.path().join("lib.rs"), "fn a() {}\n").unwrap();
432        let mut args = input("fn $NAME() {}", temp.path());
433        args.head_limit = Some(0);
434        let result = perform_ast_grep(args).await.unwrap();
435        assert_eq!(result.count, 0);
436        assert!(result.truncated);
437    }
438
439    #[tokio::test]
440    async fn nonexistent_path_returns_error() {
441        let result = perform_ast_grep(input("fn $NAME() {}", Path::new("/no/such/path/exists"))).await;
442
443        assert!(matches!(result, Err(AstGrepError::PathNotFound(_))));
444    }
445
446    #[tokio::test]
447    async fn explicit_unreadable_file_returns_error() {
448        let temp = TempDir::new().unwrap();
449        let file = temp.path().join("invalid.rs");
450        fs::write(&file, [0xff, 0xfe, 0x00, 0x01]).unwrap();
451        let result = perform_ast_grep(input("fn $NAME() {}", &file)).await;
452        assert!(matches!(result, Err(AstGrepError::ReadFailed { .. })));
453    }
454
455    #[tokio::test]
456    async fn range_is_one_based_and_byte_offsets_are_zero_based() {
457        let temp = TempDir::new().unwrap();
458        fs::write(temp.path().join("lib.rs"), "\nfn foo() {}\n").unwrap();
459        let result = perform_ast_grep(input("fn $NAME() {}", temp.path())).await.unwrap();
460        let range = &result.matches[0].range;
461
462        assert_eq!(range.start_line, 2);
463        assert_eq!(range.start_column, 1);
464        assert_eq!(range.start_byte, 1);
465    }
466
467    #[tokio::test]
468    async fn directory_search_skips_hidden_files() {
469        let temp = TempDir::new().unwrap();
470        fs::write(temp.path().join("visible.rs"), "fn visible() {}\n").unwrap();
471        fs::write(temp.path().join(".hidden.rs"), "fn hidden() {}\n").unwrap();
472        let result = perform_ast_grep(input("fn $NAME() {}", temp.path())).await.unwrap();
473
474        assert_eq!(result.count, 1);
475        assert_eq!(result.matches[0].text, "fn visible() {}");
476    }
477
478    #[tokio::test]
479    async fn directory_search_respects_gitignore() {
480        let temp = TempDir::new().unwrap();
481        fs::create_dir(temp.path().join(".git")).unwrap();
482        fs::write(temp.path().join(".gitignore"), "ignored.rs\n").unwrap();
483        fs::write(temp.path().join("included.rs"), "fn included() {}\n").unwrap();
484        fs::write(temp.path().join("ignored.rs"), "fn ignored() {}\n").unwrap();
485
486        let result = perform_ast_grep(input("fn $NAME() {}", temp.path())).await.unwrap();
487
488        assert_eq!(result.count, 1);
489        assert_eq!(result.matches[0].text, "fn included() {}");
490    }
491
492    #[tokio::test]
493    async fn includes_context_lines() {
494        let temp = TempDir::new().unwrap();
495        fs::write(temp.path().join("lib.rs"), "before\nfn foo() {}\nafter\n").unwrap();
496        let mut args = input("fn $NAME() {}", temp.path());
497        args.context_around = Some(1);
498
499        let result = perform_ast_grep(args).await.unwrap();
500
501        assert_eq!(result.matches[0].before_context.as_ref().unwrap(), &vec!["before".to_string()]);
502        assert_eq!(result.matches[0].after_context.as_ref().unwrap(), &vec!["after".to_string()]);
503    }
504
505    #[tokio::test]
506    async fn filters_matches_by_constraint() {
507        let temp = TempDir::new().unwrap();
508        fs::write(temp.path().join("lib.rs"), "use crossterm::execute;\nuse serde::Serialize;\n").unwrap();
509
510        let mut args = input("use $CRATE;", temp.path());
511        args.constraints = Some(HashMap::from([("CRATE".to_string(), "^crossterm".to_string())]));
512
513        let result = perform_ast_grep(args).await.unwrap();
514
515        assert_eq!(result.count, 1);
516        assert_eq!(result.matches[0].text, "use crossterm::execute;");
517    }
518
519    #[tokio::test]
520    async fn constraint_with_no_matches_returns_empty() {
521        let temp = TempDir::new().unwrap();
522        fs::write(temp.path().join("lib.rs"), "fn foo() {}\n").unwrap();
523
524        let mut args = input("fn $NAME() {}", temp.path());
525        args.constraints = Some(HashMap::from([("NAME".to_string(), "^bar".to_string())]));
526
527        let result = perform_ast_grep(args).await.unwrap();
528
529        assert_eq!(result.count, 0);
530    }
531
532    #[tokio::test]
533    async fn constraint_on_missing_capture_returns_empty() {
534        let temp = TempDir::new().unwrap();
535        fs::write(temp.path().join("lib.rs"), "fn foo() {}\n").unwrap();
536
537        let mut args = input("fn $NAME() {}", temp.path());
538        args.constraints = Some(HashMap::from([("MISSING".to_string(), ".*".to_string())]));
539
540        let result = perform_ast_grep(args).await.unwrap();
541
542        assert_eq!(result.count, 0);
543    }
544
545    #[tokio::test]
546    async fn invalid_constraint_regex_returns_error() {
547        let temp = TempDir::new().unwrap();
548        fs::write(temp.path().join("lib.rs"), "fn foo() {}\n").unwrap();
549
550        let mut args = input("fn $NAME() {}", temp.path());
551        args.constraints = Some(HashMap::from([("NAME".to_string(), "(".to_string())]));
552
553        let result = perform_ast_grep(args).await;
554
555        assert!(matches!(result, Err(AstGrepError::InvalidConstraintRegex { .. })));
556    }
557
558    #[tokio::test]
559    async fn multiple_constraints_all_must_match() {
560        let temp = TempDir::new().unwrap();
561        fs::write(temp.path().join("lib.rs"), "fn alpha() {}\nfn beta() {}\n").unwrap();
562
563        let mut args = input("fn $NAME() {}", temp.path());
564        args.constraints =
565            Some(HashMap::from([("NAME".to_string(), "^a".to_string()), ("NAME".to_string(), "lpha".to_string())]));
566
567        let result = perform_ast_grep(args).await.unwrap();
568
569        assert_eq!(result.count, 1);
570        assert!(result.matches[0].text.contains("alpha"));
571    }
572}
mcp_servers/coding/tools/ast_grep/mod.rs

mcp_servers/coding/tools/ast_grep/
mod.rs
