use std::collections::BTreeMap;
use std::fmt;
use serde_json::Value;
#[derive(Debug, Clone, PartialEq)]
pub struct ToolInvocation {
pub node: String,
pub agent: String,
pub reference: String,
pub name: String,
pub transport: String,
pub arguments: BTreeMap<String, Value>,
pub effects: Vec<String>,
pub result_type: String,
}
#[derive(Debug)]
pub enum ToolError {
NotAvailable(String),
Failed(String),
InvalidResult(String),
}
impl fmt::Display for ToolError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ToolError::NotAvailable(name) => write!(
f,
"no host provides the tool `{name}`; \
the artifact requires it, so the run cannot continue"
),
ToolError::Failed(message) => write!(f, "the tool failed: {message}"),
ToolError::InvalidResult(message) => {
write!(f, "the tool returned an unexpected value: {message}")
}
}
}
}
impl std::error::Error for ToolError {}
pub trait ToolHost {
fn name(&self) -> &str;
fn provides(&self, tool: &str) -> bool;
fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError>;
}
impl<H: ToolHost + ?Sized> ToolHost for Box<H> {
fn name(&self) -> &str {
(**self).name()
}
fn provides(&self, tool: &str) -> bool {
(**self).provides(tool)
}
fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError> {
(**self).call(invocation)
}
}
pub struct DenyAllTools;
impl ToolHost for DenyAllTools {
fn name(&self) -> &str {
"deny-all"
}
fn provides(&self, _tool: &str) -> bool {
false
}
fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError> {
Err(ToolError::NotAvailable(invocation.name.clone()))
}
}
#[derive(Default)]
pub struct StaticToolHost {
#[allow(clippy::type_complexity)]
handlers: BTreeMap<String, Box<dyn FnMut(&ToolInvocation) -> Result<Value, ToolError>>>,
}
impl StaticToolHost {
pub fn new() -> StaticToolHost {
StaticToolHost::default()
}
pub fn with(
mut self,
name: impl Into<String>,
handler: impl FnMut(&ToolInvocation) -> Result<Value, ToolError> + 'static,
) -> StaticToolHost {
self.handlers.insert(name.into(), Box::new(handler));
self
}
}
impl ToolHost for StaticToolHost {
fn name(&self) -> &str {
"static"
}
fn provides(&self, tool: &str) -> bool {
self.handlers.contains_key(tool)
}
fn call(&mut self, invocation: &ToolInvocation) -> Result<Value, ToolError> {
match self.handlers.get_mut(&invocation.name) {
Some(handler) => handler(invocation),
None => Err(ToolError::NotAvailable(invocation.name.clone())),
}
}
}
pub enum ApprovalMode {
Ask(Box<dyn ApprovalHandler>),
AssumeYes,
Deny,
}
impl fmt::Debug for ApprovalMode {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ApprovalMode::Ask(_) => f.write_str("Ask(..)"),
ApprovalMode::AssumeYes => f.write_str("AssumeYes"),
ApprovalMode::Deny => f.write_str("Deny"),
}
}
}
pub trait ApprovalHandler {
fn approve(&mut self, request: &ApprovalRequest) -> bool;
}
#[derive(Debug, Clone, PartialEq)]
pub struct ApprovalRequest {
pub node: String,
pub effects: Vec<String>,
pub reason: String,
}
pub struct ScriptedApprovals {
answers: Vec<bool>,
position: usize,
}
impl ScriptedApprovals {
pub fn new(answers: Vec<bool>) -> ScriptedApprovals {
ScriptedApprovals {
answers,
position: 0,
}
}
}
impl ApprovalHandler for ScriptedApprovals {
fn approve(&mut self, _request: &ApprovalRequest) -> bool {
let answer = self.answers.get(self.position).copied().unwrap_or(false);
self.position += 1;
answer
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn invocation(name: &str) -> ToolInvocation {
ToolInvocation {
node: "n0".to_string(),
agent: "test.Agent".to_string(),
reference: format!("mcp:{name}"),
name: name.to_string(),
transport: "mcp".to_string(),
arguments: BTreeMap::new(),
effects: vec!["network".to_string()],
result_type: "json".to_string(),
}
}
#[test]
fn the_default_host_denies_everything() {
let mut host = DenyAllTools;
assert!(!host.provides("web.search"));
let error = host.call(&invocation("web.search")).unwrap_err();
assert!(error.to_string().contains("no host provides"), "{error}");
}
#[test]
fn a_static_host_serves_registered_tools() {
let mut host = StaticToolHost::new().with("web.search", |_| Ok(json!(["a", "b"])));
assert!(host.provides("web.search"));
assert!(!host.provides("files.write"));
assert_eq!(
host.call(&invocation("web.search")).unwrap(),
json!(["a", "b"])
);
}
#[test]
fn an_unregistered_tool_is_reported_by_name() {
let mut host = StaticToolHost::new().with("web.search", |_| Ok(json!(null)));
let error = host.call(&invocation("files.write")).unwrap_err();
assert!(error.to_string().contains("files.write"), "{error}");
}
#[test]
fn scripted_approvals_deny_once_exhausted() {
let mut handler = ScriptedApprovals::new(vec![true]);
let request = ApprovalRequest {
node: "n0".into(),
effects: vec!["external_write".into()],
reason: "test".into(),
};
assert!(handler.approve(&request));
assert!(
!handler.approve(&request),
"an exhausted script must not keep approving"
);
}
}