ai_env 0.1.3

mod antigravity;
mod api;
mod augment;
mod claude;
mod claude_cowork;
mod codex;
mod cursor;
mod cursor_cli;
mod devin;
mod gemini;
mod github_copilot;
mod opencode;
mod replit;
mod v0;

pub use api::{AiAgent, AiEnvironment, AiNetworkPolicy};

use std::collections::HashMap;
use std::env;
use std::net::IpAddr;
use std::path::PathBuf;
use std::sync::OnceLock;

/// Returns true if the process is being run by an AI coding agent.
pub fn is_ai_agent() -> bool {
    get_environment().is_some()
}

static AGENT: OnceLock<AiAgent> = OnceLock::new();

/// Detects the AI agent running the process. Returns `Unknown` when no known
/// agent is matched.
///
/// Note that an unrecognized agent may still be present via the `AI_AGENT`
/// self-identification variable; prefer [`is_ai_agent`] or [`get_environment`]
/// to detect that case.
pub fn detect_agent() -> AiAgent {
    *AGENT.get_or_init(|| {
        let vars = env::vars().collect::<HashMap<_, _>>();

        detect_agent_from_vars(&vars)
    })
}

/// Detection markers in priority order, mirroring `@vercel/detect-agent`.
///
/// Unlike the `ci_env`/`cd_env` crates, detection isn't a flat key table:
/// several agents need value checks (`CURSOR_EXTENSION_HOST_ROLE`), multiple
/// alternative markers (Codex, Copilot), or a sub-mode gate (Claude vs its
/// Cowork mode), so the checks are spelled out explicitly. Markers are checked
/// before the `AI_AGENT` self-id so the Cowork sub-mode is detected even when
/// `AI_AGENT` is also set to the generic Claude Code string.
fn detect_agent_from_vars(vars: &HashMap<String, String>) -> AiAgent {
    let get = |key: &str| vars.get(key).map(|v| v.as_str()).filter(|v| !v.is_empty());
    let any = |keys: &[&str]| keys.iter().any(|key| get(key).is_some());

    // Cursor editor integrated terminal
    if get("CURSOR_TRACE_ID").is_some() {
        return AiAgent::Cursor;
    }

    // Cursor CLI agent (env marker or the agent-exec extension host role)
    if get("CURSOR_AGENT").is_some() || get("CURSOR_EXTENSION_HOST_ROLE") == Some("agent-exec") {
        return AiAgent::CursorCli;
    }

    if get("GEMINI_CLI").is_some() {
        return AiAgent::Gemini;
    }

    // Codex — sandboxed exec, CI, or an interactive thread
    if any(&["CODEX_SANDBOX", "CODEX_CI", "CODEX_THREAD_ID"]) {
        return AiAgent::Codex;
    }

    if get("ANTIGRAVITY_AGENT").is_some() {
        return AiAgent::Antigravity;
    }

    if get("AUGMENT_AGENT").is_some() {
        return AiAgent::Augment;
    }

    if any(&["OPENCODE_CLIENT", "OPENCODE"]) {
        return AiAgent::OpenCode;
    }

    // Claude Code, or its Cowork sub-mode
    if any(&["CLAUDECODE", "CLAUDE_CODE"]) {
        return if get("CLAUDE_CODE_IS_COWORK").is_some() {
            AiAgent::ClaudeCowork
        } else {
            AiAgent::Claude
        };
    }

    if get("REPL_ID").is_some() {
        return AiAgent::Replit;
    }

    // GitHub Copilot CLI
    if any(&[
        "COPILOT_CLI",
        "COPILOT_MODEL",
        "COPILOT_ALLOW_ALL",
        "COPILOT_GITHUB_TOKEN",
    ]) {
        return AiAgent::GithubCopilot;
    }

    // Fall back to the `AI_AGENT` self-identification string. Agents that only
    // self-identify (no boolean marker, e.g. v0) are classified by name.
    if let Some(value) = get("AI_AGENT") {
        let agent = classify_self_id(value);

        if !matches!(agent, AiAgent::Unknown) {
            return agent;
        }
    }

    if PathBuf::from("/opt/.devin").exists() {
        return AiAgent::Devin;
    }

    AiAgent::Unknown
}

/// Maps an `AI_AGENT` self-identification string to a known agent. A few names
/// are matched exactly (`v0`, `github-copilot`); the rest by name prefix
/// (e.g. `claude-code_2-1-170_agent` -> `Claude`).
fn classify_self_id(value: &str) -> AiAgent {
    let name = value.trim().to_ascii_lowercase();

    if name == "v0" {
        AiAgent::V0
    } else if name == "github-copilot" || name == "github-copilot-cli" {
        AiAgent::GithubCopilot
    } else if name.starts_with("claude") {
        AiAgent::Claude
    } else if name.starts_with("cowork") {
        AiAgent::ClaudeCowork
    } else if name.starts_with("codex") {
        AiAgent::Codex
    } else if name.starts_with("cursor") {
        AiAgent::Cursor
    } else if name.starts_with("gemini") {
        AiAgent::Gemini
    } else if name.starts_with("antigravity") {
        AiAgent::Antigravity
    } else if name.starts_with("augment") {
        AiAgent::Augment
    } else if name.starts_with("opencode") {
        AiAgent::OpenCode
    } else if name.starts_with("devin") {
        AiAgent::Devin
    } else if name.starts_with("replit") {
        AiAgent::Replit
    } else {
        AiAgent::Unknown
    }
}

static NETWORK_POLICY: OnceLock<AiNetworkPolicy> = OnceLock::new();

/// Detects the network/egress policy imposed on the process, from environment
/// variables only (no network I/O is performed). Pass the agent governing the
/// process, typically from [`detect_agent`].
///
/// Agent-specific signals are only honored when the variable belongs to the
/// given agent, so a variable leaked from another tool cannot misclassify the
/// environment:
/// - Codex — `CODEX_SANDBOX_NETWORK_DISABLED`, set on sandboxed commands with
///   network access turned off
/// - GitHub Copilot — `COPILOT_AGENT_FIREWALL_*`, the cloud agent firewall,
///   where `COPILOT_AGENT_FIREWALL_ENABLED=false` reports it disabled
///
/// Independent of the agent, `HTTPS_PROXY`-family variables pointing at a
/// host that only exists locally indicate a filtering/egress proxy (e.g. the
/// Claude Code sandbox).
///
/// A [`Filtered`][AiNetworkPolicy::Filtered] network may still allow the
/// specific hosts a program needs, so verify reachability of those hosts
/// before acting on this signal.
pub fn detect_network_policy(agent: AiAgent) -> AiNetworkPolicy {
    *NETWORK_POLICY.get_or_init(|| {
        let vars = env::vars().collect::<HashMap<_, _>>();

        detect_network_policy_from_vars(agent, &vars)
    })
}

fn detect_network_policy_from_vars(
    agent: AiAgent,
    vars: &HashMap<String, String>,
) -> AiNetworkPolicy {
    let get = |key: &str| vars.get(key).map(|v| v.as_str()).filter(|v| !v.is_empty());

    match agent {
        // Codex sandboxed commands with network access turned off
        // https://developers.openai.com/codex/environment-variables
        AiAgent::Codex
            if get("CODEX_SANDBOX_NETWORK_DISABLED").is_some_and(|v| v == "1" || v == "true") =>
        {
            return AiNetworkPolicy::Disabled;
        }
        // GitHub Copilot cloud agent firewall. The documented `_ENABLED`,
        // `_ALLOW_LIST`, and `_ALLOW_LIST_ADDITIONS` variables are accompanied
        // by undocumented `_RULESET_*` and `_LOG_FILE` internals, so match on
        // the stable prefix rather than exact names.
        // https://docs.github.com/en/copilot/how-tos/copilot-on-github/customize-copilot/customize-cloud-agent/customize-the-agent-firewall
        AiAgent::GithubCopilot
            if vars.iter().any(|(key, value)| {
                key.starts_with("COPILOT_AGENT_FIREWALL_") && !value.is_empty()
            }) =>
        {
            return if get("COPILOT_AGENT_FIREWALL_ENABLED")
                .is_some_and(|v| v == "false" || v == "0")
            {
                AiNetworkPolicy::Open
            } else {
                AiNetworkPolicy::Filtered
            };
        }
        _ => {}
    }

    // A proxy that is only reachable from inside the machine or its local
    // network implies egress is being filtered through it, as with the
    // localhost proxy of the Claude Code sandbox. This is a property of the
    // process itself rather than of any one agent, so it applies to all of
    // them. External proxies are not classified: they are just as likely
    // plain corporate infrastructure.
    for key in [
        "HTTPS_PROXY",
        "https_proxy",
        "HTTP_PROXY",
        "http_proxy",
        "ALL_PROXY",
        "all_proxy",
    ] {
        if get(key)
            .and_then(proxy_host)
            .is_some_and(is_internal_proxy_host)
        {
            return AiNetworkPolicy::Filtered;
        }
    }

    AiNetworkPolicy::Unknown
}

/// Extracts the host from a proxy URL such as `http://127.0.0.1:3129`,
/// `socks5://user:pass@proxy:1080`, or a bare `host:port` pair.
fn proxy_host(value: &str) -> Option<&str> {
    let rest = value.split_once("://").map_or(value, |(_, rest)| rest);
    let rest = rest.rsplit_once('@').map_or(rest, |(_, rest)| rest);
    let rest = match rest.find(['/', '?']) {
        Some(index) => &rest[..index],
        None => rest,
    };

    // Bracketed IPv6, e.g. `[::1]:8080`
    if let Some(v6) = rest.strip_prefix('[') {
        return v6.split(']').next();
    }

    // Bare IPv6 without a port, e.g. `::1`
    if rest.parse::<IpAddr>().is_ok() {
        return Some(rest);
    }

    let host = rest.rsplit_once(':').map_or(rest, |(host, _)| host);

    (!host.is_empty() && !host.contains(':')).then_some(host)
}

/// Whether a proxy host can only be reached from inside the machine or its
/// local network: loopback and private addresses, `localhost`, and
/// single-label hostnames (e.g. `proxy`), which only resolve locally.
fn is_internal_proxy_host(host: &str) -> bool {
    if host.eq_ignore_ascii_case("localhost") {
        return true;
    }

    if let Ok(ip) = host.parse::<IpAddr>() {
        return match ip {
            IpAddr::V4(v4) => v4.is_loopback() || v4.is_private() || v4.is_link_local(),
            IpAddr::V6(v6) => v6.is_loopback(),
        };
    }

    !host.is_empty()
        && host
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
}

/// Returns metadata about the current AI agent environment, or `None` when no
/// agent is detected.
pub fn get_environment() -> Option<AiEnvironment> {
    let environment = match detect_agent() {
        AiAgent::Antigravity => antigravity::create_environment(),
        AiAgent::Augment => augment::create_environment(),
        AiAgent::Claude => claude::create_environment(),
        AiAgent::ClaudeCowork => claude_cowork::create_environment(),
        AiAgent::Codex => codex::create_environment(),
        AiAgent::Cursor => cursor::create_environment(),
        AiAgent::CursorCli => cursor_cli::create_environment(),
        AiAgent::Devin => devin::create_environment(),
        AiAgent::Gemini => gemini::create_environment(),
        AiAgent::GithubCopilot => github_copilot::create_environment(),
        AiAgent::OpenCode => opencode::create_environment(),
        AiAgent::Replit => replit::create_environment(),
        AiAgent::V0 => v0::create_environment(),
        AiAgent::Unknown => {
            // No typed agent matched. If something still self-identified via
            // `AI_AGENT`, surface it generically; otherwise no agent is present.
            return api::self_id().map(|id| AiEnvironment {
                agent: AiAgent::Unknown,
                env_prefix: None,
                id: Some(id),
                network: detect_network_policy(AiAgent::Unknown),
                sandboxed: false,
                session_id: None,
            });
        }
    };

    Some(environment)
}

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

    fn vars(list: &[(&str, &str)]) -> HashMap<String, String> {
        list.iter()
            .map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
            .collect()
    }

    #[test]
    fn detects_claude_both_spellings() {
        assert_eq!(
            detect_agent_from_vars(&vars(&[("CLAUDECODE", "1")])),
            AiAgent::Claude
        );
        assert_eq!(
            detect_agent_from_vars(&vars(&[("CLAUDE_CODE", "1")])),
            AiAgent::Claude
        );
    }

    #[test]
    fn cowork_mode_distinguished_from_claude() {
        let env = vars(&[("CLAUDECODE", "1"), ("CLAUDE_CODE_IS_COWORK", "1")]);

        assert_eq!(detect_agent_from_vars(&env), AiAgent::ClaudeCowork);
    }

    #[test]
    fn detects_codex_via_sandbox_or_thread() {
        assert_eq!(
            detect_agent_from_vars(&vars(&[("CODEX_SANDBOX", "seatbelt")])),
            AiAgent::Codex
        );
        // CODEX_THREAD_ID marks Codex even outside a sandbox
        assert_eq!(
            detect_agent_from_vars(&vars(&[("CODEX_THREAD_ID", "th_123")])),
            AiAgent::Codex
        );
    }

    #[test]
    fn detects_cursor_editor_and_cli() {
        assert_eq!(
            detect_agent_from_vars(&vars(&[("CURSOR_TRACE_ID", "abc")])),
            AiAgent::Cursor
        );
        assert_eq!(
            detect_agent_from_vars(&vars(&[("CURSOR_AGENT", "1")])),
            AiAgent::CursorCli
        );
    }

    #[test]
    fn cursor_cli_via_extension_host_role() {
        // Only the `agent-exec` role counts, not the editor's other roles
        assert_eq!(
            detect_agent_from_vars(&vars(&[("CURSOR_EXTENSION_HOST_ROLE", "agent-exec")])),
            AiAgent::CursorCli
        );
        assert_eq!(
            detect_agent_from_vars(&vars(&[("CURSOR_EXTENSION_HOST_ROLE", "extension")])),
            AiAgent::Unknown
        );
    }

    #[test]
    fn detects_gemini() {
        let env = vars(&[("GEMINI_CLI", "1")]);

        assert_eq!(detect_agent_from_vars(&env), AiAgent::Gemini);
    }

    #[test]
    fn detects_antigravity() {
        let env = vars(&[("ANTIGRAVITY_AGENT", "1")]);

        assert_eq!(detect_agent_from_vars(&env), AiAgent::Antigravity);
    }

    #[test]
    fn detects_github_copilot() {
        assert_eq!(
            detect_agent_from_vars(&vars(&[("COPILOT_MODEL", "gpt-5")])),
            AiAgent::GithubCopilot
        );
        assert_eq!(
            detect_agent_from_vars(&vars(&[("COPILOT_GITHUB_TOKEN", "x")])),
            AiAgent::GithubCopilot
        );
    }

    #[test]
    fn detects_opencode_either_marker() {
        assert_eq!(
            detect_agent_from_vars(&vars(&[("OPENCODE", "1")])),
            AiAgent::OpenCode
        );
        assert_eq!(
            detect_agent_from_vars(&vars(&[("OPENCODE_CLIENT", "x")])),
            AiAgent::OpenCode
        );
    }

    #[test]
    fn classifies_self_id_when_no_marker() {
        assert_eq!(
            detect_agent_from_vars(&vars(&[("AI_AGENT", "claude-code_2-1-170_agent")])),
            AiAgent::Claude
        );
        // v0 only ever identifies through AI_AGENT
        assert_eq!(
            detect_agent_from_vars(&vars(&[("AI_AGENT", "v0")])),
            AiAgent::V0
        );
    }

    #[test]
    fn boolean_marker_and_self_id_agree() {
        // The real Claude Code environment sets both
        let env = vars(&[
            ("CLAUDECODE", "1"),
            ("AI_AGENT", "claude-code_2-1-170_agent"),
        ]);

        assert_eq!(detect_agent_from_vars(&env), AiAgent::Claude);
    }

    #[test]
    fn cursor_editor_outranks_claude() {
        let env = vars(&[("CURSOR_TRACE_ID", "abc"), ("CLAUDECODE", "1")]);

        assert_eq!(detect_agent_from_vars(&env), AiAgent::Cursor);
    }

    #[test]
    fn unrecognized_self_id_is_unknown() {
        // Still "present" — surfaced generically by get_environment — but not
        // a typed agent
        let env = vars(&[("AI_AGENT", "some-new-tool")]);

        assert_eq!(detect_agent_from_vars(&env), AiAgent::Unknown);
    }

    #[test]
    fn empty_values_are_ignored() {
        let env = vars(&[("CLAUDECODE", "")]);

        assert_eq!(detect_agent_from_vars(&env), AiAgent::Unknown);
    }

    #[test]
    fn no_agent() {
        let env = vars(&[("PATH", "/usr/bin")]);

        assert_eq!(detect_agent_from_vars(&env), AiAgent::Unknown);
    }

    #[test]
    fn network_disabled_for_codex_sandbox() {
        assert_eq!(
            detect_network_policy_from_vars(
                AiAgent::Codex,
                &vars(&[("CODEX_SANDBOX_NETWORK_DISABLED", "1")])
            ),
            AiNetworkPolicy::Disabled
        );
        // Only an affirmative value counts
        assert_eq!(
            detect_network_policy_from_vars(
                AiAgent::Codex,
                &vars(&[("CODEX_SANDBOX_NETWORK_DISABLED", "0")])
            ),
            AiNetworkPolicy::Unknown
        );
    }

    #[test]
    fn network_filtered_for_copilot_firewall() {
        // Undocumented runtime internals match via the prefix
        assert_eq!(
            detect_network_policy_from_vars(
                AiAgent::GithubCopilot,
                &vars(&[("COPILOT_AGENT_FIREWALL_ENABLE_RULESET_ALLOW_LIST", "true")])
            ),
            AiNetworkPolicy::Filtered
        );
        // As do the documented allowlist variables
        assert_eq!(
            detect_network_policy_from_vars(
                AiAgent::GithubCopilot,
                &vars(&[("COPILOT_AGENT_FIREWALL_ALLOW_LIST_ADDITIONS", "example.com")])
            ),
            AiNetworkPolicy::Filtered
        );
    }

    #[test]
    fn network_open_when_copilot_firewall_disabled() {
        // The explicit disable wins over other firewall variables
        let env = vars(&[
            ("COPILOT_AGENT_FIREWALL_ENABLED", "false"),
            ("COPILOT_AGENT_FIREWALL_LOG_FILE", "/tmp/fw.jsonl"),
        ]);

        assert_eq!(
            detect_network_policy_from_vars(AiAgent::GithubCopilot, &env),
            AiNetworkPolicy::Open
        );
    }

    #[test]
    fn network_unknown_for_copilot_cli_alone() {
        // The local Copilot CLI has no firewall
        let env = vars(&[("COPILOT_CLI", "1")]);

        assert_eq!(
            detect_network_policy_from_vars(AiAgent::GithubCopilot, &env),
            AiNetworkPolicy::Unknown
        );
    }

    #[test]
    fn network_vendor_signals_only_apply_to_their_own_agent() {
        // A leaked Copilot firewall variable must not classify another agent
        let copilot = vars(&[("COPILOT_AGENT_FIREWALL_ENABLE_RULESET_ALLOW_LIST", "true")]);

        assert_eq!(
            detect_network_policy_from_vars(AiAgent::Cursor, &copilot),
            AiNetworkPolicy::Unknown
        );

        // And likewise for the Codex network toggle
        let codex = vars(&[("CODEX_SANDBOX_NETWORK_DISABLED", "1")]);

        assert_eq!(
            detect_network_policy_from_vars(AiAgent::Claude, &codex),
            AiNetworkPolicy::Unknown
        );
    }

    #[test]
    fn network_filtered_for_internal_proxies() {
        for value in [
            "http://127.0.0.1:3129",
            "http://localhost:8080",
            "http://[::1]:8080",
            "socks5://user:pass@10.0.0.2:1080",
            // Single-label hosts only resolve on the local network
            "http://proxy:8080",
        ] {
            assert_eq!(
                detect_network_policy_from_vars(AiAgent::Claude, &vars(&[("HTTPS_PROXY", value)])),
                AiNetworkPolicy::Filtered,
                "value {value:?}"
            );
        }

        // Lowercase variants count too, and the proxy signal is not
        // vendor-specific, so it applies to every agent
        assert_eq!(
            detect_network_policy_from_vars(
                AiAgent::Gemini,
                &vars(&[("http_proxy", "http://127.0.0.1:3129")])
            ),
            AiNetworkPolicy::Filtered
        );
    }

    #[test]
    fn network_unknown_for_external_proxies() {
        // A fully-qualified proxy is just as likely corporate infrastructure
        for value in [
            "http://proxy.corp.example.com:8080",
            "https://8.8.8.8:443",
            "",
        ] {
            assert_eq!(
                detect_network_policy_from_vars(AiAgent::Claude, &vars(&[("HTTPS_PROXY", value)])),
                AiNetworkPolicy::Unknown,
                "value {value:?}"
            );
        }
    }

    #[test]
    fn network_disabled_outranks_other_signals() {
        let env = vars(&[
            ("CODEX_SANDBOX_NETWORK_DISABLED", "1"),
            ("HTTPS_PROXY", "http://127.0.0.1:3129"),
        ]);

        assert_eq!(
            detect_network_policy_from_vars(AiAgent::Codex, &env),
            AiNetworkPolicy::Disabled
        );
    }

    #[test]
    fn no_network_signals() {
        let env = vars(&[("PATH", "/usr/bin")]);

        assert_eq!(
            detect_network_policy_from_vars(AiAgent::Claude, &env),
            AiNetworkPolicy::Unknown
        );
    }
}