clawmates/crates/cm-api/src/microvm_client.rs

//! Drive a fleet node's microVMs from the server.
//!
//! Thin by design: the node owns the VM lifecycle (see
//! `clawmates-node::microvm`), and this is the typed way to ask it. Every call
//! is one `vm_*` op over the existing `NodeHub` request/response channel, so
//! there is no new transport, correlation or timeout machinery.
//!
//! # Not a `SandboxDriver`
//!
//! `RemoteDriver` exists to marshal `SandboxDriver` over the hub, and reusing it
//! was the plan. That trait is container-shaped — `attach_pty`, `resize_pty`,
//! argv `exec` — while a mission needs inject → run → collect. Conforming would
//! mean implementing PTY-over-vsock semantics that nothing calls, so this speaks
//! the smaller interface the mission path actually uses.
//!
//! # Timeouts
//!
//! The hub defaults to 20s, which is right for a create (measured: ~1s) and
//! badly wrong for an agent turn. `exec` therefore takes its own budget and
//! passes it to BOTH the hub and the guest, with the hub's slightly longer: if
//! the guest's own timeout fires first the reply says so, whereas a hub timeout
//! leaves us guessing whether the command is still running.

use cm_domain::NodeId;
use serde_json::{json, Value};

use crate::fleet::NodeHub;

/// Slack between the guest's deadline and the hub's, so the guest's own timeout
/// wins the race and we get a real answer rather than a transport error.
const HUB_GRACE_SECS: u64 = 30;

/// How long the hub waits for a command whose own budget is `guest_secs`.
///
/// Saturating, not `+`: a caller passing a very large budget would otherwise
/// overflow and panic in debug or wrap to a tiny timeout in release — the second
/// being far worse, since it turns a long-running agent turn into a spurious
/// transport failure.
fn hub_deadline(guest_secs: u64) -> u64 {
    guest_secs.saturating_add(HUB_GRACE_SECS)
}

pub struct MicroVm<'a> {
    hub: &'a NodeHub,
    node_id: NodeId,
    vm_id: String,
}

impl<'a> MicroVm<'a> {
    pub fn new(hub: &'a NodeHub, node_id: NodeId, vm_id: impl Into<String>) -> Self {
        Self {
            hub,
            node_id,
            vm_id: vm_id.into(),
        }
    }

    pub fn vm_id(&self) -> &str {
        &self.vm_id
    }

    /// One op, with the node's `output` string parsed back into JSON.
    ///
    /// `output` is a String on the wire (`Uplink::Result`), and a node that
    /// answered with a JSON object instead made the whole frame unparseable —
    /// the reply then vanished into the uplink's error arm and the call timed
    /// out with nothing explaining why. Parsing here, loudly, keeps that
    /// mismatch a visible error rather than a mystery timeout.
    async fn call(&self, op: &str, mut args: Value, secs: u64) -> Result<Value, String> {
        if let Some(o) = args.as_object_mut() {
            o.insert("vm_id".into(), Value::String(self.vm_id.clone()));
        }
        let out = self
            .hub
            .call_timeout(self.node_id, op, args, secs)
            .await
            .map_err(|e| format!("{op} on node {:?}: {e}", self.node_id))?;
        let body: Value = serde_json::from_str(&out.output)
            .map_err(|e| format!("{op} returned unparseable output ({e}): {}", out.output))?;
        if !out.ok {
            let why = body
                .get("error")
                .and_then(Value::as_str)
                .unwrap_or(&out.output);
            return Err(format!("{op} failed: {why}"));
        }
        Ok(body)
    }

    /// Boot the VM. Returns only once its guest agent has answered.
    ///
    /// `backend` selects the rootfs image (`missions.backend`); `None` boots the
    /// node's default. A backend whose image is not built on that node is an
    /// error naming the file — never a quiet fall back to the default, which
    /// would run a claude mission in a kimi VM and report success.
    pub async fn create(
        &self,
        vcpus: u32,
        mem_mib: u32,
        backend: Option<&str>,
    ) -> Result<Value, String> {
        // 60s, not the hub default: a create that has to copy a rootfs and boot
        // is measured near 1s, but a node under load has no reason to be fast.
        self.call(
            "vm_create",
            json!({ "vcpus": vcpus, "mem_mib": mem_mib, "backend": backend }),
            60,
        )
        .await
    }

    /// Unpack a tar inside the guest at `dest`.
    ///
    /// Takes the archive bytes rather than a path: the server holds the mission
    /// checkout, the node does not, and shipping the tar is the whole point of
    /// the inject → run → collect model.
    pub async fn inject(&self, dest: &str, tar: &[u8]) -> Result<Value, String> {
        use base64::Engine as _;
        let b64 = base64::engine::general_purpose::STANDARD.encode(tar);
        self.call("vm_inject", json!({ "dest": dest, "tar_b64": b64 }), 120)
            .await
    }

    /// Run a shell command in the guest.
    ///
    /// `Ok` means the command RAN; the exit code is in the payload. A non-zero
    /// exit is not an error here — the caller has to be able to tell "the build
    /// failed" from "we could not reach the VM", and collapsing them is the
    /// defect this codebase keeps paying for.
    /// `env` carries the provider credentials (see
    /// [`crate::mission_runtime::forwarded_provider_env`]). It is sent, never
    /// logged: this is the only channel by which a secret reaches the guest, and
    /// the guest refuses the exec rather than running a command without an entry
    /// it could not honour.
    pub async fn exec(
        &self,
        cmd: &str,
        cwd: Option<&str>,
        timeout_secs: u64,
        env: &[(String, String)],
    ) -> Result<ExecOut, String> {
        self.exec_attributed(cmd, cwd, timeout_secs, env, None, None)
            .await
    }

    /// The same exec, tagged with the run whose live output this is.
    ///
    /// When `run_id` is set the node follows `log_path` inside the guest for the
    /// life of the command and streams what it reads to the server. Probes pass
    /// `None`: they produce nothing worth streaming and have no subscriber.
    pub async fn exec_attributed(
        &self,
        cmd: &str,
        cwd: Option<&str>,
        timeout_secs: u64,
        env: &[(String, String)],
        run_id: Option<uuid::Uuid>,
        log_path: Option<&str>,
    ) -> Result<ExecOut, String> {
        let env: Option<Value> = (!env.is_empty()).then(|| {
            env.iter()
                .map(|(k, v)| (k.clone(), Value::String(v.clone())))
                .collect::<serde_json::Map<_, _>>()
                .into()
        });
        let v = self
            .call(
                "vm_exec",
                json!({
                    "cmd": cmd, "cwd": cwd, "timeout": timeout_secs, "env": env,
                    "run_id": run_id.map(|r| r.to_string()), "log_path": log_path,
                }),
                hub_deadline(timeout_secs),
            )
            .await?;
        // A guest that refused to run the command reports `ok: false` and no rc
        // — a rejected env entry, for instance. Surface its reason: falling
        // through to the missing-rc error below would hide the cause behind a
        // symptom.
        if v.get("ok").and_then(Value::as_bool) == Some(false) {
            return Err(format!(
                "vm_exec did not run: {}",
                v.get("error").and_then(Value::as_str).unwrap_or("unknown")
            ));
        }
        // A missing rc is not "success" — it means the guest did not report one,
        // which we must not read as zero.
        let rc = v
            .get("rc")
            .and_then(Value::as_i64)
            .ok_or_else(|| format!("vm_exec gave no exit code: {v}"))?;
        Ok(ExecOut {
            rc,
            stdout: v
                .get("stdout")
                .and_then(Value::as_str)
                .unwrap_or_default()
                .to_string(),
            stderr: v
                .get("stderr")
                .and_then(Value::as_str)
                .unwrap_or_default()
                .to_string(),
        })
    }

    /// Tar a path out of the guest and return the archive bytes.
    /// `exclude` names directories to leave out — build output, caches. Sent from
    /// here so the policy lives in one place: `mission_fs::transport_excludes`,
    /// the same list the delivery diff uses. Shipping `target/` blew this call's
    /// 300s budget twice, each time with the agent's work finished and stranded.
    pub async fn collect(&self, path: &str, exclude: &[&str]) -> Result<Vec<u8>, String> {
        use base64::Engine as _;
        let v = self
            .call("vm_collect", json!({ "path": path, "exclude": exclude }), 300)
            .await?;
        // The guest reports its own `ok`: a missing path is a real failure that
        // must not come back as an empty archive, which would look exactly like
        // a run that produced nothing.
        if v.get("ok").and_then(Value::as_bool) != Some(true) {
            return Err(format!(
                "vm_collect {path}: {}",
                v.get("error").and_then(Value::as_str).unwrap_or("unknown")
            ));
        }
        let b64 = v
            .get("tar_b64")
            .and_then(Value::as_str)
            .ok_or_else(|| format!("vm_collect {path} returned no archive: {v}"))?;
        base64::engine::general_purpose::STANDARD
            .decode(b64)
            .map_err(|e| format!("vm_collect {path}: undecodable archive: {e}"))
    }

    /// Stop the VM and remove everything it owned. Idempotent.
    pub async fn destroy(&self) -> Result<Value, String> {
        self.call("vm_destroy", json!({}), 60).await
    }
}

/// The result of a command that RAN. `rc != 0` is a normal outcome.
#[derive(Debug, Clone)]
pub struct ExecOut {
    pub rc: i64,
    pub stdout: String,
    pub stderr: String,
}

impl ExecOut {
    pub fn ok(&self) -> bool {
        self.rc == 0
    }
    /// One line for a log or an artifact, without dumping a whole build.
    pub fn summary(&self) -> String {
        let tail = |s: &str| {
            s.lines()
                .rev()
                .take(3)
                .collect::<Vec<_>>()
                .into_iter()
                .rev()
                .collect::<Vec<_>>()
                .join(" | ")
        };
        if self.ok() {
            format!("rc=0 {}", tail(&self.stdout))
        } else {
            format!("rc={} {}", self.rc, tail(&self.stderr))
        }
    }
}

/// VMs a node currently holds, so orphans can be reaped.
pub async fn list(hub: &NodeHub, node_id: NodeId) -> Result<Vec<String>, String> {
    let out = hub
        .call(node_id, "vm_list", json!({}))
        .await
        .map_err(|e| format!("vm_list on node {node_id:?}: {e}"))?;
    let body: Value = serde_json::from_str(&out.output)
        .map_err(|e| format!("vm_list returned unparseable output ({e}): {}", out.output))?;
    Ok(body
        .get("vms")
        .and_then(Value::as_array)
        .map(|a| {
            a.iter()
                .filter_map(|v| v.get("vm_id").and_then(Value::as_str))
                .map(str::to_string)
                .collect()
        })
        .unwrap_or_default())
}

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

    /// A command that ran and failed must be distinguishable from one that
    /// could not be reached. `rc` carries the verdict; `Err` is for transport.
    #[test]
    fn a_nonzero_exit_is_an_outcome_not_an_error() {
        let failed = ExecOut {
            rc: 3,
            stdout: String::new(),
            stderr: "boom\n".into(),
        };
        assert!(!failed.ok());
        assert!(failed.summary().starts_with("rc=3"));
        assert!(failed.summary().contains("boom"));

        let passed = ExecOut {
            rc: 0,
            stdout: "fine\n".into(),
            stderr: String::new(),
        };
        assert!(passed.ok());
        assert_eq!(passed.summary(), "rc=0 fine");
    }

    /// The summary is for logs, so it must stay short even when a build prints
    /// thousands of lines — and it must keep the LAST lines, where the error is.
    #[test]
    fn the_summary_keeps_the_tail_and_stays_short() {
        let noisy = ExecOut {
            rc: 1,
            stdout: String::new(),
            stderr: (1..=500)
                .map(|i| format!("line {i}"))
                .collect::<Vec<_>>()
                .join("\n"),
        };
        let s = noisy.summary();
        assert!(s.contains("line 500"), "the last line must survive: {s}");
        assert!(!s.contains("line 400"), "older lines must be dropped: {s}");
        assert!(s.len() < 200, "summary must stay log-sized, got {}", s.len());
    }

    /// The guest's deadline must fire before the hub's, so a slow command comes
    /// back as a reported timeout rather than an unexplained transport failure.
    #[test]
    fn the_hub_always_outlives_the_guests_own_timeout() {
        for guest in [0u64, 1, 30, 3600, 86_400] {
            assert!(
                hub_deadline(guest) > guest,
                "hub deadline for {guest}s must exceed it"
            );
        }
        // A caller passing a huge budget must not wrap to a tiny timeout, which
        // would turn a long agent turn into a spurious transport failure.
        assert!(
            hub_deadline(u64::MAX) >= u64::MAX - 1,
            "an extreme budget must saturate, not wrap"
        );
    }
}