oxicode_agent/tools/
ast_grep.rs

1/// AST-grep tool — structural code search using `sg` (ast-grep CLI).
2use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
3use async_trait::async_trait;
4use serde_json::{Value, json};
5use std::path::{Path, PathBuf};
6use std::process::Stdio;
7use tokio::io::AsyncReadExt;
8use tokio::process::Command;
9use tokio::sync::oneshot;
10
11/// Default max results when caller does not specify `limit`.
12const DEFAULT_LIMIT: usize = 50;
13
14/// AstGrepTool — wraps the `sg` (ast-grep) CLI for structural code search.
15pub struct AstGrepTool {
16    root_dir: Option<PathBuf>,
17}
18
19impl AstGrepTool {
20    /// Create with no explicit root (uses ToolContext.root() at runtime).
21    pub fn new() -> Self {
22        Self { root_dir: None }
23    }
24
25    /// Create with a specific working directory (overrides ToolContext).
26    pub fn with_cwd(cwd: PathBuf) -> Self {
27        Self {
28            root_dir: Some(cwd),
29        }
30    }
31
32    /// Resolve a `path` argument relative to the tool's effective root.
33    /// Falls back to the root directory when `path` is empty.
34    fn resolve_search_path(&self, path: &str, ctx_root: &Path) -> PathBuf {
35        if path.is_empty() {
36            ctx_root.to_path_buf()
37        } else {
38            let candidate = PathBuf::from(path);
39            if candidate.is_absolute() {
40                candidate
41            } else {
42                ctx_root.join(candidate)
43            }
44        }
45    }
46}
47
48impl Default for AstGrepTool {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54/// Run `sg run -p <pattern> --json <path>` and return the parsed match
55/// objects as a `Vec<Value>`.
56///
57/// `run` is the default ast-grep subcommand and is required when other
58/// flags follow; bare `sg -p …` is a clap usage error. We pass bare
59/// `--json` and let [`parse_sg_output`] deal with whatever style
60/// ast-grep emits — see that function's doc for the three shapes handled.
61///
62/// Returns `Err("`sg` is not installed…")` when the binary is missing,
63/// `Err(...)` when the process reports a real failure on stderr, and
64/// `Ok(matches)` on success (including the empty-Vec "no matches" case —
65/// ast-grep exits 1 with empty stdout when the pattern is well-formed
66/// but produces no hits).
67async fn run_sg(pattern: &str, target: &Path) -> Result<Vec<Value>, String> {
68    let mut child = match Command::new("sg")
69        .arg("run")
70        .arg("-p")
71        .arg(pattern)
72        .arg("--json")
73        .arg(target)
74        .stdin(Stdio::null())
75        .stdout(Stdio::piped())
76        .stderr(Stdio::piped())
77        .spawn()
78    {
79        Ok(c) => c,
80        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
81            return Err(
82                "`sg` (ast-grep CLI) is not installed or not on PATH. Install it from https://ast-grep.github.io/ to use the ast_grep tool."
83                    .to_string(),
84            );
85        }
86        Err(e) => return Err(format!("Failed to invoke `sg`: {e}")),
87    };
88
89    // SAFETY: the command was spawned with `Stdio::piped()` for stdout/stderr
90    // and the spawn succeeded (we returned early on error), so both `take()`
91    // calls cannot return None.
92    #[allow(clippy::expect_used)]
93    let mut stdout = child.stdout.take().expect("piped stdout");
94    #[allow(clippy::expect_used)]
95    let mut stderr = child.stderr.take().expect("piped stderr");
96
97    let mut stdout_buf = Vec::new();
98    let mut stderr_buf = Vec::new();
99    let (stdout_res, stderr_res) = tokio::join!(
100        stdout.read_to_end(&mut stdout_buf),
101        stderr.read_to_end(&mut stderr_buf)
102    );
103    stdout_res.map_err(|e| format!("Failed reading `sg` stdout: {e}"))?;
104    stderr_res.map_err(|e| format!("Failed reading `sg` stderr: {e}"))?;
105
106    let status = child
107        .wait()
108        .await
109        .map_err(|e| format!("Failed waiting on `sg`: {e}"))?;
110
111    // Parse stdout in a format-agnostic way: ast-grep supports three
112    // `--json` modes (`pretty`, `stream`, `compact`) and a future version
113    // could change the default. Try the array forms first (pretty emits
114    // a multi-line `[ {…}, {…} ]`, compact emits it on one line), then
115    // fall back to NDJSON (stream emits one `{…}` per line).
116    let matches = parse_sg_output(&stdout_buf).ok_or_else(|| {
117        "Failed to parse `sg` JSON output: no array or stream objects found".to_string()
118    })?;
119
120    // ast-grep exits non-zero on no matches (stdout empty) or on real
121    // failures (stderr populated). Distinguish by stderr content.
122    if !status.success() {
123        let stderr_text = String::from_utf8_lossy(&stderr_buf).trim().to_string();
124        if !stderr_text.is_empty() {
125            return Err(format!("`sg` failed: {stderr_text}"));
126        }
127    }
128
129    Ok(matches)
130}
131
132/// Parse the bytes emitted by `sg --json` regardless of output style.
133///
134/// Three documented styles:
135/// - `pretty` (the default): a multi-line JSON `[{…}, {…}, …]` array.
136/// - `compact`: a single-line JSON `[{…},{…}]` array.
137/// - `stream`: one JSON object `{…}` per line (NDJSON).
138///
139/// This helper accepts all three:
140/// 1. If the entire buffer parses as a JSON `Value::Array`, return it.
141/// 2. If it parses as a single `Value::Object`, wrap and return it.
142/// 3. Otherwise, treat each non-blank line as a separate JSON object and
143///    return the collected Vec (NDJSON).
144/// 4. If none of the above yield any matches, return `None` so the caller
145///    can decide between "no output" (success, zero matches) and a parse
146///    error.
147fn parse_sg_output(buf: &[u8]) -> Option<Vec<Value>> {
148    let trimmed = buf.iter().any(|b| !b.is_ascii_whitespace());
149    if !trimmed {
150        return Some(Vec::new());
151    }
152
153    // Try the whole-buffer shapes first.
154    if let Ok(v) = serde_json::from_slice::<Value>(buf) {
155        match v {
156            Value::Array(arr) => return Some(arr),
157            Value::Object(_) => return Some(vec![v]),
158            _ => {}
159        }
160    }
161
162    // Fall back to NDJSON (one match object per line).
163    let mut matches = Vec::new();
164    let mut parsed_any = false;
165    for line in buf.split(|b| *b == b'\n') {
166        let has_content = line.iter().any(|b| !b.is_ascii_whitespace());
167        if !has_content {
168            continue;
169        }
170        match serde_json::from_slice::<Value>(line) {
171            Ok(v) => {
172                parsed_any = true;
173                match v {
174                    Value::Array(arr) => matches.extend(arr),
175                    // `stream` style documents one OBJECT per line, but
176                    // tolerate an unbracketed single object too.
177                    Value::Object(_) => matches.push(v),
178                    _ => {}
179                }
180            }
181            Err(_) => continue,
182        }
183    }
184
185    if parsed_any { Some(matches) } else { None }
186}
187
188/// Format matches grouped by directory + file with line numbers.
189///
190/// Output shape (one section per file):
191/// ```text
192/// <relative-path>
193///   <line>:<col>-<end_col>: <trimmed text>
194///   ...
195/// ```
196/// Returns `(formatted_output, returned_count)`. Callers apply pagination
197/// BEFORE calling this — inner truncation logic was dead.
198fn format_matches(matches: &[Value], root: &Path) -> (String, usize) {
199    use std::collections::BTreeMap;
200
201    if matches.is_empty() {
202        return ("No matches found.".to_string(), 0);
203    }
204
205    // file -> Vec<(line, col, text)>
206    let mut by_file: BTreeMap<PathBuf, Vec<(usize, usize, String)>> = BTreeMap::new();
207
208    for m in matches {
209        let file = m
210            .get("file")
211            .and_then(Value::as_str)
212            .map(PathBuf::from)
213            .unwrap_or_else(|| PathBuf::from("<unknown>"));
214
215        // Range shape: { "start": {"line": N, "column": N}, "end": {...} }
216        // Older ast-grep emits `{ begin, end }` instead. Handle both.
217        let (line, col) = extract_position(m).unwrap_or((0, 0));
218
219        let text = m
220            .get("text")
221            .and_then(Value::as_str)
222            .map(str::to_string)
223            .unwrap_or_default();
224
225        // Trim trailing whitespace but keep indentation for readability.
226        let trimmed = text.lines().next().unwrap_or("").trim_end().to_string();
227
228        by_file.entry(file).or_default().push((line, col, trimmed));
229    }
230
231    let returned = matches.len();
232    let mut out = String::new();
233    out.push_str(&format!("Found {returned} match(es):\n"));
234
235    for (file, lines) in &by_file {
236        let display = file.strip_prefix(root).unwrap_or(file.as_path());
237        let display = display.to_string_lossy();
238        out.push('\n');
239        out.push_str(&format!("{display}\n"));
240        for (line, col, text) in lines {
241            if *col > 0 {
242                out.push_str(&format!("  {line}:{col}: {text}\n"));
243            } else {
244                out.push_str(&format!("  {line}: {text}\n"));
245            }
246        }
247    }
248
249    (out, returned)
250}
251
252/// Extract (line, column) from a `sg --json` match object, tolerating both
253/// the new `{ range: { start, end } }` shape and the older `{ begin, end }`
254/// shape. Lines and columns are 0-indexed in `sg` output; we display
255/// 1-indexed line numbers (line + 1) but pass columns through as-is.
256fn extract_position(m: &Value) -> Option<(usize, usize)> {
257    if let Some(range) = m.get("range").and_then(Value::as_object) {
258        let start = range.get("start").and_then(Value::as_object)?;
259        let line = start.get("line").and_then(Value::as_u64)? as usize;
260        let col = start
261            .get("column")
262            .or_else(|| start.get("col"))
263            .and_then(Value::as_u64)
264            .unwrap_or(0) as usize;
265        return Some((line + 1, col + 1));
266    }
267
268    if let Some(begin) = m.get("begin").and_then(Value::as_u64) {
269        return Some((begin as usize + 1, 1));
270    }
271
272    None
273}
274
275#[async_trait]
276impl AgentTool for AstGrepTool {
277    fn name(&self) -> &str {
278        "ast_grep"
279    }
280
281    fn label(&self) -> &str {
282        "AST Grep"
283    }
284
285    fn description(&self) -> &str {
286        "Structural code search using ast-grep. Pattern uses ast-grep pattern syntax (e.g. 'fn $NAME($$$ARGS) { $$$BODY }'). Runs `sg run -p <pattern> --json <path>` and groups results by file with line numbers. Requires the `sg` (ast-grep) CLI to be installed."
287    }
288
289    fn parameters_schema(&self) -> Value {
290        json!({
291            "type": "object",
292            "properties": {
293                "pattern": {
294                    "type": "string",
295                    "description": "AST pattern in ast-grep syntax (e.g. 'fn $NAME($$$ARGS) { $$$BODY }'). Metavariables use uppercase `$NAME`; zero-or-more use `$$$NAME`."
296                },
297                "path": {
298                    "type": "string",
299                    "description": "File, directory, or glob to search. Defaults to the workspace root."
300                },
301                "skip": {
302                    "type": "integer",
303                    "description": "Number of results to skip (for pagination).",
304                    "minimum": 0,
305                    "default": 0
306                },
307                "limit": {
308                    "type": "integer",
309                    "description": "Maximum number of results to return.",
310                    "minimum": 1,
311                    "default": 50
312                }
313            },
314            "required": ["pattern"]
315        })
316    }
317
318    async fn execute(
319        &self,
320        _tool_call_id: &str,
321        params: Value,
322        _signal: Option<oneshot::Receiver<()>>,
323        ctx: &ToolContext,
324    ) -> Result<AgentToolResult, ToolError> {
325        // ── 1. Validate pattern ───────────────────────────────────────
326        let pattern = params
327            .get("pattern")
328            .and_then(Value::as_str)
329            .ok_or_else(|| "Missing required parameter: pattern".to_string())?
330            .trim();
331
332        if pattern.is_empty() {
333            return Ok(AgentToolResult::error(
334                "Invalid pattern: must be a non-empty string",
335            ));
336        }
337
338        // ── 2. Resolve search scope ───────────────────────────────────
339        let path_arg = params.get("path").and_then(Value::as_str).unwrap_or("");
340
341        let skip = params.get("skip").and_then(Value::as_u64).unwrap_or(0) as usize;
342
343        let limit = params
344            .get("limit")
345            .and_then(Value::as_u64)
346            .unwrap_or(DEFAULT_LIMIT as u64) as usize;
347        let limit = limit.max(1);
348
349        let root = self.root_dir.as_deref().unwrap_or_else(|| ctx.root());
350        let search_path = self.resolve_search_path(path_arg, root);
351
352        // ── 3. Run `sg` and parse JSON stream ─────────────────────────
353        let all_matches = match run_sg(pattern, &search_path).await {
354            Ok(v) => v,
355            Err(msg) if msg.starts_with("`sg` is not installed") => {
356                return Ok(AgentToolResult::error(msg));
357            }
358            Err(msg) => return Ok(AgentToolResult::error(format!("ast_grep failed: {msg}"))),
359        };
360
361        let total = all_matches.len();
362
363        // Apply pagination: skip the first `skip` items, then keep at
364        // most `limit`. Truncated == "there were more matches the caller
365        // didn't see" — i.e. `total > skip + returned`. We deliberately
366        // avoid the `skip > 0` heuristic here: paginating past the first
367        // page is normal, not a truncation signal.
368        let paged: Vec<Value> = all_matches.into_iter().skip(skip).take(limit).collect();
369        let returned = paged.len();
370        let truncated = total > skip + returned;
371
372        // ── 4. Format grouped results ────────────────────────────────
373        let (body, _returned_fmt) = format_matches(&paged, root);
374        let mut result = AgentToolResult::success(body);
375        result.metadata = Some(json!({
376            "total_matches": total,
377            "returned": returned,
378            "skipped": skip,
379            "limit": limit,
380            "truncated": truncated,
381            "pattern": pattern,
382            "search_path": search_path.to_string_lossy(),
383        }));
384
385        Ok(result)
386    }
387}
oxicode_agent/tools/ast_grep.rs

oxicode_agent/tools/
ast_grep.rs
