car-server-core 0.36.0

//! On-device inference worker — process isolation for local MLX/Candle
//! generation (Parslee-ai/car-releases#74).
//!
//! A single heavy on-device generation can abort the process from the
//! Metal/MLX **C++** side — an allocation/OOM abort or a C++ exception that
//! crosses the FFI boundary *below* every Rust `catch_unwind`. In the shared
//! `car-server` daemon that kills inference for every connected client at once,
//! and (before the CarHost respawn work) left the host daemon-less.
//!
//! This module runs on-device generation in a **separate worker process** the
//! daemon owns:
//!
//! - [`WorkerOffload`] is a [`car_inference::LocalGenerationOffload`] the daemon
//!   installs on its engine. When `car-inference` resolves a request to an
//!   on-device backend it hands the request here instead of running the Metal
//!   decode loop in-process. `WorkerOffload` ships the request to the child
//!   over a line-delimited-JSON stdio protocol and returns the child's
//!   [`InferenceResult`] / [`StreamEvent`] stream.
//! - [`run_mlx_worker`] is the child's main loop: it builds a real in-process
//!   engine (which runs the actual Metal path — a worker never offloads to
//!   itself) and services requests off stdin.
//!
//! When the worker aborts mid-request the daemon observes a closed pipe / EOF:
//! the offload call returns an `Err` (the daemon fails that one RPC gracefully
//! and stays up) and the worker is dropped so the **next** call respawns it. A
//! single oversized local inference therefore takes down only its own transient
//! worker, never the daemon.
//!
//! Access is serialized by one async mutex — matching the process-wide Metal
//! device lock that already serialized in-process on-device generation, so this
//! adds no new concurrency semantics.

use std::ffi::OsString;
use std::process::Stdio;
use std::sync::Arc;

use car_inference::stream::StreamEvent;
use car_inference::tasks::generate::GenerateRequest;
use car_inference::{
    InferenceConfig, InferenceEngine, InferenceError, InferenceResult, LocalGenerationOffload,
};
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader, Lines};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::Mutex;

/// Env var the daemon sets on the worker child. `car-inference` reads it
/// (`is_offload_worker`) to (a) never offload to itself and (b) disable outcome
/// persistence so the worker can't race the daemon writing shared `~/.car`
/// files. See [`car_inference::is_offload_worker`].
pub const WORKER_ENV: &str = "CAR_INFERENCE_WORKER";

/// Daemon → worker request (one JSON object per line).
#[derive(Serialize, Deserialize)]
enum WorkerRequest {
    /// Run a non-streaming generation to completion.
    Generate { request: Box<GenerateRequest> },
    /// Run a streaming generation, emitting `Event` lines then `StreamEnd`.
    Stream { request: Box<GenerateRequest> },
}

/// Worker → daemon response (one JSON object per line).
#[derive(Serialize, Deserialize)]
enum WorkerResponse {
    /// Terminal success for a `Generate` request.
    Result(Box<InferenceResult>),
    /// One streamed event for a `Stream` request.
    Event(Box<StreamEvent>),
    /// Terminal marker for a `Stream` request that finished cleanly.
    StreamEnd,
    /// A *reported* generation error (the model/request failed, e.g. an
    /// unsupported mode or a bad prompt). The worker is still healthy — this is
    /// NOT a crash — so the daemon surfaces the error without respawning.
    Error(String),
}

/// Outcome of one worker exchange, distinguishing a live worker's reported
/// error (keep the worker) from a dead worker (respawn on the next call).
enum Exchange<T> {
    Ok(T),
    /// The worker returned an error but is still alive.
    Reported(String),
    /// The pipe broke / worker died. Drop and respawn next time.
    Dead(String),
}

/// A live worker child plus its framed pipes.
struct WorkerProc {
    child: Child,
    stdin: ChildStdin,
    stdout: Lines<BufReader<ChildStdout>>,
}

impl WorkerProc {
    fn kill(mut self) {
        // Best-effort: the child is about to be dropped; make sure it dies
        // rather than lingering as a zombie holding the Metal device.
        let _ = self.child.start_kill();
    }
}

/// Installed on the daemon's engine to route on-device generation through a
/// persistent, auto-respawning worker child. Cheap to clone (shares the child).
#[derive(Clone)]
pub struct WorkerOffload {
    inner: Arc<Mutex<Option<WorkerProc>>>,
    program: OsString,
    args: Arc<Vec<OsString>>,
}

impl WorkerOffload {
    /// Offload to a fresh `car-server --mlx-worker` child of the current
    /// executable. Fails only if the current exe path can't be resolved.
    pub fn new() -> std::io::Result<Self> {
        let exe = std::env::current_exe()?;
        Ok(Self::with_command(
            exe,
            vec![OsString::from("--mlx-worker")],
        ))
    }

    /// Construct with an explicit worker command. Used by tests to point at a
    /// stub worker (a script that speaks the same line protocol) so the
    /// spawn/exchange/respawn logic is exercised without loading a real model.
    pub fn with_command(program: impl Into<OsString>, args: Vec<OsString>) -> Self {
        Self {
            inner: Arc::new(Mutex::new(None)),
            program: program.into(),
            args: Arc::new(args),
        }
    }

    fn spawn(&self) -> Result<WorkerProc, InferenceError> {
        let mut cmd = Command::new(&self.program);
        cmd.args(self.args.iter())
            .env(WORKER_ENV, "1")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            // Inherit stderr so the worker's tracing/logs land in the daemon's
            // own stderr (Console.app under CarHost). stdout is protocol-only.
            .stderr(Stdio::inherit())
            .kill_on_drop(true);
        let mut child = cmd.spawn().map_err(|e| {
            InferenceError::InferenceFailed(format!("failed to spawn inference worker: {e}"))
        })?;
        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| InferenceError::InferenceFailed("worker stdin unavailable".into()))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| InferenceError::InferenceFailed("worker stdout unavailable".into()))?;
        Ok(WorkerProc {
            child,
            stdin,
            stdout: BufReader::new(stdout).lines(),
        })
    }

    /// Take the live worker out of `slot`, spawning one if absent. The caller
    /// owns the returned worker for the duration of one exchange and must put a
    /// healthy worker BACK; if it doesn't (a clean drop on cancellation, or an
    /// error path), `kill_on_drop` reaps the child and `slot` stays empty so the
    /// next call respawns — which is what makes an interrupted exchange safe
    /// (no half-read frame can desync a subsequent request).
    fn take_or_spawn(&self, slot: &mut Option<WorkerProc>) -> Result<WorkerProc, InferenceError> {
        if let Some(mut p) = slot.take() {
            // A worker can die while idle in the slot (a crash on the PREVIOUS
            // request that we already surfaced, or an external kill). Reap it
            // with a non-blocking wait and spawn fresh instead of handing back a
            // corpse — otherwise the next request would eat a "broken pipe"
            // failure just to discover the death. `Ok(None)` = still running.
            match p.child.try_wait() {
                Ok(None) => return Ok(p),
                _ => p.kill(),
            }
        }
        self.spawn()
    }
}

async fn write_line<W, T>(w: &mut W, value: &T) -> std::io::Result<()>
where
    W: AsyncWrite + Unpin,
    T: Serialize,
{
    let mut line = serde_json::to_string(value).map_err(std::io::Error::other)?;
    line.push('\n');
    w.write_all(line.as_bytes()).await?;
    w.flush().await
}

/// Run one non-streaming request/response against a worker. Never touches the
/// slot — the caller drops a dead worker so the borrow ends first.
async fn do_generate(proc: &mut WorkerProc, request: GenerateRequest) -> Exchange<InferenceResult> {
    let req = WorkerRequest::Generate {
        request: Box::new(request),
    };
    if let Err(e) = write_line(&mut proc.stdin, &req).await {
        return Exchange::Dead(format!("write to inference worker failed: {e}"));
    }
    match proc.stdout.next_line().await {
        Ok(Some(line)) => match serde_json::from_str::<WorkerResponse>(&line) {
            Ok(WorkerResponse::Result(ir)) => Exchange::Ok(*ir),
            Ok(WorkerResponse::Error(msg)) => Exchange::Reported(msg),
            Ok(_) => Exchange::Dead("inference worker sent an unexpected response frame".into()),
            Err(e) => Exchange::Dead(format!("inference worker sent invalid JSON: {e}")),
        },
        Ok(None) => Exchange::Dead("inference worker exited mid-request (EOF)".into()),
        Err(e) => Exchange::Dead(format!("read from inference worker failed: {e}")),
    }
}

#[async_trait::async_trait]
impl LocalGenerationOffload for WorkerOffload {
    async fn generate(&self, request: GenerateRequest) -> Result<InferenceResult, InferenceError> {
        // Hold the lock across the exchange (serializes worker access, matching
        // the in-process Metal device lock). Take the worker OUT so a cancelled
        // future (e.g. handle_infer's timeout firing) drops it — kill_on_drop
        // reaps the child and the slot stays empty, so no orphaned response
        // frame can desync the next request.
        let mut slot = self.inner.lock().await;
        let mut proc = self.take_or_spawn(&mut slot)?;
        match do_generate(&mut proc, request).await {
            Exchange::Ok(ir) => {
                *slot = Some(proc); // healthy — return it for reuse
                Ok(ir)
            }
            Exchange::Reported(msg) => {
                // A model/request error, not a crash: the worker is fine.
                *slot = Some(proc);
                Err(InferenceError::InferenceFailed(msg))
            }
            Exchange::Dead(msg) => {
                // Worker died (very likely the Metal/MLX abort we isolate for).
                // `proc` drops here (kill_on_drop), slot stays None → next call
                // respawns. Fail THIS rpc only; the daemon stays up.
                proc.kill();
                Err(InferenceError::InferenceFailed(format!(
                    "on-device inference worker crashed and was restarted; \
                     this request failed but the daemon is up: {msg}"
                )))
            }
        }
    }

    async fn stream(
        &self,
        request: GenerateRequest,
    ) -> Result<tokio::sync::mpsc::Receiver<StreamEvent>, InferenceError> {
        // Hold the lock for the whole stream via an owned guard moved into the
        // reader task (serializes streams against other worker traffic, as the
        // in-process Metal device lock did). The worker is taken OUT of the slot
        // and only put BACK on a clean `StreamEnd` — any other exit (caller
        // drops the receiver mid-stream, a worker error, EOF) leaves the worker
        // mid-protocol, so it's killed and the next request respawns rather than
        // reading leftover frames.
        let mut guard = self.inner.clone().lock_owned().await;
        let mut proc = self.take_or_spawn(&mut guard)?;
        // Send the stream request synchronously so a spawn/write failure
        // surfaces as an Err here rather than a silently-empty stream.
        {
            let req = WorkerRequest::Stream {
                request: Box::new(request),
            };
            if let Err(e) = write_line(&mut proc.stdin, &req).await {
                proc.kill(); // slot already None → next call respawns
                return Err(InferenceError::InferenceFailed(format!(
                    "write to inference worker failed: {e}"
                )));
            }
        }

        let (tx, rx) = tokio::sync::mpsc::channel::<StreamEvent>(64);
        tokio::spawn(async move {
            let mut clean = false;
            loop {
                match proc.stdout.next_line().await {
                    Ok(Some(line)) => match serde_json::from_str::<WorkerResponse>(&line) {
                        Ok(WorkerResponse::Event(ev)) => {
                            if tx.send(*ev).await.is_err() {
                                break; // caller dropped the receiver, worker is mid-stream
                            }
                        }
                        Ok(WorkerResponse::StreamEnd) => {
                            clean = true;
                            break;
                        }
                        Ok(WorkerResponse::Error(msg)) => {
                            tracing::warn!(error = %msg, "inference worker stream error");
                            // Surface as a mid-stream failure the accumulator
                            // reads as an error end (mirrors remote.rs).
                            let _ = tx.send(StreamEvent::StopReason("error".into())).await;
                            break;
                        }
                        Ok(WorkerResponse::Result(_)) | Err(_) => {
                            // Protocol break — treat as a dead worker.
                            let _ = tx.send(StreamEvent::StopReason("error".into())).await;
                            break;
                        }
                    },
                    Ok(None) | Err(_) => {
                        // Worker died mid-stream: emit an error end.
                        let _ = tx.send(StreamEvent::StopReason("error".into())).await;
                        break;
                    }
                }
            }
            if clean {
                // Clean end: return the healthy worker to the slot for reuse.
                *guard = Some(proc);
            } else {
                // Interrupted / dead: kill so the next request gets a fresh
                // worker (slot is already None). `guard` drops, unlocking.
                proc.kill();
            }
        });
        Ok(rx)
    }
}

/// The worker child's main loop. Builds a real in-process engine (persistence
/// disabled via the `CAR_INFERENCE_WORKER` env the parent set — see
/// [`car_inference::is_offload_worker`]) and services line-delimited requests
/// off stdin, writing responses to stdout. Logs go to stderr (via the tracing
/// subscriber the binary installs), keeping stdout protocol-clean.
///
/// Returns when stdin closes (the daemon dropped the worker). One request is
/// served to completion before the next is read — the daemon never pipelines,
/// and this matches the single-Metal-device serialization the in-process path
/// relied on.
pub async fn run_mlx_worker() {
    // Belt-and-suspenders: the parent sets this, but guarantee the self-offload
    // / no-persist guard holds even if a worker is launched some other way.
    std::env::set_var(WORKER_ENV, "1");

    let engine = Arc::new(InferenceEngine::new(InferenceConfig::default()));
    let mut lines = BufReader::new(tokio::io::stdin()).lines();
    let mut stdout = tokio::io::stdout();

    while let Ok(Some(line)) = lines.next_line().await {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        let req: WorkerRequest = match serde_json::from_str(line) {
            Ok(r) => r,
            Err(e) => {
                let _ = write_line(
                    &mut stdout,
                    &WorkerResponse::Error(format!("malformed worker request: {e}")),
                )
                .await;
                continue;
            }
        };
        match req {
            WorkerRequest::Generate { request } => {
                let resp = match engine.generate_tracked(*request).await {
                    Ok(ir) => WorkerResponse::Result(Box::new(ir)),
                    Err(e) => WorkerResponse::Error(e.to_string()),
                };
                if write_line(&mut stdout, &resp).await.is_err() {
                    break; // parent gone
                }
            }
            WorkerRequest::Stream { request } => {
                match engine.generate_tracked_stream(*request).await {
                    Ok(mut tracked) => {
                        let mut broke = false;
                        while let Some(ev) = tracked.events.recv().await {
                            if write_line(&mut stdout, &WorkerResponse::Event(Box::new(ev)))
                                .await
                                .is_err()
                            {
                                broke = true;
                                break;
                            }
                        }
                        if broke {
                            break;
                        }
                        if write_line(&mut stdout, &WorkerResponse::StreamEnd)
                            .await
                            .is_err()
                        {
                            break;
                        }
                    }
                    Err(e) => {
                        if write_line(&mut stdout, &WorkerResponse::Error(e.to_string()))
                            .await
                            .is_err()
                        {
                            break;
                        }
                    }
                }
            }
        }
    }
}

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

    // Externally-tagged serde shapes the stub workers below emit. Locking these
    // down as literals doubles as a wire-format regression guard.
    const RESULT_LINE: &str = r#"{"Result":{"text":"pong","tool_calls":[],"trace_id":"t","model_used":"stub","latency_ms":1,"usage":null}}"#;

    #[test]
    fn envelope_serde_round_trips() {
        // Request
        let req = WorkerRequest::Generate {
            request: Box::new(GenerateRequest {
                prompt: "hi".into(),
                ..Default::default()
            }),
        };
        let json = serde_json::to_string(&req).unwrap();
        assert!(json.starts_with(r#"{"Generate":"#), "got {json}");
        assert!(matches!(
            serde_json::from_str::<WorkerRequest>(&json).unwrap(),
            WorkerRequest::Generate { .. }
        ));

        // Responses: Result, Event(StreamEvent), StreamEnd, Error.
        assert!(matches!(
            serde_json::from_str::<WorkerResponse>(RESULT_LINE).unwrap(),
            WorkerResponse::Result(_)
        ));
        let ev = WorkerResponse::Event(Box::new(StreamEvent::TextDelta("hi".into())));
        let ev_json = serde_json::to_string(&ev).unwrap();
        assert_eq!(ev_json, r#"{"Event":{"TextDelta":"hi"}}"#);
        assert_eq!(
            serde_json::to_string(&WorkerResponse::StreamEnd).unwrap(),
            r#""StreamEnd""#
        );
        // The StopReason variant the WS-runner JSON mapping omits must survive.
        let sr = StreamEvent::StopReason("length".into());
        let sr2: StreamEvent = serde_json::from_str(&serde_json::to_string(&sr).unwrap()).unwrap();
        assert!(matches!(sr2, StreamEvent::StopReason(s) if s == "length"));
    }

    // The subprocess-driven tests use a POSIX shell stub as the "worker", so
    // they exercise the real spawn / stdio-protocol / respawn paths without a
    // model. Unix-only (Windows CI has no `sh`); the feature is macOS-primary.
    #[cfg(unix)]
    fn sh(script: &str) -> WorkerOffload {
        WorkerOffload::with_command("sh", vec![OsString::from("-c"), OsString::from(script)])
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn generate_round_trips_and_reuses_the_worker() {
        // A stub that answers every request line with one Result line and stays
        // alive — so a second call must reuse the same worker.
        let off = sh(&format!(
            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
        ));
        let r1 = off.generate(GenerateRequest::default()).await.unwrap();
        assert_eq!(r1.text, "pong");
        let r2 = off.generate(GenerateRequest::default()).await.unwrap();
        assert_eq!(r2.text, "pong");
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn worker_crash_fails_one_call_gracefully_then_respawns() {
        // A stub that exits immediately: the exchange sees a broken pipe / EOF,
        // which must surface as an Err (NOT a panic / hang) — this is the whole
        // point of car-releases#74: one bad inference fails its own RPC and the
        // daemon stays up. The next call respawns (and crashes again here), which
        // must ALSO be a clean Err, proving the slot recovered rather than
        // desyncing.
        let off = sh("exit 1");
        let e1 = off.generate(GenerateRequest::default()).await.unwrap_err();
        assert!(
            e1.to_string().contains("crashed and was restarted"),
            "unexpected error: {e1}"
        );
        let e2 = off.generate(GenerateRequest::default()).await.unwrap_err();
        assert!(
            e2.to_string().contains("crashed and was restarted"),
            "unexpected error: {e2}"
        );
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn stream_forwards_events_then_closes() {
        // A stub that emits two text deltas and a clean StreamEnd for one request.
        let off = sh("IFS= read -r line; \
             printf '%s\\n' '{\"Event\":{\"TextDelta\":\"he\"}}'; \
             printf '%s\\n' '{\"Event\":{\"TextDelta\":\"llo\"}}'; \
             printf '%s\\n' '\"StreamEnd\"'; \
             while IFS= read -r l; do :; done");
        let mut rx = off.stream(GenerateRequest::default()).await.unwrap();
        let mut got = String::new();
        while let Some(ev) = rx.recv().await {
            if let StreamEvent::TextDelta(t) = ev {
                got.push_str(&t);
            }
        }
        assert_eq!(got, "hello");
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn stream_worker_death_closes_with_error_end() {
        // A stub that emits one event then exits mid-stream (no StreamEnd): the
        // reader must surface a StopReason("error") and close, not hang.
        let off = sh("IFS= read -r line; printf '%s\\n' '{\"Event\":{\"TextDelta\":\"partial\"}}'");
        let mut rx = off.stream(GenerateRequest::default()).await.unwrap();
        let mut saw_text = false;
        let mut saw_error_end = false;
        while let Some(ev) = rx.recv().await {
            match ev {
                StreamEvent::TextDelta(t) if t == "partial" => saw_text = true,
                StreamEvent::StopReason(s) if s == "error" => saw_error_end = true,
                _ => {}
            }
        }
        assert!(saw_text, "should have forwarded the partial delta");
        assert!(
            saw_error_end,
            "mid-stream death should surface an error end"
        );
    }
}