1use std::collections::BTreeMap;
9use std::fmt;
10
11use serde_json::Value;
12
13#[derive(Debug, Clone, PartialEq)]
15pub struct ToolInvocation {
16 pub node: String,
19 pub agent: String,
25 pub reference: String,
27 pub name: String,
29 pub transport: String,
30 pub arguments: BTreeMap<String, Value>,
32 pub effects: Vec<String>,
34 pub result_type: String,
36}
37
38#[derive(Debug)]
39pub enum ToolError {
40 NotAvailable(String),
42 Failed(String),
44 InvalidResult(String),
46}
47
48impl fmt::Display for ToolError {
49 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50 match self {
51 ToolError::NotAvailable(name) => write!(
52 f,
53 "no host provides the tool `{name}`; \
54 the artifact requires it, so the run cannot continue"
55 ),
56 ToolError::Failed(message) => write!(f, "the tool failed: {message}"),
57 ToolError::InvalidResult(message) => {
58 write!(f, "the tool returned an unexpected value: {message}")
59 }
60 }
61 }
62}
63
64impl std::error::Error for ToolError {}
65
66pub trait ToolHost {
68 fn name(&self) -> &str;
69
70 fn provides(&self, tool: &str) -> bool;
73
74 fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError>;
75}
76
77impl<H: ToolHost + ?Sized> ToolHost for Box<H> {
82 fn name(&self) -> &str {
83 (**self).name()
84 }
85
86 fn provides(&self, tool: &str) -> bool {
87 (**self).provides(tool)
88 }
89
90 fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError> {
91 (**self).call(invocation)
92 }
93}
94
95pub struct DenyAllTools;
97
98impl ToolHost for DenyAllTools {
99 fn name(&self) -> &str {
100 "deny-all"
101 }
102
103 fn provides(&self, _tool: &str) -> bool {
104 false
105 }
106
107 fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError> {
108 Err(ToolError::NotAvailable(invocation.name.clone()))
109 }
110}
111
112#[derive(Default)]
115pub struct StaticToolHost {
116 #[allow(clippy::type_complexity)]
117 handlers: BTreeMap<String, Box<dyn FnMut(&ToolInvocation) -> Result<Value, ToolError>>>,
118}
119
120impl StaticToolHost {
121 pub fn new() -> StaticToolHost {
122 StaticToolHost::default()
123 }
124
125 pub fn with(
126 mut self,
127 name: impl Into<String>,
128 handler: impl FnMut(&ToolInvocation) -> Result<Value, ToolError> + 'static,
129 ) -> StaticToolHost {
130 self.handlers.insert(name.into(), Box::new(handler));
131 self
132 }
133}
134
135impl ToolHost for StaticToolHost {
136 fn name(&self) -> &str {
137 "static"
138 }
139
140 fn provides(&self, tool: &str) -> bool {
141 self.handlers.contains_key(tool)
142 }
143
144 fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError> {
145 match self.handlers.get_mut(&invocation.name) {
146 Some(handler) => handler(invocation),
147 None => Err(ToolError::NotAvailable(invocation.name.clone())),
148 }
149 }
150}
151
152pub enum HumanChannel {
160 Ask(Box<dyn Interlocutor>),
162 AssumeYes,
171 Deny,
174}
175
176impl fmt::Debug for HumanChannel {
177 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178 match self {
179 HumanChannel::Ask(_) => f.write_str("Ask(..)"),
180 HumanChannel::AssumeYes => f.write_str("AssumeYes"),
181 HumanChannel::Deny => f.write_str("Deny"),
182 }
183 }
184}
185
186pub trait Interlocutor {
188 fn approve(&mut self, request: &ApprovalRequest) -> bool;
190
191 fn consult(&mut self, request: &ConsultRequest) -> Result<String, ConsultError>;
198}
199
200#[derive(Debug, Clone, PartialEq)]
202pub struct ConsultRequest {
203 pub node: String,
204 pub index: usize,
209 pub question: String,
210 pub choices: Vec<String>,
213 pub context: Vec<(String, Value)>,
219}
220
221#[derive(Debug, Clone, PartialEq)]
223pub enum ConsultError {
224 NoChannel(String),
226 Failed(String),
228 NotAChoice {
230 answer: String,
231 choices: Vec<String>,
232 },
233}
234
235impl fmt::Display for ConsultError {
236 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
237 match self {
238 ConsultError::NoChannel(reason) => write!(f, "there is nobody to ask: {reason}"),
239 ConsultError::Failed(reason) => write!(f, "the question was not answered: {reason}"),
240 ConsultError::NotAChoice { answer, choices } => write!(
241 f,
242 "`{answer}` is not one of the choices offered ({})",
243 choices.join(", ")
244 ),
245 }
246 }
247}
248
249impl std::error::Error for ConsultError {}
250
251#[derive(Debug, Clone, PartialEq)]
252pub struct ApprovalRequest {
253 pub node: String,
254 pub effects: Vec<String>,
255 pub reason: String,
257}
258
259pub struct ScriptedApprovals {
261 answers: Vec<bool>,
262 position: usize,
263}
264
265impl ScriptedApprovals {
266 pub fn new(answers: Vec<bool>) -> ScriptedApprovals {
267 ScriptedApprovals {
268 answers,
269 position: 0,
270 }
271 }
272}
273
274impl Interlocutor for ScriptedApprovals {
275 fn approve(&mut self, _request: &ApprovalRequest) -> bool {
276 let answer = self.answers.get(self.position).copied().unwrap_or(false);
277 self.position += 1;
278 answer
279 }
280
281 fn consult(&mut self, request: &ConsultRequest) -> Result<String, ConsultError> {
282 Err(ConsultError::NoChannel(format!(
283 "scripted approvals answer gates and not questions (node `{}`)",
284 request.node
285 )))
286 }
287}
288
289pub struct ScriptedAnswers {
291 answers: Vec<String>,
292 position: usize,
293}
294
295impl ScriptedAnswers {
296 pub fn new(answers: Vec<impl Into<String>>) -> ScriptedAnswers {
297 ScriptedAnswers {
298 answers: answers.into_iter().map(Into::into).collect(),
299 position: 0,
300 }
301 }
302}
303
304impl Interlocutor for ScriptedAnswers {
305 fn approve(&mut self, _request: &ApprovalRequest) -> bool {
308 true
309 }
310
311 fn consult(&mut self, request: &ConsultRequest) -> Result<String, ConsultError> {
312 let Some(answer) = self.answers.get(self.position).cloned() else {
313 return Err(ConsultError::NoChannel(format!(
314 "the script has {} answer(s) and node `{}` asked for another",
315 self.answers.len(),
316 request.node
317 )));
318 };
319 self.position += 1;
320 Ok(answer)
321 }
322}
323
324#[cfg(test)]
325mod tests {
326 use super::*;
327 use serde_json::json;
328
329 fn invocation(name: &str) -> ToolInvocation {
330 ToolInvocation {
331 node: "n0".to_string(),
332 agent: "test.Agent".to_string(),
333 reference: format!("mcp:{name}"),
334 name: name.to_string(),
335 transport: "mcp".to_string(),
336 arguments: BTreeMap::new(),
337 effects: vec!["network".to_string()],
338 result_type: "json".to_string(),
339 }
340 }
341
342 #[test]
343 fn the_default_host_denies_everything() {
344 let mut host = DenyAllTools;
345 assert!(!host.provides("web.search"));
346 let error = host.call(&invocation("web.search")).unwrap_err();
347 assert!(error.to_string().contains("no host provides"), "{error}");
348 }
349
350 #[test]
351 fn a_static_host_serves_registered_tools() {
352 let mut host = StaticToolHost::new().with("web.search", |_| Ok(json!(["a", "b"])));
353 assert!(host.provides("web.search"));
354 assert!(!host.provides("files.write"));
355 assert_eq!(
356 host.call(&invocation("web.search")).unwrap(),
357 json!(["a", "b"])
358 );
359 }
360
361 #[test]
362 fn an_unregistered_tool_is_reported_by_name() {
363 let mut host = StaticToolHost::new().with("web.search", |_| Ok(json!(null)));
364 let error = host.call(&invocation("files.write")).unwrap_err();
365 assert!(error.to_string().contains("files.write"), "{error}");
366 }
367
368 #[test]
369 fn scripted_approvals_deny_once_exhausted() {
370 let mut handler = ScriptedApprovals::new(vec![true]);
371 let request = ApprovalRequest {
372 node: "n0".into(),
373 effects: vec!["external_write".into()],
374 reason: "test".into(),
375 };
376 assert!(handler.approve(&request));
377 assert!(
378 !handler.approve(&request),
379 "an exhausted script must not keep approving"
380 );
381 }
382}