1use std::{
14 io::{Read, Write},
15 path::{Path, PathBuf},
16 process::{Command as StdCommand, Stdio as StdStdio},
17 sync::Arc,
18 thread,
19 time::{Duration, Instant},
20};
21
22use async_trait::async_trait;
23use serde_json::{json, Value};
24use soma_provider_core::{
25 validate_provider_manifest_value, Provider, ProviderCall, ProviderCatalog, ProviderError,
26 ProviderOutput, ProviderTool,
27};
28use tokio::time::Instant as TokioInstant;
29
30use crate::{
31 error::{redact_public, SidecarError},
32 python_bridge::PYTHON_BRIDGE,
33 sidecar::{
34 collect_provider_env, output_exceeded_message, resolve_sidecar_command,
35 run_bounded_sidecar, sidecar_base_env,
36 },
37};
38
39const DEFAULT_TIMEOUT_MS: u64 = 10_000;
40const DEFAULT_MAX_INPUT_BYTES: usize = 64 * 1024;
41const DEFAULT_MAX_OUTPUT_BYTES: usize = 256 * 1024;
42
43#[derive(Clone)]
44pub struct PythonProvider {
45 path: PathBuf,
46 catalog: ProviderCatalog,
47 env_prefix: String,
48}
49
50impl PythonProvider {
51 pub fn new(path: PathBuf, catalog: ProviderCatalog, env_prefix: impl Into<String>) -> Self {
52 Self {
53 path,
54 catalog,
55 env_prefix: env_prefix.into(),
56 }
57 }
58
59 pub fn arc(
60 path: PathBuf,
61 catalog: ProviderCatalog,
62 env_prefix: impl Into<String>,
63 ) -> Arc<Self> {
64 Arc::new(Self::new(path, catalog, env_prefix))
65 }
66}
67
68#[async_trait]
69impl Provider for PythonProvider {
70 fn catalog(&self) -> ProviderCatalog {
71 self.catalog.clone()
72 }
73
74 async fn call(&self, call: ProviderCall) -> Result<ProviderOutput, ProviderError> {
75 let tool = self.tool(&call)?;
76 let runtime = PythonRuntime::from_tool(&self.catalog, tool, &call, &self.env_prefix)?;
77 let source = self.path.display().to_string();
78 let input = python_execution_payload(&self.path, &call, &runtime.env).map_err(|error| {
79 ProviderError::execution(&self.catalog.provider.name, "", error)
80 .with_provider_kind(self.catalog.provider.kind.as_str())
81 .with_source(source.clone())
82 .with_phase("input-serialization")
83 })?;
84
85 if input.len() > runtime.max_input_bytes {
86 return Err(ProviderError::validation(
87 &self.catalog.provider.name,
88 &call.action,
89 "python_input_too_large",
90 format!(
91 "Python provider input exceeds {} bytes",
92 runtime.max_input_bytes
93 ),
94 )
95 .with_provider_kind(self.catalog.provider.kind.as_str())
96 .with_source(source)
97 .with_phase("input-validation"));
98 }
99
100 let started = TokioInstant::now();
101 let sidecar = match run_bounded_sidecar(
102 &runtime.command,
103 &["-c", PYTHON_BRIDGE],
104 runtime.env,
105 &input,
106 runtime.timeout_ms,
107 runtime.max_output_bytes,
108 )
109 .await
110 {
111 Ok(sidecar) => sidecar,
112 Err(SidecarError::Timeout) => {
113 return Err(ProviderError::new(
114 "python_provider_timeout",
115 &self.catalog.provider.name,
116 Some(call.action.clone()),
117 format!("Python provider exceeded {}ms timeout", runtime.timeout_ms),
118 "Increase tool.limits.timeout_ms or fix the Python provider handler.",
119 )
120 .with_provider_kind(self.catalog.provider.kind.as_str())
121 .with_source(source)
122 .with_phase("execution"));
123 }
124 Err(error) => {
125 return Err(ProviderError::execution(
126 &self.catalog.provider.name,
127 call.action.clone(),
128 error,
129 )
130 .with_provider_kind(self.catalog.provider.kind.as_str())
131 .with_source(source)
132 .with_phase("execution"));
133 }
134 };
135 let output = sidecar.output;
136
137 tracing::debug!(
138 provider = %self.catalog.provider.name,
139 action = %call.action,
140 elapsed_ms = started.elapsed().as_millis(),
141 "Python provider sidecar completed"
142 );
143
144 if sidecar.stdout_exceeded || sidecar.stderr_exceeded {
145 let stream = if sidecar.stdout_exceeded {
146 "stdout"
147 } else {
148 "stderr"
149 };
150 return Err(ProviderError::validation(
151 &self.catalog.provider.name,
152 &call.action,
153 "python_output_too_large",
154 output_exceeded_message(stream, runtime.max_output_bytes),
155 )
156 .with_provider_kind(self.catalog.provider.kind.as_str())
157 .with_source(source)
158 .with_phase("output-validation"));
159 }
160 if !output.status.success() {
161 let stderr = String::from_utf8_lossy(&output.stderr);
162 let code = if stderr.contains("python_provider_unserializable_output") {
163 "python_provider_unserializable_output"
164 } else {
165 "python_provider_failed"
166 };
167 return Err(ProviderError::new(
168 code,
169 &self.catalog.provider.name,
170 Some(call.action),
171 format!("Python provider failed: {}", redact_public(&stderr)),
172 "Fix the Python provider handler and retry.",
173 )
174 .with_provider_kind(self.catalog.provider.kind.as_str())
175 .with_source(source)
176 .with_phase("execution"));
177 }
178
179 let value = serde_json::from_slice(&output.stdout).map_err(|error| {
180 ProviderError::validation(
181 &self.catalog.provider.name,
182 &call.action,
183 "python_invalid_json_output",
184 error.to_string(),
185 )
186 .with_provider_kind(self.catalog.provider.kind.as_str())
187 .with_source(source)
188 .with_phase("output-validation")
189 })?;
190 Ok(ProviderOutput::json(value))
191 }
192}
193
194fn python_execution_payload(
195 path: &Path,
196 call: &ProviderCall,
197 env: &[(String, String)],
198) -> Result<Vec<u8>, serde_json::Error> {
199 let mut payload = serde_json::to_value(crate::sidecar::ExecutionEnvelope::new(call))?;
200 if let Some(object) = payload.as_object_mut() {
201 let env_keys: Vec<&str> = env.iter().map(|(key, _)| key.as_str()).collect();
202 object.insert("mode".to_owned(), json!("call"));
203 object.insert("path".to_owned(), json!(path.to_path_buf()));
204 object.insert("env_keys".to_owned(), json!(env_keys));
205 }
206 serde_json::to_vec(&payload)
207}
208
209impl PythonProvider {
210 fn tool(&self, call: &ProviderCall) -> Result<&ProviderTool, ProviderError> {
211 self.catalog
212 .tools
213 .iter()
214 .find(|tool| tool.name == call.action)
215 .ok_or_else(|| {
216 ProviderError::validation(
217 &self.catalog.provider.name,
218 &call.action,
219 "unknown_python_action",
220 format!("Python provider has no action `{}`", call.action),
221 )
222 })
223 }
224}
225
226pub fn load_python_catalog(path: &Path, env_prefix: &str) -> Result<ProviderCatalog, String> {
233 let runtime = PythonRuntime::for_catalog(env_prefix);
234 let input = serde_json::to_vec(&json!({
235 "mode": "catalog",
236 "path": path,
237 }))
238 .map_err(|error| error.to_string())?;
239 let output = run_catalog_sidecar(&runtime, &input)?;
240 let value: Value = serde_json::from_slice(&output).map_err(|error| error.to_string())?;
241 validate_provider_manifest_value(&value).map_err(|error| error.to_string())
242}
243
244struct PythonRuntime {
245 command: String,
246 env: Vec<(String, String)>,
247 timeout_ms: u64,
248 max_input_bytes: usize,
249 max_output_bytes: usize,
250}
251
252impl PythonRuntime {
253 fn for_catalog(env_prefix: &str) -> Self {
254 let prefix = env_prefix.trim_matches('_').to_ascii_uppercase();
255 let timeout_var = format!("{prefix}_PYTHON_CATALOG_TIMEOUT_MS");
256 let timeout_ms = match std::env::var(&timeout_var) {
257 Ok(value) => value.parse().unwrap_or_else(|error| {
258 tracing::warn!(
259 variable = %timeout_var,
260 value,
261 error = %error,
262 "invalid provider catalog timeout env var; falling back to the default"
263 );
264 DEFAULT_TIMEOUT_MS
265 }),
266 Err(_) => DEFAULT_TIMEOUT_MS,
267 };
268 Self {
269 command: std::env::var(format!("{prefix}_PYTHON_COMMAND"))
270 .unwrap_or_else(|_| default_python_command().to_owned()),
271 env: Vec::new(),
272 timeout_ms,
273 max_input_bytes: DEFAULT_MAX_INPUT_BYTES,
274 max_output_bytes: DEFAULT_MAX_OUTPUT_BYTES,
275 }
276 }
277
278 fn from_tool(
279 catalog: &ProviderCatalog,
280 tool: &ProviderTool,
281 call: &ProviderCall,
282 env_prefix: &str,
283 ) -> Result<Self, ProviderError> {
284 let provider_meta = catalog.meta.get("python");
285 let tool_meta = tool.meta.get("python");
286 let meta_field = |key: &str| {
287 tool_meta
288 .and_then(|value| value.get(key))
289 .or_else(|| provider_meta.and_then(|value| value.get(key)))
290 };
291 let command = meta_field("command")
292 .and_then(Value::as_str)
293 .map(str::to_owned)
294 .or_else(|| {
295 std::env::var(format!(
296 "{}_PYTHON_COMMAND",
297 env_prefix.trim_matches('_').to_ascii_uppercase()
298 ))
299 .ok()
300 })
301 .unwrap_or_else(|| default_python_command().to_owned());
302 let timeout_ms = tool
303 .limits
304 .as_ref()
305 .and_then(|limits| limits.timeout_ms)
306 .or_else(|| meta_field("timeout_ms").and_then(Value::as_u64))
307 .unwrap_or(DEFAULT_TIMEOUT_MS);
308 let max_input_bytes = tool
309 .limits
310 .as_ref()
311 .and_then(|limits| limits.max_input_bytes)
312 .unwrap_or(DEFAULT_MAX_INPUT_BYTES);
313 let max_output_bytes = tool
314 .limits
315 .as_ref()
316 .and_then(|limits| limits.max_response_bytes)
317 .unwrap_or(DEFAULT_MAX_OUTPUT_BYTES);
318 Ok(Self {
319 command,
320 env: collect_provider_env(
321 &catalog.env,
322 &tool.env,
323 env_prefix,
324 &call.provider,
325 &call.action,
326 )?,
327 timeout_ms,
328 max_input_bytes,
329 max_output_bytes,
330 })
331 }
332}
333
334#[cfg(windows)]
335fn default_python_command() -> &'static str {
336 "python"
337}
338
339#[cfg(not(windows))]
340fn default_python_command() -> &'static str {
341 "python3"
342}
343
344fn run_catalog_sidecar(runtime: &PythonRuntime, input: &[u8]) -> Result<Vec<u8>, String> {
345 let mut command = StdCommand::new(resolve_sidecar_command(&runtime.command));
346 command
347 .args(["-c", PYTHON_BRIDGE])
348 .env_clear()
349 .stdin(StdStdio::piped())
350 .stdout(StdStdio::piped())
351 .stderr(StdStdio::piped());
352 for (key, value) in sidecar_base_env() {
353 command.env(key, value);
354 }
355 let mut child = command.spawn().map_err(|error| error.to_string())?;
356 let stdout = child
357 .stdout
358 .take()
359 .ok_or_else(|| "Python provider catalog stdout pipe was not captured".to_owned())?;
360 let stderr = child
361 .stderr
362 .take()
363 .ok_or_else(|| "Python provider catalog stderr pipe was not captured".to_owned())?;
364 let max_output_bytes = runtime.max_output_bytes;
365 let stdout_task = thread::spawn(move || read_bounded_sync(stdout, max_output_bytes));
366 let stderr_task = thread::spawn(move || read_bounded_sync(stderr, max_output_bytes));
367
368 if let Some(mut stdin) = child.stdin.take() {
369 stdin.write_all(input).map_err(|error| error.to_string())?;
370 }
371 let deadline = Instant::now() + Duration::from_millis(runtime.timeout_ms);
372 loop {
373 if let Some(status) = child.try_wait().map_err(|error| error.to_string())? {
374 let (stdout, stdout_exceeded) = stdout_task
375 .join()
376 .map_err(|_| "Python provider catalog stdout reader panicked".to_owned())?
377 .map_err(|error| error.to_string())?;
378 let (stderr, stderr_exceeded) = stderr_task
379 .join()
380 .map_err(|_| "Python provider catalog stderr reader panicked".to_owned())?
381 .map_err(|error| error.to_string())?;
382 if stdout_exceeded || stderr_exceeded {
383 let stream = if stdout_exceeded { "stdout" } else { "stderr" };
384 return Err(format!(
385 "Python provider catalog {}",
386 output_exceeded_message(stream, runtime.max_output_bytes)
387 ));
388 }
389 if !status.success() {
390 return Err(format!(
391 "Python provider catalog failed: {}",
392 redact_public(&String::from_utf8_lossy(&stderr))
393 ));
394 }
395 return Ok(stdout);
396 }
397 if Instant::now() >= deadline {
398 let _ = child.kill();
399 let _ = child.wait();
400 return Err(format!(
401 "Python provider catalog exceeded {}ms timeout",
402 runtime.timeout_ms
403 ));
404 }
405 std::thread::sleep(Duration::from_millis(10));
406 }
407}
408
409fn read_bounded_sync<R: Read>(
410 mut reader: R,
411 max_output_bytes: usize,
412) -> std::io::Result<(Vec<u8>, bool)> {
413 let mut bytes = Vec::new();
414 let mut exceeded = false;
415 let mut chunk = [0u8; 8192];
416 loop {
417 let read = reader.read(&mut chunk)?;
418 if read == 0 {
419 return Ok((bytes, exceeded));
420 }
421 let remaining = max_output_bytes.saturating_sub(bytes.len());
422 if remaining >= read && !exceeded {
423 bytes.extend_from_slice(&chunk[..read]);
424 } else {
425 exceeded = true;
426 if remaining > 0 {
427 bytes.extend_from_slice(&chunk[..remaining]);
428 }
429 }
430 }
431}
432
433#[cfg(test)]
434#[path = "python_tests.rs"]
435mod tests;