oxi_agent/tools/

ast_edit.rs

//! AST-edit tool — structural code rewriting via the `sg` (ast-grep) CLI.
//!
//! Like the `ast_grep` search tool but applies pattern-driven rewrites.
//! Accepts an array of `{ pat, out }` operations applied against a list of
//! files / directories / globs.
//!
//! ## Empirical CLI behavior (verified against ast-grep 0.45.0)
//!
//! - `sg` does **not** expand shell globs in positional path args:
//!   `sg -p P --json 'src/**/*.rs'` errors with `No such file or directory`.
//!   This tool expands globs via the `glob` crate before invoking `sg`.
//! - `-U` (`--update-all`) is **silently ignored** when `--json` is set: the
//!   process exits 0, prints JSON, and the file is untouched. We therefore
//!   drop `--json` for the apply pass.
//! - With `-U` and no `--json`, `sg` writes changes in place AND prints
//!   `Applied N changes` on **stderr** — so we get the replacement count
//!   without a second dry-run pass.
//! - `--json=stream` emits one JSON object per match per line on stdout,
//!   making it cheap to count and group by file.
//! - `sg` accepts any number of positional paths (directories + concrete
//!   files + globs-expanded files) in a single invocation, so we issue one
//!   process per (op × path-chunk) and never loop over individual paths.
//!
//! ## Modes
//!
//! - `dry_run=true` (default) — previews via `sg -p P -r R --json=stream`.
//!   No files are modified; counts come from stdout line count.
//! - `dry_run=false` — applies via `sg -p P -r R -U`. Counts come from the
//!   `Applied N changes` line on stderr.
use super::path_security::PathGuard;
use super::{AgentTool, AgentToolResult, ToolContext, ToolError, ToolExecutionMode};
use async_trait::async_trait;
use serde_json::{Value, json};
use std::path::{Path, PathBuf};
use std::process::Stdio;
use tokio::io::AsyncReadExt;
use tokio::process::Command;
use tokio::sync::oneshot;

/// Cap on how many paths we'll pass to a single `sg` invocation.
/// `sg` walks directories internally, so this only counts the input list.
const MAX_PATHS_PER_INVOCATION: usize = 512;

/// Cap on stdout we'll buffer for a dry-run preview (bytes). ast-grep's
/// `--json=stream` output is one line per match, so 8 MiB is generous.
const DRY_RUN_STDOUT_CAP: usize = 8 * 1024 * 1024;

/// One pattern → replacement operation supplied by the caller.
#[derive(Debug, Clone)]
struct RewriteOp {
    pat: String,
    out: String,
}

/// AstEditTool — wraps the `sg` (ast-grep) CLI for structural rewriting.
pub struct AstEditTool {
    root_dir: Option<PathBuf>,
}

impl AstEditTool {
    /// Create with no explicit root (uses ToolContext.root() at runtime).
    pub fn new() -> Self {
        Self { root_dir: None }
    }

    /// Create with a specific working directory (overrides ToolContext).
    pub fn with_cwd(cwd: PathBuf) -> Self {
        Self {
            root_dir: Some(cwd),
        }
    }

    /// Resolve a single user-supplied path string against the tool root.
    /// Absolute paths pass through; relative paths join onto the root.
    fn resolve_one(raw: &str, root: &Path) -> PathBuf {
        let candidate = PathBuf::from(raw);
        if candidate.is_absolute() {
            candidate
        } else {
            root.join(candidate)
        }
    }

    /// True if the path string contains shell-glob wildcard characters.
    fn looks_like_glob(raw: &str) -> bool {
        raw.contains('*') || raw.contains('?') || raw.contains('[')
    }

    /// Expand a list of user paths (files, dirs, or globs) into a flat list
    /// of `sg`-ready positional args: a mix of concrete file paths and
    /// directories (which `sg` walks itself).
    ///
    /// `sg` does not expand globs in positional args (verified: it errors
    /// with `No such file or directory`), so we expand them here via the
    /// `glob` crate. Directories pass through unchanged.
    fn expand_paths(raw_paths: &[String], root: &Path) -> Result<Vec<PathBuf>, ToolError> {
        let mut out: Vec<PathBuf> = Vec::new();
        let mut seen: std::collections::HashSet<PathBuf> = std::collections::HashSet::new();

        for raw in raw_paths {
            if Self::looks_like_glob(raw) {
                let candidate = Self::resolve_one(raw, root);
                let pattern_str = candidate.to_string_lossy().into_owned();
                let entries = glob::glob(&pattern_str)
                    .map_err(|e| format!("Invalid glob pattern '{}': {}", raw, e))?;
                let mut matched_any = false;
                for entry in entries {
                    let p = entry.map_err(|e| format!("Glob error for '{}': {}", raw, e))?;
                    matched_any = true;
                    // `sg` walks directories, so we can pass them through
                    // too — but only when the glob explicitly targets a
                    // directory (rare). For file globs, filter to files.
                    if p.is_dir() {
                        if seen.insert(p.clone()) {
                            out.push(p);
                        }
                    } else if p.is_file() && seen.insert(p.clone()) {
                        out.push(p);
                    }
                }
                if !matched_any {
                    return Err(format!("Glob '{}' matched no files", raw));
                }
            } else {
                let candidate = Self::resolve_one(raw, root);
                if !candidate.exists() {
                    return Err(format!("Path not found: {}", raw));
                }
                if seen.insert(candidate.clone()) {
                    out.push(candidate);
                }
            }
        }

        if out.is_empty() {
            return Err("No files matched the supplied paths/globs".to_string());
        }

        Ok(out)
    }
}

impl Default for AstEditTool {
    fn default() -> Self {
        Self::new()
    }
}

/// Parse `ops` from the input `Value`.
fn parse_ops(params: &Value) -> Result<Vec<RewriteOp>, ToolError> {
    let arr = params
        .get("ops")
        .and_then(Value::as_array)
        .ok_or_else(|| "Missing required parameter: ops (must be an array)".to_string())?;

    if arr.is_empty() {
        return Err("Parameter 'ops' must contain at least one { pat, out } entry".to_string());
    }

    let mut ops = Vec::with_capacity(arr.len());
    for (i, op) in arr.iter().enumerate() {
        let pat = op
            .get("pat")
            .and_then(Value::as_str)
            .ok_or_else(|| format!("ops[{}]: missing or non-string 'pat'", i))?
            .to_string();
        let out = op
            .get("out")
            .and_then(Value::as_str)
            .ok_or_else(|| format!("ops[{}]: missing or non-string 'out'", i))?
            .to_string();

        if pat.trim().is_empty() {
            return Err(format!("ops[{}]: 'pat' must be a non-empty string", i));
        }

        ops.push(RewriteOp { pat, out });
    }

    Ok(ops)
}

/// Parse `paths` from the input `Value`.
fn parse_paths(params: &Value) -> Result<Vec<String>, ToolError> {
    let arr = params
        .get("paths")
        .and_then(Value::as_array)
        .ok_or_else(|| "Missing required parameter: paths (must be an array)".to_string())?;

    if arr.is_empty() {
        return Err("Parameter 'paths' must contain at least one path".to_string());
    }

    let mut paths = Vec::with_capacity(arr.len());
    for (i, p) in arr.iter().enumerate() {
        let s = p
            .as_str()
            .ok_or_else(|| format!("paths[{}]: must be a string", i))?;
        paths.push(s.to_string());
    }

    Ok(paths)
}

/// Spawn `sg` for a single (op, path-chunk) and return its `(status, stdout, stderr)`.
///
/// For dry-run previews we use `--json=stream` (one JSON object per match
/// per line — cheap to count, easy to group by file). For real rewrites we
/// drop `--json` and add `-U` because ast-grep silently ignores `-U` when
/// `--json` is set; `Applied N changes` then lands on stderr.
async fn run_sg_for_op(
    op: &RewriteOp,
    paths: &[PathBuf],
    dry_run: bool,
) -> Result<(std::process::ExitStatus, Vec<u8>, Vec<u8>), String> {
    let mut cmd = Command::new("sg");
    cmd.arg("-p").arg(&op.pat).arg("-r").arg(&op.out);

    if dry_run {
        // --json=stream: one match per line, easy to count.
        cmd.arg("--json=stream");
    } else {
        // Real rewrite. ast-grep ignores -U when --json is set, so we MUST
        // omit --json here. The `Applied N changes` summary lands on stderr.
        cmd.arg("-U");
    }

    for p in paths {
        cmd.arg(p);
    }

    cmd.stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());

    let mut child = match cmd.spawn() {
        Ok(c) => c,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
            return Err(
                "`sg` (ast-grep CLI) is not installed or not on PATH. Install it from https://ast-grep.github.io/ to use the ast_edit tool."
                    .to_string(),
            );
        }
        Err(e) => return Err(format!("Failed to invoke `sg`: {e}")),
    };

    let mut stdout = child.stdout.take().expect("piped stdout");
    let mut stderr = child.stderr.take().expect("piped stderr");

    let mut stdout_buf = Vec::new();
    let mut stderr_buf = Vec::new();

    if dry_run {
        // Cap dry-run stdout so a runaway codebase doesn't OOM the agent.
        let mut limited = stdout.take(DRY_RUN_STDOUT_CAP as u64);
        let (s_res, e_res) = tokio::join!(
            AsyncReadExt::read_to_end(&mut limited, &mut stdout_buf),
            stderr.read_to_end(&mut stderr_buf)
        );
        s_res.map_err(|e| format!("Failed reading `sg` stdout: {e}"))?;
        e_res.map_err(|e| format!("Failed reading `sg` stderr: {e}"))?;
    } else {
        let (s_res, e_res) = tokio::join!(
            stdout.read_to_end(&mut stdout_buf),
            stderr.read_to_end(&mut stderr_buf)
        );
        s_res.map_err(|e| format!("Failed reading `sg` stdout: {e}"))?;
        e_res.map_err(|e| format!("Failed reading `sg` stderr: {e}"))?;
    }

    let status = child
        .wait()
        .await
        .map_err(|e| format!("Failed waiting on `sg`: {e}"))?;

    Ok((status, stdout_buf, stderr_buf))
}

/// Count and group dry-run matches from `sg --json=stream` output.
///
/// Each non-empty line is one match JSON object. We only need a count and
/// per-file breakdown, not the full payload, so malformed lines are
/// silently skipped (counted toward total if they at least parsed as JSON).
fn summarise_dry_run(stdout: &[u8]) -> (usize, std::collections::BTreeMap<PathBuf, usize>) {
    let mut total = 0usize;
    let mut by_file: std::collections::BTreeMap<PathBuf, usize> = std::collections::BTreeMap::new();

    for line in stdout.split(|b| *b == b'\n') {
        let trimmed: Vec<u8> = line
            .iter()
            .copied()
            .skip_while(|b| b.is_ascii_whitespace())
            .take_while(|b| !b.is_ascii_whitespace())
            .collect();
        if trimmed.is_empty() {
            continue;
        }
        if let Ok(v) = serde_json::from_slice::<Value>(&trimmed) {
            if let Some(file) = v.get("file").and_then(Value::as_str) {
                *by_file.entry(PathBuf::from(file)).or_insert(0) += 1;
            }
            total += 1;
        }
    }

    (total, by_file)
}

/// Parse `Applied N changes` out of `sg -U` stderr.
fn parse_applied_count(stderr: &[u8]) -> Option<usize> {
    let text = String::from_utf8_lossy(stderr);
    for line in text.lines() {
        let trimmed = line.trim();
        if let Some(rest) = trimmed.strip_prefix("Applied ") {
            let num = rest.split_whitespace().next()?;
            return num.parse::<usize>().ok();
        }
    }
    None
}

/// Split a large path list into chunks no larger than `MAX_PATHS_PER_INVOCATION`.
fn chunk_paths(paths: Vec<PathBuf>) -> Vec<Vec<PathBuf>> {
    if paths.len() <= MAX_PATHS_PER_INVOCATION {
        return vec![paths];
    }
    paths
        .chunks(MAX_PATHS_PER_INVOCATION)
        .map(|c| c.to_vec())
        .collect()
}

#[async_trait]
impl AgentTool for AstEditTool {
    fn name(&self) -> &str {
        "ast_edit"
    }

    fn label(&self) -> &str {
        "AST Edit"
    }

    fn description(&self) -> &str {
        "AST-aware structural code rewriting using ast-grep. Provide an `ops` array of `{pat, out}` pattern→replacement pairs and `paths` to files/dirs/globs to apply them to. Pattern and replacement use ast-grep syntax (e.g. pat='fn $NAME() -> i32 { $BODY }', out='fn $NAME() -> i64 { $BODY }'). Set `dry_run=true` (default) to preview matches without writing; set `dry_run=false` to apply in place. Requires the `sg` CLI on PATH. Globs are expanded by this tool before invoking ast-grep because ast-grep does not expand globs in positional path arguments."
    }

    fn parameters_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "ops": {
                    "type": "array",
                    "description": "Rewrite operations. Each entry maps an ast-grep pattern (`pat`) to a replacement template (`out`). Metavariables in `pat` (e.g. `$NAME`, `$BODY`) are interpolated into `out` by ast-grep.",
                    "items": {
                        "type": "object",
                        "properties": {
                            "pat": {
                                "type": "string",
                                "description": "AST pattern in ast-grep syntax (e.g. 'fn $NAME() -> i32 { $BODY }')."
                            },
                            "out": {
                                "type": "string",
                                "description": "Replacement template (e.g. 'fn $NAME() -> i64 { $BODY }')."
                            }
                        },
                        "required": ["pat", "out"],
                        "additionalProperties": false
                    },
                    "minItems": 1
                },
                "paths": {
                    "type": "array",
                    "description": "Files, directories, or globs to rewrite. Globs (containing `*`, `?`, or `[`) are expanded in-process before invoking ast-grep because ast-grep does not expand globs in positional path arguments.",
                    "items": { "type": "string" },
                    "minItems": 1
                },
                "dry_run": {
                    "type": "boolean",
                    "description": "When true (default), only preview matches — no files are modified. When false, apply the rewrites in place using ast-grep's `--update-all`.",
                    "default": true
                }
            },
            "required": ["ops", "paths"]
        })
    }

    fn execution_mode(&self) -> ToolExecutionMode {
        // Bulk file rewrites that go through `sg -U` mutate many files at
        // once. Defaulting to ParallelSafe would let a concurrent `edit` or
        // another `ast_edit` race on the same paths and corrupt output.
        // Force sequential execution per batch, matching eval_tool /
        // debug_tool / browse_tool. This is correct regardless of the
        // call's dry_run flag (execution_mode is static per-tool, not
        // per-call).
        ToolExecutionMode::SequentialOnly
    }

    fn intent(&self) -> Option<&str> {
        Some("Applying AST rewrites")
    }

    async fn execute(
        &self,
        _tool_call_id: &str,
        params: Value,
        _signal: Option<oneshot::Receiver<()>>,
        ctx: &ToolContext,
    ) -> Result<AgentToolResult, ToolError> {
        // ── 1. Validate inputs ───────────────────────────────────────
        let ops = parse_ops(&params)?;
        let raw_paths = parse_paths(&params)?;

        // Default dry_run = true (per spec); do NOT copy edit.rs's default.
        let dry_run = params
            .get("dry_run")
            .and_then(Value::as_bool)
            .unwrap_or(true);

        // ── 2. Resolve & expand paths under the tool's root guard ─────
        let root = self.root_dir.as_deref().unwrap_or_else(|| ctx.root());
        let guard = PathGuard::new(root);

        let expanded = Self::expand_paths(&raw_paths, root)?;
        for p in &expanded {
            guard
                .validate(p)
                .map_err(|e| format!("Path '{}' rejected: {}", p.display(), e))?;
        }

        // ── 3. Run each op × each path-chunk ─────────────────────────
        let mut total_replacements: usize = 0;
        let mut total_files_touched: std::collections::BTreeSet<PathBuf> =
            std::collections::BTreeSet::new();
        let mut per_op_summary: Vec<String> = Vec::with_capacity(ops.len());
        let mut had_error = false;
        let mut error_messages: Vec<String> = Vec::new();

        let expanded_count = expanded.len();

        // Chunk the path list once and re-use across every op — avoids
        // re-cloning the (potentially large) expanded list per op.
        let chunks = chunk_paths(expanded);

        for (op_idx, op) in ops.iter().enumerate() {
            let mut op_count: usize = 0;
            let mut op_files: std::collections::BTreeSet<PathBuf> =
                std::collections::BTreeSet::new();

            for chunk in &chunks {
                match run_sg_for_op(op, chunk, dry_run).await {
                    Ok((status, stdout, stderr)) => {
                        if !status.success() {
                            let stderr_text = String::from_utf8_lossy(&stderr).trim().to_string();
                            // Strip the noisy deprecation banner if it's the
                            // only thing on stderr.
                            let cleaned: String = stderr_text
                                .lines()
                                .filter(|l| {
                                    !l.contains("`sg` is deprecated")
                                        && !l.contains("Use `ast-grep` instead")
                                        && !l.starts_with("======")
                                        && !l.trim().is_empty()
                                })
                                .collect::<Vec<_>>()
                                .join(" ");

                            let stdout_text = String::from_utf8_lossy(&stdout).trim().to_string();

                            if !cleaned.is_empty() {
                                had_error = true;
                                error_messages.push(format!(
                                    "ops[{}] (pat='{}') failed (exit {:?}): {}",
                                    op_idx,
                                    op.pat,
                                    status.code(),
                                    cleaned
                                ));
                            } else if !stdout_text.is_empty() {
                                had_error = true;
                                error_messages.push(format!(
                                    "ops[{}] (pat='{}') failed (exit {:?}): {}",
                                    op_idx,
                                    op.pat,
                                    status.code(),
                                    stdout_text
                                ));
                            }
                            // Otherwise: non-zero exit with empty output —
                            // some sg builds signal "no matches" this way.
                            continue;
                        }

                        if dry_run {
                            let (count, by_file) = summarise_dry_run(&stdout);
                            op_count += count;
                            for (f, _n) in by_file {
                                op_files.insert(f);
                            }
                        } else if let Some(n) = parse_applied_count(&stderr) {
                            op_count += n;
                            // sg doesn't tell us which files it touched, so
                            // we assume the entire path-chunk was in scope.
                            for p in chunk {
                                op_files.insert(p.clone());
                            }
                        } else {
                            // Apply succeeded but we couldn't parse the
                            // count — surface what we know.
                            for p in chunk {
                                op_files.insert(p.clone());
                            }
                            error_messages.push(format!(
                                "ops[{}] (pat='{}'): apply succeeded but could not parse 'Applied N changes' from stderr",
                                op_idx, op.pat
                            ));
                        }
                    }
                    Err(msg) => {
                        had_error = true;
                        error_messages.push(format!("ops[{}]: {}", op_idx, msg));
                    }
                }
            }

            total_replacements += op_count;
            for f in &op_files {
                total_files_touched.insert(f.clone());
            }

            let files_label = if op_files.is_empty() {
                "0 files".to_string()
            } else {
                format!("{} location(s)", op_files.len())
            };

            let summary_line = if dry_run {
                format!(
                    "ops[{}] pat='{}': {} match(es) across {}",
                    op_idx, op.pat, op_count, files_label
                )
            } else {
                format!(
                    "ops[{}] pat='{}' → out='{}': {} replacement(s) across {}",
                    op_idx, op.pat, op.out, op_count, files_label
                )
            };
            per_op_summary.push(summary_line);
        }

        // ── 4. Format output ─────────────────────────────────────────
        let header = if dry_run {
            "AST edit preview (dry-run) — no files modified"
        } else {
            "AST edit applied"
        };

        let mut body = String::new();
        body.push_str(header);
        body.push('\n');
        body.push('\n');
        for line in &per_op_summary {
            body.push_str(line);
            body.push('\n');
        }
        body.push('\n');

        if dry_run {
            body.push_str(&format!(
                "Total: {} match(es) across {} location(s)\n",
                total_replacements,
                total_files_touched.len()
            ));
        } else {
            body.push_str(&format!(
                "Total: {} replacement(s) across {} location(s)\n",
                total_replacements,
                total_files_touched.len()
            ));
        }

        if had_error {
            body.push_str("\nErrors:\n");
            for e in &error_messages {
                body.push_str(&format!("  - {}\n", e));
            }
        }

        let trimmed_body = body.trim_end().to_string();

        let mut result = if had_error && total_replacements == 0 {
            AgentToolResult::error(trimmed_body)
        } else {
            AgentToolResult::success(trimmed_body)
        };

        result.metadata = Some(json!({
            "dry_run": dry_run,
            "ops_count": ops.len(),
            "paths_count": expanded_count,
            "total_replacements": total_replacements,
            "locations_touched": total_files_touched.len(),
            "locations": total_files_touched.iter().map(|p| p.to_string_lossy().into_owned()).collect::<Vec<_>>(),
            "errors": error_messages,
        }));

        Ok(result)
    }
}