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 ApprovalMode {
154 Ask(Box<dyn ApprovalHandler>),
156 AssumeYes,
159 Deny,
161}
162
163impl fmt::Debug for ApprovalMode {
164 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
165 match self {
166 ApprovalMode::Ask(_) => f.write_str("Ask(..)"),
167 ApprovalMode::AssumeYes => f.write_str("AssumeYes"),
168 ApprovalMode::Deny => f.write_str("Deny"),
169 }
170 }
171}
172
173pub trait ApprovalHandler {
175 fn approve(&mut self, request: &ApprovalRequest) -> bool;
176}
177
178#[derive(Debug, Clone, PartialEq)]
179pub struct ApprovalRequest {
180 pub node: String,
181 pub effects: Vec<String>,
182 pub reason: String,
184}
185
186pub struct ScriptedApprovals {
188 answers: Vec<bool>,
189 position: usize,
190}
191
192impl ScriptedApprovals {
193 pub fn new(answers: Vec<bool>) -> ScriptedApprovals {
194 ScriptedApprovals {
195 answers,
196 position: 0,
197 }
198 }
199}
200
201impl ApprovalHandler for ScriptedApprovals {
202 fn approve(&mut self, _request: &ApprovalRequest) -> bool {
203 let answer = self.answers.get(self.position).copied().unwrap_or(false);
204 self.position += 1;
205 answer
206 }
207}
208
209#[cfg(test)]
210mod tests {
211 use super::*;
212 use serde_json::json;
213
214 fn invocation(name: &str) -> ToolInvocation {
215 ToolInvocation {
216 node: "n0".to_string(),
217 agent: "test.Agent".to_string(),
218 reference: format!("mcp:{name}"),
219 name: name.to_string(),
220 transport: "mcp".to_string(),
221 arguments: BTreeMap::new(),
222 effects: vec!["network".to_string()],
223 result_type: "json".to_string(),
224 }
225 }
226
227 #[test]
228 fn the_default_host_denies_everything() {
229 let mut host = DenyAllTools;
230 assert!(!host.provides("web.search"));
231 let error = host.call(&invocation("web.search")).unwrap_err();
232 assert!(error.to_string().contains("no host provides"), "{error}");
233 }
234
235 #[test]
236 fn a_static_host_serves_registered_tools() {
237 let mut host = StaticToolHost::new().with("web.search", |_| Ok(json!(["a", "b"])));
238 assert!(host.provides("web.search"));
239 assert!(!host.provides("files.write"));
240 assert_eq!(
241 host.call(&invocation("web.search")).unwrap(),
242 json!(["a", "b"])
243 );
244 }
245
246 #[test]
247 fn an_unregistered_tool_is_reported_by_name() {
248 let mut host = StaticToolHost::new().with("web.search", |_| Ok(json!(null)));
249 let error = host.call(&invocation("files.write")).unwrap_err();
250 assert!(error.to_string().contains("files.write"), "{error}");
251 }
252
253 #[test]
254 fn scripted_approvals_deny_once_exhausted() {
255 let mut handler = ScriptedApprovals::new(vec![true]);
256 let request = ApprovalRequest {
257 node: "n0".into(),
258 effects: vec!["external_write".into()],
259 reason: "test".into(),
260 };
261 assert!(handler.approve(&request));
262 assert!(
263 !handler.approve(&request),
264 "an exhausted script must not keep approving"
265 );
266 }
267}