use std::collections::HashMap;
use std::fmt;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Instant;
use async_trait::async_trait;
use serde_json::Value;
use sha2::{Digest, Sha256};
use af_llm::Tool as LlmTool;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolSurface {
Llm,
Chassis,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolConcurrency {
Concurrent,
Exclusive,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ToolMeta {
pub surface: ToolSurface,
pub cost_units: u64,
pub timeout_secs: u64,
pub core: bool,
pub concurrency: ToolConcurrency,
}
#[derive(Debug, Clone, Default)]
pub struct CancellationToken(Arc<CancellationState>);
#[derive(Debug, Default)]
struct CancellationState {
cancelled: AtomicBool,
parent: Option<CancellationToken>,
}
impl CancellationToken {
pub fn cancel(&self) {
self.0.cancelled.store(true, Ordering::Release);
}
pub fn is_cancelled(&self) -> bool {
self.0.cancelled.load(Ordering::Acquire)
|| self
.0
.parent
.as_ref()
.is_some_and(CancellationToken::is_cancelled)
}
pub fn child(&self) -> Self {
Self(Arc::new(CancellationState {
cancelled: AtomicBool::new(false),
parent: Some(self.clone()),
}))
}
}
#[derive(Debug, Clone)]
pub struct ToolExecutionContext {
pub request: af_context::RequestContext,
pub session_id: String,
pub run_id: String,
pub step: u32,
pub call_id: String,
pub source_event_seq: u64,
pub interaction_resolution: Option<af_agent_session::InteractionResolution>,
pub cancellation: CancellationToken,
pub deadline: Instant,
}
impl Default for ToolMeta {
fn default() -> Self {
Self {
surface: ToolSurface::Llm,
cost_units: 1,
timeout_secs: 15,
core: false,
concurrency: ToolConcurrency::Exclusive,
}
}
}
#[async_trait]
pub trait Tool: Send + Sync {
fn name(&self) -> &str;
fn implementation_version(&self) -> &str {
""
}
fn description(&self) -> &str;
fn parameters(&self) -> Value;
fn output_schema(&self) -> Value;
fn meta(&self) -> ToolMeta {
ToolMeta::default()
}
async fn call(&self, args: Value) -> Result<Value, String>;
async fn call_with_context(
&self,
_context: &ToolExecutionContext,
args: Value,
) -> Result<Value, String> {
self.call(args).await
}
}
#[derive(Default, Clone)]
pub struct ToolRegistry {
tools: HashMap<String, Arc<dyn Tool>>,
wire_names: HashMap<String, String>,
validators: HashMap<String, Arc<jsonschema::Validator>>,
output_validators: HashMap<String, Arc<jsonschema::Validator>>,
}
impl ToolRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register(&mut self, tool: Arc<dyn Tool>) -> Result<&mut Self, String> {
let name = tool.name().to_string();
if self.tools.contains_key(&name) {
return Err(format!("duplicate tool '{name}'"));
}
let wire_name = model_tool_name(&name);
if self.wire_names.contains_key(&wire_name)
|| (wire_name != name && self.tools.contains_key(&wire_name))
|| self.wire_names.contains_key(&name)
{
return Err(format!(
"tool name '{name}' collides on provider name '{wire_name}'"
));
}
let validator = jsonschema::validator_for(&tool.parameters())
.map_err(|error| format!("invalid schema for tool '{name}': {error}"))?;
let output_validator = jsonschema::validator_for(&tool.output_schema())
.map_err(|error| format!("invalid output schema for tool '{name}': {error}"))?;
self.tools.insert(name.clone(), tool);
self.wire_names.insert(wire_name, name.clone());
self.validators.insert(name.clone(), Arc::new(validator));
self.output_validators
.insert(name, Arc::new(output_validator));
Ok(self)
}
pub fn extend(&mut self, other: &Self) -> Result<(), String> {
for tool in other.tools.values() {
self.register(Arc::clone(tool))?;
}
Ok(())
}
pub fn is_empty(&self) -> bool {
self.tools.is_empty()
}
pub fn len(&self) -> usize {
self.tools.len()
}
pub fn contains(&self, name: &str) -> bool {
self.tools.contains_key(name)
}
pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
self.canonical_name(name)
.and_then(|name| self.tools.get(name))
.cloned()
}
pub fn validate_arguments(&self, name: &str, arguments: &Value) -> Result<(), String> {
let name = self
.canonical_name(name)
.ok_or_else(|| self.unknown_tool_error(name))?;
self.validators
.get(name)
.expect("registered tools have compiled schemas")
.validate(arguments)
.map_err(|error| format!("invalid tool arguments: {error}"))
}
pub fn validate_output(&self, name: &str, output: &Value) -> Result<(), String> {
let name = self
.canonical_name(name)
.ok_or_else(|| self.unknown_tool_error(name))?;
self.output_validators
.get(name)
.expect("registered tools have compiled output schemas")
.validate(output)
.map_err(|error| format!("invalid tool output: {error}"))
}
pub fn names(&self) -> Vec<&str> {
let mut names = self.tools.keys().map(String::as_str).collect::<Vec<_>>();
names.sort_unstable();
names
}
pub fn filtered<'a>(&self, allowed: impl IntoIterator<Item = &'a str>) -> Self {
let mut filtered = Self::new();
for name in allowed {
if let Some(tool) = self.tools.get(name) {
filtered
.register(Arc::clone(tool))
.expect("a subset of a valid registry remains valid");
}
}
filtered
}
pub fn specs(&self) -> Vec<LlmTool> {
let mut specs = self
.tools
.values()
.filter(|tool| tool.meta().surface == ToolSurface::Llm)
.map(|t| LlmTool::function(model_tool_name(t.name()), t.description(), t.parameters()))
.collect::<Vec<_>>();
specs.sort_by(|left, right| left.function.name.cmp(&right.function.name));
specs
}
pub fn runtime_manifest(&self) -> Result<Vec<Value>, String> {
self.names()
.into_iter()
.map(|name| {
let tool = self.tools.get(name).expect("name came from this registry");
let version = tool.implementation_version().trim();
if version.is_empty() {
return Err(format!(
"tool '{name}' requires a stable implementation version"
));
}
let meta = tool.meta();
Ok(serde_json::json!({
"name": name,
"implementation_version": version,
"description": tool.description(),
"parameters": tool.parameters(),
"output_schema": tool.output_schema(),
"surface": match meta.surface { ToolSurface::Llm => "llm", ToolSurface::Chassis => "chassis" },
"timeout_secs": meta.timeout_secs,
"concurrency": match meta.concurrency { ToolConcurrency::Concurrent => "concurrent", ToolConcurrency::Exclusive => "exclusive" },
"cost_units": meta.cost_units,
"core": meta.core,
}))
})
.collect()
}
pub fn suggest_name(&self, name: &str) -> Option<&str> {
let name = name.trim();
if name.is_empty() || self.tools.contains_key(name) {
return None;
}
if let Some((canonical, _)) = self
.tools
.iter()
.find(|(canonical, _)| canonical.eq_ignore_ascii_case(name))
{
return Some(canonical);
}
let lower = name.to_ascii_lowercase();
self.tools
.keys()
.filter(|canonical| {
name.len() > canonical.len() && lower.ends_with(&canonical.to_ascii_lowercase())
})
.max_by_key(|canonical| canonical.len())
.map(String::as_str)
}
fn unknown_tool_error(&self, name: &str) -> String {
let available = if self.tools.is_empty() {
"(none registered)".to_string()
} else {
self.names().join(", ")
};
match self.suggest_name(name) {
Some(suggestion) => format!(
"unknown tool '{name}'. Did you mean '{suggestion}'? Call tools by their exact registered name. Available: {available}"
),
None => format!(
"unknown tool '{name}'. Call one of the registered tools by exact name. Available: {available}"
),
}
}
pub async fn execute_with_context(
&self,
name: &str,
context: &ToolExecutionContext,
args: Value,
) -> Result<Value, String> {
if let Some(canonical) = self.canonical_name(name) {
self.validate_arguments(canonical, &args)?;
let tool = self
.tools
.get(canonical)
.expect("resolved tool names are registered");
let result = tool.call_with_context(context, args).await;
let value = result?;
self.validate_output(canonical, &value)?;
return Ok(value);
}
Err(self.unknown_tool_error(name))
}
pub fn canonical_name<'a>(&'a self, name: &'a str) -> Option<&'a str> {
if self.tools.contains_key(name) {
return Some(name);
}
self.wire_names.get(name).map(String::as_str)
}
}
pub fn validate_json_schema(schema: &Value, value: &Value) -> Result<(), String> {
validate_json_schema_value(schema, value, "tool arguments")
}
pub fn validate_json_schema_definition(schema: &Value) -> Result<(), String> {
jsonschema::validator_for(schema)
.map(|_| ())
.map_err(|error| format!("invalid JSON schema: {error}"))
}
pub fn validate_json_schema_value(
schema: &Value,
value: &Value,
subject: &str,
) -> Result<(), String> {
let validator = jsonschema::validator_for(schema)
.map_err(|error| format!("invalid JSON schema: {error}"))?;
validator
.validate(value)
.map_err(|error| format!("invalid {subject}: {error}"))
}
pub fn model_tool_name(internal: &str) -> String {
if !internal.is_empty()
&& internal.len() <= 64
&& internal
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
{
return internal.to_string();
}
let mut prefix = internal
.bytes()
.map(|byte| {
if byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-') {
byte as char
} else {
'_'
}
})
.take(47)
.collect::<String>();
if prefix.is_empty() {
prefix.push_str("tool");
}
let digest = format!("{:x}", Sha256::digest(internal.as_bytes()));
format!("{prefix}_{}", &digest[..16])
}
pub mod support {
use serde::de::DeserializeOwned;
use serde_json::Value;
pub trait RawToolSchema {
fn parameters() -> Value;
}
pub fn extract_required<T: DeserializeOwned>(args: &Value, name: &str) -> Result<T, String> {
let value = args
.get(name)
.cloned()
.ok_or_else(|| format!("missing required argument '{name}'"))?;
serde_json::from_value(value).map_err(|error| format!("invalid argument '{name}': {error}"))
}
pub fn extract_optional<T: DeserializeOwned>(
args: &Value,
name: &str,
) -> Result<Option<T>, String> {
match args.get(name) {
None | Some(Value::Null) => Ok(None),
Some(value) => serde_json::from_value(value.clone())
.map(Some)
.map_err(|error| format!("invalid argument '{name}': {error}")),
}
}
}
impl fmt::Debug for ToolRegistry {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ToolRegistry")
.field("tools", &self.names())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn execution() -> ToolExecutionContext {
ToolExecutionContext {
request: crate::RequestContext {
tenant_id: "tenant".into(),
subject_id: "subject".into(),
roles: Default::default(),
locale: "en".into(),
request_id: "request".into(),
entitlements: Default::default(),
},
session_id: "session".into(),
run_id: "run".into(),
step: 1,
call_id: "call".into(),
source_event_seq: 1,
interaction_resolution: None,
cancellation: CancellationToken::default(),
deadline: Instant::now() + std::time::Duration::from_secs(1),
}
}
struct TestTool(&'static str);
#[async_trait]
impl Tool for TestTool {
fn name(&self) -> &str {
self.0
}
fn description(&self) -> &str {
"test"
}
fn parameters(&self) -> Value {
serde_json::json!({
"type":"object",
"required":["items","mode"],
"additionalProperties":false,
"properties":{
"items":{"type":"array","items":{"type":"object","required":["id"],"properties":{"id":{"type":"integer"}}}},
"mode":{"enum":["safe","fast"]},
"version":{"const":1},
"choice":{"oneOf":[{"type":"string"},{"type":"number"}]}
}
})
}
fn output_schema(&self) -> Value {
serde_json::json!({"type":"object"})
}
async fn call(&self, args: Value) -> Result<Value, String> {
Ok(args)
}
}
#[tokio::test]
async fn registry_rejects_duplicates_and_validates_full_schema() {
let mut registry = ToolRegistry::new();
registry
.register(Arc::new(TestTool("nested.tool")))
.unwrap();
assert!(registry
.register(Arc::new(TestTool("nested.tool")))
.is_err());
let valid = serde_json::json!({"items":[{"id":1}],"mode":"safe","version":1,"choice":"x"});
assert_eq!(
registry
.execute_with_context("nested.tool", &execution(), valid.clone())
.await
.unwrap(),
valid
);
for invalid in [
serde_json::json!({"items":[{}],"mode":"safe"}),
serde_json::json!({"items":[{"id":1}],"mode":"unsafe"}),
serde_json::json!({"items":[{"id":1}],"mode":"safe","extra":true}),
serde_json::json!({"items":[{"id":1}],"mode":"safe","version":2}),
serde_json::json!({"items":[{"id":1}],"mode":"safe","choice":true}),
] {
assert!(registry
.execute_with_context("nested.tool", &execution(), invalid)
.await
.is_err());
}
}
#[test]
fn model_names_are_provider_safe_and_reversible() {
let mut registry = ToolRegistry::new();
registry.register(Arc::new(TestTool("namespace.tool with unicode-工具-and-a-name-that-is-far-too-long-for-provider-contracts"))).unwrap();
let spec = registry.specs().pop().unwrap().function.name;
assert!(spec.len() <= 64);
assert!(spec
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')));
assert_eq!(registry.canonical_name(&spec), Some("namespace.tool with unicode-工具-and-a-name-that-is-far-too-long-for-provider-contracts"));
let mut collision = ToolRegistry::new();
collision.register(Arc::new(TestTool("a.b"))).unwrap();
assert_eq!(model_tool_name("a.b"), "a_b_2e7336dc8eba87ef");
assert!(collision
.register(Arc::new(TestTool("a_b_2e7336dc8eba87ef")))
.is_err());
}
#[test]
fn invalid_schema_is_rejected_at_registration() {
struct Invalid;
#[async_trait]
impl Tool for Invalid {
fn name(&self) -> &str {
"invalid"
}
fn description(&self) -> &str {
"invalid"
}
fn parameters(&self) -> Value {
serde_json::json!({"type":"not-a-type"})
}
fn output_schema(&self) -> Value {
serde_json::json!({"type":"object"})
}
async fn call(&self, _: Value) -> Result<Value, String> {
Ok(Value::Null)
}
}
assert!(ToolRegistry::new().register(Arc::new(Invalid)).is_err());
}
#[tokio::test]
async fn successful_output_is_validated_before_materialization() {
struct InvalidOutput;
#[async_trait]
impl Tool for InvalidOutput {
fn name(&self) -> &str {
"invalid-output"
}
fn description(&self) -> &str {
"invalid output"
}
fn parameters(&self) -> Value {
serde_json::json!({"type":"object"})
}
fn output_schema(&self) -> Value {
serde_json::json!({"type":"object"})
}
async fn call(&self, _: Value) -> Result<Value, String> {
Ok(Value::String("bad".into()))
}
}
let mut registry = ToolRegistry::new();
registry.register(Arc::new(InvalidOutput)).unwrap();
assert!(registry
.execute_with_context("invalid-output", &execution(), serde_json::json!({}))
.await
.unwrap_err()
.contains("invalid tool output"));
}
}