Skip to main content

car_server_core/
inference_worker.rs

1//! On-device inference worker — process isolation for local MLX/Candle
2//! generation (Parslee-ai/car-releases#74).
3//!
4//! A single heavy on-device generation can abort the process from the
5//! Metal/MLX **C++** side — an allocation/OOM abort or a C++ exception that
6//! crosses the FFI boundary *below* every Rust `catch_unwind`. In the shared
7//! `car-server` daemon that kills inference for every connected client at once,
8//! and (before the CarHost respawn work) left the host daemon-less.
9//!
10//! This module runs on-device generation in a **separate worker process** the
11//! daemon owns:
12//!
13//! - [`WorkerOffload`] is a [`car_inference::LocalGenerationOffload`] the daemon
14//!   installs on its engine. When `car-inference` resolves a request to an
15//!   on-device backend it hands the request here instead of running the Metal
16//!   decode loop in-process. `WorkerOffload` ships the request to the child
17//!   over a line-delimited-JSON stdio protocol and returns the child's
18//!   [`InferenceResult`] / [`StreamEvent`] stream.
19//! - [`run_mlx_worker`] is the child's main loop: it builds a real in-process
20//!   engine (which runs the actual Metal path — a worker never offloads to
21//!   itself) and services requests off stdin.
22//!
23//! When the worker aborts mid-request the daemon observes a closed pipe / EOF:
24//! the offload call returns an `Err` (the daemon fails that one RPC gracefully
25//! and stays up) and the worker is dropped so the **next** call respawns it. A
26//! single oversized local inference therefore takes down only its own transient
27//! worker, never the daemon.
28//!
29//! Access is serialized by one async mutex — matching the process-wide Metal
30//! device lock that already serialized in-process on-device generation, so this
31//! adds no new concurrency semantics.
32
33use std::collections::HashMap;
34use std::ffi::OsString;
35use std::process::Stdio;
36use std::sync::atomic::{AtomicU64, Ordering};
37use std::sync::Arc;
38
39use std::time::Duration;
40
41use car_inference::stream::StreamEvent;
42use car_inference::tasks::generate::GenerateRequest;
43use car_inference::{
44    InferenceConfig, InferenceEngine, InferenceError, InferenceResult, InferenceTerminationAck,
45    LocalGenerationOffload, LocalLoadPreflight, LocalOffloadResult, LocalOffloadStream,
46    LocalWorkerAdmission, LocalWorkerResidency,
47};
48use serde::{Deserialize, Serialize};
49use tokio::io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader, Lines};
50use tokio::process::{Child, ChildStdin, ChildStdout, Command};
51use tokio::sync::{mpsc, oneshot, Mutex};
52use tokio::time::Instant;
53
54type WorkerCancelSender = mpsc::UnboundedSender<oneshot::Sender<bool>>;
55
56struct WorkerControlRegistration {
57    inference_id: String,
58    controls: Arc<std::sync::Mutex<HashMap<String, WorkerCancelSender>>>,
59}
60
61impl WorkerControlRegistration {
62    fn install(
63        controls: Arc<std::sync::Mutex<HashMap<String, WorkerCancelSender>>>,
64        inference_id: String,
65    ) -> (Self, mpsc::UnboundedReceiver<oneshot::Sender<bool>>) {
66        let (sender, receiver) = mpsc::unbounded_channel();
67        controls
68            .lock()
69            .unwrap_or_else(std::sync::PoisonError::into_inner)
70            .insert(inference_id.clone(), sender);
71        (
72            Self {
73                inference_id,
74                controls,
75            },
76            receiver,
77        )
78    }
79}
80
81impl Drop for WorkerControlRegistration {
82    fn drop(&mut self) {
83        self.controls
84            .lock()
85            .unwrap_or_else(std::sync::PoisonError::into_inner)
86            .remove(&self.inference_id);
87    }
88}
89
90/// Env var the daemon sets on the worker child. `car-inference` reads it
91/// (`is_offload_worker`) to (a) never offload to itself and (b) disable outcome
92/// persistence so the worker can't race the daemon writing shared `~/.car`
93/// files. See [`car_inference::is_offload_worker`].
94pub const WORKER_ENV: &str = "CAR_INFERENCE_WORKER";
95
96/// Daemon → worker request (one JSON object per line).
97#[derive(Serialize, Deserialize)]
98enum WorkerRequest {
99    /// Run a non-streaming generation to completion.
100    Generate {
101        request: Box<GenerateRequest>,
102        admission: LocalWorkerAdmission,
103    },
104    /// Run a streaming generation, emitting `Event` lines then `StreamEnd`.
105    Stream {
106        request: Box<GenerateRequest>,
107        admission: LocalWorkerAdmission,
108    },
109}
110
111#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
112struct WorkerResidencyAck {
113    model_id: String,
114    measured_weights_bytes: u64,
115    retention: car_inference::backend_cache::BackendRetention,
116}
117
118/// Worker → daemon response (one JSON object per line).
119#[derive(Serialize, Deserialize)]
120enum WorkerResponse {
121    /// Terminal success for a `Generate` request.
122    Result {
123        result: Box<InferenceResult>,
124        residency: WorkerResidencyAck,
125    },
126    StreamStarted {
127        residency: WorkerResidencyAck,
128    },
129    /// One streamed event for a `Stream` request.
130    Event(Box<StreamEvent>),
131    /// Terminal marker for a `Stream` request that finished cleanly.
132    StreamEnd,
133    /// A *reported* generation error (the model/request failed, e.g. an
134    /// unsupported mode or a bad prompt). The worker is still healthy — this is
135    /// NOT a crash — so the daemon surfaces the error without respawning.
136    Error(String),
137    LocalResourceBlocked {
138        preflight: LocalLoadPreflight,
139        recovery: String,
140    },
141}
142
143/// Outcome of one worker exchange, distinguishing a live worker's reported
144/// error (keep the worker) from a dead worker (respawn on the next call).
145enum Exchange<T> {
146    Ok(T),
147    /// The worker returned an error but is still alive.
148    Reported(InferenceError),
149    /// The pipe broke / worker died. Drop and respawn next time.
150    Dead(String),
151}
152
153/// A live worker child plus its framed pipes.
154struct WorkerProc {
155    child: Child,
156    stdin: ChildStdin,
157    stdout: Lines<BufReader<ChildStdout>>,
158    policy_generation: u64,
159    state_root: Option<std::path::PathBuf>,
160}
161
162#[derive(Clone)]
163struct WorkerResident {
164    allocation_id: String,
165    coordinator: Arc<car_inference::resource_policy::LocalAdmissionCoordinator>,
166    /// Last successfully completed request. `None` means the resident has not
167    /// completed a request yet (for example, a stream that only started).
168    last_used: Option<Instant>,
169}
170
171type WorkerResidentMap = HashMap<(std::path::PathBuf, String), WorkerResident>;
172
173#[derive(Clone)]
174struct WorkerResidentOwner {
175    root: std::path::PathBuf,
176    logical_model_id: String,
177    allocation_id: String,
178    coordinator: Arc<car_inference::resource_policy::LocalAdmissionCoordinator>,
179}
180
181struct WorkerProcessGuard {
182    worker: Option<WorkerProc>,
183    residents: Vec<WorkerResidentOwner>,
184    candidate: Option<WorkerResidentOwner>,
185    resident_models: Arc<std::sync::Mutex<WorkerResidentMap>>,
186    teardown_started: bool,
187}
188
189impl WorkerProcessGuard {
190    fn new(
191        worker: WorkerProc,
192        resident_models: Arc<std::sync::Mutex<WorkerResidentMap>>,
193        candidate: Option<(std::path::PathBuf, String, String)>,
194    ) -> Self {
195        let mut residents = resident_models
196            .lock()
197            .unwrap_or_else(std::sync::PoisonError::into_inner)
198            .iter()
199            .map(|((root, logical_model_id), resident)| WorkerResidentOwner {
200                root: root.clone(),
201                logical_model_id: logical_model_id.clone(),
202                allocation_id: resident.allocation_id.clone(),
203                coordinator: resident.coordinator.clone(),
204            })
205            .collect::<Vec<_>>();
206        let mut candidate_owner = None;
207        if let Some((root, model_id, allocation_id)) =
208            candidate.filter(|(_, model_id, _)| !model_id.is_empty())
209        {
210            let root = car_inference::resource_policy::normalized_state_root_key(&root);
211            if !residents
212                .iter()
213                .any(|resident| resident.root == root && resident.logical_model_id == model_id)
214            {
215                if let Some(coordinator) =
216                    car_inference::resource_policy::local_admission_for_scope(&root)
217                {
218                    let owner = WorkerResidentOwner {
219                        root,
220                        allocation_id,
221                        logical_model_id: model_id,
222                        coordinator,
223                    };
224                    residents.push(owner.clone());
225                    candidate_owner = Some(owner);
226                }
227            }
228        }
229        Self {
230            worker: Some(worker),
231            residents,
232            candidate: candidate_owner,
233            resident_models,
234            teardown_started: false,
235        }
236    }
237
238    fn worker_mut(&mut self) -> &mut WorkerProc {
239        self.worker.as_mut().expect("worker process owned")
240    }
241
242    fn charge_candidate(&self, measured_weights_bytes: u64) {
243        if let Some(candidate) = &self.candidate {
244            candidate
245                .coordinator
246                .mark_teardown_pending_allocation_with_charge(
247                    &candidate.logical_model_id,
248                    &candidate.allocation_id,
249                    measured_weights_bytes,
250                );
251        }
252    }
253
254    fn clear_candidate(&mut self) {
255        let Some(candidate) = self.candidate.take() else {
256            return;
257        };
258        candidate
259            .coordinator
260            .finish_teardown_allocation(&candidate.logical_model_id, &candidate.allocation_id);
261        self.residents.retain(|resident| {
262            resident.root != candidate.root
263                || resident.logical_model_id != candidate.logical_model_id
264                || resident.allocation_id != candidate.allocation_id
265        });
266    }
267
268    fn charge_reported_model(
269        &mut self,
270        root: std::path::PathBuf,
271        model_id: &str,
272        allocation_id: String,
273        measured_weights_bytes: u64,
274    ) {
275        let root = car_inference::resource_policy::normalized_state_root_key(&root);
276        if self
277            .residents
278            .iter()
279            .any(|resident| resident.root == root && resident.logical_model_id == model_id)
280        {
281            return;
282        }
283        let Some(coordinator) = car_inference::resource_policy::local_admission_for_scope(&root)
284        else {
285            return;
286        };
287        coordinator.mark_teardown_pending_allocation_with_charge(
288            model_id,
289            &allocation_id,
290            measured_weights_bytes,
291        );
292        self.residents.push(WorkerResidentOwner {
293            root,
294            logical_model_id: model_id.to_string(),
295            allocation_id,
296            coordinator,
297        });
298    }
299
300    fn begin_teardown(&mut self) {
301        if self.teardown_started {
302            return;
303        }
304        self.teardown_started = true;
305        for resident in &self.residents {
306            resident.coordinator.mark_teardown_pending_allocation(
307                &resident.logical_model_id,
308                &resident.allocation_id,
309            );
310        }
311    }
312
313    fn finish_accounting(
314        residents: &[WorkerResidentOwner],
315        resident_models: &std::sync::Mutex<WorkerResidentMap>,
316    ) {
317        let mut tracked = resident_models
318            .lock()
319            .unwrap_or_else(std::sync::PoisonError::into_inner);
320        for resident in residents {
321            tracked.remove(&(resident.root.clone(), resident.logical_model_id.clone()));
322            resident
323                .coordinator
324                .finish_teardown_allocation(&resident.logical_model_id, &resident.allocation_id);
325        }
326    }
327
328    fn confirm_exited(mut self) {
329        self.worker.take();
330        Self::finish_accounting(&self.residents, &self.resident_models);
331        self.residents.clear();
332    }
333
334    fn return_to_slot(mut self, slot: &mut Option<WorkerProc>) {
335        *slot = self.worker.take();
336        self.residents.clear();
337        self.candidate = None;
338    }
339
340    async fn stop_and_confirm(mut self) -> Result<(), InferenceError> {
341        self.begin_teardown();
342        let worker = self.worker_mut();
343        match worker.child.try_wait().map_err(|error| {
344            InferenceError::InferenceFailed(format!(
345                "cannot inspect inference worker during teardown: {error}"
346            ))
347        })? {
348            Some(_) => {}
349            None => {
350                worker.child.kill().await.map_err(|error| {
351                    InferenceError::InferenceFailed(format!(
352                        "cannot stop inference worker during teardown: {error}"
353                    ))
354                })?;
355                worker.child.wait().await.map_err(|error| {
356                    InferenceError::InferenceFailed(format!(
357                        "cannot reap inference worker during teardown: {error}"
358                    ))
359                })?;
360            }
361        }
362        self.confirm_exited();
363        Ok(())
364    }
365}
366
367impl Drop for WorkerProcessGuard {
368    fn drop(&mut self) {
369        self.begin_teardown();
370        let Some(mut worker) = self.worker.take() else {
371            return;
372        };
373        let residents = std::mem::take(&mut self.residents);
374        let resident_models = self.resident_models.clone();
375        let _ = worker.child.start_kill();
376        if tokio::runtime::Handle::try_current().is_ok() {
377            tokio::spawn(async move {
378                if worker.child.wait().await.is_ok() {
379                    Self::finish_accounting(&residents, &resident_models);
380                }
381                // On wait failure the pending-teardown accounting deliberately
382                // remains, blocking an unsafe replacement.
383            });
384        } else {
385            // No executor can confirm exit. Keep fail-closed accounting and
386            // child ownership rather than permit a replacement discount.
387            std::mem::forget(worker);
388        }
389    }
390}
391
392/// How often the idle-worker reaper checks the slot.
393///
394/// A minute is far below any timescale a lingering corpse matters on, and the
395/// check is a non-blocking `try_wait` behind an uncontended lock.
396const IDLE_REAP_INTERVAL: Duration = Duration::from_secs(60);
397
398fn next_worker_allocation_scope() -> u64 {
399    static NEXT: AtomicU64 = AtomicU64::new(1);
400    NEXT.fetch_add(1, Ordering::Relaxed)
401}
402
403/// Installed on the daemon's engine to route on-device generation through a
404/// persistent, auto-respawning worker child. Cheap to clone (shares the child).
405#[derive(Clone)]
406pub struct WorkerOffload {
407    inner: Arc<Mutex<Option<WorkerProc>>>,
408    program: OsString,
409    args: Arc<Vec<OsString>>,
410    policy_generation: Arc<AtomicU64>,
411    resident_models: Arc<std::sync::Mutex<WorkerResidentMap>>,
412    /// Exact request ID → kill-and-wait rendezvous. Only an isolated worker
413    /// installs an entry, so confirmation requires exact process exit.
414    controls: Arc<std::sync::Mutex<HashMap<String, WorkerCancelSender>>>,
415    allocation_scope: u64,
416    #[cfg(test)]
417    release_delay_ms: Arc<AtomicU64>,
418}
419
420impl WorkerOffload {
421    fn allocation_id(&self, model_id: &str) -> String {
422        format!("worker:{}:{model_id}", self.allocation_scope)
423    }
424    /// Offload to a fresh `car-server --mlx-worker` child of the current
425    /// executable. Fails only if the current exe path can't be resolved.
426    pub fn new() -> std::io::Result<Self> {
427        let exe = std::env::current_exe()?;
428        Ok(Self::with_command(
429            exe,
430            vec![OsString::from("--mlx-worker")],
431        ))
432    }
433
434    /// Construct with an explicit worker command. Used by tests to point at a
435    /// stub worker (a script that speaks the same line protocol) so the
436    /// spawn/exchange/respawn logic is exercised without loading a real model.
437    pub fn with_command(program: impl Into<OsString>, args: Vec<OsString>) -> Self {
438        let me = Self {
439            inner: Arc::new(Mutex::new(None)),
440            program: program.into(),
441            args: Arc::new(args),
442            policy_generation: Arc::new(AtomicU64::new(1)),
443            resident_models: Arc::new(std::sync::Mutex::new(HashMap::new())),
444            controls: Arc::new(std::sync::Mutex::new(HashMap::new())),
445            allocation_scope: next_worker_allocation_scope(),
446            #[cfg(test)]
447            release_delay_ms: Arc::new(AtomicU64::new(0)),
448        };
449        me.spawn_idle_reaper();
450        me
451    }
452
453    /// Reap a worker that dies while sitting idle in the slot.
454    ///
455    /// [`Self::take_or_spawn`] already `try_wait`s a slotted corpse, but only
456    /// when the NEXT request arrives — so on a daemon that does no on-device
457    /// inference for days, a worker that died (a Metal abort, an external kill)
458    /// stays a zombie for exactly that long. Observed live: a `--mlx-worker`
459    /// child `<defunct>` for over two days, parented to a healthy daemon. This
460    /// makes the reaping actually proactive, which the module docs already
461    /// claimed. A dead worker is cleared from the slot rather than replaced —
462    /// the next request respawns, matching the existing lazy path.
463    fn spawn_idle_reaper(&self) {
464        if tokio::runtime::Handle::try_current().is_err() {
465            // No runtime (a sync unit test constructing the offloader). The
466            // lazy reap in take_or_spawn still applies.
467            return;
468        }
469        let inner = Arc::clone(&self.inner);
470        let resident_models = Arc::clone(&self.resident_models);
471        tokio::spawn(async move {
472            let mut tick = tokio::time::interval(IDLE_REAP_INTERVAL);
473            tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
474            loop {
475                tick.tick().await;
476                // The offloader is dropped when the daemon shuts down; once no
477                // exchange path holds it, transfer the final slotted child to
478                // the same confirmed-exit cleanup transaction before stopping.
479                if Arc::strong_count(&inner) == 1 {
480                    let mut slot = inner.lock().await;
481                    if let Some(worker) = slot.take() {
482                        drop(WorkerProcessGuard::new(
483                            worker,
484                            resident_models.clone(),
485                            None,
486                        ));
487                    }
488                    return;
489                }
490                let mut slot = inner.lock().await;
491                let Some(p) = slot.as_mut() else { continue };
492                match p.child.try_wait() {
493                    Ok(Some(_)) => {
494                        let dead = slot.take().expect("idle worker checked");
495                        WorkerProcessGuard::new(dead, resident_models.clone(), None)
496                            .confirm_exited();
497                        tracing::info!("reaped a dead idle on-device inference worker");
498                    }
499                    Ok(None) => {}
500                    Err(error) => {
501                        tracing::warn!(%error, "cannot inspect idle inference worker; retaining ownership and residency");
502                    }
503                }
504            }
505        });
506    }
507
508    fn spawn(&self) -> Result<WorkerProc, InferenceError> {
509        let mut cmd = Command::new(&self.program);
510        cmd.args(self.args.iter())
511            .env(WORKER_ENV, "1")
512            .stdin(Stdio::piped())
513            .stdout(Stdio::piped())
514            // Inherit stderr so the worker's tracing/logs land in the daemon's
515            // own stderr (Console.app under CarHost). stdout is protocol-only.
516            .stderr(Stdio::inherit())
517            .kill_on_drop(true);
518        let mut child = cmd.spawn().map_err(|e| {
519            InferenceError::InferenceFailed(format!("failed to spawn inference worker: {e}"))
520        })?;
521        let stdin = child
522            .stdin
523            .take()
524            .ok_or_else(|| InferenceError::InferenceFailed("worker stdin unavailable".into()))?;
525        let stdout = child
526            .stdout
527            .take()
528            .ok_or_else(|| InferenceError::InferenceFailed("worker stdout unavailable".into()))?;
529        Ok(WorkerProc {
530            child,
531            stdin,
532            stdout: BufReader::new(stdout).lines(),
533            policy_generation: self.policy_generation.load(Ordering::Acquire),
534            state_root: None,
535        })
536    }
537
538    /// Take the live worker out of `slot`, spawning one if absent. The caller
539    /// owns the returned worker for the duration of one exchange and must put a
540    /// healthy worker BACK; if it doesn't (a clean drop on cancellation, or an
541    /// error path), `kill_on_drop` reaps the child and `slot` stays empty so the
542    /// next call respawns — which is what makes an interrupted exchange safe
543    /// (no half-read frame can desync a subsequent request).
544    async fn take_or_spawn(
545        &self,
546        slot: &mut Option<WorkerProc>,
547    ) -> Result<WorkerProc, InferenceError> {
548        if let Some(mut p) = slot.take() {
549            // A worker can die while idle in the slot (a crash on the PREVIOUS
550            // request that we already surfaced, or an external kill). Reap it
551            // with a non-blocking wait and spawn fresh instead of handing back a
552            // corpse — otherwise the next request would eat a "broken pipe"
553            // failure just to discover the death. `Ok(None)` = still running.
554            match p.child.try_wait() {
555                Ok(None)
556                    if p.policy_generation == self.policy_generation.load(Ordering::Acquire) =>
557                {
558                    return Ok(p)
559                }
560                Ok(Some(_)) => {
561                    WorkerProcessGuard::new(p, self.resident_models.clone(), None).confirm_exited();
562                }
563                Ok(None) => {
564                    WorkerProcessGuard::new(p, self.resident_models.clone(), None)
565                        .stop_and_confirm()
566                        .await?;
567                }
568                Err(error) => {
569                    drop(WorkerProcessGuard::new(
570                        p,
571                        self.resident_models.clone(),
572                        None,
573                    ));
574                    return Err(InferenceError::InferenceFailed(format!(
575                        "cannot inspect previous inference worker; teardown retained: {error}"
576                    )));
577                }
578            }
579        }
580        self.spawn()
581    }
582
583    async fn take_or_spawn_for_scope(
584        &self,
585        slot: &mut Option<WorkerProc>,
586        state_root: &std::path::Path,
587    ) -> Result<WorkerProc, InferenceError> {
588        let state_root = car_inference::resource_policy::normalized_state_root_key(state_root);
589        let mut worker = self.take_or_spawn(slot).await?;
590        if worker
591            .state_root
592            .as_ref()
593            .is_some_and(|current| current != &state_root)
594        {
595            WorkerProcessGuard::new(worker, self.resident_models.clone(), None)
596                .stop_and_confirm()
597                .await?;
598            worker = self.spawn()?;
599        }
600        worker.state_root = Some(state_root);
601        Ok(worker)
602    }
603
604    fn remember_residency(
605        &self,
606        residency: &WorkerResidencyAck,
607        state_root: std::path::PathBuf,
608        last_used: Option<Instant>,
609    ) {
610        if residency.retention != car_inference::backend_cache::BackendRetention::Resident {
611            return;
612        }
613        let state_root = car_inference::resource_policy::normalized_state_root_key(&state_root);
614        let Some(coordinator) =
615            car_inference::resource_policy::local_admission_for_scope(&state_root)
616        else {
617            tracing::error!(root = %state_root.display(), model = %residency.model_id, "worker reported resident weights without a scoped admission owner");
618            return;
619        };
620        let allocation_id = self.allocation_id(&residency.model_id);
621        self.resident_models
622            .lock()
623            .unwrap_or_else(std::sync::PoisonError::into_inner)
624            .insert(
625                (state_root, residency.model_id.clone()),
626                WorkerResident {
627                    allocation_id,
628                    // The process allocation owns the coordinator. Dropping the
629                    // engine cannot erase machine accounting while the worker
630                    // or its asynchronous reaper still exists.
631                    coordinator,
632                    last_used,
633                },
634            );
635    }
636
637    fn mark_residency_used(
638        resident_models: &std::sync::Mutex<WorkerResidentMap>,
639        state_root: &std::path::Path,
640        model_id: &str,
641    ) {
642        let state_root = car_inference::resource_policy::normalized_state_root_key(state_root);
643        if let Some(resident) = resident_models
644            .lock()
645            .unwrap_or_else(std::sync::PoisonError::into_inner)
646            .get_mut(&(state_root, model_id.to_string()))
647        {
648            resident.last_used = Some(Instant::now());
649        }
650    }
651
652    /// Stop the one worker process represented by a locked slot. Every model
653    /// resident in that process is released as one unit.
654    async fn release_worker_locked(
655        &self,
656        slot: &mut Option<WorkerProc>,
657        model_id: &str,
658    ) -> Result<bool, InferenceError> {
659        let Some(worker) = slot.take() else {
660            return Err(InferenceError::InferenceFailed(format!(
661                "cannot confirm worker exit for resident model {model_id}: worker slot is empty"
662            )));
663        };
664        let mut release = WorkerProcessGuard::new(worker, self.resident_models.clone(), None);
665        release.begin_teardown();
666        #[cfg(test)]
667        tokio::time::sleep(Duration::from_millis(
668            self.release_delay_ms.load(Ordering::Acquire),
669        ))
670        .await;
671        let worker = release.worker_mut();
672        match worker.child.try_wait().map_err(|error| {
673            InferenceError::InferenceFailed(format!(
674                "cannot inspect worker before releasing {model_id}: {error}"
675            ))
676        })? {
677            Some(_) => {}
678            None => {
679                worker.child.kill().await.map_err(|error| {
680                    InferenceError::InferenceFailed(format!(
681                        "cannot stop worker before releasing {model_id}: {error}"
682                    ))
683                })?;
684                worker.child.wait().await.map_err(|error| {
685                    InferenceError::InferenceFailed(format!(
686                        "cannot reap worker before releasing {model_id}: {error}"
687                    ))
688                })?;
689            }
690        }
691        release.confirm_exited();
692        Ok(true)
693    }
694
695    /// Recheck the whole worker immediately after acquiring its in-flight lock.
696    /// A worker is the release unit, so one recent or in-flight resident keeps
697    /// every co-resident alive.
698    async fn release_model_if_idle(
699        &self,
700        model_id: &str,
701        idle_for: Duration,
702    ) -> Result<Vec<String>, InferenceError> {
703        let mut slot = self.inner.lock().await;
704        let resident_ids = {
705            let residents = self
706                .resident_models
707                .lock()
708                .unwrap_or_else(std::sync::PoisonError::into_inner);
709            if !residents
710                .keys()
711                .any(|(_, resident_model)| resident_model == model_id)
712            {
713                return Ok(Vec::new());
714            }
715            let now = Instant::now();
716            if residents.values().any(|resident| {
717                resident.last_used.is_none_or(|last_used| {
718                    now.checked_duration_since(last_used)
719                        .is_none_or(|elapsed| elapsed < idle_for)
720                })
721            }) {
722                return Ok(Vec::new());
723            }
724            let mut ids = residents
725                .keys()
726                .map(|(_, resident_model)| resident_model.clone())
727                .collect::<Vec<_>>();
728            ids.sort();
729            ids.dedup();
730            ids
731        };
732        self.release_worker_locked(&mut slot, model_id).await?;
733        Ok(resident_ids)
734    }
735}
736
737async fn write_line<W, T>(w: &mut W, value: &T) -> std::io::Result<()>
738where
739    W: AsyncWrite + Unpin,
740    T: Serialize,
741{
742    let mut line = serde_json::to_string(value).map_err(std::io::Error::other)?;
743    line.push('\n');
744    w.write_all(line.as_bytes()).await?;
745    w.flush().await
746}
747
748/// Run one non-streaming request/response against a worker. Never touches the
749/// slot — the caller drops a dead worker so the borrow ends first.
750async fn do_generate(
751    proc: &mut WorkerProc,
752    request: GenerateRequest,
753    admission: LocalWorkerAdmission,
754) -> Exchange<(InferenceResult, WorkerResidencyAck)> {
755    let req = WorkerRequest::Generate {
756        request: Box::new(request),
757        admission,
758    };
759    if let Err(e) = write_line(&mut proc.stdin, &req).await {
760        return Exchange::Dead(format!("write to inference worker failed: {e}"));
761    }
762    match proc.stdout.next_line().await {
763        Ok(Some(line)) => match serde_json::from_str::<WorkerResponse>(&line) {
764            Ok(WorkerResponse::Result { result, residency }) => Exchange::Ok((*result, residency)),
765            Ok(WorkerResponse::Error(msg)) => {
766                Exchange::Reported(InferenceError::InferenceFailed(msg))
767            }
768            Ok(WorkerResponse::LocalResourceBlocked {
769                preflight,
770                recovery,
771            }) => Exchange::Reported(InferenceError::LocalResourceBlocked {
772                preflight,
773                recovery,
774            }),
775            Ok(_) => Exchange::Dead("inference worker sent an unexpected response frame".into()),
776            Err(e) => Exchange::Dead(format!("inference worker sent invalid JSON: {e}")),
777        },
778        Ok(None) => Exchange::Dead("inference worker exited mid-request (EOF)".into()),
779        Err(e) => Exchange::Dead(format!("read from inference worker failed: {e}")),
780    }
781}
782
783#[async_trait::async_trait]
784impl LocalGenerationOffload for WorkerOffload {
785    async fn generate(&self, _request: GenerateRequest) -> Result<InferenceResult, InferenceError> {
786        Err(InferenceError::InferenceFailed(
787            "WorkerOffload requires admission-aware dispatch".into(),
788        ))
789    }
790
791    async fn stream(
792        &self,
793        _request: GenerateRequest,
794    ) -> Result<tokio::sync::mpsc::Receiver<StreamEvent>, InferenceError> {
795        Err(InferenceError::InferenceFailed(
796            "WorkerOffload requires admission-aware dispatch".into(),
797        ))
798    }
799
800    async fn generate_admitted(
801        &self,
802        request: GenerateRequest,
803        admission: LocalWorkerAdmission,
804    ) -> Result<LocalOffloadResult, InferenceError> {
805        // Hold the lock across the exchange (serializes worker access, matching
806        // the in-process Metal device lock). Take the worker OUT so a cancelled
807        // future (e.g. handle_infer's timeout firing) drops it — kill_on_drop
808        // reaps the child and the slot stays empty, so no orphaned response
809        // frame can desync the next request.
810        let mut slot = self.inner.lock().await;
811        let expected_model = request.model.clone().unwrap_or_default();
812        let candidate = (
813            admission.state_root.clone(),
814            expected_model.clone(),
815            self.allocation_id(&expected_model),
816        );
817        let proc = self
818            .take_or_spawn_for_scope(&mut slot, &admission.state_root)
819            .await?;
820        let mut ownership =
821            WorkerProcessGuard::new(proc, self.resident_models.clone(), Some(candidate));
822        ownership.charge_candidate(admission.measured_weights_bytes);
823        let state_root = admission.state_root.clone();
824        enum GenerateSelection {
825            Exchange(Exchange<(InferenceResult, WorkerResidencyAck)>),
826            Cancel(oneshot::Sender<bool>),
827        }
828        let inference_id = car_inference::current_inference_control_id();
829        let termination = car_inference::current_controlled_termination_token();
830        let (control_registration, mut cancel_rx) = inference_id
831            .map(|id| WorkerControlRegistration::install(self.controls.clone(), id))
832            .unzip();
833        let selected = if let Some(cancel_rx) = cancel_rx.as_mut() {
834            let worker = ownership.worker_mut();
835            tokio::select! {
836                exchange = do_generate(worker, request, admission) => GenerateSelection::Exchange(exchange),
837                request = cancel_rx.recv() => match request {
838                    Some(ack) => GenerateSelection::Cancel(ack),
839                    None => unreachable!("control registration owns the sender"),
840                },
841            }
842        } else {
843            GenerateSelection::Exchange(
844                do_generate(ownership.worker_mut(), request, admission).await,
845            )
846        };
847        let exchange = match selected {
848            GenerateSelection::Exchange(exchange) => exchange,
849            GenerateSelection::Cancel(ack) => {
850                let confirmed = ownership.stop_and_confirm().await.is_ok();
851                if confirmed {
852                    if let Some(token) = termination.as_ref() {
853                        token.confirm_exact_backend_termination();
854                    }
855                }
856                let _ = ack.send(confirmed);
857                drop(control_registration);
858                return Err(if confirmed {
859                    InferenceError::ControlledTermination
860                } else {
861                    InferenceError::InferenceFailed(
862                        "isolated inference worker termination was not confirmed".into(),
863                    )
864                });
865            }
866        };
867        drop(control_registration);
868        match exchange {
869            Exchange::Ok((ir, residency)) => {
870                if residency.model_id != expected_model {
871                    ownership.charge_reported_model(
872                        state_root,
873                        &residency.model_id,
874                        self.allocation_id(&residency.model_id),
875                        residency.measured_weights_bytes,
876                    );
877                    drop(ownership);
878                    return Err(InferenceError::InferenceFailed(format!(
879                        "local worker acknowledged model '{}' for requested '{}'",
880                        residency.model_id, expected_model
881                    )));
882                }
883                if residency.retention != car_inference::backend_cache::BackendRetention::Resident {
884                    ownership.clear_candidate();
885                }
886                self.remember_residency(&residency, state_root, Some(Instant::now()));
887                ownership.return_to_slot(&mut slot);
888                Ok(LocalOffloadResult {
889                    result: ir,
890                    residency: LocalWorkerResidency {
891                        model_id: residency.model_id,
892                        measured_weights_bytes: residency.measured_weights_bytes,
893                    },
894                    retention: residency.retention,
895                })
896            }
897            Exchange::Reported(error) => {
898                if matches!(error, InferenceError::LocalResourceBlocked { .. }) {
899                    // The worker explicitly rejected before allocating.
900                    ownership.clear_candidate();
901                    ownership.return_to_slot(&mut slot);
902                } else {
903                    // A generic error carries no retention acknowledgement.
904                    // The worker may have loaded/cache-published before the
905                    // request failed, so tear it down before releasing charge.
906                    drop(ownership);
907                }
908                Err(error)
909            }
910            Exchange::Dead(msg) => {
911                // Preserve the OS verdict before guard teardown takes the child.
912                // EOF can arrive just before wait becomes ready; keep this bounded
913                // because a protocol failure need not mean the process exited.
914                let exit = match tokio::time::timeout(
915                    Duration::from_millis(100),
916                    ownership.worker_mut().child.wait(),
917                )
918                .await
919                {
920                    Ok(Ok(status)) => format!("; worker {status}"),
921                    Ok(Err(error)) => format!("; exit status unavailable: {error}"),
922                    Err(_) => "; worker exit not yet observed".to_string(),
923                };
924                // Worker died (very likely the Metal/MLX abort we isolate for).
925                // `proc` drops here (kill_on_drop), slot stays None → next call
926                // respawns. Fail THIS rpc only; the daemon stays up.
927                drop(ownership);
928                Err(InferenceError::InferenceFailed(format!(
929                    "on-device inference worker connection failed; \
930                     this request failed but the daemon is up: {msg}{exit}. \
931                     The next request starts a new worker."
932                )))
933            }
934        }
935    }
936
937    async fn stream_admitted(
938        &self,
939        request: GenerateRequest,
940        admission: LocalWorkerAdmission,
941    ) -> Result<LocalOffloadStream, InferenceError> {
942        // Hold the lock for the whole stream via an owned guard moved into the
943        // reader task (serializes streams against other worker traffic, as the
944        // in-process Metal device lock did). The worker is taken OUT of the slot
945        // and only put BACK on a clean `StreamEnd` — any other exit (caller
946        // drops the receiver mid-stream, a worker error, EOF) leaves the worker
947        // mid-protocol, so it's killed and the next request respawns rather than
948        // reading leftover frames.
949        let mut guard = self.inner.clone().lock_owned().await;
950        let expected_model = request.model.clone().unwrap_or_default();
951        let candidate = (
952            admission.state_root.clone(),
953            expected_model.clone(),
954            self.allocation_id(&expected_model),
955        );
956        let proc = self
957            .take_or_spawn_for_scope(&mut guard, &admission.state_root)
958            .await?;
959        let mut ownership =
960            WorkerProcessGuard::new(proc, self.resident_models.clone(), Some(candidate));
961        ownership.charge_candidate(admission.measured_weights_bytes);
962        let state_root = admission.state_root.clone();
963        let inference_id = car_inference::current_inference_control_id();
964        let termination = car_inference::current_controlled_termination_token();
965        let (control_registration, mut cancel_rx) = inference_id
966            .map(|id| WorkerControlRegistration::install(self.controls.clone(), id))
967            .unzip();
968        // Send the stream request synchronously so a spawn/write failure
969        // surfaces as an Err here rather than a silently-empty stream.
970        {
971            let req = WorkerRequest::Stream {
972                request: Box::new(request),
973                admission,
974            };
975            if let Err(e) = write_line(&mut ownership.worker_mut().stdin, &req).await {
976                return Err(InferenceError::InferenceFailed(format!(
977                    "write to inference worker failed: {e}"
978                )));
979            }
980        }
981
982        enum StreamStartSelection {
983            Frame(std::io::Result<Option<String>>),
984            Cancel(oneshot::Sender<bool>),
985        }
986        // Cancellation is live as soon as the exact worker request is written,
987        // including while the child is still loading and has not emitted
988        // StreamStarted. Confirmation still requires kill + wait below.
989        let stream_start = if let Some(cancel_rx) = cancel_rx.as_mut() {
990            let worker = ownership.worker_mut();
991            tokio::select! {
992                frame = worker.stdout.next_line() => StreamStartSelection::Frame(frame),
993                request = cancel_rx.recv() => match request {
994                    Some(ack) => StreamStartSelection::Cancel(ack),
995                    None => unreachable!("control registration owns the sender"),
996                },
997            }
998        } else {
999            StreamStartSelection::Frame(ownership.worker_mut().stdout.next_line().await)
1000        };
1001        let stream_start = match stream_start {
1002            StreamStartSelection::Frame(frame) => frame,
1003            StreamStartSelection::Cancel(ack) => {
1004                let confirmed = ownership.stop_and_confirm().await.is_ok();
1005                if confirmed {
1006                    if let Some(token) = termination.as_ref() {
1007                        token.confirm_exact_backend_termination();
1008                    }
1009                }
1010                let _ = ack.send(confirmed);
1011                return Err(if confirmed {
1012                    InferenceError::ControlledTermination
1013                } else {
1014                    InferenceError::InferenceFailed(
1015                        "isolated inference worker termination before StreamStarted was not confirmed"
1016                            .into(),
1017                    )
1018                });
1019            }
1020        };
1021        let residency = match stream_start {
1022            Ok(Some(line)) => match serde_json::from_str::<WorkerResponse>(&line) {
1023                Ok(WorkerResponse::StreamStarted { residency }) => residency,
1024                Ok(WorkerResponse::LocalResourceBlocked {
1025                    preflight,
1026                    recovery,
1027                }) => {
1028                    ownership.clear_candidate();
1029                    ownership.return_to_slot(&mut guard);
1030                    return Err(InferenceError::LocalResourceBlocked {
1031                        preflight,
1032                        recovery,
1033                    });
1034                }
1035                Ok(WorkerResponse::Error(message)) => {
1036                    // No retention ACK: fail closed by tearing down the worker.
1037                    drop(ownership);
1038                    return Err(InferenceError::InferenceFailed(message));
1039                }
1040                Ok(_) => {
1041                    return Err(InferenceError::InferenceFailed(
1042                        "inference worker streamed before a successful load acknowledgement".into(),
1043                    ));
1044                }
1045                Err(error) => {
1046                    return Err(InferenceError::InferenceFailed(format!(
1047                        "inference worker sent invalid stream acknowledgement: {error}"
1048                    )));
1049                }
1050            },
1051            Ok(None) => {
1052                return Err(InferenceError::InferenceFailed(
1053                    "inference worker exited before loading the streaming model".into(),
1054                ));
1055            }
1056            Err(error) => {
1057                return Err(InferenceError::InferenceFailed(format!(
1058                    "failed reading inference worker load acknowledgement: {error}"
1059                )));
1060            }
1061        };
1062        if residency.model_id != expected_model {
1063            ownership.charge_reported_model(
1064                state_root,
1065                &residency.model_id,
1066                self.allocation_id(&residency.model_id),
1067                residency.measured_weights_bytes,
1068            );
1069            drop(ownership);
1070            return Err(InferenceError::InferenceFailed(format!(
1071                "local worker acknowledged model '{}' for requested '{}'",
1072                residency.model_id, expected_model
1073            )));
1074        }
1075        if residency.retention != car_inference::backend_cache::BackendRetention::Resident {
1076            ownership.clear_candidate();
1077        }
1078        self.remember_residency(&residency, state_root.clone(), None);
1079        let resident_models = self.resident_models.clone();
1080        let resident_model_id = residency.model_id.clone();
1081
1082        let (tx, rx) = tokio::sync::mpsc::channel::<StreamEvent>(64);
1083        tokio::spawn(async move {
1084            let _control_registration = control_registration;
1085            let mut ownership = Some(ownership);
1086            let mut clean = false;
1087            loop {
1088                enum StreamSelection {
1089                    Frame(std::io::Result<Option<String>>),
1090                    Cancel(oneshot::Sender<bool>),
1091                }
1092                let selected = if let Some(cancel_rx) = cancel_rx.as_mut() {
1093                    let worker = ownership
1094                        .as_mut()
1095                        .expect("worker ownership remains until terminal")
1096                        .worker_mut();
1097                    tokio::select! {
1098                        frame = worker.stdout.next_line() => StreamSelection::Frame(frame),
1099                        request = cancel_rx.recv() => match request {
1100                            Some(ack) => StreamSelection::Cancel(ack),
1101                            None => unreachable!("control registration owns the sender"),
1102                        },
1103                    }
1104                } else {
1105                    StreamSelection::Frame(
1106                        ownership
1107                            .as_mut()
1108                            .expect("worker ownership remains until terminal")
1109                            .worker_mut()
1110                            .stdout
1111                            .next_line()
1112                            .await,
1113                    )
1114                };
1115                let frame = match selected {
1116                    StreamSelection::Frame(frame) => frame,
1117                    StreamSelection::Cancel(ack) => {
1118                        let confirmed = ownership
1119                            .take()
1120                            .expect("worker ownership remains until terminal")
1121                            .stop_and_confirm()
1122                            .await
1123                            .is_ok();
1124                        if confirmed {
1125                            if let Some(token) = termination.as_ref() {
1126                                token.confirm_exact_backend_termination();
1127                            }
1128                        }
1129                        let _ = ack.send(confirmed);
1130                        return;
1131                    }
1132                };
1133                match frame {
1134                    Ok(Some(line)) => match serde_json::from_str::<WorkerResponse>(&line) {
1135                        Ok(WorkerResponse::Event(ev)) => {
1136                            if tx.send(*ev).await.is_err() {
1137                                break; // caller dropped the receiver, worker is mid-stream
1138                            }
1139                        }
1140                        Ok(WorkerResponse::StreamEnd) => {
1141                            clean = true;
1142                            break;
1143                        }
1144                        Ok(WorkerResponse::Error(msg)) => {
1145                            tracing::warn!(error = %msg, "inference worker stream error");
1146                            // Surface as a mid-stream failure the accumulator
1147                            // reads as an error end (mirrors remote.rs).
1148                            let _ = tx.send(StreamEvent::StopReason("error".into())).await;
1149                            break;
1150                        }
1151                        Ok(WorkerResponse::Result { .. })
1152                        | Ok(WorkerResponse::StreamStarted { .. })
1153                        | Ok(WorkerResponse::LocalResourceBlocked { .. })
1154                        | Err(_) => {
1155                            // Protocol break — treat as a dead worker.
1156                            let _ = tx.send(StreamEvent::StopReason("error".into())).await;
1157                            break;
1158                        }
1159                    },
1160                    Ok(None) | Err(_) => {
1161                        // Worker died mid-stream: emit an error end.
1162                        let _ = tx.send(StreamEvent::StopReason("error".into())).await;
1163                        break;
1164                    }
1165                }
1166            }
1167            if clean {
1168                // StreamStarted marks the resident as not-yet-used; only a
1169                // clean StreamEnd makes it eligible for later idle eviction.
1170                Self::mark_residency_used(&resident_models, &state_root, &resident_model_id);
1171                // Clean end: return the healthy worker to the slot for reuse.
1172                ownership
1173                    .take()
1174                    .expect("worker ownership remains on clean end")
1175                    .return_to_slot(&mut guard);
1176            } else {
1177                // Interrupted / dead: kill so the next request gets a fresh
1178                // worker (slot is already None). `guard` drops, unlocking.
1179                drop(ownership.take());
1180            }
1181        });
1182        Ok(LocalOffloadStream {
1183            events: rx,
1184            residency: LocalWorkerResidency {
1185                model_id: residency.model_id,
1186                measured_weights_bytes: residency.measured_weights_bytes,
1187            },
1188            retention: residency.retention,
1189        })
1190    }
1191
1192    fn refresh_resource_policy(&self, generation: u64) {
1193        self.policy_generation.store(generation, Ordering::Release);
1194        if let Ok(mut slot) = self.inner.try_lock() {
1195            if let Some(worker) = slot.take() {
1196                drop(WorkerProcessGuard::new(
1197                    worker,
1198                    self.resident_models.clone(),
1199                    None,
1200                ));
1201            }
1202        }
1203    }
1204
1205    fn resident_allocation_id(&self, model_id: &str) -> Option<String> {
1206        Some(self.allocation_id(model_id))
1207    }
1208
1209    async fn resident_models(&self) -> Vec<String> {
1210        self.resident_models
1211            .lock()
1212            .unwrap_or_else(std::sync::PoisonError::into_inner)
1213            .keys()
1214            .map(|(_, model_id)| model_id.clone())
1215            .collect()
1216    }
1217
1218    async fn release_model(&self, model_id: &str) -> Result<bool, InferenceError> {
1219        let mut slot = self.inner.lock().await;
1220        let resident = self
1221            .resident_models
1222            .lock()
1223            .unwrap_or_else(std::sync::PoisonError::into_inner)
1224            .keys()
1225            .any(|(_, resident_model)| resident_model == model_id);
1226        if !resident {
1227            return Ok(false);
1228        }
1229        self.release_worker_locked(&mut slot, model_id).await
1230    }
1231
1232    async fn release_idle_models(&self, idle_for: Duration) -> Vec<String> {
1233        // This first read chooses only a candidate. Eligibility is deliberately
1234        // decided later by release_model_if_idle while the worker's in-flight
1235        // mutex is held, immediately before process teardown.
1236        let candidate = self
1237            .resident_models
1238            .lock()
1239            .unwrap_or_else(std::sync::PoisonError::into_inner)
1240            .keys()
1241            .next()
1242            .map(|(_, model_id)| model_id.clone());
1243        let Some(candidate) = candidate else {
1244            return Vec::new();
1245        };
1246        match self.release_model_if_idle(&candidate, idle_for).await {
1247            Ok(released) => released,
1248            Err(error) => {
1249                tracing::warn!(%error, "failed to release idle on-device inference worker");
1250                Vec::new()
1251            }
1252        }
1253    }
1254
1255    async fn terminate_inference(&self, inference_id: &str) -> InferenceTerminationAck {
1256        let sender = self
1257            .controls
1258            .lock()
1259            .unwrap_or_else(std::sync::PoisonError::into_inner)
1260            .get(inference_id)
1261            .cloned();
1262        let Some(sender) = sender else {
1263            return InferenceTerminationAck::Unconfirmed;
1264        };
1265        let (ack_tx, ack_rx) = oneshot::channel();
1266        if sender.send(ack_tx).is_err() {
1267            return InferenceTerminationAck::Unconfirmed;
1268        }
1269        match ack_rx.await {
1270            Ok(true) => InferenceTerminationAck::Confirmed,
1271            Ok(false) | Err(_) => InferenceTerminationAck::Unconfirmed,
1272        }
1273    }
1274}
1275
1276/// The worker child's main loop. Builds a real in-process engine (persistence
1277/// disabled via the `CAR_INFERENCE_WORKER` env the parent set — see
1278/// [`car_inference::is_offload_worker`]) and services line-delimited requests
1279/// off stdin, writing responses to stdout. Logs go to stderr (via the tracing
1280/// subscriber the binary installs), keeping stdout protocol-clean.
1281///
1282/// Returns when stdin closes (the daemon dropped the worker). One request is
1283/// served to completion before the next is read — the daemon never pipelines,
1284/// and this matches the single-Metal-device serialization the in-process path
1285/// relied on.
1286pub async fn run_mlx_worker() {
1287    // Belt-and-suspenders: the parent sets this, but guarantee the self-offload
1288    // / no-persist guard holds even if a worker is launched some other way.
1289    std::env::set_var(WORKER_ENV, "1");
1290
1291    let mut engine: Option<(u64, std::path::PathBuf, Arc<InferenceEngine>)> = None;
1292    let mut lines = BufReader::new(tokio::io::stdin()).lines();
1293    let mut stdout = tokio::io::stdout();
1294
1295    while let Ok(Some(line)) = lines.next_line().await {
1296        let line = line.trim();
1297        if line.is_empty() {
1298            continue;
1299        }
1300        let req: WorkerRequest = match serde_json::from_str(line) {
1301            Ok(r) => r,
1302            Err(e) => {
1303                let _ = write_line(
1304                    &mut stdout,
1305                    &WorkerResponse::Error(format!("malformed worker request: {e}")),
1306                )
1307                .await;
1308                continue;
1309            }
1310        };
1311        let admission = match &req {
1312            WorkerRequest::Generate { admission, .. } | WorkerRequest::Stream { admission, .. } => {
1313                admission
1314            }
1315        };
1316        let recreate = engine.as_ref().is_none_or(|(generation, root, _)| {
1317            *generation != admission.policy_generation || root != &admission.state_root
1318        });
1319        if recreate {
1320            let mut config = InferenceConfig::default();
1321            config.state_root = admission.state_root.clone();
1322            let candidate = Arc::new(InferenceEngine::new(config));
1323            candidate.apply_local_resource_policy(admission.policy.clone());
1324            engine = Some((
1325                admission.policy_generation,
1326                admission.state_root.clone(),
1327                candidate,
1328            ));
1329        }
1330        let active_engine = Arc::clone(&engine.as_ref().expect("worker engine initialized").2);
1331        active_engine.apply_local_resource_policy(admission.policy.clone());
1332
1333        match req {
1334            WorkerRequest::Generate {
1335                request,
1336                admission: _,
1337            } => {
1338                // LOCAL_ADMISSION_BOUNDARY:worker-side-allocation
1339                // A worker owns a fresh engine, so explicit local dispatch
1340                // re-probes live memory immediately before MLX/Candle loads.
1341                // Daemon-side admission is approval to attempt, not a stale
1342                // permit that bypasses this second check.
1343                let model_id = request.model.clone().unwrap_or_default();
1344                let resp = match active_engine.generate_tracked(*request).await {
1345                    Ok(ir) => {
1346                        let measured = measured_worker_model_bytes(&active_engine, &model_id);
1347                        let retention = worker_model_retention(&active_engine, &model_id);
1348                        WorkerResponse::Result {
1349                            result: Box::new(ir),
1350                            residency: WorkerResidencyAck {
1351                                model_id,
1352                                measured_weights_bytes: measured,
1353                                retention,
1354                            },
1355                        }
1356                    }
1357                    Err(InferenceError::LocalResourceBlocked {
1358                        preflight,
1359                        recovery,
1360                    }) => WorkerResponse::LocalResourceBlocked {
1361                        preflight,
1362                        recovery,
1363                    },
1364                    Err(e) => WorkerResponse::Error(e.to_string()),
1365                };
1366                if write_line(&mut stdout, &resp).await.is_err() {
1367                    break; // parent gone
1368                }
1369            }
1370            WorkerRequest::Stream {
1371                request,
1372                admission: _,
1373            } => {
1374                let model_id = request.model.clone().unwrap_or_default();
1375                match active_engine.generate_tracked_stream(*request).await {
1376                    Ok(mut tracked) => {
1377                        let measured = measured_worker_model_bytes(&active_engine, &model_id);
1378                        let retention = worker_model_retention(&active_engine, &model_id);
1379                        if write_line(
1380                            &mut stdout,
1381                            &WorkerResponse::StreamStarted {
1382                                residency: WorkerResidencyAck {
1383                                    model_id,
1384                                    measured_weights_bytes: measured,
1385                                    retention,
1386                                },
1387                            },
1388                        )
1389                        .await
1390                        .is_err()
1391                        {
1392                            break;
1393                        }
1394                        let mut broke = false;
1395                        while let Some(ev) = tracked.events.recv().await {
1396                            if write_line(&mut stdout, &WorkerResponse::Event(Box::new(ev)))
1397                                .await
1398                                .is_err()
1399                            {
1400                                broke = true;
1401                                break;
1402                            }
1403                        }
1404                        if broke {
1405                            break;
1406                        }
1407                        if write_line(&mut stdout, &WorkerResponse::StreamEnd)
1408                            .await
1409                            .is_err()
1410                        {
1411                            break;
1412                        }
1413                    }
1414                    Err(InferenceError::LocalResourceBlocked {
1415                        preflight,
1416                        recovery,
1417                    }) => {
1418                        if write_line(
1419                            &mut stdout,
1420                            &WorkerResponse::LocalResourceBlocked {
1421                                preflight,
1422                                recovery,
1423                            },
1424                        )
1425                        .await
1426                        .is_err()
1427                        {
1428                            break;
1429                        }
1430                    }
1431                    Err(e) => {
1432                        if write_line(&mut stdout, &WorkerResponse::Error(e.to_string()))
1433                            .await
1434                            .is_err()
1435                        {
1436                            break;
1437                        }
1438                    }
1439                }
1440            }
1441        }
1442    }
1443}
1444
1445fn measured_worker_model_bytes(engine: &InferenceEngine, model_id: &str) -> u64 {
1446    let Some(schema) = engine
1447        .unified_registry
1448        .get(model_id)
1449        .or_else(|| engine.unified_registry.find_by_name(model_id))
1450    else {
1451        return 0;
1452    };
1453    car_inference::backend_cache::estimate_model_size(&engine.config.models_dir.join(&schema.name))
1454}
1455
1456fn worker_model_retention(
1457    engine: &InferenceEngine,
1458    model_id: &str,
1459) -> car_inference::backend_cache::BackendRetention {
1460    engine.local_model_retention(model_id)
1461}
1462
1463#[cfg(test)]
1464mod tests {
1465    use super::*;
1466
1467    #[test]
1468    fn local_model_preflight_worker_rechecks_before_allocation() {
1469        let source = include_str!("inference_worker.rs");
1470        assert!(source.contains("LOCAL_ADMISSION_BOUNDARY:worker-side-allocation"));
1471        assert!(source.contains("active_engine.generate_tracked(*request).await"));
1472    }
1473
1474    #[test]
1475    fn idle_backend_eviction_calls_worker_idle_release() {
1476        let source = include_str!("handler.rs");
1477        assert!(source.contains("release_idle_models"));
1478    }
1479
1480    // Externally-tagged serde shapes the stub workers below emit. Locking these
1481    // down as literals doubles as a wire-format regression guard.
1482    const RESULT_LINE: &str = r#"{"Result":{"result":{"text":"pong","tool_calls":[],"trace_id":"t","model_used":"stub","latency_ms":1,"usage":null},"residency":{"model_id":"stub","measured_weights_bytes":1,"retention":"resident"}}}"#;
1483    // These three feed only the `#[cfg(unix)]` stub-worker tests below (they
1484    // drive a `sh` script), so on Windows they are dead code under `-D warnings`.
1485    // Gate them the same way `sh` and `stub_request` already are.
1486    #[cfg(unix)]
1487    const TRANSIENT_RESULT_LINE: &str = r#"{"Result":{"result":{"text":"pong","tool_calls":[],"trace_id":"t","model_used":"stub","latency_ms":1,"usage":null},"residency":{"model_id":"stub","measured_weights_bytes":1,"retention":"transient"}}}"#;
1488    #[cfg(unix)]
1489    const STREAM_STARTED_LINE: &str = r#"{"StreamStarted":{"residency":{"model_id":"stub","measured_weights_bytes":1,"retention":"resident"}}}"#;
1490
1491    #[cfg(unix)]
1492    fn test_admission() -> LocalWorkerAdmission {
1493        LocalWorkerAdmission {
1494            policy: car_inference::ResourcePolicy::custom_gb(8.0).unwrap(),
1495            policy_generation: 1,
1496            state_root: std::path::PathBuf::from("/tmp/car-worker-root"),
1497            measured_weights_bytes: 1,
1498        }
1499    }
1500
1501    #[test]
1502    fn envelope_serde_round_trips() {
1503        // Request
1504        let admission = car_inference::LocalWorkerAdmission {
1505            policy: car_inference::ResourcePolicy::custom_gb(8.0).unwrap(),
1506            policy_generation: 7,
1507            state_root: std::path::PathBuf::from("/tmp/car-worker-root"),
1508            measured_weights_bytes: 3 * 1024 * 1024 * 1024,
1509        };
1510        let req = WorkerRequest::Generate {
1511            request: Box::new(GenerateRequest {
1512                prompt: "hi".into(),
1513                ..Default::default()
1514            }),
1515            admission: admission.clone(),
1516        };
1517        let json = serde_json::to_string(&req).unwrap();
1518        assert!(json.starts_with(r#"{"Generate":"#), "got {json}");
1519        let decoded = serde_json::from_str::<WorkerRequest>(&json).unwrap();
1520        let WorkerRequest::Generate {
1521            admission: decoded_admission,
1522            ..
1523        } = decoded
1524        else {
1525            panic!("wrong request variant")
1526        };
1527        assert_eq!(decoded_admission, admission);
1528
1529        // Responses: Result, Event(StreamEvent), StreamEnd, Error.
1530        assert!(matches!(
1531            serde_json::from_str::<WorkerResponse>(RESULT_LINE).unwrap(),
1532            WorkerResponse::Result { .. }
1533        ));
1534        let ev = WorkerResponse::Event(Box::new(StreamEvent::TextDelta("hi".into())));
1535        let ev_json = serde_json::to_string(&ev).unwrap();
1536        assert_eq!(ev_json, r#"{"Event":{"TextDelta":"hi"}}"#);
1537        assert_eq!(
1538            serde_json::to_string(&WorkerResponse::StreamEnd).unwrap(),
1539            r#""StreamEnd""#
1540        );
1541        // The StopReason variant the WS-runner JSON mapping omits must survive.
1542        let sr = StreamEvent::StopReason("length".into());
1543        let sr2: StreamEvent = serde_json::from_str(&serde_json::to_string(&sr).unwrap()).unwrap();
1544        assert!(matches!(sr2, StreamEvent::StopReason(s) if s == "length"));
1545    }
1546
1547    #[test]
1548    fn worker_protocol_preserves_structured_resource_rejection_and_residency_ack() {
1549        let preflight = car_inference::LocalLoadPreflight {
1550            model_id: "mlx/test".into(),
1551            estimate: car_inference::ModelMemoryEstimate {
1552                weights_mb: 6144,
1553                runtime_overhead_mb: 256,
1554                context_overhead_mb: 128,
1555                transient_margin_mb: 256,
1556                estimated_peak_mb: 6784,
1557                evidence: car_inference::ModelResourceEvidence::FileSystemMeasured,
1558            },
1559            configured_ceiling_mb: 4096,
1560            resident_model_mb: 0,
1561            active_reservations_mb: 0,
1562            estimated_incremental_mb: 6144,
1563            accelerator_total_mb: None,
1564            accelerator_resident_mb: None,
1565            accelerator_incremental_mb: None,
1566            live_available_mb: Some(8192),
1567            emergency_reserve_mb: 4096,
1568            verdict: car_inference::LocalLoadVerdict::ExceedsConfiguredCeiling,
1569        };
1570        let blocked = WorkerResponse::LocalResourceBlocked {
1571            preflight: preflight.clone(),
1572            recovery: "choose a smaller model".into(),
1573        };
1574        let decoded: WorkerResponse =
1575            serde_json::from_str(&serde_json::to_string(&blocked).unwrap()).unwrap();
1576        assert!(matches!(
1577            decoded,
1578            WorkerResponse::LocalResourceBlocked {
1579                preflight: actual,
1580                ..
1581            } if actual == preflight
1582        ));
1583
1584        let started = WorkerResponse::StreamStarted {
1585            residency: WorkerResidencyAck {
1586                model_id: "mlx/test".into(),
1587                measured_weights_bytes: 3 * 1024 * 1024 * 1024,
1588                retention: car_inference::backend_cache::BackendRetention::Resident,
1589            },
1590        };
1591        assert!(matches!(
1592            serde_json::from_str::<WorkerResponse>(&serde_json::to_string(&started).unwrap())
1593                .unwrap(),
1594            WorkerResponse::StreamStarted { .. }
1595        ));
1596    }
1597
1598    // The subprocess-driven tests use a POSIX shell stub as the "worker", so
1599    // they exercise the real spawn / stdio-protocol / respawn paths without a
1600    // model. Unix-only (Windows CI has no `sh`); the feature is macOS-primary.
1601    #[cfg(unix)]
1602    fn sh(script: &str) -> WorkerOffload {
1603        WorkerOffload::with_command("sh", vec![OsString::from("-c"), OsString::from(script)])
1604    }
1605
1606    #[cfg(unix)]
1607    fn stub_request() -> GenerateRequest {
1608        GenerateRequest {
1609            model: Some("stub".into()),
1610            ..Default::default()
1611        }
1612    }
1613
1614    #[cfg(unix)]
1615    #[tokio::test]
1616    async fn generate_round_trips_and_reuses_the_worker() {
1617        // A stub that answers every request line with one Result line and stays
1618        // alive — so a second call must reuse the same worker.
1619        let off = sh(&format!(
1620            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
1621        ));
1622        let r1 = off
1623            .generate_admitted(stub_request(), test_admission())
1624            .await
1625            .unwrap();
1626        assert_eq!(r1.result.text, "pong");
1627        let r2 = off
1628            .generate_admitted(stub_request(), test_admission())
1629            .await
1630            .unwrap();
1631        assert_eq!(r2.result.text, "pong");
1632    }
1633
1634    #[cfg(unix)]
1635    #[tokio::test]
1636    async fn exact_worker_cancel_confirms_only_after_kill_and_wait() {
1637        let off = sh("while IFS= read -r line; do sleep 60; done");
1638        let caller = off.clone();
1639        let generation = tokio::spawn(async move {
1640            car_inference::scope_inference_control_id("exact-worker-id".to_string(), async move {
1641                caller
1642                    .generate_admitted(stub_request(), test_admission())
1643                    .await
1644            })
1645            .await
1646        });
1647
1648        tokio::time::timeout(Duration::from_secs(2), async {
1649            loop {
1650                if off
1651                    .controls
1652                    .lock()
1653                    .unwrap_or_else(std::sync::PoisonError::into_inner)
1654                    .contains_key("exact-worker-id")
1655                {
1656                    break;
1657                }
1658                tokio::task::yield_now().await;
1659            }
1660        })
1661        .await
1662        .expect("worker request must register its exact control ID");
1663
1664        assert_eq!(
1665            off.terminate_inference("exact-worker-id").await,
1666            InferenceTerminationAck::Confirmed
1667        );
1668        assert!(generation.await.unwrap().is_err());
1669        assert!(off.inner.lock().await.is_none());
1670    }
1671
1672    #[cfg(unix)]
1673    #[tokio::test]
1674    async fn transient_worker_ack_never_creates_parent_residency() {
1675        let off = sh(&format!(
1676            "while IFS= read -r line; do printf '%s\\n' '{TRANSIENT_RESULT_LINE}'; done"
1677        ));
1678        let result = off
1679            .generate_admitted(stub_request(), test_admission())
1680            .await
1681            .unwrap();
1682
1683        assert_eq!(
1684            result.retention,
1685            car_inference::backend_cache::BackendRetention::Transient
1686        );
1687        assert!(off.resident_models().await.is_empty());
1688    }
1689
1690    /// Serializes the tests that charge the process-wide machine admission
1691    /// ledger.
1692    ///
1693    /// `scoped_local_admission` keys its coordinator per `state_root`, which
1694    /// reads as isolation but is not: the coordinator is built on
1695    /// `process_machine_admission_ledger()`, which is process-wide **by design**
1696    /// — a machine does not get more RAM because a caller used a different state
1697    /// root. So residency totals are shared across every test in this binary
1698    /// that charges concurrently.
1699    ///
1700    /// A delta measurement is not enough. An interfering charge can land between
1701    /// the baseline read and the assertion, which is why the first attempt at
1702    /// this fix passed once and then failed again. Exclusion is the only thing
1703    /// that actually holds.
1704    ///
1705    /// `cfg(unix)` to match its users: every test that takes it is unix-gated,
1706    /// so on Windows the static is genuinely unused and `-D warnings` fails the
1707    /// build. This surfaced only at release time because `check-windows` is
1708    /// skipped on pull requests and runs on tags.
1709    #[cfg(unix)]
1710    static LEDGER_GUARD: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
1711
1712    #[cfg(unix)]
1713    #[tokio::test]
1714    async fn generic_post_load_error_tears_down_before_releasing_pending_charge() {
1715        let _ledger = LEDGER_GUARD.lock().await;
1716        let fixture = tempfile::tempdir().unwrap();
1717        let admission = LocalWorkerAdmission {
1718            state_root: fixture.path().join("state"),
1719            measured_weights_bytes: 512 * 1024 * 1024,
1720            ..test_admission()
1721        };
1722        let coordinator = car_inference::resource_policy::scoped_local_admission(
1723            &admission.state_root,
1724            admission.policy.clone(),
1725            car_inference::hardware::HardwareInfo::detect(),
1726        );
1727        let resident_baseline_mb = coordinator.owned_resident_model_mb_for_testing();
1728        let off = sh(
1729            r#"while IFS= read -r line; do sleep 0.2; printf '%s\n' '{"Error":"failed after cache publication"}'; done"#,
1730        );
1731        let request = GenerateRequest {
1732            model: Some("stub".into()),
1733            ..Default::default()
1734        };
1735
1736        // Measure the delta this worker adds rather than the absolute total.
1737        // `LEDGER_GUARD` is what makes this sound — see its comment; the delta
1738        // alone is not enough, because an interfering charge can land between
1739        // the baseline read and the assertion. Both together: exclusion removes
1740        // the interference, the delta keeps the assertion honest about what
1741        // this worker is responsible for.
1742        let baseline_mb = coordinator.resident_model_mb();
1743        let caller = off.clone();
1744        let exchange =
1745            tokio::spawn(async move { caller.generate_admitted(request, admission).await });
1746        tokio::time::sleep(Duration::from_millis(30)).await;
1747        assert!(coordinator.teardown_pending("stub"));
1748        assert_eq!(
1749            coordinator.resident_model_mb().saturating_sub(baseline_mb),
1750            512,
1751            "this worker's own 512 MB charge must still be held before release"
1752        );
1753        assert_eq!(
1754            coordinator.resident_allocation_ids("stub"),
1755            vec![off.allocation_id("stub")]
1756        );
1757        assert_eq!(
1758            coordinator
1759                .owned_resident_model_mb_for_testing()
1760                .checked_sub(resident_baseline_mb),
1761            Some(512)
1762        );
1763        assert!(exchange.await.unwrap().is_err());
1764        assert!(off.inner.lock().await.is_none());
1765
1766        tokio::time::timeout(Duration::from_secs(2), async {
1767            while coordinator.teardown_pending("stub") {
1768                tokio::time::sleep(Duration::from_millis(10)).await;
1769            }
1770        })
1771        .await
1772        .expect("worker must be reaped before its pending machine charge clears");
1773        assert!(coordinator.resident_allocation_ids("stub").is_empty());
1774        assert_eq!(
1775            coordinator.owned_resident_model_mb_for_testing(),
1776            resident_baseline_mb
1777        );
1778    }
1779
1780    #[cfg(unix)]
1781    #[tokio::test]
1782    async fn mismatched_worker_ack_tears_down_requested_and_reported_allocations() {
1783        let _ledger = LEDGER_GUARD.lock().await;
1784        let fixture = tempfile::tempdir().unwrap();
1785        let admission = LocalWorkerAdmission {
1786            state_root: fixture.path().join("state"),
1787            measured_weights_bytes: 64 * 1024 * 1024,
1788            ..test_admission()
1789        };
1790        let coordinator = car_inference::resource_policy::scoped_local_admission(
1791            &admission.state_root,
1792            admission.policy.clone(),
1793            car_inference::hardware::HardwareInfo::detect(),
1794        );
1795        let off = sh(&format!(
1796            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
1797        ));
1798        let request = GenerateRequest {
1799            model: Some("requested-a".into()),
1800            ..Default::default()
1801        };
1802
1803        let error = off
1804            .generate_admitted(request, admission)
1805            .await
1806            .err()
1807            .expect("mismatched ack must fail");
1808        assert!(error.to_string().contains("acknowledged model 'stub'"));
1809        assert!(off.inner.lock().await.is_none());
1810        tokio::time::timeout(Duration::from_secs(2), async {
1811            while coordinator.teardown_pending("requested-a")
1812                || coordinator.teardown_pending("stub")
1813            {
1814                tokio::time::sleep(Duration::from_millis(10)).await;
1815            }
1816        })
1817        .await
1818        .expect("mismatch teardown must reap both exact allocation candidates");
1819        assert!(coordinator
1820            .resident_allocation_ids("requested-a")
1821            .is_empty());
1822        assert!(coordinator.resident_allocation_ids("stub").is_empty());
1823    }
1824
1825    #[cfg(unix)]
1826    #[tokio::test]
1827    async fn mismatched_stream_ack_never_returns_or_remembers_worker() {
1828        let _ledger = LEDGER_GUARD.lock().await;
1829        let fixture = tempfile::tempdir().unwrap();
1830        let admission = LocalWorkerAdmission {
1831            state_root: fixture.path().join("state"),
1832            measured_weights_bytes: 64 * 1024 * 1024,
1833            ..test_admission()
1834        };
1835        let coordinator = car_inference::resource_policy::scoped_local_admission(
1836            &admission.state_root,
1837            admission.policy.clone(),
1838            car_inference::hardware::HardwareInfo::detect(),
1839        );
1840        let off = sh(&format!(
1841            "while IFS= read -r line; do printf '%s\\n' '{STREAM_STARTED_LINE}'; done"
1842        ));
1843        let request = GenerateRequest {
1844            model: Some("requested-stream-a".into()),
1845            ..Default::default()
1846        };
1847
1848        let error = off
1849            .stream_admitted(request, admission)
1850            .await
1851            .err()
1852            .expect("mismatched stream ack must fail");
1853        assert!(error.to_string().contains("acknowledged model 'stub'"));
1854        assert!(off.inner.lock().await.is_none());
1855        tokio::time::timeout(Duration::from_secs(2), async {
1856            while coordinator.teardown_pending("requested-stream-a")
1857                || coordinator.teardown_pending("stub")
1858            {
1859                tokio::time::sleep(Duration::from_millis(10)).await;
1860            }
1861        })
1862        .await
1863        .expect("stream mismatch teardown must reap both allocation candidates");
1864        assert!(off.resident_models().await.is_empty());
1865    }
1866
1867    #[cfg(unix)]
1868    #[tokio::test]
1869    async fn replacement_worker_generation_is_charged_separately_until_old_exit() {
1870        let _ledger = LEDGER_GUARD.lock().await;
1871        let fixture = tempfile::tempdir().unwrap();
1872        let admission = LocalWorkerAdmission {
1873            state_root: fixture.path().join("state"),
1874            measured_weights_bytes: 64 * 1024 * 1024,
1875            ..test_admission()
1876        };
1877        let coordinator = car_inference::resource_policy::scoped_local_admission(
1878            &admission.state_root,
1879            admission.policy.clone(),
1880            car_inference::hardware::HardwareInfo::detect(),
1881        );
1882        let script = format!("while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done");
1883        let first = sh(&script);
1884        let second = sh(&script);
1885        assert_ne!(first.allocation_id("stub"), second.allocation_id("stub"));
1886
1887        let mut first_reservation = coordinator
1888            .reserve_measured_host("stub", 64 * 1024 * 1024, 0)
1889            .unwrap();
1890        let first_result = first
1891            .generate_admitted(
1892                GenerateRequest {
1893                    model: Some("stub".into()),
1894                    ..Default::default()
1895                },
1896                admission.clone(),
1897            )
1898            .await
1899            .unwrap();
1900        first_reservation.publish_resident_weights_as(
1901            &first.allocation_id("stub"),
1902            first_result.residency.measured_weights_bytes,
1903        );
1904        drop(first_reservation);
1905
1906        let mut second_reservation = coordinator
1907            .reserve_measured_host("stub", 64 * 1024 * 1024, 0)
1908            .unwrap();
1909        let second_result = second
1910            .generate_admitted(
1911                GenerateRequest {
1912                    model: Some("stub".into()),
1913                    ..Default::default()
1914                },
1915                admission,
1916            )
1917            .await
1918            .unwrap();
1919        second_reservation.publish_resident_weights_as(
1920            &second.allocation_id("stub"),
1921            second_result.residency.measured_weights_bytes,
1922        );
1923        drop(second_reservation);
1924        assert_eq!(coordinator.resident_allocation_ids("stub").len(), 2);
1925
1926        assert!(first.release_model("stub").await.unwrap());
1927        assert_eq!(coordinator.resident_allocation_ids("stub").len(), 1);
1928        assert!(second.release_model("stub").await.unwrap());
1929    }
1930
1931    #[cfg(unix)]
1932    #[tokio::test]
1933    async fn policy_generation_change_restarts_existing_worker() {
1934        let off = sh(r#"while IFS= read -r line; do
1935printf '{"Result":{"result":{"text":"pong","tool_calls":[],"trace_id":"t","model_used":"%s","latency_ms":1,"usage":null},"residency":{"model_id":"stub","measured_weights_bytes":1,"retention":"resident"}}}\n' "$$"
1936        done"#);
1937        let first = off
1938            .generate_admitted(stub_request(), test_admission())
1939            .await
1940            .unwrap()
1941            .result
1942            .model_used;
1943        off.refresh_resource_policy(2);
1944        let mut updated = test_admission();
1945        updated.policy_generation = 2;
1946        let second = off
1947            .generate_admitted(stub_request(), updated)
1948            .await
1949            .unwrap()
1950            .result
1951            .model_used;
1952        assert_ne!(first, second, "policy refresh must replace the old child");
1953    }
1954
1955    #[cfg(unix)]
1956    #[tokio::test]
1957    async fn targeted_release_waits_for_worker_exit_and_clears_residency() {
1958        let _ledger = LEDGER_GUARD.lock().await;
1959        let fixture = tempfile::tempdir().unwrap();
1960        let admission = LocalWorkerAdmission {
1961            state_root: fixture.path().join("state"),
1962            ..test_admission()
1963        };
1964        let coordinator = car_inference::resource_policy::scoped_local_admission(
1965            &admission.state_root,
1966            admission.policy.clone(),
1967            car_inference::hardware::HardwareInfo::detect(),
1968        );
1969        let off = sh(&format!(
1970            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
1971        ));
1972        off.generate_admitted(stub_request(), admission)
1973            .await
1974            .unwrap();
1975        coordinator.mark_resident_allocation("stub", &off.allocation_id("stub"), 1);
1976        assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
1977        assert!(off.release_model("stub").await.unwrap());
1978        assert!(off.resident_models().await.is_empty());
1979        assert!(!coordinator.is_resident("stub"));
1980        assert!(off.inner.lock().await.is_none());
1981    }
1982
1983    #[cfg(unix)]
1984    #[tokio::test(start_paused = true)]
1985    async fn release_idle_models_releases_a_resident_older_than_ttl() {
1986        let _ledger = LEDGER_GUARD.lock().await;
1987        let fixture = tempfile::tempdir().unwrap();
1988        let admission = LocalWorkerAdmission {
1989            state_root: fixture.path().join("state"),
1990            ..test_admission()
1991        };
1992        let _coordinator = car_inference::resource_policy::scoped_local_admission(
1993            &admission.state_root,
1994            admission.policy.clone(),
1995            car_inference::hardware::HardwareInfo::detect(),
1996        );
1997        let off = sh(&format!(
1998            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
1999        ));
2000        off.generate_admitted(stub_request(), admission)
2001            .await
2002            .unwrap();
2003        assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
2004
2005        tokio::time::advance(Duration::from_secs(11)).await;
2006        let released = off.release_idle_models(Duration::from_secs(10)).await;
2007
2008        assert_eq!(released, vec!["stub".to_string()]);
2009        assert!(off.resident_models().await.is_empty());
2010        assert!(off.inner.lock().await.is_none());
2011    }
2012
2013    #[cfg(unix)]
2014    #[tokio::test(start_paused = true)]
2015    async fn release_idle_models_keeps_a_recently_used_resident() {
2016        let _ledger = LEDGER_GUARD.lock().await;
2017        let fixture = tempfile::tempdir().unwrap();
2018        let admission = LocalWorkerAdmission {
2019            state_root: fixture.path().join("state"),
2020            ..test_admission()
2021        };
2022        let _coordinator = car_inference::resource_policy::scoped_local_admission(
2023            &admission.state_root,
2024            admission.policy.clone(),
2025            car_inference::hardware::HardwareInfo::detect(),
2026        );
2027        let off = sh(&format!(
2028            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2029        ));
2030        off.generate_admitted(stub_request(), admission)
2031            .await
2032            .unwrap();
2033
2034        tokio::time::advance(Duration::from_secs(9)).await;
2035        assert!(off
2036            .release_idle_models(Duration::from_secs(10))
2037            .await
2038            .is_empty());
2039        assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
2040        assert!(off.release_model("stub").await.unwrap());
2041    }
2042
2043    #[cfg(unix)]
2044    #[tokio::test(start_paused = true)]
2045    async fn release_idle_models_keeps_a_never_completed_resident() {
2046        let _ledger = LEDGER_GUARD.lock().await;
2047        let fixture = tempfile::tempdir().unwrap();
2048        let admission = LocalWorkerAdmission {
2049            state_root: fixture.path().join("state"),
2050            ..test_admission()
2051        };
2052        let _coordinator = car_inference::resource_policy::scoped_local_admission(
2053            &admission.state_root,
2054            admission.policy.clone(),
2055            car_inference::hardware::HardwareInfo::detect(),
2056        );
2057        let off = sh(&format!(
2058            "while IFS= read -r line; do printf '%s\\n' '{TRANSIENT_RESULT_LINE}'; done"
2059        ));
2060        off.generate_admitted(stub_request(), admission.clone())
2061            .await
2062            .unwrap();
2063        off.remember_residency(
2064            &WorkerResidencyAck {
2065                model_id: "stub".into(),
2066                measured_weights_bytes: 1,
2067                retention: car_inference::backend_cache::BackendRetention::Resident,
2068            },
2069            admission.state_root,
2070            None,
2071        );
2072
2073        tokio::time::advance(Duration::from_secs(60)).await;
2074        assert!(off
2075            .release_idle_models(Duration::from_secs(10))
2076            .await
2077            .is_empty());
2078        assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
2079        assert!(off.release_model("stub").await.unwrap());
2080    }
2081
2082    #[cfg(unix)]
2083    #[tokio::test(start_paused = true)]
2084    async fn release_idle_models_honors_completion_that_lands_while_sweep_waits() {
2085        let _ledger = LEDGER_GUARD.lock().await;
2086        let fixture = tempfile::tempdir().unwrap();
2087        let admission = LocalWorkerAdmission {
2088            state_root: fixture.path().join("state"),
2089            ..test_admission()
2090        };
2091        let _coordinator = car_inference::resource_policy::scoped_local_admission(
2092            &admission.state_root,
2093            admission.policy.clone(),
2094            car_inference::hardware::HardwareInfo::detect(),
2095        );
2096        let off = sh(&format!(
2097            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2098        ));
2099        off.generate_admitted(stub_request(), admission.clone())
2100            .await
2101            .unwrap();
2102        tokio::time::advance(Duration::from_secs(11)).await;
2103
2104        // Model a non-streaming turn already holding the worker lock. Its old
2105        // completion timestamp remains visible until the fresh Result arrives.
2106        let inner = off.inner.clone();
2107        let worker_guard = inner.lock().await;
2108        let sweep_offload = off.clone();
2109        let mut sweep = Box::pin(async move {
2110            sweep_offload
2111                .release_idle_models(Duration::from_secs(10))
2112                .await
2113        });
2114        tokio::select! {
2115            released = &mut sweep => panic!("sweep released while worker lock was held: {released:?}"),
2116            _ = tokio::task::yield_now() => {}
2117        }
2118        WorkerOffload::mark_residency_used(&off.resident_models, &admission.state_root, "stub");
2119        drop(worker_guard);
2120
2121        assert!(sweep.await.is_empty());
2122        assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
2123        assert!(off.release_model("stub").await.unwrap());
2124    }
2125
2126    #[cfg(unix)]
2127    #[tokio::test(start_paused = true)]
2128    async fn release_idle_models_keeps_worker_with_recent_co_resident() {
2129        let _ledger = LEDGER_GUARD.lock().await;
2130        let fixture = tempfile::tempdir().unwrap();
2131        let admission = LocalWorkerAdmission {
2132            state_root: fixture.path().join("state"),
2133            ..test_admission()
2134        };
2135        let _coordinator = car_inference::resource_policy::scoped_local_admission(
2136            &admission.state_root,
2137            admission.policy.clone(),
2138            car_inference::hardware::HardwareInfo::detect(),
2139        );
2140        let off = sh(&format!(
2141            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2142        ));
2143        off.generate_admitted(stub_request(), admission.clone())
2144            .await
2145            .unwrap();
2146        tokio::time::advance(Duration::from_secs(11)).await;
2147        off.remember_residency(
2148            &WorkerResidencyAck {
2149                model_id: "recent".into(),
2150                measured_weights_bytes: 1,
2151                retention: car_inference::backend_cache::BackendRetention::Resident,
2152            },
2153            admission.state_root,
2154            Some(Instant::now()),
2155        );
2156
2157        assert!(off
2158            .release_idle_models(Duration::from_secs(10))
2159            .await
2160            .is_empty());
2161        let mut residents = off.resident_models().await;
2162        residents.sort();
2163        assert_eq!(residents, vec!["recent".to_string(), "stub".to_string()]);
2164        assert!(off.release_model("stub").await.unwrap());
2165    }
2166
2167    #[cfg(unix)]
2168    #[tokio::test(start_paused = true)]
2169    async fn release_idle_models_releases_worker_once_when_all_residents_are_idle() {
2170        let _ledger = LEDGER_GUARD.lock().await;
2171        let fixture = tempfile::tempdir().unwrap();
2172        let admission = LocalWorkerAdmission {
2173            state_root: fixture.path().join("state"),
2174            ..test_admission()
2175        };
2176        let _coordinator = car_inference::resource_policy::scoped_local_admission(
2177            &admission.state_root,
2178            admission.policy.clone(),
2179            car_inference::hardware::HardwareInfo::detect(),
2180        );
2181        let off = sh(&format!(
2182            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2183        ));
2184        off.generate_admitted(stub_request(), admission.clone())
2185            .await
2186            .unwrap();
2187        off.remember_residency(
2188            &WorkerResidencyAck {
2189                model_id: "also-idle".into(),
2190                measured_weights_bytes: 1,
2191                retention: car_inference::backend_cache::BackendRetention::Resident,
2192            },
2193            admission.state_root,
2194            Some(Instant::now()),
2195        );
2196        tokio::time::advance(Duration::from_secs(11)).await;
2197
2198        let mut released = off.release_idle_models(Duration::from_secs(10)).await;
2199        released.sort();
2200
2201        assert_eq!(released, vec!["also-idle".to_string(), "stub".to_string()]);
2202        assert!(off.resident_models().await.is_empty());
2203        assert!(off.inner.lock().await.is_none());
2204    }
2205
2206    #[cfg(unix)]
2207    #[tokio::test]
2208    async fn cancelled_worker_release_keeps_accounting_until_background_reap() {
2209        let _ledger = LEDGER_GUARD.lock().await;
2210        let fixture = tempfile::tempdir().unwrap();
2211        let admission = LocalWorkerAdmission {
2212            state_root: fixture.path().join("state"),
2213            ..test_admission()
2214        };
2215        let coordinator = car_inference::resource_policy::scoped_local_admission(
2216            &admission.state_root,
2217            admission.policy.clone(),
2218            car_inference::hardware::HardwareInfo::detect(),
2219        );
2220        let off = sh(&format!(
2221            "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2222        ));
2223        off.generate_admitted(stub_request(), admission)
2224            .await
2225            .unwrap();
2226        coordinator.mark_resident_allocation("stub", &off.allocation_id("stub"), 1);
2227        off.release_delay_ms.store(250, Ordering::Release);
2228
2229        let release_owner = off.clone();
2230        let release = tokio::spawn(async move { release_owner.release_model("stub").await });
2231        tokio::time::sleep(Duration::from_millis(20)).await;
2232        assert!(coordinator.teardown_pending("stub"));
2233        release.abort();
2234        let _ = release.await;
2235
2236        tokio::time::timeout(Duration::from_secs(2), async {
2237            while !off.resident_models().await.is_empty() || coordinator.is_resident("stub") {
2238                tokio::time::sleep(Duration::from_millis(10)).await;
2239            }
2240        })
2241        .await
2242        .expect("cancelled release guard must retain, kill, reap, and clear accounting");
2243        assert!(off.inner.lock().await.is_none());
2244    }
2245
2246    /// The coder's Agent-build deadline (car#1512) stops local generation by
2247    /// dropping the request future under `tokio::time::timeout`. On the worker
2248    /// offload that must stop the work, not just the wait: the child is killed
2249    /// and gone, and its admission charge clears. `WorkerProcessGuard` releases
2250    /// that charge only after `wait()` confirms the exit, so a cleared charge
2251    /// is also proof of the reap.
2252    #[cfg(unix)]
2253    #[tokio::test]
2254    async fn a_dropped_worker_generation_is_killed_reaped_and_unaccounted() {
2255        let _ledger = LEDGER_GUARD.lock().await;
2256        let fixture = tempfile::tempdir().unwrap();
2257        let admission = LocalWorkerAdmission {
2258            state_root: fixture.path().join("state"),
2259            measured_weights_bytes: 256 * 1024 * 1024,
2260            ..test_admission()
2261        };
2262        let coordinator = car_inference::resource_policy::scoped_local_admission(
2263            &admission.state_root,
2264            admission.policy.clone(),
2265            car_inference::hardware::HardwareInfo::detect(),
2266        );
2267        let resident_baseline_mb = coordinator.owned_resident_model_mb_for_testing();
2268        // The stub records its pid only after reading the request, then `exec`s
2269        // into a sleep under the SAME pid: the process under test is the worker
2270        // child itself, and the exchange is genuinely in flight.
2271        let pid_file = fixture.path().join("worker.pid");
2272        let off = sh(&format!(
2273            "IFS= read -r line; echo $$ > '{tmp}'; mv '{tmp}' '{pid}'; exec sleep 60",
2274            tmp = pid_file.with_extension("tmp").display(),
2275            pid = pid_file.display(),
2276        ));
2277        let alive = |pid: u32| {
2278            std::process::Command::new("sh")
2279                .arg("-c")
2280                .arg(format!("kill -0 {pid} 2>/dev/null"))
2281                .status()
2282                .is_ok_and(|status| status.success())
2283        };
2284
2285        let mut generation = Box::pin(off.generate_admitted(stub_request(), admission));
2286        let pid = tokio::time::timeout(Duration::from_secs(5), async {
2287            tokio::select! {
2288                result = &mut generation => {
2289                    panic!("the stub never answers, yet generation ended (ok: {})", result.is_ok())
2290                }
2291                pid = async {
2292                    loop {
2293                        if let Some(pid) = std::fs::read_to_string(&pid_file)
2294                            .ok()
2295                            .and_then(|text| text.trim().parse::<u32>().ok())
2296                        {
2297                            break pid;
2298                        }
2299                        tokio::time::sleep(Duration::from_millis(10)).await;
2300                    }
2301                } => pid,
2302            }
2303        })
2304        .await
2305        .expect("the stub worker must receive the request");
2306        assert!(
2307            alive(pid),
2308            "positive control: the in-flight worker is running"
2309        );
2310        assert!(
2311            coordinator.teardown_pending("stub"),
2312            "positive control: the in-flight request holds its admission charge"
2313        );
2314
2315        // The deadline fires on a request that will never answer and drops the
2316        // future it owns: the same drop the agent-build timeout performs.
2317        assert!(
2318            tokio::time::timeout(Duration::from_millis(50), generation)
2319                .await
2320                .is_err(),
2321            "a stub that never answers must hit the deadline"
2322        );
2323
2324        tokio::time::timeout(Duration::from_secs(5), async {
2325            while alive(pid) || coordinator.teardown_pending("stub") {
2326                tokio::time::sleep(Duration::from_millis(10)).await;
2327            }
2328        })
2329        .await
2330        .expect("a dropped generation must kill and reap its worker and clear its charge");
2331        assert!(coordinator.resident_allocation_ids("stub").is_empty());
2332        assert_eq!(
2333            coordinator.owned_resident_model_mb_for_testing(),
2334            resident_baseline_mb
2335        );
2336        assert!(off.resident_models().await.is_empty());
2337        assert!(
2338            off.inner.lock().await.is_none(),
2339            "a killed worker must not return to the slot"
2340        );
2341    }
2342
2343    #[cfg(unix)]
2344    #[tokio::test]
2345    async fn worker_crash_fails_one_call_gracefully_then_respawns() {
2346        // A stub that exits immediately: the exchange sees a broken pipe / EOF,
2347        // which must surface as an Err (NOT a panic / hang) — this is the whole
2348        // point of car-releases#74: one bad inference fails its own RPC and the
2349        // daemon stays up. The next call respawns (and crashes again here), which
2350        // must ALSO be a clean Err, proving the slot recovered rather than
2351        // desyncing.
2352        let off = sh("exit 1");
2353        let e1 = off
2354            .generate_admitted(GenerateRequest::default(), test_admission())
2355            .await
2356            .err()
2357            .expect("worker crash must fail");
2358        assert!(
2359            e1.to_string().contains("worker connection failed"),
2360            "unexpected error: {e1}"
2361        );
2362        let e2 = off
2363            .generate_admitted(GenerateRequest::default(), test_admission())
2364            .await
2365            .err()
2366            .expect("worker crash must fail");
2367        assert!(
2368            e2.to_string().contains("worker connection failed"),
2369            "unexpected error: {e2}"
2370        );
2371    }
2372
2373    #[cfg(unix)]
2374    #[tokio::test]
2375    async fn worker_failure_preserves_exit_status_and_does_not_claim_a_restart() {
2376        for (script, expected) in [
2377            ("IFS= read -r line; exit 17", "exit status: 17"),
2378            ("IFS= read -r line; kill -TERM $$", "signal: 15"),
2379        ] {
2380            let off = sh(script);
2381            let error = off
2382                .generate_admitted(GenerateRequest::default(), test_admission())
2383                .await
2384                .err()
2385                .expect("worker must fail")
2386                .to_string();
2387            assert!(error.contains(expected), "{error}");
2388            assert!(
2389                error.contains("next request starts a new worker"),
2390                "{error}"
2391            );
2392            assert!(!error.contains("was restarted"), "{error}");
2393            assert!(off.inner.lock().await.is_none());
2394        }
2395    }
2396
2397    #[cfg(unix)]
2398    #[tokio::test]
2399    async fn stream_forwards_events_then_closes() {
2400        // A stub that emits two text deltas and a clean StreamEnd for one request.
2401        let off = sh(&format!(
2402            "IFS= read -r line; \
2403             printf '%s\\n' '{STREAM_STARTED_LINE}'; \
2404             printf '%s\\n' '{{\"Event\":{{\"TextDelta\":\"he\"}}}}'; \
2405             printf '%s\\n' '{{\"Event\":{{\"TextDelta\":\"llo\"}}}}'; \
2406             printf '%s\\n' '\"StreamEnd\"'; \
2407             while IFS= read -r l; do :; done"
2408        ));
2409        let mut rx = off
2410            .stream_admitted(stub_request(), test_admission())
2411            .await
2412            .unwrap()
2413            .events;
2414        let mut got = String::new();
2415        while let Some(ev) = rx.recv().await {
2416            if let StreamEvent::TextDelta(t) = ev {
2417                got.push_str(&t);
2418            }
2419        }
2420        assert_eq!(got, "hello");
2421    }
2422
2423    #[cfg(unix)]
2424    #[tokio::test]
2425    async fn stream_worker_death_closes_with_error_end() {
2426        // A stub that emits one event then exits mid-stream (no StreamEnd): the
2427        // reader must surface a StopReason("error") and close, not hang.
2428        let off = sh(&format!("IFS= read -r line; printf '%s\\n' '{STREAM_STARTED_LINE}'; printf '%s\\n' '{{\"Event\":{{\"TextDelta\":\"partial\"}}}}'"));
2429        let mut rx = off
2430            .stream_admitted(stub_request(), test_admission())
2431            .await
2432            .unwrap()
2433            .events;
2434        let mut saw_text = false;
2435        let mut saw_error_end = false;
2436        while let Some(ev) = rx.recv().await {
2437            match ev {
2438                StreamEvent::TextDelta(t) if t == "partial" => saw_text = true,
2439                StreamEvent::StopReason(s) if s == "error" => saw_error_end = true,
2440                _ => {}
2441            }
2442        }
2443        assert!(saw_text, "should have forwarded the partial delta");
2444        assert!(
2445            saw_error_end,
2446            "mid-stream death should surface an error end"
2447        );
2448    }
2449}