magi-code 0.61.0

use super::{
    ToolResult, ToolResultDisplay, ToolRuntime,
    args::{
        AST_GREP_DEFAULT_LIMIT, AST_GREP_STDERR_MAX_BYTES, AST_GREP_STDOUT_MAX_BYTES, AstGrepArgs,
    },
    contract::{metadata_key as meta, tool_name},
    fs::ExistingPathPolicy,
    process::{BoundedChildProcessLimits, run_bounded_child_process},
};
use crate::agent::cancellation::AgentCancellation;
use serde_json::json;
use std::{
    path::Path,
    process::{Command, Stdio},
    time::Duration,
};

#[cfg(unix)]
use std::os::unix::process::CommandExt;

const AST_GREP_TIMEOUT: Duration = Duration::from_secs(30);

const AST_GREP_MISSING: &str = "ast-grep is not installed. Install with: brew install ast-grep | cargo install ast-grep --locked | npm install --global @ast-grep/cli";

// ast-grep helper env is deny-by-default: PATH, HOME, TMPDIR, LANG, LC_*.
fn apply_ast_grep_env_allowlist(command: &mut Command) {
    command.env_clear();
    for (key, value) in std::env::vars_os() {
        if allowed_ast_grep_env(&key) {
            command.env(key, value);
        }
    }
}

fn allowed_ast_grep_env(key: &std::ffi::OsStr) -> bool {
    let Some(key) = key.to_str() else {
        return false;
    };
    matches!(key, "PATH" | "HOME" | "TMPDIR" | "LANG") || key.starts_with("LC_")
}

impl ToolRuntime {
    pub(super) fn ast_grep(
        &self,
        args: AstGrepArgs,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<ToolResult> {
        self.ast_grep_with_bin(args, Path::new("ast-grep"), cancellation)
    }

    pub(super) fn ast_grep_with_bin(
        &self,
        args: AstGrepArgs,
        bin: &Path,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<ToolResult> {
        let limit = args.limit.unwrap_or(AST_GREP_DEFAULT_LIMIT);
        let path_arg = args.path.clone().unwrap_or_else(|| ".".to_string());
        let resolved = self.resolve_existing_path(
            &path_arg,
            ExistingPathPolicy::ast_grep(self.ast_grep_absolute_paths),
        )?;
        cancellation.check()?;

        let mut command = Command::new(bin);
        command
            .arg("run")
            .arg("--pattern")
            .arg(&args.pattern)
            .arg("--heading=never");
        if let Some(lang) = &args.language {
            command.arg("--lang").arg(lang);
        }
        if let Some(rewrite) = &args.rewrite {
            command.arg("--rewrite").arg(rewrite);
        }
        command
            .arg("--")
            .arg(&resolved)
            .current_dir(&self.cwd)
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        apply_ast_grep_env_allowlist(&mut command);
        #[cfg(unix)]
        {
            command.process_group(0);
        }

        let child = match command.spawn() {
            Ok(child) => child,
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
                return Ok(ast_grep_failure_result(AST_GREP_MISSING));
            }
            Err(error) => return Err(error.into()),
        };

        let output = run_bounded_child_process(
            child,
            BoundedChildProcessLimits {
                stdout_max_bytes: AST_GREP_STDOUT_MAX_BYTES,
                stderr_max_bytes: AST_GREP_STDERR_MAX_BYTES,
                timeout: AST_GREP_TIMEOUT,
                poll_interval: Duration::from_millis(20),
            },
            cancellation,
        )?;
        let stdout = output.stdout;
        let stderr = output.stderr;
        let cleanup_warning = output.cleanup_warning;

        let stdout_locally_truncated = output.stdout_truncated;
        let stderr_locally_truncated = output.stderr_truncated;

        let exit_code = output
            .status
            .as_ref()
            .and_then(|status| status.code())
            .unwrap_or(-1);

        // ast-grep exits 1 for "no matches found" (like rg), which is not an error.
        let no_matches = exit_code == 1 && stdout.trim().is_empty();
        let has_cleanup_warning = cleanup_warning.is_some();

        let lines: Vec<&str> = stdout.lines().take(limit).collect();
        let returned_lines = lines.len();
        let content = lines.join("\n");
        let total_lines = stdout.lines().count();
        let truncated =
            total_lines > returned_lines || stdout_locally_truncated || stderr_locally_truncated;

        let success = !output.timed_out
            && !has_cleanup_warning
            && !cancellation.is_canceled()
            && (exit_code == 0 || no_matches);

        let mut metadata = json!({
            meta::ENGINE: "ast-grep",
            meta::EXIT_CODE: exit_code,
            meta::TIMED_OUT: output.timed_out,
            meta::MATCHES_RETURNED: returned_lines,
            meta::TRUNCATED: truncated,
        });
        if let Some(lang) = &args.language {
            metadata[meta::LANGUAGE] = json!(lang);
        }
        if args.rewrite.is_some() {
            metadata[meta::REWRITE] = json!(true);
        }
        metadata[meta::STDOUT_TRUNCATED] = json!(stdout_locally_truncated);
        metadata[meta::STDERR_TRUNCATED] = json!(stderr_locally_truncated);
        if let Some(warning) = &cleanup_warning {
            metadata[meta::CLEANUP_WARNING] = json!(warning);
        }

        let display = ToolResultDisplay::default();

        let final_content = if !stderr.is_empty() && !success && !no_matches {
            format!("{stdout}{stderr}")
        } else {
            content
        };

        Ok(ToolResult {
            tool_name: tool_name::AST_GREP.to_string(),
            success,
            content: final_content,
            metadata,
            display,
        })
    }
}

fn ast_grep_failure_result(message: &str) -> ToolResult {
    ToolResult {
        tool_name: tool_name::AST_GREP.to_string(),
        success: false,
        content: message.to_string(),
        metadata: json!({
            meta::ENGINE: "ast-grep",
        }),
        display: ToolResultDisplay::default(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agent::cancellation::AgentCancellation;
    use serde_json::json;
    use std::{
        fs,
        path::Path,
        sync::{
            Arc,
            atomic::{AtomicBool, Ordering},
        },
        thread,
        time::Duration,
    };
    use tempfile::TempDir;

    #[cfg(unix)]
    fn make_executable(path: &Path) {
        use std::os::unix::fs::PermissionsExt;
        let mut permissions = fs::metadata(path).unwrap().permissions();
        permissions.set_mode(0o755);
        fs::set_permissions(path, permissions).unwrap();
    }

    fn runtime() -> (TempDir, ToolRuntime) {
        let temp = TempDir::new().unwrap();
        let runtime = ToolRuntime::new(temp.path()).unwrap();
        (temp, runtime)
    }

    #[test]
    fn rejects_empty_pattern() {
        let (_temp, runtime) = runtime();
        let result = runtime.dispatch("ast_grep", json!({"pattern": "  "}));
        assert!(!result.success);
        assert!(result.content.contains("pattern must not be empty"));
    }

    #[test]
    fn rejects_zero_limit() {
        let (_temp, runtime) = runtime();
        let result = runtime.dispatch("ast_grep", json!({"pattern": "test", "limit": 0}));
        assert!(!result.success);
        assert!(result.content.contains("limit must be at least 1"));
    }

    #[test]
    fn rejects_excessive_limit() {
        let (_temp, runtime) = runtime();
        let result = runtime.dispatch("ast_grep", json!({"pattern": "test", "limit": 501}));
        assert!(!result.success);
        assert!(result.content.contains("limit must be at most"));
    }

    #[test]
    fn rejects_unknown_field() {
        let (_temp, runtime) = runtime();
        let result = runtime.dispatch("ast_grep", json!({"pattern": "test", "bogus": true}));
        assert!(!result.success);
    }

    #[cfg(unix)]
    #[test]
    fn missing_binary_returns_install_message() {
        let (temp, runtime) = runtime();
        let result = runtime
            .ast_grep_with_bin(
                AstGrepArgs {
                    pattern: "test".to_string(),
                    language: None,
                    path: Some(".".to_string()),
                    rewrite: None,
                    limit: None,
                },
                Path::new("/nonexistent/ast-grep-binary-12345"),
                &AgentCancellation::default(),
            )
            .unwrap();
        assert!(!result.success);
        assert!(result.content.contains("ast-grep is not installed"));
        assert_eq!(result.metadata[meta::ENGINE], "ast-grep");
        let _ = temp;
    }

    #[cfg(unix)]
    #[test]
    fn fake_binary_returns_matches() {
        let (temp, runtime) = runtime();
        let fake = temp.path().join("fake-ast-grep");
        fs::write(
            &fake,
            "#!/bin/sh\necho 'src/main.rs:5:unwrap()'\necho 'src/lib.rs:12:unwrap()'\nexit 0\n",
        )
        .unwrap();
        make_executable(&fake);

        let result = runtime
            .ast_grep_with_bin(
                AstGrepArgs {
                    pattern: "unwrap()".to_string(),
                    language: Some("rust".to_string()),

                    path: Some(".".to_string()),
                    rewrite: None,
                    limit: Some(10),
                },
                &fake,
                &AgentCancellation::default(),
            )
            .unwrap();

        assert!(result.success, "{}", result.content);
        assert!(result.content.contains("src/main.rs:5:unwrap()"));
        assert!(result.content.contains("src/lib.rs:12:unwrap()"));
        assert_eq!(result.metadata[meta::ENGINE], "ast-grep");
        assert_eq!(result.metadata[meta::EXIT_CODE], 0);
        assert_eq!(result.metadata[meta::MATCHES_RETURNED], 2);
        assert_eq!(result.metadata[meta::TRUNCATED], false);
        assert_eq!(result.metadata[meta::LANGUAGE], "rust");
    }

    #[cfg(unix)]
    #[test]
    fn fake_binary_does_not_inherit_credential_env() {
        let (temp, runtime) = runtime();
        let fake = temp.path().join("fake-ast-grep-env");
        fs::write(
            &fake,
            "#!/bin/sh\nprintf '%s' \"$MAGI_CODE_TEST_ENV_PROBE_AST\"\n",
        )
        .unwrap();
        make_executable(&fake);
        unsafe { std::env::set_var("MAGI_CODE_TEST_ENV_PROBE_AST", "leaked") };

        let result = runtime
            .ast_grep_with_bin(
                AstGrepArgs {
                    pattern: "test".to_string(),
                    language: None,
                    path: Some(".".to_string()),
                    rewrite: None,
                    limit: None,
                },
                &fake,
                &AgentCancellation::default(),
            )
            .unwrap();
        unsafe { std::env::remove_var("MAGI_CODE_TEST_ENV_PROBE_AST") };

        assert!(result.success, "{}", result.content);
        assert_eq!(result.content, "");
    }

    #[cfg(unix)]
    #[test]
    fn no_matches_exit_one_is_success() {
        let (temp, runtime) = runtime();
        let fake = temp.path().join("fake-ast-grep");
        fs::write(&fake, "#!/bin/sh\nexit 1\n").unwrap();
        make_executable(&fake);

        let result = runtime
            .ast_grep_with_bin(
                AstGrepArgs {
                    pattern: "nomatch".to_string(),
                    language: None,
                    path: Some(".".to_string()),
                    rewrite: None,
                    limit: None,
                },
                &fake,
                &AgentCancellation::default(),
            )
            .unwrap();

        assert!(result.success, "no matches should not be failure");
        assert_eq!(result.content, "");
        assert_eq!(result.metadata[meta::EXIT_CODE], 1);
        assert_eq!(result.metadata[meta::MATCHES_RETURNED], 0);
    }

    #[cfg(unix)]
    #[test]
    fn limit_caps_returned_lines_and_sets_truncated() {
        let (temp, runtime) = runtime();
        let fake = temp.path().join("fake-ast-grep");
        let lines: Vec<String> = (0..10).map(|i| format!("file:{i}:match")).collect();
        fs::write(
            &fake,
            format!("#!/bin/sh\necho '{}'\nexit 0\n", lines.join("\\n")),
        )
        .unwrap();
        make_executable(&fake);

        let result = runtime
            .ast_grep_with_bin(
                AstGrepArgs {
                    pattern: "match".to_string(),
                    language: None,
                    path: Some(".".to_string()),
                    rewrite: None,
                    limit: Some(3),
                },
                &fake,
                &AgentCancellation::default(),
            )
            .unwrap();

        assert!(result.success, "{}", result.content);
        assert_eq!(result.metadata[meta::MATCHES_RETURNED], 3);
        assert_eq!(result.metadata[meta::TRUNCATED], true);
    }

    #[cfg(unix)]
    #[test]
    fn timeout_produces_failure() {
        let (temp, runtime) = runtime();
        let fake = temp.path().join("fake-ast-grep-sleep");
        fs::write(&fake, "#!/bin/sh\necho 'match'\nsleep 60\n").unwrap();
        make_executable(&fake);

        let result = runtime
            .ast_grep_with_bin(
                AstGrepArgs {
                    pattern: "match".to_string(),
                    language: None,
                    path: Some(".".to_string()),
                    rewrite: None,
                    limit: None,
                },
                &fake,
                &AgentCancellation::default(),
            )
            .unwrap();

        // Override timeout for test speed by checking we get timed_out metadata
        // Note: this test will take ~30s due to AST_GREP_TIMEOUT
        assert!(!result.success);
        assert_eq!(result.metadata[meta::TIMED_OUT], true);
    }

    #[cfg(unix)]
    #[test]
    fn cancellation_aborts_promptly() {
        let (temp, runtime) = runtime();
        let fake = temp.path().join("fake-ast-grep-sleep");
        fs::write(&fake, "#!/bin/sh\nsleep 60\n").unwrap();
        make_executable(&fake);

        let cancel_flag = Arc::new(AtomicBool::new(false));
        let cancellation = AgentCancellation::new(Arc::clone(&cancel_flag));
        thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));
            cancel_flag.store(true, Ordering::SeqCst);
        });

        let result = runtime.ast_grep_with_bin(
            AstGrepArgs {
                pattern: "match".to_string(),
                language: None,
                path: Some(".".to_string()),
                rewrite: None,
                limit: None,
            },
            &fake,
            &cancellation,
        );

        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("prompt canceled"));
    }

    #[cfg(unix)]
    #[test]
    fn stderr_failure_included_in_content() {
        let (temp, runtime) = runtime();
        let fake = temp.path().join("fake-ast-grep-err");
        fs::write(&fake, "#!/bin/sh\necho 'bad pattern' >&2\nexit 2\n").unwrap();
        make_executable(&fake);

        let result = runtime
            .ast_grep_with_bin(
                AstGrepArgs {
                    pattern: "test".to_string(),
                    language: None,
                    path: Some(".".to_string()),
                    rewrite: None,
                    limit: None,
                },
                &fake,
                &AgentCancellation::default(),
            )
            .unwrap();

        assert!(!result.success);
        assert!(result.content.contains("bad pattern"));
        assert_eq!(result.metadata[meta::EXIT_CODE], 2);
    }

    #[test]
    fn schema_registered_in_provider_definitions() {
        let definitions = crate::tools::capability::mvp_tool_definitions_json();
        let ast_grep = definitions
            .as_array()
            .unwrap()
            .iter()
            .find(|d| d["name"] == "ast_grep")
            .expect("ast_grep definition must exist");
        let params = &ast_grep["parameters"];
        assert_eq!(params["required"], json!(["pattern"]));
        assert_eq!(params["additionalProperties"], false);
        assert!(params["properties"].get("pattern").is_some());
        assert!(params["properties"].get("language").is_some());
        assert!(params["properties"].get("path").is_some());
        assert!(params["properties"].get("rewrite").is_some());
        assert!(params["properties"].get("limit").is_some());
    }

    #[test]
    fn empty_rewrite_ignored_not_error() {
        let (_temp, runtime) = runtime();
        let result = runtime.dispatch("ast_grep", json!({"pattern": "test", "rewrite": "   "}));
        assert!(
            !result.content.contains("rewrite must not be empty"),
            "empty rewrite should be ignored, not error: {}",
            result.content
        );
    }

    #[test]
    fn dispatch_name_resolves_to_capability() {
        assert_eq!(
            crate::tools::ToolCapability::from_dispatch_name("ast_grep"),
            Some(crate::tools::ToolCapability::AstGrep)
        );
    }
}