//! vsock client for communicating with the smolvm-agent.
//!
//! This module provides a client for sending requests to the agent
//! and receiving responses.

use crate::error::{Error, Result};
use crate::platform::uds::UdsStream;
use crate::registry::{extract_registry, rewrite_image_registry, RegistryAuth};
use crate::settings::SmolSettings;
use smolvm_protocol::normalize_image_ref;
use smolvm_protocol::{
    encode_message, AgentRequest, AgentResponse, Envelope, FsNotifyEvent, ImageInfo, OverlayInfo,
    StorageStatus, FILE_TRANSFER_MAX_TOTAL, FILE_WRITE_CHUNK_SIZE, FILE_WRITE_SINGLE_SHOT_MAX,
    MAX_FRAME_SIZE, PROTOCOL_VERSION,
};
use std::io::{Read, Write};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

/// Events from a streaming exec session.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExecEvent {
    /// Standard output data.
    Stdout(Vec<u8>),
    /// Standard error data.
    Stderr(Vec<u8>),
    /// Command exited with this code.
    Exit(i32),
    /// An error occurred.
    Error(String),
}

/// One input event fed into a channel-driven interactive session
/// ([`AgentClient::interactive_session_io`]). This decouples the interactive
/// poll loop from the process's real stdin so the session can be driven by a
/// remote transport (e.g. a WebSocket terminal) instead of a local TTY.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum InteractiveInput {
    /// Bytes to forward to the command's stdin.
    Stdin(Vec<u8>),
    /// Terminal resize (PTY window change).
    Resize {
        /// New terminal width in columns.
        cols: u16,
        /// New terminal height in rows.
        rows: u16,
    },
    /// End of input — sends an empty stdin frame (EOF) to the command.
    Eof,
}

/// One output chunk produced by a channel-driven interactive session.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum InteractiveOutput {
    /// A chunk of bytes from the command's stdout.
    Stdout(Vec<u8>),
    /// A chunk of bytes from the command's stderr.
    Stderr(Vec<u8>),
}

// ============================================================================
// Socket Timeout Constants
// ============================================================================
//
// These timeouts control how long the client waits for various operations.
// They balance between allowing slow operations to complete and failing fast
// when the agent is unresponsive.

/// Default socket read timeout (30 seconds).
/// Used for most request/response operations. Long enough for the agent to
/// process requests, short enough to detect hung connections.
const DEFAULT_READ_TIMEOUT_SECS: u64 = 30;

/// Default socket write timeout (10 seconds).
/// Writes should complete quickly - if they don't, the connection is likely broken.
const DEFAULT_WRITE_TIMEOUT_SECS: u64 = 10;

/// Read timeout for image pull operations (10 minutes).
/// Image pulls can take a long time for large images over slow connections.
const IMAGE_PULL_TIMEOUT_SECS: u64 = 600;

/// Maximum silence while starting a detached container.
///
/// Archive flattening may take arbitrarily long, but the agent streams progress
/// while work is advancing. Each progress response starts a fresh socket read,
/// so this remains an inactivity timeout rather than a total operation limit.
const DETACHED_START_TIMEOUT_SECS: u64 = 120;

// (Removed INTERACTIVE_TIMEOUT_SECS — no-user-timeout execs now disable
// the socket read timeout entirely, matching interactive_session behavior.)

/// Buffer time added to user-specified timeouts (5 seconds).
/// When users specify a command timeout, we add this buffer to the socket
/// timeout to allow for protocol overhead and response transmission.
const TIMEOUT_BUFFER_SECS: u64 = 5;

/// Timeout for shutdown acknowledgment (5 seconds).
/// sync() + ack transmission is typically <100ms, but heavy writes or
/// large journals may take longer. If no ack within 5s, the VM has
/// likely already torn down — safe to proceed with SIGTERM.
const SHUTDOWN_ACK_TIMEOUT_SECS: u64 = 5;

// ============================================================================
// I/O Constants
// ============================================================================

/// Buffer size for reading stdin during interactive sessions.
const STDIN_BUF_SIZE: usize = 4096;

/// Poll timeout in milliseconds for interactive I/O loops.
/// Short enough for responsive SIGWINCH handling, long enough to avoid busy-waiting.
const POLL_TIMEOUT_MS: i32 = 100;

/// Poll interval while waiting for writer-queue capacity.
///
/// A short interval limits added stdin latency without busy-waiting.
const BACKPRESSURE_POLL_TIMEOUT_MS: i32 = 5;

/// Maximum frames buffered by the interactive writer.
///
/// When full, stdin polling stops, providing backpressure instead of buffering
/// indefinitely. 64 × 4 KiB stdin frames encode to roughly 350 KiB worst case.
const WRITER_QUEUE_FRAMES: usize = 64;

/// Write timeout for the interactive writer thread.
///
/// A short timeout lets blocked writes periodically observe shutdown.
/// `WouldBlock` is flow control and is retried from the current frame offset.
const WRITER_WRITE_TIMEOUT: Duration = Duration::from_secs(1);

/// Exit code reported when a channel-driven interactive session ends because the
/// remote peer (e.g. a WebSocket terminal) disconnected rather than the command
/// exiting. 128 + SIGINT(2), matching the shell convention for an interrupted job.
const DISCONNECT_EXIT_CODE: i32 = 130;

/// RAII guard that resets the socket read timeout on drop.
///
/// Ensures the timeout is always restored, even if the operation
/// returns early due to an error. Uses a cloned UnixStream handle
/// (shares the underlying fd) to avoid borrow conflicts.
pub struct ReadTimeoutGuard {
    stream: UdsStream,
}

impl ReadTimeoutGuard {
    /// Create a guard from a reference to the stream.
    /// Clones the underlying fd so the guard doesn't borrow the original.
    fn new(stream: &UdsStream) -> Option<Self> {
        stream.try_clone().ok().map(|s| Self { stream: s })
    }
}

impl Drop for ReadTimeoutGuard {
    fn drop(&mut self) {
        if let Err(e) = self
            .stream
            .set_read_timeout(Some(Duration::from_secs(DEFAULT_READ_TIMEOUT_SECS)))
        {
            tracing::warn!(error = %e, "failed to reset socket read timeout to default");
        }
    }
}

/// Dedicated writer for interactive-session frames.
///
/// Keeping socket writes off the session loop lets it continue draining guest
/// output while writes are blocked, avoiding the full-duplex deadlock in issue #926.
/// A single writer also lets partial writes resume without corrupting framing.
struct FrameWriter {
    /// Queued frames, bounded by [`WRITER_QUEUE_FRAMES`]. Taken (dropped) during
    /// shutdown so the thread's `recv` returns and `join` cannot deadlock.
    tx: Option<std::sync::mpsc::SyncSender<Vec<u8>>>,
    handle: Option<std::thread::JoinHandle<()>>,
    shutdown: Arc<AtomicBool>,
    /// First fatal write error, published for the session loop to pick up.
    error: Arc<Mutex<Option<std::io::Error>>>,
    /// Extra handle to the same socket, kept only to restore the default write
    /// timeout once the thread is gone.
    restore: UdsStream,
}

impl FrameWriter {
    /// Clone `stream` and start the writer thread on the clone.
    ///
    /// Note that socket options are properties of the socket, not of the
    /// descriptor, so the write timeout set here applies to the caller's handle
    /// too. That is harmless — an interactive session performs no direct writes
    /// while the thread lives — and [`Drop`] puts the default back.
    fn spawn(stream: &UdsStream) -> Result<Self> {
        let write_half = stream
            .try_clone()
            .map_err(|e| Error::agent("clone agent socket", e.to_string()))?;
        let restore = stream
            .try_clone()
            .map_err(|e| Error::agent("clone agent socket", e.to_string()))?;
        // Best effort; on failure the writer retains the inherited timeout.
        if let Err(e) = write_half.set_write_timeout(Some(WRITER_WRITE_TIMEOUT)) {
            tracing::debug!(error = %e, "could not shorten agent socket write timeout");
        }

        let (tx, rx) = std::sync::mpsc::sync_channel::<Vec<u8>>(WRITER_QUEUE_FRAMES);
        let shutdown = Arc::new(AtomicBool::new(false));
        let error = Arc::new(Mutex::new(None));

        let thread_shutdown = Arc::clone(&shutdown);
        let thread_error = Arc::clone(&error);
        let handle = std::thread::Builder::new()
            .name("smolvm-agent-writer".to_string())
            .spawn(move || {
                write_frames(write_half, rx, &thread_shutdown, &thread_error);
            })
            .map_err(|e| Error::agent("spawn writer thread", e.to_string()))?;

        Ok(Self {
            tx: Some(tx),
            handle: Some(handle),
            shutdown,
            error,
            restore,
        })
    }

    /// Queue `frame` without blocking.
    ///
    /// Returns `Ok(None)` when queued and `Ok(Some(frame))` — the frame handed
    /// back — when the queue is full, so the caller can hold it and stop reading
    /// more input. `Err` means the writer thread is gone.
    fn try_send(&self, frame: Vec<u8>) -> Result<Option<Vec<u8>>> {
        let tx = match self.tx.as_ref() {
            Some(tx) => tx,
            None => return Err(self.writer_gone_error()),
        };
        match tx.try_send(frame) {
            Ok(()) => Ok(None),
            Err(std::sync::mpsc::TrySendError::Full(frame)) => Ok(Some(frame)),
            Err(std::sync::mpsc::TrySendError::Disconnected(_)) => Err(self.writer_gone_error()),
        }
    }

    /// Take the writer's fatal error, if it has hit one.
    fn take_error(&self) -> Option<std::io::Error> {
        self.error.lock().ok().and_then(|mut slot| slot.take())
    }

    /// The error to report when the thread has exited: its own if it stored one,
    /// otherwise a generic description.
    fn writer_gone_error(&self) -> Error {
        match self.take_error() {
            Some(e) => Error::agent("send message", e.to_string()),
            None => Error::agent("send message", "agent socket writer stopped".to_string()),
        }
    }
}

impl Drop for FrameWriter {
    fn drop(&mut self) {
        self.shutdown.store(true, Ordering::Relaxed);
        // Drop the sender first: a thread parked in `recv` wakes on disconnect,
        // and one parked in `write` sees the flag at its next timeout tick.
        self.tx = None;
        if let Some(handle) = self.handle.take() {
            if handle.join().is_err() {
                tracing::warn!("agent socket writer thread panicked");
            }
        }
        if let Err(e) = self
            .restore
            .set_write_timeout(Some(Duration::from_secs(DEFAULT_WRITE_TIMEOUT_SECS)))
        {
            tracing::warn!(error = %e, "failed to reset socket write timeout to default");
        }
    }
}

/// How writing one frame ended.
enum FrameOutcome {
    /// Every byte reached the socket.
    Done,
    /// Shutdown was requested before the frame finished.
    Abandoned,
    /// The socket rejected the write.
    Failed(std::io::Error),
}

/// Body of the writer thread: drain `rx` and write each frame in full.
fn write_frames(
    stream: UdsStream,
    rx: std::sync::mpsc::Receiver<Vec<u8>>,
    shutdown: &AtomicBool,
    error: &Mutex<Option<std::io::Error>>,
) {
    let mut sink = &stream;
    // `recv` ends when every sender is dropped, i.e. at session teardown.
    while let Ok(frame) = rx.recv() {
        let mut pos = 0;
        let outcome = loop {
            if pos == frame.len() {
                break FrameOutcome::Done;
            }
            match sink.write(&frame[pos..]) {
                Ok(0) => {
                    break FrameOutcome::Failed(std::io::Error::new(
                        std::io::ErrorKind::WriteZero,
                        "agent socket accepted no bytes",
                    ))
                }
                Ok(n) => pos += n,
                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                    // The send timeout expired because the guest is not draining
                    // yet — the very condition that used to abort the session.
                    // Resume from `pos` unless we are shutting down.
                    if shutdown.load(Ordering::Relaxed) {
                        break FrameOutcome::Abandoned;
                    }
                }
                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
                Err(e) => break FrameOutcome::Failed(e),
            }
        };

        match outcome {
            FrameOutcome::Done => {
                if let Err(e) = sink.flush() {
                    store_writer_error(error, e);
                    return;
                }
            }
            FrameOutcome::Abandoned => {
                close_write_half_if_truncated(&stream, pos);
                return;
            }
            FrameOutcome::Failed(e) => {
                store_writer_error(error, e);
                close_write_half_if_truncated(&stream, pos);
                return;
            }
        }
    }
}

/// Close the write half after a partial frame.
///
/// Otherwise the peer can wait indefinitely for the remaining bytes or
/// interpret subsequent data as part of the incomplete frame.
fn close_write_half_if_truncated(stream: &UdsStream, written: usize) {
    if written == 0 {
        return;
    }
    tracing::debug!(
        written,
        "agent socket writer stopped mid-frame; closing the write half"
    );
    let _ = stream.shutdown(std::net::Shutdown::Write);
}

/// Hand as many queued frames as fit to the writer thread, in order.
///
/// Anything left in `pending` means the writer is behind; the session loop then
/// stops taking new input until this drains, which is what keeps host memory
/// bounded without ever blocking the loop's read side.
fn pump_frames(
    writer: &FrameWriter,
    pending: &mut std::collections::VecDeque<Vec<u8>>,
) -> Result<()> {
    while let Some(frame) = pending.pop_front() {
        if let Some(frame) = writer.try_send(frame)? {
            pending.push_front(frame);
            break;
        }
    }
    Ok(())
}

/// Publish the writer thread's first fatal error.
fn store_writer_error(slot: &Mutex<Option<std::io::Error>>, err: std::io::Error) {
    tracing::debug!(error = %err, "agent socket writer failed");
    if let Ok(mut slot) = slot.lock() {
        if slot.is_none() {
            *slot = Some(err);
        }
    }
}

/// Configuration for running a command interactively.
#[derive(Debug, Clone)]
pub struct RunConfig {
    /// OCI image to run.
    pub image: String,
    /// Command and arguments to execute.
    pub command: Vec<String>,
    /// Environment variables as (key, value) pairs.
    pub env: Vec<(String, String)>,
    /// Working directory inside the container.
    pub workdir: Option<String>,
    /// User to execute as inside the container.
    pub user: Option<String>,
    /// Volume mounts as (tag, guest_path, read_only) tuples.
    pub mounts: Vec<(String, String, bool)>,
    /// Timeout for command execution.
    pub timeout: Option<Duration>,
    /// Whether to allocate a TTY.
    pub tty: bool,
    /// Persistent overlay ID. If set, the overlay persists across exec sessions
    /// so filesystem changes (e.g. package installs) survive.
    pub persistent_overlay_id: Option<String>,
    /// Data to pipe to the command's stdin (non-interactive runs only). The
    /// pipe is closed after writing so the command sees EOF.
    pub stdin: Option<String>,
    /// Run as an unprivileged container (restricted caps, ro cgroup, no extra
    /// tmpfs). Default false = "VM-grade" (the microVM is the boundary).
    pub unprivileged: bool,
}

impl RunConfig {
    /// Create a new run configuration with the given image and command.
    ///
    /// The image reference is canonicalized immediately so all downstream
    /// code (cache keys, logs, protocol messages) sees the same form
    /// regardless of how the caller spelled it.
    pub fn new(image: impl Into<String>, command: Vec<String>) -> Self {
        Self {
            image: normalize_image_ref(&image.into()),
            command,
            env: Vec::new(),
            workdir: None,
            user: None,
            mounts: Vec::new(),
            timeout: None,
            tty: false,
            persistent_overlay_id: None,
            stdin: None,
            unprivileged: false,
        }
    }

    /// Set environment variables.
    pub fn with_env(mut self, env: Vec<(String, String)>) -> Self {
        self.env = env;
        self
    }

    /// Set working directory.
    pub fn with_workdir(mut self, workdir: Option<String>) -> Self {
        self.workdir = workdir;
        self
    }

    /// Set container user.
    pub fn with_user(mut self, user: Option<String>) -> Self {
        self.user = user;
        self
    }

    /// Set volume mounts.
    pub fn with_mounts(mut self, mounts: Vec<(String, String, bool)>) -> Self {
        self.mounts = mounts;
        self
    }

    /// Set timeout.
    pub fn with_timeout(mut self, timeout: Option<Duration>) -> Self {
        self.timeout = timeout;
        self
    }

    /// Enable TTY mode.
    pub fn with_tty(mut self, tty: bool) -> Self {
        self.tty = tty;
        self
    }

    /// Set stdin data piped to the command (non-interactive runs).
    pub fn with_stdin(mut self, stdin: Option<String>) -> Self {
        self.stdin = stdin;
        self
    }

    /// Set persistent overlay ID for cross-session filesystem persistence.
    pub fn with_persistent_overlay(mut self, id: Option<String>) -> Self {
        self.persistent_overlay_id = id;
        self
    }

    /// Run as an unprivileged container (defense-in-depth for untrusted code).
    pub fn with_unprivileged(mut self, unprivileged: bool) -> Self {
        self.unprivileged = unprivileged;
        self
    }
}

/// Options for pulling an OCI image.
///
/// Use `PullOptions::new()` to create with defaults, then chain methods
/// to customize behavior.
///
/// # Example
///
/// ```ignore
/// let options = PullOptions::new()
///     .oci_platform("linux/arm64")
///     .use_registry_config(true)
///     .progress(|cur, total, layer| println!("{}/{}: {}", cur, total, layer));
///
/// client.pull("alpine:latest", options)?;
/// ```
#[derive(Default)]
pub struct PullOptions<F = fn(usize, usize, &str)>
where
    F: FnMut(usize, usize, &str),
{
    /// OCI platform to pull (e.g., "linux/arm64").
    pub oci_platform: Option<String>,
    /// Explicit authentication credentials.
    pub auth: Option<RegistryAuth>,
    /// Whether to load credentials from registry config file.
    pub use_registry_config: bool,
    /// Proxy URL applied to the in-VM registry client (HTTP_PROXY/HTTPS_PROXY).
    pub proxy: Option<String>,
    /// Comma-separated NO_PROXY list of hosts/CIDRs that bypass the proxy.
    pub no_proxy: Option<String>,
    /// Progress callback: (current, total, layer_id).
    pub progress: Option<F>,
}

impl PullOptions<fn(usize, usize, &str)> {
    /// Create new pull options with defaults.
    pub fn new() -> Self {
        Self {
            oci_platform: None,
            auth: None,
            use_registry_config: false,
            proxy: None,
            no_proxy: None,
            progress: None,
        }
    }
}

impl<F: FnMut(usize, usize, &str)> PullOptions<F> {
    /// Set the target OCI platform (e.g., "linux/arm64").
    pub fn oci_platform(mut self, oci_platform: impl Into<String>) -> Self {
        self.oci_platform = Some(oci_platform.into());
        self
    }

    /// Set explicit authentication credentials.
    pub fn auth(mut self, auth: RegistryAuth) -> Self {
        self.auth = Some(auth);
        self
    }

    /// Enable loading credentials from registry config file.
    ///
    /// When enabled, loads `~/.config/smolvm/registries.toml` and
    /// automatically provides credentials for matching registries.
    /// Also applies registry mirrors if configured.
    pub fn use_registry_config(mut self, enabled: bool) -> Self {
        self.use_registry_config = enabled;
        self
    }

    /// Set the proxy URL applied to the in-VM registry client.
    pub fn proxy(mut self, proxy: impl Into<String>) -> Self {
        self.proxy = Some(proxy.into());
        self
    }

    /// Set the NO_PROXY list for the in-VM registry client.
    pub fn no_proxy(mut self, no_proxy: impl Into<String>) -> Self {
        self.no_proxy = Some(no_proxy.into());
        self
    }

    /// Set a progress callback.
    ///
    /// The callback receives (current_percent, total=100, layer_id) for each layer.
    pub fn progress<G: FnMut(usize, usize, &str)>(self, callback: G) -> PullOptions<G> {
        PullOptions {
            oci_platform: self.oci_platform,
            auth: self.auth,
            use_registry_config: self.use_registry_config,
            proxy: self.proxy,
            no_proxy: self.no_proxy,
            progress: Some(callback),
        }
    }
}

/// Raw descriptor of the process's stdin, in the portable form the terminal
/// poll loop expects. Unix: the real fd; Windows: the `STD_INPUT_HANDLE`
/// console handle cast into the portable `Fd`.
fn stdin_raw_fd() -> crate::agent::terminal::Fd {
    #[cfg(unix)]
    {
        use std::os::unix::io::AsRawFd;
        std::io::stdin().as_raw_fd()
    }
    #[cfg(not(unix))]
    {
        // SAFETY: GetStdHandle has no preconditions; it returns the process's
        // standard-input HANDLE (or INVALID_HANDLE_VALUE), which the Windows
        // poll loop interprets.
        let handle = unsafe {
            windows_sys::Win32::System::Console::GetStdHandle(
                windows_sys::Win32::System::Console::STD_INPUT_HANDLE,
            )
        };
        handle as crate::agent::terminal::Fd
    }
}

/// Check if a shutdown receive error is a benign race condition.
///
/// During shutdown the VM may tear down before the ack response is flushed,
/// causing EAGAIN, connection reset, or similar errors. These are expected
/// and don't indicate a problem — sync() has likely already completed.
fn is_benign_shutdown_error(error_str: &str) -> bool {
    error_str.contains("os error 35") // EAGAIN on macOS
        || error_str.contains("os error 11") // EAGAIN on Linux
        || error_str.contains("temporarily unavailable")
        || error_str.contains("Connection reset")
        || error_str.contains("connection reset")
}

/// Client for communicating with the smolvm-agent.
pub struct AgentClient {
    stream: UdsStream,
    /// Trace ID for correlating this client session's requests with host API calls.
    trace_id: Option<String>,
}

// ============================================================================
// Response match helpers
// ============================================================================

/// Extract typed data from an `Ok` response.
fn expect_data<T: serde::de::DeserializeOwned>(resp: AgentResponse, op: &str) -> Result<T> {
    match resp {
        AgentResponse::Ok {
            data: Some(data), ..
        } => {
            serde_json::from_value(data).map_err(|e| Error::agent("parse response", e.to_string()))
        }
        AgentResponse::Error { message, .. } => Err(Error::agent(op, message)),
        _ => Err(Error::agent(op, "unexpected response type")),
    }
}

/// Give a PTY session a `TERM` unless the caller already set one.
///
/// A shell's line editor needs `TERM` to look up the terminal's cursor-movement
/// capabilities. Without it the edit still happens — the buffer is correct — but
/// the redraw is not: the screen never catches up, so backspace looks like it
/// advances the cursor and deletes nothing. zsh's ZLE is strict here; bash's
/// readline has built-in fallbacks and hides the problem.
///
/// The host's own `TERM` is deliberately NOT forwarded. An exotic value
/// (`alacritty`, `foot`, …) often has no terminfo entry inside a slim container
/// image, which fails exactly the same way — so a value every image ships with
/// is the safer default. It also matches the `TERM` the launcher and OCI paths
/// already set. Callers wanting something else pass `TERM` in `env` explicitly.
fn with_term_default(mut env: Vec<(String, String)>, tty: bool) -> Vec<(String, String)> {
    if tty && !env.iter().any(|(key, _)| key == "TERM") {
        env.push(("TERM".to_string(), "xterm-256color".to_string()));
    }
    env
}

/// Expect an `Ok` response, ignoring any data.
fn expect_ok(resp: AgentResponse, op: &str) -> Result<()> {
    match resp {
        AgentResponse::Ok { .. } => Ok(()),
        AgentResponse::Error { message, .. } => Err(Error::agent(op, message)),
        _ => Err(Error::agent(op, "unexpected response type")),
    }
}

/// Extract exit code, stdout, stderr from a `Completed` response.
fn expect_completed(resp: AgentResponse, op: &str) -> Result<(i32, Vec<u8>, Vec<u8>)> {
    match resp {
        AgentResponse::Completed {
            exit_code,
            stdout,
            stderr,
        } => Ok((exit_code, stdout, stderr)),
        AgentResponse::Error { message, .. } => Err(Error::agent(op, message)),
        _ => Err(Error::agent(op, "unexpected response type")),
    }
}

#[cfg(test)]
impl AgentClient {
    /// Build an `AgentClient` from a pre-connected `UnixStream`.
    ///
    /// Test-only: production code must go through [`AgentClient::connect`]
    /// so socket timeouts are configured correctly. Used by the regression
    /// tests that drive the client against a `UdsStream::pair()`.
    pub(crate) fn from_stream(stream: UdsStream) -> Self {
        Self {
            stream,
            trace_id: None,
        }
    }
}

impl AgentClient {
    /// Set socket read timeout, returning an error if it fails.
    ///
    /// This is a helper to ensure timeout failures are always handled properly,
    /// preventing indefinite hangs on read operations.
    fn set_read_timeout(&self, timeout: Duration) -> Result<()> {
        self.stream.set_read_timeout(Some(timeout)).map_err(|e| {
            Error::agent(
                "set read timeout",
                format!("failed to set socket read timeout to {:?}: {}", timeout, e),
            )
        })
    }

    /// Connect to the agent via Unix socket.
    ///
    /// # Arguments
    ///
    /// * `socket_path` - Path to the vsock Unix socket
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Connection to the socket fails
    /// - Socket timeouts cannot be configured (prevents indefinite hangs)
    pub fn connect(socket_path: impl AsRef<Path>) -> Result<Self> {
        Self::connect_once(socket_path.as_ref())
    }

    /// Connect to the agent with retry logic for transient failures.
    ///
    /// This is useful when the agent might be temporarily unavailable
    /// (e.g., during high load or brief network issues).
    pub fn connect_with_retry(socket_path: impl AsRef<Path>) -> Result<Self> {
        use crate::util::{retry_with_backoff, RetryConfig};

        let path = socket_path.as_ref();

        retry_with_backoff(
            RetryConfig::for_connection(),
            "agent connect",
            || Self::connect_once(path),
            |e| {
                // Check if this is a transient error worth retrying
                let error_msg = e.to_string();
                // Connection refused/reset are transient during VM startup.
                // "No such file or directory" occurs when the vsock socket
                // file hasn't been created yet by libkrun's muxer thread —
                // transient under concurrent boot contention.
                error_msg.contains("Connection refused")
                    || error_msg.contains("connection refused")
                    || error_msg.contains("Connection reset")
                    || error_msg.contains("connection reset")
                    || error_msg.contains("Broken pipe")
                    || error_msg.contains("Resource temporarily unavailable")
                    || error_msg.contains("No such file or directory")
            },
        )
    }

    /// Connect with a short timeout, for use during startup ping probes.
    /// Uses 100ms read timeout instead of 30s to fail fast during boot.
    /// The agent completes init in ~130ms of guest uptime, so 100ms is enough
    /// to detect a ready agent without wasting time on a full 1s timeout.
    pub fn connect_with_short_timeout(socket_path: impl AsRef<Path>) -> Result<Self> {
        Self::connect_with_timeouts_ms(socket_path.as_ref(), 100, 100)
    }

    /// Connect with a moderate timeout, for state-probe "is this agent alive"
    /// checks from `machine ls` / `machine status`. 3 seconds is long enough
    /// to avoid false "unreachable" readings when the agent is momentarily
    /// busy (e.g., processing a Run request's overlayfs setup), but short
    /// enough to not make `ls` feel sluggish when the agent is truly dead.
    pub fn connect_for_state_probe(socket_path: impl AsRef<Path>) -> Result<Self> {
        Self::connect_with_timeouts_ms(socket_path.as_ref(), 3000, 3000)
    }

    /// Connect with a very short timeout for boot-time probe cycles.
    /// Uses 5ms timeout to minimize blocking between ready-marker checks.
    /// Only used in the fallback path (old agents without ready markers).
    pub fn connect_with_boot_probe_timeout(socket_path: impl AsRef<Path>) -> Result<Self> {
        Self::connect_with_timeouts_ms(socket_path.as_ref(), 5, 5)
    }

    /// Connect to a restored clone while allowing its resumed vsock transport
    /// to complete the handshake. Reopening 5 ms boot probes can fill the
    /// restored RX queue before the agent receives any ping payload.
    pub fn connect_with_clone_ready_timeout(socket_path: impl AsRef<Path>) -> Result<Self> {
        Self::connect_with_timeouts_ms(socket_path.as_ref(), 1000, 1000)
    }

    /// Internal connect implementation (single attempt).
    fn connect_once(socket_path: &Path) -> Result<Self> {
        Self::connect_with_timeouts(
            socket_path,
            DEFAULT_READ_TIMEOUT_SECS,
            DEFAULT_WRITE_TIMEOUT_SECS,
        )
    }

    /// Connect to the agent socket and configure read/write timeouts (in seconds).
    fn connect_with_timeouts(socket_path: &Path, read_secs: u64, write_secs: u64) -> Result<Self> {
        Self::connect_with_timeouts_ms(socket_path, read_secs * 1000, write_secs * 1000)
    }

    /// Connect to the agent socket and configure read/write timeouts (in milliseconds).
    fn connect_with_timeouts_ms(socket_path: &Path, read_ms: u64, write_ms: u64) -> Result<Self> {
        let stream = UdsStream::connect(socket_path)
            .map_err(|e| Error::agent("connect to agent", e.to_string()))?;

        stream
            .set_read_timeout(Some(Duration::from_millis(read_ms)))
            .map_err(|e| Error::agent("set read timeout", e.to_string()))?;
        stream
            .set_write_timeout(Some(Duration::from_millis(write_ms)))
            .map_err(|e| Error::agent("set write timeout", e.to_string()))?;

        Ok(Self {
            stream,
            trace_id: None,
        })
    }

    /// Set a trace ID for correlating this client session's requests with host API calls.
    /// All subsequent requests will include this trace_id in the Envelope.
    pub fn set_trace_id(&mut self, trace_id: String) {
        self.trace_id = Some(trace_id);
    }

    /// Encode a request wrapped in an Envelope with the current trace_id.
    fn encode_traced(&self, req: &AgentRequest) -> Result<Vec<u8>> {
        let envelope = Envelope::with_trace_id(req, self.trace_id.clone());
        encode_message(&envelope).map_err(|e| Error::agent("encode message", e.to_string()))
    }

    /// Send a request and receive a response.
    fn request(&mut self, req: &AgentRequest) -> Result<AgentResponse> {
        // Encode and send request
        let data = self.encode_traced(req)?;
        self.stream
            .write_all(&data)
            .map_err(|e| Error::agent("send message", e.to_string()))?;

        // Read response
        self.receive()
    }

    fn ping_info(&mut self) -> Result<(u32, Vec<String>)> {
        let resp = self.request(&AgentRequest::Ping)?;

        match resp {
            AgentResponse::Pong {
                version,
                capabilities,
            } => {
                if version != PROTOCOL_VERSION {
                    tracing::warn!(
                        host_version = PROTOCOL_VERSION,
                        agent_version = version,
                        "protocol version mismatch — agent may be outdated or newer than host"
                    );
                }
                Ok((version, capabilities))
            }
            AgentResponse::Error { message, .. } => Err(Error::agent("ping", message)),
            _ => Err(Error::agent("ping", "unexpected response type")),
        }
    }

    /// Ping the helper daemon and validate the protocol version.
    ///
    /// Returns the agent's protocol version. Logs a warning if the version
    /// doesn't match the host's expected version.
    pub fn ping(&mut self) -> Result<u32> {
        self.ping_info().map(|(version, _)| version)
    }

    /// Return whether the live guest agent advertises an optional feature.
    pub fn supports_capability(&mut self, capability: &str) -> Result<bool> {
        self.ping_info().map(|(_, capabilities)| {
            capabilities
                .iter()
                .any(|advertised| advertised == capability)
        })
    }

    /// Replay host-originated filesystem changes into the guest as fsnotify
    /// events so inotify-based watchers on `-v` mounts fire. Used by the host
    /// [`FsNotifyWatcher`](super::FsNotifyWatcher); a stale connection surfaces as
    /// an error the watcher treats as "VM gone, stop".
    pub fn fsnotify(&mut self, events: Vec<FsNotifyEvent>) -> Result<()> {
        if events.is_empty() {
            return Ok(());
        }
        match self.request(&AgentRequest::FsNotify { events })? {
            AgentResponse::Ok { .. } => Ok(()),
            AgentResponse::Error { message, .. } => Err(Error::agent("fsnotify", message)),
            _ => Err(Error::agent("fsnotify", "unexpected response type")),
        }
    }

    /// Pull an OCI image with the given options.
    ///
    /// This is the primary pull method. Use `PullOptions` to configure
    /// authentication, platform, and progress tracking.
    ///
    /// # Example
    ///
    /// ```ignore
    /// // Simple pull
    /// client.pull("alpine:latest", PullOptions::new())?;
    ///
    /// // Pull with registry config (loads credentials from config file)
    /// client.pull("ghcr.io/owner/repo", PullOptions::new().use_registry_config(true))?;
    ///
    /// // Pull with explicit auth and progress
    /// client.pull("private.registry/image", PullOptions::new()
    ///     .auth(RegistryAuth { username: "user".into(), password: "pass".into() })
    ///     .progress(|cur, total, layer| eprintln!("{}%", cur)))?;
    /// ```
    ///
    /// # Note
    ///
    /// This operation uses a 10-minute timeout to accommodate large images.
    pub fn pull<F: FnMut(usize, usize, &str)>(
        &mut self,
        image: &str,
        options: PullOptions<F>,
    ) -> Result<ImageInfo> {
        // Resolve effective image and auth based on options
        let (effective_image, effective_auth) = if options.use_registry_config {
            let registry_config = SmolSettings::load().unwrap_or_default().images;
            let registry = extract_registry(image);

            // Get credentials from config if not explicitly provided
            let auth = options.auth.or_else(|| {
                registry_config.get_credentials(&registry).inspect(|creds| {
                    tracing::debug!(
                        registry = %registry,
                        username = %creds.username,
                        "using configured registry credentials"
                    );
                })
            });

            // Apply mirror if configured
            let img = if let Some(mirror) = registry_config.get_mirror(&registry) {
                let mirrored = rewrite_image_registry(image, mirror);
                tracing::debug!(
                    original = %image,
                    mirrored = %mirrored,
                    mirror = %mirror,
                    "using registry mirror"
                );
                mirrored
            } else {
                image.to_string()
            };

            (img, auth)
        } else {
            (image.to_string(), options.auth)
        };

        self.pull_image_internal(
            &effective_image,
            options.oci_platform.as_deref(),
            effective_auth.as_ref(),
            options.proxy.as_deref(),
            options.no_proxy.as_deref(),
            options.progress,
        )
    }

    /// Internal implementation of image pull.
    fn pull_image_internal<F: FnMut(usize, usize, &str)>(
        &mut self,
        image: &str,
        oci_platform: Option<&str>,
        auth: Option<&RegistryAuth>,
        proxy: Option<&str>,
        no_proxy: Option<&str>,
        mut progress: Option<F>,
    ) -> Result<ImageInfo> {
        let image = normalize_image_ref(image);
        let image = image.as_str();

        // Use a long timeout for pull - large images can take minutes to download/extract.
        // The guard resets the timeout on drop (including error paths).
        self.set_read_timeout(Duration::from_secs(IMAGE_PULL_TIMEOUT_SECS))?;
        let _timeout_guard = ReadTimeoutGuard::new(&self.stream);

        // Send the pull request
        let data = self.encode_traced(&AgentRequest::Pull {
            image: image.to_string(),
            oci_platform: oci_platform.map(String::from),
            auth: auth.cloned(),
            proxy: proxy.map(String::from),
            no_proxy: no_proxy.map(String::from),
        })?;

        self.stream
            .write_all(&data)
            .map_err(|e| Error::agent("send request", e.to_string()))?;

        // Read responses - loop until we get Ok or Error (skip Progress)
        loop {
            match self.receive()? {
                AgentResponse::Progress {
                    percent,
                    layer,
                    message: _,
                } => {
                    if let Some(ref mut cb) = progress {
                        let current = percent.unwrap_or(0) as usize;
                        let layer_id = layer.as_deref().unwrap_or("");
                        cb(current, 100, layer_id);
                    }
                }
                AgentResponse::Ok { data: Some(data) } => {
                    return serde_json::from_value(data)
                        .map_err(|e| Error::agent("parse response", e.to_string()));
                }
                AgentResponse::Error { message, .. } => {
                    return Err(Error::agent("pull image", message));
                }
                _ => {
                    return Err(Error::agent("pull image", "unexpected response type"));
                }
            }
        }
    }

    // =========================================================================
    // Convenience methods for common pull patterns
    // =========================================================================

    /// Pull an OCI image with default options.
    ///
    /// Shorthand for `pull(image, PullOptions::new())`.
    pub fn pull_simple(&mut self, image: &str) -> Result<ImageInfo> {
        self.pull(image, PullOptions::new())
    }

    /// Pull an OCI image with automatic registry credential lookup.
    ///
    /// Loads credentials from `~/.config/smolvm/registries.toml` and applies
    /// registry mirrors if configured.
    ///
    /// Shorthand for `pull(image, PullOptions::new().use_registry_config(true))`.
    pub fn pull_with_registry_config(&mut self, image: &str) -> Result<ImageInfo> {
        self.pull(image, PullOptions::new().use_registry_config(true))
    }

    /// Pull an OCI image with registry config and progress callback.
    pub fn pull_with_registry_config_and_progress<F: FnMut(usize, usize, &str)>(
        &mut self,
        image: &str,
        oci_platform: Option<&str>,
        proxy: Option<&str>,
        no_proxy: Option<&str>,
        progress: F,
    ) -> Result<ImageInfo> {
        let mut opts = PullOptions::new()
            .use_registry_config(true)
            .progress(progress);
        if let Some(p) = oci_platform {
            opts = opts.oci_platform(p);
        }
        if let Some(p) = proxy {
            opts = opts.proxy(p);
        }
        if let Some(np) = no_proxy {
            opts = opts.no_proxy(np);
        }
        self.pull(image, opts)
    }

    /// Query if an image exists locally.
    pub fn query(&mut self, image: &str) -> Result<Option<ImageInfo>> {
        let resp = self.request(&AgentRequest::Query {
            image: image.to_string(),
        })?;

        match resp {
            AgentResponse::Ok { data: Some(data) } => {
                let info: ImageInfo = serde_json::from_value(data)
                    .map_err(|e| Error::agent("parse response", e.to_string()))?;
                Ok(Some(info))
            }
            AgentResponse::Error { code, .. } if code.as_deref() == Some("NOT_FOUND") => Ok(None),
            AgentResponse::Error { message, .. } => Err(Error::agent("query image", message)),
            _ => Err(Error::agent("query image", "unexpected response type")),
        }
    }

    /// List all cached images.
    pub fn list_images(&mut self) -> Result<Vec<ImageInfo>> {
        let resp = self.request(&AgentRequest::ListImages)?;
        expect_data(resp, "list images")
    }

    /// Run garbage collection.
    ///
    /// # Arguments
    ///
    /// * `dry_run` - If true, only report what would be deleted
    /// * `purge_all` - If true, delete all manifests/configs first so all layers are collected
    pub fn garbage_collect(&mut self, dry_run: bool, purge_all: bool) -> Result<u64> {
        let resp = self.request(&AgentRequest::GarbageCollect { dry_run, purge_all })?;

        match resp {
            AgentResponse::Ok { data: Some(data) } => {
                let freed = data["freed_bytes"].as_u64().unwrap_or(0);
                Ok(freed)
            }
            AgentResponse::Error { message, .. } => Err(Error::agent("garbage collect", message)),
            _ => Err(Error::agent("garbage collect", "unexpected response type")),
        }
    }

    /// Prepare an overlay filesystem for a workload.
    ///
    /// # Arguments
    ///
    /// * `image` - Image reference
    /// * `workload_id` - Unique workload identifier
    pub fn prepare_overlay(&mut self, image: &str, workload_id: &str) -> Result<OverlayInfo> {
        let resp = self.request(&AgentRequest::PrepareOverlay {
            image: image.to_string(),
            workload_id: workload_id.to_string(),
        })?;
        expect_data(resp, "prepare overlay")
    }

    /// Clean up an overlay filesystem.
    pub fn cleanup_overlay(&mut self, workload_id: &str) -> Result<()> {
        let resp = self.request(&AgentRequest::CleanupOverlay {
            workload_id: workload_id.to_string(),
        })?;
        expect_ok(resp, "cleanup overlay")
    }

    /// Format the storage disk.
    pub fn format_storage(&mut self) -> Result<()> {
        let resp = self.request(&AgentRequest::FormatStorage)?;
        expect_ok(resp, "format storage")
    }

    /// Merge `lowerdirs` (bottom -> top) into a single tar at `output` in the guest.
    ///
    /// Missing or empty entries are dropped guest-side, so callers can append a
    /// container overlay's upper dir without probing it first.
    pub fn flatten_layers(&mut self, lowerdirs: &[String], output: &str) -> Result<()> {
        let resp = self.request(&AgentRequest::FlattenLayers {
            lowerdirs: lowerdirs.to_vec(),
            output: output.to_string(),
        })?;
        expect_ok(resp, "flatten layers")
    }

    /// Get storage status.
    pub fn storage_status(&mut self) -> Result<StorageStatus> {
        let resp = self.request(&AgentRequest::StorageStatus)?;
        expect_data(resp, "storage status")
    }

    /// Test network connectivity directly from the agent (not via chroot).
    /// Used to debug TSI networking.
    pub fn network_test(&mut self, url: &str) -> Result<serde_json::Value> {
        let resp = self.request(&AgentRequest::NetworkTest {
            url: url.to_string(),
        })?;

        match resp {
            AgentResponse::Ok { data: Some(data) } => Ok(data),
            AgentResponse::Error { message, .. } => Err(Error::agent("network test", message)),
            _ => Err(Error::agent("network test", "unexpected response type")),
        }
    }

    /// Request agent shutdown.
    ///
    /// Waits for the agent to acknowledge the shutdown request before returning.
    /// This ensures the agent has called sync() to flush filesystem caches
    /// before we send SIGTERM to terminate the VM.
    ///
    /// The acknowledgment is critical for data integrity - without it, the VM
    /// may be killed before ext4 journal commits are flushed, causing layer
    /// corruption on next boot.
    pub fn shutdown(&mut self) -> Result<()> {
        // Set a timeout for shutdown acknowledgment.
        // The agent calls sync() then sends the ack — typically <100ms,
        // but heavy writes or large journals may take longer.
        // If no ack within 5s, the VM has likely already torn down.
        let _ = self
            .stream
            .set_read_timeout(Some(Duration::from_secs(SHUTDOWN_ACK_TIMEOUT_SECS)));

        let data = self.encode_traced(&AgentRequest::Shutdown)?;
        self.stream
            .write_all(&data)
            .map_err(|e| Error::agent("send shutdown", e.to_string()))?;

        // Wait for acknowledgment - this confirms sync() completed.
        // Returns Ok only when the ack is actually received, so callers
        // can distinguish "sync confirmed" from "sync unknown".
        match self.receive() {
            Ok(_) => {
                tracing::debug!("agent acknowledged shutdown (sync complete)");
                Ok(())
            }
            Err(e) => {
                let error_str = e.to_string();
                if is_benign_shutdown_error(&error_str) {
                    tracing::debug!(
                        "shutdown ack not received (connection closed) - sync may have completed"
                    );
                } else {
                    tracing::warn!(error = %e, "shutdown acknowledgment failed");
                }
                Err(Error::agent("shutdown ack", error_str))
            }
        }
    }

    // ========================================================================
    // VM-Level Exec (Direct Execution in VM)
    // ========================================================================

    /// Execute a command directly in the VM (not in a container).
    ///
    /// Runs the command in the agent's Alpine rootfs without container
    /// isolation. Returns `(exit_code, stdout_bytes, stderr_bytes)`. Output
    /// is raw bytes — binary data (image bytes, tarballs) is preserved.
    /// Callers that need a string can use `String::from_utf8_lossy(&bytes)`.
    pub fn vm_exec(
        &mut self,
        command: Vec<String>,
        env: Vec<(String, String)>,
        workdir: Option<String>,
        timeout: Option<Duration>,
        stdin_data: Option<String>,
    ) -> Result<(i32, Vec<u8>, Vec<u8>)> {
        let _timeout_guard = self.set_exec_timeout(timeout)?;
        let timeout_ms = timeout.map(|t| t.as_millis() as u64);

        let resp = self.request(&AgentRequest::VmExec {
            command,
            env,
            workdir,
            timeout_ms,
            interactive: false,
            tty: false,
            background: false,
            stdin_data,
        })?;

        expect_completed(resp, "vm exec")
    }

    /// Execute a command in the background inside the VM.
    ///
    /// Spawns the process and returns immediately with the PID.
    /// The process runs detached — stdout/stderr go to /dev/null.
    pub fn vm_exec_background(
        &mut self,
        command: Vec<String>,
        env: Vec<(String, String)>,
        workdir: Option<String>,
    ) -> Result<u32> {
        let resp = self.request(&AgentRequest::VmExec {
            command,
            env,
            workdir,
            timeout_ms: None,
            interactive: false,
            tty: false,
            background: true,
            stdin_data: None,
        })?;

        let (exit_code, stdout, _stderr) = expect_completed(resp, "vm exec background")?;
        if exit_code != 0 {
            return Err(Error::agent("vm exec background", "spawn failed"));
        }
        // PID output is always ASCII digits — lossy conversion is safe.
        let pid: u32 = String::from_utf8_lossy(&stdout)
            .trim()
            .parse()
            .map_err(|_| Error::agent("vm exec background", "invalid PID in response"))?;
        Ok(pid)
    }

    /// Raw descriptor of the underlying agent socket, in the portable form the
    /// terminal poll loop expects. Unix: the socket fd; Windows: the underlying
    /// WinSock `SOCKET` cast into the portable `Fd`.
    fn stream_raw_fd(&self) -> crate::agent::terminal::Fd {
        #[cfg(unix)]
        {
            use std::os::unix::io::AsRawFd;
            self.stream.as_raw_fd()
        }
        #[cfg(not(unix))]
        {
            self.stream.raw_socket() as crate::agent::terminal::Fd
        }
    }

    /// Run an interactive I/O session.
    ///
    /// Sends `request`, waits for `Started`, then runs the poll loop
    /// streaming stdout/stderr and forwarding stdin until `Exited`.
    fn interactive_session(&mut self, request: AgentRequest, tty: bool, op: &str) -> Result<i32> {
        use crate::agent::terminal::{
            check_sigwinch, flush_retry, get_terminal_size, install_sigwinch_handler, poll_io,
            stdin_is_tty, write_all_retry, NonBlockingStdin, RawModeGuard,
        };
        use std::io::{stderr, stdin, stdout, Read};

        // Disable socket read timeout for interactive sessions — the poll loop
        // handles readiness checking, and the session runs until the user exits.
        self.stream
            .set_read_timeout(None)
            .map_err(|e| Error::agent("set read timeout", e.to_string()))?;

        self.send(&request)?;

        // Wait for Started response
        let started = self.receive()?;
        match started {
            AgentResponse::Started => {}
            AgentResponse::Error { message, .. } => {
                return Err(Error::agent(op, message));
            }
            _ => {
                return Err(Error::agent(op, "expected Started response"));
            }
        }

        // Enable raw mode if TTY requested and stdin is a TTY
        // The guard will restore terminal settings on drop (even on panic)
        let _raw_mode = if tty && stdin_is_tty() {
            RawModeGuard::new(stdin_raw_fd())
        } else {
            None
        };

        // Route session writes through the dedicated writer.
        let writer = FrameWriter::spawn(&self.stream)?;
        let mut pending: std::collections::VecDeque<Vec<u8>> = std::collections::VecDeque::new();

        // Send initial terminal size so PTY starts at the right dimensions
        if tty {
            if let Some((cols, rows)) = get_terminal_size() {
                pending.push_back(self.encode_traced(&AgentRequest::Resize { cols, rows })?);
            }
            install_sigwinch_handler();
        }

        // Set stdin to non-blocking (guard restores on drop)
        let _nonblock_stdin = NonBlockingStdin::new()
            .map_err(|e| Error::agent("set stdin nonblocking", e.to_string()))?;

        // Socket reads remain blocking; poll() determines read readiness.
        // Outbound frames are handled by FrameWriter.
        let mut stdin_handle = stdin();
        let stdin_fd = stdin_raw_fd();
        let socket_fd = self.stream_raw_fd();
        let mut stdin_buf = [0u8; STDIN_BUF_SIZE];
        let mut stdin_eof = false;

        let exit_code = loop {
            if let Some(e) = writer.take_error() {
                return Err(Error::agent(op, format!("failed to send to agent: {}", e)));
            }
            pump_frames(&writer, &mut pending)?;

            // Stop reading stdin while the writer is backed up, bounding memory use.
            let backpressured = !pending.is_empty();
            let effective_stdin_fd = if stdin_eof || backpressured {
                -1
            } else {
                stdin_fd
            };
            let poll_timeout = if backpressured {
                BACKPRESSURE_POLL_TIMEOUT_MS
            } else {
                POLL_TIMEOUT_MS
            };
            let poll_result = poll_io(effective_stdin_fd, socket_fd, poll_timeout)
                .map_err(|e| Error::agent("poll", e.to_string()))?;

            // Check for terminal resize (SIGWINCH)
            if tty && check_sigwinch() {
                if let Some((cols, rows)) = get_terminal_size() {
                    pending.push_back(self.encode_traced(&AgentRequest::Resize { cols, rows })?);
                }
            }

            // Handle socket data FIRST — drain agent output before writing stdin
            // to prevent deadlock when send buffer is full
            if poll_result.socket_ready {
                match self.receive() {
                    Ok(AgentResponse::Stdout { data }) => {
                        write_all_retry(&mut stdout(), &data)?;
                        flush_retry(&mut stdout())?;
                    }
                    Ok(AgentResponse::Stderr { data }) => {
                        write_all_retry(&mut stderr(), &data)?;
                        flush_retry(&mut stderr())?;
                    }
                    Ok(AgentResponse::Exited { exit_code, .. }) => {
                        break exit_code;
                    }
                    Ok(AgentResponse::Error { message, .. }) => {
                        return Err(Error::agent(op, message));
                    }
                    Ok(_) => {}
                    Err(e) => {
                        // EAGAIN/WouldBlock can occur when poll() reports readiness
                        // but the data isn't available yet (common with vsock on macOS).
                        // Retry on next poll iteration instead of crashing.
                        if e.is_io()
                            && matches!(
                                e.source_io_error_kind(),
                                Some(std::io::ErrorKind::WouldBlock)
                            )
                        {
                            tracing::debug!("socket read returned EAGAIN, retrying");
                            continue;
                        }
                        return Err(e);
                    }
                }
            }

            // Socket peer closed without sending Exited — VM crashed or was killed
            if poll_result.socket_hangup && !poll_result.socket_ready {
                return Err(Error::agent(op, "connection to VM lost".to_string()));
            }

            // Handle stdin input — queue it for the writer thread
            if poll_result.stdin_ready && !stdin_eof {
                match stdin_handle.read(&mut stdin_buf) {
                    Ok(0) => {
                        stdin_eof = true;
                        pending.push_back(
                            self.encode_traced(&AgentRequest::Stdin { data: Vec::new() })?,
                        );
                    }
                    Ok(n) => {
                        pending.push_back(self.encode_traced(&AgentRequest::Stdin {
                            data: stdin_buf[..n].to_vec(),
                        })?);
                    }
                    Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
                    Err(e) => {
                        tracing::debug!(error = %e, "stdin read error, treating as EOF");
                        stdin_eof = true;
                        pending.push_back(
                            self.encode_traced(&AgentRequest::Stdin { data: Vec::new() })?,
                        );
                    }
                }
            }
        };

        Ok(exit_code)
    }

    /// Execute a command directly in the VM with interactive I/O.
    pub fn vm_exec_interactive(
        &mut self,
        command: Vec<String>,
        env: Vec<(String, String)>,
        workdir: Option<String>,
        timeout: Option<Duration>,
        tty: bool,
    ) -> Result<i32> {
        let timeout_ms = timeout.map(|t| t.as_millis() as u64);
        let env = with_term_default(env, tty);
        self.interactive_session(
            AgentRequest::VmExec {
                command,
                env,
                workdir,
                timeout_ms,
                interactive: true,
                tty,
                background: false,
                stdin_data: None,
            },
            tty,
            "vm exec interactive",
        )
    }

    /// Run a command in an image's rootfs (non-interactive).
    ///
    /// This is the non-interactive counterpart to `run_interactive()`.
    /// Both accept a `RunConfig` for consistency.
    ///
    /// # Returns
    ///
    /// A tuple of (exit_code, stdout, stderr)
    pub fn run_non_interactive(&mut self, config: RunConfig) -> Result<(i32, Vec<u8>, Vec<u8>)> {
        let _timeout_guard = self.set_exec_timeout(config.timeout)?;
        let timeout_ms = config.timeout.map(|t| t.as_millis() as u64);

        let resp = self.request(&AgentRequest::Run {
            image: config.image,
            command: config.command,
            env: config.env,
            workdir: config.workdir,
            user: config.user,
            mounts: config.mounts,
            timeout_ms,
            interactive: false,
            tty: false,
            detached: false,
            unprivileged: config.unprivileged,
            persistent_overlay_id: config.persistent_overlay_id,
            stdin_data: config.stdin,
            background: false,
        })?;

        expect_completed(resp, "run command")
    }

    /// Run a command in an image's rootfs in the background.
    ///
    /// Spawns the container and returns immediately with the crun PID.
    /// stdout/stderr go to /dev/null inside the guest. Use a persistent
    /// overlay ID so subsequent `exec` sessions see the same filesystem.
    pub fn run_background(&mut self, config: RunConfig) -> Result<u32> {
        let resp = self.request(&AgentRequest::Run {
            image: config.image,
            command: config.command,
            env: config.env,
            workdir: config.workdir,
            user: config.user,
            mounts: config.mounts,
            timeout_ms: None,
            interactive: false,
            tty: false,
            detached: false,
            unprivileged: config.unprivileged,
            persistent_overlay_id: config.persistent_overlay_id,
            stdin_data: None,
            background: true,
        })?;

        let (exit_code, stdout, _stderr) = expect_completed(resp, "run background")?;
        if exit_code != 0 {
            return Err(Error::agent("run background", "spawn failed"));
        }
        let pid: u32 = String::from_utf8_lossy(&stdout)
            .trim()
            .parse()
            .map_err(|_| Error::agent("run background", "invalid PID in response"))?;
        Ok(pid)
    }

    /// Run a command in an image's rootfs and handle streamed events as they arrive.
    ///
    /// Unlike `run_interactive`, this does not forward stdin. It is the
    /// image-backed counterpart to `vm_exec_streaming_with`.
    pub fn run_streaming_with<F>(&mut self, config: RunConfig, on_event: F) -> Result<()>
    where
        F: FnMut(ExecEvent),
    {
        let timeout_ms = config.timeout.map(|t| t.as_millis() as u64);
        // Keep a host-side deadline whenever the caller requested one. The
        // guest normally enforces the command timeout, but a wedged restored
        // VM cannot send its terminal event and must not leave the CLI blocked
        // forever on an otherwise-open vsock connection.
        let _timeout_guard = self.set_exec_timeout(config.timeout)?;

        self.send(&AgentRequest::Run {
            image: config.image,
            command: config.command,
            env: config.env,
            workdir: config.workdir,
            user: config.user,
            mounts: config.mounts,
            timeout_ms,
            interactive: true,
            tty: false,
            detached: false,
            unprivileged: config.unprivileged,
            persistent_overlay_id: config.persistent_overlay_id,
            stdin_data: None,
            background: false,
        })?;

        collect_exec_events(self, "run streaming", on_event)
    }

    /// Run a command interactively with streaming I/O.
    ///
    /// This method streams output directly to stdout/stderr and forwards stdin.
    /// It blocks until the command exits.
    ///
    /// # Arguments
    ///
    /// * `config` - Run configuration including image, command, environment, etc.
    ///
    /// # Returns
    ///
    /// The exit code of the command
    pub fn run_interactive(&mut self, config: RunConfig) -> Result<i32> {
        let timeout_ms = config.timeout.map(|t| t.as_millis() as u64);
        let tty = config.tty;
        let env = with_term_default(config.env, tty);
        self.interactive_session(
            AgentRequest::Run {
                image: config.image,
                command: config.command,
                env,
                workdir: config.workdir,
                user: config.user,
                mounts: config.mounts,
                timeout_ms,
                interactive: true,
                tty,
                detached: false,
                unprivileged: config.unprivileged,
                persistent_overlay_id: config.persistent_overlay_id,
                stdin_data: None,
                background: false,
            },
            tty,
            "run interactive",
        )
    }

    /// Start a container in detached mode and return its container ID.
    ///
    /// Sends a `Run { detached: true }` request to the agent, which starts the
    /// container in the background via `crun run --detach` and immediately
    /// returns the container ID. Subsequent `machine exec` calls against the
    /// same `persistent_overlay_id` will join this container's namespaces via
    /// `crun exec` instead of creating a new isolated container.
    ///
    /// Requires `config.persistent_overlay_id` to be set — detached containers
    /// only make sense when there is a persistent overlay to associate with.
    pub fn run_container_detached(&mut self, config: RunConfig) -> Result<String> {
        // Container startup involves overlay setup, a one-time flatten of a local
        // image archive into guest storage, and crun init, which can far exceed
        // the default 30s read timeout on first run (cold overlay, cold image).
        let _timeout_guard =
            self.set_extended_read_timeout(Duration::from_secs(DETACHED_START_TIMEOUT_SECS))?;

        self.send(&AgentRequest::Run {
            image: config.image,
            command: config.command,
            env: config.env,
            workdir: config.workdir,
            user: config.user,
            mounts: config.mounts,
            timeout_ms: None,
            interactive: false,
            tty: false,
            detached: true,
            unprivileged: config.unprivileged,
            persistent_overlay_id: config.persistent_overlay_id,
            stdin_data: None,
            background: false,
        })?;
        let resp = loop {
            match self.receive()? {
                AgentResponse::Progress { message, .. } => {
                    tracing::info!(message = %message, "detached start progress");
                }
                terminal @ (AgentResponse::Completed { .. } | AgentResponse::Error { .. }) => {
                    break terminal;
                }
                _ => {
                    return Err(Error::agent(
                        "run container detached",
                        "unexpected response type",
                    ));
                }
            }
        };
        let (exit_code, stdout, _) = expect_completed(resp, "run container detached")?;
        if exit_code != 0 {
            return Err(Error::agent(
                "run container detached",
                format!("agent returned exit code {}", exit_code),
            ));
        }
        String::from_utf8(stdout).map_err(|e| {
            Error::agent(
                "run container detached",
                format!("invalid container ID in response: {}", e),
            )
        })
    }

    /// Send stdin data to a running interactive command.
    pub fn send_stdin(&mut self, data: &[u8]) -> Result<()> {
        self.send(&AgentRequest::Stdin {
            data: data.to_vec(),
        })
    }

    /// Send a window resize event to a running interactive command.
    pub fn send_resize(&mut self, cols: u16, rows: u16) -> Result<()> {
        self.send(&AgentRequest::Resize { cols, rows })
    }

    /// Run an interactive session driven by channels instead of the process's
    /// real stdin/stdout. Input events arrive on `input`; output is delivered to
    /// `on_output`. This is the transport-agnostic counterpart to
    /// [`Self::interactive_session`] — used to bridge a VM PTY to a remote
    /// WebSocket terminal without touching the host's terminal.
    ///
    /// The loop polls only the vsock socket (input comes from the channel, not an
    /// fd) and drains pending input each iteration. When `input` disconnects
    /// (the remote peer hung up) it sends EOF once and keeps running until the
    /// command exits — a shell reading its PTY exits on EOF.
    fn interactive_session_io<F>(
        &mut self,
        request: AgentRequest,
        input: std::sync::mpsc::Receiver<InteractiveInput>,
        mut on_output: F,
        op: &str,
    ) -> Result<i32>
    where
        F: FnMut(InteractiveOutput),
    {
        use crate::agent::terminal::poll_io;

        // No socket read timeout — the poll loop handles readiness and the
        // session runs until the command exits or the peer hangs up.
        self.stream
            .set_read_timeout(None)
            .map_err(|e| Error::agent("set read timeout", e.to_string()))?;

        self.send(&request)?;
        match self.receive()? {
            AgentResponse::Started => {}
            AgentResponse::Error { message, .. } => return Err(Error::agent(op, message)),
            _ => return Err(Error::agent(op, "expected Started response")),
        }

        let socket_fd = self.stream_raw_fd();
        let mut input_eof_sent = false;

        // Route channel input through the same writer as `interactive_session`.
        let writer = FrameWriter::spawn(&self.stream)?;
        let mut pending: std::collections::VecDeque<Vec<u8>> = std::collections::VecDeque::new();

        let exit_code = loop {
            if let Some(e) = writer.take_error() {
                return Err(Error::agent(op, format!("failed to send to agent: {}", e)));
            }
            pump_frames(&writer, &mut pending)?;

            // stdin_fd = -1 → poll() ignores it; only the socket drives readiness.
            let poll_timeout = if pending.is_empty() {
                POLL_TIMEOUT_MS
            } else {
                BACKPRESSURE_POLL_TIMEOUT_MS
            };
            let poll_result = poll_io(-1, socket_fd, poll_timeout)
                .map_err(|e| Error::agent("poll", e.to_string()))?;

            // Drain agent output first (prevents deadlock when its send buffer fills).
            if poll_result.socket_ready {
                match self.receive() {
                    Ok(AgentResponse::Stdout { data }) => {
                        on_output(InteractiveOutput::Stdout(data))
                    }
                    Ok(AgentResponse::Stderr { data }) => {
                        on_output(InteractiveOutput::Stderr(data))
                    }
                    Ok(AgentResponse::Exited { exit_code, .. }) => break exit_code,
                    Ok(AgentResponse::Error { message, .. }) => {
                        return Err(Error::agent(op, message))
                    }
                    Ok(_) => {}
                    Err(e) => {
                        if e.is_io()
                            && matches!(
                                e.source_io_error_kind(),
                                Some(std::io::ErrorKind::WouldBlock)
                            )
                        {
                            continue;
                        }
                        return Err(e);
                    }
                }
            }

            if poll_result.socket_hangup && !poll_result.socket_ready {
                return Err(Error::agent(op, "connection to VM lost".to_string()));
            }

            // Forward any pending input without blocking the output path. While
            // the writer queue is backed up the events stay in their source
            // channel, so backpressure composes instead of buffering here.
            while pending.is_empty() {
                match input.try_recv() {
                    Ok(InteractiveInput::Stdin(data)) => {
                        pending.push_back(self.encode_traced(&AgentRequest::Stdin { data })?)
                    }
                    Ok(InteractiveInput::Resize { cols, rows }) => {
                        pending.push_back(self.encode_traced(&AgentRequest::Resize { cols, rows })?)
                    }
                    Ok(InteractiveInput::Eof) => {
                        if !input_eof_sent {
                            pending.push_back(
                                self.encode_traced(&AgentRequest::Stdin { data: Vec::new() })?,
                            );
                            input_eof_sent = true;
                        }
                    }
                    Err(std::sync::mpsc::TryRecvError::Empty) => break,
                    Err(std::sync::mpsc::TryRecvError::Disconnected) => {
                        // Remote peer (WebSocket client) gone. Return immediately
                        // instead of waiting for the command to exit on its own.
                        // This method runs on a DEDICATED, disposable connection,
                        // so returning drops it; the agent's interactive loop then
                        // sees the closed peer and kills the PTY child. Waiting here
                        // would pin the connection (and, on the shared client, the
                        // per-machine lock) until a command that ignores stdin EOF
                        // — a `sleep`, a daemon — finally exits.
                        return Ok(DISCONNECT_EXIT_CODE);
                    }
                }
                // Hand the event straight on, so a writer that is keeping up
                // leaves `pending` empty and this keeps draining the channel.
                pump_frames(&writer, &mut pending)?;
            }
        };

        Ok(exit_code)
    }

    /// Interactive VM exec driven by channels (remote PTY). Counterpart to
    /// [`Self::vm_exec_interactive`] that does not bind the host terminal.
    pub fn vm_exec_interactive_io<F>(
        &mut self,
        command: Vec<String>,
        env: Vec<(String, String)>,
        workdir: Option<String>,
        tty: bool,
        input: std::sync::mpsc::Receiver<InteractiveInput>,
        on_output: F,
    ) -> Result<i32>
    where
        F: FnMut(InteractiveOutput),
    {
        let env = with_term_default(env, tty);
        self.interactive_session_io(
            AgentRequest::VmExec {
                command,
                env,
                workdir,
                timeout_ms: None,
                interactive: true,
                tty,
                background: false,
                stdin_data: None,
            },
            input,
            on_output,
            "vm exec interactive (io)",
        )
    }

    /// Interactive container run driven by channels (remote PTY). Counterpart to
    /// [`Self::run_interactive`] that does not bind the host terminal.
    pub fn run_interactive_io<F>(
        &mut self,
        config: RunConfig,
        input: std::sync::mpsc::Receiver<InteractiveInput>,
        on_output: F,
    ) -> Result<i32>
    where
        F: FnMut(InteractiveOutput),
    {
        let tty = config.tty;
        let env = with_term_default(config.env, tty);
        self.interactive_session_io(
            AgentRequest::Run {
                image: config.image,
                command: config.command,
                env,
                workdir: config.workdir,
                user: config.user,
                mounts: config.mounts,
                timeout_ms: None,
                interactive: true,
                tty,
                detached: false,
                unprivileged: config.unprivileged,
                persistent_overlay_id: config.persistent_overlay_id,
                stdin_data: None,
                background: false,
            },
            input,
            on_output,
            "run interactive (io)",
        )
    }

    // ========================================================================
    // File I/O
    // ========================================================================

    /// Write a file into the VM.
    ///
    /// Transparently dispatches between single-shot and streaming
    /// based on `data.len()`:
    ///
    /// - Files ≤ [`FILE_WRITE_SINGLE_SHOT_MAX`] (1 MiB): one
    ///   [`AgentRequest::FileWrite`] message — the lowest-latency
    ///   path and what 99% of `cp` calls hit.
    /// - Files larger than that: a sequence of
    ///   [`AgentRequest::FileWriteBegin`] +
    ///   [`AgentRequest::FileWriteChunk`] messages, each under
    ///   [`MAX_FRAME_SIZE`]. This is the only correct way to upload
    ///   files whose base64-encoded form would exceed the frame
    ///   limit — without it the send blocks the socket (EAGAIN
    ///   after write timeout) and risks OOMing the guest agent.
    pub fn write_file(&mut self, path: &str, data: &[u8], mode: Option<u32>) -> Result<()> {
        self.write_file_with_progress(path, data, mode, |_| {})
    }

    /// Write a file into the VM with a progress callback.
    ///
    /// `on_progress` is called after each chunk is acked by the
    /// agent, with the running byte total. Single-shot writes (small
    /// files) call it once at the end. Callers who don't need
    /// progress should use [`Self::write_file`] which passes a no-op.
    pub fn write_file_with_progress<F: FnMut(u64)>(
        &mut self,
        path: &str,
        data: &[u8],
        mode: Option<u32>,
        mut on_progress: F,
    ) -> Result<()> {
        if data.len() <= FILE_WRITE_SINGLE_SHOT_MAX {
            let resp = self.request(&AgentRequest::FileWrite {
                path: path.to_string(),
                data: data.to_vec(),
                mode,
            })?;
            expect_ok(resp, "write file")?;
            on_progress(data.len() as u64);
            Ok(())
        } else {
            self.write_file_streaming(path, data, mode, &mut on_progress)
        }
    }

    /// Streaming file upload from a `&[u8]` slice.
    fn write_file_streaming<F: FnMut(u64)>(
        &mut self,
        path: &str,
        data: &[u8],
        mode: Option<u32>,
        on_progress: &mut F,
    ) -> Result<()> {
        self.write_file_streaming_from_reader(
            path,
            &mut std::io::Cursor::new(data),
            data.len() as u64,
            mode,
            on_progress,
        )
    }

    /// Stream a file from a [`Read`] source into the VM.
    ///
    /// Reads `FILE_WRITE_CHUNK_SIZE` bytes at a time from `reader`,
    /// sending each chunk over the protocol. Only one chunk is in
    /// memory at a time — the caller doesn't need to buffer the
    /// entire file.
    pub fn write_file_from_reader<R: std::io::Read>(
        &mut self,
        path: &str,
        reader: R,
        total_size: u64,
        mode: Option<u32>,
    ) -> Result<()> {
        self.write_file_from_reader_with_progress(path, reader, total_size, mode, |_| {})
    }

    /// Stream a file from a [`Read`] source with progress callback.
    pub fn write_file_from_reader_with_progress<R: std::io::Read, F: FnMut(u64)>(
        &mut self,
        path: &str,
        reader: R,
        total_size: u64,
        mode: Option<u32>,
        mut on_progress: F,
    ) -> Result<()> {
        if total_size <= FILE_WRITE_SINGLE_SHOT_MAX as u64 {
            // Small file: read into memory and use single-shot path.
            let mut data = Vec::with_capacity(total_size as usize);
            std::io::Read::read_to_end(&mut std::io::Read::take(reader, total_size), &mut data)
                .map_err(|e| Error::agent("read source file", e.to_string()))?;
            return self.write_file_with_progress(path, &data, mode, on_progress);
        }
        self.write_file_streaming_from_reader(
            path,
            &mut { reader },
            total_size,
            mode,
            &mut on_progress,
        )
    }

    /// Core streaming upload loop. Reads chunks from `reader` and
    /// sends them over the protocol. Only one chunk buffer is live
    /// at a time (~1 MiB).
    fn write_file_streaming_from_reader<R: std::io::Read, F: FnMut(u64)>(
        &mut self,
        path: &str,
        reader: &mut R,
        total_size: u64,
        mode: Option<u32>,
        on_progress: &mut F,
    ) -> Result<()> {
        let resp = self.request(&AgentRequest::FileWriteBegin {
            path: path.to_string(),
            mode,
            total_size,
        })?;
        expect_ok(resp, "begin streaming write")?;

        let mut buf = vec![0u8; FILE_WRITE_CHUNK_SIZE];
        let mut bytes_sent = 0u64;

        loop {
            // Fill the chunk buffer.
            let mut filled = 0;
            while filled < buf.len() {
                match reader.read(&mut buf[filled..]) {
                    Ok(0) => break,
                    Ok(n) => filled += n,
                    Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
                    Err(e) => return Err(Error::agent("read source file", e.to_string())),
                }
            }

            if filled == 0 {
                // EOF — send final empty chunk to finalize.
                let resp = self.request(&AgentRequest::FileWriteChunk {
                    data: Vec::new(),
                    done: true,
                })?;
                expect_ok(resp, "finalize streaming write")?;
                break;
            }

            bytes_sent += filled as u64;
            let done = bytes_sent >= total_size;

            let resp = self.request(&AgentRequest::FileWriteChunk {
                data: buf[..filled].to_vec(),
                done,
            })?;
            expect_ok(resp, "stream write chunk")?;
            on_progress(bytes_sent);

            if done {
                break;
            }
        }
        Ok(())
    }

    /// Read a file from the VM.
    ///
    /// Consumes the streamed `DataChunk` responses the agent emits
    /// (see `handle_streaming_file_read` in the agent). The agent
    /// sends one or more chunks, with `done: true` on the final
    /// frame — possibly empty. This method concatenates chunks and
    /// returns the full contents.
    ///
    /// Two safety bounds:
    /// - Receive timeout extended to 600 s so large files don't
    ///   spuriously fail on slow storage; a 200 MB file at 10 MB/s
    ///   would exceed the default 30 s receive timeout otherwise.
    /// - Total size capped at [`FILE_TRANSFER_MAX_TOTAL`] (4 GiB) —
    ///   symmetric with the write path. A misbehaving or compromised
    ///   guest can't OOM the host by streaming unbounded data.
    pub fn read_file(&mut self, path: &str) -> Result<Vec<u8>> {
        self.read_file_with_progress(path, |_| {})
    }

    /// Read a file from the VM with a progress callback.
    ///
    /// `on_progress` is called with the running byte total after
    /// each `DataChunk` is received. Use [`Self::read_file`] if you
    /// don't need progress.
    pub fn read_file_with_progress<F: FnMut(u64)>(
        &mut self,
        path: &str,
        on_progress: F,
    ) -> Result<Vec<u8>> {
        const FILE_READ_TIMEOUT: Duration = Duration::from_secs(600);

        let _timeout_guard = self.set_extended_read_timeout(FILE_READ_TIMEOUT)?;
        self.send_raw(&AgentRequest::FileRead {
            path: path.to_string(),
        })?;

        consume_streamed_read_with_progress(|| self.recv_raw(), on_progress)
    }

    /// Download a file from the VM directly to a local path.
    ///
    /// Unlike [`Self::read_file`] which accumulates the entire file
    /// in memory, this writes each chunk to disk as it arrives —
    /// only one 16 MiB chunk is in memory at a time.
    pub fn read_file_to_path<F: FnMut(u64)>(
        &mut self,
        guest_path: &str,
        local_path: &std::path::Path,
        on_progress: F,
    ) -> Result<u64> {
        self.read_file_to_path_capped(
            guest_path,
            local_path,
            file_transfer_max_total(),
            on_progress,
        )
    }

    /// [`Self::read_file_to_path`] with an explicit byte ceiling.
    ///
    /// The pack paths read back a layer smolvm itself asked the guest to
    /// produce, so they pass [`pack_export_max_total`] instead of the general
    /// guest-driven transfer bound.
    pub fn read_file_to_path_capped<F: FnMut(u64)>(
        &mut self,
        guest_path: &str,
        local_path: &std::path::Path,
        cap: u64,
        mut on_progress: F,
    ) -> Result<u64> {
        use std::io::Write;
        const FILE_READ_TIMEOUT: Duration = Duration::from_secs(600);

        let _timeout_guard = self.set_extended_read_timeout(FILE_READ_TIMEOUT)?;
        self.send_raw(&AgentRequest::FileRead {
            path: guest_path.to_string(),
        })?;

        let mut file = std::fs::File::create(local_path).map_err(|e| {
            Error::agent(
                "write local file",
                format!("{}: {}", local_path.display(), e),
            )
        })?;

        let mut total = 0u64;
        loop {
            match self.recv_raw()? {
                AgentResponse::DataChunk { data, done } => {
                    let next_total = total.saturating_add(data.len() as u64);
                    if next_total > cap {
                        let _ = std::fs::remove_file(local_path);
                        return Err(Error::agent(
                            "read file",
                            format!(
                                "guest streamed {} bytes, exceeding the {} byte cap",
                                next_total, cap
                            ),
                        ));
                    }
                    if !data.is_empty() {
                        file.write_all(&data)
                            .map_err(|e| Error::agent("write local file", e.to_string()))?;
                        total = next_total;
                        on_progress(total);
                    }
                    if done {
                        file.flush()
                            .map_err(|e| Error::agent("flush local file", e.to_string()))?;
                        return Ok(total);
                    }
                }
                AgentResponse::Error { message, .. } => {
                    let _ = std::fs::remove_file(local_path);
                    return Err(Error::agent("read file", message));
                }
                _ => {
                    let _ = std::fs::remove_file(local_path);
                    return Err(Error::agent("read file", "unexpected response"));
                }
            }
        }
    }

    // ========================================================================
    // Streaming Exec
    // ========================================================================

    /// Execute a command with streaming output.
    ///
    /// This compatibility wrapper buffers all events before returning. New
    /// callers that need live output should use `vm_exec_streaming_with`.
    ///
    /// Sends a VmExec request with interactive=true, tty=false. Reads
    /// Stdout/Stderr/Exited responses into a vector and blocks until the
    /// command finishes — call from a blocking context (e.g.,
    /// `spawn_blocking`).
    pub fn vm_exec_streaming(
        &mut self,
        command: Vec<String>,
        env: Vec<(String, String)>,
        workdir: Option<String>,
        timeout: Option<Duration>,
    ) -> Result<Vec<ExecEvent>> {
        let mut events = Vec::new();
        self.vm_exec_streaming_with(command, env, workdir, timeout, |event| {
            events.push(event);
        })?;
        Ok(events)
    }

    /// Execute a command with streaming output and handle events as they arrive.
    ///
    /// This is the live-output variant of `vm_exec_streaming`.
    pub fn vm_exec_streaming_with<F>(
        &mut self,
        command: Vec<String>,
        env: Vec<(String, String)>,
        workdir: Option<String>,
        timeout: Option<Duration>,
        on_event: F,
    ) -> Result<()>
    where
        F: FnMut(ExecEvent),
    {
        let timeout_ms = timeout.map(|t| t.as_millis() as u64);
        let _timeout_guard = self.set_exec_timeout(timeout)?;

        self.send(&AgentRequest::VmExec {
            command,
            env,
            workdir,
            timeout_ms,
            interactive: true,
            tty: false,
            background: false,
            stdin_data: None,
        })?;

        collect_exec_events(self, "streaming exec", on_event)
    }

    /// Low-level send without waiting for response (public).
    pub fn send_raw(&mut self, request: &AgentRequest) -> Result<()> {
        self.send(request)
    }

    /// Low-level receive a single response (public).
    pub fn recv_raw(&mut self) -> Result<AgentResponse> {
        self.receive()
    }

    /// Set a command-execution timeout and return a guard that resets it on drop.
    ///
    /// If `timeout` is Some, the socket deadline is `timeout + TIMEOUT_BUFFER_SECS`.
    /// If None, the socket read timeout is disabled entirely — the command runs
    /// until completion (or the VM dies, triggering EOF). This matches
    /// `interactive_session`'s behavior and avoids any implicit ceiling on how
    /// long a non-interactive command can run. The `ReadTimeoutGuard` restores
    /// `DEFAULT_READ_TIMEOUT_SECS` on drop so subsequent operations get the
    /// normal 30-second timeout.
    fn set_exec_timeout(&self, timeout: Option<Duration>) -> Result<Option<ReadTimeoutGuard>> {
        match timeout {
            Some(t) => {
                self.set_read_timeout(t + Duration::from_secs(TIMEOUT_BUFFER_SECS))?;
            }
            None => {
                self.stream.set_read_timeout(None).map_err(|e| {
                    Error::agent(
                        "set read timeout",
                        format!("failed to clear socket read timeout: {}", e),
                    )
                })?;
            }
        }
        Ok(ReadTimeoutGuard::new(&self.stream))
    }

    /// Set an extended read timeout and return a guard that resets it on drop.
    ///
    /// Used for long-running streaming operations (e.g., layer export) where
    /// individual chunks may take longer than the default 30s timeout.
    pub fn set_extended_read_timeout(&self, timeout: Duration) -> Result<Option<ReadTimeoutGuard>> {
        self.set_read_timeout(timeout)?;
        Ok(ReadTimeoutGuard::new(&self.stream))
    }

    /// Low-level send without waiting for response.
    fn send(&mut self, request: &AgentRequest) -> Result<()> {
        let data = self.encode_traced(request)?;
        self.stream.write_all(&data)?;
        self.stream.flush()?;
        Ok(())
    }

    /// Read exactly `buf.len()` bytes, retrying on EAGAIN/WouldBlock.
    ///
    /// Unlike `read_exact`, this never loses partially-read data on EAGAIN.
    /// On macOS, vsock sockets can spuriously return WouldBlock even in
    /// blocking mode, so we must handle it without corrupting the stream.
    ///
    /// If `propagate_initial_wouldblock` is true and WouldBlock occurs before
    /// any bytes are read, the error is propagated (preserves read timeout
    /// behavior). Once any bytes are consumed, EAGAIN is retried.
    ///
    /// # Stall protection
    ///
    /// When the socket has a read timeout configured, the retry loop is bounded
    /// by a wall-clock *idle* deadline. Without it, a peer that writes a valid
    /// length prefix and then stalls mid-frame (fewer body bytes than declared,
    /// without closing) would spin this loop at 1ms forever, pinning the host
    /// thread and the per-machine client lock and defeating every client-level
    /// timeout. The idle deadline is reset on every byte of progress, so a
    /// slow-but-steady body is never penalized — only a body that delivers *no*
    /// bytes for a full read-timeout window fails, with a `TimedOut` error.
    ///
    /// When no read timeout is configured (interactive sessions), WouldBlock is
    /// treated as the spurious macOS vsock EAGAIN it is meant to be, and the loop
    /// retries indefinitely as before — there is no deadline to enforce.
    fn read_exact_retry(
        &mut self,
        buf: &mut [u8],
        propagate_initial_wouldblock: bool,
    ) -> std::io::Result<()> {
        // Idle window: how long we tolerate zero progress before declaring a
        // stall. Derived from the socket's configured read timeout so it tracks
        // the caller's intent; `None` means blocking mode (no deadline).
        let idle_window = self.stream.read_timeout().ok().flatten();
        let mut deadline = idle_window.map(|w| std::time::Instant::now() + w);

        let mut pos = 0;
        while pos < buf.len() {
            match self.stream.read(&mut buf[pos..]) {
                Ok(0) => {
                    return Err(std::io::Error::new(
                        std::io::ErrorKind::UnexpectedEof,
                        "connection closed",
                    ));
                }
                Ok(n) => {
                    pos += n;
                    // Progress made — extend the idle deadline.
                    if let Some(w) = idle_window {
                        deadline = Some(std::time::Instant::now() + w);
                    }
                }
                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                    if pos == 0 && propagate_initial_wouldblock {
                        // No data consumed yet and caller wants timeout errors — propagate
                        return Err(e);
                    }
                    // Mid-read (or caller wants full retry). Retry, but bail out
                    // if the idle deadline has passed so a stalled mid-frame body
                    // can't busy-spin forever.
                    if let Some(d) = deadline {
                        if std::time::Instant::now() >= d {
                            return Err(std::io::Error::new(
                                std::io::ErrorKind::TimedOut,
                                "timed out reading frame body: peer stalled mid-frame",
                            ));
                        }
                    }
                    std::thread::sleep(std::time::Duration::from_millis(1));
                    continue;
                }
                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }

    /// Low-level receive a single response.
    fn receive(&mut self) -> Result<AgentResponse> {
        // Check if a read timeout is set — if so, WouldBlock before any data
        // means a real timeout and should be propagated. If no timeout (interactive
        // sessions), WouldBlock is always a spurious macOS vsock EAGAIN.
        let has_timeout = self.stream.read_timeout().ok().flatten().is_some();

        let mut header = [0u8; 4];
        self.read_exact_retry(&mut header, has_timeout)?;
        let len = u32::from_be_bytes(header) as usize;

        // Validate frame size to prevent OOM from malicious/buggy responses
        if len > MAX_FRAME_SIZE as usize {
            // Header consumed but body not read — stream is desynchronized.
            // Shut down the read half so all future reads fail immediately
            // rather than interpreting body bytes as a frame header.
            let _ = self.stream.shutdown(std::net::Shutdown::Read);
            return Err(Error::agent(
                "validate frame",
                format!(
                    "frame too large: {} bytes (max: {} bytes)",
                    len, MAX_FRAME_SIZE
                ),
            ));
        }

        let mut buf = vec![0u8; len];
        // Always retry body reads — header is already consumed so we can't
        // propagate an error without corrupting the stream.
        if let Err(e) = self.read_exact_retry(&mut buf, false) {
            // Body read failed — stream is desynchronized. Shut down the
            // read half so future reads fail cleanly.
            let _ = self.stream.shutdown(std::net::Shutdown::Read);
            return Err(e.into());
        }

        let resp: AgentResponse = serde_json::from_slice(&buf)
            .map_err(|e| Error::agent("deserialize response", e.to_string()))?;
        Ok(resp)
    }
}

/// Cumulative byte ceiling for the streaming/collect exec path.
///
/// The buffered (`Completed`) exec path caps guest output inside the guest at
/// `smolvm_agent::process::MAX_EXEC_OUTPUT` (11 MiB). The streaming path,
/// however, relays frames one at a time and the SSE handler
/// (`api::handlers::exec::exec_stream`) buffers the *entire* event vector in
/// host RAM before responding. Without a cap, a chatty or infinite guest
/// command (`yes`, `cat /dev/zero`) that emits `Stdout` frames forever and
/// never sends `Exited` grows the host `serve` process without bound → host
/// OOM → every co-tenant VM on the node is killed (cross-tenant DoS).
///
/// We mirror the buffered path's 11 MiB ceiling: once cumulative stdout+stderr
/// crosses it, we emit a truncation error event and stop relaying (terminating
/// the collect loop) instead of growing unbounded.
///
/// Semantics: this is a per-exec-session cap on the BUFFERED/relayed output of
/// the streaming exec path. Interactive PTY sessions are long-lived by design
/// and do NOT flow through here — they use
/// [`AgentClient::interactive_session_io`], which streams frame-by-frame to the
/// WebSocket without accumulating — so a legitimate long interactive session is
/// unaffected by this cap.
const MAX_STREAMING_EXEC_OUTPUT: usize = 11 * 1024 * 1024;

fn collect_exec_events<F>(client: &mut AgentClient, op: &str, on_event: F) -> Result<()>
where
    F: FnMut(ExecEvent),
{
    collect_exec_events_inner(|| client.receive(), op, MAX_STREAMING_EXEC_OUTPUT, on_event)
}

/// Cap-parameterized core of [`collect_exec_events`], pulling responses from a
/// `next_response` closure so it can be unit-tested against synthetic frame
/// sequences with a small cap (instead of booting a VM and streaming 11 MiB).
fn collect_exec_events_inner<N, F>(
    mut next_response: N,
    op: &str,
    cap: usize,
    mut on_event: F,
) -> Result<()>
where
    N: FnMut() -> Result<AgentResponse>,
    F: FnMut(ExecEvent),
{
    match next_response()? {
        AgentResponse::Started => {}
        AgentResponse::Error { message, .. } => {
            return Err(Error::agent(op, message));
        }
        _ => return Err(Error::agent(op, "expected Started")),
    }

    // Cumulative stdout+stderr bytes relayed on this session. Bounds the host
    // buffer so a guest that never sends `Exited` can't OOM the host.
    let mut total: usize = 0;
    loop {
        match next_response() {
            Ok(AgentResponse::Stdout { data }) => {
                total = total.saturating_add(data.len());
                on_event(ExecEvent::Stdout(data));
            }
            Ok(AgentResponse::Stderr { data }) => {
                total = total.saturating_add(data.len());
                on_event(ExecEvent::Stderr(data));
            }
            Ok(AgentResponse::Exited { exit_code, .. }) => {
                on_event(ExecEvent::Exit(exit_code));
                break;
            }
            Ok(AgentResponse::Error { message, .. }) => {
                on_event(ExecEvent::Error(message));
                break;
            }
            Ok(_) => {}
            Err(err) => {
                on_event(ExecEvent::Error(err.to_string()));
                break;
            }
        }

        // Cap check runs after relaying each frame (matching the buffered
        // path's "send chunk, then break at the cap" order): we relay at most
        // `cap` + one frame before terminating with a truncation signal.
        if total >= cap {
            on_event(ExecEvent::Error(format!(
                "streaming output exceeded {cap} byte cap; exec terminated (output truncated)"
            )));
            break;
        }
    }
    Ok(())
}

/// Consume streamed `DataChunk` responses, enforcing the per-transfer
/// size cap and returning the concatenated bytes.
///
/// Pulled out of `AgentClient::read_file` so it can be unit-tested
/// against synthetic response sequences without booting a VM, and
/// against a small cap so the test doesn't have to allocate 4 GiB
/// just to exercise the cap branch.
///
/// Two small variants on the same loop:
/// - [`consume_streamed_read_with_progress`]: production cap, with
///   a progress callback (used by [`AgentClient::read_file`] +
///   [`AgentClient::read_file_with_progress`]).
/// - [`consume_streamed_read_with_cap`]: parameterized cap, no
///   progress (used by tests so they can exercise the cap branch
///   with kilobytes instead of gigabytes).
///
/// Both delegate to [`consume_streamed_read_inner`].
/// The file-transfer size cap, in bytes. Defaults to the protocol's
/// [`FILE_TRANSFER_MAX_TOTAL`] (4 GiB) and can be raised at runtime with
/// `SMOLVM_FILE_TRANSFER_MAX_BYTES` (accepts a plain byte count or a suffixed
/// size like `16GiB`). This lets an operator pack a VM snapshot whose overlay
/// carries a large dependency tree (e.g. a ~5 GiB torch + CUDA-wheels env)
/// without lowering the default DoS bound for everyone else.
pub fn file_transfer_max_total() -> u64 {
    std::env::var("SMOLVM_FILE_TRANSFER_MAX_BYTES")
        .ok()
        .and_then(|s| crate::util::parse_size_bytes(s.trim()).ok())
        .unwrap_or(FILE_TRANSFER_MAX_TOTAL)
}

/// Default ceiling for pack-export reads (64 GiB).
const PACK_EXPORT_MAX_TOTAL: u64 = 64 << 30;

/// The size cap for pack-export reads, in bytes. Defaults to 64 GiB and honors
/// the same `SMOLVM_FILE_TRANSFER_MAX_BYTES` override as
/// [`file_transfer_max_total`].
///
/// Pack export reads back a layer smolvm itself asked the guest to write, and a
/// provisioned machine's flattened rootfs is routinely tens of GiB, so the
/// general transfer bound — sized for `machine cp` against a guest that chooses
/// its own payload — is the wrong ceiling here. Every pack read shares this one
/// so the routes of a single export cannot disagree.
pub fn pack_export_max_total() -> u64 {
    std::env::var("SMOLVM_FILE_TRANSFER_MAX_BYTES")
        .ok()
        .and_then(|s| crate::util::parse_size_bytes(s.trim()).ok())
        .unwrap_or(PACK_EXPORT_MAX_TOTAL)
}

fn consume_streamed_read_with_progress<F, P>(next_response: F, on_progress: P) -> Result<Vec<u8>>
where
    F: FnMut() -> Result<AgentResponse>,
    P: FnMut(u64),
{
    consume_streamed_read_inner(next_response, file_transfer_max_total(), on_progress)
}

#[cfg(test)]
fn consume_streamed_read_with_cap<F>(next_response: F, cap: u64) -> Result<Vec<u8>>
where
    F: FnMut() -> Result<AgentResponse>,
{
    consume_streamed_read_inner(next_response, cap, |_| {})
}

fn consume_streamed_read_inner<F, P>(
    mut next_response: F,
    cap: u64,
    mut on_progress: P,
) -> Result<Vec<u8>>
where
    F: FnMut() -> Result<AgentResponse>,
    P: FnMut(u64),
{
    let mut out: Vec<u8> = Vec::new();
    let mut total: u64 = 0;
    loop {
        match next_response()? {
            AgentResponse::DataChunk { data, done } => {
                // Cap *before* extending so a single oversized chunk
                // can't push us past the limit.
                let next_total = total.saturating_add(data.len() as u64);
                if next_total > cap {
                    return Err(Error::agent(
                        "read file",
                        format!(
                            "guest streamed {} bytes, exceeding the {} byte cap; \
                             use a virtiofs mount for larger files",
                            next_total, cap
                        ),
                    ));
                }
                out.extend_from_slice(&data);
                total = next_total;
                on_progress(total);
                if done {
                    return Ok(out);
                }
            }
            AgentResponse::Error { message, .. } => {
                return Err(Error::agent("read file", message));
            }
            other => {
                return Err(Error::agent(
                    "read file",
                    format!("unexpected response: {:?}", other),
                ));
            }
        }
    }
}

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

    /// Build a `DataChunk` response with `n` zero bytes.
    fn chunk(n: usize, done: bool) -> AgentResponse {
        AgentResponse::DataChunk {
            data: vec![0u8; n],
            done,
        }
    }

    /// Drive the consumer over a fixed list of responses with a
    /// small (1 KiB) cap. Tests only need to exercise the size
    /// arithmetic; the production cap of 4 GiB would need the test
    /// to allocate a Vec that big to trip — wasteful and unnecessary.
    /// The cap is parameterized on the internal helper precisely so
    /// this test can scale down.
    const TEST_CAP: u64 = 1024;

    fn drive(responses: Vec<AgentResponse>) -> Result<Vec<u8>> {
        let mut iter = responses.into_iter();
        consume_streamed_read_with_cap(
            || {
                iter.next()
                    .ok_or_else(|| Error::agent("test", "no more responses"))
            },
            TEST_CAP,
        )
    }

    #[test]
    fn read_cap_terminator_returns_full_buffer() {
        let out = drive(vec![chunk(100, false), chunk(50, true)]).unwrap();
        assert_eq!(out.len(), 150);
    }

    // Pack export reads back a layer smolvm asked the guest to produce, not a
    // guest-chosen payload, so its ceiling must never be the more restrictive
    // of the two. Holds with or without the shared override, since both caps
    // read it.
    #[test]
    fn pack_export_cap_is_never_below_the_transfer_cap() {
        assert!(pack_export_max_total() >= file_transfer_max_total());
    }

    #[test]
    fn read_cap_empty_terminator_is_valid_eof() {
        let out = drive(vec![chunk(100, false), chunk(0, true)]).unwrap();
        assert_eq!(out.len(), 100);
    }

    #[test]
    fn read_cap_rejects_single_chunk_at_or_above_limit() {
        // Single chunk that on its own pushes past the cap. We're at
        // 1024-byte cap so this chunk is 1025 bytes — trivial to
        // allocate in tests, exercises the same arithmetic that
        // would catch a 4 GiB+ chunk in production.
        let err = drive(vec![chunk(TEST_CAP as usize + 1, true)]).unwrap_err();
        let msg = format!("{}", err);
        assert!(
            msg.contains("exceeding") && msg.contains("byte cap"),
            "expected size-cap error, got: {}",
            msg
        );
    }

    #[test]
    fn read_cap_rejects_when_accumulated_chunks_exceed_limit() {
        // Two chunks under the cap individually but exceeding it
        // combined. This is the realistic exhaustion vector — a
        // misbehaving guest streaming "fine-sized" chunks forever.
        let half = (TEST_CAP / 2) as usize;
        let err = drive(vec![
            chunk(half, false),
            chunk(half, false),
            chunk(half, false), // would push past the cap
        ])
        .unwrap_err();
        assert!(format!("{}", err).contains("byte cap"));
    }

    #[test]
    fn read_cap_chunk_at_exactly_limit_is_accepted() {
        // Boundary: a chunk that lands the accumulated total at
        // exactly the cap is fine. Only > cap is rejected.
        let out = drive(vec![chunk(TEST_CAP as usize, true)]).unwrap();
        assert_eq!(out.len(), TEST_CAP as usize);
    }

    #[test]
    fn read_cap_propagates_agent_error_response() {
        let err = drive(vec![AgentResponse::Error {
            message: "no such file".to_string(),
            code: None,
        }])
        .unwrap_err();
        assert!(format!("{}", err).contains("no such file"));
    }

    #[test]
    fn read_cap_rejects_unexpected_response_type() {
        let err = drive(vec![AgentResponse::Pong {
            version: 1,
            capabilities: vec![],
        }])
        .unwrap_err();
        assert!(format!("{}", err).contains("unexpected response"));
    }
}

#[cfg(test)]
mod collect_exec_cap_tests {
    //! Regression coverage for the streaming-exec host-OOM guard.
    //!
    //! Proves the streaming/collect path (`collect_exec_events_inner`) stops
    //! and emits a truncation signal once cumulative relayed output crosses the
    //! cap, instead of buffering unbounded output from a guest that never sends
    //! `Exited` (the `yes` / `cat /dev/zero` cross-tenant DoS).
    use super::*;

    fn stdout(n: usize) -> AgentResponse {
        AgentResponse::Stdout { data: vec![0u8; n] }
    }

    /// Drive `collect_exec_events_inner` over a fixed response list with a
    /// small cap, capturing the relayed events and the terminal result.
    fn drive(cap: usize, responses: Vec<AgentResponse>) -> (Vec<ExecEvent>, Result<()>) {
        let mut iter = responses.into_iter();
        let mut events = Vec::new();
        let res = collect_exec_events_inner(
            || {
                iter.next()
                    .ok_or_else(|| Error::agent("test", "no more responses"))
            },
            "test exec",
            cap,
            |e| events.push(e),
        );
        (events, res)
    }

    const TEST_CAP: usize = 1024;

    #[test]
    fn stops_at_cumulative_cap_with_truncation_signal() {
        // An infinite chatty stream: `Started`, then far more stdout than the
        // cap, and crucially NEVER `Exited`. If the cap did not fire the loop
        // would drain all 1000 frames (and in production loop forever). Because
        // it caps, only a handful of frames are consumed before termination.
        let half = TEST_CAP / 2;
        let mut responses = vec![AgentResponse::Started];
        for _ in 0..1000 {
            responses.push(stdout(half)); // would total 500_000 bytes unbounded
        }
        let (events, res) = drive(TEST_CAP, responses);
        res.expect("collect returns Ok after capping (not the iterator-drained error)");

        // Terminated near the cap, nowhere near draining 1000 frames.
        assert!(
            events.len() < 10,
            "expected termination near the cap, got {} events",
            events.len()
        );

        // Final event is the truncation error.
        match events.last().expect("at least one event") {
            ExecEvent::Error(msg) => assert!(
                msg.contains("byte cap") && msg.contains("truncated"),
                "expected truncation error, got: {msg}"
            ),
            other => panic!("expected trailing truncation Error event, got {other:?}"),
        }

        // Relayed bytes are bounded: at most cap + one frame.
        let relayed: usize = events
            .iter()
            .map(|e| match e {
                ExecEvent::Stdout(d) | ExecEvent::Stderr(d) => d.len(),
                _ => 0,
            })
            .sum();
        assert!(
            relayed <= TEST_CAP + half,
            "relayed {relayed} bytes exceeds the cap+one-frame bound"
        );
    }

    #[test]
    fn single_oversized_frame_trips_the_cap() {
        // One frame that alone crosses the cap: relayed once, then truncated.
        let (events, res) = drive(TEST_CAP, vec![AgentResponse::Started, stdout(TEST_CAP + 1)]);
        res.unwrap();
        assert!(matches!(events.last().unwrap(), ExecEvent::Error(m) if m.contains("byte cap")));
    }

    #[test]
    fn passes_through_under_cap_and_exits_cleanly() {
        // Normal, well-behaved exec under the cap relays everything and ends on
        // `Exit` with no truncation error injected.
        let (events, res) = drive(
            TEST_CAP,
            vec![
                AgentResponse::Started,
                stdout(100),
                AgentResponse::Stderr {
                    data: b"warn".to_vec(),
                },
                AgentResponse::Exited {
                    exit_code: 0,
                    oom: false,
                },
            ],
        );
        res.unwrap();
        assert_eq!(events.last().unwrap(), &ExecEvent::Exit(0));
        assert!(
            !events.iter().any(|e| matches!(e, ExecEvent::Error(_))),
            "clean exec must not inject a truncation error"
        );
    }
}

#[cfg(test)]
mod run_background_tests {
    //! Regression test for image-backed `machine run -d`.
    //!
    //! The original bug: the CLI's image + `--detach` path pulled the image
    //! and persisted the VM record but silently dropped the command. The
    //! fix wires in `AgentClient::run_background`, which must send a
    //! `Run { background: true }` over the wire and parse the returned PID.
    //!
    //! If this test fails, the detach path either lost its `background`
    //! plumbing or stopped parsing the PID response — either way, the
    //! original "command never runs" regression is back.
    use super::*;
    use smolvm_protocol::{encode_message, AgentRequest, AgentResponse, Envelope};
    use std::io::{Read, Write};
    use std::thread;

    #[test]
    fn run_background_sends_background_true_and_returns_pid() {
        let (client_stream, mut server_stream) = UdsStream::pair().unwrap();

        // Fake agent: read one request, assert it's a background Run, respond
        // with a Completed PID. Mirrors what the real agent does in
        // `handle_run_background`.
        let server = thread::spawn(move || {
            let mut len_buf = [0u8; 4];
            server_stream.read_exact(&mut len_buf).unwrap();
            let len = u32::from_be_bytes(len_buf) as usize;

            let mut payload = vec![0u8; len];
            server_stream.read_exact(&mut payload).unwrap();

            let envelope: Envelope<AgentRequest> =
                serde_json::from_slice(&payload).expect("valid Envelope<AgentRequest>");

            match envelope.body {
                AgentRequest::Run {
                    image,
                    command,
                    persistent_overlay_id,
                    background,
                    interactive,
                    tty,
                    ..
                } => {
                    assert!(
                        background,
                        "run_background must send background: true — the image+detach CLI path \
                         depends on this field to dispatch the command inside the container"
                    );
                    assert!(!interactive, "background runs are never interactive");
                    assert!(!tty, "background runs never allocate a TTY");
                    assert_eq!(image, "docker.io/library/alpine:3.19");
                    assert_eq!(command, vec!["sh", "-c", "echo hi"]);
                    assert_eq!(
                        persistent_overlay_id,
                        Some("default".to_string()),
                        "background runs must use a persistent overlay so subsequent execs \
                         see the same filesystem"
                    );
                }
                other => panic!("expected AgentRequest::Run, got {:?}", other),
            }

            let resp = AgentResponse::Completed {
                exit_code: 0,
                stdout: b"12345".to_vec(),
                stderr: Vec::new(),
            };
            let encoded = encode_message(&resp).expect("encode response");
            server_stream.write_all(&encoded).expect("write response");
        });

        let mut client = AgentClient::from_stream(client_stream);
        let config = RunConfig::new(
            "alpine:3.19",
            vec!["sh".to_string(), "-c".to_string(), "echo hi".to_string()],
        )
        .with_persistent_overlay(Some("default".to_string()));

        let pid = client
            .run_background(config)
            .expect("run_background should succeed on a Completed response");

        assert_eq!(pid, 12345, "client must parse the PID from stdout");
        server.join().expect("server thread joined cleanly");
    }

    #[test]
    fn run_background_rejects_nonzero_exit_code() {
        // If the agent fails to spawn the container, it returns a non-zero
        // exit_code. The client must turn that into an error rather than
        // silently returning a bogus PID.
        let (client_stream, mut server_stream) = UdsStream::pair().unwrap();

        let server = thread::spawn(move || {
            let mut len_buf = [0u8; 4];
            server_stream.read_exact(&mut len_buf).unwrap();
            let len = u32::from_be_bytes(len_buf) as usize;
            let mut payload = vec![0u8; len];
            server_stream.read_exact(&mut payload).unwrap();

            let resp = AgentResponse::Completed {
                exit_code: 1,
                stdout: Vec::new(),
                stderr: b"spawn failed".to_vec(),
            };
            let encoded = encode_message(&resp).unwrap();
            server_stream.write_all(&encoded).unwrap();
        });

        let mut client = AgentClient::from_stream(client_stream);
        let config = RunConfig::new("alpine:3.19", vec!["true".to_string()])
            .with_persistent_overlay(Some("default".to_string()));

        let err = client
            .run_background(config)
            .expect_err("non-zero exit must surface as an error");
        assert!(
            format!("{}", err).contains("spawn failed")
                || format!("{}", err).contains("run background"),
            "unexpected error: {}",
            err
        );
        server.join().unwrap();
    }
}

#[cfg(test)]
mod run_container_detached_tests {
    use super::*;
    use smolvm_protocol::{decode_message, encode_message, AgentRequest, AgentResponse, Envelope};
    use std::io::{Read, Write};
    use std::thread;

    #[test]
    fn accepts_progress_before_detached_container_completion() {
        let (client_stream, mut server_stream) = UdsStream::pair().unwrap();

        let server = thread::spawn(move || {
            let mut len_buf = [0u8; 4];
            server_stream.read_exact(&mut len_buf).unwrap();
            let len = u32::from_be_bytes(len_buf) as usize;
            let mut payload = vec![0u8; len];
            server_stream.read_exact(&mut payload).unwrap();
            let envelope: Envelope<AgentRequest> =
                decode_message(&[&len_buf[..], &payload[..]].concat()).unwrap();
            assert!(matches!(
                envelope.body,
                AgentRequest::Run { detached: true, .. }
            ));

            for response in [
                AgentResponse::Progress {
                    message: "extracting flattened rootfs (64 MiB)".to_string(),
                    percent: None,
                    layer: None,
                },
                AgentResponse::Completed {
                    exit_code: 0,
                    stdout: b"container-123".to_vec(),
                    stderr: Vec::new(),
                },
            ] {
                server_stream
                    .write_all(&encode_message(&response).unwrap())
                    .unwrap();
            }
        });

        let mut client = AgentClient::from_stream(client_stream);
        let config = RunConfig::new("local-archive", Vec::new())
            .with_persistent_overlay(Some("default".to_string()));
        let id = client.run_container_detached(config).unwrap();

        assert_eq!(id, "container-123");
        server.join().unwrap();
    }
}

#[cfg(test)]
mod run_streaming_tests {
    //! Regression coverage for image-backed streaming exec.
    //!
    //! `machine exec --stream` must preserve the same execution target as
    //! buffered `machine exec`. For image-backed machines that means sending a
    //! `Run { interactive: true }` request to the agent, not `VmExec` against
    //! the bare agent rootfs.
    use super::*;
    use smolvm_protocol::{encode_message, AgentRequest, AgentResponse, Envelope};
    use std::io::{Read, Write};
    use std::thread;

    #[test]
    fn run_streaming_sends_interactive_run_and_collects_events() {
        let (client_stream, mut server_stream) = UdsStream::pair().unwrap();

        let server = thread::spawn(move || {
            let mut len_buf = [0u8; 4];
            server_stream.read_exact(&mut len_buf).unwrap();
            let len = u32::from_be_bytes(len_buf) as usize;
            let mut payload = vec![0u8; len];
            server_stream.read_exact(&mut payload).unwrap();

            let envelope: Envelope<AgentRequest> =
                serde_json::from_slice(&payload).expect("valid Envelope<AgentRequest>");
            match envelope.body {
                AgentRequest::Run {
                    image,
                    command,
                    mounts,
                    interactive,
                    tty,
                    detached,
                    background,
                    persistent_overlay_id,
                    ..
                } => {
                    assert_eq!(image, "docker.io/library/ubuntu:24.04");
                    assert_eq!(command, vec!["/bin/bash", "-lc", "echo hi"]);
                    assert_eq!(
                        mounts,
                        vec![("work".to_string(), "/work".to_string(), false)]
                    );
                    assert!(interactive, "streaming image exec must use interactive Run");
                    assert!(!tty, "plain --stream must not allocate a TTY");
                    assert!(!detached, "streaming exec must not detach");
                    assert!(!background, "streaming exec must not run in background");
                    assert_eq!(persistent_overlay_id, Some("dev".to_string()));
                }
                other => panic!("expected AgentRequest::Run, got {:?}", other),
            }

            for response in [
                AgentResponse::Started,
                AgentResponse::Stdout {
                    data: b"hi\n".to_vec(),
                },
                AgentResponse::Stderr {
                    data: b"warn\n".to_vec(),
                },
                AgentResponse::Exited {
                    exit_code: 7,
                    oom: false,
                },
            ] {
                let encoded = encode_message(&response).expect("encode response");
                server_stream.write_all(&encoded).expect("write response");
            }
        });

        let mut client = AgentClient::from_stream(client_stream);
        let config = RunConfig::new(
            "ubuntu:24.04",
            vec![
                "/bin/bash".to_string(),
                "-lc".to_string(),
                "echo hi".to_string(),
            ],
        )
        .with_mounts(vec![("work".to_string(), "/work".to_string(), false)])
        .with_persistent_overlay(Some("dev".to_string()));

        let mut events = Vec::new();
        client
            .run_streaming_with(config, |event| events.push(event))
            .expect("run_streaming_with should handle streamed events");

        assert_eq!(
            events,
            vec![
                ExecEvent::Stdout(b"hi\n".to_vec()),
                ExecEvent::Stderr(b"warn\n".to_vec()),
                ExecEvent::Exit(7),
            ]
        );
        server.join().expect("server thread joined cleanly");
    }
}

#[cfg(test)]
mod writer_thread_tests {
    //! Regression tests for issue #926.
    //!
    //! They verify nonblocking enqueueing, retry after write timeout, frame
    //! ordering, error propagation, and bounded teardown.
    use super::*;
    use std::collections::VecDeque;
    use std::io::Read;
    use std::time::{Duration, Instant};

    const FRAME_LEN: usize = 4096;

    /// Frames whose byte pattern encodes their index, so a reordered or spliced
    /// stream is detectable.
    fn make_frames(count: usize) -> Vec<Vec<u8>> {
        (0..count).map(|i| vec![i as u8; FRAME_LEN]).collect()
    }

    /// Shrink both socket buffers so a peer that stops reading wedges the writer
    /// after a few frames instead of hundreds of KiB.
    fn shrink_buffers(a: &UdsStream, b: &UdsStream) {
        for s in [a, b] {
            let _ = s.as_socket().set_send_buffer_size(FRAME_LEN);
            let _ = s.as_socket().set_recv_buffer_size(FRAME_LEN);
        }
    }

    #[test]
    fn stalled_peer_does_not_abort_or_reorder_the_stream() {
        let (client_stream, mut peer) = UdsStream::pair().unwrap();
        shrink_buffers(&client_stream, &peer);
        peer.set_read_timeout(Some(Duration::from_secs(10)))
            .expect("set peer read timeout");

        let frames = make_frames(WRITER_QUEUE_FRAMES + 16);
        let total: usize = frames.iter().map(|f| f.len()).sum();

        // Peer wedges past a full WRITER_WRITE_TIMEOUT tick before draining, so
        // the writer must survive at least one expired write mid-frame.
        let peer_thread = std::thread::spawn(move || {
            std::thread::sleep(WRITER_WRITE_TIMEOUT + Duration::from_millis(200));
            let mut received = vec![0u8; total];
            peer.read_exact(&mut received)
                .expect("peer reads all bytes");
            received
        });

        let writer = FrameWriter::spawn(&client_stream).expect("spawn writer");

        // Queue with the session loop's discipline: one held frame, retried until
        // it fits. The first WRITER_QUEUE_FRAMES sends must not block.
        let start = Instant::now();
        let mut pending: VecDeque<Vec<u8>> = frames.iter().cloned().collect();
        let mut queued = 0;
        while queued < WRITER_QUEUE_FRAMES {
            let frame = pending.pop_front().expect("frames remain");
            assert!(
                writer.try_send(frame).expect("writer alive").is_none(),
                "the first {WRITER_QUEUE_FRAMES} frames must fit in the queue"
            );
            queued += 1;
        }
        let enqueue_time = start.elapsed();
        assert!(
            enqueue_time < Duration::from_millis(500),
            "enqueueing must not block on a stalled peer (took {enqueue_time:?})"
        );

        // Remainder: exactly the held-frame retry the session loop performs.
        while let Some(frame) = pending.pop_front() {
            let mut held = Some(frame);
            while let Some(frame) = held.take() {
                held = writer.try_send(frame).expect("writer alive");
                if held.is_some() {
                    std::thread::sleep(Duration::from_millis(5));
                }
            }
        }

        let received = peer_thread.join().expect("peer thread joined");
        assert_eq!(
            received,
            frames.concat(),
            "every frame must arrive whole and in order"
        );
        assert!(
            writer.take_error().is_none(),
            "a slow peer is not a write error"
        );
        drop(writer);
    }

    #[test]
    fn teardown_does_not_hang_on_a_wedged_peer() {
        let (client_stream, peer) = UdsStream::pair().unwrap();
        shrink_buffers(&client_stream, &peer);

        let writer = FrameWriter::spawn(&client_stream).expect("spawn writer");
        // Push until the queue backs up, which means the socket buffers are full
        // and the writer thread is parked inside write() — the wedged-guest state.
        let mut wedged = false;
        for frame in make_frames(4 * WRITER_QUEUE_FRAMES) {
            if writer.try_send(frame).expect("writer alive").is_some() {
                wedged = true;
                break;
            }
        }
        assert!(
            wedged,
            "peer never applied backpressure; test is not exercising a parked writer"
        );

        let start = Instant::now();
        drop(writer);
        let elapsed = start.elapsed();
        assert!(
            elapsed < WRITER_WRITE_TIMEOUT * 3,
            "dropping the writer must not wait on the peer (took {elapsed:?})"
        );
        drop(peer);
    }

    #[test]
    fn write_error_is_published_instead_of_lost() {
        let (client_stream, peer) = UdsStream::pair().unwrap();
        shrink_buffers(&client_stream, &peer);

        let writer = FrameWriter::spawn(&client_stream).expect("spawn writer");
        drop(peer); // Connection dies mid-session: broken pipe on the next write.

        // The session loop learns of the failure one iteration late, either from
        // the error slot or from a send onto the now-dead channel.
        let deadline = Instant::now() + Duration::from_secs(5);
        let mut reported = false;
        while Instant::now() < deadline {
            if writer.take_error().is_some() {
                reported = true;
                break;
            }
            match writer.try_send(vec![0u8; FRAME_LEN]) {
                Ok(_) => std::thread::sleep(Duration::from_millis(10)),
                Err(e) => {
                    assert!(
                        e.to_string().contains("send message"),
                        "writer failure must surface as a send error, got: {e}"
                    );
                    reported = true;
                    break;
                }
            }
        }
        assert!(reported, "a dead socket must surface as an error");
    }
}

#[cfg(test)]
mod stalled_body_tests {
    //! Regression for the busy-spin-forever on a stalled mid-frame body.
    //!
    //! A peer that writes a valid 4-byte length prefix and then delivers fewer
    //! body bytes than declared (without closing the socket) used to pin the
    //! host thread in a 1ms retry loop with no wall-clock bound, defeating every
    //! client-level timeout. `read_exact_retry` now honors a wall-clock idle
    //! deadline derived from the socket read timeout, so `receive()` must return
    //! a timeout error promptly instead of hanging.
    use super::*;
    use std::io::Write;
    use std::time::{Duration, Instant};

    #[test]
    fn receive_times_out_on_stalled_mid_frame_body() {
        let (client_stream, mut server_stream) = UdsStream::pair().unwrap();

        // Short read timeout so the test is fast: the idle deadline tracks this.
        let read_timeout = Duration::from_millis(150);
        client_stream
            .set_read_timeout(Some(read_timeout))
            .expect("set client read timeout");

        // Server: declare a 64-byte body, then send only 3 bytes and STALL —
        // hold the socket open (never drop it) so the client never sees EOF.
        let server = std::thread::spawn(move || {
            let declared_len: u32 = 64;
            server_stream
                .write_all(&declared_len.to_be_bytes())
                .expect("write length prefix");
            server_stream
                .write_all(&[1u8, 2, 3])
                .expect("write partial body");
            server_stream.flush().expect("flush");
            // Stall: keep the connection open well past the client's deadline so
            // the client cannot rely on EOF to unblock.
            std::thread::sleep(Duration::from_secs(3));
            drop(server_stream);
        });

        let mut client = AgentClient::from_stream(client_stream);

        let start = Instant::now();
        let result = client.receive();
        let elapsed = start.elapsed();

        assert!(
            result.is_err(),
            "receive() must error on a stalled mid-frame body, not return Ok"
        );
        // Must return within a small multiple of the read timeout — nowhere near
        // the server's 3s hold, proving it did not busy-spin until EOF.
        assert!(
            elapsed < Duration::from_secs(2),
            "receive() should time out promptly (got {elapsed:?}); it must not \
             spin until the peer closes"
        );

        server.join().expect("server thread joined");
    }

    #[test]
    fn read_exact_retry_bounds_stalled_read_without_eof() {
        // Drive read_exact_retry directly: header-style non-propagating read of a
        // buffer larger than what the peer sends, with the peer stalling.
        let (client_stream, mut server_stream) = UdsStream::pair().unwrap();
        client_stream
            .set_read_timeout(Some(Duration::from_millis(100)))
            .expect("set read timeout");

        let server = std::thread::spawn(move || {
            server_stream.write_all(&[0xAAu8]).expect("write one byte");
            server_stream.flush().expect("flush");
            std::thread::sleep(Duration::from_secs(3));
            drop(server_stream);
        });

        let mut client = AgentClient::from_stream(client_stream);

        let start = Instant::now();
        // Ask for 8 bytes but only 1 will ever arrive; propagate=false forces the
        // retry path (the same path receive() uses for the body).
        let mut buf = [0u8; 8];
        let err = client
            .read_exact_retry(&mut buf, false)
            .expect_err("stalled read must return a bounded error");
        let elapsed = start.elapsed();

        assert_eq!(
            err.kind(),
            std::io::ErrorKind::TimedOut,
            "stalled mid-buffer read must surface a TimedOut error"
        );
        assert!(
            elapsed < Duration::from_secs(2),
            "read_exact_retry must not busy-spin until EOF (got {elapsed:?})"
        );

        server.join().expect("server thread joined");
    }
}

#[cfg(test)]
mod term_default_tests {
    use super::with_term_default;

    fn term_of(env: &[(String, String)]) -> Option<&str> {
        env.iter()
            .find(|(k, _)| k == "TERM")
            .map(|(_, v)| v.as_str())
    }

    #[test]
    fn pty_session_gets_a_term() {
        // Without this the guest shell's line editor cannot redraw, and backspace
        // appears to advance the cursor instead of deleting.
        let env = with_term_default(vec![], true);
        assert_eq!(term_of(&env), Some("xterm-256color"));
    }

    #[test]
    fn caller_supplied_term_is_kept() {
        let env = with_term_default(
            vec![("TERM".to_string(), "screen-256color".to_string())],
            true,
        );
        assert_eq!(term_of(&env), Some("screen-256color"));
        assert_eq!(env.len(), 1, "must not append a second TERM");
    }

    #[test]
    fn non_tty_exec_is_untouched() {
        // A piped exec has no terminal to describe; adding TERM there would make
        // programs emit escape sequences into what is usually captured output.
        let env = with_term_default(vec![("A".to_string(), "b".to_string())], false);
        assert_eq!(term_of(&env), None);
        assert_eq!(env.len(), 1);
    }
}
