dss/tools/dss_mcp/plugin_registry.py

"""
Dynamic Plugin Registry for DSS MCP Server

Automatically discovers and registers MCP tools from the plugins/ directory.
Plugins follow a simple contract: export TOOLS list and a handler class with execute_tool() method.
"""

import pkgutil
import importlib
import inspect
import logging
import types as python_types
from typing import List, Dict, Any, Optional
from mcp import types

logger = logging.getLogger("dss.mcp.plugins")


class PluginRegistry:
    """
    Discovers and manages dynamically loaded plugins.

    Plugin Contract:
    - Must export TOOLS: List[types.Tool] - MCP tool definitions
    - Must have a class with execute_tool(name: str, arguments: dict) method
    - Optional: PLUGIN_METADATA dict with name, version, author

    Example Plugin Structure:
    ```python
    from mcp import types

    PLUGIN_METADATA = {
        "name": "Example Plugin",
        "version": "1.0.0",
        "author": "DSS Team"
    }

    TOOLS = [
        types.Tool(
            name="example_tool",
            description="Example tool",
            inputSchema={...}
        )
    ]

    class PluginTools:
        async def execute_tool(self, name: str, arguments: dict):
            if name == "example_tool":
                return {"result": "success"}
            raise ValueError(f"Unknown tool: {name}")
    ```
    """

    def __init__(self):
        self.tools: List[types.Tool] = []
        self.handlers: Dict[str, Any] = {}  # tool_name -> handler_instance
        self.plugins: List[Dict[str, Any]] = []  # plugin metadata
        self._loaded_modules: set = set()

    def load_plugins(self, plugins_package_name: str = "dss_mcp.plugins"):
        """
        Scans the plugins directory and registers valid tool modules.

        Args:
            plugins_package_name: Fully qualified name of plugins package
                                 Default: "dss_mcp.plugins" (works when called from tools/ dir)
        """
        try:
            # Dynamically import the plugins package
            plugins_pkg = importlib.import_module(plugins_package_name)
            path = plugins_pkg.__path__
            prefix = plugins_pkg.__name__ + "."

            logger.info(f"Scanning for plugins in: {path}")

            # Iterate through all modules in the plugins directory
            for _, name, is_pkg in pkgutil.iter_modules(path, prefix):
                # Skip packages (only load .py files)
                if is_pkg:
                    continue

                # Skip template and private modules
                module_basename = name.split('.')[-1]
                if module_basename.startswith('_'):
                    logger.debug(f"Skipping private module: {module_basename}")
                    continue

                try:
                    module = importlib.import_module(name)
                    self._register_module(module)
                except Exception as e:
                    logger.error(f"Failed to load plugin module {name}: {e}", exc_info=True)

        except ImportError as e:
            logger.warning(f"Plugins package not found: {plugins_package_name} ({e})")
            logger.info("Server will run without plugins")

    def _register_module(self, module: python_types.ModuleType):
        """
        Validates and registers a single plugin module.

        Args:
            module: The imported plugin module
        """
        module_name = module.__name__

        # Check if already loaded
        if module_name in self._loaded_modules:
            logger.debug(f"Module already loaded: {module_name}")
            return

        # Contract Check 1: Must export TOOLS list
        if not hasattr(module, 'TOOLS'):
            logger.debug(f"Module {module_name} has no TOOLS export, skipping")
            return

        if not isinstance(module.TOOLS, list):
            logger.error(f"Module {module_name} TOOLS must be a list, got {type(module.TOOLS)}")
            return

        if len(module.TOOLS) == 0:
            logger.warning(f"Module {module_name} has empty TOOLS list")
            return

        # Contract Check 2: Must have a class with execute_tool method
        handler_instance = self._find_and_instantiate_handler(module)
        if not handler_instance:
            logger.warning(f"Plugin {module_name} has TOOLS but no valid handler class")
            return

        # Contract Check 3: execute_tool must be async (coroutine)
        execute_tool_method = getattr(handler_instance, 'execute_tool', None)
        if execute_tool_method and not inspect.iscoroutinefunction(execute_tool_method):
            logger.error(
                f"Plugin '{module_name}' is invalid: 'PluginTools.execute_tool' must be "
                f"an async function ('async def'). Skipping plugin."
            )
            return

        # Extract metadata
        metadata = getattr(module, 'PLUGIN_METADATA', {})
        plugin_name = metadata.get('name', module_name.split('.')[-1])
        plugin_version = metadata.get('version', 'unknown')

        # Validate tools and check for name collisions
        registered_count = 0
        for tool in module.TOOLS:
            if not hasattr(tool, 'name'):
                logger.error(f"Tool in {module_name} missing 'name' attribute")
                continue

            # Check for name collision
            if tool.name in self.handlers:
                logger.error(
                    f"Tool name collision: '{tool.name}' already registered. "
                    f"Skipping duplicate from {module_name}"
                )
                continue

            # Register tool
            self.tools.append(tool)
            self.handlers[tool.name] = handler_instance
            registered_count += 1
            logger.debug(f"Registered tool: {tool.name}")

        # Track plugin metadata
        self.plugins.append({
            "name": plugin_name,
            "version": plugin_version,
            "module": module_name,
            "tools_count": registered_count,
            "author": metadata.get('author', 'unknown')
        })

        self._loaded_modules.add(module_name)

        logger.info(
            f"Loaded plugin: {plugin_name} v{plugin_version} "
            f"({registered_count} tools from {module_name})"
        )

    def _find_and_instantiate_handler(self, module: python_types.ModuleType) -> Optional[Any]:
        """
        Finds a class implementing execute_tool and instantiates it.

        Args:
            module: The plugin module to search

        Returns:
            Instantiated handler class or None if not found
        """
        for name, obj in inspect.getmembers(module, inspect.isclass):
            # Only consider classes defined in this module (not imports)
            if obj.__module__ != module.__name__:
                continue

            # Look for execute_tool method
            if hasattr(obj, 'execute_tool'):
                try:
                    # Try to instantiate with no args
                    instance = obj()
                    logger.debug(f"Instantiated handler class: {name}")
                    return instance
                except TypeError:
                    # Try with **kwargs for flexible initialization
                    try:
                        instance = obj(**{})
                        logger.debug(f"Instantiated handler class with kwargs: {name}")
                        return instance
                    except Exception as e:
                        logger.error(
                            f"Failed to instantiate handler {name} in {module.__name__}: {e}"
                        )
                        return None
                except Exception as e:
                    logger.error(
                        f"Failed to instantiate handler {name} in {module.__name__}: {e}"
                    )
                    return None

        return None

    async def execute_tool(self, name: str, arguments: dict) -> Any:
        """
        Routes tool execution to the correct plugin handler.

        Args:
            name: Tool name
            arguments: Tool arguments

        Returns:
            Tool execution result

        Raises:
            ValueError: If tool not found in registry
        """
        if name not in self.handlers:
            raise ValueError(f"Tool '{name}' not found in plugin registry")

        handler = self.handlers[name]

        # Support both async and sync implementations
        if inspect.iscoroutinefunction(handler.execute_tool):
            return await handler.execute_tool(name, arguments)
        else:
            return handler.execute_tool(name, arguments)

    def get_all_tools(self) -> List[types.Tool]:
        """Get merged list of all plugin tools"""
        return self.tools.copy()

    def get_plugin_info(self) -> List[Dict[str, Any]]:
        """Get metadata for all loaded plugins"""
        return self.plugins.copy()

    def reload_plugins(self, plugins_package_name: str = "dss_mcp.plugins"):
        """
        Reload all plugins (useful for development).
        WARNING: This clears all registered plugins and reloads from scratch.

        Args:
            plugins_package_name: Fully qualified name of plugins package
        """
        logger.info("Reloading all plugins...")

        # Clear existing registrations
        self.tools.clear()
        self.handlers.clear()
        self.plugins.clear()
        self._loaded_modules.clear()

        # Reload
        self.load_plugins(plugins_package_name)

        logger.info(f"Plugin reload complete. Loaded {len(self.plugins)} plugins, {len(self.tools)} tools")