af_agent/
tool.rs

1//! Agent tools + registry. Port of the reusable shape of `agent_core/tools/`.
2//!
3//! A tool exposes a JSON-Schema interface to the model and an async
4//! implementation. The registry resolves the name the model emits in a
5//! `tool_call` to an implementation, exposes the LLM-facing specs, and runs
6//! the call. Domain tools live in a product
7//! crate and register here.
8
9use std::collections::HashMap;
10use std::fmt;
11use std::sync::atomic::{AtomicBool, Ordering};
12use std::sync::Arc;
13use std::time::Instant;
14
15use async_trait::async_trait;
16use serde_json::Value;
17use sha2::{Digest, Sha256};
18
19use af_llm::Tool as LlmTool;
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum ToolSurface {
23    Llm,
24    Chassis,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum ToolConcurrency {
29    Concurrent,
30    Exclusive,
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub struct ToolMeta {
35    pub surface: ToolSurface,
36    pub cost_units: u64,
37    pub timeout_secs: u64,
38    pub core: bool,
39    pub concurrency: ToolConcurrency,
40}
41
42#[derive(Debug, Clone, Default)]
43pub struct CancellationToken(Arc<CancellationState>);
44
45#[derive(Debug, Default)]
46struct CancellationState {
47    cancelled: AtomicBool,
48    parent: Option<CancellationToken>,
49}
50
51impl CancellationToken {
52    pub fn cancel(&self) {
53        self.0.cancelled.store(true, Ordering::Release);
54    }
55    pub fn is_cancelled(&self) -> bool {
56        self.0.cancelled.load(Ordering::Acquire)
57            || self
58                .0
59                .parent
60                .as_ref()
61                .is_some_and(CancellationToken::is_cancelled)
62    }
63    pub fn child(&self) -> Self {
64        Self(Arc::new(CancellationState {
65            cancelled: AtomicBool::new(false),
66            parent: Some(self.clone()),
67        }))
68    }
69}
70
71#[derive(Debug, Clone)]
72pub struct ToolExecutionContext {
73    pub request: af_context::RequestContext,
74    pub session_id: String,
75    pub run_id: String,
76    pub step: u32,
77    pub call_id: String,
78    pub source_event_seq: u64,
79    pub interaction_resolution: Option<af_agent_session::InteractionResolution>,
80    pub cancellation: CancellationToken,
81    pub deadline: Instant,
82}
83
84impl Default for ToolMeta {
85    fn default() -> Self {
86        Self {
87            surface: ToolSurface::Llm,
88            cost_units: 1,
89            timeout_secs: 15,
90            core: false,
91            concurrency: ToolConcurrency::Exclusive,
92        }
93    }
94}
95
96/// A callable tool. `parameters` is a JSON-Schema object describing the args.
97#[async_trait]
98pub trait Tool: Send + Sync {
99    fn name(&self) -> &str;
100    fn implementation_version(&self) -> &str {
101        ""
102    }
103    fn description(&self) -> &str;
104    fn parameters(&self) -> Value;
105    fn output_schema(&self) -> Value;
106    fn meta(&self) -> ToolMeta {
107        ToolMeta::default()
108    }
109
110    /// Run the tool. On success return a JSON value; on failure return a short
111    /// error string (surfaced back to the model as `{"error": ...}` so it can
112    /// recover) — mirrors the Python "success → dict / failure → {error}" rule.
113    async fn call(&self, args: Value) -> Result<Value, String>;
114    async fn call_with_context(
115        &self,
116        _context: &ToolExecutionContext,
117        args: Value,
118    ) -> Result<Value, String> {
119        self.call(args).await
120    }
121}
122
123/// Registry of tools available to an agent.
124#[derive(Default, Clone)]
125pub struct ToolRegistry {
126    tools: HashMap<String, Arc<dyn Tool>>,
127    wire_names: HashMap<String, String>,
128    validators: HashMap<String, Arc<jsonschema::Validator>>,
129    output_validators: HashMap<String, Arc<jsonschema::Validator>>,
130}
131
132impl ToolRegistry {
133    pub fn new() -> Self {
134        Self::default()
135    }
136
137    pub fn register(&mut self, tool: Arc<dyn Tool>) -> Result<&mut Self, String> {
138        let name = tool.name().to_string();
139        if self.tools.contains_key(&name) {
140            return Err(format!("duplicate tool '{name}'"));
141        }
142        let wire_name = model_tool_name(&name);
143        if self.wire_names.contains_key(&wire_name)
144            || (wire_name != name && self.tools.contains_key(&wire_name))
145            || self.wire_names.contains_key(&name)
146        {
147            return Err(format!(
148                "tool name '{name}' collides on provider name '{wire_name}'"
149            ));
150        }
151        let validator = jsonschema::validator_for(&tool.parameters())
152            .map_err(|error| format!("invalid schema for tool '{name}': {error}"))?;
153        let output_validator = jsonschema::validator_for(&tool.output_schema())
154            .map_err(|error| format!("invalid output schema for tool '{name}': {error}"))?;
155        self.tools.insert(name.clone(), tool);
156        self.wire_names.insert(wire_name, name.clone());
157        self.validators.insert(name.clone(), Arc::new(validator));
158        self.output_validators
159            .insert(name, Arc::new(output_validator));
160        Ok(self)
161    }
162
163    pub fn extend(&mut self, other: &Self) -> Result<(), String> {
164        for tool in other.tools.values() {
165            self.register(Arc::clone(tool))?;
166        }
167        Ok(())
168    }
169
170    pub fn is_empty(&self) -> bool {
171        self.tools.is_empty()
172    }
173
174    pub fn len(&self) -> usize {
175        self.tools.len()
176    }
177
178    pub fn contains(&self, name: &str) -> bool {
179        self.tools.contains_key(name)
180    }
181
182    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
183        self.canonical_name(name)
184            .and_then(|name| self.tools.get(name))
185            .cloned()
186    }
187
188    pub fn validate_arguments(&self, name: &str, arguments: &Value) -> Result<(), String> {
189        let name = self
190            .canonical_name(name)
191            .ok_or_else(|| self.unknown_tool_error(name))?;
192        self.validators
193            .get(name)
194            .expect("registered tools have compiled schemas")
195            .validate(arguments)
196            .map_err(|error| format!("invalid tool arguments: {error}"))
197    }
198
199    pub fn validate_output(&self, name: &str, output: &Value) -> Result<(), String> {
200        let name = self
201            .canonical_name(name)
202            .ok_or_else(|| self.unknown_tool_error(name))?;
203        self.output_validators
204            .get(name)
205            .expect("registered tools have compiled output schemas")
206            .validate(output)
207            .map_err(|error| format!("invalid tool output: {error}"))
208    }
209
210    /// Stable tool names for validation, diagnostics, and capability policy.
211    pub fn names(&self) -> Vec<&str> {
212        let mut names = self.tools.keys().map(String::as_str).collect::<Vec<_>>();
213        names.sort_unstable();
214        names
215    }
216
217    /// Return a registry containing only explicitly allowed tools.
218    pub fn filtered<'a>(&self, allowed: impl IntoIterator<Item = &'a str>) -> Self {
219        let mut filtered = Self::new();
220        for name in allowed {
221            if let Some(tool) = self.tools.get(name) {
222                filtered
223                    .register(Arc::clone(tool))
224                    .expect("a subset of a valid registry remains valid");
225            }
226        }
227        filtered
228    }
229
230    /// LLM-facing tool specs for the `tools` field of a completion request.
231    pub fn specs(&self) -> Vec<LlmTool> {
232        let mut specs = self
233            .tools
234            .values()
235            .filter(|tool| tool.meta().surface == ToolSurface::Llm)
236            .map(|t| LlmTool::function(model_tool_name(t.name()), t.description(), t.parameters()))
237            .collect::<Vec<_>>();
238        specs.sort_by(|left, right| left.function.name.cmp(&right.function.name));
239        specs
240    }
241
242    /// Stable implementation and execution contract persisted by profile revisions.
243    pub fn runtime_manifest(&self) -> Result<Vec<Value>, String> {
244        self.names()
245            .into_iter()
246            .map(|name| {
247                let tool = self.tools.get(name).expect("name came from this registry");
248                let version = tool.implementation_version().trim();
249                if version.is_empty() {
250                    return Err(format!(
251                        "tool '{name}' requires a stable implementation version"
252                    ));
253                }
254                let meta = tool.meta();
255                Ok(serde_json::json!({
256                    "name": name,
257                    "implementation_version": version,
258                    "description": tool.description(),
259                    "parameters": tool.parameters(),
260                    "output_schema": tool.output_schema(),
261                    "surface": match meta.surface { ToolSurface::Llm => "llm", ToolSurface::Chassis => "chassis" },
262                    "timeout_secs": meta.timeout_secs,
263                    "concurrency": match meta.concurrency { ToolConcurrency::Concurrent => "concurrent", ToolConcurrency::Exclusive => "exclusive" },
264                    "cost_units": meta.cost_units,
265                    "core": meta.core,
266                }))
267            })
268            .collect()
269    }
270
271    /// Best-effort canonical name for case mistakes or an alphabetic junk
272    /// prefix glued to a registered tool name.
273    pub fn suggest_name(&self, name: &str) -> Option<&str> {
274        let name = name.trim();
275        if name.is_empty() || self.tools.contains_key(name) {
276            return None;
277        }
278        if let Some((canonical, _)) = self
279            .tools
280            .iter()
281            .find(|(canonical, _)| canonical.eq_ignore_ascii_case(name))
282        {
283            return Some(canonical);
284        }
285
286        let lower = name.to_ascii_lowercase();
287        self.tools
288            .keys()
289            .filter(|canonical| {
290                name.len() > canonical.len() && lower.ends_with(&canonical.to_ascii_lowercase())
291            })
292            .max_by_key(|canonical| canonical.len())
293            .map(String::as_str)
294    }
295
296    fn unknown_tool_error(&self, name: &str) -> String {
297        let available = if self.tools.is_empty() {
298            "(none registered)".to_string()
299        } else {
300            self.names().join(", ")
301        };
302        match self.suggest_name(name) {
303            Some(suggestion) => format!(
304                "unknown tool '{name}'. Did you mean '{suggestion}'? Call tools by their exact registered name. Available: {available}"
305            ),
306            None => format!(
307                "unknown tool '{name}'. Call one of the registered tools by exact name. Available: {available}"
308            ),
309        }
310    }
311
312    pub async fn execute_with_context(
313        &self,
314        name: &str,
315        context: &ToolExecutionContext,
316        args: Value,
317    ) -> Result<Value, String> {
318        if let Some(canonical) = self.canonical_name(name) {
319            self.validate_arguments(canonical, &args)?;
320            let tool = self
321                .tools
322                .get(canonical)
323                .expect("resolved tool names are registered");
324            let result = tool.call_with_context(context, args).await;
325            let value = result?;
326            self.validate_output(canonical, &value)?;
327            return Ok(value);
328        }
329        Err(self.unknown_tool_error(name))
330    }
331
332    pub fn canonical_name<'a>(&'a self, name: &'a str) -> Option<&'a str> {
333        if self.tools.contains_key(name) {
334            return Some(name);
335        }
336        self.wire_names.get(name).map(String::as_str)
337    }
338}
339
340pub fn validate_json_schema(schema: &Value, value: &Value) -> Result<(), String> {
341    validate_json_schema_value(schema, value, "tool arguments")
342}
343
344pub fn validate_json_schema_definition(schema: &Value) -> Result<(), String> {
345    jsonschema::validator_for(schema)
346        .map(|_| ())
347        .map_err(|error| format!("invalid JSON schema: {error}"))
348}
349
350pub fn validate_json_schema_value(
351    schema: &Value,
352    value: &Value,
353    subject: &str,
354) -> Result<(), String> {
355    let validator = jsonschema::validator_for(schema)
356        .map_err(|error| format!("invalid JSON schema: {error}"))?;
357    validator
358        .validate(value)
359        .map_err(|error| format!("invalid {subject}: {error}"))
360}
361
362pub fn model_tool_name(internal: &str) -> String {
363    if !internal.is_empty()
364        && internal.len() <= 64
365        && internal
366            .bytes()
367            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
368    {
369        return internal.to_string();
370    }
371    let mut prefix = internal
372        .bytes()
373        .map(|byte| {
374            if byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-') {
375                byte as char
376            } else {
377                '_'
378            }
379        })
380        .take(47)
381        .collect::<String>();
382    if prefix.is_empty() {
383        prefix.push_str("tool");
384    }
385    let digest = format!("{:x}", Sha256::digest(internal.as_bytes()));
386    format!("{prefix}_{}", &digest[..16])
387}
388
389pub mod support {
390    use serde::de::DeserializeOwned;
391    use serde_json::Value;
392
393    pub trait RawToolSchema {
394        fn parameters() -> Value;
395    }
396
397    pub fn extract_required<T: DeserializeOwned>(args: &Value, name: &str) -> Result<T, String> {
398        let value = args
399            .get(name)
400            .cloned()
401            .ok_or_else(|| format!("missing required argument '{name}'"))?;
402        serde_json::from_value(value).map_err(|error| format!("invalid argument '{name}': {error}"))
403    }
404
405    pub fn extract_optional<T: DeserializeOwned>(
406        args: &Value,
407        name: &str,
408    ) -> Result<Option<T>, String> {
409        match args.get(name) {
410            None | Some(Value::Null) => Ok(None),
411            Some(value) => serde_json::from_value(value.clone())
412                .map(Some)
413                .map_err(|error| format!("invalid argument '{name}': {error}")),
414        }
415    }
416}
417
418impl fmt::Debug for ToolRegistry {
419    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
420        f.debug_struct("ToolRegistry")
421            .field("tools", &self.names())
422            .finish()
423    }
424}
425
426#[cfg(test)]
427mod tests {
428    use super::*;
429
430    fn execution() -> ToolExecutionContext {
431        ToolExecutionContext {
432            request: crate::RequestContext {
433                tenant_id: "tenant".into(),
434                subject_id: "subject".into(),
435                roles: Default::default(),
436                locale: "en".into(),
437                request_id: "request".into(),
438                entitlements: Default::default(),
439            },
440            session_id: "session".into(),
441            run_id: "run".into(),
442            step: 1,
443            call_id: "call".into(),
444            source_event_seq: 1,
445            interaction_resolution: None,
446            cancellation: CancellationToken::default(),
447            deadline: Instant::now() + std::time::Duration::from_secs(1),
448        }
449    }
450
451    struct TestTool(&'static str);
452    #[async_trait]
453    impl Tool for TestTool {
454        fn name(&self) -> &str {
455            self.0
456        }
457        fn description(&self) -> &str {
458            "test"
459        }
460        fn parameters(&self) -> Value {
461            serde_json::json!({
462                "type":"object",
463                "required":["items","mode"],
464                "additionalProperties":false,
465                "properties":{
466                    "items":{"type":"array","items":{"type":"object","required":["id"],"properties":{"id":{"type":"integer"}}}},
467                    "mode":{"enum":["safe","fast"]},
468                    "version":{"const":1},
469                    "choice":{"oneOf":[{"type":"string"},{"type":"number"}]}
470                }
471            })
472        }
473        fn output_schema(&self) -> Value {
474            serde_json::json!({"type":"object"})
475        }
476        async fn call(&self, args: Value) -> Result<Value, String> {
477            Ok(args)
478        }
479    }
480
481    #[tokio::test]
482    async fn registry_rejects_duplicates_and_validates_full_schema() {
483        let mut registry = ToolRegistry::new();
484        registry
485            .register(Arc::new(TestTool("nested.tool")))
486            .unwrap();
487        assert!(registry
488            .register(Arc::new(TestTool("nested.tool")))
489            .is_err());
490        let valid = serde_json::json!({"items":[{"id":1}],"mode":"safe","version":1,"choice":"x"});
491        assert_eq!(
492            registry
493                .execute_with_context("nested.tool", &execution(), valid.clone())
494                .await
495                .unwrap(),
496            valid
497        );
498        for invalid in [
499            serde_json::json!({"items":[{}],"mode":"safe"}),
500            serde_json::json!({"items":[{"id":1}],"mode":"unsafe"}),
501            serde_json::json!({"items":[{"id":1}],"mode":"safe","extra":true}),
502            serde_json::json!({"items":[{"id":1}],"mode":"safe","version":2}),
503            serde_json::json!({"items":[{"id":1}],"mode":"safe","choice":true}),
504        ] {
505            assert!(registry
506                .execute_with_context("nested.tool", &execution(), invalid)
507                .await
508                .is_err());
509        }
510    }
511
512    #[test]
513    fn model_names_are_provider_safe_and_reversible() {
514        let mut registry = ToolRegistry::new();
515        registry.register(Arc::new(TestTool("namespace.tool with unicode-工具-and-a-name-that-is-far-too-long-for-provider-contracts"))).unwrap();
516        let spec = registry.specs().pop().unwrap().function.name;
517        assert!(spec.len() <= 64);
518        assert!(spec
519            .bytes()
520            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')));
521        assert_eq!(registry.canonical_name(&spec), Some("namespace.tool with unicode-工具-and-a-name-that-is-far-too-long-for-provider-contracts"));
522
523        let mut collision = ToolRegistry::new();
524        collision.register(Arc::new(TestTool("a.b"))).unwrap();
525        assert_eq!(model_tool_name("a.b"), "a_b_2e7336dc8eba87ef");
526        assert!(collision
527            .register(Arc::new(TestTool("a_b_2e7336dc8eba87ef")))
528            .is_err());
529    }
530
531    #[test]
532    fn invalid_schema_is_rejected_at_registration() {
533        struct Invalid;
534        #[async_trait]
535        impl Tool for Invalid {
536            fn name(&self) -> &str {
537                "invalid"
538            }
539            fn description(&self) -> &str {
540                "invalid"
541            }
542            fn parameters(&self) -> Value {
543                serde_json::json!({"type":"not-a-type"})
544            }
545            fn output_schema(&self) -> Value {
546                serde_json::json!({"type":"object"})
547            }
548            async fn call(&self, _: Value) -> Result<Value, String> {
549                Ok(Value::Null)
550            }
551        }
552        assert!(ToolRegistry::new().register(Arc::new(Invalid)).is_err());
553    }
554
555    #[tokio::test]
556    async fn successful_output_is_validated_before_materialization() {
557        struct InvalidOutput;
558        #[async_trait]
559        impl Tool for InvalidOutput {
560            fn name(&self) -> &str {
561                "invalid-output"
562            }
563            fn description(&self) -> &str {
564                "invalid output"
565            }
566            fn parameters(&self) -> Value {
567                serde_json::json!({"type":"object"})
568            }
569            fn output_schema(&self) -> Value {
570                serde_json::json!({"type":"object"})
571            }
572            async fn call(&self, _: Value) -> Result<Value, String> {
573                Ok(Value::String("bad".into()))
574            }
575        }
576        let mut registry = ToolRegistry::new();
577        registry.register(Arc::new(InvalidOutput)).unwrap();
578        assert!(registry
579            .execute_with_context("invalid-output", &execution(), serde_json::json!({}))
580            .await
581            .unwrap_err()
582            .contains("invalid tool output"));
583    }
584}
af_agent/tool.rs

af_agent/
tool.rs
