Coverage for agentos/server/mcp_server.py: 37%

1""" 

2AgentOS v0.40 MCP Server — 将AgentOS暴露为MCP Server。 

3支持工具列表、资源、提示模板的MCP协议暴露。 

4""" 

5 

6from __future__ import annotations 

7 

8import json 

9import logging 

10from dataclasses import dataclass, field 

11from typing import Optional, Callable, Any 

12 

13 

14logger = logging.getLogger(__name__) 

15 

16 

17@dataclass 

18class MCPServerConfig: 

19 """MCP 服务端配置。""" 

20 name: str = "AgentOS-MCP-Server" 

21 version: str = "0.40.0" 

22 transport: str = "stdio" # stdio | sse | streamable-http 

23 host: str = "0.0.0.0" 

24 port: int = 9000 

25 

26 

27@dataclass 

28class MCPTool: 

29 """MCP工具定义。""" 

30 name: str 

31 description: str 

32 input_schema: dict 

33 handler: Callable 

34 annotations: dict = field(default_factory=dict) 

35 

36 

37@dataclass 

38class MCPResource: 

39 """MCP资源定义。""" 

40 uri: str 

41 name: str 

42 description: str = "" 

43 mime_type: str = "text/plain" 

44 handler: Callable | None = None 

45 

46 

47@dataclass 

48class MCPPrompt: 

49 """MCP提示模板。""" 

50 name: str 

51 description: str = "" 

52 arguments: list[dict] = field(default_factory=list) 

53 template: str = "" 

54 

55 

56class MCPServer: 

57 """MCP Server核心 — 将AgentOS能力以MCP协议暴露。""" 

58 

59 JSONRPC_VERSION = "2.0" 

60 

61 def __init__(self, config: MCPServerConfig | None = None): 

62 self.config = config or MCPServerConfig() 

63 self._tools: dict[str, MCPTool] = {} 

64 self._resources: dict[str, MCPResource] = {} 

65 self._prompts: dict[str, MCPPrompt] = {} 

66 self._initialized = False 

67 self._session_id: str | None = None 

68 

69 # ── 注册 ────────────────────────────────────── 

70 

71 def register_tool(self, tool: MCPTool): 

72 self._tools[tool.name] = tool 

73 logger.info(f"MCP tool registered: {tool.name}") 

74 

75 def register_resource(self, resource: MCPResource): 

76 self._resources[resource.uri] = resource 

77 

78 def register_prompt(self, prompt: MCPPrompt): 

79 self._prompts[prompt.name] = prompt 

80 

81 # ── 协议处理 ────────────────────────────────── 

82 

83 def handle_request(self, raw: dict) -> dict: 

84 """处理MCP JSON-RPC请求。""" 

85 method = raw.get("method", "") 

86 params = raw.get("params", {}) 

87 req_id = raw.get("id") 

88 

89 try: 

90 if method == "initialize": 

91 result = self._handle_initialize(params) 

92 elif method == "notifications/initialized": 

93 self._initialized = True 

94 return {} 

95 elif method == "tools/list": 

96 result = self._handle_tools_list() 

97 elif method == "tools/call": 

98 result = self._handle_tool_call(params) 

99 elif method == "resources/list": 

100 result = self._handle_resources_list() 

101 elif method == "resources/read": 

102 result = self._handle_resource_read(params) 

103 elif method == "prompts/list": 

104 result = self._handle_prompts_list() 

105 elif method == "prompts/get": 

106 result = self._handle_prompt_get(params) 

107 else: 

108 return self._error(req_id, -32601, f"Method not found: {method}") 

109 

110 return self._success(req_id, result) 

111 except Exception as e: 

112 logger.exception(f"MCP handler error: {e}") 

113 return self._error(req_id, -32603, str(e)) 

114 

115 def _success(self, req_id, result) -> dict: 

116 if req_id is None: 

117 return {} 

118 return {"jsonrpc": self.JSONRPC_VERSION, "id": req_id, "result": result} 

119 

120 def _error(self, req_id, code: int, message: str) -> dict: 

121 if req_id is None: 

122 return {} 

123 return {"jsonrpc": self.JSONRPC_VERSION, "id": req_id, "error": {"code": code, "message": message}} 

124 

125 # ── 方法实现 ────────────────────────────────── 

126 

127 def _handle_initialize(self, params: dict) -> dict: 

128 client_info = params.get("clientInfo", {}) 

129 self._session_id = params.get("sessionId") 

130 return { 

131 "protocolVersion": "2024-11-05", 

132 "capabilities": { 

133 "tools": {"listChanged": True}, 

134 "resources": {"subscribe": False, "listChanged": False}, 

135 "prompts": {"listChanged": False}, 

136 }, 

137 "serverInfo": {"name": self.config.name, "version": self.config.version}, 

138 } 

139 

140 def _handle_tools_list(self) -> dict: 

141 tools = [] 

142 for t in self._tools.values(): 

143 tools.append({ 

144 "name": t.name, 

145 "description": t.description, 

146 "inputSchema": t.input_schema, 

147 "annotations": t.annotations, 

148 }) 

149 return {"tools": tools} 

150 

151 def _handle_tool_call(self, params: dict) -> dict: 

152 name = params.get("name", "") 

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

154 tool = self._tools.get(name) 

155 if not tool: 

156 raise ValueError(f"Tool not found: {name}") 

157 

158 result = tool.handler(arguments) 

159 content = [] 

160 if isinstance(result, dict): 

161 if "content" in result: 

162 content = result["content"] 

163 else: 

164 content = [{"type": "text", "text": json.dumps(result, ensure_ascii=False)}] 

165 elif isinstance(result, str): 

166 content = [{"type": "text", "text": result}] 

167 else: 

168 content = [{"type": "text", "text": str(result)}] 

169 

170 return {"content": content} 

171 

172 def _handle_resources_list(self) -> dict: 

173 resources = [] 

174 for r in self._resources.values(): 

175 resources.append({"uri": r.uri, "name": r.name, "description": r.description, "mimeType": r.mime_type}) 

176 return {"resources": resources} 

177 

178 def _handle_resource_read(self, params: dict) -> dict: 

179 uri = params.get("uri", "") 

180 resource = self._resources.get(uri) 

181 if not resource: 

182 raise ValueError(f"Resource not found: {uri}") 

183 text = resource.handler() if resource.handler else "" 

184 return {"contents": [{"uri": uri, "mimeType": resource.mime_type, "text": text}]} 

185 

186 def _handle_prompts_list(self) -> dict: 

187 prompts = [] 

188 for p in self._prompts.values(): 

189 prompts.append({"name": p.name, "description": p.description, "arguments": p.arguments}) 

190 return {"prompts": prompts} 

191 

192 def _handle_prompt_get(self, params: dict) -> dict: 

193 name = params.get("name", "") 

194 prompt = self._prompts.get(name) 

195 if not prompt: 

196 raise ValueError(f"Prompt not found: {name}") 

197 return {"description": prompt.description, "messages": [{"role": "user", "content": {"type": "text", "text": prompt.template}}]} 

198 

199 def list_tools(self) -> list[dict]: 

200 """Compliance-facing tool listing method.""" 

201 return [ 

202 {"name": t.name, "description": t.description, "inputSchema": t.input_schema} 

203 for t in self._tools.values() 

204 ] 

205 

206 # ── 统计 ────────────────────────────────────── 

207 

208 def stats(self) -> dict: 

209 return {"tools": len(self._tools), "resources": len(self._resources), "prompts": len(self._prompts), "transport": self.config.transport} 

210 

211 

212class MCPClient: 

213 """MCP客户端 — AgentOS中Agent连接到外部MCP Server。""" 

214 

215 def __init__(self, server_url: str = "", transport: str = "stdio"): 

216 self.server_url = server_url 

217 self.transport = transport 

218 self._tools: list[dict] = [] 

219 self._connected = False 

220 

221 async def connect(self): 

222 # 模拟连接（实际生产环境用mcp SDK） 

223 self._connected = True 

224 logger.info(f"MCP client connected to {self.server_url}") 

225 

226 async def connect_server(self, config) -> bool: 

227 """Compliance-facing: connect to a MCP server.""" 

228 self.server_url = getattr(config, "command", "") 

229 self._connected = True 

230 return True 

231 

232 async def call_tool(self, name: str, arguments: dict = {}) -> dict: 

233 logger.info(f"MCP client calling tool: {name}") 

234 return {"result": f"Simulated call to {name}"} 

235 

236 async def list_tools(self) -> list[dict]: 

237 return self._tools 

238 

239 def disconnect(self): 

240 self._connected = False