Skip to main content

cpex_hosts_python/
worker.rs

1// Location: ./crates/cpex-hosts-python/src/worker.rs
2// Copyright 2026
3// SPDX-License-Identifier: Apache-2.0
4// Authors: Ted Habeck
5//
6// Worker client — launches `worker.py` in the plugin's venv and speaks the
7// newline-delimited JSON protocol to it.
8//
9// The async analog of `cpex/framework/isolated/venv_comm.py`. Where that uses
10// two threads and a `Queue` per request, this uses two tokio tasks and a
11// one-shot channel per request:
12//
13// ```text
14//   send_task ──> outbound channel ──> writer task ──> child stdin
15//                      registers a one-shot in `pending`
16//   child stdout ──> reader task ──> demux on request_id ──> that one-shot
17// ```
18//
19// Demuxing on `request_id` rather than assuming request/response ordering is
20// what lets several hook invocations be in flight against one worker at once —
21// the executor dispatches parallel-phase plugins concurrently, and a
22// FIFO-coupled client would serialize them.
23//
24// # Failure modes, kept distinct
25//
26// A hung worker (`Timeout`) and a dead worker (`WorkerDied`) call for different
27// operator responses, so they do not collapse into one error. Process death is
28// detected by reader EOF, at which point every pending one-shot is dropped —
29// each in-flight `send_task` then resolves to `WorkerDied` instead of waiting
30// out a timeout that can never be satisfied.
31
32use std::collections::{BTreeSet, HashMap};
33use std::path::{Path, PathBuf};
34use std::process::Stdio;
35use std::sync::Arc;
36
37use serde_json::Value;
38use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt, BufReader};
39use tokio::process::{Child, Command};
40use tokio::sync::{Mutex, mpsc, oneshot};
41
42use crate::error::HostError;
43
44/// How long to wait for a worker to exit after a `shutdown` task before
45/// killing it. Matches `venv_comm.py`'s `stop_worker` grace window.
46const SHUTDOWN_GRACE_SECS: u64 = 5;
47
48/// Task type for a hook invocation, as `worker.py` dispatches on it.
49pub const TASK_LOAD_AND_RUN_HOOK: &str = "load_and_run_hook";
50
51/// Task type that asks the worker to exit.
52pub const TASK_SHUTDOWN: &str = "shutdown";
53
54/// Task type that asks the worker which wire-protocol features it implements.
55///
56/// Answered before any plugin code loads. A worker predating the handshake
57/// replies `{"status":"error","message":"task type not supported."}`, which the
58/// host reads as "no features" rather than as a transport fault — see
59/// [`WorkerClient::capabilities`].
60pub const TASK_CAPABILITIES: &str = "capabilities";
61
62/// Worker feature name for reading the task's `credential` field.
63pub const FEATURE_CREDENTIAL: &str = "credential";
64
65/// Worker feature name for delivering the task's `extensions` field to a hook.
66pub const FEATURE_EXTENSIONS: &str = "extensions";
67
68/// Worker feature name for serving PPE's `cmf.*` hooks natively, with the CMF
69/// `MessagePayload` and the `cmf.*` hook names.
70///
71/// A worker that does not claim it gets the legacy payload shim instead (see
72/// `crate::legacy::cmf_shim`). No released worker claims it yet.
73pub const FEATURE_CMF_HOOKS: &str = "cmf_hooks";
74
75/// What a worker reported it can do, from the spawn-time handshake.
76///
77/// An empty feature set is a valid, meaningful answer: it is what an old worker
78/// that does not know the handshake reports, and it must make the host fail
79/// closed rather than assume support.
80#[derive(Debug, Clone, Default, PartialEq, Eq)]
81pub struct WorkerCapabilities {
82    /// The worker's declared protocol version, or 0 if it reported none.
83    pub protocol_version: u64,
84
85    /// Feature names the worker claims to implement.
86    pub features: BTreeSet<String>,
87
88    /// The installed `cpex` distribution version, for diagnostics only.
89    ///
90    /// Never gated on: two builds shipped as 0.1.2 with different worker
91    /// protocols, so the version string cannot carry a decision. `features` is
92    /// the contract.
93    pub cpex_version: Option<String>,
94}
95
96impl WorkerCapabilities {
97    /// The conservative answer for a worker that does not know the handshake.
98    pub fn none() -> Self {
99        Self::default()
100    }
101
102    /// Whether the worker claims the named feature.
103    pub fn has(&self, feature: &str) -> bool {
104        self.features.contains(feature)
105    }
106}
107
108/// Map of in-flight request ids to the channel awaiting each response.
109///
110/// Carries a `Result` rather than a bare `Value` so the reader task can hand a
111/// waiting caller a typed failure — an oversized response frame is a real
112/// answer to that request ("it cannot be delivered"), and reporting it beats
113/// letting the caller sit out the full timeout for a response that was already
114/// discarded.
115type Pending = Arc<Mutex<HashMap<String, oneshot::Sender<Result<Value, HostError>>>>>;
116
117/// How many recent stderr lines to retain for diagnosis.
118///
119/// A worker that dies at import time (a bad dependency pin, a missing module)
120/// says so on stderr and nowhere else. Without this, the host reports only
121/// "worker process died" and the actual cause is lost — the tail is small
122/// enough to be cheap and long enough to hold a Python traceback.
123const STDERR_RETAIN_LINES: usize = 40;
124
125/// Ring buffer of the worker's most recent stderr lines.
126type StderrTail = Arc<Mutex<std::collections::VecDeque<String>>>;
127
128/// Cap on one retained stderr line, independent of `max_content_size`.
129///
130/// Stderr is diagnostic, so an over-long line is truncated rather than treated
131/// as a protocol violation — but the retained copy still has to be bounded, or
132/// a worker that writes one enormous line drives host memory on its own.
133/// Sized to hold a Python traceback frame comfortably.
134const STDERR_MAX_LINE_BYTES: usize = 8 * 1024;
135
136/// Appended to a stderr line that was cut at `STDERR_MAX_LINE_BYTES`, so an
137/// operator reading the tail can tell truncation from a worker that really
138/// stopped mid-sentence.
139const TRUNCATION_MARKER: &str = "… [truncated by host]";
140
141/// Outcome of one bounded line read.
142enum BoundedLine {
143    /// A complete line, within the limit.
144    Line(String),
145
146    /// The line exceeded the limit. Carries the bytes read up to the cap, for
147    /// the diagnostic paths that want a prefix; the rest was discarded as it
148    /// arrived, so nothing beyond the cap was ever buffered.
149    Oversized { prefix: String, scanned: usize },
150
151    /// Stream closed.
152    Eof,
153}
154
155/// Read one `\n`-delimited line, refusing to buffer more than `limit` bytes.
156///
157/// `AsyncBufReadExt::lines` grows its `String` without bound, which is what
158/// makes an unbounded inbound stream a host memory problem regardless of the
159/// outbound `max_content_size` check. This reads through the `BufRead` fill
160/// buffer instead: once `limit` bytes have accumulated without a newline, the
161/// remainder of the line is consumed and dropped rather than retained, so the
162/// stream stays framed (the next read starts at a real line boundary) while
163/// peak memory stays at `limit` plus one fill buffer.
164async fn read_bounded_line<R>(reader: &mut R, limit: usize) -> std::io::Result<BoundedLine>
165where
166    R: AsyncBufRead + Unpin,
167{
168    let mut buffer = Vec::new();
169    let mut scanned = 0usize;
170    let mut over = false;
171
172    loop {
173        let available = reader.fill_buf().await?;
174        if available.is_empty() {
175            // EOF. A trailing fragment with no newline is still a line.
176            if scanned == 0 {
177                return Ok(BoundedLine::Eof);
178            }
179            break;
180        }
181
182        let (chunk, consumed, done) = match available.iter().position(|b| *b == b'\n') {
183            Some(index) => (&available[..index], index + 1, true),
184            None => {
185                let all = available.len();
186                (available, all, false)
187            },
188        };
189
190        scanned += chunk.len();
191        if !over {
192            // Only retain up to the cap; everything past it is consumed and
193            // dropped, which is what keeps this bounded.
194            let room = limit.saturating_sub(buffer.len());
195            let keep = chunk.len().min(room);
196            buffer.extend_from_slice(&chunk[..keep]);
197            if scanned > limit {
198                over = true;
199            }
200        }
201
202        reader.consume(consumed);
203        if done {
204            break;
205        }
206    }
207
208    // Trim a `\r` from CRLF, matching what `lines()` does.
209    if buffer.last() == Some(&b'\r') {
210        buffer.pop();
211    }
212
213    let text = String::from_utf8_lossy(&buffer).into_owned();
214    if over {
215        Ok(BoundedLine::Oversized {
216            prefix: text,
217            scanned,
218        })
219    } else {
220        Ok(BoundedLine::Line(text))
221    }
222}
223
224/// Set once the reader task observes that the worker's stdout has closed.
225///
226/// Needed because clearing `pending` only rescues requests that were *already*
227/// registered when the worker died. A send issued afterwards would otherwise
228/// register a fresh one-shot into a map nobody is reading, succeed at writing
229/// into an mpsc channel whose consumer has exited, and then wait out the full
230/// timeout for a response that can never arrive. Checking this flag turns that
231/// into an immediate `WorkerDied`.
232type Alive = Arc<std::sync::atomic::AtomicBool>;
233
234/// A live `worker.py` subprocess and the plumbing that talks to it.
235pub struct WorkerClient {
236    /// Outbound task lines, drained by the writer task.
237    outbound: mpsc::UnboundedSender<String>,
238
239    /// Response channels keyed by request id.
240    pending: Pending,
241
242    /// False once the reader task has seen the worker's stdout close.
243    alive: Alive,
244
245    /// Recent stderr, for explaining a death.
246    stderr_tail: StderrTail,
247
248    /// The child handle, kept so shutdown can wait on and kill it. `None`
249    /// once shutdown has consumed it.
250    child: Mutex<Option<Child>>,
251
252    /// Cap on one serialized task, enforced before the write.
253    max_content_size: usize,
254
255    /// Per-invocation response timeout.
256    timeout_secs: u64,
257
258    /// What the spawn-time handshake reported, once it has run.
259    ///
260    /// Held on the client rather than the plugin so a respawned worker starts
261    /// with no report and has to be asked again: the replacement may come from
262    /// a rebuilt venv with a different `worker.py`.
263    reported: std::sync::OnceLock<WorkerCapabilities>,
264
265    /// Whether the one probe retry this worker gets has been spent. See
266    /// [`Self::claim_probe_retry`].
267    probe_retry_spent: std::sync::atomic::AtomicBool,
268}
269
270/// Manual rather than derived: the outbound channel carries serialized task
271/// lines, which on a credential-bearing hook include the plaintext token. A
272/// derive would put those in any debug dump of the client.
273impl std::fmt::Debug for WorkerClient {
274    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
275        f.debug_struct("WorkerClient")
276            .field("max_content_size", &self.max_content_size)
277            .field("timeout_secs", &self.timeout_secs)
278            .finish_non_exhaustive()
279    }
280}
281
282impl WorkerClient {
283    /// Launch `worker.py` with the venv's interpreter and start the I/O tasks.
284    ///
285    /// `cwd` becomes the child's working directory, which is where a plugin's
286    /// relative-path side effects land.
287    pub async fn spawn(
288        python: &Path,
289        script_path: &Path,
290        cwd: Option<&Path>,
291        max_content_size: usize,
292        timeout_secs: u64,
293    ) -> Result<Self, HostError> {
294        if !python.exists() {
295            return Err(HostError::WorkerStart {
296                message: format!("venv interpreter not found at {}", python.display()),
297            });
298        }
299
300        let mut command = Command::new(python);
301        command
302            .arg(script_path)
303            .stdin(Stdio::piped())
304            .stdout(Stdio::piped())
305            .stderr(Stdio::piped())
306            // Without this, a killed host leaves orphaned workers behind.
307            .kill_on_drop(true);
308
309        if let Some(dir) = cwd {
310            command.current_dir(dir);
311        }
312
313        let mut child = command.spawn().map_err(|e| HostError::WorkerStart {
314            message: format!(
315                "could not spawn {} {}: {e}",
316                python.display(),
317                script_path.display()
318            ),
319        })?;
320
321        let stdin = child.stdin.take().ok_or_else(|| HostError::WorkerStart {
322            message: "worker stdin was not piped".into(),
323        })?;
324        let stdout = child.stdout.take().ok_or_else(|| HostError::WorkerStart {
325            message: "worker stdout was not piped".into(),
326        })?;
327        let stderr = child.stderr.take().ok_or_else(|| HostError::WorkerStart {
328            message: "worker stderr was not piped".into(),
329        })?;
330
331        let pending: Pending = Arc::new(Mutex::new(HashMap::new()));
332        let alive: Alive = Arc::new(std::sync::atomic::AtomicBool::new(true));
333        let stderr_tail: StderrTail = Arc::new(Mutex::new(std::collections::VecDeque::new()));
334        let (outbound, outbound_rx) = mpsc::unbounded_channel::<String>();
335
336        tokio::spawn(writer_loop(stdin, outbound_rx));
337        // The inbound cap is the same `max_content_size` the outbound path
338        // enforces: one configured limit on the size of a protocol frame,
339        // applied in both directions.
340        tokio::spawn(reader_loop(
341            stdout,
342            Arc::clone(&pending),
343            Arc::clone(&alive),
344            max_content_size,
345        ));
346        tokio::spawn(stderr_loop(stderr, Arc::clone(&stderr_tail)));
347
348        Ok(Self {
349            outbound,
350            pending,
351            alive,
352            stderr_tail,
353            child: Mutex::new(Some(child)),
354            max_content_size,
355            timeout_secs,
356            reported: std::sync::OnceLock::new(),
357            probe_retry_spent: std::sync::atomic::AtomicBool::new(false),
358        })
359    }
360
361    /// Claim this worker's single retry of a failed capabilities probe.
362    ///
363    /// True exactly once per worker. A probe that timed out or hit a worker
364    /// still starting leaves no report, and with no report a `cmf.*` hook is
365    /// served through the legacy shim for the worker's whole life. One retry
366    /// on first use recovers from a transient failure; bounding it to one
367    /// keeps a probe that keeps failing from adding a round trip to every
368    /// call.
369    pub fn claim_probe_retry(&self) -> bool {
370        !self
371            .probe_retry_spent
372            .swap(true, std::sync::atomic::Ordering::SeqCst)
373    }
374
375    /// The handshake report from the last successful [`Self::capabilities`]
376    /// call, or `None` if the worker was never asked.
377    pub fn reported_capabilities(&self) -> Option<&WorkerCapabilities> {
378        self.reported.get()
379    }
380
381    /// Build a `WorkerDied` error, appending recent stderr when there is any.
382    ///
383    /// The stderr tail is what turns "worker process died" into something
384    /// actionable — an import error, a missing dependency, a traceback.
385    async fn died(&self, context: &str) -> HostError {
386        let tail = self.stderr_tail.lock().await;
387        if tail.is_empty() {
388            return HostError::WorkerDied {
389                message: context.to_string(),
390            };
391        }
392
393        let recent: Vec<&str> = tail.iter().map(String::as_str).collect();
394        HostError::WorkerDied {
395            message: format!("{context}; recent worker stderr: {}", recent.join(" | ")),
396        }
397    }
398
399    /// The worker's retained stderr lines, oldest first.
400    ///
401    /// Exposed so a test can assert the credential channel never surfaces here:
402    /// this buffer is the one place the host copies worker stderr, and the one
403    /// place it is echoed onward (into a `WorkerDied` message).
404    pub async fn stderr_lines(&self) -> Vec<String> {
405        self.stderr_tail.lock().await.iter().cloned().collect()
406    }
407
408    /// Whether the worker's stdout is still open.
409    ///
410    /// A `false` here is definitive — the process has closed its output. A
411    /// `true` is only as fresh as the reader task's last observation, so a
412    /// send can still race a death; that path is covered by the dropped-sender
413    /// branch in `send_task`.
414    pub fn is_alive(&self) -> bool {
415        self.alive.load(std::sync::atomic::Ordering::Acquire)
416    }
417
418    /// Send a task and await its response.
419    ///
420    /// The caller supplies the task body; this method injects the `request_id`
421    /// and handles registration, the size check, the write, and the timeout.
422    pub async fn send_task(&self, mut task: Value) -> Result<Value, HostError> {
423        // Checked up front: once the worker is known dead, waiting out the
424        // timeout tells the caller nothing it does not already know.
425        if !self.is_alive() {
426            return Err(self.died("worker process is no longer running").await);
427        }
428
429        let request_id = uuid::Uuid::new_v4().to_string();
430
431        {
432            let obj = task.as_object_mut().ok_or_else(|| HostError::Protocol {
433                message: "task must be a JSON object".into(),
434            })?;
435            obj.insert("request_id".into(), Value::String(request_id.clone()));
436        }
437
438        let line = serde_json::to_string(&task).map_err(|e| HostError::Protocol {
439            message: format!("could not serialize task: {e}"),
440        })?;
441
442        // Checked before registering or writing, so an oversized task costs
443        // nothing and cannot desync the worker's line-oriented reader.
444        if line.len() > self.max_content_size {
445            return Err(HostError::TaskTooLarge {
446                size: line.len(),
447                limit: self.max_content_size,
448            });
449        }
450
451        let (tx, rx) = oneshot::channel();
452        self.pending.lock().await.insert(request_id.clone(), tx);
453
454        // Registration happens before the write so a fast response cannot
455        // arrive before there is anywhere to route it.
456        if self.outbound.send(format!("{line}\n")).is_err() {
457            self.pending.lock().await.remove(&request_id);
458            return Err(self
459                .died("writer task is gone — the worker is no longer accepting tasks")
460                .await);
461        }
462
463        let response =
464            match tokio::time::timeout(std::time::Duration::from_secs(self.timeout_secs), rx).await
465            {
466                Ok(Ok(Ok(response))) => response,
467                // The reader task delivered a typed failure for this request —
468                // an over-limit response frame it had to discard.
469                Ok(Ok(Err(e))) => return Err(e),
470                // The sender was dropped without a value: the reader task saw EOF
471                // and cleared `pending`, so the worker is gone.
472                Ok(Err(_)) => {
473                    return Err(self
474                        .died("worker exited while the request was in flight")
475                        .await);
476                },
477                Err(_) => {
478                    self.pending.lock().await.remove(&request_id);
479                    return Err(HostError::Timeout {
480                        timeout_secs: self.timeout_secs,
481                    });
482                },
483            };
484
485        interpret_response(response)
486    }
487
488    /// Ask the worker which wire-protocol features it implements.
489    ///
490    /// # Why an unsupported task type is not an error here
491    ///
492    /// A worker predating the handshake answers the unknown `capabilities` task
493    /// with `{"status":"error","message":"task type not supported."}`, which
494    /// `send_task` surfaces as [`HostError::WorkerError`]. That is the *expected*
495    /// reply from an old worker, not a transport fault, so it maps to an empty
496    /// feature set — the conservative answer that makes the caller fail closed
497    /// for any plugin needing a gated field.
498    ///
499    /// Everything else (a timeout, a dead worker, a malformed frame) stays an
500    /// error: those mean the probe itself did not complete, and treating them as
501    /// "no features" would report a specific, wrong diagnosis for a worker that
502    /// may well be current.
503    pub async fn capabilities(&self) -> Result<WorkerCapabilities, HostError> {
504        let response = match self
505            .send_task(serde_json::json!({ "task_type": TASK_CAPABILITIES }))
506            .await
507        {
508            Ok(response) => response,
509            Err(HostError::WorkerError { message }) => {
510                tracing::info!(
511                    reply = %message,
512                    "worker does not implement the capabilities handshake; \
513                     treating it as supporting no gated wire features"
514                );
515                let none = WorkerCapabilities::none();
516                let _ = self.reported.set(none.clone());
517                return Ok(none);
518            },
519            Err(e) => return Err(e),
520        };
521
522        let features = response
523            .get("features")
524            .and_then(Value::as_array)
525            .map(|list| {
526                list.iter()
527                    .filter_map(Value::as_str)
528                    .map(str::to_string)
529                    .collect::<BTreeSet<String>>()
530            })
531            .unwrap_or_default();
532
533        let reported = WorkerCapabilities {
534            protocol_version: response
535                .get("protocol_version")
536                .and_then(Value::as_u64)
537                .unwrap_or(0),
538            features,
539            cpex_version: response
540                .get("cpex_version")
541                .and_then(Value::as_str)
542                .map(str::to_string),
543        };
544        let _ = self.reported.set(reported.clone());
545        Ok(reported)
546    }
547
548    /// Ask the worker to exit, then wait out the grace window and kill it.
549    ///
550    /// Idempotent: a second call after the child has been reaped is a no-op,
551    /// which matters because `shutdown` may run from both the plugin lifecycle
552    /// and a drop path.
553    pub async fn shutdown(&self) -> Result<(), HostError> {
554        let Some(mut child) = self.child.lock().await.take() else {
555            return Ok(());
556        };
557
558        // `venv_comm.py` sends a fixed "shutdown" request id rather than a
559        // UUID, and the worker echoes it back; nothing awaits the reply.
560        let task = serde_json::json!({ "task_type": TASK_SHUTDOWN, "request_id": TASK_SHUTDOWN });
561        if let Ok(line) = serde_json::to_string(&task) {
562            let _ = self.outbound.send(format!("{line}\n"));
563        }
564
565        match tokio::time::timeout(
566            std::time::Duration::from_secs(SHUTDOWN_GRACE_SECS),
567            child.wait(),
568        )
569        .await
570        {
571            Ok(Ok(_status)) => Ok(()),
572            Ok(Err(e)) => Err(HostError::WorkerDied {
573                message: format!("could not wait on the worker process: {e}"),
574            }),
575            Err(_) => {
576                // A worker that ignores the shutdown task (or is wedged inside
577                // plugin code) must not hold teardown open.
578                tracing::warn!("worker did not exit within {SHUTDOWN_GRACE_SECS}s; killing it");
579                child.kill().await.map_err(|e| HostError::WorkerDied {
580                    message: format!("could not kill the unresponsive worker: {e}"),
581                })?;
582                Ok(())
583            },
584        }
585    }
586}
587
588/// Turn a worker response envelope into a result.
589///
590/// The worker signals failure with `{"status": "error", "message": ...}`; any
591/// other envelope is the hook result. `request_id` is stripped — it is
592/// transport bookkeeping, and `venv_comm.py` removes it too before returning.
593fn interpret_response(mut response: Value) -> Result<Value, HostError> {
594    if let Some(obj) = response.as_object_mut() {
595        obj.remove("request_id");
596
597        if obj.get("status").and_then(Value::as_str) == Some("error") {
598            let message = obj
599                .get("message")
600                .and_then(Value::as_str)
601                .unwrap_or("worker reported an error with no message")
602                .to_string();
603            return Err(HostError::WorkerError { message });
604        }
605    }
606    Ok(response)
607}
608
609/// Drain outbound task lines into the child's stdin.
610///
611/// Ends when the channel closes (client dropped) or the pipe breaks (worker
612/// gone). Either way the reader loop is what surfaces the failure to callers.
613async fn writer_loop(
614    mut stdin: tokio::process::ChildStdin,
615    mut outbound: mpsc::UnboundedReceiver<String>,
616) {
617    while let Some(line) = outbound.recv().await {
618        if stdin.write_all(line.as_bytes()).await.is_err() {
619            tracing::debug!("worker stdin closed; writer stopping");
620            break;
621        }
622        if stdin.flush().await.is_err() {
623            tracing::debug!("worker stdin flush failed; writer stopping");
624            break;
625        }
626    }
627}
628
629/// Read response lines and route each to the one-shot for its request id.
630///
631/// Each frame is read under `limit` (the configured `max_content_size`), so a
632/// hostile or runaway worker cannot grow host memory by never emitting a
633/// newline. An over-limit frame is dropped rather than buffered: the waiting
634/// caller gets a typed `ResponseTooLarge`, and because the rest of the line is
635/// consumed to its newline the stream stays framed for the next response.
636///
637/// On EOF, `pending` is drained and every sender dropped, which resolves each
638/// in-flight `send_task` to `WorkerDied` rather than leaving it to time out.
639async fn reader_loop(
640    stdout: tokio::process::ChildStdout,
641    pending: Pending,
642    alive: Alive,
643    limit: usize,
644) {
645    let mut reader = BufReader::new(stdout);
646
647    loop {
648        match read_bounded_line(&mut reader, limit).await {
649            Ok(BoundedLine::Line(line)) => {
650                let line = line.trim();
651                if line.is_empty() {
652                    continue;
653                }
654
655                let Ok(response) = serde_json::from_str::<Value>(line) else {
656                    // A plugin that printed to stdout lands here. Skipping the
657                    // line keeps the stream in sync; the request still times
658                    // out or gets its real response later.
659                    tracing::warn!("worker emitted a non-JSON stdout line; ignoring it");
660                    continue;
661                };
662
663                let Some(request_id) = response.get("request_id").and_then(Value::as_str) else {
664                    tracing::warn!("worker response had no request_id; dropping it");
665                    continue;
666                };
667
668                match pending.lock().await.remove(request_id) {
669                    Some(tx) => {
670                        // Receiver gone means the caller already timed out.
671                        let _ = tx.send(Ok(response));
672                    },
673                    None => {
674                        // The fixed "shutdown" id lands here by design, as does
675                        // any late response whose caller has given up.
676                        tracing::debug!(request_id, "response for an unknown request id");
677                    },
678                }
679            },
680            Ok(BoundedLine::Oversized { prefix, scanned }) => {
681                // The frame is gone — only `limit` bytes of it were ever kept.
682                // Recovering the request id from that prefix is best-effort
683                // (JSON field order is not guaranteed), so the fallback is to
684                // fail every in-flight request: something the worker sent could
685                // not be delivered, and guessing which caller it belonged to
686                // would be worse than telling all of them.
687                tracing::error!(
688                    scanned,
689                    limit,
690                    "worker response frame exceeded max_content_size; \
691                     the frame was discarded rather than buffered"
692                );
693
694                let request_id = request_id_from_prefix(&prefix);
695                let mut pending = pending.lock().await;
696                match request_id.and_then(|id| pending.remove(&id).map(|tx| (id, tx))) {
697                    Some((id, tx)) => {
698                        tracing::warn!(
699                            request_id = %id,
700                            "failing the request whose response was over the size limit"
701                        );
702                        let _ = tx.send(Err(response_too_large(scanned, limit)));
703                    },
704                    None => {
705                        // No recoverable id. Fail everyone rather than leave a
706                        // caller waiting on a response that no longer exists.
707                        for (id, tx) in pending.drain() {
708                            tracing::warn!(
709                                request_id = %id,
710                                "failing an in-flight request: an over-limit response frame was \
711                                 discarded and its owner could not be identified"
712                            );
713                            let _ = tx.send(Err(response_too_large(scanned, limit)));
714                        }
715                    },
716                }
717            },
718            Ok(BoundedLine::Eof) => {
719                tracing::debug!("worker stdout closed");
720                break;
721            },
722            Err(e) => {
723                tracing::warn!("error reading worker stdout: {e}");
724                break;
725            },
726        }
727    }
728
729    // Order matters: mark dead *before* clearing, so a send that races this
730    // teardown either sees the flag and fails fast, or finds its sender
731    // dropped. Either way it does not wait out the timeout.
732    alive.store(false, std::sync::atomic::Ordering::Release);
733
734    // Dropping the senders is what converts a dead worker into an immediate
735    // error for everyone still waiting.
736    pending.lock().await.clear();
737}
738
739/// The error a caller gets when its response frame was over the inbound cap.
740///
741/// A dedicated constructor because the failure needs one stable shape and one
742/// stable wording wherever the reader raises it. `HostError::TaskTooLarge` is
743/// deliberately not reused: its `Display` says "serialized task", which would
744/// point an operator at the *outbound* check and at their own payload rather
745/// than at the worker's reply.
746fn response_too_large(size: usize, limit: usize) -> HostError {
747    HostError::ResponseTooLarge { size, limit }
748}
749
750/// Best-effort recovery of a `request_id` from a truncated response frame.
751///
752/// The frame is not parseable JSON — it was cut mid-value — so this scans for
753/// the `"request_id"` key textually. `worker.py` emits the id early enough that
754/// this usually succeeds, which lets the reader fail exactly the one caller
755/// whose response was dropped instead of every in-flight request. A miss is
756/// expected and handled by the caller, not an error.
757fn request_id_from_prefix(prefix: &str) -> Option<String> {
758    const KEY: &str = "\"request_id\"";
759    let after_key = &prefix[prefix.find(KEY)? + KEY.len()..];
760    let after_colon = &after_key[after_key.find(':')? + 1..];
761    let opening = after_colon.find('"')?;
762    let rest = &after_colon[opening + 1..];
763    // A truncated frame can end mid-id; without a closing quote there is no
764    // way to know the value is complete, so treat it as unrecoverable.
765    let closing = rest.find('"')?;
766    Some(rest[..closing].to_string())
767}
768
769/// Log the worker's stderr and retain a bounded tail for diagnosis.
770///
771/// Logged at debug: the worker's own logging goes here, and a credential-bearing
772/// hook must never surface token material through it. `worker.py` scrubs its
773/// side (it holds the plaintext only to scrub it back out of anything the
774/// plugin logs, raises, or returns); this side never parses these lines or
775/// echoes them anywhere except into a `WorkerDied` message, which is the one
776/// place an operator genuinely needs them.
777///
778/// # Bounding
779///
780/// Stderr gets a *softer* policy than stdout. It carries no framing and no
781/// request routing — it is the channel that explains a death, so discarding it
782/// or killing the worker over it would destroy the diagnostic the operator
783/// needs precisely when things are going wrong. Each line is therefore
784/// truncated at `STDERR_MAX_LINE_BYTES` with an explicit marker (so an operator
785/// can tell a host cut from a worker that stopped mid-sentence) rather than
786/// raising an error, and the `STDERR_RETAIN_LINES` ring already caps the total.
787/// Together those give a hard ceiling of roughly
788/// `STDERR_RETAIN_LINES * STDERR_MAX_LINE_BYTES` on retained stderr, with peak
789/// read memory bounded per line — so a worker that spews stderr forever, in
790/// however few newlines, cannot grow the host.
791async fn stderr_loop(stderr: tokio::process::ChildStderr, tail: StderrTail) {
792    let mut reader = BufReader::new(stderr);
793
794    loop {
795        let line = match read_bounded_line(&mut reader, STDERR_MAX_LINE_BYTES).await {
796            Ok(BoundedLine::Line(line)) => line,
797            Ok(BoundedLine::Oversized { prefix, scanned }) => {
798                tracing::warn!(
799                    scanned,
800                    limit = STDERR_MAX_LINE_BYTES,
801                    "worker stderr line exceeded the per-line limit; retaining a truncated prefix"
802                );
803                format!("{prefix}{TRUNCATION_MARKER}")
804            },
805            Ok(BoundedLine::Eof) => break,
806            Err(e) => {
807                tracing::debug!("error reading worker stderr: {e}");
808                break;
809            },
810        };
811
812        tracing::debug!(target: "cpex_hosts_python::worker_stderr", "{line}");
813
814        let mut tail = tail.lock().await;
815        if tail.len() == STDERR_RETAIN_LINES {
816            tail.pop_front();
817        }
818        tail.push_back(line);
819    }
820}
821
822/// Resolve `worker.py` inside a venv's site-packages.
823///
824/// The worker ships with the `cpex` framework, which the plugin's requirements
825/// pull into the venv — so the host does not carry its own copy. Both
826/// `lib/pythonX.Y/site-packages` (Unix) and `Lib/site-packages` (Windows) are
827/// searched because the interpreter version is not known up front.
828pub fn resolve_worker_script(venv_path: &Path, relative: &str) -> Result<PathBuf, HostError> {
829    // An absolute override (used by tests and by hosts that vendor a worker)
830    // bypasses venv resolution entirely.
831    let relative_path = Path::new(relative);
832    if relative_path.is_absolute() {
833        return Ok(relative_path.to_path_buf());
834    }
835
836    let mut candidates = Vec::new();
837
838    let unix_lib = venv_path.join("lib");
839    if let Ok(entries) = std::fs::read_dir(&unix_lib) {
840        for entry in entries.flatten() {
841            candidates.push(entry.path().join("site-packages").join(relative_path));
842        }
843    }
844    candidates.push(
845        venv_path
846            .join("Lib")
847            .join("site-packages")
848            .join(relative_path),
849    );
850    // Last resort: a worker sitting directly under the venv.
851    candidates.push(venv_path.join(relative_path));
852
853    candidates
854        .iter()
855        .find(|p| p.exists())
856        .cloned()
857        .ok_or_else(|| HostError::WorkerStart {
858            message: format!(
859                "could not find '{relative}' in the venv at {} — is `cpex` installed in it? \
860                 (a plugin's requirements file must pull in the framework that ships worker.py)",
861                venv_path.display()
862            ),
863        })
864}
865
866#[cfg(test)]
867mod tests {
868    use super::*;
869    use crate::testing::{TempDir, skip_without_python3};
870
871    /// A stub worker that speaks the same protocol as `worker.py`.
872    ///
873    /// Deterministic stand-in for the real thing: no venv, no framework
874    /// import, and each behavior selectable per task so the failure modes are
875    /// reachable without contriving a real plugin to misbehave.
876    ///
877    /// Recognized `task_type` values:
878    ///   - `echo`      — reply with `{"status":"success","echo":<payload>}`
879    ///   - `slow`      — sleep `delay` seconds, then echo (drives the timeout)
880    ///   - `fail`      — reply `{"status":"error","message":<message>}`
881    ///   - `die`       — exit immediately without replying
882    ///   - `noisy`     — print a non-JSON line first, then reply
883    ///   - `huge`      — reply with a padded response of `size` bytes
884    ///   - `unframed`  — stream `size` bytes with no newline at all, then reply
885    ///   - `loud`      — write a `size`-byte stderr line, then reply
886    ///   - `deaf`      — ignore `shutdown` and keep running (drives the kill)
887    ///   - `shutdown`  — reply and exit, unless previously put in `deaf` mode
888    const STUB_WORKER: &str = r#"
889import json, sys, time
890
891deaf = False
892while True:
893    line = sys.stdin.readline()
894    if not line:
895        break
896    line = line.strip()
897    if not line:
898        continue
899    task = json.loads(line)
900    kind = task.get("task_type")
901    rid = task.get("request_id", "unknown")
902
903    if kind == "shutdown":
904        if deaf:
905            # Acknowledge but refuse to exit, so the host must kill us.
906            print(json.dumps({"status": "success", "request_id": rid}), flush=True)
907            while True:
908                time.sleep(0.05)
909        print(json.dumps({"status": "success", "message": "Shutting down", "request_id": rid}), flush=True)
910        break
911    if kind == "deaf":
912        deaf = True
913        print(json.dumps({"status": "success", "request_id": rid}), flush=True)
914        continue
915    if kind == "die":
916        sys.exit(3)
917    if kind == "slow":
918        time.sleep(float(task.get("delay", 1.0)))
919        print(json.dumps({"status": "success", "echo": task.get("payload"), "request_id": rid}), flush=True)
920        continue
921    if kind == "fail":
922        print(json.dumps({"status": "error", "message": task.get("message", "stub failure"), "request_id": rid}), flush=True)
923        continue
924    if kind == "noisy":
925        print("this is not JSON and must not desync the stream", flush=True)
926        print(json.dumps({"status": "success", "echo": task.get("payload"), "request_id": rid}), flush=True)
927        continue
928    if kind == "huge":
929        # request_id first, so the host can still name the caller from the
930        # truncated prefix it kept.
931        size = int(task.get("size", 1024))
932        print(json.dumps({"request_id": rid, "status": "success", "pad": "P" * size}), flush=True)
933        continue
934    if kind == "unframed":
935        # No newline at all: the pathological case for a line reader.
936        size = int(task.get("size", 1024))
937        sys.stdout.write("U" * size)
938        sys.stdout.flush()
939        print(json.dumps({"request_id": rid, "status": "success"}), flush=True)
940        continue
941    if kind == "loud":
942        size = int(task.get("size", 1024))
943        sys.stderr.write("L" * size + "\n")
944        sys.stderr.flush()
945        print(json.dumps({"status": "success", "echo": task.get("payload"), "request_id": rid}), flush=True)
946        continue
947
948    print(json.dumps({"status": "success", "echo": task.get("payload"), "request_id": rid}), flush=True)
949"#;
950
951    /// Write the stub to disk and spawn a client against it.
952    async fn stub_client(
953        dir: &TempDir,
954        max_content_size: usize,
955        timeout_secs: u64,
956    ) -> WorkerClient {
957        let script = dir.path().join("stub_worker.py");
958        std::fs::write(&script, STUB_WORKER).unwrap();
959
960        let python = which_python3();
961        WorkerClient::spawn(
962            &python,
963            &script,
964            Some(dir.path()),
965            max_content_size,
966            timeout_secs,
967        )
968        .await
969        .expect("stub worker spawns")
970    }
971
972    /// Absolute path to python3, since `WorkerClient::spawn` requires an
973    /// existing interpreter path rather than a PATH lookup.
974    fn which_python3() -> PathBuf {
975        let out = std::process::Command::new("sh")
976            .args(["-c", "command -v python3"])
977            .output()
978            .expect("run command -v");
979        PathBuf::from(String::from_utf8_lossy(&out.stdout).trim())
980    }
981
982    #[tokio::test]
983    async fn a_task_round_trips_to_its_response() {
984        if skip_without_python3("a_task_round_trips_to_its_response") {
985            return;
986        }
987        let dir = TempDir::new();
988        let client = stub_client(&dir, 1_000_000, 10).await;
989
990        let response = client
991            .send_task(serde_json::json!({ "task_type": "echo", "payload": {"tool": "search"} }))
992            .await
993            .expect("the stub replies");
994
995        assert_eq!(response["echo"]["tool"], "search");
996        assert!(
997            response.get("request_id").is_none(),
998            "request_id is transport bookkeeping and must be stripped"
999        );
1000    }
1001
1002    #[tokio::test]
1003    async fn concurrent_sends_each_receive_their_own_response() {
1004        // Demux correctness. The stub answers the slower task second, so a
1005        // client that assumed FIFO ordering would hand each caller the other's
1006        // response — this test fails loudly on that.
1007        if skip_without_python3("concurrent_sends_each_receive_their_own_response") {
1008            return;
1009        }
1010        let dir = TempDir::new();
1011        let client = Arc::new(stub_client(&dir, 1_000_000, 10).await);
1012
1013        let slow = {
1014            let client = Arc::clone(&client);
1015            tokio::spawn(async move {
1016                client
1017                    .send_task(serde_json::json!({
1018                        "task_type": "slow", "delay": 0.4, "payload": "slow-one"
1019                    }))
1020                    .await
1021            })
1022        };
1023
1024        // Give the slow task a head start so it is genuinely outstanding.
1025        tokio::time::sleep(std::time::Duration::from_millis(80)).await;
1026
1027        let fast = client
1028            .send_task(serde_json::json!({ "task_type": "echo", "payload": "fast-one" }))
1029            .await
1030            .expect("the fast task replies");
1031
1032        let slow = slow.await.unwrap().expect("the slow task replies");
1033
1034        assert_eq!(fast["echo"], "fast-one");
1035        assert_eq!(slow["echo"], "slow-one");
1036    }
1037
1038    #[tokio::test]
1039    async fn many_concurrent_sends_all_resolve_correctly() {
1040        // Scales the demux check past two, where an off-by-one in the pending
1041        // map would still look fine.
1042        if skip_without_python3("many_concurrent_sends_all_resolve_correctly") {
1043            return;
1044        }
1045        let dir = TempDir::new();
1046        let client = Arc::new(stub_client(&dir, 1_000_000, 20).await);
1047
1048        let mut handles = Vec::with_capacity(16);
1049        for i in 0..16 {
1050            let client = Arc::clone(&client);
1051            handles.push(tokio::spawn(async move {
1052                let response = client
1053                    .send_task(serde_json::json!({ "task_type": "echo", "payload": i }))
1054                    .await
1055                    .expect("each task replies");
1056                (i, response["echo"].as_i64().unwrap())
1057            }));
1058        }
1059
1060        for handle in handles {
1061            let (sent, echoed) = handle.await.unwrap();
1062            assert_eq!(
1063                sent, echoed,
1064                "response {echoed} was delivered to the caller that sent {sent}"
1065            );
1066        }
1067    }
1068
1069    #[tokio::test]
1070    async fn an_oversized_task_is_rejected_before_it_is_sent() {
1071        if skip_without_python3("an_oversized_task_is_rejected_before_it_is_sent") {
1072            return;
1073        }
1074        let dir = TempDir::new();
1075        let client = stub_client(&dir, 256, 10).await;
1076
1077        let err = client
1078            .send_task(serde_json::json!({ "task_type": "echo", "payload": "x".repeat(1024) }))
1079            .await
1080            .expect_err("a task over the cap must be rejected");
1081
1082        match err {
1083            HostError::TaskTooLarge { size, limit } => {
1084                assert!(size > limit);
1085                assert_eq!(limit, 256);
1086            },
1087            other => panic!("expected TaskTooLarge, got {other:?}"),
1088        }
1089
1090        // The client must still be usable — rejecting a task locally neither
1091        // wrote to the pipe nor desynced the worker.
1092        let response = client
1093            .send_task(serde_json::json!({ "task_type": "echo", "payload": "small" }))
1094            .await
1095            .expect("the client survives a rejected task");
1096        assert_eq!(response["echo"], "small");
1097    }
1098
1099    // --- inbound bounds ------------------------------------------------------
1100
1101    #[tokio::test]
1102    async fn an_oversized_response_frame_is_rejected_rather_than_buffered() {
1103        // The reviewer's finding: `max_content_size` was enforced only outbound,
1104        // so a worker could return a frame of any size and the host would grow
1105        // its heap to hold it. Before the fix this test hangs to the timeout
1106        // (the giant frame is buffered, parsed, and delivered as a success);
1107        // after it, the caller gets a prompt protocol error.
1108        if skip_without_python3("an_oversized_response_frame_is_rejected_rather_than_buffered") {
1109            return;
1110        }
1111        let dir = TempDir::new();
1112        // Big enough for the outbound task, far smaller than the reply.
1113        let client = stub_client(&dir, 4096, 10).await;
1114
1115        let err = client
1116            .send_task(serde_json::json!({ "task_type": "huge", "size": 300_000 }))
1117            .await
1118            .expect_err("a response over the cap cannot be delivered");
1119
1120        match err {
1121            HostError::ResponseTooLarge { size, limit } => {
1122                assert_eq!(limit, 4096, "the error carries the configured cap");
1123                assert!(
1124                    size >= limit,
1125                    "the reported size must be at least the cap it broke: {size}"
1126                );
1127                let message = err.to_string();
1128                assert!(
1129                    message.contains("4096") && message.contains("discarded"),
1130                    "the error should name the limit and say the frame was dropped: {message}"
1131                );
1132                assert!(
1133                    !message.contains("PPPP"),
1134                    "the error must not echo the frame it refused: {message}"
1135                );
1136            },
1137            other => panic!("expected a ResponseTooLarge error, got {other:?}"),
1138        }
1139
1140        // Framing survived: the discarded frame was consumed to its newline, so
1141        // the next response lands on a real line boundary.
1142        let response = client
1143            .send_task(serde_json::json!({ "task_type": "echo", "payload": "still-framed" }))
1144            .await
1145            .expect("the stream is still in sync after a dropped frame");
1146        assert_eq!(response["echo"], "still-framed");
1147    }
1148
1149    #[tokio::test]
1150    async fn an_unterminated_response_stream_does_not_grow_without_bound() {
1151        // The worst case for a line reader: bytes with no newline in sight.
1152        // `lines()` would buffer all of them; `read_bounded_line` stops
1153        // retaining at the cap and discards the rest as it arrives.
1154        if skip_without_python3("an_unterminated_response_stream_does_not_grow_without_bound") {
1155            return;
1156        }
1157        let dir = TempDir::new();
1158        let client = stub_client(&dir, 4096, 10).await;
1159
1160        let err = client
1161            .send_task(serde_json::json!({ "task_type": "unframed", "size": 500_000 }))
1162            .await
1163            .expect_err("an unframed flood cannot answer the request");
1164        assert!(
1165            matches!(err, HostError::ResponseTooLarge { .. }),
1166            "got {err:?}"
1167        );
1168
1169        // The stub appended a real response after the flood; the reader must
1170        // have resynced on that newline rather than staying wedged.
1171        let response = client
1172            .send_task(serde_json::json!({ "task_type": "echo", "payload": "resynced" }))
1173            .await
1174            .expect("the reader recovers at the next newline");
1175        assert_eq!(response["echo"], "resynced");
1176    }
1177
1178    #[tokio::test]
1179    async fn oversized_stderr_is_truncated_rather_than_killing_the_worker() {
1180        // Stderr is diagnostic, so the policy differs from stdout: bound it and
1181        // mark the cut, but keep the worker alive — this is the channel that
1182        // explains a failure, and losing the worker over a verbose traceback
1183        // would trade a diagnostic for an outage.
1184        if skip_without_python3("oversized_stderr_is_truncated_rather_than_killing_the_worker") {
1185            return;
1186        }
1187        let dir = TempDir::new();
1188        let client = stub_client(&dir, 1_000_000, 10).await;
1189
1190        let response = client
1191            .send_task(serde_json::json!({
1192                "task_type": "loud", "size": 200_000, "payload": "survived"
1193            }))
1194            .await
1195            .expect("a loud worker is still a working worker");
1196        assert_eq!(response["echo"], "survived");
1197        assert!(client.is_alive());
1198
1199        let lines = client.stderr_lines().await;
1200        let loud = lines
1201            .iter()
1202            .find(|line| line.starts_with("LLL"))
1203            .expect("the loud line was retained");
1204
1205        assert!(
1206            loud.ends_with(TRUNCATION_MARKER),
1207            "a cut line must say so, or an operator reads it as the worker stopping mid-sentence"
1208        );
1209        assert!(
1210            loud.len() <= STDERR_MAX_LINE_BYTES + TRUNCATION_MARKER.len(),
1211            "the retained line is bounded: {} bytes",
1212            loud.len()
1213        );
1214        // And the prefix is genuinely useful, not an empty stub.
1215        assert!(loud.len() > STDERR_MAX_LINE_BYTES / 2);
1216    }
1217
1218    #[tokio::test]
1219    async fn total_retained_stderr_stays_bounded_under_a_flood() {
1220        // Per-line truncation alone is not enough: many long lines must also be
1221        // capped, which is what the retain ring does. Together they put a hard
1222        // ceiling on stderr memory.
1223        if skip_without_python3("total_retained_stderr_stays_bounded_under_a_flood") {
1224            return;
1225        }
1226        let dir = TempDir::new();
1227        let client = stub_client(&dir, 1_000_000, 20).await;
1228
1229        for _ in 0..(STDERR_RETAIN_LINES * 2) {
1230            client
1231                .send_task(serde_json::json!({
1232                    "task_type": "loud", "size": 20_000, "payload": "ok"
1233                }))
1234                .await
1235                .expect("the worker keeps serving");
1236        }
1237
1238        let lines = client.stderr_lines().await;
1239        assert!(
1240            lines.len() <= STDERR_RETAIN_LINES,
1241            "the ring must cap line count: {}",
1242            lines.len()
1243        );
1244        let total: usize = lines.iter().map(String::len).sum();
1245        assert!(
1246            total <= STDERR_RETAIN_LINES * (STDERR_MAX_LINE_BYTES + TRUNCATION_MARKER.len()),
1247            "retained stderr must have a hard ceiling: {total} bytes"
1248        );
1249    }
1250
1251    // --- bounded line reading (no subprocess needed) -------------------------
1252
1253    #[tokio::test]
1254    async fn a_bounded_read_returns_whole_lines_under_the_limit() {
1255        let mut reader = BufReader::new(&b"first\nsecond\n"[..]);
1256
1257        for expected in ["first", "second"] {
1258            match read_bounded_line(&mut reader, 64).await.unwrap() {
1259                BoundedLine::Line(line) => assert_eq!(line, expected),
1260                other => panic!("expected a line, got {}", describe(&other)),
1261            }
1262        }
1263        assert!(matches!(
1264            read_bounded_line(&mut reader, 64).await.unwrap(),
1265            BoundedLine::Eof
1266        ));
1267    }
1268
1269    #[tokio::test]
1270    async fn a_bounded_read_reports_the_scanned_size_and_keeps_only_the_cap() {
1271        let input = format!("{}\nafter\n", "x".repeat(500));
1272        let mut reader = BufReader::new(input.as_bytes());
1273
1274        match read_bounded_line(&mut reader, 100).await.unwrap() {
1275            BoundedLine::Oversized { prefix, scanned } => {
1276                assert_eq!(prefix.len(), 100, "nothing past the cap is retained");
1277                assert_eq!(scanned, 500, "the true size is still reported");
1278            },
1279            other => panic!("expected Oversized, got {}", describe(&other)),
1280        }
1281
1282        // Resync: the oversized line was consumed through its newline.
1283        match read_bounded_line(&mut reader, 100).await.unwrap() {
1284            BoundedLine::Line(line) => assert_eq!(line, "after"),
1285            other => panic!("expected the next line, got {}", describe(&other)),
1286        }
1287    }
1288
1289    #[tokio::test]
1290    async fn a_line_exactly_at_the_limit_is_not_oversized() {
1291        // Off-by-one guard: the cap is the largest *acceptable* size, matching
1292        // the outbound `line.len() > max_content_size` comparison.
1293        let input = "xxxxx\n";
1294        let mut reader = BufReader::new(input.as_bytes());
1295        match read_bounded_line(&mut reader, 5).await.unwrap() {
1296            BoundedLine::Line(line) => assert_eq!(line, "xxxxx"),
1297            other => panic!(
1298                "5 bytes under a 5-byte cap is fine, got {}",
1299                describe(&other)
1300            ),
1301        }
1302    }
1303
1304    #[tokio::test]
1305    async fn a_trailing_fragment_without_a_newline_is_still_a_line() {
1306        let mut reader = BufReader::new(&b"no trailing newline"[..]);
1307        match read_bounded_line(&mut reader, 64).await.unwrap() {
1308            BoundedLine::Line(line) => assert_eq!(line, "no trailing newline"),
1309            other => panic!("expected a line, got {}", describe(&other)),
1310        }
1311    }
1312
1313    #[tokio::test]
1314    async fn a_bounded_read_trims_crlf_like_lines_does() {
1315        let mut reader = BufReader::new(&b"windows\r\n"[..]);
1316        match read_bounded_line(&mut reader, 64).await.unwrap() {
1317            BoundedLine::Line(line) => assert_eq!(line, "windows"),
1318            other => panic!("expected a line, got {}", describe(&other)),
1319        }
1320    }
1321
1322    /// Name a `BoundedLine` for a panic message, since it has no `Debug`.
1323    fn describe(line: &BoundedLine) -> &'static str {
1324        match line {
1325            BoundedLine::Line(_) => "Line",
1326            BoundedLine::Oversized { .. } => "Oversized",
1327            BoundedLine::Eof => "Eof",
1328        }
1329    }
1330
1331    #[test]
1332    fn a_request_id_is_recovered_from_a_truncated_frame() {
1333        // What lets the reader fail exactly the one caller whose response was
1334        // dropped instead of every in-flight request.
1335        assert_eq!(
1336            request_id_from_prefix(r#"{"request_id": "abc-123", "status": "succ"#),
1337            Some("abc-123".to_string())
1338        );
1339        assert_eq!(
1340            request_id_from_prefix(r#"{"status":"success","request_id":"xyz","pad":"PPP"#),
1341            Some("xyz".to_string())
1342        );
1343    }
1344
1345    #[test]
1346    fn an_unrecoverable_request_id_is_reported_as_absent() {
1347        // Each of these must be a clean `None` — a wrong guess would fail the
1348        // wrong caller, which is worse than failing all of them.
1349        for prefix in [
1350            r#"{"status":"success","pad":"PPPP"#, // id not in the prefix
1351            r#"{"request_id": "cut-off-mid-val"#, // no closing quote
1352            r#"{"request_id""#,                   // no colon
1353            "",                                   // nothing at all
1354        ] {
1355            assert_eq!(
1356                request_id_from_prefix(prefix),
1357                None,
1358                "should not guess an id from {prefix:?}"
1359            );
1360        }
1361    }
1362
1363    #[test]
1364    fn the_over_limit_error_names_the_response_not_the_task() {
1365        // Reusing `TaskTooLarge` would point an operator at their own outbound
1366        // payload for a fault in the worker's reply.
1367        let err = response_too_large(9_000, 4_096);
1368        let message = err.to_string();
1369        assert!(matches!(
1370            err,
1371            HostError::ResponseTooLarge {
1372                size: 9_000,
1373                limit: 4_096
1374            }
1375        ));
1376        assert_eq!(err.code(), "response_too_large");
1377        assert!(message.contains("response"), "{message}");
1378        assert!(!message.contains("serialized task"), "{message}");
1379        assert!(
1380            message.contains("9000") && message.contains("4096"),
1381            "{message}"
1382        );
1383    }
1384
1385    #[tokio::test]
1386    async fn no_response_within_the_timeout_yields_a_timeout_error() {
1387        if skip_without_python3("no_response_within_the_timeout_yields_a_timeout_error") {
1388            return;
1389        }
1390        let dir = TempDir::new();
1391        let client = stub_client(&dir, 1_000_000, 1).await;
1392
1393        let err = client
1394            .send_task(serde_json::json!({ "task_type": "slow", "delay": 5.0 }))
1395            .await
1396            .expect_err("a task that outruns the timeout must error");
1397
1398        assert!(
1399            matches!(err, HostError::Timeout { timeout_secs: 1 }),
1400            "got {err:?}"
1401        );
1402    }
1403
1404    #[tokio::test]
1405    async fn a_worker_error_response_becomes_a_worker_error() {
1406        if skip_without_python3("a_worker_error_response_becomes_a_worker_error") {
1407            return;
1408        }
1409        let dir = TempDir::new();
1410        let client = stub_client(&dir, 1_000_000, 10).await;
1411
1412        let err = client
1413            .send_task(serde_json::json!({ "task_type": "fail", "message": "plugin blew up" }))
1414            .await
1415            .expect_err("status=error must not be returned as success");
1416
1417        match err {
1418            HostError::WorkerError { message } => assert!(message.contains("plugin blew up")),
1419            other => panic!("expected WorkerError, got {other:?}"),
1420        }
1421    }
1422
1423    #[tokio::test]
1424    async fn worker_death_mid_flight_resolves_rather_than_hangs() {
1425        // The important one: a dead worker must not leave a caller waiting out
1426        // a timeout that can never be satisfied. The generous 30s client
1427        // timeout means a hang here would show up as a test that runs for 30s
1428        // and then reports Timeout instead of WorkerDied.
1429        if skip_without_python3("worker_death_mid_flight_resolves_rather_than_hangs") {
1430            return;
1431        }
1432        let dir = TempDir::new();
1433        let client = stub_client(&dir, 1_000_000, 30).await;
1434
1435        let started = std::time::Instant::now();
1436        let err = client
1437            .send_task(serde_json::json!({ "task_type": "die" }))
1438            .await
1439            .expect_err("a worker that exits without replying must error");
1440
1441        assert!(
1442            matches!(err, HostError::WorkerDied { .. }),
1443            "expected WorkerDied (not Timeout), got {err:?}"
1444        );
1445        assert!(
1446            started.elapsed() < std::time::Duration::from_secs(5),
1447            "the error must arrive on worker death, not on timeout expiry"
1448        );
1449    }
1450
1451    #[tokio::test]
1452    async fn sends_after_worker_death_fail_fast() {
1453        if skip_without_python3("sends_after_worker_death_fail_fast") {
1454            return;
1455        }
1456        let dir = TempDir::new();
1457        let client = stub_client(&dir, 1_000_000, 30).await;
1458
1459        let _ = client
1460            .send_task(serde_json::json!({ "task_type": "die" }))
1461            .await;
1462
1463        // A send issued *after* the death must also fail fast. Clearing
1464        // `pending` on EOF only rescues already-registered requests; without an
1465        // explicit liveness flag this send would register into an unwatched map
1466        // and wait out the full 30s timeout.
1467        let started = std::time::Instant::now();
1468        let err = client
1469            .send_task(serde_json::json!({ "task_type": "echo", "payload": "anyone home?" }))
1470            .await
1471            .expect_err("a send to a dead worker cannot succeed");
1472
1473        assert!(
1474            matches!(err, HostError::WorkerDied { .. }),
1475            "expected WorkerDied, got {err:?}"
1476        );
1477        assert!(
1478            started.elapsed() < std::time::Duration::from_secs(2),
1479            "a send to a known-dead worker must fail immediately, not on timeout: {:?}",
1480            started.elapsed()
1481        );
1482        assert!(!client.is_alive());
1483    }
1484
1485    #[tokio::test]
1486    async fn a_non_json_stdout_line_does_not_desync_the_stream() {
1487        // Plugins print. A stray line must be skipped, not consumed as if it
1488        // were the response to the pending request.
1489        if skip_without_python3("a_non_json_stdout_line_does_not_desync_the_stream") {
1490            return;
1491        }
1492        let dir = TempDir::new();
1493        let client = stub_client(&dir, 1_000_000, 10).await;
1494
1495        let response = client
1496            .send_task(serde_json::json!({ "task_type": "noisy", "payload": "after-the-noise" }))
1497            .await
1498            .expect("the real response still arrives");
1499        assert_eq!(response["echo"], "after-the-noise");
1500    }
1501
1502    #[tokio::test]
1503    async fn shutdown_terminates_a_live_worker() {
1504        if skip_without_python3("shutdown_terminates_a_live_worker") {
1505            return;
1506        }
1507        let dir = TempDir::new();
1508        let client = stub_client(&dir, 1_000_000, 10).await;
1509
1510        client
1511            .send_task(serde_json::json!({ "task_type": "echo", "payload": "alive" }))
1512            .await
1513            .expect("worker is up");
1514
1515        client.shutdown().await.expect("a cooperative worker exits");
1516        // Idempotent: teardown can be reached more than once.
1517        client
1518            .shutdown()
1519            .await
1520            .expect("a second shutdown is a no-op");
1521    }
1522
1523    #[tokio::test]
1524    async fn a_worker_that_ignores_shutdown_is_killed_after_the_grace_window() {
1525        if skip_without_python3("a_worker_that_ignores_shutdown_is_killed_after_the_grace_window") {
1526            return;
1527        }
1528        let dir = TempDir::new();
1529        let client = stub_client(&dir, 1_000_000, 10).await;
1530
1531        // Put the stub into a mode where it acknowledges shutdown and then
1532        // spins forever.
1533        client
1534            .send_task(serde_json::json!({ "task_type": "deaf" }))
1535            .await
1536            .expect("stub accepts deaf mode");
1537
1538        let started = std::time::Instant::now();
1539        client
1540            .shutdown()
1541            .await
1542            .expect("an unresponsive worker is killed, not waited on forever");
1543
1544        let elapsed = started.elapsed();
1545        assert!(
1546            elapsed >= std::time::Duration::from_secs(SHUTDOWN_GRACE_SECS),
1547            "the grace window should be honored before killing: {elapsed:?}"
1548        );
1549        assert!(
1550            elapsed < std::time::Duration::from_secs(SHUTDOWN_GRACE_SECS + 5),
1551            "the kill must happen promptly once the window expires: {elapsed:?}"
1552        );
1553    }
1554
1555    #[tokio::test]
1556    async fn a_death_error_carries_the_workers_stderr() {
1557        // A worker that dies during startup — a bad dependency pin, a missing
1558        // module — explains itself on stderr and nowhere else. Without the
1559        // tail, the host reports a bare "worker process died" and the operator
1560        // has nothing to act on.
1561        if skip_without_python3("a_death_error_carries_the_workers_stderr") {
1562            return;
1563        }
1564        let dir = TempDir::new();
1565        let script = dir.path().join("import_error.py");
1566        std::fs::write(
1567            &script,
1568            "import sys\nprint('ImportError: cannot import name McpError', file=sys.stderr, flush=True)\nsys.exit(1)\n",
1569        )
1570        .unwrap();
1571
1572        let client =
1573            WorkerClient::spawn(&which_python3(), &script, Some(dir.path()), 1_000_000, 30)
1574                .await
1575                .expect("the process starts, then exits");
1576
1577        let err = client
1578            .send_task(serde_json::json!({ "task_type": "echo" }))
1579            .await
1580            .expect_err("a dead worker cannot answer");
1581
1582        let message = err.to_string();
1583        assert!(matches!(err, HostError::WorkerDied { .. }), "got {err:?}");
1584        assert!(
1585            message.contains("McpError"),
1586            "the death error should carry the stderr that explains it: {message}"
1587        );
1588    }
1589
1590    #[tokio::test]
1591    async fn spawning_against_a_missing_interpreter_errors() {
1592        let err = WorkerClient::spawn(
1593            Path::new("/nonexistent/bin/python"),
1594            Path::new("/tmp/worker.py"),
1595            None,
1596            1024,
1597            5,
1598        )
1599        .await
1600        .expect_err("a missing interpreter cannot be launched");
1601
1602        assert!(matches!(err, HostError::WorkerStart { .. }), "got {err:?}");
1603    }
1604
1605    // --- response interpretation (no subprocess needed) ---------------------
1606
1607    #[test]
1608    fn an_error_envelope_becomes_a_worker_error() {
1609        let err = interpret_response(serde_json::json!({
1610            "status": "error", "message": "boom", "request_id": "abc"
1611        }))
1612        .expect_err("status=error is a failure");
1613        match err {
1614            HostError::WorkerError { message } => assert_eq!(message, "boom"),
1615            other => panic!("expected WorkerError, got {other:?}"),
1616        }
1617    }
1618
1619    #[test]
1620    fn an_error_envelope_without_a_message_still_errors() {
1621        let err = interpret_response(serde_json::json!({ "status": "error" }))
1622            .expect_err("a message-less error is still an error");
1623        assert!(matches!(err, HostError::WorkerError { .. }));
1624    }
1625
1626    #[test]
1627    fn a_success_envelope_is_returned_without_its_request_id() {
1628        let response = interpret_response(serde_json::json!({
1629            "status": "success", "continue_processing": true, "request_id": "abc"
1630        }))
1631        .expect("success passes through");
1632        assert!(response.get("request_id").is_none());
1633        assert_eq!(response["continue_processing"], true);
1634    }
1635
1636    #[test]
1637    fn a_result_envelope_with_no_status_field_is_success() {
1638        // The hook path returns `model_dump()` of a result model, which has no
1639        // `status` key at all — treating "no status" as an error would fail
1640        // every successful invocation.
1641        let response = interpret_response(serde_json::json!({
1642            "continue_processing": false, "violation": {"code": "pii"}
1643        }))
1644        .expect("a bare result model is not an error");
1645        assert_eq!(response["violation"]["code"], "pii");
1646    }
1647
1648    // --- worker script resolution ------------------------------------------
1649
1650    #[test]
1651    fn an_absolute_script_path_bypasses_venv_resolution() {
1652        let resolved = resolve_worker_script(Path::new("/venv"), "/opt/custom/worker.py").unwrap();
1653        assert_eq!(resolved, PathBuf::from("/opt/custom/worker.py"));
1654    }
1655
1656    #[test]
1657    fn the_worker_script_is_found_in_unix_site_packages() {
1658        let dir = TempDir::new();
1659        let site = dir
1660            .path()
1661            .join("lib")
1662            .join("python3.12")
1663            .join("site-packages");
1664        let worker = site.join("cpex").join("framework").join("isolated");
1665        std::fs::create_dir_all(&worker).unwrap();
1666        std::fs::write(worker.join("worker.py"), "# stub").unwrap();
1667
1668        let resolved =
1669            resolve_worker_script(dir.path(), "cpex/framework/isolated/worker.py").unwrap();
1670        assert_eq!(resolved, worker.join("worker.py"));
1671    }
1672
1673    #[test]
1674    fn a_missing_worker_script_names_the_venv_and_the_likely_cause() {
1675        let dir = TempDir::new();
1676        let err = resolve_worker_script(dir.path(), "cpex/framework/isolated/worker.py")
1677            .expect_err("an empty venv has no worker");
1678
1679        let message = err.to_string();
1680        assert!(message.contains("worker.py"));
1681        assert!(
1682            message.contains("cpex"),
1683            "the message should point at the missing framework: {message}"
1684        );
1685    }
1686}