Coverage for merco/mcp/manager.py: 68%

1"""MCP Server lifecycle manager — connect, discover tools, register.""" 

2import asyncio 

3import logging 

4import time 

5from .config import MCPServerConfig 

6from .tool import MCPServerTool 

7 

8logger = logging.getLogger("merco.mcp") 

9 

10# Optional mcp imports 

11_MCP_AVAILABLE = False 

12try: 

13 from mcp import ClientSession, StdioServerParameters 

14 from mcp.client.stdio import stdio_client 

15 _MCP_AVAILABLE = True 

16except ImportError: 

17 pass 

18 

19 

20class MCPServerManager: 

21 def __init__(self, tool_registry, hooks=None): 

22 self._registry = tool_registry 

23 self._hooks = hooks 

24 self._servers: dict[str, dict] = {} # name → {config, tools: [MCPServerTool]} 

25 self._original_config: dict = {} 

26 

27 async def load_config(self, servers_config: dict) -> None: 

28 """Load config from merco.json. Connect enabled, skip if already connected.""" 

29 self._original_config = servers_config 

30 if not _MCP_AVAILABLE: 

31 logger.warning("mcp package not installed — skipping MCP") 

32 return 

33 for name, data in servers_config.items(): 

34 if name in self._servers: 

35 continue 

36 cfg = MCPServerConfig.from_dict(name, data) 

37 if not cfg.enabled: 

38 continue 

39 await self.connect(name, cfg) 

40 

41 async def connect(self, name: str, config: MCPServerConfig) -> bool: 

42 """Connect to MCP server + discover tools + register them.""" 

43 if not _MCP_AVAILABLE: 

44 return False 

45 try: 

46 if config.command: 

47 tools = await self._connect_stdio(config) 

48 elif config.url: 

49 tools = await self._connect_http(config) 

50 else: 

51 logger.warning("MCP '%s': no command or url", name) 

52 return False 

53 

54 # Unregister old tools if reconnecting 

55 await self._unregister_tools(name) 

56 

57 # Register each tool 

58 server_tools = [] 

59 for spec in tools: 

60 tool = MCPServerTool(spec, name, self._call_tool) 

61 self._registry.register(tool) 

62 server_tools.append(tool) 

63 

64 self._servers[name] = {"config": config, "tools": server_tools} 

65 if self._hooks: 

66 await self._hooks.emit("mcp.connect", server=name, tools=len(tools)) 

67 logger.info("MCP '%s': %d tools registered", name, len(tools)) 

68 return True 

69 except Exception as e: 

70 logger.warning("MCP '%s' connection failed: %s", name, e) 

71 return False 

72 

73 async def _connect_stdio(self, config: MCPServerConfig) -> list[dict]: 

74 params = StdioServerParameters( 

75 command=config.command, args=config.args, env=config.env 

76 ) 

77 async with stdio_client(params) as (read, write): 

78 async with ClientSession(read, write) as session: 

79 await session.initialize() 

80 result = await session.list_tools() 

81 return [t.model_dump() for t in result.tools] 

82 

83 async def _connect_http(self, config: MCPServerConfig) -> list[dict]: 

84 # StreamableHTTP support — use mcp.client.streamable_http if available 

85 try: 

86 from mcp.client.streamable_http import streamablehttp_client 

87 except ImportError: 

88 raise ImportError("mcp HTTP transport requires mcp>=1.0 with streamable_http") 

89 async with streamablehttp_client(config.url, headers=config.headers) as (read, write): 

90 async with ClientSession(read, write) as session: 

91 await session.initialize() 

92 result = await session.list_tools() 

93 return [t.model_dump() for t in result.tools] 

94 

95 async def _unregister_tools(self, name: str) -> None: 

96 if name in self._servers: 

97 for tool in self._servers[name]["tools"]: 

98 self._registry.unregister(tool.name) 

99 

100 async def disconnect(self, name: str) -> None: 

101 await self._unregister_tools(name) 

102 self._servers.pop(name, None) 

103 

104 async def shutdown(self): 

105 """关闭所有 MCP 连接。""" 

106 for name in list(self._servers.keys()): 

107 await self.disconnect(name) 

108 

109 async def reload(self) -> None: 

110 for name in list(self._servers.keys()): 

111 await self.disconnect(name) 

112 await self.load_config(self._original_config) 

113 

114 async def _call_tool(self, tool_name: str, arguments: dict) -> dict: 

115 # Find which server owns this tool, call via MCP session 

116 for name, state in self._servers.items(): 

117 for tool in state["tools"]: 

118 if tool.name == tool_name: 

119 t0 = time.monotonic() 

120 try: 

121 if state["config"].command: 

122 result = await self._call_stdio_tool(state["config"], tool_name, arguments) 

123 else: 

124 result = await self._call_http_tool(state["config"], tool_name, arguments) 

125 if self._hooks: 

126 await self._hooks.emit("mcp.tool_call", server=name, tool=tool_name, 

127 duration=time.monotonic()-t0) 

128 return result 

129 except Exception as e: 

130 if self._hooks: 

131 await self._hooks.emit("mcp.error", server=name, tool=tool_name, error=str(e)) 

132 raise 

133 return {"error": f"Tool '{tool_name}' not found in any MCP server", "isError": True} 

134 

135 async def _call_stdio_tool(self, config: MCPServerConfig, tool_name: str, arguments: dict) -> dict: 

136 params = StdioServerParameters( 

137 command=config.command, args=config.args, env=config.env 

138 ) 

139 async with stdio_client(params) as (read, write): 

140 async with ClientSession(read, write) as session: 

141 await session.initialize() 

142 result = await session.call_tool(tool_name, arguments) 

143 return result.model_dump() 

144 

145 async def _call_http_tool(self, config: MCPServerConfig, tool_name: str, arguments: dict) -> dict: 

146 try: 

147 from mcp.client.streamable_http import streamablehttp_client 

148 except ImportError: 

149 raise ImportError("mcp HTTP transport requires mcp>=1.0 with streamable_http") 

150 async with streamablehttp_client(config.url, headers=config.headers) as (read, write): 

151 async with ClientSession(read, write) as session: 

152 await session.initialize() 

153 result = await session.call_tool(tool_name, arguments) 

154 return result.model_dump() 

155 

156 def status(self) -> dict: 

157 return { 

158 name: { 

159 "connected": True, 

160 "tools_count": len(state["tools"]), 

161 "enabled": state["config"].enabled, 

162 } 

163 for name, state in self._servers.items() 

164 }