blank/backend/mcp_client.py

"""
MCP 客户端 - 工具发现和调用（支持 SSE 响应）
"""
import json
import httpx
import asyncio
import re
from typing import Dict, List, Any, Optional
from config import MCP_SERVERS
from debug_logger import log_debug


def parse_sse_response(text: str) -> str:
    """
    解析 SSE 响应，提取 JSON 数据

    SSE 格式:
    event: message
    data: {...json...}

    Args:
        text: SSE 响应文本

    Returns:
        提取的 JSON 字符串
    """
    # 查找 data: 行并提取 JSON
    for line in text.split('\n'):
        if line.startswith('data:'):
            data_content = line[5:].strip()
            if data_content:
                return data_content
    return text


class MCPClient:
    """MCP 客户端，负责工具发现和调用"""

    def __init__(self, server_id: str = None, session_id: str = None, auth_token: str = None):
        self.server_id = server_id or "novel-translator"
        self.server_config = MCP_SERVERS.get(self.server_id, {})
        self.session_id = session_id
        self.auth_token = auth_token  # JWT token for authenticated MCPs
        self.base_url = self.server_config.get("url", "")

    async def discover_tools(self) -> List[Dict[str, Any]]:
        """从 MCP 服务器发现可用工具"""
        if not self.base_url:
            return []

        try:
            async with httpx.AsyncClient(timeout=30.0, verify=False) as client:
                url = f"{self.base_url}"
                headers = {
                    "Content-Type": "application/json",
                    "Accept": "application/json, text/event-stream"
                }

                if self.session_id:
                    headers["X-Session-ID"] = self.session_id

                # 添加 JWT token 认证
                if self.auth_token:
                    headers["Authorization"] = f"Bearer {self.auth_token}"

                payload = {
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "tools/list"
                }

                response = await client.post(url, json=payload, headers=headers)

                if response.status_code == 200:
                    # 解析 SSE 响应
                    json_text = parse_sse_response(response.text)
                    result = json.loads(json_text)

                    if "result" in result:
                        return result["result"].get("tools", [])

                return []

        except Exception as e:
            print(f"MCP 工具发现失败: {e}")
            import traceback
            traceback.print_exc()
            return []

    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """调用 MCP 工具"""
        if not self.base_url:
            return {"error": "MCP 服务器未配置"}

        try:
            async with httpx.AsyncClient(timeout=120.0, verify=False) as client:
                url = f"{self.base_url}"
                headers = {
                    "Content-Type": "application/json",
                    "Accept": "application/json, text/event-stream"
                }

                if self.session_id:
                    headers["X-Session-ID"] = self.session_id

                # 添加 JWT token 认证
                if self.auth_token:
                    headers["Authorization"] = f"Bearer {self.auth_token}"

                # DEBUG: 打印调用详情
                print(f"[DEBUG MCPClient.call_tool] tool_name: {tool_name}")
                print(f"[DEBUG MCPClient.call_tool]   server_id: {self.server_id}")
                print(f"[DEBUG MCPClient.call_tool]   url: {url}")
                print(f"[DEBUG MCPClient.call_tool]   auth_token present: {bool(self.auth_token)}")
                if self.auth_token:
                    print(f"[DEBUG MCPClient.call_tool]   auth_token: {self.auth_token[:30]}...")
                    print(f"[DEBUG MCPClient.call_tool]   Authorization header: {headers['Authorization'][:50]}...")
                else:
                    print(f"[DEBUG MCPClient.call_tool]   NO auth_token!")

                payload = {
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/call",
                    "params": {
                        "name": tool_name,
                        "arguments": arguments
                    }
                }

                response = await client.post(url, json=payload, headers=headers)

                print(f"[DEBUG MCPClient.call_tool]   response status: {response.status_code}")

                if response.status_code == 200:
                    # 解析 SSE 响应
                    json_text = parse_sse_response(response.text)
                    result = json.loads(json_text)

                    if "result" in result:
                        content_list = result["result"].get("content", [])
                        text_results = []
                        for item in content_list:
                            if item.get("type") == "text":
                                text_results.append(item.get("text", ""))
                        return {
                            "success": True,
                            "result": "\n".join(text_results),
                            "raw": result["result"]
                        }
                    elif "error" in result:
                        return {"error": result["error"].get("message", "Unknown error")}

                return {"error": f"工具调用失败: {response.status_code}"}

        except Exception as e:
            print(f"MCP 工具调用失败: {e}")
            import traceback
            traceback.print_exc()
            return {"error": str(e)}

    @staticmethod
    async def get_all_tools_async(session_id: str = None) -> List[Dict[str, Any]]:
        """获取所有已配置 MCP 服务器的工具列表（异步版本）"""
        all_tools = []

        for server_id in MCP_SERVERS.keys():
            if not MCP_SERVERS[server_id].get("enabled", False):
                continue

            client = MCPClient(server_id, session_id)
            try:
                tools = await client.discover_tools()
                for tool in tools:
                    tool["_server_id"] = server_id
                    all_tools.append(tool)
            except Exception as e:
                print(f"发现 {server_id} 工具失败: {e}")

        return all_tools

    @staticmethod
    async def get_all_tools_with_tokens_async(
        session_id: str = None,
        mcp_tokens: dict = None,
        enabled_mcp_list: list = None  # 新增：前端传递的已启用 MCP 列表
    ) -> List[Dict[str, Any]]:
        """
        获取所有已配置 MCP 服务器的工具列表（带 token 认证）

        Args:
            session_id: 会话 ID
            mcp_tokens: MCP token 映射
            enabled_mcp_list: 前端传递的已启用 MCP 列表（优先级高于 config.py 中的配置）
                            - None: 前端未传递，使用 config.py fallback
                            - []: 前端明确禁用所有 MCP
                            - [xxx]: 前端指定启用的 MCP 列表
        """
        all_tools = []

        # DEBUG: 打印接收到的启用列表
        print(f"[DEBUG MCPClient.get_all_tools_with_tokens_async] enabled_mcp_list: {enabled_mcp_list}")
        log_debug("mcp_client.get_all_tools_with_tokens_async_start", {
            "enabled_mcp_list": enabled_mcp_list,
            "enabled_mcp_list_type": str(type(enabled_mcp_list)),
            "enabled_mcp_list_is_none": enabled_mcp_list is None,
            "enabled_mcp_list_len": len(enabled_mcp_list) if enabled_mcp_list is not None else "N/A"
        })

        for server_id in MCP_SERVERS.keys():
            # 优先使用前端传递的启用列表
            log_debug("mcp_client.checking_server", {
                "server_id": server_id,
                "enabled_mcp_list": enabled_mcp_list,
                "is_none": enabled_mcp_list is None
            })
            if enabled_mcp_list is not None:
                # 前端传递了启用列表（可能是空数组），只处理列表中的 MCP
                in_list = server_id in enabled_mcp_list
                log_debug("mcp_client.server_in_list_check", {
                    "server_id": server_id,
                    "in_enabled_list": in_list,
                    "will_skip": not in_list
                })
                if server_id not in enabled_mcp_list:
                    print(f"[DEBUG MCPClient] Skipping {server_id} (not in enabled_mcp_list from frontend)")
                    continue
            else:
                # 前端未传递启用列表（None），使用配置文件中的 enabled 状态作为 fallback
                if not MCP_SERVERS[server_id].get("enabled", False):
                    continue

            # 获取该服务器的 token
            auth_token = None
            if mcp_tokens and server_id in mcp_tokens:
                auth_token = mcp_tokens[server_id]

            print(f"[DEBUG MCPClient] Discovering tools for {server_id} (has_token: {bool(auth_token)})")

            client = MCPClient(server_id, session_id, auth_token)
            try:
                tools = await client.discover_tools()
                for tool in tools:
                    tool["_server_id"] = server_id
                    all_tools.append(tool)
            except Exception as e:
                print(f"发现 {server_id} 工具失败: {e}")

        return all_tools

    @staticmethod
    def get_all_tools(session_id: str = None) -> List[Dict[str, Any]]:
        """获取所有已配置 MCP 服务器的工具列表（同步版本）"""
        return asyncio.run(MCPClient.get_all_tools_async(session_id))