1#![cfg_attr(not(test), deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable))]
9
10pub mod artifact;
11pub mod artifact_types;
12pub mod chart;
13mod fs;
14pub mod images;
15pub mod preview;
16pub mod sandbox;
17pub mod undo;
18
19pub use artifact::{CreateArtifact, ReadArtifact, UpdateArtifact};
20pub use artifact_types::{ArtifactHooks, ArtifactType, ArtifactTypes};
21pub use chart::MakeChart;
22pub use fs::{DEFAULT_TRASH_DIR, DeleteFile, ListFiles, MoveFile, MoveFiles, ReadFile, WriteFile};
23pub use preview::Preview;
24pub use undo::{Journal, UndoOp};
25
26use std::collections::BTreeMap;
27use std::path::{Path, PathBuf};
28use std::sync::Arc;
29use std::sync::atomic::{AtomicBool, Ordering};
30
31use serde::de::DeserializeOwned;
32use serde_json::Value;
33
34#[derive(Debug, Clone, PartialEq)]
36pub struct ToolSpec {
37 pub name: String,
39 pub description: String,
40 pub parameters: Value,
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
46pub enum Access {
47 Read,
49 Write,
51 Network,
53 Execute,
55}
56
57pub struct ToolContext<'a> {
59 pub workspace: &'a Path,
61 cancel: &'a AtomicBool,
62 journal: Option<&'a Journal>,
63 chat: Option<&'a str>,
64 asker: Option<&'a dyn Asker>,
65 env: Option<&'a EnvPlan>,
66 host: Option<&'a (dyn std::any::Any + Sync)>,
67 views: Option<&'a std::sync::Mutex<Vec<byob_protocol::View>>>,
68}
69
70#[derive(Debug, Clone, Default, PartialEq, Eq)]
74pub struct EnvPlan {
75 pub container: Option<bool>,
78 pub project: Option<bool>,
81 pub image: Option<String>,
83 pub providers: Option<EnvProviders>,
86}
87
88#[derive(Debug, Clone, Default, PartialEq, Eq)]
90pub struct EnvProviders {
91 pub off: Vec<String>,
93 pub external: Vec<CommandEnv>,
95}
96
97#[derive(Debug, Clone, Default, PartialEq, Eq)]
101pub struct CommandEnv {
102 pub id: String,
104 pub label: String,
106 pub detect: Vec<String>,
107 pub inputs: Vec<String>,
109 pub program: PathBuf,
110 pub args: Vec<String>,
111 pub dir: PathBuf,
113}
114
115impl<'a> ToolContext<'a> {
116 pub fn new(workspace: &'a Path, cancel: &'a AtomicBool) -> Self {
117 Self { workspace, cancel, journal: None, chat: None, asker: None, env: None, host: None, views: None }
118 }
119
120 #[must_use]
122 pub fn with_views(mut self, views: &'a std::sync::Mutex<Vec<byob_protocol::View>>) -> Self {
123 self.views = Some(views);
124 self
125 }
126
127 pub fn attach_view(&self, view: byob_protocol::View) -> bool {
129 let Some(v) = self.views else { return false };
130 v.lock().unwrap_or_else(std::sync::PoisonError::into_inner).push(view);
131 true
132 }
133
134 #[must_use]
138 pub fn with_host(mut self, host: &'a (dyn std::any::Any + Sync)) -> Self {
139 self.host = Some(host);
140 self
141 }
142
143 pub fn host<T: 'static>(&self) -> Option<&'a T> {
145 self.host.and_then(|h| (h as &dyn std::any::Any).downcast_ref::<T>())
146 }
147
148 #[must_use]
150 pub fn with_env(mut self, plan: &'a EnvPlan) -> Self {
151 self.env = Some(plan);
152 self
153 }
154
155 pub fn env(&self) -> EnvPlan {
157 self.env.cloned().unwrap_or_default()
158 }
159
160 #[must_use]
162 pub fn with_asker(mut self, asker: &'a dyn Asker) -> Self {
163 self.asker = Some(asker);
164 self
165 }
166
167 pub fn ask(&self, q: &Question) -> Answer {
170 if self.cancelled() {
171 return Answer::AutoNo;
172 }
173 self.asker.map_or(Answer::AutoNo, |a| a.ask(q))
174 }
175
176 #[must_use]
178 pub fn with_chat(mut self, chat: &'a str) -> Self {
179 self.chat = Some(chat);
180 self
181 }
182
183 pub fn chat(&self) -> Option<&'a str> {
184 self.chat
185 }
186
187 #[must_use]
189 pub fn with_journal(mut self, journal: &'a Journal) -> Self {
190 self.journal = Some(journal);
191 self
192 }
193
194 pub fn journal(&self) -> Option<&'a Journal> {
198 self.journal
199 }
200
201 pub fn record(&self, op: UndoOp) {
203 if let Some(j) = self.journal {
204 j.record(op);
205 }
206 }
207
208 pub fn cancelled(&self) -> bool {
209 self.cancel.load(Ordering::Relaxed)
210 }
211
212 pub fn cancel_flag(&self) -> &AtomicBool {
214 self.cancel
215 }
216}
217
218#[derive(Debug, Clone, PartialEq, Eq)]
220pub enum FileChange {
221 Created(PathBuf),
222 Modified(PathBuf),
223 Deleted(PathBuf),
224 Moved { from: PathBuf, to: PathBuf },
225}
226
227#[derive(Debug, Clone, PartialEq, Eq)]
231pub struct Question {
232 pub title: String,
234 pub text: String,
236 pub warning: Option<String>,
238 pub access: Access,
241 pub always_ask: bool,
243}
244
245#[derive(Debug, Clone, Copy, PartialEq, Eq)]
247pub enum Answer {
248 Yes,
250 No,
252 AutoYes,
254 AutoNo,
256}
257
258impl Answer {
259 pub fn allowed(self) -> bool {
260 matches!(self, Self::Yes | Self::AutoYes)
261 }
262
263 pub fn by_user(self) -> bool {
265 matches!(self, Self::Yes | Self::No)
266 }
267}
268
269pub trait Asker: Send + Sync {
271 fn ask(&self, q: &Question) -> Answer;
272}
273
274pub trait EnvControl: Send + Sync {
278 fn describe(&self, ws: &Path, plan: &EnvPlan) -> byob_protocol::WorkspaceEnv;
280 fn apply(&self, ws: &Path, action: &byob_protocol::EnvAction) -> Result<String, String>;
282 fn support(&self) -> byob_protocol::EnvSupport;
284 fn check_containers(&self) -> Result<String, String>;
286 fn forget(&self, ws: &Path) -> Result<String, String>;
289}
290
291#[derive(Debug, Clone, Default, PartialEq)]
293pub struct ToolOutput {
294 pub text: String,
295 pub failed: bool,
297 pub changes: Vec<FileChange>,
298}
299
300impl ToolOutput {
301 pub fn ok(text: impl Into<String>) -> Self {
302 Self { text: text.into(), ..Default::default() }
303 }
304
305 pub fn err(text: impl Into<String>) -> Self {
306 Self { text: text.into(), failed: true, ..Default::default() }
307 }
308}
309
310pub trait Tool: Send + Sync {
311 fn spec(&self) -> ToolSpec;
312
313 fn title(&self, args: &Value) -> String;
315
316 fn pending_title(&self, _args: &Value) -> String {
319 format!("Run {}", self.spec().name)
320 }
321
322 fn preview(&self, args: &Value, _workspace: &Path) -> Option<Preview> {
325 Some(preview::args(args))
326 }
327
328 fn access(&self, args: &Value) -> Access;
330
331 fn always_ask(&self, _args: &Value) -> Option<String> {
335 None
336 }
337
338 fn run(&self, args: &Value, ctx: &ToolContext<'_>) -> ToolOutput;
340
341 fn concurrent(&self) -> bool {
345 false
346 }
347}
348
349pub fn parse_args<T: DeserializeOwned>(args: &Value) -> Result<T, String> {
351 serde_json::from_value(args.clone()).map_err(|e| format!("invalid arguments: {e}"))
352}
353
354#[derive(Default, Clone)]
356pub struct Registry {
357 tools: BTreeMap<String, Arc<dyn Tool>>,
358}
359
360impl Registry {
361 pub fn register(&mut self, tool: Arc<dyn Tool>) {
363 self.tools.insert(tool.spec().name, tool);
364 }
365
366 pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
367 self.tools.get(name)
368 }
369
370 pub fn remove(&mut self, name: &str) -> Option<Arc<dyn Tool>> {
372 self.tools.remove(name)
373 }
374
375 pub fn specs(&self) -> Vec<ToolSpec> {
376 self.tools.values().map(|t| t.spec()).collect()
377 }
378
379 pub fn len(&self) -> usize {
380 self.tools.len()
381 }
382
383 pub fn is_empty(&self) -> bool {
384 self.tools.is_empty()
385 }
386}
387
388pub fn register_builtins(reg: &mut Registry) {
392 register_builtins_with(reg, &ArtifactTypes::builtin());
393}
394
395pub fn register_builtins_with(reg: &mut Registry, types: &ArtifactTypes) {
397 reg.register(Arc::new(ListFiles));
398 reg.register(Arc::new(ReadFile));
399 reg.register(Arc::new(WriteFile));
400 reg.register(Arc::new(MoveFile));
401 reg.register(Arc::new(MoveFiles));
402 reg.register(Arc::new(DeleteFile::default()));
403 reg.register(Arc::new(MakeChart));
404 reg.register(Arc::new(CreateArtifact { types: types.clone() }));
405 reg.register(Arc::new(UpdateArtifact { types: types.clone() }));
406 reg.register(Arc::new(ReadArtifact { types: types.clone() }));
407}
408
409pub fn builtin() -> Registry {
411 let mut reg = Registry::default();
412 register_builtins(&mut reg);
413 reg
414}
415
416#[cfg(test)]
417mod tests {
418 use super::*;
419
420 struct Echo;
421
422 impl Tool for Echo {
423 fn spec(&self) -> ToolSpec {
424 ToolSpec {
425 name: "echo".into(),
426 description: "Echo text".into(),
427 parameters: serde_json::json!({"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}),
428 }
429 }
430 fn title(&self, _args: &Value) -> String {
431 "Echo".into()
432 }
433 fn access(&self, _args: &Value) -> Access {
434 Access::Read
435 }
436 fn run(&self, args: &Value, _ctx: &ToolContext<'_>) -> ToolOutput {
437 #[derive(serde::Deserialize)]
438 struct Args {
439 text: String,
440 }
441 match parse_args::<Args>(args) {
442 Ok(a) => ToolOutput::ok(a.text),
443 Err(e) => ToolOutput::err(e),
444 }
445 }
446 }
447
448 #[test]
449 fn registry_dispatch_and_bad_args() {
450 let mut reg = Registry::default();
451 reg.register(Arc::new(Echo));
452 let cancel = AtomicBool::new(false);
453 let ctx = ToolContext::new(Path::new("/tmp"), &cancel);
454 let tool = reg.get("echo").unwrap();
455 assert_eq!(tool.run(&serde_json::json!({"text":"hi"}), &ctx), ToolOutput::ok("hi"));
456 assert!(tool.run(&serde_json::json!({"nope":1}), &ctx).failed);
457 assert_eq!(reg.specs().len(), 1);
458 }
459
460 #[test]
461 fn builtins_are_valid() {
462 let reg = builtin();
463 assert_eq!(reg.len(), 10);
464 for spec in reg.specs() {
465 assert!(spec.name.len() <= 64 && spec.name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'));
466 assert_eq!(spec.parameters["type"], "object");
467 let tool = reg.get(&spec.name).unwrap();
468 let expected = if matches!(spec.name.as_str(), "list_files" | "read_file" | "make_chart" | "create_artifact" | "update_artifact" | "read_artifact")
469 {
470 Access::Read
471 } else {
472 Access::Write
473 };
474 assert_eq!(tool.access(&serde_json::json!({})), expected, "{}", spec.name);
475 }
476 }
477
478 #[test]
479 fn default_pending_title_and_preview() {
480 assert_eq!(Echo.pending_title(&serde_json::json!({})), "Run echo");
481 let p = Echo.preview(&serde_json::json!({"text": "hi"}), Path::new("/tmp"));
482 assert_eq!(p, Some(Preview::Args { fields: vec![("text".into(), "hi".into())], code: None, more: 0 }));
483 }
484
485 #[test]
486 fn access_is_ordered_by_danger() {
487 assert!(Access::Read < Access::Write && Access::Write < Access::Network && Access::Network < Access::Execute);
488 }
489}