soma_provider_adapters/
python_protocol.rs

1use std::path::{Path, PathBuf};
2
3use serde::{Deserialize, Serialize, de::DeserializeOwned};
4use serde_json::Value;
5use soma_provider_core::{ProviderCall, ProviderSurface};
6use thiserror::Error;
7
8pub(crate) const PYTHON_WORKER_SCHEMA_VERSION: u32 = 1;
9pub(crate) const ONE_SHOT_REQUEST_ID: u64 = 0;
10pub(crate) const PYTHON_PROTOCOL_HEADROOM_BYTES: usize = 1024;
11
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
13#[serde(tag = "mode", rename_all = "snake_case")]
14pub(crate) enum PythonWorkerRequest {
15    Catalog {
16        schema_version: u32,
17        request_id: u64,
18        path: PathBuf,
19    },
20    Call {
21        schema_version: u32,
22        request_id: u64,
23        path: PathBuf,
24        #[serde(default)]
25        env_keys: Vec<String>,
26        provider: String,
27        action: String,
28        params: Value,
29        surface: ProviderSurface,
30        snapshot_id: String,
31    },
32}
33
34impl PythonWorkerRequest {
35    pub(crate) fn catalog(path: &Path) -> Self {
36        Self::Catalog {
37            schema_version: PYTHON_WORKER_SCHEMA_VERSION,
38            request_id: ONE_SHOT_REQUEST_ID,
39            path: path.to_path_buf(),
40        }
41    }
42
43    pub(crate) fn call(path: &Path, call: &ProviderCall, env_keys: Vec<String>) -> Self {
44        Self::Call {
45            schema_version: PYTHON_WORKER_SCHEMA_VERSION,
46            request_id: ONE_SHOT_REQUEST_ID,
47            path: path.to_path_buf(),
48            env_keys,
49            provider: call.provider.clone(),
50            action: call.action.clone(),
51            params: call.params.clone(),
52            surface: call.surface,
53            snapshot_id: call.snapshot_id.clone(),
54        }
55    }
56
57    fn schema_version(&self) -> u32 {
58        match self {
59            Self::Catalog { schema_version, .. } | Self::Call { schema_version, .. } => {
60                *schema_version
61            }
62        }
63    }
64
65    fn request_id(&self) -> u64 {
66        match self {
67            Self::Catalog { request_id, .. } | Self::Call { request_id, .. } => *request_id,
68        }
69    }
70
71    fn mode(&self) -> &'static str {
72        match self {
73            Self::Catalog { .. } => "catalog",
74            Self::Call { .. } => "call",
75        }
76    }
77}
78
79#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
80#[serde(tag = "mode", rename_all = "snake_case")]
81pub(crate) enum PythonWorkerResponse {
82    Catalog {
83        schema_version: u32,
84        request_id: u64,
85        catalog: Value,
86    },
87    Call {
88        schema_version: u32,
89        request_id: u64,
90        output: Value,
91    },
92}
93
94impl PythonWorkerResponse {
95    fn schema_version(&self) -> u32 {
96        match self {
97            Self::Catalog { schema_version, .. } | Self::Call { schema_version, .. } => {
98                *schema_version
99            }
100        }
101    }
102
103    fn request_id(&self) -> u64 {
104        match self {
105            Self::Catalog { request_id, .. } | Self::Call { request_id, .. } => *request_id,
106        }
107    }
108
109    fn mode(&self) -> &'static str {
110        match self {
111            Self::Catalog { .. } => "catalog",
112            Self::Call { .. } => "call",
113        }
114    }
115}
116
117/// Current major version of the persistent Python runner protocol.
118pub const PYTHON_RUNNER_PROTOCOL_MAJOR: u16 = 1;
119/// Current minor version of the persistent Python runner protocol.
120pub const PYTHON_RUNNER_PROTOCOL_MINOR: u16 = 0;
121/// Maximum JSON payload accepted by the length-prefixed control channel.
122pub const PYTHON_RUNNER_MAX_FRAME_BYTES: usize = 8 * 1024 * 1024;
123
124/// A separately negotiated runner protocol version.
125#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
126pub struct PythonRunnerProtocolVersion {
127    pub major: u16,
128    pub minor: u16,
129}
130
131impl PythonRunnerProtocolVersion {
132    #[must_use]
133    pub const fn current() -> Self {
134        Self {
135            major: PYTHON_RUNNER_PROTOCOL_MAJOR,
136            minor: PYTHON_RUNNER_PROTOCOL_MINOR,
137        }
138    }
139
140    pub fn negotiate(self, peer: Self) -> Result<Self, PythonProtocolError> {
141        if self.major != peer.major {
142            return Err(PythonProtocolError::ProtocolMajorMismatch {
143                host: self.major,
144                worker: peer.major,
145            });
146        }
147        Ok(Self {
148            major: self.major,
149            minor: self.minor.min(peer.minor),
150        })
151    }
152}
153
154/// Optional protocol behavior advertised during the worker handshake.
155#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
156#[serde(rename_all = "kebab-case")]
157pub enum PythonRunnerFeature {
158    Describe,
159    Invoke,
160    Cancel,
161    Health,
162    Drain,
163    Shutdown,
164    HostCalls,
165}
166
167/// Python implementation details supplied by the worker, never inferred by the host.
168#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
169pub struct PythonRuntimeIdentity {
170    pub implementation: String,
171    pub version: String,
172}
173
174/// Worker-initiated handshake sent before any host request.
175#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
176pub struct PythonRunnerHello {
177    pub protocol: PythonRunnerProtocolVersion,
178    pub sdk_version: String,
179    pub python: PythonRuntimeIdentity,
180    #[serde(default)]
181    pub features: Vec<PythonRunnerFeature>,
182}
183
184/// Host-to-worker messages on the persistent control channel.
185#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
186#[serde(tag = "type", rename_all = "snake_case")]
187pub enum PythonRunnerHostMessage {
188    Initialize {
189        protocol: PythonRunnerProtocolVersion,
190        #[serde(default)]
191        features: Vec<PythonRunnerFeature>,
192        generation_id: String,
193    },
194    Request {
195        #[serde(flatten)]
196        request: PythonRunnerHostRequest,
197    },
198    HostReply {
199        request_id: u64,
200        #[serde(default, skip_serializing_if = "Option::is_none")]
201        result: Option<Value>,
202        #[serde(default, skip_serializing_if = "Option::is_none")]
203        error: Option<PythonRunnerError>,
204    },
205}
206
207/// Worker-to-host messages on the persistent control channel.
208#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
209#[serde(tag = "type", rename_all = "snake_case")]
210pub enum PythonRunnerWorkerMessage {
211    Hello {
212        protocol: PythonRunnerProtocolVersion,
213        sdk_version: String,
214        python: PythonRuntimeIdentity,
215        #[serde(default)]
216        features: Vec<PythonRunnerFeature>,
217    },
218    Ready {
219        protocol: PythonRunnerProtocolVersion,
220        #[serde(default)]
221        features: Vec<PythonRunnerFeature>,
222        generation_id: String,
223    },
224    Reply {
225        #[serde(flatten)]
226        reply: PythonRunnerReply,
227    },
228    HostCall {
229        #[serde(flatten)]
230        call: PythonRunnerHostCall,
231    },
232}
233
234/// Host-side lifecycle for one at-most-once invocation.
235#[derive(Debug, Clone, Copy, PartialEq, Eq)]
236pub enum PythonRequestState {
237    Queued,
238    Written,
239    Accepted,
240    Terminal,
241    TimedOut,
242    WorkerLost,
243}
244
245impl PythonRequestState {
246    #[must_use]
247    pub const fn is_pending(self) -> bool {
248        matches!(self, Self::Queued | Self::Written | Self::Accepted)
249    }
250}
251
252/// Return the deterministic feature intersection in host preference order.
253#[must_use]
254pub fn negotiate_runner_features(
255    host: &[PythonRunnerFeature],
256    worker: &[PythonRunnerFeature],
257) -> Vec<PythonRunnerFeature> {
258    host.iter()
259        .copied()
260        .filter(|feature| worker.contains(feature))
261        .collect()
262}
263
264/// Trace context propagated with one invocation.
265#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
266pub struct PythonTraceContext {
267    pub traceparent: String,
268    #[serde(default, skip_serializing_if = "Option::is_none")]
269    pub tracestate: Option<String>,
270}
271
272/// Authenticated actor and scopes captured when dispatch begins.
273#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
274pub struct PythonActorContext {
275    pub actor_id: String,
276    #[serde(default)]
277    pub scopes: Vec<String>,
278}
279
280/// Complete at-most-once invocation envelope.
281#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
282pub struct PythonInvocationRequest {
283    pub invocation_id: String,
284    /// Surface request identity; unlike the numeric frame id this remains
285    /// stable across runtime boundaries.
286    pub request_id: String,
287    pub provider: String,
288    pub action: String,
289    pub arguments: Value,
290    pub surface: ProviderSurface,
291    pub snapshot_id: String,
292    pub deadline_unix_ms: u64,
293    #[serde(default, skip_serializing_if = "Option::is_none")]
294    pub trace: Option<PythonTraceContext>,
295    #[serde(default, skip_serializing_if = "Option::is_none")]
296    pub actor: Option<PythonActorContext>,
297    pub cancellation_token_id: String,
298    pub generation_id: String,
299}
300
301/// Requests sent by the Soma host over the dedicated runner control channel.
302#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
303#[serde(tag = "method", rename_all = "snake_case")]
304pub enum PythonRunnerHostRequest {
305    Describe {
306        request_id: u64,
307        path: PathBuf,
308        generation_id: String,
309    },
310    Invoke {
311        request_id: u64,
312        invocation: Box<PythonInvocationRequest>,
313    },
314    Cancel {
315        request_id: u64,
316        invocation_id: String,
317        cancellation_token_id: String,
318    },
319    Health {
320        request_id: u64,
321    },
322    Drain {
323        request_id: u64,
324    },
325    Shutdown {
326        request_id: u64,
327    },
328}
329
330/// Worker lifecycle visible to the supervisor.
331#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
332#[serde(rename_all = "snake_case")]
333pub enum PythonWorkerHealth {
334    Starting,
335    Ready,
336    Draining,
337    Unhealthy,
338}
339
340/// Stable invocation lifecycle used to decide whether retry is safe.
341#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
342#[serde(rename_all = "snake_case")]
343pub enum PythonInvocationState {
344    Pending,
345    Accepted,
346    Running,
347    Completed,
348    Cancelled,
349    Indeterminate,
350}
351
352/// Worker replies to host requests.
353#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
354#[serde(tag = "status", rename_all = "snake_case")]
355pub enum PythonRunnerReply {
356    Ok {
357        request_id: u64,
358        result: Value,
359    },
360    Accepted {
361        request_id: u64,
362        invocation_id: String,
363        state: PythonInvocationState,
364    },
365    Health {
366        request_id: u64,
367        health: PythonWorkerHealth,
368        generation_id: String,
369    },
370    Error {
371        request_id: u64,
372        error: PythonRunnerError,
373    },
374}
375
376/// Brokered calls a worker may request from the host.
377#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
378#[serde(tag = "method")]
379pub enum PythonRunnerHostCall {
380    #[serde(rename = "host.http")]
381    Http {
382        request_id: u64,
383        invocation_id: String,
384        request: Value,
385    },
386    #[serde(rename = "host.secret")]
387    Secret {
388        request_id: u64,
389        invocation_id: String,
390        name: String,
391    },
392    #[serde(rename = "host.state.get")]
393    StateGet {
394        request_id: u64,
395        invocation_id: String,
396        key: String,
397    },
398    #[serde(rename = "host.state.put")]
399    StatePut {
400        request_id: u64,
401        invocation_id: String,
402        key: String,
403        value: Value,
404    },
405    #[serde(rename = "host.log")]
406    Log {
407        request_id: u64,
408        invocation_id: String,
409        level: String,
410        message: String,
411        #[serde(default)]
412        fields: Value,
413    },
414    #[serde(rename = "host.metric")]
415    Metric {
416        request_id: u64,
417        invocation_id: String,
418        name: String,
419        value: serde_json::Number,
420        #[serde(default)]
421        attributes: Value,
422    },
423    #[serde(rename = "host.progress")]
424    Progress {
425        request_id: u64,
426        invocation_id: String,
427        current: u64,
428        #[serde(default, skip_serializing_if = "Option::is_none")]
429        total: Option<u64>,
430        #[serde(default, skip_serializing_if = "Option::is_none")]
431        message: Option<String>,
432    },
433    #[serde(rename = "host.cancelled")]
434    Cancelled {
435        request_id: u64,
436        invocation_id: String,
437    },
438}
439
440/// Stable public error identifiers for Python setup and execution.
441#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
442#[serde(rename_all = "snake_case")]
443pub enum PythonRunnerErrorCode {
444    PythonRuntimeMissing,
445    PythonVersionUnsupported,
446    PythonDependencyResolutionFailed,
447    PythonWorkerStartFailed,
448    PythonProtocolMismatch,
449    PythonCatalogTimeout,
450    PythonImportFailed,
451    PythonSchemaInvalid,
452    PythonPolicyDenied,
453    PythonCallTimeout,
454    PythonCallCancelled,
455    PythonWorkerCrashed,
456    PythonOutputTooLarge,
457    PythonInvalidOutput,
458    PythonNativeAbiMismatch,
459}
460
461/// Stable stage at which a runner error occurred.
462#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
463#[serde(rename_all = "snake_case")]
464pub enum PythonRunnerErrorPhase {
465    Runtime,
466    DependencyResolution,
467    WorkerStartup,
468    Protocol,
469    Catalog,
470    Import,
471    Schema,
472    Policy,
473    Invocation,
474    NativeBinding,
475}
476
477/// Redacted error envelope safe to cross the private runner protocol.
478#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
479pub struct PythonRunnerError {
480    pub code: PythonRunnerErrorCode,
481    pub phase: PythonRunnerErrorPhase,
482    #[serde(default, skip_serializing_if = "Option::is_none")]
483    pub provider: Option<String>,
484    #[serde(default, skip_serializing_if = "Option::is_none")]
485    pub source: Option<String>,
486    #[serde(default, skip_serializing_if = "Option::is_none")]
487    pub generation_id: Option<String>,
488    #[serde(default, skip_serializing_if = "Option::is_none")]
489    pub action: Option<String>,
490    pub retryable: bool,
491    pub public_message: String,
492}
493
494#[derive(Debug, Error)]
495pub enum PythonProtocolError {
496    #[error("invalid Python worker JSON: {0}")]
497    Json(#[from] serde_json::Error),
498    #[error("unsupported Python worker schema version {actual}; expected {expected}")]
499    UnsupportedSchemaVersion { expected: u32, actual: u32 },
500    #[error("unexpected Python worker request id {actual}; expected {expected}")]
501    UnexpectedRequestId { expected: u64, actual: u64 },
502    #[error("unexpected Python worker response mode {actual}; expected {expected}")]
503    UnexpectedResponseMode {
504        expected: &'static str,
505        actual: &'static str,
506    },
507    #[error("Python runner protocol major mismatch: host {host}, worker {worker}")]
508    ProtocolMajorMismatch { host: u16, worker: u16 },
509    #[error("Python runner frame header is incomplete: got {actual} bytes; expected 4")]
510    FrameHeaderTooShort { actual: usize },
511    #[error("Python runner frame payload is {actual} bytes; limit is {limit}")]
512    FrameTooLarge { limit: usize, actual: usize },
513    #[error("Python runner frame length mismatch: declared {declared} bytes; got {actual}")]
514    FrameLengthMismatch { declared: usize, actual: usize },
515}
516
517/// Encode one control message as a four-byte big-endian length plus UTF-8 JSON.
518pub fn encode_runner_frame<T: Serialize>(message: &T) -> Result<Vec<u8>, PythonProtocolError> {
519    let payload = serde_json::to_vec(message)?;
520    if payload.len() > PYTHON_RUNNER_MAX_FRAME_BYTES {
521        return Err(PythonProtocolError::FrameTooLarge {
522            limit: PYTHON_RUNNER_MAX_FRAME_BYTES,
523            actual: payload.len(),
524        });
525    }
526
527    let mut frame = Vec::with_capacity(4 + payload.len());
528    frame.extend_from_slice(&(payload.len() as u32).to_be_bytes());
529    frame.extend_from_slice(&payload);
530    Ok(frame)
531}
532
533/// Decode exactly one complete control frame.
534pub fn decode_runner_frame<T: DeserializeOwned>(bytes: &[u8]) -> Result<T, PythonProtocolError> {
535    if bytes.len() < 4 {
536        return Err(PythonProtocolError::FrameHeaderTooShort {
537            actual: bytes.len(),
538        });
539    }
540
541    let declared = u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize;
542    if declared > PYTHON_RUNNER_MAX_FRAME_BYTES {
543        return Err(PythonProtocolError::FrameTooLarge {
544            limit: PYTHON_RUNNER_MAX_FRAME_BYTES,
545            actual: declared,
546        });
547    }
548
549    let payload = &bytes[4..];
550    if payload.len() != declared {
551        return Err(PythonProtocolError::FrameLengthMismatch {
552            declared,
553            actual: payload.len(),
554        });
555    }
556
557    Ok(serde_json::from_slice(payload)?)
558}
559
560pub(crate) fn encode_python_request(
561    request: &PythonWorkerRequest,
562) -> Result<Vec<u8>, PythonProtocolError> {
563    validate_schema_version(request.schema_version())?;
564    let mut encoded = serde_json::to_vec(request)?;
565    encoded.push(b'\n');
566    Ok(encoded)
567}
568
569pub(crate) fn decode_python_response(
570    bytes: &[u8],
571) -> Result<PythonWorkerResponse, PythonProtocolError> {
572    let response: PythonWorkerResponse = serde_json::from_slice(bytes)?;
573    validate_schema_version(response.schema_version())?;
574    Ok(response)
575}
576
577pub(crate) fn validate_python_response(
578    request: &PythonWorkerRequest,
579    response: &PythonWorkerResponse,
580) -> Result<(), PythonProtocolError> {
581    validate_schema_version(request.schema_version())?;
582    validate_schema_version(response.schema_version())?;
583    if response.request_id() != request.request_id() {
584        return Err(PythonProtocolError::UnexpectedRequestId {
585            expected: request.request_id(),
586            actual: response.request_id(),
587        });
588    }
589    if response.mode() != request.mode() {
590        return Err(PythonProtocolError::UnexpectedResponseMode {
591            expected: request.mode(),
592            actual: response.mode(),
593        });
594    }
595    Ok(())
596}
597
598fn validate_schema_version(actual: u32) -> Result<(), PythonProtocolError> {
599    if actual == PYTHON_WORKER_SCHEMA_VERSION {
600        return Ok(());
601    }
602    Err(PythonProtocolError::UnsupportedSchemaVersion {
603        expected: PYTHON_WORKER_SCHEMA_VERSION,
604        actual,
605    })
606}
607
608#[cfg(test)]
609#[path = "python_protocol_tests.rs"]
610mod tests;
soma_provider_adapters/python_protocol.rs

soma_provider_adapters/
python_protocol.rs
