vtcode_core/tools/
ast_grep.rs

1//! AST-grep integration for VTCode
2//!
3//! This module provides integration with the ast-grep CLI tool for
4//! syntax-aware code search, transformation, linting, and refactoring.
5
6use anyhow::{Context, Result};
7use serde_json::{Value, json};
8use std::collections::HashMap;
9use tokio;
10
11/// AST-grep engine for syntax-aware code operations
12pub struct AstGrepEngine {
13    /// Path to the ast-grep executable
14    sgrep_path: String,
15}
16
17impl AstGrepEngine {
18    /// Create a new AST-grep engine
19    pub fn new() -> Result<Self> {
20        // Try to find ast-grep in PATH
21        let sgrep_path = if cfg!(target_os = "windows") {
22            "ast-grep.exe"
23        } else {
24            "ast-grep"
25        };
26
27        // Verify ast-grep is available
28        let output = std::process::Command::new(sgrep_path)
29            .arg("--version")
30            .output()
31            .context("Failed to execute ast-grep")?;
32
33        if !output.status.success() {
34            return Err(anyhow::anyhow!(
35                "ast-grep not found or not working properly"
36            ));
37        }
38
39        Ok(Self {
40            sgrep_path: sgrep_path.to_string(),
41        })
42    }
43
44    /// Map language names to ast-grep language identifiers
45    fn map_language(language: &str) -> &str {
46        match language.to_lowercase().as_str() {
47            "rust" => "rust",
48            "rs" => "rust",
49            "python" => "python",
50            "py" => "python",
51            "javascript" => "javascript",
52            "js" => "javascript",
53            "typescript" => "typescript",
54            "ts" => "typescript",
55            "tsx" => "tsx",
56            "go" => "go",
57            "golang" => "go",
58            "java" => "java",
59            "cpp" => "cpp",
60            "c++" => "cpp",
61            "c" => "c",
62            "html" => "html",
63            "css" => "css",
64            "json" => "json",
65            "yaml" => "yaml",
66            "yml" => "yaml",
67            _ => language,
68        }
69    }
70
71    /// Search code using AST-grep patterns
72    pub async fn search(
73        &self,
74        pattern: &str,
75        path: &str,
76        language: Option<&str>,
77        context_lines: Option<usize>,
78        max_results: Option<usize>,
79    ) -> Result<Value> {
80        let sgrep_path = self.sgrep_path.clone();
81        let pattern = pattern.to_string();
82        let path = path.to_string();
83        let language = language.map(|s| s.to_string());
84        let _context_lines = context_lines.unwrap_or(0);
85        let _max_results = max_results.unwrap_or(100);
86
87        let handle = tokio::task::spawn_blocking(move || {
88            let mut cmd = std::process::Command::new(&sgrep_path);
89            cmd.arg("run")
90                .arg("--pattern")
91                .arg(&pattern)
92                .arg("--json")
93                .arg("--context")
94                .arg(&_context_lines.to_string())
95                .arg(&path);
96
97            if let Some(lang) = language {
98                cmd.arg("--lang").arg(Self::map_language(&lang));
99            }
100
101            cmd.output()
102        });
103
104        let output = handle
105            .await
106            .context("Failed to spawn ast-grep search task")?
107            .context("Failed to execute ast-grep search")?;
108
109        if !output.status.success() {
110            let stderr = String::from_utf8_lossy(&output.stderr);
111            return Err(anyhow::anyhow!("ast-grep search failed: {}", stderr));
112        }
113
114        let stdout = String::from_utf8_lossy(&output.stdout);
115        let results: Value =
116            serde_json::from_str(&stdout).context("Failed to parse ast-grep search results")?;
117
118        Ok(json!({ "success": true, "matches": results }))
119    }
120
121    /// Transform code using AST-grep patterns
122    pub async fn transform(
123        &self,
124        pattern: &str,
125        replacement: &str,
126        path: &str,
127        language: Option<&str>,
128        preview_only: bool,
129        update_all: bool,
130    ) -> Result<Value> {
131        let sgrep_path = self.sgrep_path.clone();
132        let pattern = pattern.to_string();
133        let replacement = replacement.to_string();
134        let path = path.to_string();
135        let language = language.map(|s| s.to_string());
136
137        let handle = tokio::task::spawn_blocking(move || {
138            let mut cmd = std::process::Command::new(&sgrep_path);
139            cmd.arg("run")
140                .arg("--pattern")
141                .arg(&pattern)
142                .arg("--rewrite")
143                .arg(&replacement)
144                .arg("--json")
145                .arg(&path);
146
147            if let Some(lang) = language {
148                cmd.arg("--lang").arg(Self::map_language(&lang));
149            }
150
151            // Note: We can't use --interactive and --json together
152            // For preview, we'll just show the matches without applying changes
153            if update_all && !preview_only {
154                cmd.arg("--update-all");
155            }
156
157            cmd.output()
158        });
159
160        let output = handle
161            .await
162            .context("Failed to spawn ast-grep transform task")?
163            .context("Failed to execute ast-grep transform")?;
164
165        if !output.status.success() {
166            let stderr = String::from_utf8_lossy(&output.stderr);
167            return Err(anyhow::anyhow!("ast-grep transform failed: {}", stderr));
168        }
169
170        let stdout = String::from_utf8_lossy(&output.stdout);
171        let results: Value =
172            serde_json::from_str(&stdout).context("Failed to parse ast-grep transform results")?;
173
174        Ok(json!({ "success": true, "changes": results }))
175    }
176
177    /// Lint code using AST-grep rules with custom rules
178    pub async fn lint(
179        &self,
180        path: &str,
181        language: Option<&str>,
182        severity_filter: Option<&str>,
183        custom_rules: Option<Vec<HashMap<String, Value>>>,
184    ) -> Result<Value> {
185        let sgrep_path = self.sgrep_path.clone();
186        let path = path.to_string();
187        let language = language.map(|s| s.to_string());
188        let _severity_filter = severity_filter.map(|s| s.to_string());
189        let _custom_rules = custom_rules.clone();
190
191        let handle = tokio::task::spawn_blocking(move || {
192            let mut cmd = std::process::Command::new(&sgrep_path);
193            cmd.arg("run")
194                .arg("--pattern")
195                .arg("// TODO: $$")
196                .arg("--json")
197                .arg(&path);
198
199            if let Some(lang) = language {
200                cmd.arg("--lang").arg(Self::map_language(&lang));
201            }
202
203            cmd.output()
204        });
205
206        let output = handle
207            .await
208            .context("Failed to spawn ast-grep lint task")?
209            .context("Failed to execute ast-grep lint")?;
210
211        // Even if the command fails (e.g., no project config), we'll still return results
212        // if we got any output
213        let stdout = String::from_utf8_lossy(&output.stdout);
214        if !stdout.trim().is_empty() {
215            if let Ok(results) = serde_json::from_str::<Value>(&stdout) {
216                return Ok(json!({ "success": true, "issues": results }));
217            }
218        }
219
220        // If we couldn't parse JSON output, return a more generic response
221        if !output.status.success() {
222            let stderr = String::from_utf8_lossy(&output.stderr);
223            // If it's a project config error, we'll return an empty result instead of failing
224            if stderr.contains("No ast-grep project configuration") {
225                return Ok(json!({
226                    "success": true,
227                    "issues": [],
228                    "warning": "No ast-grep project configuration found. Linting may be limited."
229                }));
230            }
231            return Err(anyhow::anyhow!("ast-grep lint failed: {}", stderr));
232        }
233
234        Ok(json!({ "success": true, "issues": [] }))
235    }
236
237    /// Get refactoring suggestions using AST-grep
238    pub async fn refactor(
239        &self,
240        path: &str,
241        language: Option<&str>,
242        refactor_type: &str,
243    ) -> Result<Value> {
244        // Different refactoring suggestions based on type
245        let (pattern, replacement) = match refactor_type {
246            "extract_function" => (
247                "function $func($) { $$ }",
248                "// TODO: Extract function $func to separate module\nfunction $func($) { $$ }",
249            ),
250            "remove_console_logs" => ("console.log($$)", ""),
251            "simplify_conditions" => ("if ($cond) { true } else { false }", "$cond"),
252            "extract_constants" => (
253                "$NUM", // Simple number extraction
254                "const MY_CONSTANT = $NUM;\n// TODO: Replace $NUM with MY_CONSTANT",
255            ),
256            "modernize_syntax" => ("var $VAR = $$", "let $VAR = $$"),
257            _ => ("$$", "// TODO: Consider refactoring this code"),
258        };
259
260        let sgrep_path = self.sgrep_path.clone();
261        let path = path.to_string();
262        let language = language.map(|s| s.to_string());
263        let pattern = pattern.to_string();
264        let replacement = replacement.to_string();
265
266        let handle = tokio::task::spawn_blocking(move || {
267            let mut cmd = std::process::Command::new(&sgrep_path);
268            cmd.arg("run")
269                .arg("--pattern")
270                .arg(&pattern)
271                .arg("--rewrite")
272                .arg(&replacement)
273                .arg("--json")
274                .arg(&path);
275
276            if let Some(lang) = language {
277                cmd.arg("--lang").arg(Self::map_language(&lang));
278            }
279
280            cmd.output()
281        });
282
283        let output = handle
284            .await
285            .context("Failed to spawn ast-grep refactor task")?
286            .context("Failed to execute ast-grep refactor")?;
287
288        if !output.status.success() {
289            let stderr = String::from_utf8_lossy(&output.stderr);
290            return Err(anyhow::anyhow!("ast-grep refactor failed: {}", stderr));
291        }
292
293        let stdout = String::from_utf8_lossy(&output.stdout);
294        let results: Value =
295            serde_json::from_str(&stdout).context("Failed to parse ast-grep refactor results")?;
296
297        Ok(json!({ "success": true, "suggestions": results }))
298    }
299
300    /// Run a custom ast-grep command with full options
301    pub async fn run_custom(
302        &self,
303        pattern: &str,
304        path: &str,
305        language: Option<&str>,
306        rewrite: Option<&str>,
307        context_lines: Option<usize>,
308        max_results: Option<usize>,
309        interactive: bool,
310        update_all: bool,
311    ) -> Result<Value> {
312        let sgrep_path = self.sgrep_path.clone();
313        let pattern = pattern.to_string();
314        let path = path.to_string();
315        let language = language.map(|s| s.to_string());
316        let rewrite = rewrite.map(|s| s.to_string());
317        let context_lines = context_lines.unwrap_or(0);
318        let max_results = max_results.unwrap_or(100);
319
320        let handle = tokio::task::spawn_blocking(move || {
321            let mut cmd = std::process::Command::new(&sgrep_path);
322            cmd.arg("run")
323                .arg("--pattern")
324                .arg(&pattern)
325                .arg("--json")
326                .arg("--context")
327                .arg(context_lines.to_string())
328                .arg("--max-results")
329                .arg(max_results.to_string())
330                .arg(&path);
331
332            if let Some(lang) = language {
333                cmd.arg("--lang").arg(Self::map_language(&lang));
334            }
335
336            if let Some(rewrite_pattern) = rewrite {
337                cmd.arg("--rewrite").arg(&rewrite_pattern);
338            }
339
340            if interactive {
341                cmd.arg("--interactive");
342            } else if update_all {
343                cmd.arg("--update-all");
344            }
345
346            cmd.output()
347        });
348
349        let output = handle
350            .await
351            .context("Failed to spawn ast-grep custom task")?
352            .context("Failed to execute ast-grep custom command")?;
353
354        if !output.status.success() {
355            let stderr = String::from_utf8_lossy(&output.stderr);
356            return Err(anyhow::anyhow!(
357                "ast-grep custom command failed: {}",
358                stderr
359            ));
360        }
361
362        let stdout = String::from_utf8_lossy(&output.stdout);
363        let results: Value =
364            serde_json::from_str(&stdout).context("Failed to parse ast-grep custom results")?;
365
366        Ok(json!({ "success": true, "results": results }))
367    }
368}
vtcode_core/tools/ast_grep.rs

vtcode_core/tools/
ast_grep.rs
