agent-manager/src/shell_ai.rs

//! am ai — langage naturel → action shell (issues #96 #97).
//!
//! The Shell AI command turns a natural-language request into a shell
//! action using AIChat (the `aichat` catalog agent) as the generation
//! engine:
//!
//! - conversational mode (no `--exec`): `aichat -f <ctx> "<prompt>"` is
//!   spawned in the foreground, exactly as the user would run it;
//! - exec mode (`--exec`): aichat is asked to generate the command with
//!   `--code` (it prints only the command and never runs it — aichat is
//!   invoked without `--execute`), then the generated command is
//!   classified (safe / risky), a confidence score
//!   is estimated, the configured safety policy is applied (dry-run by
//!   default), and only then — after confirmation when required — is the
//!   command executed through the user's shell.
//!
//! Nothing is ever executed without an explicit confirmation: the default
//! policy is `dry-run`, overridden by `--yes` (execute without asking),
//! the `settings.shell_ai.default_safety` setting (dry-run | confirm |
//! auto) or the explicit `--dry-run` flag (never execute, show only).

use crate::app::App;
use crate::commands::{require_agent, resolve_exec};
use anyhow::{anyhow, bail, Context, Result};
use std::collections::BTreeMap;
use std::process::Command;

/// Safety policies of `am ai --exec` (issue #97).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SafetyMode {
    /// Show the generated command only — nothing is executed unless the
    /// user passes `--yes`.
    DryRun,
    /// Execute safe commands directly; ask before risky ones.
    Confirm,
    /// Execute everything without asking (explicitly enabled by the user).
    Auto,
}

impl SafetyMode {
    pub fn as_str(&self) -> &'static str {
        match self {
            SafetyMode::DryRun => "dry-run",
            SafetyMode::Confirm => "confirm",
            SafetyMode::Auto => "auto",
        }
    }
}

/// Risk class of a generated shell command.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Risk {
    Safe,
    Risky,
}

impl Risk {
    pub fn as_str(&self) -> &'static str {
        match self {
            Risk::Safe => "safe",
            Risk::Risky => "risky",
        }
    }
}

/// Built-in substrings that mark a command as risky (issue #97). The user
/// can extend this list with `settings.shell_ai.risky_patterns`.
pub const BUILTIN_RISKY_PATTERNS: &[&str] = &[
    "rm -rf",
    "rm -fr",
    "rm -r -f",
    "dd if=",
    "mkfs.",
    "fdisk",
    "parted",
    "gdisk",
    ":(){",
    "chmod -R 777",
    "chmod 777 /",
    "chown -R",
    "> /dev/sd",
    ">/dev/sd",
    "sudo rm",
    "git push --force",
    "git push -f",
    "drop database",
    "drop table",
    "truncate table",
    "shutdown",
    "reboot",
    "poweroff",
    "kill -9",
    "pkill -9",
    "killall",
    "init 0",
    "init 6",
    "mv / ",
    "rm /",
    "format c:",
    "del /f /s",
    "rd /s",
    "cipher /w",
    "curl ... | sh",
    "curl ... | bash",
    "wget ... | sh",
];

/// Classify a generated command against the built-in patterns plus the
/// user-configured ones (`settings.shell_ai.risky_patterns`).
pub fn classify(cmd: &str, extra_patterns: &[String]) -> Risk {
    let lower = cmd.to_lowercase();
    let hits = BUILTIN_RISKY_PATTERNS
        .iter()
        .any(|p| lower.contains(&p.to_lowercase()))
        || extra_patterns.iter().any(|p| lower.contains(&p.to_lowercase()));
    if hits {
        Risk::Risky
    } else {
        Risk::Safe
    }
}

/// A coarse confidence heuristic (0-100) shown to the user. Risky or
/// compound commands are scored lower; simple, read-only commands score
/// higher. It is an estimate, never a guarantee.
pub fn estimate_certainty(cmd: &str, risk: Risk) -> u8 {
    let mut score: i32 = 92;
    if risk == Risk::Risky {
        score -= 25;
    }
    // Compound commands chain several effects — harder to predict.
    for sep in ["&&", "||", ";\n", "\n", " | ", " 2>"] {
        if cmd.contains(sep) {
            score -= 6;
        }
    }
    // Redirections and pipes move data around.
    if cmd.contains('>') || cmd.contains('<') {
        score -= 5;
    }
    if cmd.contains("sudo") {
        score -= 8;
    }
    score.clamp(40, 99) as u8
}

/// Decide what to do with a generated command under the effective policy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Decision {
    /// Print and abort (dry-run or declined confirmation).
    Abort,
    /// Ask the user (y/N) before executing.
    Ask,
    /// Execute directly.
    Execute,
}

/// Apply the safety policy: explicit `--dry-run` flag > `--yes` flag >
/// configured default mode (fallback: dry-run).
pub fn decide(
    mode: SafetyMode,
    risk: Risk,
    dry_run_flag: bool,
    yes_flag: bool,
) -> Decision {
    if dry_run_flag {
        return Decision::Abort;
    }
    if yes_flag {
        return Decision::Execute;
    }
    match mode {
        SafetyMode::Auto => Decision::Execute,
        SafetyMode::Confirm => {
            if risk == Risk::Risky {
                Decision::Ask
            } else {
                Decision::Execute
            }
        }
        SafetyMode::DryRun => Decision::Abort,
    }
}

/// Extract the shell command from aichat's `--dry-run` output: strips a
/// fenced code block if present, otherwise uses the trimmed output as-is.
pub fn extract_command(stdout: &str) -> Option<String> {
    let trimmed = stdout.trim();
    if trimmed.is_empty() {
        return None;
    }
    // Fenced block: ```bash ... ``` (or ```sh, ```powershell, ```cmd ...)
    let mut lines = trimmed.lines();
    if lines.next().is_some_and(|l| l.trim_start().starts_with("```")) {
        let body: Vec<&str> = lines
            .take_while(|l| !l.trim().starts_with("```"))
            .collect();
        let cmd = body.join("\n").trim().to_string();
        if !cmd.is_empty() {
            return Some(cmd);
        }
    }
    Some(trimmed.to_string())
}

/// Map a provider of the registry to the environment variables AIChat
/// understands, and resolve the model to pass with `--model`. The token is
/// resolved from the OS keyring via the standard `@secret` mechanism and
/// never appears on the command line (issue #89).
pub fn provider_env(
    app: &App,
    provider: Option<&str>,
    model: Option<&str>,
) -> Result<(Vec<String>, BTreeMap<String, String>)> {
    // Without any flag, aichat keeps its own configuration (its config file
    // and API keys) — the registry is only applied on explicit request.
    if provider.is_none() && model.is_none() {
        return Ok((Vec::new(), BTreeMap::new()));
    }
    let Some((pname, def, resolved_model)) =
        crate::providers::resolve_for(&app.config, None, None, provider, model)
    else {
        if provider.is_some() {
            let known = crate::providers::names(&app.config);
            let hint = if known.is_empty() {
                "aucun (am providers add <nom> --base-url <url>)".to_string()
            } else {
                known.join(", ")
            };
            bail!(crate::i18n::tr_fmt_in(
                app.lang(),
                "provider '{}' inconnu — providers enregistrés: {}",
                &[&provider.unwrap_or(""), &hint]
            ));
        }
        return Ok((Vec::new(), BTreeMap::new()));
    };
    let mut args: Vec<String> = Vec::new();
    let mut env: BTreeMap<String, String> = BTreeMap::new();
    // Known providers map to AIChat's standard API-key variables; custom
    // provider names fall back to the openai-compatible channel.
    let key_var = match pname {
        "anthropic" => Some("ANTHROPIC_API_KEY"),
        "openai" => Some("OPENAI_API_KEY"),
        "deepseek" => Some("DEEPSEEK_API_KEY"),
        "google" => Some("GOOGLE_API_KEY"),
        "groq" => Some("GROQ_API_KEY"),
        "openrouter" => Some("OPENROUTER_API_KEY"),
        "xai" | "grok" => Some("XAI_API_KEY"),
        "ollama" => None, // local — no key; aichat knows its default endpoint
        _ => Some("OPENAI_API_KEY"),
    };
    if let Some(var) = key_var {
        env.insert(var.to_string(), "@secret".to_string());
        if !matches!(pname, "openai" | "ollama" | "anthropic" | "deepseek" | "google" | "groq" | "openrouter" | "xai" | "grok") {
            // Custom endpoints ride the openai-compatible channel.
            env.insert("OPENAI_API_BASE".to_string(), def.base_url.clone());
        }
    }
    if let Some(m) = resolved_model {
        args.push("--model".to_string());
        args.push(m.to_string());
    }
    let warnings = crate::secrets::resolve_env_secrets(
        &crate::secrets::store(),
        "aichat",
        Some(pname),
        &mut env,
    );
    for w in warnings {
        app.log.warn(&w);
    }
    Ok((args, env))
}

/// Generate a shell command for `prompt` using aichat in dry-run mode.
/// Returns the generated command (never executed by aichat).
pub fn generate(
    app: &App,
    prompt: &str,
    files: &[String],
    provider: Option<&str>,
    model: Option<&str>,
    role_args: &[String],
) -> Result<String> {
    let agent = require_agent(app, "aichat")?;
    // aichat >= 0.30 renamed --exec to --execute AND changed --dry-run to
    // skip the API call entirely (it echoes the request), so the old
    // `--execute --dry-run` no longer yields the generated command. Plain
    // chat with `--code` extracts and prints only the command (same output
    // shape, works on every aichat version). aichat NEVER executes here —
    // am enforces the safety policy before running the command itself.
    let mut extra_args: Vec<String> = vec!["--code".to_string()];
    extra_args.extend(role_args.iter().cloned());
    for f in files {
        extra_args.push("-f".to_string());
        extra_args.push(f.clone());
    }
    extra_args.push(prompt.to_string());
    let (p_args, p_env) = provider_env(app, provider, model)?;
    extra_args.extend(p_args);
    let exec = resolve_exec(app, agent, &extra_args, &p_env)?;
    app.log.verbose(&format!(
        "shell-ai: {} {} (dry-run generation)",
        exec.program,
        exec.args.join(" ")
    ));
    let (prog, prefix) = crate::runner::resolve_program(&exec.program);
    let mut full_args = prefix;
    full_args.extend(exec.args.iter().cloned());
    let out = Command::new(&prog)
        .args(&full_args)
        .envs(&exec.env)
        .output()
        .with_context(|| format!("failed to run {}", exec.program))?;
    if !out.status.success() {
        let err = String::from_utf8_lossy(&out.stderr);
        let err = err.trim();
        bail!(crate::i18n::tr_fmt_in(
            app.lang(),
            "aichat n'a pas généré de commande (exit {}){}",
            &[
                &out.status.code().unwrap_or(-1).to_string(),
                &if err.is_empty() {
                    String::new()
                } else {
                    format!(": {err}")
                },
            ]
        ));
    }
    let stdout = String::from_utf8_lossy(&out.stdout);
    extract_command(&stdout).ok_or_else(|| {
        anyhow!(crate::i18n::tr_in(
            app.lang(),
            "aichat n'a retourné aucune commande (dry-run)"
        ))
    })
}

/// Execute a generated command through the user's shell (default_shell >
/// $SHELL > $COMSPEC > cmd). The command is passed as a single argument.
pub fn execute(app: &App, command: &str) -> Result<i32> {
    let spec = crate::shell::resolve_default(
        app.config.settings.default_shell.as_deref(),
        std::env::var_os("SHELL"),
        std::env::var_os("COMSPEC"),
    );
    app.log.verbose(&format!(
        "shell-ai: executing via {} {}",
        spec.program,
        spec.args.join(" ")
    ));
    let (prog, prefix) = crate::runner::resolve_program(spec.program);
    let mut args = prefix;
    args.extend(spec.args.iter().map(|s| s.to_string()));
    args.push(command.to_string());
    let status = Command::new(&prog)
        .args(&args)
        .status()
        .with_context(|| format!("failed to run {}", spec.program))?;
    Ok(status.code().unwrap_or(1))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn classifies_safe_and_risky() {
        assert_eq!(classify("ls -la", &[]), Risk::Safe);
        assert_eq!(classify("echo hello", &[]), Risk::Safe);
        assert_eq!(classify("find . -name '*.json'", &[]), Risk::Safe);
        assert_eq!(classify("rm -rf /tmp/x", &[]), Risk::Risky);
        assert_eq!(classify("sudo rm -rf /var/log", &[]), Risk::Risky);
        assert_eq!(classify("dd if=/dev/zero of=/dev/sda", &[]), Risk::Risky);
        assert_eq!(classify("git push --force origin main", &[]), Risk::Risky);
        assert_eq!(classify("drop database prod;", &[]), Risk::Risky);
    }

    #[test]
    fn classifies_extra_user_patterns() {
        let extra = vec!["killall node".to_string()];
        assert_eq!(classify("killall node", &extra), Risk::Risky);
        assert_eq!(classify("killall nodejs", &extra), Risk::Risky);
        assert_eq!(classify("node --version", &extra), Risk::Safe);
    }

    #[test]
    fn certainty_stays_in_bounds_and_penalizes_risk() {
        let safe = estimate_certainty("ls -la", Risk::Safe);
        let risky = estimate_certainty("rm -rf /", Risk::Risky);
        assert!(safe > risky, "{safe} should be > {risky}");
        assert!((40..=99).contains(&safe));
        assert!((40..=99).contains(&risky));
        let compound = estimate_certainty("rm -rf /tmp/a && echo ok", Risk::Risky);
        assert!(compound < risky || compound == risky);
    }

    #[test]
    fn extract_command_handles_fences_and_plain() {
        assert_eq!(
            extract_command("```bash\nrm *.tmp\n```"),
            Some("rm *.tmp".to_string())
        );
        assert_eq!(
            extract_command("```sh\necho hello\n```\n"),
            Some("echo hello".to_string())
        );
        assert_eq!(
            extract_command("rm *.tmp\n"),
            Some("rm *.tmp".to_string())
        );
        assert_eq!(extract_command("   \n  "), None);
    }

    #[test]
    fn decision_matrix() {
        // dry-run (default): abort, even with --yes overriding to execute.
        assert_eq!(
            decide(SafetyMode::DryRun, Risk::Risky, false, false),
            Decision::Abort
        );
        assert_eq!(
            decide(SafetyMode::DryRun, Risk::Safe, false, false),
            Decision::Abort
        );
        assert_eq!(
            decide(SafetyMode::DryRun, Risk::Risky, false, true),
            Decision::Execute
        );
        // explicit --dry-run wins over --yes.
        assert_eq!(
            decide(SafetyMode::Auto, Risk::Risky, true, true),
            Decision::Abort
        );
        // confirm: ask only for risky.
        assert_eq!(
            decide(SafetyMode::Confirm, Risk::Safe, false, false),
            Decision::Execute
        );
        assert_eq!(
            decide(SafetyMode::Confirm, Risk::Risky, false, false),
            Decision::Ask
        );
        // auto: execute everything.
        assert_eq!(
            decide(SafetyMode::Auto, Risk::Risky, false, false),
            Decision::Execute
        );
    }

    #[test]
    fn safety_mode_parsing_falls_back_to_dry_run() {
        use crate::config::ShellAiSettings;
        let s = ShellAiSettings {
            default_safety: Some("confirm".to_string()),
            risky_patterns: vec![],
        };
        assert_eq!(s.safety_mode(), SafetyMode::Confirm);
        let s = ShellAiSettings {
            default_safety: Some("n'importe quoi".to_string()),
            risky_patterns: vec![],
        };
        assert_eq!(s.safety_mode(), SafetyMode::DryRun);
        let s = ShellAiSettings {
            default_safety: None,
            risky_patterns: vec![],
        };
        assert_eq!(s.safety_mode(), SafetyMode::DryRun);
    }
}