1use bincode::de::Decoder;
43use bincode::enc::Encoder;
44use bincode::error::{DecodeError, EncodeError};
45use bincode::{Decode, Encode};
46use cu29::prelude::*;
47use cu29::reflect::Reflect;
48use cu29_value::{Value as CuValue, py_to_value, to_value, value_to_py};
49use minicbor::data::{IanaTag, Int as CborInt, Type as CborType};
50use minicbor::{Decoder as CborDecoder, Encoder as CborEncoder};
51use serde::de::DeserializeOwned;
52use serde::{Deserialize, Serialize};
53use std::collections::BTreeMap;
54use std::ffi::CString;
55use std::io::{BufReader, Read, Write};
56use std::marker::PhantomData;
57use std::path::{Path, PathBuf};
58use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
59use std::sync::OnceLock;
60
61use pyo3::prelude::*;
62use pyo3::types::{PyAny, PyModule, PyTuple};
63
64const DEFAULT_SCRIPT_PATH: &str = "python/task.py";
65const MAX_CBOR_FRAME_BYTES: usize = 16 * 1024 * 1024;
66const PYTHON_BOOTSTRAP: &str = include_str!("bootstrap.py");
67const PYTHON_COMMAND_CANDIDATES: &[&str] = &["python3", "python"];
68
69pub trait PyTaskState:
75 Default
76 + Clone
77 + core::fmt::Debug
78 + Serialize
79 + DeserializeOwned
80 + Reflect
81 + TypePath
82 + GetTypeRegistration
83 + cu29::bevy_reflect::Typed
84{
85}
86
87impl<T> PyTaskState for T where
88 T: Default
89 + Clone
90 + core::fmt::Debug
91 + Serialize
92 + DeserializeOwned
93 + Reflect
94 + TypePath
95 + GetTypeRegistration
96 + cu29::bevy_reflect::Typed
97{
98}
99
100pub trait PyInputSpec {
105 type Input<'m>: CuMsgPack
106 where
107 Self: 'm;
108 type Owned: Clone + Serialize + DeserializeOwned;
109
110 fn to_owned(input: &Self::Input<'_>) -> Self::Owned;
111}
112
113pub trait PyOutputSpec {
118 type Output<'m>: CuMsgPayload
119 where
120 Self: 'm;
121 type Owned: Clone + Serialize + DeserializeOwned;
122
123 fn to_owned(output: &Self::Output<'_>) -> Self::Owned;
124 fn replace_output(output: &mut Self::Output<'_>, owned: Self::Owned);
125}
126
127impl PyInputSpec for () {
128 type Input<'m>
129 = ()
130 where
131 Self: 'm;
132 type Owned = ();
133
134 fn to_owned(_input: &Self::Input<'_>) -> Self::Owned {}
135}
136
137impl<T> PyInputSpec for (T,)
138where
139 T: CuMsgPayload,
140{
141 type Input<'m>
142 = input_msg!(T)
143 where
144 Self: 'm;
145 type Owned = CuMsg<T>;
146
147 fn to_owned(input: &Self::Input<'_>) -> Self::Owned {
148 input.clone()
149 }
150}
151
152macro_rules! impl_py_input_spec_tuple {
153 ($first_ty:ident => $first_var:ident, $second_ty:ident => $second_var:ident $(, $ty:ident => $var:ident)* $(,)?) => {
154 impl<$first_ty, $second_ty $(, $ty)*> PyInputSpec for ($first_ty, $second_ty $(, $ty)*)
155 where
156 $first_ty: CuMsgPayload,
157 $second_ty: CuMsgPayload,
158 $($ty: CuMsgPayload),*
159 {
160 type Input<'m>
161 = input_msg!('m, $first_ty, $second_ty $(, $ty)*)
162 where
163 Self: 'm;
164 type Owned = (CuMsg<$first_ty>, CuMsg<$second_ty> $(, CuMsg<$ty>)*);
165
166 fn to_owned(input: &Self::Input<'_>) -> Self::Owned {
167 #[allow(non_snake_case)]
168 let ($first_var, $second_var $(, $var)*) = *input;
169 ($first_var.clone(), $second_var.clone() $(, $var.clone())*)
170 }
171 }
172 };
173}
174
175impl_py_input_spec_tuple!(T1 => v1, T2 => v2);
176impl_py_input_spec_tuple!(T1 => v1, T2 => v2, T3 => v3);
177impl_py_input_spec_tuple!(T1 => v1, T2 => v2, T3 => v3, T4 => v4);
178impl_py_input_spec_tuple!(T1 => v1, T2 => v2, T3 => v3, T4 => v4, T5 => v5);
179impl_py_input_spec_tuple!(T1 => v1, T2 => v2, T3 => v3, T4 => v4, T5 => v5, T6 => v6);
180impl_py_input_spec_tuple!(
181 T1 => v1,
182 T2 => v2,
183 T3 => v3,
184 T4 => v4,
185 T5 => v5,
186 T6 => v6,
187 T7 => v7
188);
189impl_py_input_spec_tuple!(
190 T1 => v1,
191 T2 => v2,
192 T3 => v3,
193 T4 => v4,
194 T5 => v5,
195 T6 => v6,
196 T7 => v7,
197 T8 => v8
198);
199impl_py_input_spec_tuple!(
200 T1 => v1,
201 T2 => v2,
202 T3 => v3,
203 T4 => v4,
204 T5 => v5,
205 T6 => v6,
206 T7 => v7,
207 T8 => v8,
208 T9 => v9
209);
210impl_py_input_spec_tuple!(
211 T1 => v1,
212 T2 => v2,
213 T3 => v3,
214 T4 => v4,
215 T5 => v5,
216 T6 => v6,
217 T7 => v7,
218 T8 => v8,
219 T9 => v9,
220 T10 => v10
221);
222impl_py_input_spec_tuple!(
223 T1 => v1,
224 T2 => v2,
225 T3 => v3,
226 T4 => v4,
227 T5 => v5,
228 T6 => v6,
229 T7 => v7,
230 T8 => v8,
231 T9 => v9,
232 T10 => v10,
233 T11 => v11
234);
235impl_py_input_spec_tuple!(
236 T1 => v1,
237 T2 => v2,
238 T3 => v3,
239 T4 => v4,
240 T5 => v5,
241 T6 => v6,
242 T7 => v7,
243 T8 => v8,
244 T9 => v9,
245 T10 => v10,
246 T11 => v11,
247 T12 => v12
248);
249
250pub type PyUnaryTask<In, State, Out> = PyTask<(In,), State, (Out,)>;
252
253impl PyOutputSpec for () {
254 type Output<'m>
255 = ()
256 where
257 Self: 'm;
258 type Owned = ();
259
260 fn to_owned(_output: &Self::Output<'_>) -> Self::Owned {}
261
262 fn replace_output(_output: &mut Self::Output<'_>, _owned: Self::Owned) {}
263}
264
265impl<T> PyOutputSpec for (T,)
266where
267 T: CuMsgPayload,
268{
269 type Output<'m>
270 = output_msg!(T)
271 where
272 Self: 'm;
273 type Owned = PyCuMsg<T>;
274
275 fn to_owned(output: &Self::Output<'_>) -> Self::Owned {
276 PyCuMsg::from_output(output)
277 }
278
279 fn replace_output(output: &mut Self::Output<'_>, owned: Self::Owned) {
280 *output = owned.into_output();
281 }
282}
283
284macro_rules! impl_py_output_spec_tuple {
285 ($first_ty:ident => $first_var:ident, $second_ty:ident => $second_var:ident $(, $ty:ident => $var:ident)* $(,)?) => {
286 impl<$first_ty, $second_ty $(, $ty)*> PyOutputSpec for ($first_ty, $second_ty $(, $ty)*)
287 where
288 $first_ty: CuMsgPayload,
289 $second_ty: CuMsgPayload,
290 $($ty: CuMsgPayload),*
291 {
292 type Output<'m>
293 = output_msg!($first_ty, $second_ty $(, $ty)*)
294 where
295 Self: 'm;
296 type Owned = (PyCuMsg<$first_ty>, PyCuMsg<$second_ty> $(, PyCuMsg<$ty>)*);
297
298 fn to_owned(output: &Self::Output<'_>) -> Self::Owned {
299 #[allow(non_snake_case)]
300 let ($first_var, $second_var $(, $var)*) = output;
301 (
302 PyCuMsg::from_output($first_var),
303 PyCuMsg::from_output($second_var)
304 $(, PyCuMsg::from_output($var))*
305 )
306 }
307
308 fn replace_output(output: &mut Self::Output<'_>, owned: Self::Owned) {
309 #[allow(non_snake_case)]
310 let ($first_var, $second_var $(, $var)*) = owned;
311 *output = (
312 $first_var.into_output(),
313 $second_var.into_output()
314 $(, $var.into_output())*
315 );
316 }
317 }
318 };
319}
320
321impl_py_output_spec_tuple!(T1 => v1, T2 => v2);
322impl_py_output_spec_tuple!(T1 => v1, T2 => v2, T3 => v3);
323impl_py_output_spec_tuple!(T1 => v1, T2 => v2, T3 => v3, T4 => v4);
324impl_py_output_spec_tuple!(T1 => v1, T2 => v2, T3 => v3, T4 => v4, T5 => v5);
325impl_py_output_spec_tuple!(T1 => v1, T2 => v2, T3 => v3, T4 => v4, T5 => v5, T6 => v6);
326impl_py_output_spec_tuple!(
327 T1 => v1,
328 T2 => v2,
329 T3 => v3,
330 T4 => v4,
331 T5 => v5,
332 T6 => v6,
333 T7 => v7
334);
335impl_py_output_spec_tuple!(
336 T1 => v1,
337 T2 => v2,
338 T3 => v3,
339 T4 => v4,
340 T5 => v5,
341 T6 => v6,
342 T7 => v7,
343 T8 => v8
344);
345impl_py_output_spec_tuple!(
346 T1 => v1,
347 T2 => v2,
348 T3 => v3,
349 T4 => v4,
350 T5 => v5,
351 T6 => v6,
352 T7 => v7,
353 T8 => v8,
354 T9 => v9
355);
356impl_py_output_spec_tuple!(
357 T1 => v1,
358 T2 => v2,
359 T3 => v3,
360 T4 => v4,
361 T5 => v5,
362 T6 => v6,
363 T7 => v7,
364 T8 => v8,
365 T9 => v9,
366 T10 => v10
367);
368impl_py_output_spec_tuple!(
369 T1 => v1,
370 T2 => v2,
371 T3 => v3,
372 T4 => v4,
373 T5 => v5,
374 T6 => v6,
375 T7 => v7,
376 T8 => v8,
377 T9 => v9,
378 T10 => v10,
379 T11 => v11
380);
381impl_py_output_spec_tuple!(
382 T1 => v1,
383 T2 => v2,
384 T3 => v3,
385 T4 => v4,
386 T5 => v5,
387 T6 => v6,
388 T7 => v7,
389 T8 => v8,
390 T9 => v9,
391 T10 => v10,
392 T11 => v11,
393 T12 => v12
394);
395
396#[derive(Debug, Clone, Copy, PartialEq, Eq, Reflect)]
397pub enum PyTaskMode {
398 Process,
405 Embedded,
411}
412
413impl PyTaskMode {
414 fn parse(value: &str) -> Option<Self> {
415 match value.trim().to_ascii_lowercase().as_str() {
416 "process" => Some(Self::Process),
417 "embedded" => Some(Self::Embedded),
418 _ => None,
419 }
420 }
421}
422
423#[derive(Serialize)]
424struct ProcessRequest<I, S, O> {
425 ctx: PyTaskContextSnapshot,
426 input: I,
427 state: S,
428 output: O,
429}
430
431impl<I, S, O> ProcessRequest<I, S, O> {
432 fn as_child_request(&self) -> ChildRequest<'_, I, S, O> {
433 ChildRequest::Process {
434 ctx: &self.ctx,
435 input: &self.input,
436 state: &self.state,
437 output: &self.output,
438 }
439 }
440}
441
442#[derive(Serialize)]
443struct StateRequest<S> {
444 ctx: PyTaskContextSnapshot,
445 state: S,
446}
447
448impl<S> StateRequest<S> {
449 fn as_child_request(&self, hook: StateHook) -> ChildRequest<'_, (), S, ()> {
450 match hook {
451 StateHook::Start => ChildRequest::Start {
452 ctx: &self.ctx,
453 state: &self.state,
454 },
455 StateHook::Stop => ChildRequest::Stop {
456 ctx: &self.ctx,
457 state: &self.state,
458 },
459 }
460 }
461}
462
463#[derive(Clone, Debug, Serialize)]
464struct PyTaskContextSnapshot {
465 now_ns: u64,
466 recent_ns: u64,
467 cl_id: u64,
468 task_id: Option<&'static str>,
469 task_index: Option<usize>,
470}
471
472impl PyTaskContextSnapshot {
473 fn from_cu_context(ctx: &CuContext) -> Self {
474 Self {
475 now_ns: ctx.now().as_nanos(),
476 recent_ns: ctx.recent().as_nanos(),
477 cl_id: ctx.cl_id(),
478 task_id: ctx.task_id(),
479 task_index: ctx.task_index(),
480 }
481 }
482}
483
484#[derive(Serialize, Deserialize)]
485struct ProcessResult<S, O> {
486 state: S,
487 output: O,
488}
489
490#[derive(Serialize, Deserialize)]
491struct StateResult<S> {
492 state: S,
493}
494
495#[derive(Debug, Clone, Serialize, Deserialize)]
496#[serde(bound(serialize = "T: Serialize", deserialize = "T: DeserializeOwned"))]
497#[doc(hidden)]
498pub struct PyCuMsg<T>
499where
500 T: CuMsgPayload,
501{
502 payload: Option<T>,
503 tov: Tov,
504 metadata: CuMsgMetadata,
505 #[serde(default, rename = "__cu_payload_present__")]
506 payload_present: bool,
507 #[serde(
508 default,
509 rename = "__cu_payload_template__",
510 skip_serializing_if = "Option::is_none"
511 )]
512 payload_template: Option<T>,
513}
514
515impl<T> PyCuMsg<T>
516where
517 T: CuMsgPayload,
518{
519 fn from_output(output: &CuMsg<T>) -> Self {
520 let payload = output.payload().cloned();
521 let payload_present = payload.is_some();
522 Self {
523 payload,
524 tov: output.tov,
525 metadata: output.metadata.clone(),
526 payload_present,
527 payload_template: (!payload_present).then(T::default),
528 }
529 }
530
531 fn into_output(self) -> CuMsg<T> {
532 let mut output = CuMsg::new(self.payload);
533 output.tov = self.tov;
534 output.metadata = self.metadata;
535 output
536 }
537}
538
539#[derive(Serialize)]
540#[serde(tag = "kind", rename_all = "snake_case")]
541enum ChildRequest<'a, I, S, O> {
542 Start {
543 ctx: &'a PyTaskContextSnapshot,
544 state: &'a S,
545 },
546 Stop {
547 ctx: &'a PyTaskContextSnapshot,
548 state: &'a S,
549 },
550 Process {
551 ctx: &'a PyTaskContextSnapshot,
552 input: &'a I,
553 state: &'a S,
554 output: &'a O,
555 },
556 Shutdown,
557}
558
559#[derive(Deserialize)]
560#[serde(tag = "kind", rename_all = "snake_case")]
561enum ChildResponse<S, O> {
562 Ready {
563 #[serde(default)]
564 cbor2_accelerated: bool,
565 },
566 State {
567 state: S,
568 },
569 Result {
570 state: S,
571 output: O,
572 },
573 Error {
574 message: String,
575 },
576}
577
578#[derive(Copy, Clone)]
579enum StateHook {
580 Start,
581 Stop,
582}
583
584#[derive(Reflect)]
621#[reflect(no_field_bounds, from_reflect = false, type_path = false)]
622pub struct PyTask<I, S, O>
623where
624 I: PyInputSpec + 'static,
625 S: PyTaskState + 'static,
626 O: PyOutputSpec + 'static,
627{
628 mode: PyTaskMode,
629 script: String,
630 state: S,
631 #[reflect(ignore)]
632 backend: Option<PythonBackend>,
633 #[reflect(ignore)]
634 marker: PhantomData<fn() -> (I, O)>,
635}
636
637impl<I, S, O> TypePath for PyTask<I, S, O>
638where
639 I: PyInputSpec + 'static,
640 S: PyTaskState + 'static,
641 O: PyOutputSpec + 'static,
642{
643 fn type_path() -> &'static str {
644 "cu_python_task::PyTask"
645 }
646
647 fn short_type_path() -> &'static str {
648 "PyTask"
649 }
650
651 fn type_ident() -> Option<&'static str> {
652 Some("PyTask")
653 }
654
655 fn crate_name() -> Option<&'static str> {
656 Some("cu_python_task")
657 }
658
659 fn module_path() -> Option<&'static str> {
660 Some("cu_python_task")
661 }
662}
663
664impl<I, S, O> Freezable for PyTask<I, S, O>
665where
666 I: PyInputSpec + 'static,
667 S: PyTaskState + 'static,
668 O: PyOutputSpec + 'static,
669{
670 fn freeze<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
671 let bytes = minicbor_serde::to_vec(&self.state)
672 .map_err(|e| EncodeError::OtherString(e.to_string()))?;
673 Encode::encode(&bytes, encoder)
674 }
675
676 fn thaw<D: Decoder>(&mut self, decoder: &mut D) -> Result<(), DecodeError> {
677 let bytes = <Vec<u8> as Decode<D::Context>>::decode(decoder)?;
678 self.state = minicbor_serde::from_slice(&bytes)
679 .map_err(|e| DecodeError::OtherString(e.to_string()))?;
680 Ok(())
681 }
682}
683
684impl<I, S, O> CuTask for PyTask<I, S, O>
685where
686 I: PyInputSpec + 'static,
687 S: PyTaskState + 'static,
688 O: PyOutputSpec + 'static,
689{
690 type Resources<'r> = ();
691 type Input<'m> = I::Input<'m>;
692 type Output<'m> = O::Output<'m>;
693
694 fn new(config: Option<&ComponentConfig>, _resources: Self::Resources<'_>) -> CuResult<Self> {
695 let mode = parse_mode(config)?;
696 let script = resolve_script_path(config)?;
697 if !script.is_file() {
698 return Err(CuError::from(format!(
699 "Python task script '{}' does not exist",
700 script.display()
701 )));
702 }
703 Ok(Self {
704 mode,
705 script: script.display().to_string(),
706 state: S::default(),
707 backend: None,
708 marker: PhantomData,
709 })
710 }
711
712 fn start(&mut self, ctx: &CuContext) -> CuResult<()> {
713 if self.backend.is_none() {
714 let mut backend = PythonBackend::launch(self.mode, Path::new(&self.script))?;
715 backend.start_hook(ctx, &mut self.state)?;
716 self.backend = Some(backend);
717 }
718 Ok(())
719 }
720
721 fn process<'i, 'o>(
722 &mut self,
723 ctx: &CuContext,
724 input: &Self::Input<'i>,
725 output: &mut Self::Output<'o>,
726 ) -> CuResult<()> {
727 if self.backend.is_none() {
728 self.start(ctx)?;
729 }
730
731 let request = ProcessRequest {
732 ctx: PyTaskContextSnapshot::from_cu_context(ctx),
733 input: I::to_owned(input),
734 state: self.state.clone(),
735 output: O::to_owned(output),
736 };
737
738 let backend = self
739 .backend
740 .as_mut()
741 .expect("backend initialized immediately above");
742 let ProcessResult {
743 state,
744 output: new_output,
745 } = backend.process(ctx, &request)?;
746 self.state = state;
747 O::replace_output(output, new_output);
748 Ok(())
749 }
750
751 fn stop(&mut self, ctx: &CuContext) -> CuResult<()> {
752 if let Some(mut backend) = self.backend.take() {
753 backend.stop(ctx, &mut self.state)?;
754 }
755 Ok(())
756 }
757}
758
759fn parse_mode(config: Option<&ComponentConfig>) -> CuResult<PyTaskMode> {
760 let Some(config) = config else {
761 return Ok(PyTaskMode::Process);
762 };
763 let Some(raw) = config.get::<String>("mode")? else {
764 return Ok(PyTaskMode::Process);
765 };
766 PyTaskMode::parse(&raw).ok_or_else(|| {
767 CuError::from(format!(
768 "Unsupported Python task mode '{raw}', expected 'process' or 'embedded'"
769 ))
770 })
771}
772
773fn resolve_script_path(config: Option<&ComponentConfig>) -> CuResult<PathBuf> {
774 let raw = if let Some(config) = config {
775 config.get::<String>("script")?
776 } else {
777 None
778 };
779
780 let path = raw
781 .map(PathBuf::from)
782 .unwrap_or_else(|| PathBuf::from(DEFAULT_SCRIPT_PATH));
783 absolutize_path(&path)
784}
785
786fn absolutize_path(path: &Path) -> CuResult<PathBuf> {
787 if path.is_absolute() {
788 return Ok(path.to_path_buf());
789 }
790 let cwd = std::env::current_dir()
791 .map_err(|e| CuError::new_with_cause("Failed to read current working directory", e))?;
792 Ok(cwd.join(path))
793}
794
795enum PythonBackend {
796 Process(ProcessBackend),
797 Embedded(EmbeddedBackend),
798}
799
800impl PythonBackend {
801 fn launch(mode: PyTaskMode, script: &Path) -> CuResult<Self> {
802 match mode {
803 PyTaskMode::Process => Ok(Self::Process(ProcessBackend::start(script)?)),
804 PyTaskMode::Embedded => Ok(Self::Embedded(EmbeddedBackend::start(script)?)),
805 }
806 }
807
808 fn start_hook<S>(&mut self, ctx: &CuContext, state: &mut S) -> CuResult<()>
809 where
810 S: Serialize + DeserializeOwned + Clone,
811 {
812 match self {
813 Self::Process(backend) => backend.start_hook(ctx, state),
814 Self::Embedded(backend) => backend.start_hook(ctx, state),
815 }
816 }
817
818 fn process<I, S, O>(
819 &mut self,
820 ctx: &CuContext,
821 request: &ProcessRequest<I, S, O>,
822 ) -> CuResult<ProcessResult<S, O>>
823 where
824 I: Serialize,
825 S: Serialize + DeserializeOwned,
826 O: Serialize + DeserializeOwned,
827 {
828 match self {
829 Self::Process(backend) => backend.process(ctx, request),
830 Self::Embedded(backend) => backend.process(ctx, request),
831 }
832 }
833
834 fn stop<S>(&mut self, ctx: &CuContext, state: &mut S) -> CuResult<()>
835 where
836 S: Serialize + DeserializeOwned + Clone,
837 {
838 match self {
839 Self::Process(backend) => backend.stop(ctx, state),
840 Self::Embedded(backend) => backend.stop(ctx, state),
841 }
842 }
843}
844
845struct ProcessBackend {
846 child: Child,
847 stdin: ChildStdin,
848 stdout: BufReader<ChildStdout>,
849}
850
851impl ProcessBackend {
852 fn start(script: &Path) -> CuResult<Self> {
853 let python = python_command()?;
854 let mut child = Command::new(python)
855 .arg("-u")
856 .arg("-c")
857 .arg(PYTHON_BOOTSTRAP)
858 .arg(script)
859 .stdin(Stdio::piped())
860 .stdout(Stdio::piped())
861 .stderr(Stdio::inherit())
862 .spawn()
863 .map_err(|e| CuError::new_with_cause("Failed to spawn Python task process", e))?;
864
865 let stdin = child
866 .stdin
867 .take()
868 .ok_or_else(|| CuError::from("Python task process did not expose stdin"))?;
869 let stdout = child
870 .stdout
871 .take()
872 .ok_or_else(|| CuError::from("Python task process did not expose stdout"))?;
873
874 let mut backend = Self {
875 child,
876 stdin,
877 stdout: BufReader::new(stdout),
878 };
879 backend.wait_ready()?;
880 Ok(backend)
881 }
882
883 fn wait_ready(&mut self) -> CuResult<()> {
884 match read_cbor_frame::<_, ChildResponse<(), ()>>(&mut self.stdout) {
885 Ok(ChildResponse::Ready { cbor2_accelerated }) => {
886 if !cbor2_accelerated {
887 warning!(
888 "cu_python_task process mode is using pure-Python cbor2; install the C extension backend for better performance"
889 );
890 }
891 Ok(())
892 }
893 Ok(ChildResponse::Error { message }) => Err(CuError::from(format!(
894 "Python task process failed during startup:\n{message}"
895 ))),
896 Ok(ChildResponse::Result { .. }) => Err(CuError::from(
897 "Unexpected result frame while waiting for Python task process startup",
898 )),
899 Ok(ChildResponse::State { .. }) => Err(CuError::from(
900 "Unexpected state-only frame while waiting for Python task process startup",
901 )),
902 Err(read_error) => {
903 if let Some(status) = self
904 .child
905 .try_wait()
906 .map_err(|e| CuError::new_with_cause("Failed to poll Python task process", e))?
907 {
908 Err(CuError::from(format!(
909 "Python task process exited during startup with status {status}"
910 )))
911 } else {
912 Err(read_error)
913 }
914 }
915 }
916 }
917
918 fn process<I, S, O>(
919 &mut self,
920 _ctx: &CuContext,
921 request: &ProcessRequest<I, S, O>,
922 ) -> CuResult<ProcessResult<S, O>>
923 where
924 I: Serialize,
925 S: Serialize + DeserializeOwned,
926 O: Serialize + DeserializeOwned,
927 {
928 let child_request = request.as_child_request();
929 write_cbor_frame(&mut self.stdin, &child_request)?;
930 match read_cbor_frame::<_, ChildResponse<S, O>>(&mut self.stdout)? {
931 ChildResponse::Result { state, output } => Ok(ProcessResult { state, output }),
932 ChildResponse::State { .. } => Err(CuError::from(
933 "Python task process unexpectedly sent a state-only response while processing",
934 )),
935 ChildResponse::Error { message } => Err(CuError::from(format!(
936 "Python task process raised an exception:\n{message}"
937 ))),
938 ChildResponse::Ready { .. } => Err(CuError::from(
939 "Python task process unexpectedly sent a second ready frame",
940 )),
941 }
942 }
943
944 fn start_hook<S>(&mut self, ctx: &CuContext, state: &mut S) -> CuResult<()>
945 where
946 S: Serialize + DeserializeOwned + Clone,
947 {
948 self.call_state_hook(StateHook::Start, ctx, state)
949 }
950
951 fn call_state_hook<S>(
952 &mut self,
953 hook: StateHook,
954 ctx: &CuContext,
955 state: &mut S,
956 ) -> CuResult<()>
957 where
958 S: Serialize + DeserializeOwned + Clone,
959 {
960 let request = StateRequest {
961 ctx: PyTaskContextSnapshot::from_cu_context(ctx),
962 state: state.clone(),
963 };
964 let child_request = request.as_child_request(hook);
965 write_cbor_frame(&mut self.stdin, &child_request)?;
966 match read_cbor_frame::<_, ChildResponse<S, ()>>(&mut self.stdout)? {
967 ChildResponse::State { state: new_state } => {
968 *state = new_state;
969 Ok(())
970 }
971 ChildResponse::Result { .. } => Err(CuError::from(
972 "Python task process unexpectedly sent a process response for a state-only hook",
973 )),
974 ChildResponse::Error { message } => Err(CuError::from(format!(
975 "Python task process raised an exception:\n{message}"
976 ))),
977 ChildResponse::Ready { .. } => Err(CuError::from(
978 "Python task process unexpectedly sent a ready frame for a state-only hook",
979 )),
980 }
981 }
982
983 fn stop<S>(&mut self, ctx: &CuContext, state: &mut S) -> CuResult<()>
984 where
985 S: Serialize + DeserializeOwned + Clone,
986 {
987 self.call_state_hook(StateHook::Stop, ctx, state)?;
988 let _ = write_cbor_frame(&mut self.stdin, &ChildRequest::<(), (), ()>::Shutdown);
989 self.child
990 .wait()
991 .map_err(|e| CuError::new_with_cause("Failed to wait for Python task process", e))?;
992 Ok(())
993 }
994}
995
996#[pyclass(name = "CuContext")]
997struct PyEmbeddedContext {
998 ctx: CuContext,
999}
1000
1001#[pymethods]
1002impl PyEmbeddedContext {
1003 #[getter]
1004 fn cl_id(&self) -> u64 {
1005 self.ctx.cl_id()
1006 }
1007
1008 #[getter]
1009 fn task_id(&self) -> Option<String> {
1010 self.ctx.task_id().map(str::to_string)
1011 }
1012
1013 #[getter]
1014 fn task_index(&self) -> Option<usize> {
1015 self.ctx.task_index()
1016 }
1017
1018 #[getter]
1019 fn now_ns(&self) -> u64 {
1020 self.ctx.now().as_nanos()
1021 }
1022
1023 #[getter]
1024 fn recent_ns(&self) -> u64 {
1025 self.ctx.recent().as_nanos()
1026 }
1027
1028 fn now(&self) -> u64 {
1029 self.ctx.now().as_nanos()
1030 }
1031
1032 fn recent(&self) -> u64 {
1033 self.ctx.recent().as_nanos()
1034 }
1035}
1036
1037struct EmbeddedBackend {
1038 call_process: Py<PyAny>,
1039 call_state_hook: Py<PyAny>,
1040 process_fn: Py<PyAny>,
1041 start_fn: Option<Py<PyAny>>,
1042 stop_fn: Option<Py<PyAny>>,
1043}
1044
1045impl EmbeddedBackend {
1046 fn start(script: &Path) -> CuResult<Self> {
1047 Python::initialize();
1048 Python::attach(|py| {
1049 let code = CString::new(PYTHON_BOOTSTRAP).map_err(|e| {
1050 CuError::new_with_cause("Invalid embedded Python bootstrap code", e)
1051 })?;
1052 let filename = CString::new("cu_python_task_bootstrap.py")
1053 .map_err(|e| CuError::new_with_cause("Invalid bootstrap filename", e))?;
1054 let module_name = CString::new("cu_python_task_bootstrap")
1055 .map_err(|e| CuError::new_with_cause("Invalid bootstrap module name", e))?;
1056
1057 let module = PyModule::from_code(
1058 py,
1059 code.as_c_str(),
1060 filename.as_c_str(),
1061 module_name.as_c_str(),
1062 )
1063 .map_err(python_error)?;
1064 let load_task_functions = module
1065 .getattr("load_task_functions")
1066 .map_err(python_error)?;
1067 let call_process = module.getattr("call_process").map_err(python_error)?;
1068 let call_state_hook = module
1069 .getattr("call_optional_state_hook")
1070 .map_err(python_error)?;
1071 let functions = load_task_functions
1072 .call1((script.display().to_string(),))
1073 .map_err(python_error)?;
1074 let functions = functions
1075 .cast::<PyTuple>()
1076 .map_err(|e| CuError::from(format!("Embedded Python task error: {e}")))?;
1077 let process_fn = functions.get_item(0).map_err(python_error)?;
1078 let start_fn = functions.get_item(1).map_err(python_error)?;
1079 let stop_fn = functions.get_item(2).map_err(python_error)?;
1080 Ok(Self {
1081 call_process: call_process.unbind(),
1082 call_state_hook: call_state_hook.unbind(),
1083 process_fn: process_fn.unbind(),
1084 start_fn: (!start_fn.is_none()).then(|| start_fn.unbind()),
1085 stop_fn: (!stop_fn.is_none()).then(|| stop_fn.unbind()),
1086 })
1087 })
1088 }
1089
1090 fn start_hook<S>(&self, ctx: &CuContext, state: &mut S) -> CuResult<()>
1091 where
1092 S: Serialize + DeserializeOwned + Clone,
1093 {
1094 self.call_state_hook(self.start_fn.as_ref(), ctx, state)
1095 }
1096
1097 fn process<I, S, O>(
1098 &mut self,
1099 ctx: &CuContext,
1100 request: &ProcessRequest<I, S, O>,
1101 ) -> CuResult<ProcessResult<S, O>>
1102 where
1103 I: Serialize,
1104 S: Serialize + DeserializeOwned,
1105 O: Serialize + DeserializeOwned,
1106 {
1107 let request_value = to_value(request)
1108 .map_err(|e| CuError::new_with_cause("Failed to encode embedded Python request", e))?;
1109 let response_value = Python::attach(|py| {
1110 let call_process = self.call_process.bind(py);
1111 let process_fn = self.process_fn.bind(py);
1112 let request_py = value_to_py(&request_value, py).map_err(python_error)?;
1113 let live_ctx =
1114 Py::new(py, PyEmbeddedContext { ctx: ctx.clone() }).map_err(python_error)?;
1115 call_process
1116 .call1((process_fn, request_py, live_ctx))
1117 .map_err(python_error)
1118 .and_then(|response| py_to_value(&response).map_err(python_error))
1119 })?;
1120 response_value
1121 .deserialize_into::<ProcessResult<S, O>>()
1122 .map_err(|e| CuError::new_with_cause("Failed to decode embedded Python response", e))
1123 }
1124
1125 fn call_state_hook<S>(
1126 &self,
1127 hook_fn: Option<&Py<PyAny>>,
1128 ctx: &CuContext,
1129 state: &mut S,
1130 ) -> CuResult<()>
1131 where
1132 S: Serialize + DeserializeOwned + Clone,
1133 {
1134 let request = StateRequest {
1135 ctx: PyTaskContextSnapshot::from_cu_context(ctx),
1136 state: state.clone(),
1137 };
1138 let request_value = to_value(&request).map_err(|e| {
1139 CuError::new_with_cause("Failed to encode embedded Python hook request", e)
1140 })?;
1141 let response_value = Python::attach(|py| {
1142 let call_state_hook = self.call_state_hook.bind(py);
1143 let request_py = value_to_py(&request_value, py).map_err(python_error)?;
1144 let live_ctx =
1145 Py::new(py, PyEmbeddedContext { ctx: ctx.clone() }).map_err(python_error)?;
1146 let hook_py = hook_fn.map_or_else(|| py.None(), |hook_fn| hook_fn.clone_ref(py));
1147 call_state_hook
1148 .call1((hook_py, request_py, live_ctx))
1149 .map_err(python_error)
1150 .and_then(|response| py_to_value(&response).map_err(python_error))
1151 })?;
1152 let StateResult { state: new_state } = response_value
1153 .deserialize_into::<StateResult<S>>()
1154 .map_err(|e| {
1155 CuError::new_with_cause("Failed to decode embedded Python hook response", e)
1156 })?;
1157 *state = new_state;
1158 Ok(())
1159 }
1160
1161 fn stop<S>(&self, ctx: &CuContext, state: &mut S) -> CuResult<()>
1162 where
1163 S: Serialize + DeserializeOwned + Clone,
1164 {
1165 self.call_state_hook(self.stop_fn.as_ref(), ctx, state)?;
1166 Ok(())
1167 }
1168}
1169
1170fn python_error(error: PyErr) -> CuError {
1171 CuError::from(format!("Embedded Python task error: {error}"))
1172}
1173
1174fn python_command() -> CuResult<&'static str> {
1175 static PYTHON_COMMAND: OnceLock<Option<&'static str>> = OnceLock::new();
1176 match PYTHON_COMMAND.get_or_init(detect_python_command) {
1177 Some(command) => Ok(*command),
1178 None => Err(CuError::from(
1179 "Could not find a usable Python interpreter (`python3` or `python`) in PATH",
1180 )),
1181 }
1182}
1183
1184fn detect_python_command() -> Option<&'static str> {
1185 PYTHON_COMMAND_CANDIDATES.iter().copied().find(|command| {
1186 Command::new(command)
1187 .arg("--version")
1188 .stdin(Stdio::null())
1189 .stdout(Stdio::null())
1190 .stderr(Stdio::null())
1191 .status()
1192 .map(|status| status.success())
1193 .unwrap_or(false)
1194 })
1195}
1196
1197fn write_cbor_frame<W, T>(writer: &mut W, value: &T) -> CuResult<()>
1198where
1199 W: Write,
1200 T: Serialize,
1201{
1202 let _guard = enable_shared_handle_serialization();
1203 let value = to_value(value)
1204 .map_err(|e| CuError::new_with_cause("Failed to encode Python task CBOR value", e))?;
1205 let payload = encode_wire_value(&value)?;
1206 if payload.len() > MAX_CBOR_FRAME_BYTES {
1207 return Err(CuError::from(format!(
1208 "Python task frame exceeded {} bytes",
1209 MAX_CBOR_FRAME_BYTES
1210 )));
1211 }
1212 let len = u32::try_from(payload.len())
1213 .map_err(|_| CuError::from("Python task frame length exceeded u32"))?;
1214 writer
1215 .write_all(&len.to_le_bytes())
1216 .map_err(|e| CuError::new_with_cause("Failed to write Python task frame length", e))?;
1217 writer
1218 .write_all(&payload)
1219 .map_err(|e| CuError::new_with_cause("Failed to write Python task frame payload", e))?;
1220 writer
1221 .flush()
1222 .map_err(|e| CuError::new_with_cause("Failed to flush Python task frame", e))
1223}
1224
1225fn read_cbor_frame<R, T>(reader: &mut R) -> CuResult<T>
1226where
1227 R: Read,
1228 T: DeserializeOwned,
1229{
1230 let mut len_bytes = [0_u8; 4];
1231 reader
1232 .read_exact(&mut len_bytes)
1233 .map_err(|e| CuError::new_with_cause("Failed to read Python task frame length", e))?;
1234 let len = u32::from_le_bytes(len_bytes) as usize;
1235 if len > MAX_CBOR_FRAME_BYTES {
1236 return Err(CuError::from(format!(
1237 "Python task frame exceeded {} bytes",
1238 MAX_CBOR_FRAME_BYTES
1239 )));
1240 }
1241
1242 let mut payload = vec![0_u8; len];
1243 reader
1244 .read_exact(&mut payload)
1245 .map_err(|e| CuError::new_with_cause("Failed to read Python task frame payload", e))?;
1246 let value = decode_wire_value(&payload)?;
1247 value
1248 .deserialize_into::<T>()
1249 .map_err(|e| CuError::new_with_cause("Failed to decode Python task CBOR frame", e))
1250}
1251
1252fn encode_wire_value(value: &CuValue) -> CuResult<Vec<u8>> {
1253 let mut encoder = CborEncoder::new(Vec::new());
1254 encode_wire_value_into(&mut encoder, value)
1255 .map_err(|e| CuError::new_with_cause("Failed to serialize Python task CBOR frame", e))?;
1256 Ok(encoder.into_writer())
1257}
1258
1259fn encode_wire_value_into<W>(
1260 encoder: &mut CborEncoder<W>,
1261 value: &CuValue,
1262) -> Result<(), minicbor::encode::Error<W::Error>>
1263where
1264 W: minicbor::encode::Write,
1265{
1266 match value {
1267 CuValue::Bool(v) => {
1268 encoder.bool(*v)?;
1269 }
1270 CuValue::U8(v) => {
1271 encoder.u8(*v)?;
1272 }
1273 CuValue::U16(v) => {
1274 encoder.u16(*v)?;
1275 }
1276 CuValue::U32(v) => {
1277 encoder.u32(*v)?;
1278 }
1279 CuValue::U64(v) => {
1280 encoder.u64(*v)?;
1281 }
1282 CuValue::U128(v) => {
1283 if let Ok(v) = u64::try_from(*v) {
1284 encoder.u64(v)?;
1285 } else {
1286 encode_bignum(encoder, IanaTag::PosBignum, *v)?;
1287 }
1288 }
1289 CuValue::I8(v) => {
1290 encoder.i8(*v)?;
1291 }
1292 CuValue::I16(v) => {
1293 encoder.i16(*v)?;
1294 }
1295 CuValue::I32(v) => {
1296 encoder.i32(*v)?;
1297 }
1298 CuValue::I64(v) => {
1299 encoder.i64(*v)?;
1300 }
1301 CuValue::I128(v) => {
1302 if let Ok(v) = CborInt::try_from(*v) {
1303 encoder.int(v)?;
1304 } else {
1305 encode_bignum(encoder, IanaTag::NegBignum, (-1_i128 - *v) as u128)?;
1306 }
1307 }
1308 CuValue::F32(v) => {
1309 encoder.f32(*v)?;
1310 }
1311 CuValue::F64(v) => {
1312 encoder.f64(*v)?;
1313 }
1314 CuValue::Char(v) => {
1315 encoder.char(*v)?;
1316 }
1317 CuValue::String(v) => {
1318 encoder.str(v)?;
1319 }
1320 CuValue::Unit | CuValue::Option(None) => {
1321 encoder.null()?;
1322 }
1323 CuValue::Option(Some(v)) | CuValue::Newtype(v) => {
1324 encode_wire_value_into(encoder, v)?;
1325 }
1326 CuValue::Seq(values) => {
1327 encoder.array(values.len() as u64)?;
1328 for value in values {
1329 encode_wire_value_into(encoder, value)?;
1330 }
1331 }
1332 CuValue::Map(values) => {
1333 encoder.map(values.len() as u64)?;
1334 for (key, value) in values {
1335 encode_wire_value_into(encoder, key)?;
1336 encode_wire_value_into(encoder, value)?;
1337 }
1338 }
1339 CuValue::Bytes(value) => {
1340 encoder.bytes(value)?;
1341 }
1342 CuValue::CuTime(value) => {
1343 encoder.u64(value.0)?;
1344 }
1345 }
1346 Ok(())
1347}
1348
1349fn encode_bignum<W>(
1350 encoder: &mut CborEncoder<W>,
1351 tag: IanaTag,
1352 value: u128,
1353) -> Result<(), minicbor::encode::Error<W::Error>>
1354where
1355 W: minicbor::encode::Write,
1356{
1357 let bytes = value.to_be_bytes();
1358 let start = bytes
1359 .iter()
1360 .position(|byte| *byte != 0)
1361 .unwrap_or(bytes.len() - 1);
1362 encoder.tag(tag)?.bytes(&bytes[start..])?;
1363 Ok(())
1364}
1365
1366fn decode_wire_value(payload: &[u8]) -> CuResult<CuValue> {
1367 let mut decoder = CborDecoder::new(payload);
1368 let value = decode_wire_value_from(&mut decoder)
1369 .map_err(|e| CuError::new_with_cause("Failed to decode Python task CBOR frame", e))?;
1370 if decoder.position() != payload.len() {
1371 return Err(CuError::from(
1372 "Failed to decode Python task CBOR frame: trailing data",
1373 ));
1374 }
1375 Ok(canonicalize_wire_value(value))
1376}
1377
1378fn decode_wire_value_from(
1379 decoder: &mut CborDecoder<'_>,
1380) -> Result<CuValue, minicbor::decode::Error> {
1381 match decoder.datatype()? {
1382 CborType::Bool => Ok(CuValue::Bool(decoder.bool()?)),
1383 CborType::Null => {
1384 decoder.null()?;
1385 Ok(CuValue::Option(None))
1386 }
1387 CborType::U8 => Ok(CuValue::U8(decoder.u8()?)),
1388 CborType::U16 => Ok(CuValue::U16(decoder.u16()?)),
1389 CborType::U32 => Ok(CuValue::U32(decoder.u32()?)),
1390 CborType::U64 => Ok(CuValue::U64(decoder.u64()?)),
1391 CborType::I8 => Ok(CuValue::I8(decoder.i8()?)),
1392 CborType::I16 => Ok(CuValue::I16(decoder.i16()?)),
1393 CborType::I32 => Ok(CuValue::I32(decoder.i32()?)),
1394 CborType::I64 => Ok(CuValue::I64(decoder.i64()?)),
1395 CborType::Int => decode_cbor_int(decoder.int()?),
1396 CborType::F32 => Ok(CuValue::F32(decoder.f32()?)),
1397 CborType::F64 => Ok(CuValue::F64(decoder.f64()?)),
1398 CborType::Bytes | CborType::BytesIndef => Ok(CuValue::Bytes(read_byte_string(decoder)?)),
1399 CborType::String | CborType::StringIndef => Ok(CuValue::String(read_text_string(decoder)?)),
1400 CborType::Array | CborType::ArrayIndef => decode_array(decoder),
1401 CborType::Map | CborType::MapIndef => decode_map(decoder),
1402 CborType::Tag => decode_tagged_value(decoder),
1403 other => Err(minicbor::decode::Error::message(format!(
1404 "unsupported Python task CBOR type {other}"
1405 ))),
1406 }
1407}
1408
1409fn decode_cbor_int(value: CborInt) -> Result<CuValue, minicbor::decode::Error> {
1410 if let Ok(value) = u64::try_from(value) {
1411 Ok(CuValue::U64(value))
1412 } else if let Ok(value) = i64::try_from(value) {
1413 Ok(CuValue::I64(value))
1414 } else {
1415 Ok(CuValue::I128(i128::from(value)))
1416 }
1417}
1418
1419fn decode_array(decoder: &mut CborDecoder<'_>) -> Result<CuValue, minicbor::decode::Error> {
1420 let len = decoder.array()?;
1421 let mut values = Vec::new();
1422 match len {
1423 Some(len) => {
1424 values.reserve(len as usize);
1425 for _ in 0..len {
1426 values.push(decode_wire_value_from(decoder)?);
1427 }
1428 }
1429 None => loop {
1430 if decoder.datatype()? == CborType::Break {
1431 decoder.skip()?;
1432 break;
1433 }
1434 values.push(decode_wire_value_from(decoder)?);
1435 },
1436 }
1437 Ok(CuValue::Seq(values))
1438}
1439
1440fn decode_map(decoder: &mut CborDecoder<'_>) -> Result<CuValue, minicbor::decode::Error> {
1441 let len = decoder.map()?;
1442 let mut values = BTreeMap::new();
1443 match len {
1444 Some(len) => {
1445 for _ in 0..len {
1446 let key = decode_wire_value_from(decoder)?;
1447 let value = decode_wire_value_from(decoder)?;
1448 values.insert(key, value);
1449 }
1450 }
1451 None => loop {
1452 if decoder.datatype()? == CborType::Break {
1453 decoder.skip()?;
1454 break;
1455 }
1456 let key = decode_wire_value_from(decoder)?;
1457 let value = decode_wire_value_from(decoder)?;
1458 values.insert(key, value);
1459 },
1460 }
1461 Ok(CuValue::Map(values))
1462}
1463
1464fn decode_tagged_value(decoder: &mut CborDecoder<'_>) -> Result<CuValue, minicbor::decode::Error> {
1465 match IanaTag::try_from(decoder.tag()?) {
1466 Ok(IanaTag::PosBignum) => {
1467 let value = decode_bignum(decoder)?;
1468 if let Ok(value) = u64::try_from(value) {
1469 Ok(CuValue::U64(value))
1470 } else {
1471 Ok(CuValue::U128(value))
1472 }
1473 }
1474 Ok(IanaTag::NegBignum) => {
1475 let value = decode_bignum(decoder)?;
1476 if value <= i64::MAX as u128 {
1477 Ok(CuValue::I64(-1 - value as i64))
1478 } else if let Ok(value) = i128::try_from(value) {
1479 Ok(CuValue::I128(-1 - value))
1480 } else {
1481 Err(minicbor::decode::Error::message(
1482 "Python task CBOR negative bignum exceeded i128",
1483 ))
1484 }
1485 }
1486 Ok(other) => Err(minicbor::decode::Error::message(format!(
1487 "unsupported Python task CBOR tag {}",
1488 u64::from(other)
1489 ))),
1490 Err(tag) => Err(minicbor::decode::Error::message(format!(
1491 "unsupported Python task CBOR tag {tag}",
1492 ))),
1493 }
1494}
1495
1496fn decode_bignum(decoder: &mut CborDecoder<'_>) -> Result<u128, minicbor::decode::Error> {
1497 let bytes = read_byte_string(decoder)?;
1498 let first = bytes
1499 .iter()
1500 .position(|byte| *byte != 0)
1501 .unwrap_or(bytes.len());
1502 let bytes = &bytes[first..];
1503 if bytes.len() > 16 {
1504 return Err(minicbor::decode::Error::message(
1505 "Python task CBOR bignum exceeded u128",
1506 ));
1507 }
1508 let mut buffer = [0_u8; 16];
1509 buffer[16 - bytes.len()..].copy_from_slice(bytes);
1510 Ok(u128::from_be_bytes(buffer))
1511}
1512
1513fn read_byte_string(decoder: &mut CborDecoder<'_>) -> Result<Vec<u8>, minicbor::decode::Error> {
1514 match decoder.datatype()? {
1515 CborType::Bytes => Ok(decoder.bytes()?.to_vec()),
1516 CborType::BytesIndef => {
1517 let mut bytes = Vec::new();
1518 for chunk in decoder.bytes_iter()? {
1519 bytes.extend_from_slice(chunk?);
1520 }
1521 Ok(bytes)
1522 }
1523 other => Err(minicbor::decode::Error::message(format!(
1524 "expected byte string, found {other}"
1525 ))),
1526 }
1527}
1528
1529fn read_text_string(decoder: &mut CborDecoder<'_>) -> Result<String, minicbor::decode::Error> {
1530 match decoder.datatype()? {
1531 CborType::String => Ok(decoder.str()?.to_owned()),
1532 CborType::StringIndef => {
1533 let mut text = String::new();
1534 for chunk in decoder.str_iter()? {
1535 text.push_str(chunk?);
1536 }
1537 Ok(text)
1538 }
1539 other => Err(minicbor::decode::Error::message(format!(
1540 "expected text string, found {other}"
1541 ))),
1542 }
1543}
1544
1545fn canonicalize_wire_value(value: CuValue) -> CuValue {
1546 match value {
1547 CuValue::Option(None) => CuValue::Unit,
1548 CuValue::Option(Some(value)) | CuValue::Newtype(value) => canonicalize_wire_value(*value),
1549 CuValue::Seq(values) => CuValue::Seq(
1550 values
1551 .into_iter()
1552 .map(canonicalize_wire_value)
1553 .collect::<Vec<_>>(),
1554 ),
1555 CuValue::Map(values) => CuValue::Map(
1556 values
1557 .into_iter()
1558 .map(|(key, value)| (canonicalize_wire_value(key), canonicalize_wire_value(value)))
1559 .collect::<BTreeMap<_, _>>(),
1560 ),
1561 other => other,
1562 }
1563}
1564
1565#[cfg(test)]
1566mod tests {
1567 use super::*;
1568 use cu29::prelude::{CuContext, RobotClock};
1569 use std::sync::{Mutex, MutexGuard, OnceLock};
1570 use tempfile::TempDir;
1571
1572 #[derive(
1573 Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Encode, Decode, Reflect,
1574 )]
1575 struct TestPayload {
1576 value: i32,
1577 flag: bool,
1578 }
1579
1580 #[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1581 struct SeenState {
1582 seen: bool,
1583 }
1584
1585 #[derive(Debug, Clone, Serialize, Deserialize)]
1586 struct SharedHandleInput {
1587 data: CuHandle<CuSharedMemoryBuffer<u8>>,
1588 }
1589
1590 const TEST_TASK_IDS: &[&str] = &["py"];
1591
1592 fn cwd_lock() -> MutexGuard<'static, ()> {
1593 static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
1594 LOCK.get_or_init(|| Mutex::new(()))
1595 .lock()
1596 .expect("cwd lock poisoned")
1597 }
1598
1599 struct CurrentDirGuard {
1600 original: PathBuf,
1601 }
1602
1603 impl CurrentDirGuard {
1604 fn switch_to(path: &Path) -> Self {
1605 let original = std::env::current_dir().expect("cwd");
1606 std::env::set_current_dir(path).expect("switch cwd");
1607 Self { original }
1608 }
1609 }
1610
1611 impl Drop for CurrentDirGuard {
1612 fn drop(&mut self) {
1613 std::env::set_current_dir(&self.original).expect("restore cwd");
1614 }
1615 }
1616
1617 #[test]
1618 fn mode_defaults_to_process() {
1619 assert_eq!(parse_mode(None).unwrap(), PyTaskMode::Process);
1620 }
1621
1622 #[test]
1623 fn mode_parses_embedded() {
1624 let mut cfg = ComponentConfig::new();
1625 cfg.set("mode", "embedded".to_string());
1626 assert_eq!(parse_mode(Some(&cfg)).unwrap(), PyTaskMode::Embedded);
1627 }
1628
1629 #[test]
1630 fn invalid_mode_is_rejected() {
1631 let mut cfg = ComponentConfig::new();
1632 cfg.set("mode", "nope".to_string());
1633 let err = parse_mode(Some(&cfg)).expect_err("mode should fail");
1634 assert!(err.to_string().contains("Unsupported Python task mode"));
1635 }
1636
1637 #[test]
1638 fn default_script_path_uses_python_task_py_from_cwd() {
1639 let _lock = cwd_lock();
1640 let temp_dir = TempDir::new().expect("temp dir");
1641 let _cwd = CurrentDirGuard::switch_to(temp_dir.path());
1642 let cwd = std::env::current_dir().expect("cwd after switch");
1643 let resolved = resolve_script_path(None).expect("resolve");
1644 assert_eq!(resolved, cwd.join(DEFAULT_SCRIPT_PATH));
1645 }
1646
1647 #[test]
1648 fn relative_script_paths_are_absolutized() {
1649 let _lock = cwd_lock();
1650 let temp_dir = TempDir::new().expect("temp dir");
1651 let _cwd = CurrentDirGuard::switch_to(temp_dir.path());
1652 let cwd = std::env::current_dir().expect("cwd after switch");
1653
1654 let mut cfg = ComponentConfig::new();
1655 cfg.set("script", "scripts/example.py".to_string());
1656 let resolved = resolve_script_path(Some(&cfg)).expect("resolve");
1657
1658 assert_eq!(resolved, cwd.join("scripts/example.py"));
1659 }
1660
1661 fn write_test_script(contents: &str) -> (TempDir, PathBuf) {
1662 let temp_dir = TempDir::new().expect("temp dir");
1663 let script = temp_dir.path().join("task.py");
1664 std::fs::write(&script, contents).expect("write script");
1665 (temp_dir, script)
1666 }
1667
1668 fn test_context(now_ns: u64, cl_id: u64) -> CuContext {
1669 let (clock, mock) = RobotClock::mock();
1670 mock.set_value(now_ns);
1671 let mut ctx = CuContext::builder(clock)
1672 .cl_id(cl_id)
1673 .task_ids(TEST_TASK_IDS)
1674 .build();
1675 ctx.set_current_task(0);
1676 ctx
1677 }
1678
1679 fn test_context_snapshot(now_ns: u64, cl_id: u64) -> PyTaskContextSnapshot {
1680 PyTaskContextSnapshot {
1681 now_ns,
1682 recent_ns: now_ns,
1683 cl_id,
1684 task_id: Some("py"),
1685 task_index: Some(0),
1686 }
1687 }
1688
1689 fn stop_process_backend<S>(backend: &mut ProcessBackend, ctx: &CuContext, state: &mut S)
1690 where
1691 S: Serialize + DeserializeOwned + Clone,
1692 {
1693 backend.stop(ctx, state).expect("stop backend");
1694 }
1695
1696 fn stop_embedded_backend<S>(backend: &mut EmbeddedBackend, ctx: &CuContext, state: &mut S)
1697 where
1698 S: Serialize + DeserializeOwned + Clone,
1699 {
1700 backend.stop(ctx, state).expect("stop backend");
1701 }
1702
1703 #[test]
1704 fn process_backend_keeps_absent_output_payload_when_python_does_not_touch_it() {
1705 let ctx = test_context(1, 7);
1706 let (_temp_dir, script) =
1707 write_test_script("def process(ctx, inp, state, output):\n state['seen'] = True\n");
1708 let mut backend = ProcessBackend::start(&script).expect("start backend");
1709 let output = PyCuMsg::from_output(&CuMsg::<TestPayload>::default());
1710
1711 let mut result: ProcessResult<SeenState, PyCuMsg<TestPayload>> = backend
1712 .process(
1713 &ctx,
1714 &ProcessRequest {
1715 ctx: PyTaskContextSnapshot::from_cu_context(&ctx),
1716 input: (),
1717 state: SeenState::default(),
1718 output,
1719 },
1720 )
1721 .expect("process request");
1722 stop_process_backend(&mut backend, &ctx, &mut result.state);
1723
1724 assert!(result.state.seen);
1725 assert!(result.output.payload.is_none());
1726 }
1727
1728 #[test]
1729 fn process_backend_supports_attribute_writes_on_absent_output_payload() {
1730 let ctx = test_context(41, 7);
1731 let (_temp_dir, script) = write_test_script(
1732 "def process(ctx, inp, state, output):\n output.payload.value = 41\n output.payload.flag = True\n",
1733 );
1734 let mut backend = ProcessBackend::start(&script).expect("start backend");
1735 let output = PyCuMsg::from_output(&CuMsg::<TestPayload>::default());
1736
1737 let mut result: ProcessResult<(), PyCuMsg<TestPayload>> = backend
1738 .process(
1739 &ctx,
1740 &ProcessRequest {
1741 ctx: PyTaskContextSnapshot::from_cu_context(&ctx),
1742 input: (),
1743 state: (),
1744 output,
1745 },
1746 )
1747 .expect("process request");
1748 stop_process_backend(&mut backend, &ctx, &mut result.state);
1749
1750 assert_eq!(
1751 result.output.payload,
1752 Some(TestPayload {
1753 value: 41,
1754 flag: true,
1755 })
1756 );
1757 }
1758
1759 #[test]
1760 fn process_backend_preserves_rebound_scalar_state() {
1761 let ctx = test_context(41, 7);
1762 let (_temp_dir, script) = write_test_script(
1763 "def process(ctx, inp, current_state, out):\n current_state = current_state + 1\n",
1764 );
1765 let mut backend = ProcessBackend::start(&script).expect("start backend");
1766 let output = PyCuMsg::from_output(&CuMsg::<TestPayload>::default());
1767
1768 let mut result: ProcessResult<u64, PyCuMsg<TestPayload>> = backend
1769 .process(
1770 &ctx,
1771 &ProcessRequest {
1772 ctx: PyTaskContextSnapshot::from_cu_context(&ctx),
1773 input: (),
1774 state: 41_u64,
1775 output,
1776 },
1777 )
1778 .expect("process request");
1779 stop_process_backend(&mut backend, &ctx, &mut result.state);
1780
1781 assert_eq!(result.state, 42);
1782 assert!(result.output.payload.is_none());
1783 }
1784
1785 #[test]
1786 fn process_backend_receives_ctx_snapshot() {
1787 let rust_ctx = test_context(1, 2);
1788 let request_ctx = test_context_snapshot(41, 7);
1789 let (_temp_dir, script) = write_test_script(
1790 "def process(ctx, inp, state, output):\n output.payload.value = ctx.now()\n output.payload.flag = (ctx.cl_id == 7 and ctx.task_id == 'py' and ctx.task_index == 0)\n",
1791 );
1792 let mut backend = ProcessBackend::start(&script).expect("start backend");
1793 let output = PyCuMsg::from_output(&CuMsg::<TestPayload>::default());
1794
1795 let mut result: ProcessResult<(), PyCuMsg<TestPayload>> = backend
1796 .process(
1797 &rust_ctx,
1798 &ProcessRequest {
1799 ctx: request_ctx,
1800 input: (),
1801 state: (),
1802 output,
1803 },
1804 )
1805 .expect("process request");
1806 stop_process_backend(&mut backend, &rust_ctx, &mut result.state);
1807
1808 assert_eq!(
1809 result.output.payload,
1810 Some(TestPayload {
1811 value: 41,
1812 flag: true,
1813 })
1814 );
1815 }
1816
1817 #[test]
1818 fn process_backend_optional_start_and_stop_hooks_are_noops() {
1819 let ctx = test_context(41, 7);
1820 let (_temp_dir, script) =
1821 write_test_script("def process(ctx, inp, state, output):\n pass\n");
1822 let mut backend = ProcessBackend::start(&script).expect("start backend");
1823 let mut state = 11_u64;
1824
1825 backend.start_hook(&ctx, &mut state).expect("start hook");
1826 assert_eq!(state, 11);
1827
1828 stop_process_backend(&mut backend, &ctx, &mut state);
1829 assert_eq!(state, 11);
1830 }
1831
1832 #[test]
1833 fn process_backend_runs_optional_start_and_stop_hooks() {
1834 let start_ctx = test_context(40, 2);
1835 let stop_ctx = test_context(1, 5);
1836 let (_temp_dir, script) = write_test_script(
1837 "def start(ctx, state):\n state = ctx.now() + ctx.cl_id\n\n\
1838def process(ctx, inp, state, output):\n pass\n\n\
1839def stop(ctx, state):\n state = state + ctx.cl_id\n",
1840 );
1841 let mut backend = ProcessBackend::start(&script).expect("start backend");
1842 let mut state = 0_u64;
1843
1844 backend
1845 .start_hook(&start_ctx, &mut state)
1846 .expect("start hook");
1847 assert_eq!(state, 42);
1848
1849 stop_process_backend(&mut backend, &stop_ctx, &mut state);
1850 assert_eq!(state, 47);
1851 }
1852
1853 #[test]
1854 fn process_backend_shared_handles_are_exposed_without_copying() {
1855 let ctx = test_context(41, 7);
1856 let (_temp_dir, script) = write_test_script(
1857 "def process(ctx, inp, state, output):\n\
1858 \x20\x20\x20\x20view = inp.data.memoryview()\n\
1859 \x20\x20\x20\x20output.payload.value = view[0] + view[1] + view[2] + view[3]\n\
1860 \x20\x20\x20\x20output.payload.flag = True\n\
1861 \x20\x20\x20\x20view[0] = 9\n",
1862 );
1863 let pool = CuSharedMemoryPool::<u8>::new("py_shm_test", 1, 4).expect("shared pool");
1864 let handle = pool.acquire().expect("pooled handle");
1865 handle.with_inner_mut(|inner| inner.copy_from_slice(&[1, 2, 3, 4]));
1866
1867 let mut backend = ProcessBackend::start(&script).expect("start backend");
1868 let output = PyCuMsg::from_output(&CuMsg::<TestPayload>::default());
1869 let mut result: ProcessResult<(), PyCuMsg<TestPayload>> = backend
1870 .process(
1871 &ctx,
1872 &ProcessRequest {
1873 ctx: PyTaskContextSnapshot::from_cu_context(&ctx),
1874 input: SharedHandleInput {
1875 data: handle.clone(),
1876 },
1877 state: (),
1878 output,
1879 },
1880 )
1881 .expect("process request");
1882 stop_process_backend(&mut backend, &ctx, &mut result.state);
1883
1884 assert_eq!(
1885 result.output.payload,
1886 Some(TestPayload {
1887 value: 10,
1888 flag: true,
1889 })
1890 );
1891 let observed = handle.with_inner(|inner| inner[0]);
1892 assert_eq!(observed, 9);
1893 }
1894
1895 #[test]
1896 fn cbor_wire_frames_round_trip_128_bit_integers() {
1897 let mut buffer = Vec::new();
1898 let large_u128 = u128::from(u64::MAX) + 99;
1899 let large_i128 = i128::from(i64::MIN) - 99;
1900
1901 write_cbor_frame(&mut buffer, &large_u128).expect("serialize u128");
1902 let decoded_u128 = read_cbor_frame::<_, u128>(&mut std::io::Cursor::new(&buffer))
1903 .expect("deserialize u128");
1904 assert_eq!(decoded_u128, large_u128);
1905
1906 buffer.clear();
1907
1908 write_cbor_frame(&mut buffer, &large_i128).expect("serialize i128");
1909 let decoded_i128 = read_cbor_frame::<_, i128>(&mut std::io::Cursor::new(&buffer))
1910 .expect("deserialize i128");
1911 assert_eq!(decoded_i128, large_i128);
1912 }
1913
1914 #[test]
1915 fn embedded_backend_supports_attribute_writes_on_absent_output_payload() {
1916 let ctx = test_context(41, 7);
1917 let (_temp_dir, script) = write_test_script(
1918 "def process(ctx, inp, state, output):\n output.payload.value = 41\n output.payload.flag = True\n",
1919 );
1920 let mut backend = EmbeddedBackend::start(&script).expect("start backend");
1921 let output = PyCuMsg::from_output(&CuMsg::<TestPayload>::default());
1922
1923 let mut result: ProcessResult<(), PyCuMsg<TestPayload>> = backend
1924 .process(
1925 &ctx,
1926 &ProcessRequest {
1927 ctx: PyTaskContextSnapshot::from_cu_context(&ctx),
1928 input: (),
1929 state: (),
1930 output,
1931 },
1932 )
1933 .expect("process request");
1934 stop_embedded_backend(&mut backend, &ctx, &mut result.state);
1935
1936 assert_eq!(
1937 result.output.payload,
1938 Some(TestPayload {
1939 value: 41,
1940 flag: true,
1941 })
1942 );
1943 }
1944
1945 #[test]
1946 fn embedded_backend_preserves_rebound_scalar_state() {
1947 let ctx = test_context(41, 7);
1948 let (_temp_dir, script) = write_test_script(
1949 "def process(ctx, inp, current_state, out):\n current_state = current_state + 1\n",
1950 );
1951 let mut backend = EmbeddedBackend::start(&script).expect("start backend");
1952 let output = PyCuMsg::from_output(&CuMsg::<TestPayload>::default());
1953
1954 let mut result: ProcessResult<u64, PyCuMsg<TestPayload>> = backend
1955 .process(
1956 &ctx,
1957 &ProcessRequest {
1958 ctx: PyTaskContextSnapshot::from_cu_context(&ctx),
1959 input: (),
1960 state: 41_u64,
1961 output,
1962 },
1963 )
1964 .expect("process request");
1965 stop_embedded_backend(&mut backend, &ctx, &mut result.state);
1966
1967 assert_eq!(result.state, 42);
1968 assert!(result.output.payload.is_none());
1969 }
1970
1971 #[test]
1972 fn embedded_backend_receives_live_ctx() {
1973 let ctx = test_context(41, 7);
1974 let request_ctx = test_context_snapshot(1, 2);
1975 let (_temp_dir, script) = write_test_script(
1976 "def process(ctx, inp, state, output):\n output.payload.value = ctx.now()\n output.payload.flag = (ctx.cl_id == 7 and ctx.task_id == 'py' and ctx.task_index == 0)\n",
1977 );
1978 let mut backend = EmbeddedBackend::start(&script).expect("start backend");
1979 let output = PyCuMsg::from_output(&CuMsg::<TestPayload>::default());
1980
1981 let mut result: ProcessResult<(), PyCuMsg<TestPayload>> = backend
1982 .process(
1983 &ctx,
1984 &ProcessRequest {
1985 ctx: request_ctx,
1986 input: (),
1987 state: (),
1988 output,
1989 },
1990 )
1991 .expect("process request");
1992 stop_embedded_backend(&mut backend, &ctx, &mut result.state);
1993
1994 assert_eq!(
1995 result.output.payload,
1996 Some(TestPayload {
1997 value: 41,
1998 flag: true,
1999 })
2000 );
2001 }
2002
2003 #[test]
2004 fn embedded_backend_optional_start_and_stop_hooks_are_noops() {
2005 let ctx = test_context(41, 7);
2006 let (_temp_dir, script) =
2007 write_test_script("def process(ctx, inp, state, output):\n pass\n");
2008 let mut backend = EmbeddedBackend::start(&script).expect("start backend");
2009 let mut state = 11_u64;
2010
2011 backend.start_hook(&ctx, &mut state).expect("start hook");
2012 assert_eq!(state, 11);
2013
2014 stop_embedded_backend(&mut backend, &ctx, &mut state);
2015 assert_eq!(state, 11);
2016 }
2017
2018 #[test]
2019 fn embedded_backend_runs_optional_start_and_stop_hooks() {
2020 let start_ctx = test_context(40, 2);
2021 let stop_ctx = test_context(1, 5);
2022 let (_temp_dir, script) = write_test_script(
2023 "def start(ctx, state):\n state = ctx.now() + ctx.cl_id\n\n\
2024def process(ctx, inp, state, output):\n pass\n\n\
2025def stop(ctx, state):\n state = state + ctx.cl_id\n",
2026 );
2027 let mut backend = EmbeddedBackend::start(&script).expect("start backend");
2028 let mut state = 0_u64;
2029
2030 backend
2031 .start_hook(&start_ctx, &mut state)
2032 .expect("start hook");
2033 assert_eq!(state, 42);
2034
2035 stop_embedded_backend(&mut backend, &stop_ctx, &mut state);
2036 assert_eq!(state, 47);
2037 }
2038}