1"""

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

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

4"""

5

6 from __future__ import annotations

7

8 import json

9 import logging

10 from dataclasses import dataclass, field

11 from typing import Optional, Callable, Any

12

13

14 logger = logging.getLogger(__name__)

15

16

17@dataclass

18 class 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

28 class 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

38 class 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

48 class MCPPrompt:

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

50 name: str

51 description: str = ""

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

53 template: str = ""

54

55

56 class 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

212 class 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