Skip to main content

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 af_context::{RunId, SessionId, ToolCallId};
10use std::collections::HashMap;
11use std::fmt;
12use std::future::Future;
13use std::pin::Pin;
14use std::sync::atomic::{AtomicBool, Ordering};
15use std::sync::Arc;
16use std::time::Instant;
17
18use async_trait::async_trait;
19use serde_json::Value;
20use sha2::{Digest, Sha256};
21
22use af_llm::Tool as LlmTool;
23
24/// Who may call a tool.
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum ToolSurface {
27    /// Advertised to the model.
28    Llm,
29    /// Only the host or platform.
30    Chassis,
31}
32
33/// Whether a tool may run alongside others in one step.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum ToolConcurrency {
36    /// May run in parallel.
37    Concurrent,
38    /// Runs alone; acts as a barrier.
39    Exclusive,
40}
41
42/// User-visible risk class used by the conversation approval policy.
43#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
44#[serde(rename_all = "snake_case")]
45pub enum ToolRiskClass {
46    /// Missing metadata; execution must fail closed.
47    #[default]
48    Unclassified,
49    /// No external or durable mutation.
50    Read,
51    /// Reversible or bounded mutation permitted by `low_risk_auto`.
52    LowRiskWrite,
53    /// Mutation that always requires explicit confirmation.
54    ConfirmedWrite,
55    /// Money movement or equivalent funds authority.
56    Funds,
57}
58
59impl ToolRiskClass {
60    /// Stable snake_case manifest value.
61    pub const fn as_str(self) -> &'static str {
62        match self {
63            Self::Unclassified => "unclassified",
64            Self::Read => "read",
65            Self::LowRiskWrite => "low_risk_write",
66            Self::ConfirmedWrite => "confirmed_write",
67            Self::Funds => "funds",
68        }
69    }
70}
71
72/// Component that performs a tool's effect after runtime authorization.
73#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
74#[serde(rename_all = "snake_case")]
75pub enum ToolExecutionOwner {
76    /// Missing metadata; execution must fail closed.
77    #[default]
78    Unclassified,
79    /// The shared Factory runtime owns the operation.
80    Runtime,
81    /// The consuming product host owns the operation.
82    ProductHost,
83}
84
85impl ToolExecutionOwner {
86    /// Stable snake_case manifest value.
87    pub const fn as_str(self) -> &'static str {
88        match self {
89            Self::Unclassified => "unclassified",
90            Self::Runtime => "runtime",
91            Self::ProductHost => "product_host",
92        }
93    }
94}
95
96/// Static execution metadata for a tool.
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98pub struct ToolMeta {
99    /// Who may call it.
100    pub surface: ToolSurface,
101    /// Provider-neutral cost units settled for this operation.
102    pub cost_units: u64,
103    /// Per-call timeout.
104    pub timeout_secs: u64,
105    /// Whether the tool is part of the platform core (never narrowed away).
106    pub core: bool,
107    /// Concurrency class.
108    pub concurrency: ToolConcurrency,
109    /// Explicit approval risk classification.
110    pub risk_class: ToolRiskClass,
111    /// Component that owns execution after authorization.
112    pub execution_owner: ToolExecutionOwner,
113    /// Mutating tools must cross the runtime's normal confirmation hook even
114    /// when a profile forgot to list them explicitly.
115    pub requires_confirmation: bool,
116}
117
118/// Actual usage reported by one completed tool execution.
119#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
120pub struct ToolUsage {
121    /// Model input tokens consumed inside the tool.
122    pub prompt_tokens: u64,
123    /// Model output tokens consumed inside the tool.
124    pub completion_tokens: u64,
125    /// Provider-neutral non-token cost units.
126    pub cost_units: u64,
127}
128
129/// Hierarchical cancellation flag. Cancelling a parent cancels every child;
130/// [`cancelled`](Self::cancelled) resolves without polling.
131#[derive(Debug, Clone, Default)]
132pub struct CancellationToken(Arc<CancellationState>);
133
134#[derive(Debug, Default)]
135struct CancellationState {
136    cancelled: AtomicBool,
137    notify: tokio::sync::Notify,
138    parent: Option<CancellationToken>,
139}
140
141impl CancellationToken {
142    /// Cancel this token and every child.
143    pub fn cancel(&self) {
144        self.0.cancelled.store(true, Ordering::Release);
145        self.0.notify.notify_waiters();
146    }
147    /// Whether this token or an ancestor was cancelled.
148    pub fn is_cancelled(&self) -> bool {
149        self.0.cancelled.load(Ordering::Acquire)
150            || self
151                .0
152                .parent
153                .as_ref()
154                .is_some_and(CancellationToken::is_cancelled)
155    }
156    /// A child token cancelled with its parent.
157    pub fn child(&self) -> Self {
158        Self(Arc::new(CancellationState {
159            cancelled: AtomicBool::new(false),
160            notify: tokio::sync::Notify::new(),
161            parent: Some(self.clone()),
162        }))
163    }
164    /// Resolves once this token or any ancestor is cancelled.
165    pub fn cancelled(&self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
166        Box::pin(async move {
167            let notified = self.0.notify.notified();
168            tokio::pin!(notified);
169            notified.as_mut().enable();
170            if self.is_cancelled() {
171                return;
172            }
173            match &self.0.parent {
174                Some(parent) => tokio::select! {
175                    _ = notified => {}
176                    _ = parent.cancelled() => {}
177                },
178                None => notified.await,
179            }
180        })
181    }
182}
183
184/// Identity, position and controls for one tool execution.
185#[derive(Debug, Clone)]
186pub struct ToolExecutionContext {
187    /// The originating request.
188    pub request: af_context::RequestContext,
189    /// Session this record belongs to.
190    pub session_id: SessionId,
191    /// Run this record belongs to.
192    pub run_id: RunId,
193    /// 1-based step number inside the Turn.
194    pub step: u32,
195    /// Tool call this record refers to.
196    pub call_id: ToolCallId,
197    /// Sequence of the Session event this was derived from.
198    pub source_event_seq: u64,
199    /// Explicit approval risk classification of the registered tool.
200    pub risk_class: ToolRiskClass,
201    /// Component that owns execution after authorization.
202    pub execution_owner: ToolExecutionOwner,
203    /// Resolution of the interaction that unblocked this call, if any.
204    pub interaction_resolution: Option<af_agent_session::InteractionResolution>,
205    /// Cancelled when the Run or step is abandoned.
206    pub cancellation: CancellationToken,
207    /// Latest time by which the work must finish.
208    pub deadline: Instant,
209}
210
211/// A chassis tool call derived from a successful tool result.
212#[derive(Debug, Clone, PartialEq)]
213pub struct ToolCompletionAction {
214    /// Registered tool to call through the normal authorization boundary.
215    pub tool: String,
216    /// Arguments validated against that tool's schema.
217    pub arguments: Value,
218}
219
220impl Default for ToolMeta {
221    fn default() -> Self {
222        Self {
223            surface: ToolSurface::Llm,
224            cost_units: 1,
225            timeout_secs: 15,
226            core: false,
227            concurrency: ToolConcurrency::Exclusive,
228            risk_class: ToolRiskClass::Unclassified,
229            execution_owner: ToolExecutionOwner::Unclassified,
230            requires_confirmation: false,
231        }
232    }
233}
234
235/// A callable tool. `parameters` is a JSON-Schema object describing the args.
236#[async_trait]
237pub trait Tool: Send + Sync {
238    /// Stable internal name.
239    fn name(&self) -> &str;
240    /// Implementation version pinned by Profile revisions.
241    fn implementation_version(&self) -> &str {
242        ""
243    }
244    /// Model-facing description.
245    fn description(&self) -> &str;
246    /// JSON Schema for the arguments.
247    fn parameters(&self) -> Value;
248    /// JSON Schema for the result.
249    fn output_schema(&self) -> Value;
250    /// Execution metadata.
251    fn meta(&self) -> ToolMeta {
252        ToolMeta::default()
253    }
254    /// Whether an authenticated, private response artifact may be created in
255    /// `confirm_changes` mode without an extra prompt. Never use for external
256    /// side effects, funds, or writes not already explicit in the user request.
257    fn auto_approve_response_artifact(&self) -> bool {
258        false
259    }
260    /// Whether successful calls with identical normalized arguments may be reused in one Run.
261    /// The runtime additionally requires [`ToolRiskClass::Read`].
262    fn reuse_within_run(&self) -> bool {
263        false
264    }
265    /// Whether results larger than the inline threshold may use the host artifact store.
266    fn artifact_eligible(&self) -> bool {
267        false
268    }
269    /// Whether the tool remains model-visible while the runtime produces a conclusion.
270    fn available_during_conclusion(&self) -> bool {
271        false
272    }
273    /// Return the actual usage for a successful result. Tools with dynamic
274    /// provider usage can override this; static tools inherit `meta.cost_units`.
275    fn usage(&self, _result: &Value) -> Option<ToolUsage> {
276        None
277    }
278
279    /// Run the tool. On success return a JSON value; on failure return a short
280    /// error string (surfaced back to the model as `{"error": ...}` so it can
281    /// recover) — mirrors the Python "success → dict / failure → {error}" rule.
282    async fn call(&self, args: Value) -> Result<Value, String>;
283    /// Like [`call`](Self::call) with the execution context; the default ignores the context.
284    async fn call_with_context(
285        &self,
286        _context: &ToolExecutionContext,
287        args: Value,
288    ) -> Result<Value, String> {
289        self.call(args).await
290    }
291    /// Optionally request a chassis action after a successful validated result.
292    /// The runtime persists and authorizes the returned call before proceeding.
293    fn completion_action(&self, _result: &Value) -> Option<ToolCompletionAction> {
294        None
295    }
296    /// Optional bounded, domain-shaped view of a validated large result for the model.
297    /// The runtime persists the complete result before invoking this projection.
298    fn project_artifact_result(&self, _result: &Value) -> Result<Option<Value>, String> {
299        Ok(None)
300    }
301}
302
303/// Registry of tools available to an agent.
304#[derive(Default, Clone)]
305pub struct ToolRegistry {
306    tools: HashMap<String, Arc<dyn Tool>>,
307    wire_names: HashMap<String, String>,
308    validators: HashMap<String, Arc<jsonschema::Validator>>,
309    output_validators: HashMap<String, Arc<jsonschema::Validator>>,
310}
311
312impl ToolRegistry {
313    /// An empty registry.
314    pub fn new() -> Self {
315        Self::default()
316    }
317
318    /// Register a tool, compiling its schemas; duplicate or invalid tools are rejected.
319    pub fn register(&mut self, tool: Arc<dyn Tool>) -> Result<&mut Self, String> {
320        let name = tool.name().to_string();
321        if self.tools.contains_key(&name) {
322            return Err(format!("duplicate tool '{name}'"));
323        }
324        let wire_name = model_tool_name(&name);
325        if self.wire_names.contains_key(&wire_name)
326            || (wire_name != name && self.tools.contains_key(&wire_name))
327            || self.wire_names.contains_key(&name)
328        {
329            return Err(format!(
330                "tool name '{name}' collides on provider name '{wire_name}'"
331            ));
332        }
333        let validator = jsonschema::validator_for(&tool.parameters())
334            .map_err(|error| format!("invalid schema for tool '{name}': {error}"))?;
335        let output_validator = jsonschema::validator_for(&tool.output_schema())
336            .map_err(|error| format!("invalid output schema for tool '{name}': {error}"))?;
337        self.tools.insert(name.clone(), tool);
338        self.wire_names.insert(wire_name, name.clone());
339        self.validators.insert(name.clone(), Arc::new(validator));
340        self.output_validators
341            .insert(name, Arc::new(output_validator));
342        Ok(self)
343    }
344
345    /// Merge another registry; conflicting names are rejected.
346    pub fn extend(&mut self, other: &Self) -> Result<(), String> {
347        for tool in other.tools.values() {
348            self.register(Arc::clone(tool))?;
349        }
350        Ok(())
351    }
352
353    /// Whether no tools are registered.
354    pub fn is_empty(&self) -> bool {
355        self.tools.is_empty()
356    }
357
358    /// Number of registered tools.
359    pub fn len(&self) -> usize {
360        self.tools.len()
361    }
362
363    /// Whether `name` (internal or model form) is registered.
364    pub fn contains(&self, name: &str) -> bool {
365        self.tools.contains_key(name)
366    }
367
368    /// Look up a tool by internal or model name.
369    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
370        self.canonical_name(name)
371            .and_then(|name| self.tools.get(name))
372            .cloned()
373    }
374
375    /// Validate arguments against the tool's schema.
376    pub fn validate_arguments(&self, name: &str, arguments: &Value) -> Result<(), String> {
377        let name = self
378            .canonical_name(name)
379            .ok_or_else(|| self.unknown_tool_error(name))?;
380        self.validators
381            .get(name)
382            .ok_or_else(|| self.unknown_tool_error(name))?
383            .validate(arguments)
384            .map_err(|error| format!("invalid tool arguments: {error}"))
385    }
386
387    /// Resolve a provider-returned name, repairing it only when the arguments
388    /// uniquely match one registered tool schema. Authorization still uses the
389    /// resolved tool's own policy.
390    pub fn resolve_call_name<'a>(
391        &'a self,
392        name: &str,
393        arguments: &Value,
394    ) -> Result<&'a str, String> {
395        let original_error = self.validate_arguments(name, arguments).err();
396        if original_error.is_none() {
397            let canonical = self
398                .canonical_name(name)
399                .ok_or_else(|| self.unknown_tool_error(name))?;
400            return self
401                .tools
402                .get_key_value(canonical)
403                .map(|(registered, _)| registered.as_str())
404                .ok_or_else(|| self.unknown_tool_error(name));
405        }
406        let argument_keys = arguments.as_object().map(|object| object.keys());
407        let mut matches = self
408            .tools
409            .iter()
410            .filter(|(candidate, tool)| {
411                self.validate_arguments(candidate, arguments).is_ok()
412                    && argument_keys.clone().is_some_and(|keys| {
413                        let schema = tool.parameters();
414                        let properties = schema.get("properties").and_then(Value::as_object);
415                        keys.into_iter().all(|key| {
416                            properties.is_some_and(|properties| properties.contains_key(key))
417                        })
418                    })
419            })
420            .map(|(candidate, _)| candidate);
421        let candidate = matches.next();
422        match (candidate, matches.next()) {
423            (Some(candidate), None) => Ok(candidate),
424            _ => Err(original_error.expect("invalid original call has an error")),
425        }
426    }
427
428    /// Validate a result against the tool's output schema.
429    pub fn validate_output(&self, name: &str, output: &Value) -> Result<(), String> {
430        let name = self
431            .canonical_name(name)
432            .ok_or_else(|| self.unknown_tool_error(name))?;
433        self.output_validators
434            .get(name)
435            .ok_or_else(|| self.unknown_tool_error(name))?
436            .validate(output)
437            .map_err(|error| format!("invalid tool output: {error}"))
438    }
439
440    /// Resolve a completed tool's actual usage.
441    pub fn usage(&self, name: &str, output: &Value) -> Option<ToolUsage> {
442        self.get(name).and_then(|tool| tool.usage(output))
443    }
444
445    /// Stable tool names for validation, diagnostics, and capability policy.
446    pub fn names(&self) -> Vec<&str> {
447        let mut names = self.tools.keys().map(String::as_str).collect::<Vec<_>>();
448        names.sort_unstable();
449        names
450    }
451
452    /// Tools whose metadata requires confirmation.
453    pub fn confirmation_required_names(&self) -> impl Iterator<Item = &str> {
454        self.tools
455            .values()
456            .filter(|tool| tool.meta().requires_confirmation)
457            .map(|tool| tool.name())
458    }
459
460    /// Return a registry containing only explicitly allowed tools.
461    pub fn filtered<'a>(&self, allowed: impl IntoIterator<Item = &'a str>) -> Self {
462        let mut filtered = Self::new();
463        for name in allowed {
464            if let Some(tool) = self.tools.get(name) {
465                // A subset of an already-validated registry cannot introduce a
466                // conflict, so a failure here is unreachable and safely ignored.
467                let _ = filtered.register(Arc::clone(tool));
468            }
469        }
470        filtered
471    }
472
473    /// LLM-facing tool specs for the `tools` field of a completion request.
474    pub fn specs(&self) -> Vec<LlmTool> {
475        let mut specs = self
476            .tools
477            .values()
478            .filter(|tool| tool.meta().surface == ToolSurface::Llm)
479            .map(|t| LlmTool::function(model_tool_name(t.name()), t.description(), t.parameters()))
480            .collect::<Vec<_>>();
481        specs.sort_by(|left, right| left.function.name.cmp(&right.function.name));
482        specs
483    }
484
485    /// Stable implementation and execution contract persisted by profile revisions.
486    pub fn runtime_manifest(&self) -> Result<Vec<Value>, String> {
487        self.names()
488            .into_iter()
489            .map(|name| {
490                let tool = self
491                    .tools
492                    .get(name)
493                    .ok_or_else(|| self.unknown_tool_error(name))?;
494                let version = tool.implementation_version().trim();
495                if version.is_empty() {
496                    return Err(format!(
497                        "tool '{name}' requires a stable implementation version"
498                    ));
499                }
500                let meta = tool.meta();
501                if meta.risk_class == ToolRiskClass::Unclassified
502                    || meta.execution_owner == ToolExecutionOwner::Unclassified
503                {
504                    return Err(format!(
505                        "tool '{name}' requires explicit risk_class and execution_owner"
506                    ));
507                }
508                Ok(serde_json::json!({
509                    "name": name,
510                    "implementation_version": version,
511                    "description": tool.description(),
512                    "parameters": tool.parameters(),
513                    "output_schema": tool.output_schema(),
514                    "surface": match meta.surface { ToolSurface::Llm => "llm", ToolSurface::Chassis => "chassis" },
515                    "timeout_secs": meta.timeout_secs,
516                    "concurrency": match meta.concurrency { ToolConcurrency::Concurrent => "concurrent", ToolConcurrency::Exclusive => "exclusive" },
517                    "cost_units": meta.cost_units,
518                    "core": meta.core,
519                    "risk_class": meta.risk_class.as_str(),
520                    "execution_owner": meta.execution_owner.as_str(),
521                    "requires_confirmation": meta.requires_confirmation,
522                    "reuse_within_run": tool.reuse_within_run(),
523                    "artifact_eligible": tool.artifact_eligible(),
524                    "available_during_conclusion": tool.available_during_conclusion(),
525                }))
526            })
527            .collect()
528    }
529
530    /// Best-effort canonical name for case mistakes or an alphabetic junk
531    /// prefix glued to a registered tool name.
532    pub fn suggest_name(&self, name: &str) -> Option<&str> {
533        let name = name.trim();
534        if name.is_empty() || self.tools.contains_key(name) {
535            return None;
536        }
537        if let Some((canonical, _)) = self
538            .tools
539            .iter()
540            .find(|(canonical, _)| canonical.eq_ignore_ascii_case(name))
541        {
542            return Some(canonical);
543        }
544
545        let lower = name.to_ascii_lowercase();
546        self.tools
547            .keys()
548            .filter(|canonical| {
549                let canonical = canonical.to_ascii_lowercase();
550                let prefix_len = lower.len().checked_sub(canonical.len());
551                let plural_prefix_len = lower
552                    .strip_suffix('s')
553                    .and_then(|singular| singular.len().checked_sub(canonical.len()))
554                    .filter(|_| lower[..lower.len() - 1].ends_with(&canonical));
555                prefix_len
556                    .filter(|_| lower.ends_with(&canonical))
557                    .or(plural_prefix_len)
558                    .is_some_and(|len| {
559                        len > 0
560                            && name
561                                .as_bytes()
562                                .get(..len)
563                                .is_some_and(|prefix| prefix.iter().all(u8::is_ascii_alphabetic))
564                    })
565            })
566            .max_by_key(|canonical| canonical.len())
567            .map(String::as_str)
568    }
569
570    fn unknown_tool_error(&self, name: &str) -> String {
571        let available = if self.tools.is_empty() {
572            "(none registered)".to_string()
573        } else {
574            self.names().join(", ")
575        };
576        match self.suggest_name(name) {
577            Some(suggestion) => format!(
578                "unknown tool '{name}'. Did you mean '{suggestion}'? Call tools by their exact registered name. Available: {available}"
579            ),
580            None => format!(
581                "unknown tool '{name}'. Call one of the registered tools by exact name. Available: {available}"
582            ),
583        }
584    }
585
586    /// Validate, execute and validate the result of one tool call.
587    pub async fn execute_with_context(
588        &self,
589        name: &str,
590        context: &ToolExecutionContext,
591        args: Value,
592    ) -> Result<Value, String> {
593        if let Some(canonical) = self.canonical_name(name) {
594            self.validate_arguments(canonical, &args)?;
595            let tool = self
596                .tools
597                .get(canonical)
598                .ok_or_else(|| self.unknown_tool_error(canonical))?;
599            let result = tool.call_with_context(context, args).await;
600            let value = result?;
601            self.validate_output(canonical, &value)?;
602            return Ok(value);
603        }
604        Err(self.unknown_tool_error(name))
605    }
606
607    /// Internal name for an internal or model-facing name.
608    pub fn canonical_name<'a>(&'a self, name: &'a str) -> Option<&'a str> {
609        if self.tools.contains_key(name) {
610            return Some(name);
611        }
612        self.wire_names
613            .get(name)
614            .map(String::as_str)
615            .or_else(|| self.suggest_name(name))
616    }
617}
618
619/// Validate `value` against `schema`.
620pub fn validate_json_schema(schema: &Value, value: &Value) -> Result<(), String> {
621    validate_json_schema_value(schema, value, "tool arguments")
622}
623
624/// Check that `schema` is a valid JSON Schema.
625pub fn validate_json_schema_definition(schema: &Value) -> Result<(), String> {
626    jsonschema::validator_for(schema)
627        .map(|_| ())
628        .map_err(|error| format!("invalid JSON schema: {error}"))
629}
630
631/// Validate `value` against `schema`, labelling errors with `subject`.
632pub fn validate_json_schema_value(
633    schema: &Value,
634    value: &Value,
635    subject: &str,
636) -> Result<(), String> {
637    let validator = jsonschema::validator_for(schema)
638        .map_err(|error| format!("invalid JSON schema: {error}"))?;
639    validator
640        .validate(value)
641        .map_err(|error| format!("invalid {subject}: {error}"))
642}
643
644/// Model-facing tool name (dots become underscores).
645pub fn model_tool_name(internal: &str) -> String {
646    if !internal.is_empty()
647        && internal.len() <= 64
648        && internal
649            .bytes()
650            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
651    {
652        return internal.to_string();
653    }
654    let mut prefix = internal
655        .bytes()
656        .map(|byte| {
657            if byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-') {
658                byte as char
659            } else {
660                '_'
661            }
662        })
663        .take(47)
664        .collect::<String>();
665    if prefix.is_empty() {
666        prefix.push_str("tool");
667    }
668    let digest = format!("{:x}", Sha256::digest(internal.as_bytes()));
669    format!("{prefix}_{}", &digest[..16])
670}
671
672/// Helpers for implementing [`Tool`].
673pub mod support {
674    use serde::de::DeserializeOwned;
675    use serde_json::Value;
676
677    /// Provides the argument schema for a typed tool.
678    pub trait RawToolSchema {
679        /// JSON Schema for the arguments.
680        fn parameters() -> Value;
681    }
682
683    /// Deserialize a required argument.
684    pub fn extract_required<T: DeserializeOwned>(args: &Value, name: &str) -> Result<T, String> {
685        let value = args
686            .get(name)
687            .cloned()
688            .ok_or_else(|| format!("missing required argument '{name}'"))?;
689        serde_json::from_value(value).map_err(|error| format!("invalid argument '{name}': {error}"))
690    }
691
692    /// Deserialize an optional argument.
693    pub fn extract_optional<T: DeserializeOwned>(
694        args: &Value,
695        name: &str,
696    ) -> Result<Option<T>, String> {
697        match args.get(name) {
698            None | Some(Value::Null) => Ok(None),
699            Some(value) => serde_json::from_value(value.clone())
700                .map(Some)
701                .map_err(|error| format!("invalid argument '{name}': {error}")),
702        }
703    }
704}
705
706impl fmt::Debug for ToolRegistry {
707    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
708        f.debug_struct("ToolRegistry")
709            .field("tools", &self.names())
710            .finish()
711    }
712}
713
714#[cfg(test)]
715mod tests {
716    use super::*;
717
718    fn execution() -> ToolExecutionContext {
719        ToolExecutionContext {
720            request: crate::RequestContext {
721                tenant_id: "tenant".parse().unwrap(),
722                subject_id: "subject".parse().unwrap(),
723                roles: Default::default(),
724                locale: "en".into(),
725                request_id: "request".parse().unwrap(),
726                entitlements: Default::default(),
727            },
728            session_id: "session".parse().unwrap(),
729            run_id: "run".parse().unwrap(),
730            step: 1,
731            call_id: "call".parse().unwrap(),
732            source_event_seq: 1,
733            risk_class: ToolRiskClass::Read,
734            execution_owner: ToolExecutionOwner::Runtime,
735            interaction_resolution: None,
736            cancellation: CancellationToken::default(),
737            deadline: Instant::now() + std::time::Duration::from_secs(1),
738        }
739    }
740
741    struct TestTool(&'static str);
742    #[async_trait]
743    impl Tool for TestTool {
744        fn name(&self) -> &str {
745            self.0
746        }
747        fn description(&self) -> &str {
748            "test"
749        }
750        fn parameters(&self) -> Value {
751            serde_json::json!({
752                "type":"object",
753                "required":["items","mode"],
754                "additionalProperties":false,
755                "properties":{
756                    "items":{"type":"array","items":{"type":"object","required":["id"],"properties":{"id":{"type":"integer"}}}},
757                    "mode":{"enum":["safe","fast"]},
758                    "version":{"const":1},
759                    "choice":{"oneOf":[{"type":"string"},{"type":"number"}]}
760                }
761            })
762        }
763        fn output_schema(&self) -> Value {
764            serde_json::json!({"type":"object"})
765        }
766        fn implementation_version(&self) -> &str {
767            "1"
768        }
769        async fn call(&self, args: Value) -> Result<Value, String> {
770            Ok(args)
771        }
772    }
773
774    #[test]
775    fn runtime_manifest_rejects_missing_execution_metadata() {
776        let mut registry = ToolRegistry::new();
777        registry
778            .register(Arc::new(TestTool("unclassified")))
779            .unwrap();
780        assert!(registry
781            .runtime_manifest()
782            .unwrap_err()
783            .contains("explicit risk_class and execution_owner"));
784    }
785
786    #[tokio::test]
787    async fn registry_rejects_duplicates_and_validates_full_schema() {
788        let mut registry = ToolRegistry::new();
789        registry
790            .register(Arc::new(TestTool("nested.tool")))
791            .unwrap();
792        assert!(registry
793            .register(Arc::new(TestTool("nested.tool")))
794            .is_err());
795        let valid = serde_json::json!({"items":[{"id":1}],"mode":"safe","version":1,"choice":"x"});
796        assert_eq!(
797            registry
798                .execute_with_context("nested.tool", &execution(), valid.clone())
799                .await
800                .unwrap(),
801            valid
802        );
803        for invalid in [
804            serde_json::json!({"items":[{}],"mode":"safe"}),
805            serde_json::json!({"items":[{"id":1}],"mode":"unsafe"}),
806            serde_json::json!({"items":[{"id":1}],"mode":"safe","extra":true}),
807            serde_json::json!({"items":[{"id":1}],"mode":"safe","version":2}),
808            serde_json::json!({"items":[{"id":1}],"mode":"safe","choice":true}),
809        ] {
810            assert!(registry
811                .execute_with_context("nested.tool", &execution(), invalid)
812                .await
813                .is_err());
814        }
815    }
816
817    #[test]
818    fn provider_name_mismatch_repairs_only_a_unique_schema_match() {
819        struct CoinTool(&'static str);
820        struct PermissiveTool;
821        #[async_trait]
822        impl Tool for CoinTool {
823            fn name(&self) -> &str {
824                self.0
825            }
826            fn description(&self) -> &str {
827                "coin"
828            }
829            fn parameters(&self) -> Value {
830                serde_json::json!({"type":"object","required":["coin"],"additionalProperties":false,"properties":{"coin":{"type":"string"}}})
831            }
832            fn output_schema(&self) -> Value {
833                serde_json::json!({"type":"object"})
834            }
835            fn implementation_version(&self) -> &str {
836                "1"
837            }
838            async fn call(&self, args: Value) -> Result<Value, String> {
839                Ok(args)
840            }
841        }
842        #[async_trait]
843        impl Tool for PermissiveTool {
844            fn name(&self) -> &str {
845                "permissive"
846            }
847            fn description(&self) -> &str {
848                "accepts undeclared extra arguments"
849            }
850            fn parameters(&self) -> Value {
851                serde_json::json!({"type":"object","properties":{"query":{"type":"string"}}})
852            }
853            fn output_schema(&self) -> Value {
854                serde_json::json!({"type":"object"})
855            }
856            async fn call(&self, args: Value) -> Result<Value, String> {
857                Ok(args)
858            }
859        }
860        let mut registry = ToolRegistry::new();
861        registry.register(Arc::new(TestTool("wrong"))).unwrap();
862        registry.register(Arc::new(PermissiveTool)).unwrap();
863        registry.register(Arc::new(CoinTool("candles"))).unwrap();
864        let args = serde_json::json!({"coin":"BTC"});
865        assert_eq!(
866            registry.resolve_call_name("wrong", &args).unwrap(),
867            "candles"
868        );
869        registry.register(Arc::new(CoinTool("snapshot"))).unwrap();
870        assert!(registry.resolve_call_name("wrong", &args).is_err());
871    }
872
873    #[test]
874    fn model_names_are_provider_safe_and_reversible() {
875        let mut registry = ToolRegistry::new();
876        registry.register(Arc::new(TestTool("namespace.tool with unicode-工具-and-a-name-that-is-far-too-long-for-provider-contracts"))).unwrap();
877        let spec = registry.specs().pop().unwrap().function.name;
878        assert!(spec.len() <= 64);
879        assert!(spec
880            .bytes()
881            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')));
882        assert_eq!(registry.canonical_name(&spec), Some("namespace.tool with unicode-工具-and-a-name-that-is-far-too-long-for-provider-contracts"));
883
884        let mut collision = ToolRegistry::new();
885        collision.register(Arc::new(TestTool("a.b"))).unwrap();
886        assert_eq!(model_tool_name("a.b"), "a_b_2e7336dc8eba87ef");
887        assert!(collision
888            .register(Arc::new(TestTool("a_b_2e7336dc8eba87ef")))
889            .is_err());
890    }
891
892    #[tokio::test]
893    async fn provider_junk_prefixes_resolve_once_at_the_registry_boundary() {
894        let mut registry = ToolRegistry::new();
895        registry
896            .register(Arc::new(TestTool("analyze_wallet")))
897            .unwrap();
898        let args = serde_json::json!({"items":[{"id":1}],"mode":"safe"});
899
900        assert_eq!(
901            registry
902                .execute_with_context("Notebookanalyze_wallet", &execution(), args.clone())
903                .await
904                .unwrap(),
905            args
906        );
907        assert!(registry
908            .execute_with_context("Listanalyze_wallets", &execution(), args.clone())
909            .await
910            .is_ok());
911        assert!(registry
912            .execute_with_context("namespace.analyze_wallet", &execution(), args)
913            .await
914            .is_err());
915    }
916
917    #[test]
918    fn invalid_schema_is_rejected_at_registration() {
919        struct Invalid;
920        #[async_trait]
921        impl Tool for Invalid {
922            fn name(&self) -> &str {
923                "invalid"
924            }
925            fn description(&self) -> &str {
926                "invalid"
927            }
928            fn parameters(&self) -> Value {
929                serde_json::json!({"type":"not-a-type"})
930            }
931            fn output_schema(&self) -> Value {
932                serde_json::json!({"type":"object"})
933            }
934            async fn call(&self, _: Value) -> Result<Value, String> {
935                Ok(Value::Null)
936            }
937        }
938        assert!(ToolRegistry::new().register(Arc::new(Invalid)).is_err());
939    }
940
941    #[tokio::test]
942    async fn successful_output_is_validated_before_materialization() {
943        struct InvalidOutput;
944        #[async_trait]
945        impl Tool for InvalidOutput {
946            fn name(&self) -> &str {
947                "invalid-output"
948            }
949            fn description(&self) -> &str {
950                "invalid output"
951            }
952            fn parameters(&self) -> Value {
953                serde_json::json!({"type":"object"})
954            }
955            fn output_schema(&self) -> Value {
956                serde_json::json!({"type":"object"})
957            }
958            async fn call(&self, _: Value) -> Result<Value, String> {
959                Ok(Value::String("bad".into()))
960            }
961        }
962        let mut registry = ToolRegistry::new();
963        registry.register(Arc::new(InvalidOutput)).unwrap();
964        assert!(registry
965            .execute_with_context("invalid-output", &execution(), serde_json::json!({}))
966            .await
967            .unwrap_err()
968            .contains("invalid tool output"));
969    }
970}