awoooi/apps/api/src/plugins/mcp/registry.py

"""
MCP Provider Registry - ADR-015 模組化架構
==========================================

Provider 註冊中心，實現依賴注入 (DI) 模式：
1. 統一管理所有 MCP Tool Providers
2. 支援動態註冊/反註冊
3. 支援健康檢查

@see docs/adr/ADR-015-mcp-modular-architecture.md
"""

import structlog

from src.plugins.mcp.interfaces import MCPTool, MCPToolProvider, MCPToolResult

logger = structlog.get_logger(__name__)


class AuditedMCPToolProvider(MCPToolProvider):
    """Provider wrapper that writes every MCP tool call to the audit subsystem."""

    def __init__(self, provider: MCPToolProvider) -> None:
        # __provider 使用 Python name mangling（_AuditedMCPToolProvider__provider）
        # 防止 caller 透過 wrapper._provider 直接存取 inner provider（ADR-116 封裝要求）
        self.__provider = provider  # noqa: SLF001 — intentional name mangling

    @property
    def name(self) -> str:
        return self.__provider.name

    @property
    def enabled(self) -> bool:
        return self.__provider.enabled

    async def list_tools(self) -> list[MCPTool]:
        return await self.__provider.list_tools()

    async def execute(
        self,
        tool_name: str,
        parameters: dict,
    ) -> MCPToolResult:
        from src.services.mcp_audit_service import monotonic_ms, record_mcp_call

        audit_context = parameters.get("_mcp_audit") if isinstance(parameters, dict) else None
        provider_parameters = {
            key: value for key, value in parameters.items()
            if key != "_mcp_audit"
        }
        started = monotonic_ms()
        result: MCPToolResult | None = None
        try:
            result = await self.__provider.execute(tool_name, provider_parameters)
            return result
        finally:
            duration_ms = monotonic_ms() - started
            await record_mcp_call(
                mcp_server=self.name,
                tool_name=tool_name,
                input_params=parameters,
                output_result=result.output if result else None,
                duration_ms=duration_ms,
                success=bool(result.success) if result else False,
                error_message=result.error if result else "provider_exception",
                session_id=audit_context.get("session_id") if isinstance(audit_context, dict) else None,
                flywheel_node=audit_context.get("flywheel_node") if isinstance(audit_context, dict) else None,
                incident_id=audit_context.get("incident_id") if isinstance(audit_context, dict) else None,
                agent_role=audit_context.get("agent_role") if isinstance(audit_context, dict) else None,
            )

    async def health_check(self) -> bool:
        return await self.__provider.health_check()


class ProviderRegistry:
    """
    MCP Tool Provider 註冊中心

    使用方式:
        # 註冊
        registry = ProviderRegistry()
        registry.register(K8sProvider())
        registry.register(SignOzProvider())

        # 取得
        k8s = registry.get("kubernetes")
        await k8s.execute("kubectl_get", {...})

        # 列出所有
        for provider in registry.all():
            print(provider.name)
    """

    def __init__(self) -> None:
        self._providers: dict[str, MCPToolProvider] = {}

    def register(self, provider: MCPToolProvider) -> None:
        """
        註冊 Provider

        Args:
            provider: MCPToolProvider 實例

        Raises:
            ValueError: 如果 Provider 名稱已存在
        """
        if provider.name in self._providers:
            raise ValueError(f"Provider '{provider.name}' already registered")

        self._providers[provider.name] = AuditedMCPToolProvider(provider)
        logger.info(
            "provider_registered",
            name=provider.name,
            enabled=provider.enabled,
        )

    def unregister(self, name: str) -> bool:
        """
        反註冊 Provider

        Args:
            name: Provider 名稱

        Returns:
            bool: 是否成功反註冊
        """
        if name in self._providers:
            del self._providers[name]
            logger.info("provider_unregistered", name=name)
            return True
        return False

    def get(self, name: str) -> MCPToolProvider | None:
        """
        取得 Provider

        Args:
            name: Provider 名稱

        Returns:
            MCPToolProvider | None: Provider 實例，若不存在則返回 None
        """
        provider = self._providers.get(name)
        if provider and not provider.enabled:
            logger.warning("provider_disabled", name=name)
            return None
        return provider

    def all(self) -> list[MCPToolProvider]:
        """
        取得所有已啟用的 Providers

        Returns:
            list[MCPToolProvider]: 已啟用的 Provider 列表
        """
        return [p for p in self._providers.values() if p.enabled]

    def names(self) -> list[str]:
        """
        取得所有已啟用的 Provider 名稱

        Returns:
            list[str]: Provider 名稱列表
        """
        return [p.name for p in self.all()]

    async def health_check_all(self) -> dict[str, bool]:
        """
        檢查所有 Provider 健康狀態

        Returns:
            dict[str, bool]: {provider_name: is_healthy}
        """
        results = {}
        for provider in self.all():
            try:
                results[provider.name] = await provider.health_check()
            except Exception as e:
                logger.warning(
                    "provider_health_check_failed",
                    name=provider.name,
                    error=str(e),
                )
                results[provider.name] = False
        return results

    def __contains__(self, name: str) -> bool:
        return name in self._providers

    def __len__(self) -> int:
        return len(self._providers)


# =============================================================================
# Global Registry Singleton
# =============================================================================

_registry: ProviderRegistry | None = None


def get_provider_registry() -> ProviderRegistry:
    """
    取得全域 Provider Registry

    Returns:
        ProviderRegistry: 單例實例
    """
    global _registry
    if _registry is None:
        _registry = ProviderRegistry()
    return _registry


def register_provider(provider: MCPToolProvider) -> None:
    """
    便捷函數: 註冊 Provider 到全域 Registry
    """
    get_provider_registry().register(provider)


def get_provider(name: str) -> MCPToolProvider | None:
    """
    便捷函數: 從全域 Registry 取得 Provider
    """
    return get_provider_registry().get(name)