oxi_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    let mut stdout = child.stdout.take().expect("piped stdout");
90    let mut stderr = child.stderr.take().expect("piped stderr");
91
92    let mut stdout_buf = Vec::new();
93    let mut stderr_buf = Vec::new();
94    let (stdout_res, stderr_res) = tokio::join!(
95        stdout.read_to_end(&mut stdout_buf),
96        stderr.read_to_end(&mut stderr_buf)
97    );
98    stdout_res.map_err(|e| format!("Failed reading `sg` stdout: {e}"))?;
99    stderr_res.map_err(|e| format!("Failed reading `sg` stderr: {e}"))?;
100
101    let status = child
102        .wait()
103        .await
104        .map_err(|e| format!("Failed waiting on `sg`: {e}"))?;
105
106    // Parse stdout in a format-agnostic way: ast-grep supports three
107    // `--json` modes (`pretty`, `stream`, `compact`) and a future version
108    // could change the default. Try the array forms first (pretty emits
109    // a multi-line `[ {…}, {…} ]`, compact emits it on one line), then
110    // fall back to NDJSON (stream emits one `{…}` per line).
111    let matches = parse_sg_output(&stdout_buf).ok_or_else(|| {
112        "Failed to parse `sg` JSON output: no array or stream objects found".to_string()
113    })?;
114
115    // ast-grep exits non-zero on no matches (stdout empty) or on real
116    // failures (stderr populated). Distinguish by stderr content.
117    if !status.success() {
118        let stderr_text = String::from_utf8_lossy(&stderr_buf).trim().to_string();
119        if !stderr_text.is_empty() {
120            return Err(format!("`sg` failed: {stderr_text}"));
121        }
122    }
123
124    Ok(matches)
125}
126
127/// Parse the bytes emitted by `sg --json` regardless of output style.
128///
129/// Three documented styles:
130/// - `pretty` (the default): a multi-line JSON `[{…}, {…}, …]` array.
131/// - `compact`: a single-line JSON `[{…},{…}]` array.
132/// - `stream`: one JSON object `{…}` per line (NDJSON).
133///
134/// This helper accepts all three:
135/// 1. If the entire buffer parses as a JSON `Value::Array`, return it.
136/// 2. If it parses as a single `Value::Object`, wrap and return it.
137/// 3. Otherwise, treat each non-blank line as a separate JSON object and
138///    return the collected Vec (NDJSON).
139/// 4. If none of the above yield any matches, return `None` so the caller
140///    can decide between "no output" (success, zero matches) and a parse
141///    error.
142fn parse_sg_output(buf: &[u8]) -> Option<Vec<Value>> {
143    let trimmed = buf.iter().any(|b| !b.is_ascii_whitespace());
144    if !trimmed {
145        return Some(Vec::new());
146    }
147
148    // Try the whole-buffer shapes first.
149    if let Ok(v) = serde_json::from_slice::<Value>(buf) {
150        match v {
151            Value::Array(arr) => return Some(arr),
152            Value::Object(_) => return Some(vec![v]),
153            _ => {}
154        }
155    }
156
157    // Fall back to NDJSON (one match object per line).
158    let mut matches = Vec::new();
159    let mut parsed_any = false;
160    for line in buf.split(|b| *b == b'\n') {
161        let has_content = line.iter().any(|b| !b.is_ascii_whitespace());
162        if !has_content {
163            continue;
164        }
165        match serde_json::from_slice::<Value>(line) {
166            Ok(v) => {
167                parsed_any = true;
168                match v {
169                    Value::Array(arr) => matches.extend(arr),
170                    // `stream` style documents one OBJECT per line, but
171                    // tolerate an unbracketed single object too.
172                    Value::Object(_) => matches.push(v),
173                    _ => {}
174                }
175            }
176            Err(_) => continue,
177        }
178    }
179
180    if parsed_any { Some(matches) } else { None }
181}
182
183/// Format matches grouped by directory + file with line numbers.
184///
185/// Output shape (one section per file):
186/// ```text
187/// <relative-path>
188///   <line>:<col>-<end_col>: <trimmed text>
189///   ...
190/// ```
191/// Returns `(formatted_output, returned_count)`. Callers apply pagination
192/// BEFORE calling this — inner truncation logic was dead.
193fn format_matches(matches: &[Value], root: &Path) -> (String, usize) {
194    use std::collections::BTreeMap;
195
196    if matches.is_empty() {
197        return ("No matches found.".to_string(), 0);
198    }
199
200    // file -> Vec<(line, col, text)>
201    let mut by_file: BTreeMap<PathBuf, Vec<(usize, usize, String)>> = BTreeMap::new();
202
203    for m in matches {
204        let file = m
205            .get("file")
206            .and_then(Value::as_str)
207            .map(PathBuf::from)
208            .unwrap_or_else(|| PathBuf::from("<unknown>"));
209
210        // Range shape: { "start": {"line": N, "column": N}, "end": {...} }
211        // Older ast-grep emits `{ begin, end }` instead. Handle both.
212        let (line, col) = extract_position(m).unwrap_or((0, 0));
213
214        let text = m
215            .get("text")
216            .and_then(Value::as_str)
217            .map(str::to_string)
218            .unwrap_or_default();
219
220        // Trim trailing whitespace but keep indentation for readability.
221        let trimmed = text.lines().next().unwrap_or("").trim_end().to_string();
222
223        by_file.entry(file).or_default().push((line, col, trimmed));
224    }
225
226    let returned = matches.len();
227    let mut out = String::new();
228    out.push_str(&format!("Found {returned} match(es):\n"));
229
230    for (file, lines) in &by_file {
231        let display = file.strip_prefix(root).unwrap_or(file.as_path());
232        let display = display.to_string_lossy();
233        out.push('\n');
234        out.push_str(&format!("{display}\n"));
235        for (line, col, text) in lines {
236            if *col > 0 {
237                out.push_str(&format!("  {line}:{col}: {text}\n"));
238            } else {
239                out.push_str(&format!("  {line}: {text}\n"));
240            }
241        }
242    }
243
244    (out, returned)
245}
246
247/// Extract (line, column) from a `sg --json` match object, tolerating both
248/// the new `{ range: { start, end } }` shape and the older `{ begin, end }`
249/// shape. Lines and columns are 0-indexed in `sg` output; we display
250/// 1-indexed line numbers (line + 1) but pass columns through as-is.
251fn extract_position(m: &Value) -> Option<(usize, usize)> {
252    if let Some(range) = m.get("range").and_then(Value::as_object) {
253        let start = range.get("start").and_then(Value::as_object)?;
254        let line = start.get("line").and_then(Value::as_u64)? as usize;
255        let col = start
256            .get("column")
257            .or_else(|| start.get("col"))
258            .and_then(Value::as_u64)
259            .unwrap_or(0) as usize;
260        return Some((line + 1, col + 1));
261    }
262
263    if let Some(begin) = m.get("begin").and_then(Value::as_u64) {
264        return Some((begin as usize + 1, 1));
265    }
266
267    None
268}
269
270#[async_trait]
271impl AgentTool for AstGrepTool {
272    fn name(&self) -> &str {
273        "ast_grep"
274    }
275
276    fn label(&self) -> &str {
277        "AST Grep"
278    }
279
280    fn description(&self) -> &str {
281        "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."
282    }
283
284    fn parameters_schema(&self) -> Value {
285        json!({
286            "type": "object",
287            "properties": {
288                "pattern": {
289                    "type": "string",
290                    "description": "AST pattern in ast-grep syntax (e.g. 'fn $NAME($$$ARGS) { $$$BODY }'). Metavariables use uppercase `$NAME`; zero-or-more use `$$$NAME`."
291                },
292                "path": {
293                    "type": "string",
294                    "description": "File, directory, or glob to search. Defaults to the workspace root."
295                },
296                "skip": {
297                    "type": "integer",
298                    "description": "Number of results to skip (for pagination).",
299                    "minimum": 0,
300                    "default": 0
301                },
302                "limit": {
303                    "type": "integer",
304                    "description": "Maximum number of results to return.",
305                    "minimum": 1,
306                    "default": 50
307                }
308            },
309            "required": ["pattern"]
310        })
311    }
312
313    async fn execute(
314        &self,
315        _tool_call_id: &str,
316        params: Value,
317        _signal: Option<oneshot::Receiver<()>>,
318        ctx: &ToolContext,
319    ) -> Result<AgentToolResult, ToolError> {
320        // ── 1. Validate pattern ───────────────────────────────────────
321        let pattern = params
322            .get("pattern")
323            .and_then(Value::as_str)
324            .ok_or_else(|| "Missing required parameter: pattern".to_string())?
325            .trim();
326
327        if pattern.is_empty() {
328            return Ok(AgentToolResult::error(
329                "Invalid pattern: must be a non-empty string",
330            ));
331        }
332
333        // ── 2. Resolve search scope ───────────────────────────────────
334        let path_arg = params.get("path").and_then(Value::as_str).unwrap_or("");
335
336        let skip = params.get("skip").and_then(Value::as_u64).unwrap_or(0) as usize;
337
338        let limit = params
339            .get("limit")
340            .and_then(Value::as_u64)
341            .unwrap_or(DEFAULT_LIMIT as u64) as usize;
342        let limit = limit.max(1);
343
344        let root = self.root_dir.as_deref().unwrap_or_else(|| ctx.root());
345        let search_path = self.resolve_search_path(path_arg, root);
346
347        // ── 3. Run `sg` and parse JSON stream ─────────────────────────
348        let all_matches = match run_sg(pattern, &search_path).await {
349            Ok(v) => v,
350            Err(msg) if msg.starts_with("`sg` is not installed") => {
351                return Ok(AgentToolResult::error(msg));
352            }
353            Err(msg) => return Ok(AgentToolResult::error(format!("ast_grep failed: {msg}"))),
354        };
355
356        let total = all_matches.len();
357
358        // Apply pagination: skip the first `skip` items, then keep at
359        // most `limit`. Truncated == "there were more matches the caller
360        // didn't see" — i.e. `total > skip + returned`. We deliberately
361        // avoid the `skip > 0` heuristic here: paginating past the first
362        // page is normal, not a truncation signal.
363        let paged: Vec<Value> = all_matches.into_iter().skip(skip).take(limit).collect();
364        let returned = paged.len();
365        let truncated = total > skip + returned;
366
367        // ── 4. Format grouped results ────────────────────────────────
368        let (body, _returned_fmt) = format_matches(&paged, root);
369        let mut result = AgentToolResult::success(body);
370        result.metadata = Some(json!({
371            "total_matches": total,
372            "returned": returned,
373            "skipped": skip,
374            "limit": limit,
375            "truncated": truncated,
376            "pattern": pattern,
377            "search_path": search_path.to_string_lossy(),
378        }));
379
380        Ok(result)
381    }
382}
oxi_agent/tools/ast_grep.rs

oxi_agent/tools/
ast_grep.rs
