wasm-agent 0.2.0

// SPDX-License-Identifier: MIT
//! Tool registry and dispatch for WASM-safe function signatures.

use std::collections::{BTreeMap, HashMap};
use crate::error::AgentError;
use crate::types::ToolResult;

/// A tool handler: takes a JSON string input, returns a [`ToolResult`].
///
/// Uses `Box<dyn Fn>` for WASM-compatible dispatch: no async, no `Send`
/// requirement in single-threaded WASM environments.
pub type ToolFn = Box<dyn Fn(&str) -> ToolResult>;

/// Metadata describing a tool for system-prompt generation and introspection.
#[derive(Debug, Clone)]
pub struct ToolSpec {
    /// Unique name used to invoke the tool from the ReAct loop.
    pub name: String,
    /// Human-readable description included in the agent's system prompt.
    pub description: String,
    /// JSON Schema string describing the tool's input format.
    pub input_schema: String,
}

impl ToolSpec {
    /// Creates a new [`ToolSpec`].
    ///
    /// # Arguments
    /// * `name`: Unique tool name (must be non-empty).
    /// * `description`: Human-readable description.
    /// * `schema`: JSON Schema for the tool's input.
    pub fn new(name: impl Into<String>, description: impl Into<String>, schema: impl Into<String>) -> Self {
        Self { name: name.into(), description: description.into(), input_schema: schema.into() }
    }
}

/// Registry of available tools, keyed by name.
pub struct ToolRegistry {
    // Sorted, so tools_prompt() is identical on every run (a prompt that
    // changes order between runs defeats provider prompt caching).
    specs: BTreeMap<String, ToolSpec>,
    handlers: HashMap<String, ToolFn>,
}

impl ToolRegistry {
    /// Creates an empty registry.
    pub fn new() -> Self {
        Self { specs: BTreeMap::new(), handlers: HashMap::new() }
    }

    /// Registers a tool with its specification and handler.
    ///
    /// # Errors
    /// Returns [`AgentError::InvalidToolSignature`] if the tool name is empty.
    pub fn register(&mut self, spec: ToolSpec, handler: ToolFn) -> Result<(), AgentError> {
        if spec.name.is_empty() {
            return Err(AgentError::InvalidToolSignature("tool name cannot be empty".into()));
        }
        self.specs.insert(spec.name.clone(), spec.clone());
        self.handlers.insert(spec.name, handler);
        Ok(())
    }

    /// Dispatches a tool call by name with the given JSON input string.
    ///
    /// # Errors
    /// Returns [`AgentError::ToolNotFound`] if no tool with that name is registered.
    pub fn dispatch(&self, tool_name: &str, input: &str) -> Result<ToolResult, AgentError> {
        let handler = self.handlers.get(tool_name)
            .ok_or_else(|| AgentError::ToolNotFound { name: tool_name.to_string() })?;
        Ok(handler(input))
    }

    /// Dispatches a native tool call (see [`crate::native`]). Its arguments
    /// arrive as JSON. For a tool with an object schema they are passed on
    /// unchanged; for a tool with any other schema the API saw
    /// `{"input": ...}`, so that value is unwrapped (a JSON string is passed as
    /// plain text). Arguments that are not valid JSON produce a failed
    /// [`ToolResult`] the model can see and correct, not an error.
    ///
    /// # Errors
    /// [`AgentError::ToolNotFound`] if no tool with that name is registered.
    pub fn dispatch_call(&self, call: &crate::native::ToolCall) -> Result<ToolResult, AgentError> {
        let spec = self.specs.get(&call.name).ok_or_else(|| AgentError::ToolNotFound { name: call.name.clone() })?;
        let wrapped = crate::native::api_schema(&spec.input_schema).map(|(_, w)| w).unwrap_or(true);
        if !wrapped {
            return self.dispatch(&call.name, &call.arguments);
        }
        let input = match serde_json::from_str::<serde_json::Value>(&call.arguments) {
            Ok(v) => match v.get("input") {
                Some(serde_json::Value::String(s)) => s.clone(),
                Some(other) => other.to_string(),
                None => {
                    return Ok(ToolResult {
                        tool_name: call.name.clone(),
                        output: "missing \"input\" argument".into(),
                        success: false,
                    })
                }
            },
            Err(e) => {
                return Ok(ToolResult { tool_name: call.name.clone(), output: format!("arguments are not valid JSON: {e}"), success: false })
            }
        };
        self.dispatch(&call.name, &input)
    }

    /// The `tools` array for OpenAI Chat Completions (and OpenAI-compatible
    /// APIs: Mistral, Groq, DeepSeek, OpenRouter, Ollama's `/v1`), sorted by name.
    ///
    /// # Errors
    /// [`AgentError::InvalidToolSignature`] if a tool's schema is not JSON.
    pub fn openai_tools(&self) -> Result<Vec<serde_json::Value>, AgentError> {
        self.specs
            .values()
            .map(|s| {
                let (params, _) = crate::native::api_schema(&s.input_schema)?;
                Ok(serde_json::json!({"type": "function", "function": {"name": s.name, "description": s.description, "parameters": params}}))
            })
            .collect()
    }

    /// The `tools` array for Anthropic Messages, sorted by name.
    ///
    /// # Errors
    /// [`AgentError::InvalidToolSignature`] if a tool's schema is not JSON.
    pub fn anthropic_tools(&self) -> Result<Vec<serde_json::Value>, AgentError> {
        self.specs
            .values()
            .map(|s| {
                let (schema, _) = crate::native::api_schema(&s.input_schema)?;
                Ok(serde_json::json!({"name": s.name, "description": s.description, "input_schema": schema}))
            })
            .collect()
    }

    /// Register a tool whose input is a Rust type: the JSON Schema comes from
    /// [schemars](https://crates.io/crates/schemars) and the input is parsed
    /// with serde before `f` runs. Input that does not match the type gives
    /// the model a failed result explaining why, so it can retry. Needs the
    /// `schemars` feature.
    ///
    /// ```
    /// use schemars::JsonSchema;
    /// use serde::Deserialize;
    ///
    /// #[derive(Deserialize, JsonSchema)]
    /// struct Add { a: i64, b: i64 }
    ///
    /// let mut tools = wasm_agent::ToolRegistry::new();
    /// tools.register_typed("add", "Adds two integers", |x: Add| Ok((x.a + x.b).to_string()))?;
    /// assert_eq!(tools.dispatch("add", r#"{"a": 2, "b": 3}"#)?.output, "5");
    /// assert!(!tools.dispatch("add", r#"{"a": "two"}"#)?.success);
    /// # Ok::<(), wasm_agent::AgentError>(())
    /// ```
    ///
    /// # Errors
    /// [`AgentError::InvalidToolSignature`] if the name is empty.
    #[cfg(feature = "schemars")]
    #[cfg_attr(docsrs, doc(cfg(feature = "schemars")))]
    pub fn register_typed<A, F>(&mut self, name: &str, description: &str, f: F) -> Result<(), AgentError>
    where
        A: serde::de::DeserializeOwned + schemars::JsonSchema + 'static,
        F: Fn(A) -> Result<String, String> + 'static,
    {
        let schema = serde_json::to_string(&schemars::schema_for!(A)).map_err(|e| AgentError::Serialization(e.to_string()))?;
        let tool = name.to_string();
        self.register(
            ToolSpec::new(name, description, schema),
            Box::new(move |input: &str| match serde_json::from_str::<A>(input) {
                Ok(args) => match f(args) {
                    Ok(output) => ToolResult { tool_name: tool.clone(), output, success: true },
                    Err(output) => ToolResult { tool_name: tool.clone(), output, success: false },
                },
                Err(e) => ToolResult { tool_name: tool.clone(), output: format!("invalid arguments: {e}"), success: false },
            }),
        )
    }

    /// Returns the spec for a registered tool, or `None` if not found.
    pub fn spec(&self, name: &str) -> Option<&ToolSpec> { self.specs.get(name) }

    /// Returns the number of registered tools.
    pub fn tool_count(&self) -> usize { self.specs.len() }

    /// Returns the names of all registered tools, sorted.
    pub fn tool_names(&self) -> Vec<&str> { self.specs.keys().map(|s| s.as_str()).collect() }

    /// Generates a system-prompt snippet listing all registered tools, sorted
    /// by name, with each tool's input schema when it has one.
    pub fn tools_prompt(&self) -> String {
        let mut lines = vec!["Available tools:".to_string()];
        for spec in self.specs.values() {
            let schema = spec.input_schema.trim();
            if schema.is_empty() || schema == "{}" {
                lines.push(format!("- {}: {}", spec.name, spec.description));
            } else {
                lines.push(format!("- {}: {} (input: {})", spec.name, spec.description, schema));
            }
        }
        lines.join("\n")
    }
}

impl Default for ToolRegistry {
    fn default() -> Self { Self::new() }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn echo_tool() -> (ToolSpec, ToolFn) {
        let spec = ToolSpec::new("echo", "Echoes input back", r#"{"type":"string"}"#);
        let handler: ToolFn = Box::new(|input: &str| ToolResult {
            tool_name: "echo".into(),
            output: format!("echo: {input}"),
            success: true,
        });
        (spec, handler)
    }

    #[test]
    fn test_registry_register_and_dispatch_ok() {
        let mut reg = ToolRegistry::new();
        let (spec, handler) = echo_tool();
        reg.register(spec, handler).unwrap();
        let result = reg.dispatch("echo", "hello").unwrap();
        assert!(result.success);
        assert_eq!(result.output, "echo: hello");
    }

    #[test]
    fn test_registry_dispatch_unknown_tool_returns_error() {
        let reg = ToolRegistry::new();
        let err = reg.dispatch("nonexistent", "").unwrap_err();
        assert!(matches!(err, AgentError::ToolNotFound { .. }));
    }

    #[test]
    fn test_registry_register_empty_name_returns_error() {
        let mut reg = ToolRegistry::new();
        let spec = ToolSpec::new("", "bad", "{}");
        let err = reg.register(spec, Box::new(|_| ToolResult {
            tool_name: "".into(), output: "".into(), success: false,
        })).unwrap_err();
        assert!(matches!(err, AgentError::InvalidToolSignature(_)));
    }

    #[test]
    fn test_registry_tool_count_increments() {
        let mut reg = ToolRegistry::new();
        assert_eq!(reg.tool_count(), 0);
        let (spec, handler) = echo_tool();
        reg.register(spec, handler).unwrap();
        assert_eq!(reg.tool_count(), 1);
    }

    #[test]
    fn test_registry_tools_prompt_contains_tool_name() {
        let mut reg = ToolRegistry::new();
        let (spec, handler) = echo_tool();
        reg.register(spec, handler).unwrap();
        assert!(reg.tools_prompt().contains("echo"));
    }

    #[test]
    fn test_registry_spec_retrieval_present_and_absent() {
        let mut reg = ToolRegistry::new();
        let (spec, handler) = echo_tool();
        reg.register(spec, handler).unwrap();
        assert!(reg.spec("echo").is_some());
        assert!(reg.spec("missing").is_none());
    }

    #[test]
    fn test_registry_tool_names_lists_all() {
        let mut reg = ToolRegistry::new();
        let (spec, handler) = echo_tool();
        reg.register(spec, handler).unwrap();
        let names = reg.tool_names();
        assert!(names.contains(&"echo"));
    }

    #[test]
    fn test_registry_default_is_empty() {
        let reg = ToolRegistry::default();
        assert_eq!(reg.tool_count(), 0);
    }

    #[test]
    fn test_tools_prompt_is_sorted_and_stable() {
        // 0.1.x iterated a HashMap, so the tool order changed between runs.
        let mut reg = ToolRegistry::new();
        for name in ["zeta", "alpha", "mid"] {
            reg.register(
                ToolSpec::new(name, "d", r#"{"type":"string"}"#),
                Box::new(|_: &str| ToolResult { tool_name: String::new(), output: String::new(), success: true }),
            )
            .unwrap();
        }
        let p = reg.tools_prompt();
        let a = p.find("- alpha").unwrap();
        let m = p.find("- mid").unwrap();
        let z = p.find("- zeta").unwrap();
        assert!(a < m && m < z);
        assert!(p.contains(r#"(input: {"type":"string"})"#));
        assert_eq!(reg.tool_names(), vec!["alpha", "mid", "zeta"]);
    }

    #[test]
    fn test_dispatch_call_unwraps_non_object_schemas() {
        let reg = registry_with_echo_schema(r#"{"type":"string"}"#);
        let call = |args: &str| crate::native::ToolCall { id: None, name: "echo".into(), arguments: args.into() };
        assert_eq!(reg.dispatch_call(&call(r#"{"input":"hi"}"#)).unwrap().output, "hi");
        assert_eq!(reg.dispatch_call(&call(r#"{"input":[1,2]}"#)).unwrap().output, "[1,2]");
        assert!(!reg.dispatch_call(&call("{bad")).unwrap().success);
        assert!(!reg.dispatch_call(&call("{}")).unwrap().success);
        let obj = registry_with_echo_schema(r#"{"type":"object","properties":{"q":{"type":"string"}}}"#);
        assert_eq!(obj.dispatch_call(&call(r#"{"q":"x"}"#)).unwrap().output, r#"{"q":"x"}"#);
        assert!(matches!(reg.dispatch_call(&crate::native::ToolCall { id: None, name: "nope".into(), arguments: "{}".into() }), Err(AgentError::ToolNotFound { .. })));
    }

    #[test]
    fn test_openai_and_anthropic_tool_definitions() {
        let reg = registry_with_echo_schema(r#"{"type":"string"}"#);
        let o = reg.openai_tools().unwrap();
        assert_eq!(o[0]["type"], "function");
        assert_eq!(o[0]["function"]["parameters"]["type"], "object");
        assert_eq!(o[0]["function"]["parameters"]["required"][0], "input");
        let a = reg.anthropic_tools().unwrap();
        assert_eq!(a[0]["input_schema"]["properties"]["input"]["type"], "string");
        let bad = registry_with_echo_schema("{not json");
        assert!(bad.openai_tools().is_err());
    }

    fn registry_with_echo_schema(schema: &str) -> ToolRegistry {
        let mut reg = ToolRegistry::new();
        reg.register(
            ToolSpec::new("echo", "Echoes", schema),
            Box::new(|i: &str| ToolResult { tool_name: "echo".into(), output: i.into(), success: true }),
        )
        .unwrap();
        reg
    }
}