use std::collections::BTreeMap;
use std::fmt::Debug;
use async_trait::async_trait;
use serde::de::DeserializeOwned;
use serde_json::Value;
use super::{Disposition, ToolClient, ToolError, ToolId};
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum ToolFailure {
#[error("did not happen: {0}")]
DidNotHappen(String),
#[error("outcome unknown: {0}")]
InDoubt(String),
#[error("landed and failed: {0}")]
Landed(String),
}
impl ToolFailure {
#[must_use]
pub const fn disposition(&self) -> Disposition {
match self {
Self::DidNotHappen(_) => Disposition::DidNotHappen,
Self::InDoubt(_) => Disposition::InDoubt,
Self::Landed(_) => Disposition::Landed,
}
}
fn at(self, tool: ToolId) -> ToolError {
let detail = match &self {
Self::DidNotHappen(d) | Self::InDoubt(d) | Self::Landed(d) => d.clone(),
};
match self {
Self::DidNotHappen(_) => ToolError::Refused { tool, detail },
Self::InDoubt(_) => ToolError::TimedOut { tool, detail },
Self::Landed(_) => ToolError::ToolFailed { tool, detail },
}
}
}
#[async_trait]
pub trait Tool: DeserializeOwned + schemars::JsonSchema + Send + Sync + 'static {
const SERVER: &'static str;
const NAME: &'static str;
fn mutates() -> bool {
true
}
async fn call(self) -> Result<Value, ToolFailure>;
}
struct Registered {
description: String,
schema: Value,
mutates: bool,
#[allow(clippy::type_complexity)]
invoke: Box<
dyn Fn(Value) -> futures_core::future::BoxFuture<'static, Result<Value, ToolError>>
+ Send
+ Sync,
>,
}
#[derive(Default)]
pub struct ToolBox {
tools: BTreeMap<ToolId, Registered>,
}
impl Debug for ToolBox {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ToolBox")
.field("tools", &self.tools.keys().collect::<Vec<_>>())
.finish()
}
}
fn close_objects(node: &mut Value) {
if node.get("type") == Some(&serde_json::json!("object"))
&& let Some(map) = node.as_object_mut()
{
map.entry("additionalProperties")
.or_insert_with(|| serde_json::json!(false));
}
for key in ["properties", "$defs", "definitions"] {
if let Some(children) = node.get_mut(key).and_then(Value::as_object_mut) {
for child in children.values_mut() {
close_objects(child);
}
}
}
if let Some(items) = node.get_mut("items") {
close_objects(items);
}
}
impl ToolBox {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with<T: Tool>(mut self) -> Self {
let mut schema = serde_json::to_value(schemars::schema_for!(T))
.expect("a schemars-generated schema must serialize to JSON");
close_objects(&mut schema);
let description = schema
.get("description")
.and_then(Value::as_str)
.unwrap_or_default()
.to_owned();
let id = ToolId::new(T::SERVER, T::NAME);
assert!(
!self.tools.contains_key(&id),
"typed tool '{id}' was registered twice — silently replacing one body with \
another makes registration order decide which implementation runs"
);
self.tools.insert(
id,
Registered {
description,
schema,
mutates: T::mutates(),
invoke: Box::new(|args: Value| {
Box::pin(async move {
let parsed: T =
serde_json::from_value(args).map_err(|e| ToolError::Malformed {
tool: ToolId::new(T::SERVER, T::NAME),
detail: format!("arguments do not match the declared shape: {e}"),
})?;
parsed
.call()
.await
.map_err(|failure| failure.at(ToolId::new(T::SERVER, T::NAME)))
})
}),
},
);
self
}
pub fn ids(&self) -> impl Iterator<Item = &ToolId> {
self.tools.keys()
}
pub fn servers(&self) -> impl Iterator<Item = &str> {
let mut seen: Option<&str> = None;
self.tools.keys().filter_map(move |id| {
let server = id.server.as_str();
if seen == Some(server) {
return None;
}
seen = Some(server);
Some(server)
})
}
#[cfg(feature = "manifest")]
pub fn check_against(
&self,
manifest: &crate::manifest::Manifest,
remote_servers: &std::collections::BTreeSet<String>,
) -> Result<(), Vec<String>> {
let granted: BTreeMap<ToolId, &crate::manifest::ToolGrant> = manifest
.spec
.tools
.iter()
.filter_map(|g| ToolId::parse(&g.reference).map(|id| (id, g)))
.collect();
let mut problems = Vec::new();
for (id, registered) in &self.tools {
if registered.mutates && granted.get(id).is_some_and(|g| !g.mutates) {
problems.push(format!(
"'{id}' declares that it mutates and the manifest grants it as \
read-only — that exemption lets model-chosen arguments reach \
something that changes the world"
));
}
if let Some(grant) = granted.get(id)
&& grant.arguments.is_some()
{
problems.push(format!(
"'{id}' repeats its argument schema in the manifest — typed tools \
derive it from the Rust argument type, so the second copy can only \
drift; remove `arguments` from this grant"
));
}
}
for id in self.tools.keys() {
if !granted.contains_key(id) {
problems.push(format!(
"'{id}' is implemented but the manifest grants no such tool — \
this binary can do something its declaration does not admit"
));
}
}
for id in granted.keys() {
if id.server == super::AGENT_SERVER {
continue;
}
if !self.tools.contains_key(id) && !remote_servers.contains(&id.server) {
problems.push(format!(
"'{id}' is granted but nothing implements it and no transport is \
wired for server '{}' — the model will be offered a tool that \
fails when chosen",
id.server
));
}
}
if problems.is_empty() {
Ok(())
} else {
Err(problems)
}
}
#[must_use]
pub fn declared(&self, id: &ToolId) -> Option<(&str, &Value, bool)> {
self.tools
.get(id)
.map(|t| (t.description.as_str(), &t.schema, t.mutates))
}
}
#[async_trait]
impl ToolClient for ToolBox {
async fn call(
&self,
tool: &ToolId,
arguments: &Value,
_provenance: Option<&crate::core::Provenance>,
) -> Result<Value, ToolError> {
let Some(registered) = self.tools.get(tool) else {
return Err(ToolError::Refused {
tool: tool.clone(),
detail: "this box offers no such tool".to_owned(),
});
};
(registered.invoke)(arguments.clone()).await
}
fn destination(&self, _tool: &ToolId) -> crate::tools::Destination {
crate::tools::Destination::Local
}
}
#[cfg(test)]
mod tests {
use super::{Disposition, Tool, ToolBox, ToolClient, ToolFailure, ToolId};
use serde_json::{Value, json};
#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
struct Refuses {
account: String,
}
#[async_trait::async_trait]
impl Tool for Refuses {
const SERVER: &'static str = "ledger";
const NAME: &'static str = "read";
fn mutates() -> bool {
false
}
async fn call(self) -> Result<Value, ToolFailure> {
Err(ToolFailure::DidNotHappen(format!(
"no account {}",
self.account
)))
}
}
#[test]
#[cfg(feature = "providers")]
fn a_typed_tools_schema_can_constrain_generation() {
let box_ = ToolBox::new().with::<Refuses>();
let (_, schema, _) = box_
.declared(&ToolId::new("ledger", "read"))
.expect("the tool was registered");
assert_eq!(
crate::model::strict_schema_problem(schema),
None,
"schema: {schema}"
);
}
#[test]
fn a_failure_keeps_its_disposition_and_gains_the_calling_tools_identity() {
for (failure, expected) in [
(
ToolFailure::DidNotHappen("nope".into()),
Disposition::DidNotHappen,
),
(ToolFailure::InDoubt("unknown".into()), Disposition::InDoubt),
(ToolFailure::Landed("failed".into()), Disposition::Landed),
] {
assert_eq!(failure.disposition(), expected);
let id = ToolId::new("ledger", "read");
let error = failure.clone().at(id.clone());
assert_eq!(
error.disposition(),
expected,
"attaching an identity changed what the runtime concludes about \
{failure:?}"
);
}
}
#[tokio::test]
async fn a_body_that_refuses_is_reported_as_not_having_happened() {
let box_ = ToolBox::new().with::<Refuses>();
let error = box_
.call(
&ToolId::new("ledger", "read"),
&json!({ "account": "AC-9" }),
None,
)
.await
.expect_err("the body refused");
assert_eq!(error.disposition(), Disposition::DidNotHappen);
assert!(
error.to_string().contains("AC-9"),
"the body's own detail did not survive: {error}"
);
}
}