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

byob_plugins/
python.rs
