telos_agent/tools/builtin/
python.rs

1use std::path::{Path, PathBuf};
2
3use async_trait::async_trait;
4use serde_json::{Value, json};
5use tokio::io::AsyncReadExt;
6use tokio::process::Command;
7use tokio::time::{Duration, timeout};
8
9use crate::error::AgentError;
10use crate::tools::api::{PermissionDecision, Tool, ToolContext, ToolDefinition, ToolOutput};
11
12use super::{display_relative, optional_usize_any, required_string};
13
14const DEFAULT_TIMEOUT_MS: u64 = 120_000;
15const MAX_OUTPUT_CHARS: usize = 20_000;
16
17pub struct PythonScriptTool;
18
19#[async_trait]
20impl Tool for PythonScriptTool {
21    fn definition(&self) -> ToolDefinition {
22        ToolDefinition {
23            name: "PythonScript".into(),
24            description:
25                "Run a Python script with Telos' isolated bundled Python runtime. Optional packages install into .telos/python/packages, never into the system Python."
26                    .into(),
27            input_schema: json!({
28                "type": "object",
29                "properties": {
30                    "script": { "type": "string" },
31                    "packages": {
32                        "type": "array",
33                        "items": { "type": "string" },
34                        "description": "Optional pip package specifiers to install into the workspace-local Telos Python package directory before running."
35                    },
36                    "timeout_ms": {
37                        "type": "integer",
38                        "description": "Maximum runtime in milliseconds. Defaults to 120000."
39                    },
40                    "artifacts": {
41                        "type": "array",
42                        "items": { "type": "string" },
43                        "description": "Workspace-relative output files expected from the script."
44                    }
45                },
46                "required": ["script"]
47            }),
48        }
49    }
50
51    fn prompt_text(&self) -> Option<&'static str> {
52        Some(
53            "Use PythonScript for data processing, office documents, web parsing, and Playwright automation when Python libraries are helpful. \
54Keep scripts self-contained, write outputs under the workspace, and list expected output files in `artifacts`. \
55Only request `packages` when the bundled runtime does not already provide the dependency.",
56        )
57    }
58
59    async fn validate(&self, arguments: &Value, _context: &ToolContext) -> Result<(), AgentError> {
60        required_string(arguments, "script")?;
61        package_specs(arguments)?;
62        string_array(arguments, "artifacts")?;
63        Ok(())
64    }
65
66    async fn check_permission(
67        &self,
68        arguments: &Value,
69        _context: &ToolContext,
70    ) -> Result<PermissionDecision, AgentError> {
71        let packages = package_specs(arguments)?;
72        let reason = if packages.is_empty() {
73            "Python script execution requires approval".to_string()
74        } else {
75            format!(
76                "Python script execution requests workspace-local package installation: {}",
77                packages.join(", ")
78            )
79        };
80        Ok(PermissionDecision::Ask { reason })
81    }
82
83    async fn invoke(
84        &self,
85        arguments: Value,
86        context: ToolContext,
87    ) -> Result<ToolOutput, AgentError> {
88        let runtime = PythonRuntime::resolve(&context)?;
89        let script = required_string(&arguments, "script")?;
90        let packages = package_specs(&arguments)?;
91        let artifacts = string_array(&arguments, "artifacts")?.unwrap_or_default();
92        let timeout_ms = optional_usize_any(&arguments, &["timeout_ms"])
93            .unwrap_or(DEFAULT_TIMEOUT_MS as usize) as u64;
94
95        let run_dir = context
96            .cwd
97            .join(".telos")
98            .join("python")
99            .join("runs")
100            .join(format!("{}-{}", context.session_id, context.turn_id));
101        tokio::fs::create_dir_all(&run_dir).await.map_err(python_io_error)?;
102        let script_path = run_dir.join("script.py");
103        tokio::fs::write(&script_path, script).await.map_err(python_io_error)?;
104
105        let package_dir = context.cwd.join(".telos").join("python").join("packages");
106        if !packages.is_empty() {
107            tokio::fs::create_dir_all(&package_dir).await.map_err(python_io_error)?;
108            emit_progress(
109                &context,
110                "installing Python packages into workspace",
111                json!({ "packages": packages.clone(), "target": display_relative(&context.cwd, &package_dir) }),
112            );
113            let install_output =
114                runtime.run_pip_install(&context, &packages, &package_dir, timeout_ms).await?;
115            if !install_output.status.success() {
116                return Err(AgentError::ToolExecution {
117                    tool: "PythonScript".into(),
118                    message: format_process_failure("pip install", &install_output),
119                });
120            }
121        }
122
123        emit_progress(
124            &context,
125            "running Python script",
126            json!({
127                "python": runtime.python_exe,
128                "script": display_relative(&context.cwd, &script_path),
129                "timeout_ms": timeout_ms
130            }),
131        );
132        let output = runtime.run_script(&context, &script_path, &package_dir, timeout_ms).await?;
133
134        let stdout = String::from_utf8_lossy(&output.stdout).to_string();
135        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
136        if !output.status.success() {
137            return Err(AgentError::ToolExecution {
138                tool: "PythonScript".into(),
139                message: format_process_failure("python script", &output),
140            });
141        }
142
143        Ok(ToolOutput::json(json!({
144            "success": true,
145            "status": output.status.code(),
146            "python": runtime.python_exe,
147            "script_path": script_path,
148            "script_relative_path": display_relative(&context.cwd, &script_path),
149            "stdout": trim_large_output(&stdout),
150            "stderr": trim_large_output(&stderr),
151            "artifacts": collect_artifacts(&context.cwd, &artifacts).await?,
152        })))
153    }
154}
155
156#[derive(Debug, Clone)]
157struct PythonRuntime {
158    python_exe: PathBuf,
159    python_home: Option<PathBuf>,
160    playwright_browsers: Option<PathBuf>,
161}
162
163impl PythonRuntime {
164    fn resolve(context: &ToolContext) -> Result<Self, AgentError> {
165        if let Some(path) = context.env.get("TELOS_PYTHON_EXE").filter(|value| !value.is_empty()) {
166            let python_exe = PathBuf::from(path);
167            if python_exe.exists() {
168                return Ok(Self {
169                    python_exe,
170                    python_home: env_path(&context.env, "TELOS_PYTHON_HOME"),
171                    playwright_browsers: env_path(&context.env, "PLAYWRIGHT_BROWSERS_PATH")
172                        .or_else(|| env_path(&context.env, "TELOS_PLAYWRIGHT_BROWSERS_PATH")),
173                });
174            }
175        }
176
177        for root in runtime_roots(context) {
178            if let Some(runtime) = Self::from_root(&root) {
179                return Ok(runtime);
180            }
181        }
182
183        Err(AgentError::Config(
184            "bundled Python runtime not found. Set TELOS_PYTHON_EXE or bundle resources/runtimes/python-<platform>/.".into(),
185        ))
186    }
187
188    fn from_root(root: &Path) -> Option<Self> {
189        let platform = runtime_platform();
190        let candidates = [
191            root.join("runtimes").join(format!("python-{platform}")),
192            root.join("runtimes").join("python").join(platform),
193            root.join("python"),
194        ];
195        for python_home in candidates {
196            let python_exe = python_home.join(python_executable_name());
197            if python_exe.exists() {
198                let playwright_browsers =
199                    [root.join("playwright-browsers"), python_home.join("playwright-browsers")]
200                        .into_iter()
201                        .find(|path| path.exists());
202                return Some(Self { python_exe, python_home: None, playwright_browsers });
203            }
204        }
205        None
206    }
207
208    async fn run_pip_install(
209        &self,
210        context: &ToolContext,
211        packages: &[String],
212        package_dir: &Path,
213        timeout_ms: u64,
214    ) -> Result<std::process::Output, AgentError> {
215        let mut command = self.base_command(context, Some(package_dir));
216        command
217            .arg("-m")
218            .arg("pip")
219            .arg("install")
220            .arg("--disable-pip-version-check")
221            .arg("--target")
222            .arg(package_dir)
223            .args(packages)
224            .current_dir(&context.cwd);
225        run_child("PythonScript", command, timeout_ms).await
226    }
227
228    async fn run_script(
229        &self,
230        context: &ToolContext,
231        script_path: &Path,
232        package_dir: &Path,
233        timeout_ms: u64,
234    ) -> Result<std::process::Output, AgentError> {
235        let mut command = self.base_command(context, Some(package_dir));
236        command.arg(script_path).current_dir(&context.cwd);
237        run_child("PythonScript", command, timeout_ms).await
238    }
239
240    fn base_command(&self, context: &ToolContext, package_dir: Option<&Path>) -> Command {
241        let mut command = Command::new(&self.python_exe);
242        command.env_clear().envs(context.env.iter());
243        command.env("PYTHONNOUSERSITE", "1");
244        command.env("PIP_DISABLE_PIP_VERSION_CHECK", "1");
245        if let Some(home) = &self.python_home {
246            command.env("PYTHONHOME", home);
247        }
248        if let Some(path) = &self.playwright_browsers {
249            command.env("PLAYWRIGHT_BROWSERS_PATH", path);
250        }
251        if let Some(package_dir) = package_dir {
252            let mut pythonpath = package_dir.as_os_str().to_os_string();
253            if let Some(existing) = context.env.get("PYTHONPATH").filter(|value| !value.is_empty())
254            {
255                pythonpath.push(if cfg!(windows) { ";" } else { ":" });
256                pythonpath.push(existing);
257            }
258            command.env("PYTHONPATH", pythonpath);
259        }
260        hide_console_window(&mut command);
261        command.kill_on_drop(true);
262        command
263    }
264}
265
266fn runtime_roots(context: &ToolContext) -> Vec<PathBuf> {
267    let mut roots = Vec::new();
268    if let Some(path) = env_path(&context.env, "TELOS_RUNTIME_DIR") {
269        roots.push(path);
270    }
271    roots.push(context.cwd.join("desktop").join("src-tauri").join("resources"));
272    roots.push(context.cwd.join("resources"));
273    if let Ok(exe) = std::env::current_exe()
274        && let Some(parent) = exe.parent()
275    {
276        roots.push(parent.join("resources"));
277        roots.push(parent.join("../Resources").join("resources"));
278        roots.push(parent.join("../../Resources").join("resources"));
279    }
280    roots
281}
282
283fn env_path(env: &std::collections::HashMap<String, String>, key: &str) -> Option<PathBuf> {
284    env.get(key).filter(|value| !value.trim().is_empty()).map(PathBuf::from)
285}
286
287fn runtime_platform() -> &'static str {
288    match (std::env::consts::OS, std::env::consts::ARCH) {
289        ("windows", "x86_64") => "windows-x64",
290        ("windows", "aarch64") => "windows-arm64",
291        ("macos", "x86_64") => "macos-x64",
292        ("macos", "aarch64") => "macos-arm64",
293        ("linux", "x86_64") => "linux-x64",
294        ("linux", "aarch64") => "linux-arm64",
295        _ => "unknown",
296    }
297}
298
299fn python_executable_name() -> &'static str {
300    if cfg!(windows) { "python.exe" } else { "bin/python3" }
301}
302
303async fn run_child(
304    tool: &str,
305    mut command: Command,
306    timeout_ms: u64,
307) -> Result<std::process::Output, AgentError> {
308    command.stdout(std::process::Stdio::piped()).stderr(std::process::Stdio::piped());
309    let mut child = command
310        .spawn()
311        .map_err(|err| AgentError::ToolExecution { tool: tool.into(), message: err.to_string() })?;
312    let mut stdout = child.stdout.take().expect("stdout was piped");
313    let mut stderr = child.stderr.take().expect("stderr was piped");
314
315    let stdout_task = tokio::spawn(async move {
316        let mut buf = Vec::new();
317        stdout.read_to_end(&mut buf).await.map(|_| buf)
318    });
319    let stderr_task = tokio::spawn(async move {
320        let mut buf = Vec::new();
321        stderr.read_to_end(&mut buf).await.map(|_| buf)
322    });
323
324    let status = timeout(Duration::from_millis(timeout_ms.max(1)), child.wait())
325        .await
326        .map_err(|_| AgentError::ToolExecution {
327            tool: tool.into(),
328            message: format!("Python command timed out after {timeout_ms}ms"),
329        })?
330        .map_err(|err| AgentError::ToolExecution { tool: tool.into(), message: err.to_string() })?;
331    let stdout = stdout_task
332        .await
333        .map_err(|err| AgentError::ToolExecution { tool: tool.into(), message: err.to_string() })?
334        .map_err(|err| AgentError::ToolExecution { tool: tool.into(), message: err.to_string() })?;
335    let stderr = stderr_task
336        .await
337        .map_err(|err| AgentError::ToolExecution { tool: tool.into(), message: err.to_string() })?
338        .map_err(|err| AgentError::ToolExecution { tool: tool.into(), message: err.to_string() })?;
339    Ok(std::process::Output { status, stdout, stderr })
340}
341
342async fn collect_artifacts(cwd: &Path, artifacts: &[String]) -> Result<Vec<Value>, AgentError> {
343    let mut values = Vec::new();
344    for artifact in artifacts {
345        let path = super::resolve_workspace_path(cwd, artifact)?;
346        let path = super::canonicalize_within_cwd(cwd, &path).await?;
347        let metadata = tokio::fs::metadata(&path).await.map_err(python_io_error)?;
348        values.push(json!({
349            "path": path,
350            "relative_path": display_relative(cwd, &path),
351            "bytes": metadata.len(),
352        }));
353    }
354    Ok(values)
355}
356
357fn package_specs(arguments: &Value) -> Result<Vec<String>, AgentError> {
358    let Some(packages) = string_array(arguments, "packages")? else {
359        return Ok(Vec::new());
360    };
361    for package in &packages {
362        if package.trim().is_empty() || package.contains(['\0', '\n', '\r']) {
363            return Err(AgentError::Validation("invalid Python package specifier".into()));
364        }
365    }
366    Ok(packages)
367}
368
369fn string_array(arguments: &Value, key: &str) -> Result<Option<Vec<String>>, AgentError> {
370    let Some(value) = arguments.get(key) else {
371        return Ok(None);
372    };
373    let Some(array) = value.as_array() else {
374        return Err(AgentError::Validation(format!("`{key}` must be an array of strings")));
375    };
376    let mut values = Vec::with_capacity(array.len());
377    for item in array {
378        let Some(value) = item.as_str() else {
379            return Err(AgentError::Validation(format!("`{key}` must be an array of strings")));
380        };
381        values.push(value.to_string());
382    }
383    Ok(Some(values))
384}
385
386fn emit_progress(context: &ToolContext, message: &str, data: Value) {
387    if let Some(tx) = &context.progress {
388        let _ = tx.send(crate::tools::api::ToolProgress {
389            tool_call_id: context.tool_call_id.clone(),
390            message: message.into(),
391            data: Some(data),
392        });
393    }
394}
395
396fn format_process_failure(action: &str, output: &std::process::Output) -> String {
397    let stdout = String::from_utf8_lossy(&output.stdout);
398    let stderr = String::from_utf8_lossy(&output.stderr);
399    format!(
400        "{action} failed with exit code {:?}\nstdout:\n{}\nstderr:\n{}",
401        output.status.code(),
402        trim_large_output(&stdout),
403        trim_large_output(&stderr)
404    )
405}
406
407fn trim_large_output(output: &str) -> String {
408    if output.chars().count() <= MAX_OUTPUT_CHARS {
409        return output.to_string();
410    }
411    let preview = output.chars().take(MAX_OUTPUT_CHARS).collect::<String>();
412    format!("{preview}\n<truncated output after {MAX_OUTPUT_CHARS} chars>")
413}
414
415fn python_io_error(err: std::io::Error) -> AgentError {
416    AgentError::ToolExecution { tool: "PythonScript".into(), message: err.to_string() }
417}
418
419#[cfg_attr(not(windows), allow(unused_variables))]
420fn hide_console_window(command: &mut Command) {
421    #[cfg(windows)]
422    {
423        const CREATE_NO_WINDOW: u32 = 0x0800_0000;
424        command.creation_flags(CREATE_NO_WINDOW);
425    }
426}
427
428#[cfg(test)]
429mod tests {
430    use super::*;
431    use serde_json::json;
432
433    #[test]
434    fn package_specs_rejects_control_characters() {
435        let result = package_specs(&json!({ "packages": ["ok", "bad\npkg"] }));
436        assert!(result.is_err());
437    }
438
439    #[test]
440    fn runtime_platform_is_known_on_supported_targets() {
441        assert_ne!(runtime_platform(), "unknown");
442    }
443}
telos_agent/tools/builtin/python.rs

telos_agent/tools/builtin/
python.rs
