soma_provider_adapters/python/
supervisor_state.rs

1use std::{
2    collections::BTreeMap,
3    sync::{
4        Arc, Mutex, OnceLock, Weak,
5        atomic::{AtomicBool, Ordering},
6    },
7    time::Duration,
8};
9
10use serde::Serialize;
11use tokio::sync::Semaphore;
12
13use crate::python_protocol::{PythonProtocolError, PythonRunnerErrorCode};
14
15use super::PythonSupervisorError;
16
17/// One redacted, structured line emitted by a persistent Python worker.
18#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
19pub struct PythonWorkerLogEntry {
20    pub sequence: u64,
21    pub stream: &'static str,
22    pub message: String,
23}
24
25/// Per-call identity, policy context, and deadline for a Python invocation.
26#[derive(Debug, Clone, Copy)]
27pub struct PythonInvocationOptions<'a> {
28    pub surface: soma_provider_core::ProviderSurface,
29    pub snapshot_id: &'a str,
30    pub timeout: Duration,
31    pub context: &'a soma_provider_core::ProviderInvocationContext,
32}
33
34impl std::fmt::Display for PythonSupervisorError {
35    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        formatter.write_str(&self.message)
37    }
38}
39
40impl std::error::Error for PythonSupervisorError {}
41
42pub(super) fn start_error() -> PythonSupervisorError {
43    PythonSupervisorError::new(
44        "python_worker_start_failed",
45        "Python worker could not be started",
46    )
47}
48
49pub(super) fn worker_budget(group: &str, limit: usize) -> Arc<Semaphore> {
50    let limit = limit.max(1);
51    let key = (group.to_owned(), limit);
52    let mut budgets = WORKER_BUDGETS
53        .get_or_init(|| Mutex::new(BTreeMap::new()))
54        .lock()
55        .expect("Python worker budget lock should not be poisoned");
56    budgets.retain(|_, budget| budget.strong_count() != 0);
57    if let Some(budget) = budgets.get(&key).and_then(Weak::upgrade) {
58        return budget;
59    }
60    let budget = Arc::new(Semaphore::new(limit));
61    budgets.insert(key, Arc::downgrade(&budget));
62    budget
63}
64
65pub(super) fn candidate_budget(limit: usize) -> Arc<Semaphore> {
66    shared_budget(limit, &CANDIDATE_BUDGETS)
67}
68
69type BudgetMap = Mutex<BTreeMap<usize, Weak<Semaphore>>>;
70type WorkerBudgetMap = Mutex<BTreeMap<(String, usize), Weak<Semaphore>>>;
71static WORKER_BUDGETS: OnceLock<WorkerBudgetMap> = OnceLock::new();
72static CANDIDATE_BUDGETS: OnceLock<BudgetMap> = OnceLock::new();
73
74#[cfg(test)]
75pub(super) fn worker_budget_keys_are_live() -> bool {
76    WORKER_BUDGETS
77        .get()
78        .expect("worker budgets")
79        .lock()
80        .expect("worker budget lock")
81        .values()
82        .all(|budget| budget.strong_count() != 0)
83}
84
85fn shared_budget(limit: usize, budgets: &'static OnceLock<BudgetMap>) -> Arc<Semaphore> {
86    let limit = limit.max(1);
87    let mut budgets = budgets
88        .get_or_init(|| Mutex::new(BTreeMap::new()))
89        .lock()
90        .expect("Python worker budget lock should not be poisoned");
91    budgets.retain(|_, budget| budget.strong_count() != 0);
92    if let Some(budget) = budgets.get(&limit).and_then(Weak::upgrade) {
93        return budget;
94    }
95    let budget = Arc::new(Semaphore::new(limit));
96    budgets.insert(limit, Arc::downgrade(&budget));
97    budget
98}
99
100pub(super) struct BusyGuard<'a> {
101    busy: &'a AtomicBool,
102    discard_worker: &'a AtomicBool,
103    completed: bool,
104}
105
106impl<'a> BusyGuard<'a> {
107    pub(super) fn new(busy: &'a AtomicBool, discard_worker: &'a AtomicBool) -> Self {
108        Self {
109            busy,
110            discard_worker,
111            completed: false,
112        }
113    }
114
115    pub(super) fn complete(&mut self) {
116        self.completed = true;
117        self.busy.store(false, Ordering::Release);
118    }
119}
120
121impl Drop for BusyGuard<'_> {
122    fn drop(&mut self) {
123        if !self.completed {
124            self.discard_worker.store(true, Ordering::Release);
125            self.busy.store(false, Ordering::Release);
126        }
127    }
128}
129
130pub(super) fn map_worker_error(code: PythonRunnerErrorCode) -> PythonSupervisorError {
131    match code {
132        PythonRunnerErrorCode::PythonCallTimeout => PythonSupervisorError::new(
133            "python_provider_timeout",
134            "Python provider exceeded its timeout",
135        ),
136        PythonRunnerErrorCode::PythonCallCancelled => PythonSupervisorError::new(
137            "python_provider_cancelled",
138            "Python provider invocation was cancelled",
139        ),
140        PythonRunnerErrorCode::PythonOutputTooLarge => PythonSupervisorError::new(
141            "python_output_too_large",
142            "Python provider output exceeded its limit",
143        ),
144        PythonRunnerErrorCode::PythonPolicyDenied => PythonSupervisorError::new(
145            "python_policy_denied",
146            "Python provider host capability was denied",
147        ),
148        _ => PythonSupervisorError::new(
149            "python_provider_failed",
150            "Python provider invocation failed",
151        ),
152    }
153}
154
155pub(super) fn invalid_output() -> PythonSupervisorError {
156    PythonSupervisorError::new(
157        "python_invalid_output",
158        "Python provider produced invalid output",
159    )
160}
161
162pub(super) fn protocol_error() -> PythonSupervisorError {
163    PythonSupervisorError::new(
164        "python_protocol_mismatch",
165        "Python worker violated the runner protocol",
166    )
167}
168
169impl From<PythonProtocolError> for PythonSupervisorError {
170    fn from(_: PythonProtocolError) -> Self {
171        protocol_error()
172    }
173}
174
175impl From<std::io::Error> for PythonSupervisorError {
176    fn from(_: std::io::Error) -> Self {
177        PythonSupervisorError::new("python_worker_crashed", "Python worker exited unexpectedly")
178    }
179}
soma_provider_adapters/python/supervisor_state.rs

soma_provider_adapters/python/
supervisor_state.rs
