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).
This tool uses ast-grep for efficient and semantic code search and transformation using AST patterns.
AST-grep understands code structure, making it superior to regex for code transformations.
Parameters:
- `pattern`: The AST pattern to search for (required)
- `paths`: Optional array of file paths or glob patterns to search within (default: current directory)
- `language`: Optional language of the code (e.g., 'rust', 'javascript', 'python', 'typescript', 'go', 'java', 'c', 'cpp', 'php')
- `rewrite`: Optional rewrite pattern to apply for refactoring transformations
- `json_output`: Optional boolean to get output in JSON format (default: false)
- `context`: Optional number of lines of context to show around matches (default: 0)
- `update_all`: Optional boolean to apply rewrites to all matches without confirmation (default: false)
Note: Response size is controlled by global mcp_response_tokens_threshold setting.
Use more specific patterns to reduce output size if responses are truncated.
Pattern Syntax:
- Use metavariables like $NAME, $ARGS, $BODY for flexible matching
- Use $$$ for matching any number of statements/expressions
- Use $_ for anonymous wildcards (non-capturing)
- Patterns match AST structure, not text
Meta Variables:
- `$VAR` - matches single AST node (like `$NAME`, `$VALUE`)
- `$$$` - matches zero or more nodes (like `$$$ARGS`, `$$$BODY`)
- `$_` - anonymous wildcard, doesn't capture content
- Same-named variables must match identical content
Advanced Patterns:
- Structural matching: `if ($COND) { $$$BODY }` finds all if statements
- Method chains: `$OBJ.$METHOD1().$METHOD2($$$)` finds chained calls
- Nested expressions: `console.log($$$)` matches even in `func(console.log(x))`
Common Examples by Language:
**JavaScript/TypeScript:**
- Function calls: `console.log($$$)` or `$OBJ.$METHOD($$$)`
- Function definitions: `function $NAME($ARGS) { $$$ }`
- Arrow functions: `($ARGS) => $BODY`
- Variable declarations: `const $VAR = $VALUE`
- Import statements: `import $NAME from '$PATH'`
**PHP:**
- Function calls: `$NAME($$$)`
- Method calls: `$OBJ->$METHOD($$$)`
- Class definitions: `class $NAME { $$$ }`
- Variable assignments: `$$VAR = $VALUE`
**Rust:**
- Function calls: `println!($$$)` or `$NAME($$$)`
- Function definitions: `fn $NAME($ARGS) { $$$ }`
- Struct definitions: `struct $NAME { $$$ }`
- Use statements: `use $PATH;`
**Python:**
- Function calls: `print($$$)` or `$OBJ.$METHOD($$$)`
- Function definitions: `def $NAME($ARGS): $$$`
- Class definitions: `class $NAME: $$$`
- Import statements: `import $NAME`
Rewrite Examples:
- Rename functions: pattern `old_func($ARGS)` → rewrite `new_func($ARGS)`
- Add visibility: pattern `fn $NAME($ARGS)` → rewrite `pub fn $NAME($ARGS)`
- Modernize JS: pattern `var $NAME = $VALUE` → rewrite `const $NAME = $VALUE`
- Update method calls: pattern `$OBJ.oldMethod($ARGS)` → rewrite `$OBJ.newMethod($ARGS)`
Usage Examples:
- Find console logs: `{\"pattern\": \"console.log($$$)\", \"language\": \"javascript\"}`
- Rename function: `{\"pattern\": \"oldFunc($ARGS)\", \"rewrite\": \"newFunc($ARGS)\", \"language\": \"javascript\"}`
- Find PHP classes: `{\"pattern\": \"class $NAME\", \"language\": \"php\", \"paths\": [\"src/**/*.php\"]}`
- Search with context: `{\"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 {
crate::utils::glob::expand_glob_patterns_filtered(file_paths, None)
} 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,
})
}