mcphub/core/tool_registry.py

"""
Tool Registry - Central tool management for MCP

Manages tool definitions with type safety and validation.
Part of Option B clean architecture refactoring.
"""

import logging
from collections.abc import Callable
from typing import Any

from pydantic import BaseModel, ConfigDict, Field

logger = logging.getLogger(__name__)

class ToolDefinition(BaseModel):
    """
    Type-safe tool definition.

    Represents a single MCP tool with all necessary metadata.

    Attributes:
        name: Unique tool identifier (e.g., "wordpress_get_post")
        description: Human-readable tool description
        input_schema: JSON Schema for tool parameters
        handler: Async function that executes the tool
        required_scope: Required API key scope ("read", "write", "admin")
        plugin_type: Plugin type this tool belongs to (e.g., "wordpress")
    """

    name: str = Field(..., description="Unique tool identifier")
    description: str = Field(..., description="Tool description")
    input_schema: dict[str, Any] = Field(
        default_factory=lambda: {"type": "object", "properties": {}},
        description="JSON Schema for parameters",
    )
    handler: Callable = Field(..., description="Async handler function")
    required_scope: str = Field(
        default="read", description="Required API key scope (read/write/admin)"
    )
    plugin_type: str = Field(..., description="Plugin type (wordpress, gitea, etc)")

    model_config = ConfigDict(arbitrary_types_allowed=True)  # Allow Callable type

class ToolRegistry:
    """
    Central registry for all MCP tools.

    Manages tool registration, retrieval, and validation.
    Ensures unique tool names and provides filtering by plugin type.

    Examples:
        >>> registry = ToolRegistry()
        >>> tool = ToolDefinition(
        ...     name="wordpress_list_posts",
        ...     description="List WordPress posts",
        ...     handler=async_handler_func,
        ...     plugin_type="wordpress"
        ... )
        >>> registry.register(tool)
        >>> all_tools = registry.get_all()
        >>> wp_tools = registry.get_by_plugin_type("wordpress")
    """

    def __init__(self):
        """Initialize empty tool registry."""
        self.tools: dict[str, ToolDefinition] = {}
        self.logger = logging.getLogger("ToolRegistry")
        self.logger.info("ToolRegistry initialized")

    def register(self, tool: ToolDefinition) -> None:
        """
        Register a tool in the registry.

        Args:
            tool: Tool definition to register

        Raises:
            ValueError: If tool name already exists

        Examples:
            >>> registry.register(tool_definition)
        """
        if tool.name in self.tools:
            raise ValueError(f"Tool '{tool.name}' already registered")

        self.tools[tool.name] = tool
        self.logger.debug(f"Registered tool: {tool.name} ({tool.plugin_type})")

    def register_many(self, tools: list[ToolDefinition]) -> int:
        """
        Register multiple tools at once.

        Args:
            tools: List of tool definitions

        Returns:
            Number of tools successfully registered

        Examples:
            >>> count = registry.register_many([tool1, tool2, tool3])
            >>> print(f"Registered {count} tools")
        """
        count = 0
        for tool in tools:
            try:
                self.register(tool)
                count += 1
            except ValueError as e:
                self.logger.warning(f"Skipped duplicate tool: {e}")
            except Exception as e:
                self.logger.error(f"Failed to register tool: {e}", exc_info=True)

        self.logger.info(f"Registered {count}/{len(tools)} tools")
        return count

    def get_all(self) -> list[ToolDefinition]:
        """
        Get all registered tools.

        Returns:
            List of all tool definitions

        Examples:
            >>> tools = registry.get_all()
            >>> print(f"Total tools: {len(tools)}")
        """
        return list(self.tools.values())

    def get_by_name(self, name: str) -> ToolDefinition | None:
        """
        Get a tool by its name.

        Args:
            name: Tool name

        Returns:
            Tool definition if found, None otherwise

        Examples:
            >>> tool = registry.get_by_name("wordpress_get_post")
            >>> if tool:
            ...     print(f"Found: {tool.description}")
        """
        return self.tools.get(name)

    def get_by_plugin_type(self, plugin_type: str) -> list[ToolDefinition]:
        """
        Get all tools for a specific plugin type.

        Args:
            plugin_type: Plugin type (e.g., "wordpress", "gitea")

        Returns:
            List of tools for the specified plugin type

        Examples:
            >>> wp_tools = registry.get_by_plugin_type("wordpress")
            >>> print(f"WordPress tools: {len(wp_tools)}")
        """
        return [tool for tool in self.tools.values() if tool.plugin_type == plugin_type]

    def get_count(self) -> int:
        """
        Get total number of registered tools.

        Returns:
            Count of registered tools

        Examples:
            >>> count = registry.get_count()
        """
        return len(self.tools)

    def get_count_by_plugin(self) -> dict[str, int]:
        """
        Get tool counts grouped by plugin type.

        Returns:
            Dictionary mapping plugin type to tool count

        Examples:
            >>> counts = registry.get_count_by_plugin()
            >>> print(counts)  # {'wordpress': 95, 'gitea': 50}
        """
        counts = {}
        for tool in self.tools.values():
            plugin_type = tool.plugin_type
            counts[plugin_type] = counts.get(plugin_type, 0) + 1
        return counts

    def clear(self) -> None:
        """
        Clear all registered tools.

        Primarily for testing purposes.

        Examples:
            >>> registry.clear()
            >>> assert registry.get_count() == 0
        """
        self.tools.clear()
        self.logger.info("Tool registry cleared")

    def __repr__(self) -> str:
        """String representation of registry."""
        counts = self.get_count_by_plugin()
        counts_str = ", ".join(f"{k}: {v}" for k, v in counts.items())
        return f"ToolRegistry(total={self.get_count()}, {counts_str})"

# Singleton instance
_tool_registry: ToolRegistry | None = None

def get_tool_registry() -> ToolRegistry:
    """
    Get the singleton tool registry instance.

    Returns:
        Global ToolRegistry instance

    Examples:
        >>> registry = get_tool_registry()
        >>> registry.register(tool)
    """
    global _tool_registry
    if _tool_registry is None:
        _tool_registry = ToolRegistry()
    return _tool_registry