app/mcp/plugin_registry.py

"""
插件MCP工具动态注册器 - 扫描插件并注册MCP工具

扩展功能：
- 工具可用性绑定插件生命周期
- 最近错误记录
- 元数据完整性检查
"""

import importlib.util
import sys
import logging
import functools
import inspect
from pathlib import Path
from typing import Dict, Any, List, TYPE_CHECKING, Callable, Optional
from datetime import datetime

if TYPE_CHECKING:
    from mcp.server import FastMCP

logger = logging.getLogger(__name__)


class MCPToolAvailability:
    """MCP 工具可用状态常量"""

    AVAILABLE = "available"
    UNAVAILABLE = "unavailable"
    INCOMPLETE = "incomplete"


class PluginMCPRegistry:
    """插件MCP工具动态注册器

    扩展功能：
    - 工具可用性绑定插件生命周期
    - 最近错误记录
    - 元数据完整性检查
    """

    def __init__(self, mcp_server: "FastMCP"):
        self.mcp = mcp_server
        self.registered_tools: Dict[str, Dict[str, Any]] = {}
        self._tool_wrappers: Dict[str, Callable] = {}

    def _create_tool_wrapper(self, tool_name: str, original_func: Callable) -> Callable:
        """创建工具调用包装器，保持原始函数签名用于 schema 生成。

        包装器在调用前查询插件状态和依赖状态，不可用时直接返回结构化错误。
        捕获参数错误和运行异常，写入最近错误记录。
        """
        @functools.wraps(original_func)
        async def wrapper(*args, **kwargs):
            # 检查工具是否在注册表中
            tool_info = self.registered_tools.get(tool_name)
            if not tool_info:
                return {
                    "success": False,
                    "error": f"工具 {tool_name} 不存在",
                    "error_code": "TOOL_NOT_FOUND",
                }

            # 检查工具可用性
            availability = tool_info.get("availability", MCPToolAvailability.AVAILABLE)
            if availability == MCPToolAvailability.UNAVAILABLE:
                reason = tool_info.get("unavailable_reason", "工具不可用")
                return {
                    "success": False,
                    "error": f"工具 {tool_name} 不可用：{reason}",
                    "error_code": "TOOL_UNAVAILABLE",
                    "unavailable_reason": reason,
                }

            # 记录调用时间
            tool_info["last_called_at"] = datetime.now()

            try:
                # 调用原始函数
                result = await original_func(*args, **kwargs)
                return result
            except TypeError as e:
                # 参数错误
                error_msg = f"参数错误：{str(e)}"
                tool_info["last_error"] = {
                    "error_code": "INVALID_PARAMETERS",
                    "message": error_msg,
                    "occurred_at": datetime.now().isoformat(),
                }
                logger.error(f"工具 {tool_name} 参数错误: {e}")
                return {
                    "success": False,
                    "error": error_msg,
                    "error_code": "INVALID_PARAMETERS",
                }
            except Exception as e:
                # 运行时错误
                error_msg = f"运行错误：{str(e)}"
                tool_info["last_error"] = {
                    "error_code": "RUNTIME_ERROR",
                    "message": error_msg,
                    "occurred_at": datetime.now().isoformat(),
                }
                logger.error(f"工具 {tool_name} 运行错误: {e}")
                return {
                    "success": False,
                    "error": error_msg,
                    "error_code": "RUNTIME_ERROR",
                }

        # 保持原始函数签名
        wrapper.__signature__ = inspect.signature(original_func)
        return wrapper

    def scan_and_register_all(self, plugins_dir: Path) -> int:
        """扫描所有插件并注册MCP工具。"""
        total = 0
        for plugin_dir in plugins_dir.iterdir():
            if plugin_dir.is_dir() and not plugin_dir.name.startswith('_'):
                count = self.register_plugin_tools(plugin_dir)
                total += count
                if count > 0:
                    logger.info(f"插件 {plugin_dir.name} 注册了 {count} 个MCP工具")
        return total

    def register_plugin_tools_by_name(self, plugins_dir: Path, plugin_names: List[str]) -> int:
        """仅注册指定插件的 MCP 工具"""
        total = 0
        for plugin_name in plugin_names:
            plugin_dir = plugins_dir / plugin_name
            if plugin_dir.is_dir() and not plugin_dir.name.startswith('_'):
                count = self.register_plugin_tools(plugin_dir)
                total += count
                if count > 0:
                    logger.info(f"插件 {plugin_name} 注册了 {count} 个MCP工具")
        return total

    def register_plugin_tools(self, plugin_dir: Path) -> int:
        """注册单个插件的MCP工具。"""
        mcp_file = plugin_dir / "mcp.py"
        if not mcp_file.exists():
            return 0

        plugin_name = plugin_dir.name
        module = self._import_module(mcp_file, plugin_name)
        if not module:
            return 0

        tools = self._scan_tools(module)
        count = 0
        for tool_def in tools:
            # 直接使用 decorator 中定义的名称（约定：{插件名}-{工具名}）
            full_name = tool_def['name']
            try:
                wrapped_func = self._create_tool_wrapper(full_name, tool_def['func'])
                self._tool_wrappers[full_name] = wrapped_func

                self.mcp.tool(
                    name=full_name,
                    title=tool_def.get('title'),
                    description=tool_def.get('description'),
                )(wrapped_func)

                # 检查元数据完整性
                incomplete_reasons = []
                if not tool_def.get('title'):
                    incomplete_reasons.append("缺少标题")
                if not tool_def.get('description'):
                    incomplete_reasons.append("缺少描述")

                availability = MCPToolAvailability.AVAILABLE
                if incomplete_reasons:
                    availability = MCPToolAvailability.INCOMPLETE

                self.registered_tools[full_name] = {
                    "plugin": plugin_name,
                    "tool": tool_def['name'],
                    "input_model": tool_def.get('input_model'),
                    "output_model": tool_def.get('output_model'),
                    "title": tool_def.get('title'),
                    "description": tool_def.get('description'),
                    # 新增字段
                    "availability": availability,
                    "unavailable_reason": None,
                    "incomplete_reasons": incomplete_reasons if incomplete_reasons else None,
                    "last_error": None,
                    "last_called_at": None,
                    "registered_at": datetime.now().isoformat(),
                }
                count += 1
            except Exception as e:
                logger.error(f"注册工具 {full_name} 失败: {e}")

        return count

    def _import_module(self, mcp_file: Path, plugin_name: str):
        """动态导入模块。"""
        try:
            spec = importlib.util.spec_from_file_location(
                f"plugins.{plugin_name}.mcp", mcp_file
            )
            if not spec or not spec.loader:
                return None
            module = importlib.util.module_from_spec(spec)
            sys.modules[spec.name] = module
            spec.loader.exec_module(module)
            return module
        except Exception as e:
            logger.error(f"导入模块 {mcp_file} 失败: {e}")
            return None

    def _scan_tools(self, module) -> List[Dict]:
        """扫描模块中的工具定义。"""
        tools = []
        for attr_name in dir(module):
            attr = getattr(module, attr_name)
            if callable(attr) and hasattr(attr, '_mcp_tool_def'):
                tools.append(attr._mcp_tool_def)
        return tools

    def get_tool_docs(self, plugin_name: str = None) -> Dict:
        """获取工具文档。"""
        if plugin_name:
            return {
                name: info for name, info in self.registered_tools.items()
                if info['plugin'] == plugin_name
            }
        return self.registered_tools.copy()

    def unregister_plugin_tools(self, plugin_name: str) -> int:
        """注销插件的MCP工具。"""
        tools_to_remove = [
            name for name, info in self.registered_tools.items()
            if info['plugin'] == plugin_name
        ]
        for tool_name in tools_to_remove:
            del self.registered_tools[tool_name]
        logger.info(f"已注销插件 {plugin_name} 的 {len(tools_to_remove)} 个工具")
        return len(tools_to_remove)

    def get_plugin_tools(self, plugin_name: str) -> List[str]:
        """获取插件的所有工具名称。"""
        return [
            name for name, info in self.registered_tools.items()
            if info['plugin'] == plugin_name
        ]

    def set_tool_availability(
        self,
        tool_name: str,
        availability: str,
        unavailable_reason: Optional[str] = None,
    ) -> bool:
        """设置工具可用状态

        Args:
            tool_name: 工具全名
            availability: 可用状态
            unavailable_reason: 不可用原因

        Returns:
            是否设置成功
        """
        tool_info = self.registered_tools.get(tool_name)
        if not tool_info:
            return False

        tool_info["availability"] = availability
        tool_info["unavailable_reason"] = unavailable_reason
        return True

    def set_plugin_tools_availability(
        self,
        plugin_name: str,
        availability: str,
        unavailable_reason: Optional[str] = None,
    ) -> int:
        """设置插件所有工具的可用状态

        Args:
            plugin_name: 插件名
            availability: 可用状态
            unavailable_reason: 不可用原因

        Returns:
            更新的工具数量
        """
        count = 0
        for tool_name, tool_info in self.registered_tools.items():
            if tool_info["plugin"] == plugin_name:
                tool_info["availability"] = availability
                tool_info["unavailable_reason"] = unavailable_reason
                count += 1
        return count

    def get_mcp_status(self) -> Dict[str, Any]:
        """获取 MCP 服务状态

        Returns:
            包含服务状态、工具总数、按插件分组、最近错误的字典
        """
        total_tools = len(self.registered_tools)
        available_tools = sum(
            1 for info in self.registered_tools.values()
            if info.get("availability") == MCPToolAvailability.AVAILABLE
        )
        unavailable_tools = sum(
            1 for info in self.registered_tools.values()
            if info.get("availability") == MCPToolAvailability.UNAVAILABLE
        )
        incomplete_tools = sum(
            1 for info in self.registered_tools.values()
            if info.get("availability") == MCPToolAvailability.INCOMPLETE
        )

        # 按插件分组
        groups = {}
        for tool_name, tool_info in self.registered_tools.items():
            plugin_name = tool_info["plugin"]
            if plugin_name not in groups:
                groups[plugin_name] = {
                    "plugin_name": plugin_name,
                    "tools": [],
                }
            groups[plugin_name]["tools"].append({
                "name": tool_name,
                "title": tool_info.get("title"),
                "description": tool_info.get("description"),
                "availability": tool_info.get("availability"),
                "unavailable_reason": tool_info.get("unavailable_reason"),
                "incomplete_reasons": tool_info.get("incomplete_reasons"),
                "last_error": tool_info.get("last_error"),
                "last_called_at": tool_info.get("last_called_at"),
            })

        # 收集最近错误
        recent_errors = []
        for tool_name, tool_info in self.registered_tools.items():
            if tool_info.get("last_error"):
                recent_errors.append({
                    "tool_name": tool_name,
                    "plugin_name": tool_info["plugin"],
                    **tool_info["last_error"],
                })

        return {
            "service_available": True,
            "total_tools": total_tools,
            "available_tools": available_tools,
            "unavailable_tools": unavailable_tools,
            "incomplete_tools": incomplete_tools,
            "groups": list(groups.values()),
            "recent_errors": recent_errors,
        }


# 全局单例
_plugin_registry: Optional[PluginMCPRegistry] = None


def get_plugin_registry() -> PluginMCPRegistry:
    """获取全局插件注册器实例。"""
    global _plugin_registry
    if _plugin_registry is None:
        raise RuntimeError("PluginMCPRegistry 未初始化，请先调用 init_plugin_registry(mcp)")
    return _plugin_registry


def init_plugin_registry(mcp_server: "FastMCP") -> PluginMCPRegistry:
    """初始化全局插件注册器。"""
    global _plugin_registry
    _plugin_registry = PluginMCPRegistry(mcp_server)
    return _plugin_registry