09. MCP协议与AI工具生态系统¶

"""
MCP核心组件架构

1. MCP Server: 提供工具、资源和能力的后端服务
2. MCP Client: 与Server通信的客户端组件
3. Protocol: 定义通信规范
4. Transport: 数据传输层（stdio/sse/http）
"""

from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass
from enum import Enum
import json

class MCPError(Exception):
    """MCP协议错误基类"""
    pass

class ToolNotFoundError(MCPError):
    """工具未找到错误"""
    pass

class PermissionDeniedError(MCPError):
    """权限拒绝错误"""
    pass

@dataclass  # @dataclass自动生成__init__等方法
class MCPTool:
    """MCP工具定义"""
    name: str
    description: str
    input_schema: Dict[str, Any]  # JSON Schema
    output_schema: Optional[Dict[str, Any]] = None  # Optional表示值可以为None
    annotations: Optional[Dict[str, Any]] = None

    def to_dict(self) -> Dict:
        """转换为字典表示"""
        return {
            "name": self.name,
            "description": self.description,
            "inputSchema": self.input_schema,
            "outputSchema": self.output_schema,
            "annotations": self.annotations
        }

@dataclass
class MCPResource:
    """MCP资源定义"""
    uri: str
    name: str
    description: str
    mime_type: Optional[str] = None

    def to_dict(self) -> Dict:
        return {
            "uri": self.uri,
            "name": self.name,
            "description": self.description,
            "mimeType": self.mime_type
        }

@dataclass
class MCPCapability:
    """MCP能力声明"""
    tools: bool = False
    resources: bool = False
    prompts: bool = False
    logging: bool = False

    def to_dict(self) -> Dict:
        result = {}
        if self.tools:
            result["tools"] = {}
        if self.resources:
            result["resources"] = {}
        if self.prompts:
            result["prompts"] = {}
        if self.logging:
            result["logging"] = {}
        return result

class MCPMessage:
    """MCP协议消息基类"""

    def __init__(self, jsonrpc: str = "2.0"):
        self.jsonrpc = jsonrpc

class MCPRequest(MCPMessage):
    """MCP请求消息"""

    def __init__(self, method: str, params: Dict = None, request_id: str = None):
        super().__init__()  # super()调用父类方法
        self.id = request_id or self._generate_id()
        self.method = method
        self.params = params or {}

    def to_dict(self) -> Dict:
        return {
            "jsonrpc": self.jsonrpc,
            "id": self.id,
            "method": self.method,
            "params": self.params
        }

    def _generate_id(self) -> str:
        import uuid
        return str(uuid.uuid4())

class MCPResponse(MCPMessage):
    """MCP响应消息"""

    def __init__(self, request_id: str, result: Any = None, error: Dict = None):
        super().__init__()
        self.id = request_id
        self.result = result
        self.error = error

    def to_dict(self) -> Dict:
        result = {
            "jsonrpc": self.jsonrpc,
            "id": self.id
        }
        if self.error:
            result["error"] = self.error
        else:
            result["result"] = self.result
        return result

    @classmethod  # 类方法：cls指向类本身，cls(...)等价于MCPResponse(...)，是工厂方法模式
    def success(cls, request_id: str, result: Any) -> "MCPResponse":
        return cls(request_id, result=result)  # cls()创建新实例，等价于MCPResponse()

    @classmethod  # @classmethod接收类作为第一个参数
    def error(cls, request_id: str, code: int, message: str, data: Any = None) -> "MCPResponse":
        error = {"code": code, "message": message}
        if data:
            error["data"] = data
        return cls(request_id, error=error)

class MCPNotification(MCPMessage):
    """MCP通知消息（无响应）"""

    def __init__(self, method: str, params: Dict = None):
        super().__init__()
        self.method = method
        self.params = params or {}

    def to_dict(self) -> Dict:
        return {
            "jsonrpc": self.jsonrpc,
            "method": self.method,
            "params": self.params
        }

# MCP协议方法定义
class MCPMethods:
    """MCP标准方法"""

    # 初始化
    INITIALIZE = "initialize"
    INITIALIZED = "initialized"

    # 工具
    TOOLS_LIST = "tools/list"
    TOOLS_CALL = "tools/call"

    # 资源
    RESOURCES_LIST = "resources/list"
    RESOURCES_READ = "resources/read"
    RESOURCES_SUBSCRIBE = "resources/subscribe"

    # 提示
    PROMPTS_LIST = "prompts/list"
    PROMPTS_GET = "prompts/get"

    # 能力
    COMPLETION_COMPLETE = "completion/complete"

    # 日志
    LOGGING_SET_LEVEL = "logging/setLevel"

    # 生命周期
    PING = "ping"

    # 通知
    NOTIFICATIONS_PROGRESS = "notifications/progress"
    NOTIFICATIONS_MESSAGE = "notifications/message"
    NOTIFICATIONS_RESOURCES_UPDATED = "notifications/resources/updated"

import asyncio  # Python标准异步库
from typing import AsyncIterator
import sys

class MCPTransport(ABC):  # ABC抽象基类，定义接口规范
    """MCP传输层抽象基类"""

    @abstractmethod  # @abstractmethod子类必须实现此方法
    async def read(self) -> AsyncIterator[str]:  # async def定义协程函数
        """读取消息流"""
        pass

    @abstractmethod
    async def write(self, message: str) -> None:
        """写入消息"""
        pass

    @abstractmethod
    async def close(self) -> None:
        """关闭传输"""
        pass

class StdioTransport(MCPTransport):
    """标准输入输出传输"""

    def __init__(self, stdin=None, stdout=None):
        self.stdin = stdin or sys.stdin
        self.stdout = stdout or sys.stdout
        self._closed = False

    async def read(self) -> AsyncIterator[str]:
        """从stdin读取JSON-RPC消息"""
        while not self._closed:
            try:  # try/except捕获异常，防止程序崩溃
                line = await asyncio.get_event_loop().run_in_executor(  # await等待异步操作完成
                    None, self.stdin.readline
                )
                if not line:
                    break
                line = line.strip()
                if line:
                    yield line  # yield产出值，函数变为生成器
            except Exception as e:
                print(f"Read error: {e}", file=sys.stderr)
                break

    async def write(self, message: str) -> None:
        """向stdout写入消息"""
        if not self._closed:
            self.stdout.write(message + "\n")
            self.stdout.flush()

    async def close(self) -> None:
        """关闭传输"""
        self._closed = True

class SSETransport(MCPTransport):
    """Server-Sent Events传输"""

    def __init__(self, endpoint_url: str):
        self.endpoint_url = endpoint_url
        self.session = None
        self._closed = False

    async def connect(self):
        """建立SSE连接"""
        import aiohttp
        self.session = aiohttp.ClientSession()
        self.response = await self.session.get(
            self.endpoint_url,
            headers={"Accept": "text/event-stream"}
        )

    async def read(self) -> AsyncIterator[str]:
        """读取SSE事件"""
        if not self.response:
            raise RuntimeError("Not connected")

        async for line in self.response.content:
            line = line.decode("utf-8").strip()
            if line.startswith("data: "):
                data = line[6:]
                if data:
                    yield data

    async def write(self, message: str) -> None:
        """通过POST发送消息"""
        if not self.session:
            raise RuntimeError("Not connected")

        async with self.session.post(  # async with异步上下文管理器
            self.endpoint_url,
            json=json.loads(message)  # json.loads将JSON字符串→Python对象
        ) as response:
            await response.text()

    async def close(self) -> None:
        """关闭连接"""
        self._closed = True
        if self.response:
            self.response.close()
        if self.session:
            await self.session.close()

class WebSocketTransport(MCPTransport):
    """WebSocket传输"""

    def __init__(self, ws_url: str):
        self.ws_url = ws_url
        self.websocket = None
        self._closed = False

    async def connect(self):
        """建立WebSocket连接"""
        import websockets
        self.websocket = await websockets.connect(self.ws_url)

    async def read(self) -> AsyncIterator[str]:
        """读取WebSocket消息"""
        if not self.websocket:
            raise RuntimeError("Not connected")

        async for message in self.websocket:
            if self._closed:
                break
            yield message

    async def write(self, message: str) -> None:
        """发送WebSocket消息"""
        if not self.websocket:
            raise RuntimeError("Not connected")

        await self.websocket.send(message)

    async def close(self) -> None:
        """关闭WebSocket"""
        self._closed = True
        if self.websocket:
            await self.websocket.close()

class MCPServer:
    """MCP服务器实现"""

    def __init__(
        self,
        name: str,
        version: str,
        capabilities: MCPCapability = None
    ):
        self.name = name
        self.version = version
        self.capabilities = capabilities or MCPCapability()

        # 注册的工具
        self._tools: Dict[str, MCPTool] = {}
        self._tool_handlers: Dict[str, Callable] = {}

        # 注册的资源
        self._resources: Dict[str, MCPResource] = {}
        self._resource_handlers: Dict[str, Callable] = {}

        # 提示模板
        self._prompts: Dict[str, Dict] = {}

        # 客户端信息
        self._client_info: Optional[Dict] = None
        self._client_capabilities: Optional[Dict] = None

        # 传输层
        self._transport: Optional[MCPTransport] = None

    def register_tool(
        self,
        name: str,
        description: str,
        input_schema: Dict,
        handler: Callable,
        output_schema: Dict = None,
        annotations: Dict = None
    ):
        """注册工具"""
        tool = MCPTool(
            name=name,
            description=description,
            input_schema=input_schema,
            output_schema=output_schema,
            annotations=annotations
        )
        self._tools[name] = tool
        self._tool_handlers[name] = handler
        self.capabilities.tools = True

    def register_resource(
        self,
        uri: str,
        name: str,
        description: str,
        handler: Callable,
        mime_type: str = None
    ):
        """注册资源"""
        resource = MCPResource(
            uri=uri,
            name=name,
            description=description,
            mime_type=mime_type
        )
        self._resources[uri] = resource
        self._resource_handlers[uri] = handler
        self.capabilities.resources = True

    def register_prompt(
        self,
        name: str,
        description: str,
        template: str,
        arguments: List[Dict] = None
    ):
        """注册提示模板"""
        self._prompts[name] = {
            "name": name,
            "description": description,
            "template": template,
            "arguments": arguments or []
        }
        self.capabilities.prompts = True

    async def handle_request(self, request: MCPRequest) -> MCPResponse:
        """处理请求"""

        method = request.method
        params = request.params

        try:
            if method == MCPMethods.INITIALIZE:
                return await self._handle_initialize(params, request.id)

            elif method == MCPMethods.TOOLS_LIST:
                return await self._handle_tools_list(request.id)

            elif method == MCPMethods.TOOLS_CALL:
                return await self._handle_tools_call(params, request.id)

            elif method == MCPMethods.RESOURCES_LIST:
                return await self._handle_resources_list(request.id)

            elif method == MCPMethods.RESOURCES_READ:
                return await self._handle_resources_read(params, request.id)

            elif method == MCPMethods.PROMPTS_LIST:
                return await self._handle_prompts_list(request.id)

            elif method == MCPMethods.PROMPTS_GET:
                return await self._handle_prompts_get(params, request.id)

            elif method == MCPMethods.PING:
                return MCPResponse.success(request.id, {})

            else:
                return MCPResponse.error(
                    request.id,
                    code=-32601,
                    message=f"Method not found: {method}"
                )

        except Exception as e:
            return MCPResponse.error(
                request.id,
                code=-32603,
                message=str(e)
            )

    async def _handle_initialize(
        self,
        params: Dict,
        request_id: str
    ) -> MCPResponse:
        """处理初始化请求"""

        self._client_info = params.get("clientInfo")
        self._client_capabilities = params.get("capabilities", {})
        protocol_version = params.get("protocolVersion", "2024-11-05")

        result = {
            "protocolVersion": protocol_version,
            "capabilities": self.capabilities.to_dict(),
            "serverInfo": {
                "name": self.name,
                "version": self.version
            }
        }

        return MCPResponse.success(request_id, result)

    async def _handle_tools_list(self, request_id: str) -> MCPResponse:
        """处理工具列表请求"""
        tools = [tool.to_dict() for tool in self._tools.values()]
        return MCPResponse.success(request_id, {"tools": tools})

    async def _handle_tools_call(
        self,
        params: Dict,
        request_id: str
    ) -> MCPResponse:
        """处理工具调用请求"""

        tool_name = params.get("name")
        arguments = params.get("arguments", {})

        if tool_name not in self._tools:
            return MCPResponse.error(
                request_id,
                code=-32602,
                message=f"Tool not found: {tool_name}"
            )

        handler = self._tool_handlers[tool_name]

        try:
            result = await handler(**arguments)  # **arguments将字典解包为关键字参数，如{"city":"北京"}变为handler(city="北京")
            return MCPResponse.success(request_id, {
                "content": [{"type": "text", "text": str(result)}]
            })
        except Exception as e:
            return MCPResponse.error(
                request_id,
                code=-32603,
                message=f"Tool execution failed: {str(e)}"
            )

    async def _handle_resources_list(self, request_id: str) -> MCPResponse:
        """处理资源列表请求"""
        resources = [res.to_dict() for res in self._resources.values()]
        return MCPResponse.success(request_id, {"resources": resources})

    async def _handle_resources_read(
        self,
        params: Dict,
        request_id: str
    ) -> MCPResponse:
        """处理资源读取请求"""

        uri = params.get("uri")

        if uri not in self._resources:
            return MCPResponse.error(
                request_id,
                code=-32602,
                message=f"Resource not found: {uri}"
            )

        handler = self._resource_handlers[uri]

        try:
            content = await handler()
            return MCPResponse.success(request_id, {
                "contents": [{
                    "uri": uri,
                    "mimeType": self._resources[uri].mime_type,
                    "text": content
                }]
            })
        except Exception as e:
            return MCPResponse.error(
                request_id,
                code=-32603,
                message=f"Resource read failed: {str(e)}"
            )

    async def _handle_prompts_list(self, request_id: str) -> MCPResponse:
        """处理提示列表请求"""
        prompts = list(self._prompts.values())
        return MCPResponse.success(request_id, {"prompts": prompts})

    async def _handle_prompts_get(
        self,
        params: Dict,
        request_id: str
    ) -> MCPResponse:
        """处理提示获取请求"""

        name = params.get("name")
        arguments = params.get("arguments", {})

        if name not in self._prompts:
            return MCPResponse.error(
                request_id,
                code=-32602,
                message=f"Prompt not found: {name}"
            )

        prompt_def = self._prompts[name]
        template = prompt_def["template"]

        # 填充模板
        try:
            filled = template.format(**arguments)
        except KeyError as e:
            return MCPResponse.error(
                request_id,
                code=-32602,
                message=f"Missing argument: {e}"
            )

        return MCPResponse.success(request_id, {
            "description": prompt_def["description"],
            "messages": [{
                "role": "user",
                "content": {"type": "text", "text": filled}
            }]
        })

    async def run(self, transport: MCPTransport):
        """运行服务器"""
        self._transport = transport

        async for message_str in transport.read():
            try:
                message = json.loads(message_str)

                # 处理请求
                if "method" in message and "id" in message:
                    request = MCPRequest(
                        method=message["method"],
                        params=message.get("params", {}),
                        request_id=message["id"]
                    )
                    response = await self.handle_request(request)
                    await transport.write(json.dumps(response.to_dict()))  # json.dumps将Python对象→JSON字符串

                # 处理通知
                elif "method" in message:
                    # 处理通知（无需响应）
                    pass

            except json.JSONDecodeError as e:
                error_response = MCPResponse.error(
                    None,
                    code=-32700,
                    message=f"Parse error: {str(e)}"
                )
                await transport.write(json.dumps(error_response.to_dict()))

class FileSystemMCPServer(MCPServer):
    """文件系统MCP服务器示例"""

    def __init__(self, root_path: str = "."):
        super().__init__(
            name="filesystem-server",
            version="1.0.0",
            capabilities=MCPCapability(tools=True, resources=True)
        )
        self.root_path = Path(root_path).resolve()

        self._setup_tools()
        self._setup_resources()

    def _setup_tools(self):
        """设置文件系统工具"""

        # 读取文件工具
        self.register_tool(
            name="read_file",
            description="读取文件内容",
            input_schema={
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "文件路径（相对root）"
                    }
                },
                "required": ["path"]
            },
            handler=self._read_file
        )

        # 写入文件工具
        self.register_tool(
            name="write_file",
            description="写入文件内容",
            input_schema={
                "type": "object",
                "properties": {
                    "path": {"type": "string"},
                    "content": {"type": "string"}
                },
                "required": ["path", "content"]
            },
            handler=self._write_file,
            annotations={
                "destructive": ["write"],
                "idempotent": False
            }
        )

        # 列出目录工具
        self.register_tool(
            name="list_directory",
            description="列出目录内容",
            input_schema={
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "目录路径（相对root）"
                    }
                },
                "required": ["path"]
            },
            handler=self._list_directory
        )

        # 搜索文件工具
        self.register_tool(
            name="search_files",
            description="搜索文件",
            input_schema={
                "type": "object",
                "properties": {
                    "pattern": {
                        "type": "string",
                        "description": "搜索模式（glob）"
                    }
                },
                "required": ["pattern"]
            },
            handler=self._search_files
        )

        # 获取文件信息工具
        self.register_tool(
            name="get_file_info",
            description="获取文件元数据",
            input_schema={
                "type": "object",
                "properties": {
                    "path": {"type": "string"}
                },
                "required": ["path"]
            },
            handler=self._get_file_info
        )

    def _setup_resources(self):
        """设置文件系统资源"""

        # 根目录资源
        self.register_resource(
            uri="file:///root",
            name="Root Directory",
            description="文件系统根目录",
            handler=self._read_root_resource,
            mime_type="application/json"
        )

    def _resolve_path(self, path: str) -> Path:
        """解析并验证路径"""
        resolved = (self.root_path / path).resolve()

        # 安全检查：确保路径在root_path内
        try:
            resolved.relative_to(self.root_path)
        except ValueError:
            raise PermissionError(f"Access denied: {path}")

        return resolved

    async def _read_file(self, path: str) -> str:
        """读取文件"""
        file_path = self._resolve_path(path)

        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {path}")

        if not file_path.is_file():
            raise ValueError(f"Not a file: {path}")

        return file_path.read_text(encoding="utf-8")

    async def _write_file(self, path: str, content: str) -> str:
        """写入文件"""
        file_path = self._resolve_path(path)

        # 确保父目录存在
        file_path.parent.mkdir(parents=True, exist_ok=True)

        file_path.write_text(content, encoding="utf-8")
        return f"Successfully wrote to {path}"

    async def _list_directory(self, path: str) -> str:
        """列出目录"""
        dir_path = self._resolve_path(path)

        if not dir_path.exists():
            raise FileNotFoundError(f"Directory not found: {path}")

        if not dir_path.is_dir():
            raise ValueError(f"Not a directory: {path}")

        entries = []
        for entry in dir_path.iterdir():
            entry_type = "directory" if entry.is_dir() else "file"
            entries.append({
                "name": entry.name,
                "type": entry_type,
                "size": entry.stat().st_size if entry.is_file() else None
            })

        return json.dumps(entries, indent=2)

    async def _search_files(self, pattern: str) -> str:
        """搜索文件"""
        matches = list(self.root_path.glob(pattern))

        results = []
        for match in matches:
            try:
                # 确保在root_path内
                match.relative_to(self.root_path)
                results.append(str(match.relative_to(self.root_path)))
            except ValueError:
                continue

        return json.dumps(results, indent=2)

    async def _get_file_info(self, path: str) -> str:
        """获取文件信息"""
        file_path = self._resolve_path(path)

        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {path}")

        stat = file_path.stat()

        info = {
            "name": file_path.name,
            "path": str(file_path.relative_to(self.root_path)),
            "type": "directory" if file_path.is_dir() else "file",
            "size": stat.st_size,
            "created": datetime.fromtimestamp(stat.st_ctime).isoformat(),
            "modified": datetime.fromtimestamp(stat.st_mtime).isoformat(),
            "permissions": oct(stat.st_mode)[-3:]
        }

        return json.dumps(info, indent=2)

    async def _read_root_resource(self) -> str:
        """读取根目录资源"""
        entries = []
        for entry in self.root_path.iterdir():
            entries.append({
                "name": entry.name,
                "type": "directory" if entry.is_dir() else "file"
            })

        return json.dumps({
            "root": str(self.root_path),
            "entries": entries
        }, indent=2)

# 运行服务器
async def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description="File System MCP Server")
    parser.add_argument("--root", default=".", help="Root directory path")
    args = parser.parse_args()

    server = FileSystemMCPServer(root_path=args.root)
    transport = StdioTransport()

    print(f"File System MCP Server starting...", file=sys.stderr)
    print(f"Root path: {server.root_path}", file=sys.stderr)

    await server.run(transport)

if __name__ == "__main__":
    asyncio.run(main())  # 创建事件循环运行顶层协程

class MCPClient:
    """MCP客户端实现"""

    def __init__(self, transport: MCPTransport):
        self.transport = transport
        self._pending_requests: Dict[str, asyncio.Future] = {}
        self._server_capabilities: Optional[Dict] = None
        self._server_info: Optional[Dict] = None
        self._initialized = False
        self._request_counter = 0

    async def initialize(
        self,
        client_name: str,
        client_version: str,
        capabilities: Dict = None
    ) -> Dict:
        """初始化连接"""

        request = MCPRequest(
            method=MCPMethods.INITIALIZE,
            params={
                "protocolVersion": "2024-11-05",
                "capabilities": capabilities or {},
                "clientInfo": {
                    "name": client_name,
                    "version": client_version
                }
            }
        )

        response = await self._send_request(request)

        if response.error:
            raise MCPError(f"Initialization failed: {response.error}")

        self._server_capabilities = response.result.get("capabilities", {})
        self._server_info = response.result.get("serverInfo", {})
        self._initialized = True

        # 发送initialized通知
        await self._send_notification(MCPMethods.INITIALIZED, {})

        return response.result

    async def list_tools(self) -> List[MCPTool]:
        """获取可用工具列表"""
        if not self._initialized:
            raise RuntimeError("Client not initialized")

        request = MCPRequest(method=MCPMethods.TOOLS_LIST)
        response = await self._send_request(request)

        if response.error:
            raise MCPError(f"Failed to list tools: {response.error}")

        tools_data = response.result.get("tools", [])
        return [
            MCPTool(
                name=t["name"],
                description=t["description"],
                input_schema=t["inputSchema"],
                output_schema=t.get("outputSchema"),
                annotations=t.get("annotations")
            )
            for t in tools_data
        ]

    async def call_tool(
        self,
        tool_name: str,
        arguments: Dict
    ) -> Dict:
        """调用工具"""
        if not self._initialized:
            raise RuntimeError("Client not initialized")

        request = MCPRequest(
            method=MCPMethods.TOOLS_CALL,
            params={
                "name": tool_name,
                "arguments": arguments
            }
        )

        response = await self._send_request(request)

        if response.error:
            raise MCPError(f"Tool call failed: {response.error}")

        return response.result

    async def list_resources(self) -> List[MCPResource]:
        """获取可用资源列表"""
        if not self._initialized:
            raise RuntimeError("Client not initialized")

        request = MCPRequest(method=MCPMethods.RESOURCES_LIST)
        response = await self._send_request(request)

        if response.error:
            raise MCPError(f"Failed to list resources: {response.error}")

        resources_data = response.result.get("resources", [])
        return [
            MCPResource(
                uri=r["uri"],
                name=r["name"],
                description=r["description"],
                mime_type=r.get("mimeType")
            )
            for r in resources_data
        ]

    async def read_resource(self, uri: str) -> Dict:
        """读取资源"""
        if not self._initialized:
            raise RuntimeError("Client not initialized")

        request = MCPRequest(
            method=MCPMethods.RESOURCES_READ,
            params={"uri": uri}
        )

        response = await self._send_request(request)

        if response.error:
            raise MCPError(f"Resource read failed: {response.error}")

        return response.result

    async def _send_request(self, request: MCPRequest) -> MCPResponse:
        """发送请求并等待响应"""

        future = asyncio.get_event_loop().create_future()
        self._pending_requests[request.id] = future

        try:
            await self.transport.write(json.dumps(request.to_dict()))

            # 等待响应（带超时）
            response_data = await asyncio.wait_for(future, timeout=30.0)

            return MCPResponse(
                request_id=response_data.get("id"),
                result=response_data.get("result"),
                error=response_data.get("error")
            )

        except asyncio.TimeoutError:
            raise MCPError(f"Request timeout: {request.method}")

        finally:
            self._pending_requests.pop(request.id, None)

    async def _send_notification(self, method: str, params: Dict):
        """发送通知"""
        notification = MCPNotification(method=method, params=params)
        await self.transport.write(json.dumps(notification.to_dict()))

    async def _handle_incoming(self):
        """处理传入消息"""
        async for message_str in self.transport.read():
            try:
                message = json.loads(message_str)

                # 处理响应
                if "id" in message and "result" in message:
                    request_id = message["id"]
                    if request_id in self._pending_requests:
                        future = self._pending_requests.pop(request_id)
                        if not future.done():
                            future.set_result(message)

                # 处理错误响应
                elif "id" in message and "error" in message:
                    request_id = message["id"]
                    if request_id in self._pending_requests:
                        future = self._pending_requests.pop(request_id)
                        if not future.done():
                            future.set_result(message)

                # 处理服务器通知
                elif "method" in message:
                    await self._handle_notification(message)

            except json.JSONDecodeError:
                print(f"Failed to parse message: {message_str}", file=sys.stderr)

    async def _handle_notification(self, notification: Dict):
        """处理服务器通知"""
        method = notification.get("method")
        params = notification.get("params", {})

        if method == MCPMethods.NOTIFICATIONS_PROGRESS:
            # 处理进度通知
            print(f"Progress: {params.get('message', '')}")

        elif method == MCPMethods.NOTIFICATIONS_MESSAGE:
            # 处理消息通知
            level = params.get("level", "info")
            message = params.get("message", "")
            print(f"[{level.upper()}] {message}")

    async def close(self):
        """关闭客户端"""
        await self.transport.close()

class MCPEnhancedLLM:
    """集成MCP的增强型LLM"""

    def __init__(self, llm_client, mcp_client: MCPClient):
        self.llm = llm_client
        self.mcp = mcp_client
        self._tools: List[MCPTool] = []

    async def initialize(self):
        """初始化并获取工具"""
        await self.mcp.initialize(
            client_name="mcp-enhanced-llm",
            client_version="1.0.0"
        )
        self._tools = await self.mcp.list_tools()

    async def chat(self, user_message: str) -> str:
        """对话并可能调用工具"""

        # 构建系统提示
        system_prompt = self._build_system_prompt()

        # 第一次LLM调用
        response = await self.llm.generate(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message}
            ],
            tools=self._format_tools_for_llm()
        )

        # 检查是否需要调用工具
        if response.tool_calls:
            # 执行工具调用
            tool_results = await self._execute_tool_calls(response.tool_calls)

            # 第二次LLM调用（带工具结果）
            final_response = await self.llm.generate(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_message},
                    {"role": "assistant", "content": response.content},
                    *self._format_tool_results(tool_results)
                ]
            )

            return final_response.content

        return response.content

    def _build_system_prompt(self) -> str:
        """构建系统提示"""
        tools_desc = "\n".join([
            f"- {tool.name}: {tool.description}"
            for tool in self._tools
        ])

        return f"""You are an AI assistant with access to external tools.

Available tools:
{tools_desc}

When you need to use a tool, respond with a tool call in the appropriate format.
After receiving tool results, incorporate them into your response."""

    def _format_tools_for_llm(self) -> List[Dict]:
        """将MCP工具格式化为LLM工具格式"""
        return [
            {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.input_schema
                }
            }
            for tool in self._tools
        ]

    async def _execute_tool_calls(self, tool_calls: List[Dict]) -> List[Dict]:
        """执行工具调用"""
        results = []

        for call in tool_calls:
            tool_name = call["function"]["name"]
            arguments = json.loads(call["function"]["arguments"])

            try:
                result = await self.mcp.call_tool(tool_name, arguments)
                results.append({
                    "tool_call_id": call["id"],
                    "role": "tool",
                    "name": tool_name,
                    "content": json.dumps(result)
                })
            except Exception as e:
                results.append({
                    "tool_call_id": call["id"],
                    "role": "tool",
                    "name": tool_name,
                    "content": f"Error: {str(e)}"
                })

        return results

    def _format_tool_results(self, results: List[Dict]) -> List[Dict]:
        """格式化工具结果为消息"""
        return [
            {
                "role": "tool",
                "tool_call_id": r["tool_call_id"],
                "content": r["content"]
            }
            for r in results
        ]