1use 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
44const SHUTDOWN_GRACE_SECS: u64 = 5;
47
48pub const TASK_LOAD_AND_RUN_HOOK: &str = "load_and_run_hook";
50
51pub const TASK_SHUTDOWN: &str = "shutdown";
53
54pub const TASK_CAPABILITIES: &str = "capabilities";
61
62pub const FEATURE_CREDENTIAL: &str = "credential";
64
65pub const FEATURE_EXTENSIONS: &str = "extensions";
67
68pub const FEATURE_CMF_HOOKS: &str = "cmf_hooks";
74
75#[derive(Debug, Clone, Default, PartialEq, Eq)]
81pub struct WorkerCapabilities {
82 pub protocol_version: u64,
84
85 pub features: BTreeSet<String>,
87
88 pub cpex_version: Option<String>,
94}
95
96impl WorkerCapabilities {
97 pub fn none() -> Self {
99 Self::default()
100 }
101
102 pub fn has(&self, feature: &str) -> bool {
104 self.features.contains(feature)
105 }
106}
107
108type Pending = Arc<Mutex<HashMap<String, oneshot::Sender<Result<Value, HostError>>>>>;
116
117const STDERR_RETAIN_LINES: usize = 40;
124
125type StderrTail = Arc<Mutex<std::collections::VecDeque<String>>>;
127
128const STDERR_MAX_LINE_BYTES: usize = 8 * 1024;
135
136const TRUNCATION_MARKER: &str = "… [truncated by host]";
140
141enum BoundedLine {
143 Line(String),
145
146 Oversized { prefix: String, scanned: usize },
150
151 Eof,
153}
154
155async 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 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 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 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
224type Alive = Arc<std::sync::atomic::AtomicBool>;
233
234pub struct WorkerClient {
236 outbound: mpsc::UnboundedSender<String>,
238
239 pending: Pending,
241
242 alive: Alive,
244
245 stderr_tail: StderrTail,
247
248 child: Mutex<Option<Child>>,
251
252 max_content_size: usize,
254
255 timeout_secs: u64,
257
258 reported: std::sync::OnceLock<WorkerCapabilities>,
264
265 probe_retry_spent: std::sync::atomic::AtomicBool,
268}
269
270impl 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 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 .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 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 pub fn claim_probe_retry(&self) -> bool {
370 !self
371 .probe_retry_spent
372 .swap(true, std::sync::atomic::Ordering::SeqCst)
373 }
374
375 pub fn reported_capabilities(&self) -> Option<&WorkerCapabilities> {
378 self.reported.get()
379 }
380
381 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 pub async fn stderr_lines(&self) -> Vec<String> {
405 self.stderr_tail.lock().await.iter().cloned().collect()
406 }
407
408 pub fn is_alive(&self) -> bool {
415 self.alive.load(std::sync::atomic::Ordering::Acquire)
416 }
417
418 pub async fn send_task(&self, mut task: Value) -> Result<Value, HostError> {
423 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 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 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 Ok(Ok(Err(e))) => return Err(e),
470 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 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 pub async fn shutdown(&self) -> Result<(), HostError> {
554 let Some(mut child) = self.child.lock().await.take() else {
555 return Ok(());
556 };
557
558 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 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
588fn 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
609async 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
629async 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 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 let _ = tx.send(Ok(response));
672 },
673 None => {
674 tracing::debug!(request_id, "response for an unknown request id");
677 },
678 }
679 },
680 Ok(BoundedLine::Oversized { prefix, scanned }) => {
681 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 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 alive.store(false, std::sync::atomic::Ordering::Release);
733
734 pending.lock().await.clear();
737}
738
739fn response_too_large(size: usize, limit: usize) -> HostError {
747 HostError::ResponseTooLarge { size, limit }
748}
749
750fn 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 let closing = rest.find('"')?;
766 Some(rest[..closing].to_string())
767}
768
769async 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
822pub fn resolve_worker_script(venv_path: &Path, relative: &str) -> Result<PathBuf, HostError> {
829 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 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 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 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 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 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 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 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 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 #[tokio::test]
1102 async fn an_oversized_response_frame_is_rejected_rather_than_buffered() {
1103 if skip_without_python3("an_oversized_response_frame_is_rejected_rather_than_buffered") {
1109 return;
1110 }
1111 let dir = TempDir::new();
1112 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 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 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 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 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 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 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 #[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 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 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 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 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 for prefix in [
1350 r#"{"status":"success","pad":"PPPP"#, r#"{"request_id": "cut-off-mid-val"#, r#"{"request_id""#, "", ] {
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 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 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 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 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 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 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 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 #[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 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 #[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}