Files
taiji-AI-PAD/services/mcp-server/mcp_protocol.py
T
xiaohei ff14abe1de feat: 实现MCP Server函数工具调用和沙箱安全机制
✅ 主要功能:
- 实现函数注册表 (function_registry.py)
  - 16个内置安全函数(数学、字符串、日期、JSON、哈希等)
  - 函数元数据管理
  - 函数查询和列表

- 实现沙箱执行器 (sandbox_executor.py)
  - 超时控制 (默认5秒)
  - 参数验证和资源限制
  - 安全的执行环境

- 实现 _execute_function_tool 方法
  - 函数白名单验证
  - 参数类型验证
  - 沙箱执行
  - 完善的错误处理

🔒 安全特性:
- 只允许执行注册表中的函数
- 禁止网络和文件系统访问
- 超时和资源限制
- 参数验证和类型检查

📚 文档:
- 添加实现方案文档
- 添加实现完成文档

完成度: 95% (核心功能完成,单元测试待完成)
版本: v1.2.0
2025-12-22 05:30:55 +00:00

689 lines
25 KiB
Python

"""
MCP (Model Context Protocol) 协议处理器
实现MCP协议的核心功能,包括工具管理、资源管理和代理通信
"""
import json
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional, Union
import asyncio
import logging
import traceback
import redis.asyncio as redis
import nats
import httpx
from schemas import (
MCPRequest, MCPResponse, MCPError,
ToolDefinition, ToolResult, ExecutionResult
)
from function_registry import get_function_registry
from sandbox_executor import get_sandbox_executor
logger = logging.getLogger(__name__)
class MCPProtocolHandler:
"""MCP协议处理器"""
def __init__(
self,
redis_client: redis.Redis,
nats_client: nats.NATS,
litellm_url: str = "http://litellm-gateway:4000"
):
self.redis = redis_client
self.nats = nats_client
self.litellm_url = litellm_url
self.http_client = httpx.AsyncClient(timeout=30.0)
# MCP协议版本
self.protocol_version = "2024-11-05"
# 函数注册表和沙箱执行器
self.function_registry = get_function_registry()
self.sandbox_executor = get_sandbox_executor(
timeout=5.0,
max_memory_mb=100,
allow_network=False,
allow_filesystem=False
)
# 支持的MCP方法
self.supported_methods = {
"initialize",
"tools/list",
"tools/call",
"resources/list",
"resources/read",
"prompts/list",
"prompts/get",
"completion/complete",
"logging/setLevel"
}
# 工具注册表
self._tools_registry: Dict[str, ToolDefinition] = {}
# 资源注册表
self._resources_registry: Dict[str, Dict[str, Any]] = {}
# 会话管理
self._sessions: Dict[str, Dict[str, Any]] = {}
async def execute_request(self, agent_id: str, request: MCPRequest) -> ExecutionResult:
"""执行MCP请求"""
execution_id = str(uuid.uuid4())
started_at = datetime.utcnow()
logger.info(f"开始执行MCP请求: {execution_id}, method: {request.method}")
try:
# 验证方法是否支持
if request.method not in self.supported_methods:
raise ValueError(f"不支持的MCP方法: {request.method}")
# 发布执行开始事件
await self._publish_execution_event(
"execution.started",
{
"execution_id": execution_id,
"agent_id": agent_id,
"method": request.method,
"timestamp": started_at.isoformat()
}
)
# 执行具体方法
result = await self._dispatch_method(agent_id, request)
completed_at = datetime.utcnow()
execution_time = (completed_at - started_at).total_seconds() * 1000
# 发布执行完成事件
await self._publish_execution_event(
"execution.completed",
{
"execution_id": execution_id,
"agent_id": agent_id,
"method": request.method,
"execution_time": execution_time,
"success": True,
"timestamp": completed_at.isoformat()
}
)
return ExecutionResult(
execution_id=execution_id,
success=True,
result=result,
execution_time=execution_time,
started_at=started_at,
completed_at=completed_at
)
except Exception as e:
completed_at = datetime.utcnow()
execution_time = (completed_at - started_at).total_seconds() * 1000
error_msg = str(e)
logger.error(f"MCP请求执行失败: {execution_id}, error: {error_msg}")
logger.error(traceback.format_exc())
# 发布执行失败事件
await self._publish_execution_event(
"execution.failed",
{
"execution_id": execution_id,
"agent_id": agent_id,
"method": request.method,
"execution_time": execution_time,
"error": error_msg,
"timestamp": completed_at.isoformat()
}
)
return ExecutionResult(
execution_id=execution_id,
success=False,
error=error_msg,
execution_time=execution_time,
started_at=started_at,
completed_at=completed_at
)
async def _dispatch_method(self, agent_id: str, request: MCPRequest) -> Any:
"""分发MCP方法调用"""
method = request.method
params = request.params or {}
if method == "initialize":
return await self._handle_initialize(params)
elif method == "tools/list":
return await self._handle_tools_list(agent_id, params)
elif method == "tools/call":
return await self._handle_tools_call(agent_id, params)
elif method == "resources/list":
return await self._handle_resources_list(agent_id, params)
elif method == "resources/read":
return await self._handle_resources_read(agent_id, params)
elif method == "prompts/list":
return await self._handle_prompts_list(agent_id, params)
elif method == "prompts/get":
return await self._handle_prompts_get(agent_id, params)
elif method == "completion/complete":
return await self._handle_completion_complete(agent_id, params)
elif method == "logging/setLevel":
return await self._handle_logging_set_level(params)
else:
raise ValueError(f"未实现的方法: {method}")
async def _handle_initialize(self, params: Dict[str, Any]) -> Dict[str, Any]:
"""处理初始化请求"""
client_info = params.get("clientInfo", {})
protocol_version = params.get("protocolVersion")
logger.info(f"MCP客户端初始化: {client_info}")
return {
"protocolVersion": self.protocol_version,
"capabilities": {
"tools": {
"listChanged": True
},
"resources": {
"subscribe": True,
"listChanged": True
},
"prompts": {
"listChanged": True
},
"completion": {
"argument": True
},
"logging": {}
},
"serverInfo": {
"name": "taiji-AI-PAD MCP Server",
"version": "1.0.0"
}
}
async def _handle_tools_list(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
"""处理工具列表请求"""
try:
# 从Redis获取Agent的工具列表
agent_tools_key = f"agent:{agent_id}:tools"
tool_names = await self.redis.smembers(agent_tools_key)
tools = []
for tool_name in tool_names:
tool_info = await self._get_tool_info(tool_name)
if tool_info:
tools.append({
"name": tool_info["name"],
"description": tool_info["description"],
"inputSchema": tool_info.get("schema", {})
})
return {"tools": tools}
except Exception as e:
logger.error(f"获取工具列表失败: {e}")
return {"tools": []}
async def _handle_tools_call(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
"""处理工具调用请求"""
tool_name = params.get("name")
arguments = params.get("arguments", {})
if not tool_name:
raise ValueError("工具名称不能为空")
logger.info(f"调用工具: {tool_name}, arguments: {arguments}")
try:
# 验证Agent是否有权限使用该工具
agent_tools_key = f"agent:{agent_id}:tools"
if not await self.redis.sismember(agent_tools_key, tool_name):
raise ValueError(f"Agent {agent_id} 无权限使用工具 {tool_name}")
# 获取工具信息
tool_info = await self._get_tool_info(tool_name)
if not tool_info:
raise ValueError(f"工具 {tool_name} 不存在")
# 执行工具调用
result = await self._execute_tool(tool_name, tool_info, arguments)
return {
"content": [
{
"type": "text",
"text": json.dumps(result.result) if result.success else f"错误: {result.error}"
}
],
"isError": not result.success
}
except Exception as e:
logger.error(f"工具调用失败: {e}")
return {
"content": [
{
"type": "text",
"text": f"工具调用失败: {str(e)}"
}
],
"isError": True
}
async def _handle_resources_list(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
"""处理资源列表请求"""
try:
# 从Redis获取Agent的资源列表
agent_resources_key = f"agent:{agent_id}:resources"
resource_names = await self.redis.smembers(agent_resources_key)
resources = []
for resource_name in resource_names:
resource_info = await self._get_resource_info(resource_name)
if resource_info:
resources.append({
"uri": resource_info["uri"],
"name": resource_info["name"],
"description": resource_info.get("description"),
"mimeType": resource_info.get("mimeType", "application/json")
})
return {"resources": resources}
except Exception as e:
logger.error(f"获取资源列表失败: {e}")
return {"resources": []}
async def _handle_resources_read(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
"""处理资源读取请求"""
uri = params.get("uri")
if not uri:
raise ValueError("资源URI不能为空")
try:
# 验证权限
agent_resources_key = f"agent:{agent_id}:resources"
# 这里应该根据URI找到资源名称
resource_name = uri.split("/")[-1] # 简化处理
if not await self.redis.sismember(agent_resources_key, resource_name):
raise ValueError(f"Agent {agent_id} 无权限访问资源 {uri}")
# 读取资源内容
content = await self._read_resource_content(uri)
return {
"contents": [
{
"uri": uri,
"mimeType": "application/json",
"text": json.dumps(content)
}
]
}
except Exception as e:
logger.error(f"读取资源失败: {e}")
raise ValueError(f"无法读取资源 {uri}: {str(e)}")
async def _handle_prompts_list(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
"""处理提示词列表请求"""
# 获取Agent相关的提示词
prompts = [
{
"name": "system_prompt",
"description": "系统提示词",
"arguments": [
{
"name": "context",
"description": "上下文信息",
"required": False
}
]
}
]
return {"prompts": prompts}
async def _handle_prompts_get(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
"""处理获取提示词请求"""
name = params.get("name")
arguments = params.get("arguments", {})
if name == "system_prompt":
# 构建系统提示词
agent_info = await self._get_agent_info(agent_id)
prompt = f"""你是 {agent_info.get('name', 'AI助手')}。
角色定义: {agent_info.get('role', '通用助手')}
目标: {agent_info.get('goal', '帮助用户完成任务')}
可用工具: {', '.join(agent_info.get('tools', []))}
请根据用户的请求,选择合适的工具来完成任务。"""
return {
"description": "Agent系统提示词",
"messages": [
{
"role": "system",
"content": {
"type": "text",
"text": prompt
}
}
]
}
raise ValueError(f"未知的提示词: {name}")
async def _handle_completion_complete(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
"""处理补全请求"""
ref = params.get("ref", {})
argument = params.get("argument", {})
# 根据参考信息生成补全建议
completions = []
if ref.get("type") == "tool":
tool_name = ref.get("name")
if tool_name:
tool_info = await self._get_tool_info(tool_name)
if tool_info and "schema" in tool_info:
# 基于工具schema生成参数建议
schema = tool_info["schema"]
properties = schema.get("properties", {})
for prop_name, prop_info in properties.items():
completions.append({
"type": "text",
"text": prop_name,
"insertText": f'"{prop_name}": ""'
})
return {"completion": {"values": completions}}
async def _handle_logging_set_level(self, params: Dict[str, Any]) -> Dict[str, Any]:
"""处理设置日志级别请求"""
level = params.get("level", "info")
# 设置日志级别
numeric_level = getattr(logging, level.upper(), logging.INFO)
logging.getLogger().setLevel(numeric_level)
logger.info(f"日志级别已设置为: {level}")
return {"success": True}
async def _execute_tool(
self,
tool_name: str,
tool_info: Dict[str, Any],
arguments: Dict[str, Any]
) -> ToolResult:
"""执行工具调用"""
start_time = datetime.utcnow()
try:
# 根据工具类型执行不同的逻辑
category = tool_info.get("category", "api")
if category == "api":
result = await self._execute_api_tool(tool_info, arguments)
elif category == "function":
result = await self._execute_function_tool(tool_info, arguments)
elif category == "llm":
result = await self._execute_llm_tool(tool_info, arguments)
else:
raise ValueError(f"不支持的工具类型: {category}")
execution_time = (datetime.utcnow() - start_time).total_seconds() * 1000
return ToolResult(
success=True,
result=result,
execution_time=execution_time,
cost=tool_info.get("cost_per_call", 0.0)
)
except Exception as e:
execution_time = (datetime.utcnow() - start_time).total_seconds() * 1000
return ToolResult(
success=False,
error=str(e),
execution_time=execution_time,
cost=tool_info.get("cost_per_call", 0.0)
)
async def _execute_api_tool(self, tool_info: Dict[str, Any], arguments: Dict[str, Any]) -> Any:
"""执行API工具调用"""
endpoint = tool_info.get("endpoint")
method = tool_info.get("method", "POST")
headers = tool_info.get("headers", {})
timeout = tool_info.get("timeout", 30)
if not endpoint:
raise ValueError("API端点不能为空")
# 发送HTTP请求
response = await self.http_client.request(
method=method,
url=endpoint,
json=arguments,
headers=headers,
timeout=timeout
)
response.raise_for_status()
return response.json()
async def _execute_function_tool(self, tool_info: Dict[str, Any], arguments: Dict[str, Any]) -> Any:
"""
执行函数工具调用
安全机制:
1. 只允许执行注册表中的函数
2. 在沙箱环境中执行
3. 参数验证和资源限制
4. 超时控制
"""
function_name = tool_info.get("name") or tool_info.get("function_name")
if not function_name:
raise ValueError("函数名称不能为空")
# 检查函数是否在注册表中
function_info = self.function_registry.get(function_name)
if not function_info:
raise ValueError(
f"函数 {function_name} 未在注册表中。"
f"可用函数: {', '.join(self.function_registry.list_all()[:10])}"
)
# 验证参数
self._validate_function_arguments(function_info, arguments)
# 在沙箱中执行函数
try:
result = await self.sandbox_executor.execute(
func=function_info["func"],
arguments=arguments,
function_name=function_name
)
logger.info(f"函数 {function_name} 执行成功")
return result
except TimeoutError as e:
logger.error(f"函数 {function_name} 执行超时: {e}")
raise
except Exception as e:
logger.error(f"函数 {function_name} 执行失败: {e}")
raise RuntimeError(f"函数执行失败: {str(e)}")
def _validate_function_arguments(
self,
function_info: Dict[str, Any],
arguments: Dict[str, Any]
):
"""验证函数参数"""
expected_params = {
param["name"]: param
for param in function_info.get("parameters", [])
}
# 检查必需参数
for param_name, param_info in expected_params.items():
if param_info.get("required", True):
if param_name not in arguments:
raise ValueError(f"缺少必需参数: {param_name}")
# 检查参数类型(基本验证)
for param_name, param_value in arguments.items():
if param_name not in expected_params:
logger.warning(f"未知参数: {param_name},将被忽略")
continue
param_info = expected_params[param_name]
expected_type = param_info.get("type", "any")
# 基本类型检查
if expected_type == "number" and not isinstance(param_value, (int, float)):
try:
float(param_value) # 尝试转换
except (ValueError, TypeError):
raise ValueError(
f"参数 {param_name} 应该是数字类型,但得到 {type(param_value).__name__}"
)
elif expected_type == "string" and not isinstance(param_value, str):
raise ValueError(
f"参数 {param_name} 应该是字符串类型,但得到 {type(param_value).__name__}"
)
elif expected_type == "object" and not isinstance(param_value, (dict, list)):
raise ValueError(
f"参数 {param_name} 应该是对象类型,但得到 {type(param_value).__name__}"
)
async def _execute_llm_tool(self, tool_info: Dict[str, Any], arguments: Dict[str, Any]) -> Any:
"""执行LLM工具调用"""
# 调用LiteLLM网关
payload = {
"model": arguments.get("model", "gpt-3.5-turbo"),
"messages": arguments.get("messages", []),
"temperature": arguments.get("temperature", 0.7),
"max_tokens": arguments.get("max_tokens", 150)
}
response = await self.http_client.post(
f"{self.litellm_url}/chat/completions",
json=payload,
headers={"Authorization": "Bearer sk-taiji-master-key"}
)
response.raise_for_status()
return response.json()
async def _get_tool_info(self, tool_name: str) -> Optional[Dict[str, Any]]:
"""获取工具信息"""
tool_key = f"tool:{tool_name}"
tool_data = await self.redis.get(tool_key)
if tool_data:
return json.loads(tool_data)
return None
async def _get_resource_info(self, resource_name: str) -> Optional[Dict[str, Any]]:
"""获取资源信息"""
resource_key = f"resource:{resource_name}"
resource_data = await self.redis.get(resource_key)
if resource_data:
return json.loads(resource_data)
return None
async def _get_agent_info(self, agent_id: str) -> Dict[str, Any]:
"""获取Agent信息"""
agent_key = f"agent:{agent_id}"
agent_data = await self.redis.get(agent_key)
if agent_data:
return json.loads(agent_data)
return {}
async def _read_resource_content(self, uri: str) -> Any:
"""读取资源内容"""
# 这里可以根据URI类型读取不同的资源
# 例如:文件、数据库、API等
if uri.startswith("file://"):
# 读取文件
file_path = uri[7:] # 移除file://前缀
with open(file_path, 'r') as f:
return f.read()
elif uri.startswith("http://") or uri.startswith("https://"):
# 读取HTTP资源
response = await self.http_client.get(uri)
response.raise_for_status()
return response.json()
else:
raise ValueError(f"不支持的资源类型: {uri}")
async def _publish_execution_event(self, event_type: str, data: Dict[str, Any]):
"""发布执行事件"""
try:
if self.nats:
await self.nats.publish(
f"mcp.{event_type}",
json.dumps(data).encode()
)
except Exception as e:
logger.error(f"发布事件失败: {e}")
async def register_tool(self, tool_definition: ToolDefinition) -> bool:
"""注册工具"""
try:
tool_key = f"tool:{tool_definition.name}"
tool_data = tool_definition.dict()
await self.redis.setex(
tool_key,
3600, # 1小时过期
json.dumps(tool_data)
)
logger.info(f"工具注册成功: {tool_definition.name}")
return True
except Exception as e:
logger.error(f"工具注册失败: {e}")
return False
async def register_agent_tool(self, agent_id: str, tool_name: str) -> bool:
"""为Agent注册工具"""
try:
agent_tools_key = f"agent:{agent_id}:tools"
await self.redis.sadd(agent_tools_key, tool_name)
await self.redis.expire(agent_tools_key, 3600)
logger.info(f"Agent {agent_id} 工具注册成功: {tool_name}")
return True
except Exception as e:
logger.error(f"Agent工具注册失败: {e}")
return False
async def close(self):
"""清理资源"""
try:
await self.http_client.aclose()
except Exception as e:
logger.error(f"资源清理失败: {e}")