use super::super::{McpFunction, McpToolCall, McpToolResult};
use anyhow::{anyhow, Result};
use serde_json::{json, Value};
fn group_ast_grep_output(output: &str) -> String {
let lines: Vec<&str> = output.lines().collect();
let mut result = Vec::new();
let mut current_file = String::new();
let mut file_lines = Vec::new();
for line in lines {
if let Some(colon_pos) = line.find(':') {
let filename = &line[..colon_pos];
let rest = &line[colon_pos + 1..];
if filename != current_file {
if !file_lines.is_empty() {
result.push(format!("{}:\n{}", current_file, file_lines.join("\n")));
file_lines.clear();
}
current_file = filename.to_string();
}
file_lines.push(rest.to_string());
} else {
if !file_lines.is_empty() {
result.push(format!("{}:\n{}", current_file, file_lines.join("\n")));
file_lines.clear();
current_file.clear();
}
result.push(line.to_string());
}
}
if !file_lines.is_empty() {
result.push(format!("{}:\n{}", current_file, file_lines.join("\n")));
}
if result.is_empty() {
output.to_string()
} else {
result.join("\n\n")
}
}
pub fn get_ast_grep_function() -> McpFunction {
McpFunction {
name: "ast_grep".to_string(),
description: "Search and refactor code using AST patterns with ast-grep (sg).
AST-grep understands code structure — superior to regex for code transformations.
Pattern Syntax:
- `$VAR` — matches exactly ONE AST node (not multiple args/params)
- `$$$VAR` — matches ZERO or more nodes; use inside `()`, `{}`, `[]` for lists
- `$_` — anonymous wildcard, matches one node without capturing
- Same-named variables must match identical content
CRITICAL — Patterns are structurally exact:
Every element you include must be present; every element you omit must be absent.
Optional syntax (return types, async, pub, decorators, type annotations, throws) must be in the pattern to match code that has it — and absent to match code that lacks it.
When unsure, make TWO calls: one with the optional part, one without.
- WRONG: `fn $NAME($ARGS)` — `$ARGS` is one node, misses multi-param functions
- RIGHT: `fn $NAME($$$ARGS)` — matches zero or more parameters
- `fn $NAME($$$ARGS) { $$$BODY }` does NOT match `fn foo() -> Bar {}` — return type absent from pattern
- `function $NAME($$$ARGS) {}` does NOT match `async function foo() {}` — async absent from pattern
Common Examples by Language:
**JavaScript/TypeScript:**
- Function calls: `$OBJ.$METHOD($$$)` or `console.log($$$)`
- Regular function: `function $NAME($$$ARGS) { $$$ }`
- Async function (separate pattern): `async function $NAME($$$ARGS) { $$$ }`
- Arrow function: `($$$ARGS) => $BODY`
- Variable: `const $VAR = $VALUE`
- Import: `import $NAME from '$PATH'`
**PHP:**
- Function calls: `$NAME($$$)` / method calls: `$OBJ->$METHOD($$$)`
- Class: `class $NAME { $$$ }`
**Rust:**
- Fn without return: `fn $NAME($$$ARGS) { $$$BODY }`
- Fn with return (separate pattern): `fn $NAME($$$ARGS) -> $RET { $$$BODY }`
- Struct: `struct $NAME { $$$ }` / macro: `println!($$$)`
**Python:**
- Fn without annotation: `def $NAME($$$ARGS): $$$`
- Fn with return annotation (separate pattern): `def $NAME($$$ARGS) -> $RET: $$$`
- Class: `class $NAME: $$$`
Examples:
- `{\"pattern\": \"console.log($$$)\", \"language\": \"javascript\"}`
- `{\"pattern\": \"oldFunc($$$ARGS)\", \"rewrite\": \"newFunc($$$ARGS)\", \"language\": \"javascript\"}`
- `{\"pattern\": \"class $NAME\", \"language\": \"php\", \"paths\": [\"src/**/*.php\"]}`
- `{\"pattern\": \"$OBJ.oldMethod($$$ARGS)\", \"rewrite\": \"$OBJ.newMethod($$$ARGS)\"}`
- `{\"pattern\": \"TODO\", \"context\": 2}`
".to_string(),
parameters: json!({
"type": "object",
"required": ["pattern"],
"properties": {
"pattern": {
"type": "string",
"description": "The AST pattern to search for. Use metavariables ($NAME, $$$) to match code structure, not text content"
},
"paths": {
"type": "array",
"items": {"type": "string"},
"description": "Optional array of file paths or glob patterns to search within (default: current directory)"
},
"language": {
"type": "string",
"description": "Optional language of the code (e.g., 'rust', 'javascript', 'python', 'typescript', 'go', 'java', 'c', 'cpp', 'php')"
},
"rewrite": {
"type": "string",
"description": "Optional rewrite pattern to apply for refactoring transformations"
},
"json_output": {
"type": "boolean",
"default": false,
"description": "Optional boolean to get output in JSON format (default: false)"
},
"context": {
"type": "integer",
"default": 0,
"description": "Optional number of lines of context to show around matches (default: 0)"
},
"update_all": {
"type": "boolean",
"default": false,
"description": "Optional boolean to apply rewrites to all matches without confirmation (default: false)"
}
}
}),
}
}
pub async fn execute_ast_grep_command(call: &McpToolCall) -> Result<McpToolResult> {
use tokio::process::Command as TokioCommand;
let pattern = match call.parameters.get("pattern") {
Some(Value::String(p)) => {
if p.trim().is_empty() {
return Ok(McpToolResult::error(
call.tool_name.clone(),
call.tool_id.clone(),
"Pattern parameter cannot be empty".to_string(),
));
}
p.clone()
}
Some(_) => {
return Ok(McpToolResult::error(
call.tool_name.clone(),
call.tool_id.clone(),
"Pattern parameter must be a string".to_string(),
));
}
None => {
return Ok(McpToolResult::error(
call.tool_name.clone(),
call.tool_id.clone(),
"Missing required 'pattern' parameter".to_string(),
));
}
};
let paths = call
.parameters
.get("paths")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|item| item.as_str().map(|s| s.to_string()))
.collect::<Vec<String>>()
});
let language = call
.parameters
.get("language")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let rewrite = call
.parameters
.get("rewrite")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let json_output = call
.parameters
.get("json_output")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let context = call
.parameters
.get("context")
.and_then(|v| v.as_i64())
.unwrap_or(0);
let update_all = call
.parameters
.get("update_all")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let mut cmd = TokioCommand::new("sg");
cmd.arg("-p");
cmd.arg(&pattern);
if let Some(lang) = &language {
cmd.arg("-l");
cmd.arg(lang);
}
if let Some(rewrite_pattern) = &rewrite {
cmd.arg("--rewrite");
cmd.arg(rewrite_pattern);
if update_all {
cmd.arg("--update-all");
}
}
if json_output {
cmd.arg("--json");
}
if context > 0 {
cmd.arg("-A");
cmd.arg(context.to_string());
cmd.arg("-B");
cmd.arg(context.to_string());
}
let expanded_paths_result = if let Some(file_paths) = &paths {
let has_glob_patterns = file_paths
.iter()
.any(|p| p.contains('*') || p.contains('?'));
if has_glob_patterns {
crate::utils::glob::expand_glob_patterns_filtered(file_paths, None)
} else {
let mut verified_paths = Vec::new();
for path in file_paths {
let path_obj = std::path::Path::new(path);
if path_obj.exists() {
verified_paths.push(path.clone());
} else {
return Ok(McpToolResult::error(
call.tool_name.clone(),
call.tool_id.clone(),
format!("Path does not exist: {}", path),
));
}
}
Ok(verified_paths)
}
} else {
Ok(vec![])
};
let actual_file_paths = match expanded_paths_result {
Ok(expanded_paths) => {
match (expanded_paths.is_empty(), &paths) {
(true, Some(provided_paths)) => {
let paths_str = provided_paths.join(", ");
return Ok(McpToolResult::error(
call.tool_name.clone(),
call.tool_id.clone(),
format!("No files found matching the specified paths: [{}]. Please verify the file paths exist and are not in ignored directories.", paths_str),
));
}
(true, None) => {
cmd.arg(".");
vec![".".to_string()]
}
(false, _) => {
for path in &expanded_paths {
cmd.arg(path);
}
expanded_paths
}
}
}
Err(e) => {
return Ok(McpToolResult::error(
call.tool_name.clone(),
call.tool_id.clone(),
format!("Failed to expand glob patterns: {e}"),
));
}
};
cmd.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.stdin(std::process::Stdio::null())
.kill_on_drop(true);
let mut debug_args = vec!["-p".to_string(), pattern.clone()];
if let Some(lang) = &language {
debug_args.push("-l".to_string());
debug_args.push(lang.clone());
}
if let Some(rewrite_pattern) = &rewrite {
debug_args.push("--rewrite".to_string());
debug_args.push(rewrite_pattern.clone());
if update_all {
debug_args.push("--update-all".to_string());
}
}
if json_output {
debug_args.push("--json".to_string());
}
if context > 0 {
debug_args.push("-A".to_string());
debug_args.push(context.to_string());
debug_args.push("-B".to_string());
debug_args.push(context.to_string());
}
let file_count = actual_file_paths.len();
if file_count <= 5 {
for path in &actual_file_paths {
debug_args.push(path.clone());
}
} else {
for path in actual_file_paths.iter().take(3) {
debug_args.push(path.clone());
}
debug_args.push(format!("... and {} more files", file_count - 3));
}
crate::log_debug!(
"Executing ast-grep command: sg {:?} (targeting {} files)",
debug_args,
file_count
);
let child = cmd
.spawn()
.map_err(|e| anyhow!("Failed to spawn ast-grep command: {}", e))?;
let result = child.wait_with_output().await;
let output = match result.map_err(|e| anyhow!("AST-grep command execution failed: {}", e)) {
Ok(output) => {
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
let grouped_output = group_ast_grep_output(&stdout);
let final_output = grouped_output;
let combined = if stderr.is_empty() {
final_output
} else if final_output.is_empty() {
stderr
} else {
format!("{}\n\nError: {}", final_output, stderr)
};
let status_code = output.status.code().unwrap_or(-1);
let success = output.status.success();
let operation_type = if rewrite.is_some() {
"rewrite"
} else {
"search"
};
let result = json!({
"success": success,
"output": combined,
"code": status_code,
"operation": operation_type,
"parameters": {
"pattern": pattern,
"paths": paths,
"language": language,
"rewrite": rewrite,
"json_output": json_output,
"context": context,
"update_all": update_all
},
"message": if success {
format!("AST-grep {operation_type} executed successfully with exit code {status_code}")
} else {
format!("AST-grep {operation_type} failed with exit code {status_code}")
}
});
result
}
Err(e) => json!({
"success": false,
"output": format!("Failed to execute ast-grep command: {e}"),
"code": -1,
"operation": if rewrite.is_some() { "rewrite" } else { "search" },
"parameters": {
"pattern": pattern,
"paths": paths,
"language": language,
"rewrite": rewrite,
"json_output": json_output,
"context": context,
"update_all": update_all
},
"message": format!("Failed to execute ast-grep command: {}", e)
}),
};
Ok(McpToolResult {
tool_name: "ast_grep".to_string(),
tool_id: call.tool_id.clone(),
result: output,
})
}