af_agent/

tool.rs

//! Agent tools + registry. Port of the reusable shape of `agent_core/tools/`.
//!
//! A tool exposes a JSON-Schema interface to the model and an async
//! implementation. The registry resolves the name the model emits in a
//! `tool_call` to an implementation, exposes the LLM-facing specs, and runs
//! the call. Domain tools live in a product
//! crate and register here.

use af_context::{RunId, SessionId, ToolCallId};
use std::collections::HashMap;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
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;

/// Who may call a tool.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolSurface {
    /// Advertised to the model.
    Llm,
    /// Only the host or platform.
    Chassis,
}

/// Whether a tool may run alongside others in one step.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolConcurrency {
    /// May run in parallel.
    Concurrent,
    /// Runs alone; acts as a barrier.
    Exclusive,
}

/// Static execution metadata for a tool.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ToolMeta {
    /// Who may call it.
    pub surface: ToolSurface,
    /// Provider-neutral cost units settled for this operation.
    pub cost_units: u64,
    /// Per-call timeout.
    pub timeout_secs: u64,
    /// Whether the tool is part of the platform core (never narrowed away).
    pub core: bool,
    /// Concurrency class.
    pub concurrency: ToolConcurrency,
    /// Mutating tools must cross the runtime's normal confirmation hook even
    /// when a profile forgot to list them explicitly.
    pub requires_confirmation: bool,
}

/// Hierarchical cancellation flag. Cancelling a parent cancels every child;
/// [`cancelled`](Self::cancelled) resolves without polling.
#[derive(Debug, Clone, Default)]
pub struct CancellationToken(Arc<CancellationState>);

#[derive(Debug, Default)]
struct CancellationState {
    cancelled: AtomicBool,
    notify: tokio::sync::Notify,
    parent: Option<CancellationToken>,
}

impl CancellationToken {
    /// Cancel this token and every child.
    pub fn cancel(&self) {
        self.0.cancelled.store(true, Ordering::Release);
        self.0.notify.notify_waiters();
    }
    /// Whether this token or an ancestor was cancelled.
    pub fn is_cancelled(&self) -> bool {
        self.0.cancelled.load(Ordering::Acquire)
            || self
                .0
                .parent
                .as_ref()
                .is_some_and(CancellationToken::is_cancelled)
    }
    /// A child token cancelled with its parent.
    pub fn child(&self) -> Self {
        Self(Arc::new(CancellationState {
            cancelled: AtomicBool::new(false),
            notify: tokio::sync::Notify::new(),
            parent: Some(self.clone()),
        }))
    }
    /// Resolves once this token or any ancestor is cancelled.
    pub fn cancelled(&self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
        Box::pin(async move {
            let notified = self.0.notify.notified();
            tokio::pin!(notified);
            notified.as_mut().enable();
            if self.is_cancelled() {
                return;
            }
            match &self.0.parent {
                Some(parent) => tokio::select! {
                    _ = notified => {}
                    _ = parent.cancelled() => {}
                },
                None => notified.await,
            }
        })
    }
}

/// Identity, position and controls for one tool execution.
#[derive(Debug, Clone)]
pub struct ToolExecutionContext {
    /// The originating request.
    pub request: af_context::RequestContext,
    /// Session this record belongs to.
    pub session_id: SessionId,
    /// Run this record belongs to.
    pub run_id: RunId,
    /// 1-based step number inside the Turn.
    pub step: u32,
    /// Tool call this record refers to.
    pub call_id: ToolCallId,
    /// Sequence of the Session event this was derived from.
    pub source_event_seq: u64,
    /// Resolution of the interaction that unblocked this call, if any.
    pub interaction_resolution: Option<af_agent_session::InteractionResolution>,
    /// Cancelled when the Run or step is abandoned.
    pub cancellation: CancellationToken,
    /// Latest time by which the work must finish.
    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,
            requires_confirmation: false,
        }
    }
}

/// A callable tool. `parameters` is a JSON-Schema object describing the args.
#[async_trait]
pub trait Tool: Send + Sync {
    /// Stable internal name.
    fn name(&self) -> &str;
    /// Implementation version pinned by Profile revisions.
    fn implementation_version(&self) -> &str {
        ""
    }
    /// Model-facing description.
    fn description(&self) -> &str;
    /// JSON Schema for the arguments.
    fn parameters(&self) -> Value;
    /// JSON Schema for the result.
    fn output_schema(&self) -> Value;
    /// Execution metadata.
    fn meta(&self) -> ToolMeta {
        ToolMeta::default()
    }

    /// Run the tool. On success return a JSON value; on failure return a short
    /// error string (surfaced back to the model as `{"error": ...}` so it can
    /// recover) — mirrors the Python "success → dict / failure → {error}" rule.
    async fn call(&self, args: Value) -> Result<Value, String>;
    /// Like [`call`](Self::call) with the execution context; the default ignores the context.
    async fn call_with_context(
        &self,
        _context: &ToolExecutionContext,
        args: Value,
    ) -> Result<Value, String> {
        self.call(args).await
    }
}

/// Registry of tools available to an agent.
#[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 {
    /// An empty registry.
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a tool, compiling its schemas; duplicate or invalid tools are rejected.
    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)
    }

    /// Merge another registry; conflicting names are rejected.
    pub fn extend(&mut self, other: &Self) -> Result<(), String> {
        for tool in other.tools.values() {
            self.register(Arc::clone(tool))?;
        }
        Ok(())
    }

    /// Whether no tools are registered.
    pub fn is_empty(&self) -> bool {
        self.tools.is_empty()
    }

    /// Number of registered tools.
    pub fn len(&self) -> usize {
        self.tools.len()
    }

    /// Whether `name` (internal or model form) is registered.
    pub fn contains(&self, name: &str) -> bool {
        self.tools.contains_key(name)
    }

    /// Look up a tool by internal or model name.
    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
        self.canonical_name(name)
            .and_then(|name| self.tools.get(name))
            .cloned()
    }

    /// Validate arguments against the tool's schema.
    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)
            .ok_or_else(|| self.unknown_tool_error(name))?
            .validate(arguments)
            .map_err(|error| format!("invalid tool arguments: {error}"))
    }

    /// Validate a result against the tool's output schema.
    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)
            .ok_or_else(|| self.unknown_tool_error(name))?
            .validate(output)
            .map_err(|error| format!("invalid tool output: {error}"))
    }

    /// Stable tool names for validation, diagnostics, and capability policy.
    pub fn names(&self) -> Vec<&str> {
        let mut names = self.tools.keys().map(String::as_str).collect::<Vec<_>>();
        names.sort_unstable();
        names
    }

    /// Tools whose metadata requires confirmation.
    pub fn confirmation_required_names(&self) -> impl Iterator<Item = &str> {
        self.tools
            .values()
            .filter(|tool| tool.meta().requires_confirmation)
            .map(|tool| tool.name())
    }

    /// Return a registry containing only explicitly allowed tools.
    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) {
                // A subset of an already-validated registry cannot introduce a
                // conflict, so a failure here is unreachable and safely ignored.
                let _ = filtered.register(Arc::clone(tool));
            }
        }
        filtered
    }

    /// LLM-facing tool specs for the `tools` field of a completion request.
    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
    }

    /// Stable implementation and execution contract persisted by profile revisions.
    pub fn runtime_manifest(&self) -> Result<Vec<Value>, String> {
        self.names()
            .into_iter()
            .map(|name| {
                let tool = self
                    .tools
                    .get(name)
                    .ok_or_else(|| self.unknown_tool_error(name))?;
                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,
                    "requires_confirmation": meta.requires_confirmation,
                }))
            })
            .collect()
    }

    /// Best-effort canonical name for case mistakes or an alphabetic junk
    /// prefix glued to a registered tool name.
    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| {
                let canonical = canonical.to_ascii_lowercase();
                let prefix_len = lower.len().checked_sub(canonical.len());
                let plural_prefix_len = lower
                    .strip_suffix('s')
                    .and_then(|singular| singular.len().checked_sub(canonical.len()))
                    .filter(|_| lower[..lower.len() - 1].ends_with(&canonical));
                prefix_len
                    .filter(|_| lower.ends_with(&canonical))
                    .or(plural_prefix_len)
                    .is_some_and(|len| {
                        len > 0
                            && name
                                .as_bytes()
                                .get(..len)
                                .is_some_and(|prefix| prefix.iter().all(u8::is_ascii_alphabetic))
                    })
            })
            .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}"
            ),
        }
    }

    /// Validate, execute and validate the result of one tool call.
    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)
                .ok_or_else(|| self.unknown_tool_error(canonical))?;
            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))
    }

    /// Internal name for an internal or model-facing 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)
            .or_else(|| self.suggest_name(name))
    }
}

/// Validate `value` against `schema`.
pub fn validate_json_schema(schema: &Value, value: &Value) -> Result<(), String> {
    validate_json_schema_value(schema, value, "tool arguments")
}

/// Check that `schema` is a valid JSON Schema.
pub fn validate_json_schema_definition(schema: &Value) -> Result<(), String> {
    jsonschema::validator_for(schema)
        .map(|_| ())
        .map_err(|error| format!("invalid JSON schema: {error}"))
}

/// Validate `value` against `schema`, labelling errors with `subject`.
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}"))
}

/// Model-facing tool name (dots become underscores).
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])
}

/// Helpers for implementing [`Tool`].
pub mod support {
    use serde::de::DeserializeOwned;
    use serde_json::Value;

    /// Provides the argument schema for a typed tool.
    pub trait RawToolSchema {
        /// JSON Schema for the arguments.
        fn parameters() -> Value;
    }

    /// Deserialize a required argument.
    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}"))
    }

    /// Deserialize an optional argument.
    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".parse().unwrap(),
                subject_id: "subject".parse().unwrap(),
                roles: Default::default(),
                locale: "en".into(),
                request_id: "request".parse().unwrap(),
                entitlements: Default::default(),
            },
            session_id: "session".parse().unwrap(),
            run_id: "run".parse().unwrap(),
            step: 1,
            call_id: "call".parse().unwrap(),
            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());
    }

    #[tokio::test]
    async fn provider_junk_prefixes_resolve_once_at_the_registry_boundary() {
        let mut registry = ToolRegistry::new();
        registry
            .register(Arc::new(TestTool("analyze_wallet")))
            .unwrap();
        let args = serde_json::json!({"items":[{"id":1}],"mode":"safe"});

        assert_eq!(
            registry
                .execute_with_context("Notebookanalyze_wallet", &execution(), args.clone())
                .await
                .unwrap(),
            args
        );
        assert!(registry
            .execute_with_context("Listanalyze_wallets", &execution(), args.clone())
            .await
            .is_ok());
        assert!(registry
            .execute_with_context("namespace.analyze_wallet", &execution(), args)
            .await
            .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"));
    }
}