hatchery_node/

m4a_accounts.rs

//! Per-session m4a accounts assigned by the node at spawn (owner model: session = nick = one key
//! = one role). When `HATCHERY_M4A_AGENT_BIN` and `HATCHERY_M4A_ACCOUNTS_ROOT` are set, every PTY
//! spawn gets its own account: the node runs `m4a-agent account provision` (registry and vaults
//! stay under the node's accounts root, outside the session sandbox), reads the fresh capability
//! from a one-shot 0600 file it removes at once, and injects ONLY `M4A_NICK` and
//! `M4A_SESSION_CAP` into that session's child environment. The capability is never logged.

use std::ffi::OsString;
use std::path::{Path, PathBuf};
use std::process::Command;

use gate4agent_catalog::EnvMutation;

/// g4a provider id -> the provider name `m4a-agent` accepts. g4a's built-in catalog is
/// `claude`, `codex`, `grok`, `kimi`; anything else (custom agents) gets no account.
pub fn m4a_provider(g4a_agent_id: &str) -> Option<&'static str> {
    match g4a_agent_id {
        "claude" => Some("claude-code"),
        "codex" => Some("codex"),
        "grok" => Some("grok"),
        "kimi" => Some("kimi-code"),
        _ => None,
    }
}

pub const AGENT_BIN_ENV: &str = "HATCHERY_M4A_AGENT_BIN";
pub const ACCOUNTS_ROOT_ENV: &str = "HATCHERY_M4A_ACCOUNTS_ROOT";

#[derive(Debug, Clone)]
pub struct M4aSessionAccounts {
    bin: PathBuf,
    base: PathBuf,
}

/// One session's account as handed to its child: nick and capability only.
pub struct SessionAccount {
    pub nick: String,
    /// The provider session key the account is bound to (also given to the session as
    /// `M4A_SESSION_ID`, so its own inbox commands address the right account).
    pub session_id: String,
    cap: String,
}

impl Drop for SessionAccount {
    fn drop(&mut self) {
        // Best effort wipe of the capability copy held by the node.
        let bytes = unsafe { self.cap.as_bytes_mut() };
        bytes.iter_mut().for_each(|b| *b = 0);
    }
}

impl std::fmt::Debug for SessionAccount {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SessionAccount").field("nick", &self.nick).finish_non_exhaustive()
    }
}

impl SessionAccount {
    /// The only variables the session gets from the node for m4a.
    /// The one-time `/m4a/attach` reply (the capability travels only in this body).
    pub fn attach_reply(&self) -> serde_json::Value {
        serde_json::json!({"nick": self.nick, "session_id": self.session_id, "cap": self.cap})
    }

    pub fn env_mutations(&self) -> Vec<EnvMutation> {
        vec![
            EnvMutation { key: OsString::from("M4A_NICK"), value: Some(OsString::from(&self.nick)) },
            EnvMutation { key: OsString::from("M4A_SESSION_CAP"), value: Some(OsString::from(&self.cap)) },
            EnvMutation { key: OsString::from("M4A_SESSION_ID"), value: Some(OsString::from(&self.session_id)) },
        ]
    }
}

impl M4aSessionAccounts {
    /// Fallback identity channel: `<root>/sessions/<nick>.session` (owner-only), the same three
    /// keys as [`SessionAccount::env_mutations`], for a provider whose tools run in a process
    /// that did not inherit the session environment (`m4a-agent --session-file <path> ...`).
    pub fn session_file_path(&self, nick: &str) -> Option<PathBuf> {
        let ok = !nick.is_empty() && nick.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_');
        ok.then(|| self.base.join("sessions").join(format!("{nick}.session")))
    }

    pub fn write_session_file(&self, account: &SessionAccount) -> Result<PathBuf, String> {
        let path = self.session_file_path(&account.nick).ok_or("invalid nick")?;
        let dir = path.parent().ok_or("no dir")?;
        std::fs::create_dir_all(dir).map_err(|e| format!("sessions dir: {}", e.kind()))?;
        let mut body = String::new();
        for m in account.env_mutations() {
            if let Some(v) = m.value {
                body.push_str(&format!("{}={}\n", m.key.to_string_lossy(), v.to_string_lossy()));
            }
        }
        let mut opts = std::fs::OpenOptions::new();
        opts.write(true).create(true).truncate(true);
        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt;
            opts.mode(0o600);
        }
        use std::io::Write;
        opts.open(&path)
            .and_then(|mut f| f.write_all(body.as_bytes()))
            .map_err(|e| format!("session file: {}", e.kind()))?;
        body.zeroize_best_effort();
        Ok(path)
    }

    /// Provider flags that give one spawned session its own m4a MCP server (`m4a-agent
    /// --session-file <file> mcp`), without touching any user config. Empty for providers with
    /// no verified per-session MCP flag (grok TUI, kimi PTY): those use attach (`/m4a/attach`).
    pub fn mcp_launch_args(&self, g4a_provider: &str, nick: &str) -> Result<Vec<OsString>, String> {
        let session_file = self.session_file_path(nick).ok_or("invalid nick")?;
        let bin = self.bin.to_string_lossy().into_owned();
        let sf = session_file.to_string_lossy().into_owned();
        let provider = g4a_provider.to_owned();
        match g4a_provider {
            "claude" => {
                let config = serde_json::json!({"mcpServers": {"m4a": {
                    "type": "stdio",
                    "command": bin,
                    "args": ["--session-file", sf, "mcp", "--provider", provider],
                }}});
                let path = self.base.join("sessions").join(format!("{nick}.mcp.json"));
                std::fs::create_dir_all(self.base.join("sessions")).map_err(|e| format!("sessions dir: {}", e.kind()))?;
                std::fs::write(&path, config.to_string()).map_err(|e| format!("mcp config: {}", e.kind()))?;
                Ok(vec![OsString::from("--mcp-config"), path.into_os_string()])
            }
            "codex" => {
                // TOML literal strings: no escapes, no double quotes (Windows wrapper safe).
                if bin.contains('\'') || sf.contains('\'') {
                    return Err("path not representable as a TOML literal string".into());
                }
                Ok(vec![
                    OsString::from("-c"),
                    OsString::from(format!("mcp_servers.m4a.command='{bin}'")),
                    OsString::from("-c"),
                    OsString::from(format!("mcp_servers.m4a.args=['--session-file','{sf}','mcp','--provider','codex']")),
                ])
            }
            _ => Ok(Vec::new()),
        }
    }

    pub fn remove_session_file(&self, nick: &str) {
        if let Some(path) = self.session_file_path(nick) {
            let _ = std::fs::remove_file(path);
        }
        let _ = std::fs::remove_file(self.base.join("sessions").join(format!("{nick}.mcp.json")));
    }

    pub fn new(bin: impl Into<PathBuf>, base: impl Into<PathBuf>) -> Self {
        Self { bin: bin.into(), base: base.into() }
    }

    /// The node's letter-wake token (`<root>/.node-wake-token`, 0600), created on first use.
    pub fn wake_token(&self) -> Result<String, String> {
        let path = self.base.join(".node-wake-token");
        if let Ok(t) = std::fs::read_to_string(&path) {
            if t.trim().len() >= 32 {
                return Ok(t.trim().to_owned());
            }
        }
        std::fs::create_dir_all(&self.base).map_err(|e| format!("accounts root: {}", e.kind()))?;
        let mut raw = [0u8; 24];
        getrandom(&mut raw)?;
        let token: String = raw.iter().map(|b| format!("{b:02x}")).collect();
        let mut opts = std::fs::OpenOptions::new();
        opts.write(true).create(true).truncate(true);
        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt;
            opts.mode(0o600);
        }
        use std::io::Write;
        opts.open(&path).and_then(|mut f| f.write_all(token.as_bytes())).map_err(|e| format!("wake token: {}", e.kind()))?;
        Ok(token)
    }

    pub fn from_env() -> Option<Self> {
        let get = |k: &str| std::env::var_os(k).filter(|v| !v.is_empty()).map(PathBuf::from);
        Some(Self::new(get(AGENT_BIN_ENV)?, get(ACCOUNTS_ROOT_ENV)?))
    }

    pub fn base(&self) -> &Path {
        &self.base
    }

    /// The account of `(provider, session_key)`, created if missing, with a fresh capability
    /// (a respawn of the same session invalidates the previous process's key).
    pub fn provision(&self, provider: &str, session_key: &str) -> Result<SessionAccount, String> {
        std::fs::create_dir_all(&self.base).map_err(|e| format!("accounts root: {}", e.kind()))?;
        let mut rand = [0u8; 12];
        getrandom(&mut rand)?;
        let tag: String = rand.iter().map(|b| format!("{b:02x}")).collect();
        let cap_file = self.base.join(format!(".spawn-{tag}.cap"));
        let out = Command::new(&self.bin)
            .args(["account", "provision", "--provider", provider, "--session", session_key, "--rotate", "--cap-out"])
            .arg(&cap_file)
            .env("M4A_ACCOUNTS_ROOT", &self.base)
            .env_remove("M4A_SESSION_CAP")
            .env_remove("M4A_PRODUCT_INVITE")
            .output()
            .map_err(|e| format!("m4a-agent: {}", e.kind()))?;
        let cap = std::fs::read_to_string(&cap_file);
        let _ = std::fs::remove_file(&cap_file);
        if !out.status.success() {
            return Err(format!("m4a-agent account provision failed ({})", out.status.code().unwrap_or(-1)));
        }
        let v: serde_json::Value = serde_json::from_slice(&out.stdout).map_err(|_| "m4a-agent answer is not json".to_string())?;
        let nick = v["nick"].as_str().ok_or("m4a-agent answer has no nick")?.to_string();
        let cap = cap.map_err(|_| "no capability issued".to_string())?;
        if cap.trim().is_empty() || cap.contains('\n') {
            return Err("capability is malformed".into());
        }
        Ok(SessionAccount { nick, session_id: session_key.to_owned(), cap })
    }
}

fn getrandom(buf: &mut [u8]) -> Result<(), String> {
    use std::io::Read;
    std::fs::File::open("/dev/urandom")
        .and_then(|mut f| f.read_exact(buf))
        .or_else(|_| {
            // Non-unix fallback: time + address entropy is enough for a unique temp name.
            let t = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_nanos()).unwrap_or(0);
            for (i, b) in buf.iter_mut().enumerate() {
                *b = (t >> ((i % 16) * 8)) as u8 ^ (i as u8);
            }
            Ok::<(), std::io::Error>(())
        })
        .map_err(|e| e.to_string())
}

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

    fn fake_agent(dir: &Path) -> PathBuf {
        // Stands in for m4a-agent: one nick per --session, the cap written to --cap-out.
        let p = dir.join("fake-m4a-agent");
        std::fs::write(
            &p,
            "#!/bin/sh\nwhile [ $# -gt 0 ]; do case \"$1\" in --session) s=$2; shift;; --cap-out) o=$2; shift;; esac; shift; done\n\
             umask 077; printf 'cap-%s-%s' \"$s\" \"$$\" > \"$o\"; printf '{\"nick\":\"codex-%s\",\"cap_issued\":true}\\n' \"$s\"\n",
        )
        .unwrap();
        use std::os::unix::fs::PermissionsExt;
        std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o755)).unwrap();
        p
    }

    #[test]
    fn each_spawn_gets_its_own_account_and_only_nick_and_cap_are_injected() {
        let dir = std::env::temp_dir().join(format!("hatch-m4a-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let acc = M4aSessionAccounts::new(fake_agent(&dir), dir.join("accounts"));
        let a = acc.provision("codex", "inc-1").unwrap();
        let b = acc.provision("codex", "inc-2").unwrap();
        assert_ne!(a.nick, b.nick);
        let env = a.env_mutations();
        let keys: Vec<_> = env.iter().map(|m| m.key.to_string_lossy().into_owned()).collect();
        assert_eq!(keys, ["M4A_NICK", "M4A_SESSION_CAP", "M4A_SESSION_ID"]);
        assert_eq!(env[2].value.as_deref(), Some(std::ffi::OsStr::new("inc-1")));
        assert_ne!(env[1].value, b.env_mutations()[1].value);
        // The one-shot cap file is gone and the cap is not in Debug output.
        let left: Vec<_> = std::fs::read_dir(acc.base()).unwrap().filter_map(Result::ok).filter(|e| e.file_name().to_string_lossy().ends_with(".cap")).collect();
        assert!(left.is_empty());
        assert!(!format!("{a:?}").contains("cap-"));
    }

    #[test]
    fn every_builtin_g4a_provider_maps_to_an_m4a_provider() {
        for spec in gate4agent_catalog::builtin_specs() {
            let id = spec.id.as_str().to_string();
            assert!(m4a_provider(&id).is_some(), "g4a provider {id} has no m4a mapping");
        }
        assert_eq!(m4a_provider("custom"), None);
    }

    #[cfg(unix)]
    #[test]
    fn wake_token_is_private_stable_and_long() {
        use std::os::unix::fs::PermissionsExt;
        let dir = std::env::temp_dir().join(format!("node-wake-token-{}", std::process::id()));
        let a = M4aSessionAccounts::new("/bin/false", &dir);
        let t1 = a.wake_token().unwrap();
        assert_eq!(t1.len(), 48);
        assert_eq!(a.wake_token().unwrap(), t1);
        let mode = std::fs::metadata(dir.join(".node-wake-token")).unwrap().permissions().mode() & 0o777;
        assert_eq!(mode, 0o600);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn a_failing_agent_injects_nothing() {
        let acc = M4aSessionAccounts::new("/nonexistent/m4a-agent", std::env::temp_dir().join("hatch-m4a-none"));
        assert!(acc.provision("codex", "x").is_err());
    }
}

trait ZeroizeBestEffort {
    fn zeroize_best_effort(&mut self);
}

impl ZeroizeBestEffort for String {
    fn zeroize_best_effort(&mut self) {
        // SAFETY: zero bytes keep the buffer valid UTF-8.
        unsafe { self.as_bytes_mut() }.iter_mut().for_each(|b| *b = 0);
    }
}

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

    #[test]
    fn mcp_args_are_session_scoped_and_carry_no_capability() {
        let root = std::env::temp_dir().join(format!("m4a-mcp-args-{}", std::process::id()));
        let a = M4aSessionAccounts::new("m4a-agent", &root);
        let c = a.mcp_launch_args("claude", "claude-1").unwrap();
        assert_eq!(c[0], OsString::from("--mcp-config"));
        let cfg = std::fs::read_to_string(&c[1]).unwrap();
        assert!(cfg.contains("--session-file") && cfg.contains("claude-1.session"));
        let x = a.mcp_launch_args("codex", "codex-1").unwrap();
        assert_eq!(x.len(), 4);
        assert!(x.iter().all(|v| !v.to_string_lossy().contains('"')));
        assert!(a.mcp_launch_args("grok", "grok-1").unwrap().is_empty());
        a.remove_session_file("claude-1");
        assert!(!std::path::Path::new(&c[1]).exists());
        let _ = std::fs::remove_dir_all(root);
    }

    #[test]
    fn session_file_path_refuses_path_like_nicks() {
        let a = M4aSessionAccounts::new("bin", std::env::temp_dir().join("m4a-sf-test"));
        assert!(a.session_file_path("grok-1").unwrap().ends_with("sessions/grok-1.session"));
        assert!(a.session_file_path("../x").is_none());
        assert!(a.session_file_path("").is_none());
    }
}