Skip to main content

cu_python_task/
lib.rs

1//! Python-backed Copper tasks.
2//!
3//! This crate exists for rapid algorithm prototyping, not for production realtime
4//! robotics.
5//!
6//! A [`PyTask`] lets Copper keep ownership of scheduling, logging, replay, and task
7//! lifecycle while delegating one task's algorithm body to a Python function:
8//!
9//! ```python
10//! def process(ctx, input, state, output):
11//!     ...
12//! ```
13//!
14//! Two execution modes are available:
15//!
16//! - [`PyTaskMode::Process`]: spawn a separate interpreter and exchange
17//!   length-prefixed CBOR frames over stdin/stdout. This avoids putting the GIL
18//!   inside the Copper process, but adds another serialization layer, extra
19//!   copying, allocations, IPC overhead, and scheduler jitter. If a payload
20//!   contains `CuHandle<CuSharedMemoryBuffer<T>>`, the handle is exported by
21//!   descriptor so Python can read or write the underlying shared-memory bytes
22//!   without copying them through CBOR.
23//! - [`PyTaskMode::Embedded`]: call Python in-process through PyO3. This avoids
24//!   the external CBOR transport, but executes under the GIL inside the Copper
25//!   process and still allocates and converts values on every call.
26//!
27//! Both modes are fundamentally at odds with Copper's design center: predictable
28//! low-latency execution with minimal allocation on the realtime path. Using
29//! Python here will increase latency, jitter, and allocation pressure enough to
30//! ruin the realtime characteristics of the stack. Compared to a native Rust
31//! Copper task, the performance is abysmal.
32//!
33//! The intended workflow is narrow:
34//!
35//! - prototype one task quickly in Python
36//! - stabilize the algorithm
37//! - rewrite it in Rust, optionally using an LLM-assisted translation as a first
38//!   draft
39//!
40//! Do not treat Python tasks as a normal production integration path.
41
42use 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
69/// State carried across invocations of a [`PyTask`].
70///
71/// The state must be serializable because both backends turn it into an owned,
72/// mutable Python value on every `process(...)` call and then deserialize the
73/// result back into Rust.
74pub 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
100/// Describes how a Python-backed task receives its Copper inputs.
101///
102/// Python cannot safely borrow Copper messages directly, so every call converts
103/// the runtime input into an owned representation first.
104pub 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
113/// Describes how a Python-backed task exposes Copper outputs.
114///
115/// Outputs are converted into mutable owned values before calling Python, then
116/// written back into Copper message slots after the Python function returns.
117pub 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
250/// Convenience alias for the common one-input, one-output case.
251pub 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    /// Run the task in a separate Python interpreter process.
399    ///
400    /// This keeps the GIL out of the Copper process, but every cycle pays for
401    /// CBOR serialization, copies, allocations, IPC, and process scheduling.
402    /// Shared-memory-backed `CuHandle` buffers are exported by descriptor
403    /// automatically on this path.
404    Process,
405    /// Run the task inside the Copper process through PyO3.
406    ///
407    /// This removes the external transport layer, but now the GIL and Python
408    /// runtime are inside the same process as Copper. This mode is unsupported
409    /// on macOS in this workspace.
410    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/// A Copper task whose `process(...)` implementation is provided by a Python
585/// script.
586///
587/// The configured script must define:
588///
589/// ```python
590/// def process(ctx, input, state, output):
591///     ...
592/// ```
593///
594/// It may also define optional lifecycle hooks:
595///
596/// ```python
597/// def start(ctx, state):
598///     ...
599///
600/// def stop(ctx, state):
601///     ...
602/// ```
603///
604/// Missing `start`/`stop` hooks are treated as no-ops.
605///
606/// `ctx` is passed as the first argument. In embedded mode it is a live PyO3
607/// wrapper over the current Copper context; in process mode it is a per-call
608/// snapshot exposing the same clock/task metadata API.
609///
610/// Configuration keys:
611///
612/// - `script`: path to the Python file, default `python/task.py`; relative
613///   paths are resolved against the process current working directory
614/// - `mode`: `"process"` or `"embedded"`, default `"process"`
615///
616/// `input`, `state`, and `output` are all converted into owned values before the
617/// Python call. That makes the API flexible, but it also means allocations and
618/// copies happen on every call. This is the core reason the type is suitable for
619/// prototyping only and strongly discouraged on a realtime path.
620#[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}