1use std::collections::BTreeMap;
12use std::ffi::OsString;
13use std::fs;
14use std::hash::Hasher;
15use std::path::{Path, PathBuf};
16use std::time::{Duration, UNIX_EPOCH};
17
18use byob_sandbox::{Sandbox, TempDir};
19use serde_json::{Map, Value, json};
20
21use crate::manifest::{MANIFEST, ParseOptions, parse_manifest_str, plugin_id_ok, setting_key_ok, tool_name_ok};
22use crate::python::{ScriptRun, run_script_sandboxed, run_script_unsandboxed};
23
24pub const RUNNER_MARK: &str = "# --- byob runner: calls run() with the JSON arguments on stdin. Keep this at the end. ---";
26
27const RUNNER: &str = r#"
28if __name__ == "__main__":
29 import inspect as _inspect, json as _json, sys as _sys
30 _raw = _sys.stdin.read().strip()
31 _args = _json.loads(_raw) if _raw else {}
32 if not isinstance(_args, dict):
33 _args = {}
34 _params = _inspect.signature(run).parameters
35 if not any(p.kind == p.VAR_KEYWORD for p in _params.values()):
36 _args = {k: v for k, v in _args.items() if k in _params}
37 _out = run(**_args)
38 if _out is None:
39 pass
40 elif isinstance(_out, str):
41 print(_out)
42 else:
43 print(_json.dumps(_out, ensure_ascii=False, indent=1, default=str))
44"#;
45
46const SCHEMA_SHIM: &str = r#"
48import ast, json, sys
49
50SIMPLE = {"int": "integer", "float": "number", "str": "string", "bool": "boolean",
51 "list": "array", "tuple": "array", "set": "array", "dict": "object",
52 "List": "array", "Tuple": "array", "Set": "array", "Sequence": "array", "Dict": "object", "Mapping": "object"}
53
54def name_of(n):
55 if isinstance(n, ast.Name):
56 return n.id
57 if isinstance(n, ast.Attribute):
58 return n.attr
59 return None
60
61def schema(n):
62 if n is None:
63 return None
64 if isinstance(n, ast.Constant) and isinstance(n.value, str):
65 try:
66 return schema(ast.parse(n.value, mode="eval").body)
67 except SyntaxError:
68 return {"type": "string"}
69 if isinstance(n, ast.Constant) and n.value is None:
70 return {"type": "null"}
71 if isinstance(n, ast.BinOp) and isinstance(n.op, ast.BitOr):
72 parts = [schema(n.left), schema(n.right)]
73 parts = [p for p in parts if p and p.get("type") != "null"]
74 return parts[0] if len(parts) == 1 else ({"anyOf": parts} if parts else {})
75 nm = name_of(n)
76 if nm is not None:
77 if nm in SIMPLE:
78 return {"type": SIMPLE[nm]}
79 return {}
80 if isinstance(n, ast.Subscript):
81 base = name_of(n.value)
82 inner = n.slice
83 args = list(inner.elts) if isinstance(inner, ast.Tuple) else [inner]
84 if base in ("Optional",):
85 return schema(args[0])
86 if base == "Union":
87 parts = [schema(a) for a in args]
88 parts = [p for p in parts if p and p.get("type") != "null"]
89 return parts[0] if len(parts) == 1 else ({"anyOf": parts} if parts else {})
90 if base == "Literal":
91 values = [a.value for a in args if isinstance(a, ast.Constant)]
92 out = {"enum": values}
93 kinds = {type(v).__name__ for v in values}
94 if len(kinds) == 1 and kinds.pop() in SIMPLE:
95 out["type"] = SIMPLE[type(values[0]).__name__]
96 return out
97 if base in ("list", "List", "Sequence", "set", "Set", "tuple", "Tuple"):
98 item = schema(args[0]) if args else None
99 return {"type": "array", "items": item} if item else {"type": "array"}
100 if base in ("dict", "Dict", "Mapping"):
101 return {"type": "object"}
102 if base == "Annotated":
103 out = schema(args[0]) or {}
104 for a in args[1:]:
105 if isinstance(a, ast.Constant) and isinstance(a.value, str):
106 out = dict(out, description=a.value)
107 return out
108 return {}
109
110def from_default(v):
111 for t, s in ((bool, "boolean"), (int, "integer"), (float, "number"), (str, "string"), (list, "array"), (dict, "object")):
112 if isinstance(v, t):
113 return {"type": s}
114 return {}
115
116def arg_docs(doc):
117 out, on = {}, False
118 for line in (doc or "").splitlines():
119 s = line.strip()
120 if s.lower().rstrip(":") in ("args", "arguments", "parameters", "params"):
121 on = True
122 continue
123 if on and s and ":" in s:
124 k, _, v = s.partition(":")
125 k = k.split("(")[0].strip().lstrip("-* ").strip()
126 if k.isidentifier() and v.strip():
127 out[k] = v.strip()
128 elif on and s and s.endswith(":"):
129 on = False
130 return out
131
132def main():
133 code = json.loads(sys.stdin.read()).get("code", "")
134 try:
135 tree = ast.parse(code)
136 except SyntaxError as e:
137 print(json.dumps({"error": "line %s: %s" % (e.lineno, e.msg)}))
138 return
139 fns = [n for n in tree.body if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)) and n.name == "run"]
140 if not fns:
141 print(json.dumps({"error": "the code must define a top-level function: def run(arg: type = default, ...)"}))
142 return
143 fn = fns[-1]
144 if isinstance(fn, ast.AsyncFunctionDef):
145 print(json.dumps({"error": "run must be a plain def, not async def"}))
146 return
147 doc = ast.get_docstring(fn) or ""
148 docs = arg_docs(doc)
149 a = fn.args
150 positional = list(a.posonlyargs) + list(a.args)
151 defaults = [None] * (len(positional) - len(a.defaults)) + list(a.defaults)
152 pairs = list(zip(positional, defaults)) + list(zip(a.kwonlyargs, a.kw_defaults))
153 props, required = {}, []
154 for arg, default in pairs:
155 if arg.arg in ("self", "cls"):
156 continue
157 s = schema(arg.annotation)
158 if default is not None:
159 try:
160 value = ast.literal_eval(default)
161 if s in (None, {}):
162 s = from_default(value)
163 if value is not None:
164 s = dict(s or {}, default=value)
165 except Exception:
166 pass
167 else:
168 required.append(arg.arg)
169 s = dict(s or {"type": "string"})
170 if arg.arg in docs and "description" not in s:
171 s["description"] = docs[arg.arg]
172 props[arg.arg] = s
173 out = {"type": "object", "properties": props}
174 if required:
175 out["required"] = required
176 summary = doc.strip().split("\n\n")[0].strip()
177 print(json.dumps({"params": out, "doc": summary}))
178
179main()
180"#;
181
182#[derive(Debug, Clone, PartialEq)]
184pub struct Derived {
185 pub params: Value,
187 pub doc: String,
189}
190
191pub const MAX_CODE_BYTES: usize = 100_000;
193
194pub fn derive_schema(code: &str, python: &[OsString], sandbox: Option<&Sandbox>) -> Result<Derived, String> {
198 if code.len() > MAX_CODE_BYTES {
199 return Err(format!("the code is too long ({} KB; at most {} KB)", code.len() / 1024, MAX_CODE_BYTES / 1024));
200 }
201 let tmp = TempDir::new("byob-schema-")?;
202 let script = tmp.path().join("byob_schema.py");
203 fs::write(&script, SCHEMA_SHIM).map_err(|e| format!("couldn't write the schema helper: {e}"))?;
204 let run = ScriptRun {
205 program: python.to_vec(),
206 script,
207 args: Vec::new(),
208 cwd: tmp.path().to_path_buf(),
209 stdin: serde_json::to_vec(&json!({ "code": code })).unwrap_or_default(),
210 env: BTreeMap::new(),
211 timeout: Duration::from_secs(20),
212 max_output: 256 * 1024,
213 };
214 let out = match sandbox {
215 Some(sb) => run_script_sandboxed(&run, sb, tmp.path(), false, &|| false),
216 None => run_script_unsandboxed(&run, &|| false),
217 }
218 .map_err(|e| format!("couldn't run Python to read the code: {e}"))?;
219 if !out.success() {
220 return Err(format!("couldn't read the code's signature: {}", last_lines(&out.stderr, 10)));
221 }
222 let v: Value = serde_json::from_str(out.stdout.trim()).map_err(|e| format!("couldn't read the code's signature ({e})"))?;
223 if let Some(e) = v.get("error").and_then(Value::as_str) {
224 return Err(e.to_owned());
225 }
226 let params = v.get("params").cloned().filter(Value::is_object).ok_or("couldn't read the code's signature")?;
227 Ok(Derived { params, doc: v.get("doc").and_then(Value::as_str).unwrap_or("").to_owned() })
228}
229
230pub fn last_lines(text: &str, n: usize) -> String {
232 let lines: Vec<&str> = text.trim_end().lines().collect();
233 let start = lines.len().saturating_sub(n);
234 lines.get(start..).unwrap_or_default().join("\n")
235}
236
237pub fn tool_script(code: &str) -> String {
239 let own = code.split(RUNNER_MARK).next().unwrap_or(code).trim_end();
240 format!("{own}\n\n\n{RUNNER_MARK}{RUNNER}")
241}
242
243pub fn script_code(script: &str) -> &str {
245 script.split(RUNNER_MARK).next().unwrap_or(script).trim_end()
246}
247
248#[derive(Debug, Clone, PartialEq, Eq)]
251pub struct WantedSetting {
252 pub key: String,
253 pub label: String,
254 pub secret: bool,
255}
256
257impl WantedSetting {
258 pub fn parse(v: &Value) -> Result<Self, String> {
261 let (key, label, secret) = match v {
262 Value::String(k) => (k.trim().to_owned(), None, None),
263 Value::Object(o) => (
264 o.get("key").and_then(Value::as_str).unwrap_or("").trim().to_owned(),
265 o.get("label").and_then(Value::as_str).map(str::to_owned),
266 o.get("secret").and_then(Value::as_bool),
267 ),
268 _ => return Err("each setting is a name like \"API_KEY\" or {\"key\": \"API_KEY\", \"secret\": true}".into()),
269 };
270 if !setting_key_ok(&key) {
271 return Err(format!("setting {key:?} must be an environment variable name in capitals, like WEATHER_API_KEY"));
272 }
273 let looks_secret = ["KEY", "TOKEN", "SECRET", "PASSWORD"].iter().any(|w| key.contains(w));
274 let label = label.filter(|l| !l.trim().is_empty()).unwrap_or_else(|| key.replace('_', " ").to_lowercase());
275 Ok(Self { secret: secret.unwrap_or(looks_secret), key, label })
276 }
277}
278
279#[derive(Debug, Clone, PartialEq)]
281pub struct ToolDecl {
282 pub name: String,
283 pub description: String,
284 pub params: Value,
285 pub run: String,
287 pub network: bool,
289 pub settings: Vec<WantedSetting>,
290}
291
292pub fn new_manifest(id: &str, name: &str, description: &str) -> Value {
294 json!({ "id": id, "name": name, "version": "0.1.0", "description": description })
295}
296
297pub fn read_manifest_json(dir: &Path, id: &str) -> Result<Value, String> {
299 let file = dir.join(MANIFEST);
300 if !file.is_file() {
301 return Ok(new_manifest(id, id, "Tools made in a chat"));
302 }
303 let text = fs::read_to_string(&file).map_err(|e| format!("{MANIFEST}: {e}"))?;
304 serde_json::from_str(&text).map_err(|e| format!("{MANIFEST} isn't valid JSON: {e}"))
305}
306
307pub fn write_manifest(dir: &Path, manifest: &Value) -> Result<(), String> {
310 let text = serde_json::to_string_pretty(manifest).map_err(|e| e.to_string())?;
311 parse_manifest_str(&text, ParseOptions { builtin: false, file: Some(MANIFEST) }).map_err(|e| e.to_string())?;
312 fs::create_dir_all(dir).map_err(|e| format!("couldn't create {}: {e}", dir.display()))?;
313 fs::write(dir.join(MANIFEST), text + "\n").map_err(|e| format!("couldn't write {MANIFEST}: {e}"))
314}
315
316pub fn upsert_tool(dir: &Path, id: &str, tool: &ToolDecl) -> Result<bool, String> {
319 if !plugin_id_ok(id) {
320 return Err(format!("plugin {id:?}: ids are lowercase letters, digits and dashes, like \"trip-planner\""));
321 }
322 if !tool_name_ok(&tool.name) {
323 return Err(format!("tool name {:?} must start with a letter and use only letters, digits and _ (at most 64)", tool.name));
324 }
325 let mut m = read_manifest_json(dir, id)?;
326 let obj = m.as_object_mut().ok_or(format!("{MANIFEST} must be a JSON object"))?;
327 let entry = json!({ "name": tool.name, "description": tool.description, "params": tool.params, "run": tool.run });
328 let tools = obj.entry("tools").or_insert_with(|| Value::Array(Vec::new()));
329 let list = tools.as_array_mut().ok_or(format!("{MANIFEST}: tools must be a list"))?;
330 let replaced = match list.iter_mut().find(|t| t.get("name").and_then(Value::as_str) == Some(tool.name.as_str())) {
331 Some(t) => {
332 *t = entry;
333 true
334 }
335 None => {
336 list.push(entry);
337 false
338 }
339 };
340 if tool.network {
341 let perms = obj.entry("permissions").or_insert_with(|| Value::Object(Map::new()));
342 if let Some(p) = perms.as_object_mut() {
343 p.insert("network".into(), Value::Bool(true));
344 }
345 }
346 if !tool.settings.is_empty() {
347 let settings = obj.entry("settings").or_insert_with(|| Value::Array(Vec::new()));
348 let list = settings.as_array_mut().ok_or(format!("{MANIFEST}: settings must be a list"))?;
349 for s in &tool.settings {
350 if !list.iter().any(|x| x.get("key").and_then(Value::as_str) == Some(s.key.as_str())) {
351 list.push(json!({ "key": s.key, "label": s.label, "secret": s.secret, "required": true }));
352 }
353 }
354 }
355 write_manifest(dir, &m)?;
356 Ok(replaced)
357}
358
359pub fn remove_tool(dir: &Path, name: &str) -> Result<bool, String> {
361 let file = dir.join(MANIFEST);
362 if !file.is_file() {
363 return Ok(false);
364 }
365 let mut m = read_manifest_json(dir, "x")?;
366 let Some(list) = m.get_mut("tools").and_then(Value::as_array_mut) else { return Ok(false) };
367 let before = list.len();
368 list.retain(|t| t.get("name").and_then(Value::as_str) != Some(name));
369 if list.len() == before {
370 return Ok(false);
371 }
372 write_manifest(dir, &m)?;
373 Ok(true)
374}
375
376const MAX_FINGERPRINT_FILES: usize = 2000;
378
379pub fn fingerprint(dir: &Path) -> u64 {
383 let mut h = Fnv::default();
384 fingerprint_into(&mut h, dir, 3);
385 h.finish()
386}
387
388pub fn fingerprint_into(h: &mut Fnv, dir: &Path, depth: usize) {
390 let mut stack: Vec<(PathBuf, usize)> = vec![(dir.to_path_buf(), 0)];
391 let mut seen = 0usize;
392 while let Some((d, level)) = stack.pop() {
393 let Ok(rd) = fs::read_dir(&d) else { continue };
394 let mut entries: Vec<_> = rd.flatten().collect();
395 entries.sort_by_key(|e| e.file_name());
396 for e in entries {
397 let name = e.file_name();
398 let name = name.to_string_lossy();
399 if name.starts_with('.') || name == "__pycache__" {
400 continue;
401 }
402 seen += 1;
403 if seen > MAX_FINGERPRINT_FILES {
404 return;
405 }
406 let Ok(meta) = e.metadata() else { continue };
407 h.write(e.path().as_os_str().as_encoded_bytes());
408 if meta.is_dir() {
409 h.write_u8(1);
410 if level < depth {
411 stack.push((e.path(), level + 1));
412 }
413 } else {
414 h.write_u64(meta.len());
415 let t = meta.modified().ok().and_then(|t| t.duration_since(UNIX_EPOCH).ok()).map_or(0, |d| d.as_nanos());
416 h.write_u128(t);
417 }
418 }
419 }
420}
421
422#[derive(Debug, Clone, Copy)]
424pub struct Fnv(u64);
425
426impl Default for Fnv {
427 fn default() -> Self {
428 Self(0xcbf2_9ce4_8422_2325)
429 }
430}
431
432impl Hasher for Fnv {
433 fn finish(&self) -> u64 {
434 self.0
435 }
436
437 fn write(&mut self, bytes: &[u8]) {
438 for b in bytes {
439 self.0 ^= u64::from(*b);
440 self.0 = self.0.wrapping_mul(0x0100_0000_01b3);
441 }
442 }
443}
444
445#[cfg(test)]
446mod tests {
447 use super::*;
448 use crate::discover::load_plugin;
449
450 fn python() -> Vec<OsString> {
451 crate::python::default_python()
452 }
453
454 #[test]
455 fn schemas_come_from_type_hints_defaults_and_docstrings() {
456 let code = r#"
457from typing import Literal, Optional
458def run(city: str, days: int = 3, units: Literal["metric", "imperial"] = "metric", tags: list[str] = [], verbose=False, note: Optional[str] = None, ratio: "float" = 0.5):
459 """Weather for a city.
460
461 Args:
462 city: The city name, e.g. Paris
463 days (int): How many days
464 """
465 return city
466"#;
467 let d = derive_schema(code, &python(), None).unwrap();
468 assert_eq!(d.doc, "Weather for a city.");
469 let p = &d.params;
470 assert_eq!(p["required"], json!(["city"]));
471 assert_eq!(p["properties"]["city"], json!({"type": "string", "description": "The city name, e.g. Paris"}));
472 assert_eq!(p["properties"]["days"], json!({"type": "integer", "default": 3, "description": "How many days"}));
473 assert_eq!(p["properties"]["units"], json!({"enum": ["metric", "imperial"], "type": "string", "default": "metric"}));
474 assert_eq!(p["properties"]["tags"], json!({"type": "array", "items": {"type": "string"}, "default": []}));
475 assert_eq!(p["properties"]["verbose"], json!({"type": "boolean", "default": false}));
476 assert_eq!(p["properties"]["note"], json!({"type": "string"}));
477 assert_eq!(p["properties"]["ratio"], json!({"type": "number", "default": 0.5}));
478 }
479
480 #[test]
481 fn bad_code_says_why_without_running_it() {
482 let marker = std::env::temp_dir().join(format!("byob-schema-ran-{}", std::process::id()));
483 let code = format!("open({:?}, 'w').write('ran')\ndef run(x: int):\n return x\n", marker.display().to_string());
484 assert!(derive_schema(&code, &python(), None).is_ok());
485 assert!(!marker.exists(), "the code must only be parsed");
486 let e = derive_schema("def run(x:\n", &python(), None).unwrap_err();
487 assert!(e.starts_with("line 1"), "{e}");
488 assert!(derive_schema("def other(): pass", &python(), None).unwrap_err().contains("def run("));
489 }
490
491 #[test]
492 fn scripts_get_one_runner_and_run_with_json_arguments() {
493 let code = "def run(a: int, b: int = 2):\n return {'sum': a + b}\n";
494 let script = tool_script(code);
495 assert_eq!(tool_script(&script), script, "the footer is replaced, not stacked");
496 assert_eq!(script_code(&script), code.trim_end());
497 let d = tempfile::tempdir().unwrap();
498 let file = d.path().join("t.py");
499 fs::write(&file, &script).unwrap();
500 let run = ScriptRun {
501 program: python(),
502 script: file,
503 args: vec![],
504 cwd: d.path().to_path_buf(),
505 stdin: br#"{"a": 1, "extra": "ignored"}"#.to_vec(),
506 env: BTreeMap::new(),
507 timeout: Duration::from_secs(20),
508 max_output: 10_000,
509 };
510 let out = run_script_unsandboxed(&run, &|| false).unwrap();
511 assert!(out.success(), "{}", out.stderr);
512 assert_eq!(serde_json::from_str::<Value>(&out.stdout).unwrap(), json!({"sum": 3}));
513 }
514
515 #[test]
516 fn manifests_gain_and_replace_tools() {
517 let d = tempfile::tempdir().unwrap();
518 let dir = d.path().join("scratch");
519 let decl = ToolDecl {
520 name: "weather".into(),
521 description: "Weather".into(),
522 params: json!({"type": "object", "properties": {}}),
523 run: "tools/weather.py".into(),
524 network: true,
525 settings: vec![WantedSetting::parse(&json!("WEATHER_API_KEY")).unwrap()],
526 };
527 assert!(!upsert_tool(&dir, "scratch", &decl).unwrap());
528 assert!(upsert_tool(&dir, "scratch", &ToolDecl { description: "Better".into(), network: false, ..decl.clone() }).unwrap());
529 fs::create_dir_all(dir.join("tools")).unwrap();
530 fs::write(dir.join("tools/weather.py"), "").unwrap();
531 let p = load_plugin(&dir, false).unwrap().unwrap();
532 assert_eq!(p.manifest.tools.len(), 1);
533 assert_eq!(p.manifest.tools[0].description, "Better");
534 assert!(p.manifest.permissions.network, "network is never turned off again");
535 assert!(p.manifest.settings[0].secret && p.manifest.settings[0].key == "WEATHER_API_KEY");
536 assert!(upsert_tool(&dir, "Bad Id", &decl).is_err());
537 assert!(upsert_tool(&dir, "scratch", &ToolDecl { name: "no spaces".into(), ..decl.clone() }).is_err());
538 let bad = ToolDecl { run: "../escape.py".into(), ..decl };
540 assert!(upsert_tool(&dir, "scratch", &bad).unwrap_err().contains("byob-plugin.json"));
541 assert!(remove_tool(&dir, "weather").unwrap());
542 assert!(!remove_tool(&dir, "weather").unwrap());
543 }
544
545 #[test]
546 fn settings_parse() {
547 assert!(WantedSetting::parse(&json!("API_TOKEN")).unwrap().secret);
548 assert!(!WantedSetting::parse(&json!("CITY")).unwrap().secret);
549 assert!(!WantedSetting::parse(&json!({"key": "API_KEY", "secret": false})).unwrap().secret);
550 assert!(WantedSetting::parse(&json!("lower")).is_err());
551 assert!(WantedSetting::parse(&json!(3)).is_err());
552 }
553
554 #[test]
555 fn fingerprints_follow_edits() {
556 let d = tempfile::tempdir().unwrap();
557 let empty = fingerprint(d.path());
558 fs::write(d.path().join("a.py"), "x").unwrap();
559 let one = fingerprint(d.path());
560 assert_ne!(empty, one);
561 assert_eq!(one, fingerprint(d.path()));
562 fs::write(d.path().join("a.py"), "xy").unwrap();
563 assert_ne!(one, fingerprint(d.path()));
564 let two = fingerprint(d.path());
565 fs::create_dir(d.path().join("__pycache__")).unwrap();
566 fs::write(d.path().join("__pycache__/a.pyc"), "junk").unwrap();
567 assert_eq!(two, fingerprint(d.path()), "caches don't count");
568 }
569}