Trae-Agent中的MCP逻辑

MCP (Model Context Protocol) 逻辑全面解析

什么是MCP？

MCP (Model Context Protocol) 是一个标准化协议，允许AI模型与外部工具/服务进行交互。在Trae Agent中，MCP用于扩展Agent的能力，让它可以调用外部MCP服务器提供的工具。

---

MCP架构概览

┌─────────────────────────────────────────────────────────────┐
│                     Trae Agent                              │
│                                                             │
│  ┌─────────────┐    ┌─────────────┐    ┌─────────────┐    │
│  │ TraeAgent   │───▶│ MCPClient   │───▶│ MCPTool     │    │
│  │ (Agent层)   │    │ (MCP客户端)  │    │ (工具封装)   │    │
│  └─────────────┘    └─────────────┘    └─────────────┘    │
│         │                  │                  │             │
└─────────│──────────────────│──────────────────│─────────────┘
          │                  │                  │
          ▼                  ▼                  ▼
    ┌─────────────────────────────────────────────────────┐
    │                   MCP 生态                         │
    │  ┌──────────────┐  ┌──────────────┐                │
    │  │ MCP Server 1 │  │ MCP Server 2 │  ...           │
    │  │ (Filesystem) │  │ (GitHub)     │                │
    │  └──────────────┘  └──────────────┘                │
    └─────────────────────────────────────────────────────┘

---

MCP工作流程

1. 配置阶段

文件: trae_agent/utils/config.py (第120-140行)

@dataclass
class MCPServerConfig:
    # 标准IO传输 (最常用)
    command: str | None = None      # 命令，如 "npx"
    args: list[str] | None = None   # 参数，如 ["-y", "@modelcontextprotocol/server-filesystem"]
    env: dict[str, str] | None = None  # 环境变量
    cwd: str | None = None          # 工作目录

    # HTTP传输
    url: str | None = None
    
    # WebSocket传输
    http_url: str | None = None
    tcp: str | None = None
    
    # 通用配置
    timeout: int | None = None
    trust: bool | None = None
    description: str | None = None

YAML配置示例:

mcp_servers:
  filesystem:
    command: "npx"
    args: ["-y", "@modelcontextprotocol/server-filesystem", "/path/to/allowed_directory"]
    cwd: "/home/user"
    description: "File system access"
  
  github:
    command: "uvx"
    args: ["mcp-server-github"]

allow_mcp_servers:
  - filesystem
  - github

---

2. 初始化阶段

文件: trae_agent/agent/trae_agent.py (第65-88行)

async def initialise_mcp(self):
    """初始化MCP工具"""
    await self.discover_mcp_tools()
    
    # 将MCP工具添加到Agent的工具列表
    if self.mcp_tools:
        self._tools.extend(self.mcp_tools)

async def discover_mcp_tools(self):
    """发现并连接MCP服务器"""
    if self.mcp_servers_config:
        for mcp_server_name, mcp_server_config in self.mcp_servers_config.items():
            # 检查是否允许使用该MCP服务器
            if self.allow_mcp_servers is None:
                continue
            if mcp_server_name not in self.allow_mcp_servers:
                continue
            
            # 创建MCP客户端并连接
            mcp_client = MCPClient()
            await mcp_client.connect_and_discover(
                mcp_server_name,
                mcp_server_config,
                self.mcp_tools,  # 接收发现的工具
                self._model_config.model_provider.provider,  # 模型提供商
            )
            
            # 保存客户端引用，用于清理
            self.mcp_clients.append(mcp_client)

---

3. 连接阶段

文件: trae_agent/utils/mcp_client.py (第40-90行)

async def connect_and_discover(self, mcp_server_name, mcp_server_config, mcp_tools_container, model_provider):
    """连接MCP服务器并发现工具"""
    
    # 1. 创建传输参数 (StdioServerParameters)
    params = StdioServerParameters(
        command=mcp_server_config.command,
        args=mcp_server_config.args,
        env=mcp_server_config.env,
        cwd=mcp_server_config.cwd,
    )
    
    # 2. 创建标准IO客户端
    transport = await self.exit_stack.enter_async_context(stdio_client(params))
    
    # 3. 连接到服务器
    await self.connect_to_server(mcp_server_name, transport)
    
    # 4. 发现可用工具
    mcp_tools = await self.list_tools()
    
    # 5. 为每个工具创建MCPTool封装
    for tool in mcp_tools.tools:
        mcp_tool = MCPTool(self, tool, model_provider)
        mcp_tools_container.append(mcp_tool)

async def connect_to_server(self, mcp_server_name, transport):
    """连接到MCP服务器"""
    if self.get_mcp_server_status(mcp_server_name) != MCPServerStatus.CONNECTED:
        self.update_mcp_server_status(mcp_server_name, MCPServerStatus.CONNECTING)
        
        stdio, write = transport
        self.session = await self.exit_stack.enter_async_context(
            ClientSession(stdio, write)
        )
        await self.session.initialize()
        
        self.update_mcp_server_status(mcp_server_name, MCPServerStatus.CONNECTED)

---

4. 工具调用阶段

文件: trae_agent/tools/mcp_tool.py (第30-58行)

class MCPTool(Tool):
    def __init__(self, client, tool: mcp.types.Tool, model_provider: str | None = None):
        super().__init__(model_provider)
        self.client = client  # MCPClient实例
        self.tool = tool      # MCP工具定义

    @override
    def get_name(self) -> str:
        return self.tool.name  # 工具名称

    @override
    def get_description(self) -> str:
        return self.tool.description  # 工具描述

    @override
    def get_parameters(self) -> list[ToolParameter]:
        # 将MCP工具的输入schema转换为Trae Agent的参数格式
        inputSchema = self.tool.inputSchema
        required = inputSchema.get("required", [])
        properties = inputSchema.get("properties", {})
        
        for name, prop in properties.items():
            tool_para = ToolParameter(
                name=name,
                type=prop["type"],
                items=prop.get("items", None),
                description=prop["description"],
                required=name in required,
            )
            parameters.append(tool_para)
        return parameters

    @override
    async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
        try:
            # 调用MCP服务器的工具
            output = await self.client.call_tool(self.get_name(), arguments)
            
            if output.isError:
                return ToolExecResult(output=None, error=output.content[0].text)
            else:
                return ToolExecResult(output=output.content[0].text)
        except Exception as e:
            return ToolExecResult(error=f"Error running mcp tool: {e}", error_code=-1)

---

5. 清理阶段

文件: trae_agent/agent/base_agent.py (第195-200行)

async def execute_task(self) -> AgentExecution:
    try:
        # 执行任务
        ...
    finally:
        # 确保清理MCP客户端，防止异步上下文泄漏
        with contextlib.suppress(Exception):
            await self.cleanup_mcp_clients()

async def cleanup_mcp_clients(self) -> None:
    """清理所有MCP客户端"""
    for client in self.mcp_clients:
        await client.cleanup(mcp_server_name)
    self.mcp_clients.clear()

---

完整数据流

1. 配置文件 (trae_config.yaml)
   │
   ├── mcp_servers:
   │     ├── filesystem: {command: npx, args: [...], cwd: ...}
   │     └── github: {command: uvx, args: [...]}
   │
   └── allow_mcp_servers: [filesystem, github]
          │
          ▼
2. Config加载 (config.py)
   │
   ├── MCPServerConfig(command, args, env, cwd, ...)
   │
   ▼
3. TraeAgent初始化 (trae_agent.py)
   │
   ├── for mcp_server_name, mcp_server_config in mcp_servers_config:
   │
   ▼
4. MCPClient连接 (mcp_client.py)
   │
   ├── StdioServerParameters(command, args, env, cwd)
   │──▶ stdio_client() ──▶ ClientSession
   │
   ▼
5. 发现工具 (mcp_client.py)
   │
   ├── session.list_tools() ──▶ 获取工具列表
   │
   ▼
6. MCPTool封装 (mcp_tool.py)
   │
   ├── for tool in mcp_tools:
   │     MCPTool(client, tool)
   │
   ▼
7. Agent执行 (base_agent.py)
   │
   ├── Agent调用LLM
   ├── LLM返回工具调用
   ├── Agent执行MCPTool.execute()
   ├── MCPTool调用 self.client.call_tool(name, args)
   ├── MCPClient调用 session.call_tool(name, args)
   ├── MCP服务器执行并返回结果
   │
   ▼
8. 清理 (base_agent.py)
   │
   ├── for client in mcp_clients:
   │     await client.cleanup()
   │
   ▼
9. 任务完成

---

MCP服务器状态管理

class MCPServerStatus(Enum):
    DISCONNECTED = "disconnected"   # 未连接或出错
    CONNECTING = "connecting"       # 正在连接
    CONNECTED = "connected"         # 已连接并可用

---

核心要点

阶段	关键类	职责
配置	MCPServerConfig	定义MCP服务器连接参数
连接	MCPClient	管理与MCP服务器的连接
工具	MCPTool	将MCP工具封装为Trae Agent工具
执行	ClientSession	通过MCP协议调用远程工具
清理	cleanup()	释放资源，防止泄漏

---

MCP优势

1. 标准化协议 - 统一的工具调用方式
2. 动态发现 - 自动发现服务器提供的工具
3. 解耦 - Agent不需要知道工具的具体实现
4. 可扩展 - 轻松添加新的MCP服务器
5. 隔离 - MCP工具在独立进程中运行

---

关键文件索引

文件	职责	重要性
utils/config.py	MCP配置定义	⭐⭐⭐
utils/mcp_client.py	MCP客户端连接管理	⭐⭐⭐
tools/mcp_tool.py	MCP工具封装	⭐⭐⭐
agent/trae_agent.py	MCP初始化和发现	⭐⭐⭐
agent/base_agent.py	MCP清理逻辑	⭐⭐

---

使用示例

添加Filesystem MCP服务器

# trae_config.yaml
agents:
  trae_agent:
    mcp_servers:
      filesystem:
        command: "npx"
        args: ["-y", "@modelcontextprotocol/server-filesystem", "/Users/xiaoniuchongchong/projects"]
        description: "访问项目文件"
    
    allow_mcp_servers:
      - filesystem

添加GitHub MCP服务器

agents:
  trae_agent:
    mcp_servers:
      github:
        command: "uvx"
        args: ["mcp-server-github", "--github-token", "ghp_xxx"]
        description: "GitHub API访问"
    
    allow_mcp_servers:
      - github

---

最后更新: 2025-03-12