Skip to main content

byob_plugins/
python.rs

1//! A plugin's declared Python tools (`tools[].run`): the model's arguments go to the
2//! script as JSON on stdin; what it prints (plain text, or `{"text": …, "files": […]}`)
3//! goes back. Settings arrive as environment variables, plus `HOB_CHAT_DIR` and
4//! `HOB_PLUGIN_DIR` (and `BYOB_*` aliases).
5//!
6//! Given a [`Sandbox`] ([`PythonTool::with_sandbox`]), scripts run confined by
7//! `byob-sandbox` ([`run_script_sandboxed`]: workspace + temp writable, plugin folder
8//! readable, network only if the plugin declares `permissions.network`). Without one
9//! they're a plain child process ([`run_script_unsandboxed`]). Either way: timeout,
10//! bounded output, process-group kill, and [`Access::Execute`] so the permission gate asks.
11
12use std::collections::BTreeMap;
13use std::ffi::OsString;
14use std::fs;
15use std::path::{Path, PathBuf};
16use std::process::Command;
17use std::sync::Arc;
18use std::time::{Duration, SystemTime};
19
20use byob_sandbox::{Limits, Policy, RunOptions, Sandbox, TempDir};
21use byob_tools::{Access, FileChange, Tool, ToolContext, ToolOutput, ToolSpec};
22use serde_json::Value;
23
24use crate::PluginRef;
25use crate::discover::Plugin;
26use crate::manifest::{ManifestTool, ToolKind, path_inside};
27use crate::names::{ToolNames, namespaced};
28
29/// hob's limit for a plugin tool.
30pub const SCRIPT_TIMEOUT: Duration = Duration::from_secs(60);
31/// Most bytes kept from each of stdout and stderr.
32pub const MAX_SCRIPT_OUTPUT: usize = 256 * 1024;
33/// Most files compared before/after a run to spot new ones.
34const MAX_SCANNED_FILES: usize = 5000;
35
36/// The interpreter a Python tool runs with, as a command prefix.
37pub fn default_python() -> Vec<OsString> {
38    vec![OsString::from(if cfg!(windows) { "python" } else { "python3" })]
39}
40
41/// One script run.
42#[derive(Debug, Clone)]
43pub struct ScriptRun {
44    /// Interpreter and its leading arguments, e.g. `["python3"]`; the script path follows.
45    pub program: Vec<OsString>,
46    pub script: PathBuf,
47    /// Extra arguments after the script.
48    pub args: Vec<OsString>,
49    pub cwd: PathBuf,
50    pub stdin: Vec<u8>,
51    /// Added to the inherited environment.
52    pub env: BTreeMap<String, String>,
53    pub timeout: Duration,
54    /// Per stream; the rest is read and dropped.
55    pub max_output: usize,
56}
57
58#[derive(Debug, Clone, Default, PartialEq, Eq)]
59pub struct ScriptOutcome {
60    pub stdout: String,
61    pub stderr: String,
62    /// `None` when killed or ended by a signal.
63    pub exit_code: Option<i32>,
64    pub timed_out: bool,
65    pub stopped: bool,
66    /// Some output was dropped for being over `max_output`.
67    pub truncated: bool,
68}
69
70impl ScriptOutcome {
71    pub fn success(&self) -> bool {
72        self.exit_code == Some(0) && !self.timed_out && !self.stopped
73    }
74}
75
76/// Run a script as a plain child process: **no sandbox**. Kills it (and its process
77/// group) on timeout or when `cancelled()` turns true. `Err` only when it couldn't start.
78pub fn run_script_unsandboxed(run: &ScriptRun, cancelled: &dyn Fn() -> bool) -> Result<ScriptOutcome, String> {
79    let Some((program, lead)) = run.program.split_first() else { return Err("no interpreter configured".into()) };
80    let mut cmd = Command::new(program);
81    cmd.args(lead).arg(&run.script).args(&run.args).current_dir(&run.cwd).envs(&run.env);
82    execute(cmd, run, cancelled)
83}
84
85/// Run a script in `sandbox`: read the system, the plugin folder and the interpreter;
86/// read+write `run.cwd` (the workspace) and a private temp folder; network only if
87/// `network`. The environment is reset to a minimal one plus `run.env`.
88pub fn run_script_sandboxed(
89    run: &ScriptRun,
90    sandbox: &Sandbox,
91    plugin_dir: &Path,
92    network: bool,
93    cancelled: &dyn Fn() -> bool,
94) -> Result<ScriptOutcome, String> {
95    let Some((program, lead)) = run.program.split_first() else { return Err("no interpreter configured".into()) };
96    let cwd = run.cwd.canonicalize().map_err(|e| format!("{}: {e}", run.cwd.display()))?;
97    let tmp = TempDir::new("byob-plugin-")?;
98    let mut policy = Policy::for_workspace(&cwd, tmp.path())
99        .network(network)
100        .limits(Limits::for_timeout(run.timeout.as_secs()))
101        .allow_read(plugin_dir)
102        .allow_path_dirs(std::env::var_os("PATH").as_deref())
103        // Where the OS can't confine, run anyway (as before the sandbox existed).
104        .allow_unsandboxed(true);
105    for p in byob_sandbox::interpreter_read_paths(program) {
106        policy = policy.allow_read(p);
107    }
108    let mut args: Vec<OsString> = lead.to_vec();
109    args.push(run.script.clone().into_os_string());
110    args.extend(run.args.iter().cloned());
111    let mut cmd = sandbox.command(&policy, program, &args);
112    cmd.current_dir(&cwd);
113    byob_sandbox::clean_env(&mut cmd, tmp.path());
114    cmd.envs(&run.env);
115    let out = execute(cmd, run, cancelled)?;
116    drop(tmp);
117    Ok(out)
118}
119
120fn execute(cmd: Command, run: &ScriptRun, cancelled: &dyn Fn() -> bool) -> Result<ScriptOutcome, String> {
121    let opts = RunOptions { stdin: run.stdin.clone(), timeout: run.timeout, max_output: run.max_output };
122    let o = byob_sandbox::run(cmd, &opts, cancelled)?;
123    if let Some(e) = o.sandbox_error() {
124        return Err(format!("the sandbox couldn't start: {e}"));
125    }
126    Ok(ScriptOutcome {
127        stdout: String::from_utf8_lossy(&o.stdout.bytes()).into_owned(),
128        stderr: String::from_utf8_lossy(&o.stderr.bytes()).into_owned(),
129        exit_code: o.exit_code,
130        timed_out: o.timed_out,
131        stopped: o.cancelled,
132        truncated: o.stdout.truncated() || o.stderr.truncated(),
133    })
134}
135
136/// `"Converted {amount} {from}"` + args → `"Converted 5 EUR"`; unknown keys stay.
137pub fn fill_title(template: &str, args: &Value) -> String {
138    let mut out = String::new();
139    let mut rest = template;
140    while let Some(open) = rest.find('{') {
141        out.push_str(&rest[..open]);
142        let after = &rest[open + 1..];
143        let key_len = after.find(|c: char| !(c.is_alphanumeric() || c == '_')).unwrap_or(after.len());
144        if key_len > 0 && after[key_len..].starts_with('}') {
145            let key = &after[..key_len];
146            match args.get(key) {
147                Some(Value::String(s)) => out.push_str(s),
148                Some(v) => out.push_str(&v.to_string()),
149                None => out.push_str(&rest[open..open + key_len + 2]),
150            }
151            rest = &after[key_len + 1..];
152        } else {
153            out.push('{');
154            rest = after;
155        }
156    }
157    out.push_str(rest);
158    out.chars().take(120).collect()
159}
160
161/// Files under `dir` (two levels down, dot files skipped) with their modification times.
162fn snapshot(dir: &Path) -> BTreeMap<PathBuf, Option<SystemTime>> {
163    let mut m = BTreeMap::new();
164    let mut stack = vec![(dir.to_path_buf(), 0u8)];
165    while let Some((d, depth)) = stack.pop() {
166        let Ok(rd) = fs::read_dir(&d) else { continue };
167        for e in rd.flatten() {
168            if m.len() >= MAX_SCANNED_FILES {
169                return m;
170            }
171            if e.file_name().to_string_lossy().starts_with('.') {
172                continue;
173            }
174            let Ok(meta) = e.metadata() else { continue };
175            if meta.is_dir() {
176                if depth < 2 {
177                    stack.push((e.path(), depth + 1));
178                }
179            } else {
180                m.insert(e.path(), meta.modified().ok());
181            }
182        }
183    }
184    m
185}
186
187/// A plugin's Python tool.
188pub struct PythonTool {
189    name: String,
190    plugin: PluginRef,
191    decl: ManifestTool,
192    script: String,
193    interpreter: Vec<OsString>,
194    sandbox: Option<Sandbox>,
195    timeout: Duration,
196}
197
198impl PythonTool {
199    /// `decl` must be a script tool ([`ToolKind::Script`]); `None` otherwise.
200    pub fn new(name: String, plugin: PluginRef, decl: ManifestTool, interpreter: Vec<OsString>) -> Option<Self> {
201        let ToolKind::Script(script) = &decl.kind else { return None };
202        let script = script.clone();
203        Some(Self { name, plugin, decl, script, interpreter, sandbox: None, timeout: SCRIPT_TIMEOUT })
204    }
205
206    /// Run the script in `sandbox` (see [`run_script_sandboxed`]); `None`: unsandboxed.
207    pub fn with_sandbox(mut self, sandbox: Option<Sandbox>) -> Self {
208        self.sandbox = sandbox;
209        self
210    }
211
212    pub fn with_timeout(mut self, timeout: Duration) -> Self {
213        self.timeout = timeout;
214        self
215    }
216
217    pub fn plugin(&self) -> &PluginRef {
218        &self.plugin
219    }
220
221    fn finish(&self, r: &ScriptOutcome, workspace: &Path, before: &BTreeMap<PathBuf, Option<SystemTime>>) -> ToolOutput {
222        let name = &self.plugin.name;
223        if r.stopped {
224            return ToolOutput::err("Stopped by the user.");
225        }
226        if r.timed_out {
227            return ToolOutput::err(format!("{name} took longer than {} seconds and was stopped.", self.timeout.as_secs()));
228        }
229        if r.exit_code != Some(0) {
230            let output = [r.stdout.trim(), r.stderr.trim()].iter().filter(|s| !s.is_empty()).copied().collect::<Vec<_>>().join("\n");
231            return ToolOutput::err(format!("{name} ran into a problem:\n{output}"));
232        }
233        let out = r.stdout.trim();
234        let mut text = out.to_string();
235        let mut listed: Vec<PathBuf> = Vec::new();
236        if out.starts_with('{')
237            && let Ok(Value::Object(j)) = serde_json::from_str::<Value>(out)
238        {
239            if let Some(t) = j.get("text").and_then(Value::as_str) {
240                text = t.to_string();
241            }
242            if let Some(files) = j.get("files").and_then(Value::as_array) {
243                listed = files.iter().filter_map(Value::as_str).filter(|f| path_inside(f)).map(|f| workspace.join(f)).filter(|p| p.is_file()).collect();
244            }
245        }
246        if r.truncated {
247            text.push_str("\n…(output cut: it was too long)");
248        }
249        let after = snapshot(workspace);
250        let mut changes: Vec<FileChange> = after
251            .iter()
252            .filter(|(f, m)| before.get(*f) != Some(*m))
253            .map(|(f, _)| if before.contains_key(f) { FileChange::Modified(f.clone()) } else { FileChange::Created(f.clone()) })
254            .collect();
255        for f in listed {
256            if !changes.iter().any(|c| matches!(c, FileChange::Created(p) | FileChange::Modified(p) if *p == f)) {
257                changes.push(FileChange::Modified(f));
258            }
259        }
260        if text.is_empty() {
261            text = "(no result)".into();
262        }
263        ToolOutput { text, failed: false, changes }
264    }
265}
266
267impl Tool for PythonTool {
268    fn spec(&self) -> ToolSpec {
269        ToolSpec { name: self.name.clone(), description: self.decl.description.clone(), parameters: self.decl.params.clone() }
270    }
271
272    fn title(&self, args: &Value) -> String {
273        match &self.decl.title {
274            Some(t) => fill_title(t, args),
275            None => format!("Used {}", self.plugin.name),
276        }
277    }
278
279    fn access(&self, _args: &Value) -> Access {
280        Access::Execute
281    }
282
283    fn run(&self, args: &Value, ctx: &ToolContext<'_>) -> ToolOutput {
284        if ctx.cancelled() {
285            return ToolOutput::err("Stopped by the user.");
286        }
287        let script = self.plugin.dir.join(&self.script);
288        if !path_inside(&self.script) || !script.is_file() {
289            return ToolOutput::err(format!("{}: {} is missing", self.plugin.name, self.script));
290        }
291        let args = if args.is_object() { args.clone() } else { Value::Object(serde_json::Map::new()) };
292        let chat = ctx.workspace.to_string_lossy().into_owned();
293        let dir = self.plugin.dir.to_string_lossy().into_owned();
294        let mut env = self.plugin.env.clone();
295        for (k, v) in [("HOB_CHAT_DIR", &chat), ("BYOB_CHAT_DIR", &chat), ("HOB_PLUGIN_DIR", &dir), ("BYOB_PLUGIN_DIR", &dir)] {
296            env.insert(k.into(), v.clone());
297        }
298        let run = ScriptRun {
299            program: self.interpreter.clone(),
300            script,
301            args: Vec::new(),
302            cwd: ctx.workspace.to_path_buf(),
303            stdin: serde_json::to_vec(&args).unwrap_or_else(|_| b"{}".to_vec()),
304            env,
305            timeout: self.timeout,
306            max_output: MAX_SCRIPT_OUTPUT,
307        };
308        let before = snapshot(ctx.workspace);
309        let result = match &self.sandbox {
310            Some(sb) => run_script_sandboxed(&run, sb, &self.plugin.dir, self.plugin.network, &|| ctx.cancelled()),
311            None => run_script_unsandboxed(&run, &|| ctx.cancelled()),
312        };
313        match result {
314            Ok(r) => self.finish(&r, ctx.workspace, &before),
315            Err(e) => ToolOutput::err(format!("{} couldn't run: {e}", self.plugin.name)),
316        }
317    }
318}
319
320/// The Python tools `plugin` declares, named as declared (so the plugin's `prompt` still
321/// matches) or `<plugin>__<tool>` when that name is taken. Tools whose script is missing
322/// are left out (the plugin's `problems` say so).
323pub fn python_tools(plugin: &Plugin, settings: &BTreeMap<String, String>, interpreter: &[OsString], names: &mut ToolNames) -> Vec<Arc<PythonTool>> {
324    python_tools_in(plugin, settings, interpreter, names, None)
325}
326
327/// [`python_tools`] whose scripts run in `sandbox` (when given).
328pub fn python_tools_in(
329    plugin: &Plugin,
330    settings: &BTreeMap<String, String>,
331    interpreter: &[OsString],
332    names: &mut ToolNames,
333    sandbox: Option<&Sandbox>,
334) -> Vec<Arc<PythonTool>> {
335    let r = PluginRef::new(plugin, settings);
336    plugin
337        .manifest
338        .tools
339        .iter()
340        .filter(|t| matches!(&t.kind, ToolKind::Script(s) if plugin.dir.join(s).is_file()))
341        .filter_map(|t| {
342            let name = names.claim(&t.name, &namespaced(plugin.id(), &t.name));
343            PythonTool::new(name, r.clone(), t.clone(), interpreter.to_vec()).map(|t| Arc::new(t.with_sandbox(sandbox.cloned())))
344        })
345        .collect()
346}
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351    use serde_json::json;
352
353    #[test]
354    fn titles() {
355        let a = json!({"amount": 5, "from": "EUR", "x": null});
356        assert_eq!(fill_title("Converted {amount} {from} {to} {x}", &a), "Converted 5 EUR {to} null");
357        assert_eq!(fill_title("{ not a key } {", &a), "{ not a key } {");
358        assert_eq!(fill_title(&"é".repeat(200), &a).chars().count(), 120);
359    }
360
361    #[test]
362    fn missing_interpreter_is_an_error() {
363        let run = ScriptRun {
364            program: vec![OsString::from("/definitely/not/an/interpreter")],
365            script: PathBuf::from("x"),
366            args: vec![],
367            cwd: std::env::temp_dir(),
368            stdin: vec![],
369            env: BTreeMap::new(),
370            timeout: SCRIPT_TIMEOUT,
371            max_output: 10,
372        };
373        assert!(run_script_unsandboxed(&run, &|| false).unwrap_err().starts_with("couldn't start"));
374        let empty = ScriptRun { program: vec![], ..run };
375        assert!(run_script_unsandboxed(&empty, &|| false).is_err());
376    }
377}