supermachine 0.4.3

//! Host-side framing for the in-guest exec agent.
//!
//! Mirrors the wire protocol implemented in
//! `crates/supermachine-guest-agent`. See
//! `docs/design/exec-2026-05-03.md` for the protocol overview.
//!
//! Public surface (re-exported from the crate root):
//!
//! - [`crate::Vm::exec`] / [`crate::Vm::exec_builder`] (in `api.rs`)
//! - [`ExecBuilder`] — chainable spawn config
//! - [`ExecChild`] — handle to a running guest process
//! - [`ExecStdin`] / [`ExecStdout`] / [`ExecStderr`] — std-style stdio
//!
//! Internals:
//!
//! - `spawn` dials the host-side `<vsock_mux>-exec.sock`, sends
//!   the REQUEST frame, then spawns a demux thread that reads
//!   frames and pumps STDOUT/STDERR bytes into pipe pairs (so the
//!   caller's [`std::io::Read`] just works), and posts the EXIT
//!   status onto a channel.

use std::collections::BTreeMap;
use std::fs::File;
use std::io::{self, Read, Write};
use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
use std::os::unix::net::UnixStream;
use std::path::Path;
use std::path::PathBuf;
use std::process::ExitStatus;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use serde::Serialize;

// Wire frame types — must match the agent's constants.
const FRAME_REQUEST: u8 = 0xff;
const FRAME_CONTROL: u8 = 0xfe;
const FRAME_STDIN: u8 = 0;
const FRAME_STDOUT: u8 = 1;
const FRAME_STDERR: u8 = 2;
const FRAME_RESIZE: u8 = 3;
const FRAME_SIGNAL: u8 = 4;
const FRAME_EXIT: u8 = 5;
const FRAME_ERROR: u8 = 6;

/// Configurable spawn for [`crate::Vm::exec`]. Built via
/// [`crate::Vm::exec_builder`] or constructed inline by `Vm::exec`.
///
/// ```no_run
/// # use supermachine::{Image, Vm, VmConfig};
/// let image = Image::from_snapshot("path/to/snapshot")?;
/// let vm = Vm::start(&image, &VmConfig::new())?;
/// let mut child = vm.exec_builder()
///     .argv(["sh", "-c", "echo hi && exit 7"])
///     .env("FOO", "bar")
///     .spawn()?;
/// let status = child.wait()?;
/// assert_eq!(status.code(), Some(7));
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
pub struct ExecBuilder {
    exec_path: PathBuf,
    argv: Vec<String>,
    env: BTreeMap<String, String>,
    cwd: Option<String>,
    tty: bool,
    cols: Option<u16>,
    rows: Option<u16>,
    timeout: Option<Duration>,
    /// Files to atomically write inside the guest before exec.
    /// Eliminates a separate `write_file` round-trip.
    stage_files: Vec<StageFile>,
    /// Additional argvs to run sequentially after `argv` succeeds
    /// (`&&` semantics). Stops on first non-zero exit. Removes the
    /// need for a `sh -c "X && Y"` wrapper and its extra fork.
    chain: Vec<Vec<String>>,
}

#[derive(Clone)]
struct StageFile {
    path: String,
    data: Vec<u8>,
    mode: Option<u32>,
}

impl ExecBuilder {
    /// Construct a builder that dials a specific exec-socket
    /// path. Most embedders should use [`crate::Vm::exec_builder`]
    /// instead; this lower-level constructor is for tooling that
    /// already has a router daemon running and just wants to
    /// dial a known socket (e.g. the `supermachine exec` CLI).
    pub fn new(exec_path: PathBuf) -> Self {
        Self {
            exec_path,
            argv: Vec::new(),
            env: BTreeMap::new(),
            cwd: None,
            tty: false,
            cols: None,
            rows: None,
            timeout: None,
            stage_files: Vec::new(),
            chain: Vec::new(),
        }
    }

    /// The argv to execute in the guest. First element is the
    /// program (resolved via PATH inside the guest).
    pub fn argv<I, S>(mut self, argv: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.argv = argv.into_iter().map(Into::into).collect();
        self
    }

    /// Set an environment variable for the spawned process.
    /// Repeatable. The agent merges these on top of its own
    /// environment, so `PATH`, `HOME`, etc. are inherited unless
    /// you explicitly override them.
    pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.env.insert(key.into(), value.into());
        self
    }

    /// Working directory for the spawned process inside the guest.
    pub fn cwd(mut self, path: impl Into<String>) -> Self {
        self.cwd = Some(path.into());
        self
    }

    /// Allocate a pseudo-terminal. stdin/stdout pass through the
    /// pty master; stderr is empty (the pty merges fd 1 + fd 2).
    /// Use this for interactive shells.
    pub fn tty(mut self, on: bool) -> Self {
        self.tty = on;
        self
    }

    /// Initial terminal size for tty mode. Has no effect without
    /// `.tty(true)`. Use [`ExecChild::resize`] to change later.
    pub fn winsize(mut self, cols: u16, rows: u16) -> Self {
        self.cols = Some(cols);
        self.rows = Some(rows);
        self
    }

    /// Hard wall-clock timeout. If the process is still running
    /// when the deadline fires, the watchdog sends SIGKILL and
    /// the resulting [`ExecOutcome`] has `timed_out = true`.
    /// Only honored by [`ExecBuilder::output`] (the streaming
    /// [`ExecBuilder::spawn`] path leaves cancellation to the
    /// caller, who already has [`ExecChild::signal`]).
    ///
    /// Reliable: `timed_out = true` is always set when the
    /// deadline fires.
    ///
    /// Process group: the agent runs `setpgid(0, 0)` in the
    /// child after fork, then forwards SIGNAL to the entire
    /// process group via `kill(-pid, sig)`. Forking children
    /// (e.g. `sh -c "rustc x.rs && /tmp/x"`) get killed
    /// transitively — no orphaned processes blocking the EXIT
    /// frame. Verified to take ~500 ms end-to-end on a 500 ms
    /// timeout across `sleep`, `sh -c sleep`, busy loops, and
    /// shell pipelines.
    pub fn timeout(mut self, d: Duration) -> Self {
        self.timeout = Some(d);
        self
    }

    /// Stage `bytes` at `path` inside the guest before exec runs.
    /// Atomic (write+rename). Folded into the same vsock RPC as
    /// the exec request — eliminates the separate `Vm::write_file`
    /// round-trip. Repeatable for multiple files.
    ///
    /// Use when your workload starts with a known-small input file
    /// (source code, a test config) — replace
    /// `vm.write_file(p, b)?; vm.exec_builder().argv([…]).output()?`
    /// with `vm.exec_builder().stage_file(p, b).argv([…]).output()?`.
    pub fn stage_file(mut self, path: impl Into<String>, bytes: impl Into<Vec<u8>>) -> Self {
        self.stage_files.push(StageFile {
            path: path.into(),
            data: bytes.into(),
            mode: None,
        });
        self
    }

    /// Like [`Self::stage_file`] but also sets the file mode.
    pub fn stage_file_mode(
        mut self,
        path: impl Into<String>,
        bytes: impl Into<Vec<u8>>,
        mode: u32,
    ) -> Self {
        self.stage_files.push(StageFile {
            path: path.into(),
            data: bytes.into(),
            mode: Some(mode),
        });
        self
    }

    /// Append a second (third, …) argv to run after the previous
    /// one succeeds (`&&` semantics). Stops on first non-zero
    /// exit. Removes the need for a `sh -c "X && Y"` wrapper —
    /// the agent forks each command directly in the guest.
    pub fn chain<I, S>(mut self, argv: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.chain.push(argv.into_iter().map(Into::into).collect());
        self
    }

    /// Dial the agent, send the REQUEST, and return the
    /// [`ExecChild`] handle. Use this when you want to stream
    /// stdin/stdout/stderr yourself; for a one-call "run, drain,
    /// collect" pattern use [`ExecBuilder::output`].
    pub fn spawn(self) -> io::Result<ExecChild> {
        if self.argv.is_empty() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "exec: argv is empty",
            ));
        }
        spawn(self)
    }

    /// Run to completion, collecting stdout + stderr + exit
    /// status into one [`ExecOutcome`]. Blocks until the process
    /// exits (or the configured [`ExecBuilder::timeout`] fires
    /// and the process is killed).
    ///
    /// Mirrors [`std::process::Command::output`]. Use this for
    /// the common "exec a command, look at the result" case;
    /// reach for [`ExecBuilder::spawn`] only when you need to
    /// stream stdio.
    ///
    /// ```no_run
    /// # use std::time::Duration;
    /// # use supermachine::{Image, Vm, VmConfig};
    /// let image = Image::from_snapshot("path/to/snapshot")?;
    /// let vm = Vm::start(&image, &VmConfig::new())?;
    /// let out = vm.exec_builder()
    ///     .argv(["sh", "-c", "echo hi; echo err >&2; exit 7"])
    ///     .timeout(Duration::from_secs(30))
    ///     .output()?;
    /// assert_eq!(out.status.code(), Some(7));
    /// assert_eq!(out.stdout, b"hi\n");
    /// assert_eq!(out.stderr, b"err\n");
    /// assert!(!out.timed_out);
    /// # Ok::<(), Box<dyn std::error::Error>>(())
    /// ```
    pub fn output(self) -> io::Result<ExecOutcome> {
        let timeout = self.timeout;
        let t0 = Instant::now();
        let mut child = self.spawn()?;
        // Close stdin immediately so workloads that block on
        // `read(0)` (cat, anything reading stdin) see EOF and
        // exit. Drop sends an empty STDIN frame to the agent;
        // the agent's input thread observes that as "EOF for
        // current child" without exiting, so SIGNAL frames from
        // our watchdog still reach the agent later.
        drop(child.stdin());
        let mut stdout = child.stdout();
        let mut stderr = child.stderr();
        // Drain both pipes in parallel so a chatty stderr never
        // backpressures stdout (or vice versa) and deadlocks.
        let stdout_handle = stdout.take().map(|mut s| {
            thread::spawn(move || {
                let mut buf = Vec::new();
                let _ = s.read_to_end(&mut buf);
                buf
            })
        });
        let stderr_handle = stderr.take().map(|mut s| {
            thread::spawn(move || {
                let mut buf = Vec::new();
                let _ = s.read_to_end(&mut buf);
                buf
            })
        });
        // Optional watchdog: after `timeout`, send SIGKILL.
        // Cancelled if the process exits cleanly first by closing
        // the cancel channel.
        let (cancel_tx, cancel_rx) = channel::<()>();
        let timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
        let watchdog = if let Some(d) = timeout {
            let timed_out = Arc::clone(&timed_out);
            // Snapshot the socket writer so the watchdog can SIGNAL
            // without owning the ExecChild.
            let sock_w = Arc::clone(&child.sock_w);
            Some(thread::spawn(move || {
                if cancel_rx.recv_timeout(d).is_err() {
                    timed_out.store(true, std::sync::atomic::Ordering::SeqCst);
                    let r = send_frame(&sock_w, FRAME_SIGNAL, &[9u8]);
                    if std::env::var_os("SUPERMACHINE_TIMEOUT_TRACE").is_some() {
                        eprintln!("[exec.timeout] fired at +{:?}, send_frame={:?}", d, r);
                    }
                }
            }))
        } else {
            None
        };
        let exit = child.wait_with_rss();
        // Stop the watchdog (if any).
        let _ = cancel_tx.send(());
        if let Some(h) = watchdog {
            let _ = h.join();
        }
        let stdout_bytes = stdout_handle.map(|h| h.join().unwrap_or_default()).unwrap_or_default();
        let stderr_bytes = stderr_handle.map(|h| h.join().unwrap_or_default()).unwrap_or_default();
        let (status, peak_rss_kib) = exit?;
        Ok(ExecOutcome {
            status,
            stdout: stdout_bytes,
            stderr: stderr_bytes,
            duration: t0.elapsed(),
            peak_rss_kib,
            timed_out: timed_out.load(std::sync::atomic::Ordering::SeqCst),
        })
    }
}

/// Result of [`ExecBuilder::output`]. Mirrors the shape of
/// [`std::process::Output`] with two additions: `duration` (host-
/// side wall-clock) and `timed_out` (true iff the wall-clock
/// timeout fired and the process was SIGKILL'd).
///
/// `peak_rss_kib` is `None` today; a future supermachine-agent
/// rev will populate it via `getrusage(RUSAGE_CHILDREN)`.
#[derive(Debug)]
pub struct ExecOutcome {
    /// Exit status, exactly as [`std::process::Output::status`].
    pub status: ExitStatus,
    /// Captured stdout bytes (verbatim — caller decodes if needed).
    pub stdout: Vec<u8>,
    /// Captured stderr bytes (verbatim).
    pub stderr: Vec<u8>,
    /// Host-side wall-clock duration from [`ExecBuilder::output`]
    /// being called to the process exiting (or being killed).
    pub duration: Duration,
    /// Peak resident set size of the process and its children, in
    /// KiB. `None` until the guest agent gains rusage support.
    pub peak_rss_kib: Option<u64>,
    /// `true` iff the wall-clock timeout from
    /// [`ExecBuilder::timeout`] fired and the process was killed
    /// with SIGKILL.
    pub timed_out: bool,
}

impl ExecOutcome {
    /// `true` when the exit status is success AND the process
    /// wasn't forcibly killed by the timeout.
    pub fn success(&self) -> bool {
        self.status.success() && !self.timed_out
    }
}

/// JSON shape for the REQUEST frame. Fields match the agent's
/// `ExecRequest` struct.
#[derive(Serialize)]
struct RequestPayload<'a> {
    argv: &'a [String],
    env: &'a BTreeMap<String, String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    cwd: Option<&'a str>,
    tty: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    cols: Option<u16>,
    #[serde(skip_serializing_if = "Option::is_none")]
    rows: Option<u16>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    stage_files: Vec<StageFilePayload<'a>>,
    #[serde(skip_serializing_if = "<[_]>::is_empty")]
    chain: &'a [Vec<String>],
}

#[derive(Serialize)]
struct StageFilePayload<'a> {
    path: &'a str,
    data_b64: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    mode: Option<u32>,
}

/// A running guest process. Hold this until you call [`wait`] or
/// drop it; dropping closes the agent connection, which kills the
/// child via the agent's `SIGHUP` propagation.
///
/// [`wait`]: ExecChild::wait
pub struct ExecChild {
    /// Writer half of the unix socket to the agent. Wrapped in a
    /// mutex so STDIN frames don't interleave with RESIZE/SIGNAL
    /// frames coming from other places.
    sock_w: Arc<Mutex<UnixStream>>,
    /// stdout pipe read end (caller side).
    stdout_r: Option<OwnedFd>,
    /// stderr pipe read end (caller side).
    stderr_r: Option<OwnedFd>,
    /// stdin pipe write end (caller side). `None` if already closed
    /// via [`ExecStdin::close`].
    stdin: Option<ExecStdin>,
    /// Demux thread join handle.
    demux: Option<JoinHandle<()>>,
    /// EXIT status receiver, posted by the demux thread.
    exit_rx: Receiver<DemuxExit>,
}

/// Outcome reported by the demux thread.
enum DemuxExit {
    /// Agent sent EXIT frame. `peak_rss_kib` is `None` for older
    /// agents that only send the 4-byte exit_code payload, and
    /// `Some(kib)` for v0.2.x+ agents that include rusage.
    Status { code: u32, peak_rss_kib: Option<u64> },
    /// Agent closed the socket without sending EXIT.
    EofBeforeExit,
    /// Agent sent ERROR frame.
    Error(String),
    /// Demux I/O failed.
    Io(io::Error),
}

impl ExecChild {
    /// Take ownership of the stdin handle. Calling this twice
    /// returns `None` the second time.
    pub fn stdin(&mut self) -> Option<ExecStdin> {
        self.stdin.take()
    }

    /// Take ownership of stdout. Returns a `Read`er. Useful when
    /// you want to spawn your own pump thread.
    pub fn stdout(&mut self) -> Option<ExecStdout> {
        self.stdout_r.take().map(ExecStdout::from_fd)
    }

    /// Take ownership of stderr. Empty in tty mode.
    pub fn stderr(&mut self) -> Option<ExecStderr> {
        self.stderr_r.take().map(ExecStderr::from_fd)
    }

    /// Send a signal (SIGTERM/SIGINT/...) to the guest process.
    pub fn signal(&self, signum: i32) -> io::Result<()> {
        let payload = [signum as u8];
        send_frame(&self.sock_w, FRAME_SIGNAL, &payload)
    }

    /// Resize the pty (tty mode only). No-op in pipe mode (the
    /// agent ignores RESIZE frames if no pty is attached).
    pub fn resize(&self, cols: u16, rows: u16) -> io::Result<()> {
        let mut payload = [0u8; 4];
        payload[0..2].copy_from_slice(&cols.to_be_bytes());
        payload[2..4].copy_from_slice(&rows.to_be_bytes());
        send_frame(&self.sock_w, FRAME_RESIZE, &payload)
    }

    /// Block until the guest process exits. Returns the exit
    /// status (or a synthesized one if the agent disconnected).
    pub fn wait(self) -> io::Result<ExitStatus> {
        Ok(self.wait_with_rss()?.0)
    }

    /// Like [`ExecChild::wait`] but also surfaces the peak RSS
    /// the agent reports (`Some(kib)` for v0.2.x+ agents,
    /// `None` for older ones). Used internally by
    /// [`ExecBuilder::output`] to populate
    /// [`ExecOutcome::peak_rss_kib`].
    pub fn wait_with_rss(mut self) -> io::Result<(ExitStatus, Option<u64>)> {
        // Drop stdin if the caller didn't already; the child needs
        // EOF to exit cleanly for any program that reads stdin.
        self.stdin.take();

        let exit = self
            .exit_rx
            .recv()
            .map_err(|_| io::Error::new(io::ErrorKind::Other, "exec: demux thread died"))?;
        if let Some(h) = self.demux.take() {
            let _ = h.join();
        }
        match exit {
            DemuxExit::Status { code, peak_rss_kib } => {
                Ok((synthesize_exit(code), peak_rss_kib))
            }
            DemuxExit::EofBeforeExit => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "exec: agent closed connection before sending EXIT",
            )),
            DemuxExit::Error(msg) => Err(io::Error::new(io::ErrorKind::Other, msg)),
            DemuxExit::Io(e) => Err(e),
        }
    }
}

impl Drop for ExecChild {
    fn drop(&mut self) {
        // Closing stdin lets the child exit naturally for typical
        // programs (cat, sh, etc.). Closing the whole socket is a
        // harder kill — agent will SIGHUP the child.
        self.stdin.take();
        // Demux thread will exit when the socket reader sees EOF
        // or the channel is dropped. We don't join here because
        // Drop must not block; spurious drops are fine.
    }
}

/// Stdin handle. Bytes you write here get framed into STDIN
/// frames and sent to the agent, which writes them to the child's
/// stdin. Drop or call [`close`] for EOF.
///
/// [`close`]: ExecStdin::close
pub struct ExecStdin {
    sock_w: Arc<Mutex<UnixStream>>,
    closed: bool,
}

impl ExecStdin {
    fn new(sock_w: Arc<Mutex<UnixStream>>) -> Self {
        Self { sock_w, closed: false }
    }

    /// Send EOF on stdin. The agent half-closes the child's stdin
    /// pipe; programs that read until EOF (cat, sh, ...) will
    /// then exit.
    pub fn close(mut self) -> io::Result<()> {
        if self.closed {
            return Ok(());
        }
        self.closed = true;
        send_frame(&self.sock_w, FRAME_STDIN, &[])
    }
}

impl Write for ExecStdin {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }
        send_frame(&self.sock_w, FRAME_STDIN, buf)?;
        Ok(buf.len())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Drop for ExecStdin {
    fn drop(&mut self) {
        if !self.closed {
            // Best-effort EOF; ignore errors — the socket may be
            // gone if the child already exited.
            let _ = send_frame(&self.sock_w, FRAME_STDIN, &[]);
        }
    }
}

/// Stdout handle. Wraps the host-side read end of a pipe; the
/// demux thread fills it from STDOUT frames.
pub struct ExecStdout {
    file: File,
}

impl ExecStdout {
    fn from_fd(fd: OwnedFd) -> Self {
        // SAFETY: fd is an owned pipe read end.
        let file = unsafe { File::from_raw_fd(fd.into_raw_fd()) };
        Self { file }
    }
}

impl Read for ExecStdout {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.file.read(buf)
    }
}

/// Stderr handle. Same shape as [`ExecStdout`]; empty in tty mode.
pub struct ExecStderr {
    file: File,
}

impl ExecStderr {
    fn from_fd(fd: OwnedFd) -> Self {
        let file = unsafe { File::from_raw_fd(fd.into_raw_fd()) };
        Self { file }
    }
}

impl Read for ExecStderr {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.file.read(buf)
    }
}

/// Send a CONTROL frame to the in-guest exec agent and wait for
/// its single-frame ack. Used for `signal` (forward TERM/KILL to
/// the workload) and any future control actions.
///
/// Returns Ok if the agent ack'd success; Err with the agent's
/// reported reason otherwise. Network/socket errors surface as
/// `io::Error`; remote-side failures (e.g. "no workload pid yet"
/// during the bake-time window) surface as
/// `io::ErrorKind::Other` with the agent's message.
pub fn send_control(exec_path: &Path, action_json: &serde_json::Value) -> io::Result<()> {
    let _ = send_control_with_ack(exec_path, action_json, None)?;
    Ok(())
}

/// Like [`send_control`] but returns the parsed ack JSON for
/// callers that need ack fields (e.g. `read_file` reads the
/// `data_b64` field). `read_timeout` overrides the default 5s
/// receive timeout — pass a generous value when the agent might
/// take a while (large file reads).
///
/// `pub(crate)` because the JSON shape isn't stabilized for
/// outside consumers.
pub(crate) fn send_control_with_ack(
    exec_path: &Path,
    action_json: &serde_json::Value,
    read_timeout: Option<std::time::Duration>,
) -> io::Result<serde_json::Value> {
    use std::io::Read;
    let mut sock = UnixStream::connect(exec_path)?;
    sock.set_read_timeout(read_timeout.or(Some(std::time::Duration::from_secs(5))))?;
    sock.set_write_timeout(Some(std::time::Duration::from_secs(5)))?;
    // Serialize the JSON body once.
    let body = serde_json::to_vec(action_json)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("encode CONTROL: {e}")))?;
    let mut header = [0u8; 5];
    header[0] = FRAME_CONTROL;
    header[1..5].copy_from_slice(&(body.len() as u32).to_be_bytes());
    sock.write_all(&header)?;
    if !body.is_empty() {
        sock.write_all(&body)?;
    }

    // Read the ack frame. Bumped to 16 MiB to accommodate
    // read_file responses (base64-encoded ≈ 12 MiB raw bytes).
    let mut ack_hdr = [0u8; 5];
    sock.read_exact(&mut ack_hdr)?;
    let kind = ack_hdr[0];
    let len = u32::from_be_bytes([ack_hdr[1], ack_hdr[2], ack_hdr[3], ack_hdr[4]]) as usize;
    if len > 16 * 1024 * 1024 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("CONTROL ack frame too large: {len}"),
        ));
    }
    let mut ack_body = vec![0u8; len];
    if !ack_body.is_empty() {
        sock.read_exact(&mut ack_body)?;
    }

    if kind != FRAME_CONTROL {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("expected CONTROL ack frame, got {kind:#x}"),
        ));
    }
    let ack: serde_json::Value = serde_json::from_slice(&ack_body)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, format!("CONTROL ack JSON: {e}")))?;
    if ack.get("ok").and_then(|v| v.as_bool()) == Some(true) {
        return Ok(ack);
    }
    let msg = ack
        .get("error")
        .and_then(|v| v.as_str())
        .unwrap_or("CONTROL: agent reported failure")
        .to_owned();
    Err(io::Error::new(io::ErrorKind::Other, msg))
}

// ---------- spawn ----------

fn spawn(builder: ExecBuilder) -> io::Result<ExecChild> {
    let sock = UnixStream::connect(&builder.exec_path)?;
    let sock_r = sock.try_clone()?;
    let sock_w = Arc::new(Mutex::new(sock));

    // Send REQUEST.
    let stage_payload: Vec<StageFilePayload> = builder
        .stage_files
        .iter()
        .map(|s| StageFilePayload {
            path: &s.path,
            data_b64: crate::api::b64_encode(&s.data),
            mode: s.mode,
        })
        .collect();
    let payload = RequestPayload {
        argv: &builder.argv,
        env: &builder.env,
        cwd: builder.cwd.as_deref(),
        tty: builder.tty,
        cols: builder.cols,
        rows: builder.rows,
        stage_files: stage_payload,
        chain: &builder.chain,
    };
    let json = serde_json::to_vec(&payload)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("exec: encode REQUEST: {e}")))?;
    send_frame(&sock_w, FRAME_REQUEST, &json)?;

    // Pipes for stdout/stderr.
    let (stdout_r, stdout_w) = pipe()?;
    let (stderr_r, stderr_w) = pipe()?;

    let (exit_tx, exit_rx) = channel::<DemuxExit>();
    let demux = thread::Builder::new()
        .name("supermachine-exec-demux".into())
        .spawn(move || {
            demux_loop(sock_r, stdout_w, stderr_w, exit_tx);
        })
        .map_err(|e| io::Error::new(io::ErrorKind::Other, format!("exec: demux thread: {e}")))?;

    let stdin = ExecStdin::new(sock_w.clone());

    Ok(ExecChild {
        sock_w,
        stdout_r: Some(stdout_r),
        stderr_r: Some(stderr_r),
        stdin: Some(stdin),
        demux: Some(demux),
        exit_rx,
    })
}

fn demux_loop(
    mut sock: UnixStream,
    stdout_w: OwnedFd,
    stderr_w: OwnedFd,
    exit_tx: Sender<DemuxExit>,
) {
    // Use raw fds so we can write straight to the pipe with libc.
    // Keep the OwnedFd alive so dropping closes the host side
    // (which surfaces EOF to the caller's Read).
    let stdout_fd = stdout_w.as_raw_fd();
    let stderr_fd = stderr_w.as_raw_fd();
    loop {
        let (kind, payload) = match read_frame(&mut sock) {
            Ok(f) => f,
            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
                let _ = exit_tx.send(DemuxExit::EofBeforeExit);
                return;
            }
            Err(e) => {
                let _ = exit_tx.send(DemuxExit::Io(e));
                return;
            }
        };
        match kind {
            FRAME_STDOUT => {
                let _ = pipe_write(stdout_fd, &payload);
            }
            FRAME_STDERR => {
                let _ = pipe_write(stderr_fd, &payload);
            }
            FRAME_EXIT => {
                if payload.len() >= 4 {
                    let code = u32::from_be_bytes([payload[0], payload[1], payload[2], payload[3]]);
                    let peak_rss_kib = if payload.len() >= 12 {
                        Some(u64::from_be_bytes([
                            payload[4], payload[5], payload[6], payload[7],
                            payload[8], payload[9], payload[10], payload[11],
                        ]))
                    } else {
                        None
                    };
                    let _ = exit_tx.send(DemuxExit::Status { code, peak_rss_kib });
                } else {
                    let _ = exit_tx.send(DemuxExit::Status { code: 0, peak_rss_kib: None });
                }
                return;
            }
            FRAME_ERROR => {
                let msg = String::from_utf8_lossy(&payload).into_owned();
                let _ = exit_tx.send(DemuxExit::Error(msg));
                return;
            }
            _ => {
                // STDIN / RESIZE / SIGNAL / REQUEST shouldn't come
                // from the agent. Ignore unknown frames so the
                // protocol can grow additive types without
                // breaking older hosts.
            }
        }
    }
}

// ---------- frame I/O ----------

fn send_frame(sock: &Arc<Mutex<UnixStream>>, kind: u8, payload: &[u8]) -> io::Result<()> {
    let mut header = [0u8; 5];
    header[0] = kind;
    header[1..5].copy_from_slice(&(payload.len() as u32).to_be_bytes());
    let mut g = sock.lock().map_err(|_| {
        io::Error::new(io::ErrorKind::Other, "exec: socket mutex poisoned")
    })?;
    g.write_all(&header)?;
    if !payload.is_empty() {
        g.write_all(payload)?;
    }
    Ok(())
}

fn read_frame(sock: &mut UnixStream) -> io::Result<(u8, Vec<u8>)> {
    let mut header = [0u8; 5];
    sock.read_exact(&mut header)?;
    let kind = header[0];
    let len = u32::from_be_bytes([header[1], header[2], header[3], header[4]]) as usize;
    // Ceiling generous enough that the agent can stream large
    // outputs without us tripping a guard, but tight enough that
    // a malformed length doesn't allocate gigabytes.
    if len > 16 * 1024 * 1024 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("exec: frame len {len} > 16 MiB"),
        ));
    }
    let mut payload = vec![0u8; len];
    if !payload.is_empty() {
        sock.read_exact(&mut payload)?;
    }
    Ok((kind, payload))
}

// ---------- pipe helpers ----------

fn pipe() -> io::Result<(OwnedFd, OwnedFd)> {
    let mut fds = [0 as RawFd; 2];
    let r = unsafe { libc::pipe(fds.as_mut_ptr()) };
    if r < 0 {
        return Err(io::Error::last_os_error());
    }
    Ok((unsafe { OwnedFd::from_raw_fd(fds[0]) }, unsafe {
        OwnedFd::from_raw_fd(fds[1])
    }))
}

fn pipe_write(fd: RawFd, mut buf: &[u8]) -> io::Result<()> {
    while !buf.is_empty() {
        let n = unsafe {
            libc::write(fd, buf.as_ptr() as *const libc::c_void, buf.len())
        };
        if n < 0 {
            let e = io::Error::last_os_error();
            if e.raw_os_error() == Some(libc::EINTR) {
                continue;
            }
            return Err(e);
        }
        buf = &buf[n as usize..];
    }
    Ok(())
}

// ---------- exit status synthesis ----------

#[cfg(unix)]
fn synthesize_exit(code: u32) -> ExitStatus {
    // ExitStatus on unix wraps a wait status int. The agent
    // already collapsed signal/exit into a u32 (signal -> 128+sig).
    // We re-encode as if WIFEXITED with that exit code: high byte
    // = exit code, low byte = 0. Calling `.code()` then returns
    // `Some(code)` which matches what callers expect.
    use std::os::unix::process::ExitStatusExt;
    ExitStatus::from_raw(((code as i32) & 0xff) << 8)
}