coro_core/tools/

registry.rs

//! Tool registry for managing available tools

use crate::tools::{Tool, ToolExecutor};
use std::collections::HashMap;

/// Registry for managing tool creation and registration
pub struct ToolRegistry {
    factories: HashMap<String, Box<dyn ToolFactory>>,
}

/// Factory trait for creating tools
pub trait ToolFactory: Send + Sync {
    /// Create a new instance of the tool
    fn create(&self) -> Box<dyn Tool>;

    /// Get the name of the tool this factory creates
    fn tool_name(&self) -> &str;

    /// Get the description of the tool this factory creates
    fn tool_description(&self) -> &str;
}

impl ToolRegistry {
    /// Create a new tool registry
    pub fn new() -> Self {
        Self {
            factories: HashMap::new(),
        }
    }

    /// Register a tool factory
    pub fn register_factory(&mut self, factory: Box<dyn ToolFactory>) {
        self.factories
            .insert(factory.tool_name().to_string(), factory);
    }

    /// Create a tool by name
    pub fn create_tool(&self, name: &str) -> Option<Box<dyn Tool>> {
        self.factories.get(name).map(|factory| factory.create())
    }

    /// List all available tool names
    pub fn list_tools(&self) -> Vec<&str> {
        self.factories.keys().map(|s| s.as_str()).collect()
    }

    /// Get tool information
    pub fn get_tool_info(&self, name: &str) -> Option<(&str, &str)> {
        self.factories
            .get(name)
            .map(|factory| (factory.tool_name(), factory.tool_description()))
    }

    /// Create a tool executor with the specified tools
    pub fn create_executor(&self, tool_names: &[String]) -> ToolExecutor {
        let mut executor = ToolExecutor::new();

        for name in tool_names {
            if let Some(tool) = self.create_tool(name) {
                executor.register_tool(tool);
            }
        }

        executor
    }

    /// Create a tool executor with all available tools
    pub fn create_executor_with_all(&self) -> ToolExecutor {
        let mut executor = ToolExecutor::new();

        for factory in self.factories.values() {
            executor.register_tool(factory.create());
        }

        executor
    }
}

impl Default for ToolRegistry {
    fn default() -> Self {
        let mut registry = Self::new();

        // Register built-in tools
        registry.register_factory(Box::new(crate::tools::builtin::ThinkingToolFactory));
        registry.register_factory(Box::new(crate::tools::builtin::TaskDoneToolFactory));
        registry.register_factory(Box::new(crate::tools::builtin::McpToolFactory));

        registry
    }
}

/// Macro to help implement tool factories
#[macro_export]
macro_rules! impl_tool_factory {
    ($factory:ident, $tool:ident, $name:expr, $description:expr) => {
        pub struct $factory;

        impl $crate::tools::ToolFactory for $factory {
            fn create(&self) -> Box<dyn $crate::tools::Tool> {
                Box::new($tool::new())
            }

            fn tool_name(&self) -> &str {
                $name
            }

            fn tool_description(&self) -> &str {
                $description
            }
        }
    };
}

#[cfg(test)]
mod tests {
    use crate::tools::registry::ToolRegistry;

    #[test]
    fn test_default_registry_has_all_tools() {
        let registry = ToolRegistry::default();
        let tools = registry.list_tools();

        // Expected tools based on Python version
        let expected_tools = vec!["sequentialthinking", "task_done", "mcp_tool"];

        println!("Available tools: {:?}", tools);

        // Check that all expected tools are registered
        for expected_tool in &expected_tools {
            assert!(
                tools.contains(expected_tool),
                "Tool '{}' is not registered in the default registry",
                expected_tool
            );
        }

        // Check that we have the expected number of tools
        assert_eq!(
            tools.len(),
            expected_tools.len(),
            "Expected {} tools, but found {}. Tools: {:?}",
            expected_tools.len(),
            tools.len(),
            tools
        );
    }

    #[test]
    fn test_tool_creation() {
        let registry = ToolRegistry::default();

        // Test creating each tool
        let tools_to_test = vec!["sequentialthinking", "task_done", "mcp_tool"];

        for tool_name in tools_to_test {
            let tool = registry.create_tool(tool_name);
            assert!(tool.is_some(), "Failed to create tool '{}'", tool_name);

            let tool = tool.unwrap();
            assert_eq!(
                tool.name(),
                tool_name,
                "Tool name mismatch for '{}'",
                tool_name
            );

            // Verify tool has a description
            assert!(
                !tool.description().is_empty(),
                "Tool '{}' has empty description",
                tool_name
            );

            // Verify tool has parameters schema
            let schema = tool.parameters_schema();
            assert!(
                schema.is_object(),
                "Tool '{}' parameters schema is not an object",
                tool_name
            );
        }
    }

    #[test]
    fn test_tool_info() {
        let registry = ToolRegistry::default();

        for tool_name in registry.list_tools() {
            let info = registry.get_tool_info(tool_name);
            assert!(
                info.is_some(),
                "Failed to get info for tool '{}'",
                tool_name
            );

            let (name, description) = info.unwrap();
            assert_eq!(name, tool_name);
            assert!(!description.is_empty());
        }
    }

    #[test]
    fn test_executor_creation() {
        let registry = ToolRegistry::default();

        // Test creating executor with specific tools
        let tool_names = vec![
            "bash".to_string(),
            "str_replace_based_edit_tool".to_string(),
        ];
        let _executor = registry.create_executor(&tool_names);

        // The executor should have the requested tools
        // Note: We can't easily test this without exposing internal state
        // This test mainly ensures the method doesn't panic

        // Test creating executor with all tools
        let _all_executor = registry.create_executor_with_all();
    }

    #[test]
    fn test_tool_examples() {
        let registry = ToolRegistry::default();

        for tool_name in registry.list_tools() {
            let tool = registry.create_tool(tool_name).unwrap();
            let examples = tool.examples();

            // Each tool should have at least one example
            assert!(!examples.is_empty(), "Tool '{}' has no examples", tool_name);

            // Verify example structure
            for (i, example) in examples.iter().enumerate() {
                assert!(
                    !example.description.is_empty(),
                    "Tool '{}' example {} has empty description",
                    tool_name,
                    i
                );

                assert!(
                    example.parameters.is_object(),
                    "Tool '{}' example {} parameters is not an object",
                    tool_name,
                    i
                );

                assert!(
                    !example.expected_result.is_empty(),
                    "Tool '{}' example {} has empty expected result",
                    tool_name,
                    i
                );
            }
        }
    }

    #[test]
    fn test_tool_parameter_schemas() {
        let registry = ToolRegistry::default();

        for tool_name in registry.list_tools() {
            let tool = registry.create_tool(tool_name).unwrap();
            let schema = tool.parameters_schema();

            // Schema should be an object
            assert!(
                schema.is_object(),
                "Tool '{}' schema is not an object",
                tool_name
            );

            let schema_obj = schema.as_object().unwrap();

            // Should have type property
            if let Some(type_val) = schema_obj.get("type") {
                assert_eq!(
                    type_val.as_str(),
                    Some("object"),
                    "Tool '{}' schema type is not 'object'",
                    tool_name
                );
            }

            // Should have properties
            if let Some(properties) = schema_obj.get("properties") {
                assert!(
                    properties.is_object(),
                    "Tool '{}' schema properties is not an object",
                    tool_name
                );

                let props = properties.as_object().unwrap();
                assert!(
                    !props.is_empty(),
                    "Tool '{}' has no properties in schema",
                    tool_name
                );
            }
        }
    }
}