ironflow-core 6.0.0

//! Tool registry: stores tools and converts them to the format of each HTTP API.

use std::collections::{HashMap, HashSet};
use std::fmt;

use serde_json::{Value, json};

use super::tool_trait::{Tool, ToolError, ToolOutput};

/// Registry of client-side tools available to the HTTP agent provider.
///
/// Converts registered tools to the tool format of each HTTP API (OpenAI,
/// Anthropic, Gemini) and routes tool call execution to the correct tool
/// implementation.
///
/// # Examples
///
/// ```no_run
/// use ironflow_core::providers::http::tools::ToolRegistry;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let registry = ToolRegistry::new();
/// // registry.register(my_tool);
/// let tools_json = registry.to_openai_tools();
/// # Ok(())
/// # }
/// ```
pub struct ToolRegistry {
    tools: Vec<Box<dyn Tool>>,
    index: HashMap<String, usize>,
    connectors: HashSet<String>,
}

impl fmt::Debug for ToolRegistry {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ToolRegistry")
            .field(
                "tools",
                &self
                    .tools
                    .iter()
                    .map(|tool| tool.name())
                    .collect::<Vec<_>>(),
            )
            .field("connectors", &self.connectors)
            .finish()
    }
}

impl ToolRegistry {
    /// Create an empty registry.
    pub fn new() -> Self {
        Self {
            tools: Vec::new(),
            index: HashMap::new(),
            connectors: HashSet::new(),
        }
    }

    /// Register a tool. Panics if a tool with the same name is already registered.
    ///
    /// # Panics
    ///
    /// Panics if a tool with the same name already exists in the registry.
    pub fn register(mut self, tool: impl Tool + 'static) -> Self {
        let name = tool.name().to_string();
        assert!(
            !self.index.contains_key(&name),
            "tool '{}' already registered",
            name
        );
        let idx = self.tools.len();
        self.tools.push(Box::new(tool));
        self.index.insert(name, idx);
        self
    }

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

    /// Returns `true` if no tools are registered.
    pub fn is_empty(&self) -> bool {
        self.tools.is_empty()
    }

    /// Names of the registered tools, in registration order.
    ///
    /// # Examples
    ///
    /// ```
    /// use ironflow_core::providers::http::tools::ToolRegistry;
    ///
    /// assert!(ToolRegistry::new().tool_names().is_empty());
    /// ```
    pub fn tool_names(&self) -> Vec<&str> {
        self.tools.iter().map(|tool| tool.name()).collect()
    }

    /// Convert all registered tools to the OpenAI `tools` array format.
    ///
    /// Each tool is represented as:
    /// ```json
    /// {
    ///   "type": "function",
    ///   "function": {
    ///     "name": "tool_name",
    ///     "description": "tool description",
    ///     "parameters": { ... json schema ... }
    ///   }
    /// }
    /// ```
    pub fn to_openai_tools(&self) -> Vec<Value> {
        self.tools
            .iter()
            .map(|tool| {
                json!({
                    "type": "function",
                    "function": {
                        "name": tool.name(),
                        "description": tool.description(),
                        "parameters": tool.parameters_schema()
                    }
                })
            })
            .collect()
    }

    /// Convert all registered tools to the Anthropic Messages API `tools` format.
    ///
    /// Each tool is represented as:
    /// ```json
    /// {
    ///   "name": "tool_name",
    ///   "description": "tool description",
    ///   "input_schema": { ... json schema ... }
    /// }
    /// ```
    ///
    /// # Examples
    ///
    /// ```
    /// use ironflow_core::providers::http::tools::ToolRegistry;
    ///
    /// assert!(ToolRegistry::new().to_anthropic_tools().is_empty());
    /// ```
    pub fn to_anthropic_tools(&self) -> Vec<Value> {
        self.tools
            .iter()
            .map(|tool| {
                json!({
                    "name": tool.name(),
                    "description": tool.description(),
                    "input_schema": tool.parameters_schema()
                })
            })
            .collect()
    }

    /// Convert all registered tools to function declarations with a raw JSON
    /// schema, the base of the Gemini `functionDeclarations` format.
    ///
    /// Each tool is represented as:
    /// ```json
    /// {
    ///   "name": "tool_name",
    ///   "description": "tool description",
    ///   "parameters": { ... json schema ... }
    /// }
    /// ```
    ///
    /// # Examples
    ///
    /// ```
    /// use ironflow_core::providers::http::tools::ToolRegistry;
    ///
    /// assert!(ToolRegistry::new().to_function_declarations().is_empty());
    /// ```
    pub fn to_function_declarations(&self) -> Vec<Value> {
        self.tools
            .iter()
            .map(|tool| {
                json!({
                    "name": tool.name(),
                    "description": tool.description(),
                    "parameters": tool.parameters_schema()
                })
            })
            .collect()
    }

    /// Execute a tool by name with the given input.
    ///
    /// Returns `None` if no tool with that name is registered.
    pub async fn execute(&self, name: &str, input: Value) -> Option<Result<ToolOutput, ToolError>> {
        let idx = self.index.get(name)?;
        let tool = &self.tools[*idx];
        Some(tool.execute(input).await)
    }

    /// Check if a tool with the given name is registered.
    pub fn has_tool(&self, name: &str) -> bool {
        self.index.contains_key(name)
    }

    /// Returns whether the tool registered under `name` is read-only.
    ///
    /// Returns `false` (not read-only) for a name that is not registered,
    /// so unknown tool calls act as barriers rather than being parallelized.
    pub fn is_read_only(&self, name: &str) -> bool {
        self.index
            .get(name)
            .is_some_and(|&idx| self.tools[idx].read_only())
    }

    /// Register a connector name for MCP prefix routing.
    ///
    /// Called by `register_mcp_tools` to track which prefixes are valid for
    /// tool call routing.
    pub fn register_connector(mut self, name: &str) -> Self {
        self.connectors.insert(name.to_string());
        self
    }

    /// Returns the set of registered MCP connector prefixes.
    pub fn connectors(&self) -> &HashSet<String> {
        &self.connectors
    }
}

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

#[cfg(test)]
mod tests {
    use std::future::Future;
    use std::pin::Pin;

    use serde_json::json;

    use super::*;

    struct AddTool;

    impl Tool for AddTool {
        fn name(&self) -> &str {
            "add"
        }
        fn description(&self) -> &str {
            "Adds two numbers"
        }
        fn parameters_schema(&self) -> Value {
            json!({
                "type": "object",
                "properties": {
                    "a": {"type": "number"},
                    "b": {"type": "number"}
                },
                "required": ["a", "b"]
            })
        }
        fn execute(
            &self,
            input: Value,
        ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
            Box::pin(async move {
                let a = input
                    .get("a")
                    .and_then(|v| v.as_f64())
                    .ok_or_else(|| ToolError::new("missing 'a'"))?;
                let b = input
                    .get("b")
                    .and_then(|v| v.as_f64())
                    .ok_or_else(|| ToolError::new("missing 'b'"))?;
                Ok(ToolOutput::success(format!("{}", a + b)))
            })
        }
    }

    struct EchoTool;

    impl Tool for EchoTool {
        fn name(&self) -> &str {
            "echo"
        }
        fn description(&self) -> &str {
            "Echoes input"
        }
        fn parameters_schema(&self) -> Value {
            json!({"type": "object", "properties": {"msg": {"type": "string"}}})
        }
        fn execute(
            &self,
            input: Value,
        ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
            Box::pin(async move {
                let msg = input.get("msg").and_then(|v| v.as_str()).unwrap_or("empty");
                Ok(ToolOutput::success(msg))
            })
        }
    }

    struct ReadOnlyTool;

    impl Tool for ReadOnlyTool {
        fn name(&self) -> &str {
            "lookup"
        }
        fn description(&self) -> &str {
            "Reads without side effects"
        }
        fn parameters_schema(&self) -> Value {
            json!({"type": "object", "properties": {}})
        }
        fn execute(
            &self,
            _input: Value,
        ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
            Box::pin(async { Ok(ToolOutput::success("found")) })
        }
        fn read_only(&self) -> bool {
            true
        }
    }

    #[test]
    fn registry_new_is_empty() {
        let registry = ToolRegistry::new();
        assert!(registry.is_empty());
        assert_eq!(registry.len(), 0);
    }

    #[test]
    fn registry_register_increments_count() {
        let registry = ToolRegistry::new().register(AddTool).register(EchoTool);
        assert_eq!(registry.len(), 2);
        assert!(!registry.is_empty());
    }

    #[test]
    fn registry_has_tool() {
        let registry = ToolRegistry::new().register(AddTool);
        assert!(registry.has_tool("add"));
        assert!(!registry.has_tool("echo"));
    }

    #[test]
    #[should_panic(expected = "tool 'add' already registered")]
    fn registry_duplicate_panics() {
        ToolRegistry::new().register(AddTool).register(AddTool);
    }

    #[test]
    fn tool_names_keep_registration_order() {
        let registry = ToolRegistry::new().register(EchoTool).register(AddTool);
        assert_eq!(registry.tool_names(), vec!["echo", "add"]);
    }

    #[test]
    fn to_openai_tools_format() {
        let registry = ToolRegistry::new().register(AddTool);
        let tools = registry.to_openai_tools();

        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0]["type"], "function");
        assert_eq!(tools[0]["function"]["name"], "add");
        assert_eq!(tools[0]["function"]["description"], "Adds two numbers");
        assert_eq!(tools[0]["function"]["parameters"]["type"], "object");
        assert!(tools[0]["function"]["parameters"]["properties"]["a"].is_object());
    }

    #[test]
    fn to_openai_tools_multiple() {
        let registry = ToolRegistry::new().register(AddTool).register(EchoTool);
        let tools = registry.to_openai_tools();

        assert_eq!(tools.len(), 2);
        assert_eq!(tools[0]["function"]["name"], "add");
        assert_eq!(tools[1]["function"]["name"], "echo");
    }

    #[test]
    fn to_anthropic_tools_format() {
        let registry = ToolRegistry::new().register(AddTool);
        let tools = registry.to_anthropic_tools();

        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0]["name"], "add");
        assert_eq!(tools[0]["description"], "Adds two numbers");
        assert_eq!(tools[0]["input_schema"]["type"], "object");
        assert!(tools[0].get("type").is_none());
        assert!(tools[0].get("function").is_none());
    }

    #[test]
    fn to_function_declarations_format() {
        let registry = ToolRegistry::new().register(AddTool);
        let declarations = registry.to_function_declarations();

        assert_eq!(declarations.len(), 1);
        assert_eq!(declarations[0]["name"], "add");
        assert_eq!(declarations[0]["description"], "Adds two numbers");
        assert_eq!(declarations[0]["parameters"]["type"], "object");
        assert!(declarations[0]["parameters"]["properties"]["a"].is_object());
    }

    #[test]
    fn provider_tool_formats_keep_registration_order() {
        let registry = ToolRegistry::new().register(EchoTool).register(AddTool);

        let anthropic = registry.to_anthropic_tools();
        assert_eq!(anthropic[0]["name"], "echo");
        assert_eq!(anthropic[1]["name"], "add");

        let declarations = registry.to_function_declarations();
        assert_eq!(declarations[0]["name"], "echo");
        assert_eq!(declarations[1]["name"], "add");
    }

    #[tokio::test]
    async fn execute_existing_tool() {
        let registry = ToolRegistry::new().register(AddTool);
        let result = registry
            .execute("add", json!({"a": 3, "b": 4}))
            .await
            .expect("tool should exist")
            .expect("tool should succeed");

        assert_eq!(result.content, "7");
        assert!(!result.is_error);
    }

    #[tokio::test]
    async fn execute_nonexistent_tool_returns_none() {
        let registry = ToolRegistry::new().register(AddTool);
        let result = registry.execute("nonexistent", json!({})).await;
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn execute_tool_error_propagates() {
        let registry = ToolRegistry::new().register(AddTool);
        let result = registry
            .execute("add", json!({"a": 1}))
            .await
            .expect("tool should exist");

        assert!(result.is_err());
    }

    #[test]
    fn is_read_only_true_for_read_only_tool() {
        let registry = ToolRegistry::new().register(ReadOnlyTool);
        assert!(registry.is_read_only("lookup"));
    }

    #[test]
    fn is_read_only_false_for_default_tool() {
        let registry = ToolRegistry::new().register(AddTool);
        assert!(!registry.is_read_only("add"));
    }

    #[test]
    fn is_read_only_false_for_unknown_tool() {
        let registry = ToolRegistry::new().register(ReadOnlyTool);
        assert!(!registry.is_read_only("nonexistent"));
    }
}