返回 JoyAI-Echo
1 """MCP client: connects to MCP servers and wraps their tools as native nanobot tools."""
2
3 import asyncio
4 import os
5 import shutil
6 from contextlib import AsyncExitStack
7 from typing import Any
8
9 import httpx
10 from loguru import logger
11
12 from nanobot.agent.tools.base import Tool
13 from nanobot.agent.tools.registry import ToolRegistry
14
15 # Transient connection errors that warrant a single retry.
16 # These typically happen when an MCP server restarts or a network
17 # connection is interrupted between calls.
18 _TRANSIENT_EXC_NAMES: frozenset[str] = frozenset((
19 "ClosedResourceError",
20 "BrokenResourceError",
21 "EndOfStream",
22 "BrokenPipeError",
23 "ConnectionResetError",
24 "ConnectionRefusedError",
25 "ConnectionAbortedError",
26 "ConnectionError",
27 ))
28
29 _WINDOWS_SHELL_LAUNCHERS: frozenset[str] = frozenset(("npx", "npm", "pnpm", "yarn", "bunx"))
30
31
32 def _is_transient(exc: BaseException) -> bool:
33 """Check if an exception looks like a transient connection error."""
34 return type(exc).__name__ in _TRANSIENT_EXC_NAMES
35
36
37 def _windows_command_basename(command: str) -> str:
38 """Return the lowercase basename for a Windows command or path."""
39 return command.replace("\\", "/").rsplit("/", maxsplit=1)[-1].lower()
40
41
42 def _normalize_windows_stdio_command(
43 command: str,
44 args: list[str] | None,
45 env: dict[str, str] | None,
46 ) -> tuple[str, list[str], dict[str, str] | None]:
47 """Wrap Windows shell launchers so MCP stdio servers start reliably."""
48 normalized_args = list(args or [])
49 if os.name != "nt":
50 return command, normalized_args, env
51
52 basename = _windows_command_basename(command)
53 if basename in {"cmd", "cmd.exe", "powershell", "powershell.exe", "pwsh", "pwsh.exe"}:
54 return command, normalized_args, env
55
56 if basename.endswith((".exe", ".com")):
57 return command, normalized_args, env
58
59 resolved = shutil.which(command, path=(env or {}).get("PATH")) or command
60 resolved_basename = _windows_command_basename(resolved)
61 should_wrap = (
62 basename in _WINDOWS_SHELL_LAUNCHERS
63 or basename.endswith((".cmd", ".bat"))
64 or resolved_basename.endswith((".cmd", ".bat"))
65 )
66 if not should_wrap:
67 return command, normalized_args, env
68
69 comspec = (env or {}).get("COMSPEC") or os.environ.get("COMSPEC") or "cmd.exe"
70 return comspec, ["/d", "/c", command, *normalized_args], env
71
72
73 def _extract_nullable_branch(options: Any) -> tuple[dict[str, Any], bool] | None:
74 """Return the single non-null branch for nullable unions."""
75 if not isinstance(options, list):
76 return None
77
78 non_null: list[dict[str, Any]] = []
79 saw_null = False
80 for option in options:
81 if not isinstance(option, dict):
82 return None
83 if option.get("type") == "null":
84 saw_null = True
85 continue
86 non_null.append(option)
87
88 if saw_null and len(non_null) == 1:
89 return non_null[0], True
90 return None
91
92
93 def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]:
94 """Normalize only nullable JSON Schema patterns for tool definitions."""
95 if not isinstance(schema, dict):
96 return {"type": "object", "properties": {}}
97
98 normalized = dict(schema)
99
100 raw_type = normalized.get("type")
101 if isinstance(raw_type, list):
102 non_null = [item for item in raw_type if item != "null"]
103 if "null" in raw_type and len(non_null) == 1:
104 normalized["type"] = non_null[0]
105 normalized["nullable"] = True
106
107 for key in ("oneOf", "anyOf"):
108 nullable_branch = _extract_nullable_branch(normalized.get(key))
109 if nullable_branch is not None:
110 branch, _ = nullable_branch
111 merged = {k: v for k, v in normalized.items() if k != key}
112 merged.update(branch)
113 normalized = merged
114 normalized["nullable"] = True
115 break
116
117 if "properties" in normalized and isinstance(normalized["properties"], dict):
118 normalized["properties"] = {
119 name: _normalize_schema_for_openai(prop) if isinstance(prop, dict) else prop
120 for name, prop in normalized["properties"].items()
121 }
122
123 if "items" in normalized and isinstance(normalized["items"], dict):
124 normalized["items"] = _normalize_schema_for_openai(normalized["items"])
125
126 if normalized.get("type") != "object":
127 return normalized
128
129 normalized.setdefault("properties", {})
130 normalized.setdefault("required", [])
131 return normalized
132
133
134 class MCPToolWrapper(Tool):
135 """Wraps a single MCP server tool as a nanobot Tool."""
136
137 def __init__(self, session, server_name: str, tool_def, tool_timeout: int = 30):
138 self._session = session
139 self._original_name = tool_def.name
140 self._name = f"mcp_{server_name}_{tool_def.name}"
141 self._description = tool_def.description or tool_def.name
142 raw_schema = tool_def.inputSchema or {"type": "object", "properties": {}}
143 self._parameters = _normalize_schema_for_openai(raw_schema)
144 self._tool_timeout = tool_timeout
145
146 @property
147 def name(self) -> str:
148 return self._name
149
150 @property
151 def description(self) -> str:
152 return self._description
153
154 @property
155 def parameters(self) -> dict[str, Any]:
156 return self._parameters
157
158 async def execute(self, **kwargs: Any) -> str:
159 from mcp import types
160
161 for attempt in range(2): # At most 1 retry
162 try:
163 result = await asyncio.wait_for(
164 self._session.call_tool(self._original_name, arguments=kwargs),
165 timeout=self._tool_timeout,
166 )
167 except asyncio.TimeoutError:
168 logger.warning(
169 "MCP tool '{}' timed out after {}s", self._name, self._tool_timeout
170 )
171 return f"(MCP tool call timed out after {self._tool_timeout}s)"
172 except asyncio.CancelledError:
173 # MCP SDK's anyio cancel scopes can leak CancelledError on timeout/failure.
174 # Re-raise only if our task was externally cancelled (e.g. /stop).
175 task = asyncio.current_task()
176 if task is not None and task.cancelling() > 0:
177 raise
178 logger.warning("MCP tool '{}' was cancelled by server/SDK", self._name)
179 return "(MCP tool call was cancelled)"
180 except Exception as exc:
181 if _is_transient(exc):
182 if attempt == 0:
183 logger.warning(
184 "MCP tool '{}' hit transient error ({}), retrying once...",
185 self._name,
186 type(exc).__name__,
187 )
188 await asyncio.sleep(1) # Brief backoff before retry
189 continue
190 # Second transient failure — give up with retry-specific message
191 logger.error(
192 "MCP tool '{}' failed after retry: {}: {}",
193 self._name,
194 type(exc).__name__,
195 exc,
196 )
197 return f"(MCP tool call failed after retry: {type(exc).__name__})"
198 logger.exception(
199 "MCP tool '{}' failed: {}: {}",
200 self._name,
201 type(exc).__name__,
202 exc,
203 )
204 return f"(MCP tool call failed: {type(exc).__name__})"
205 else:
206 # Success — extract result
207 parts = []
208 for block in result.content:
209 if isinstance(block, types.TextContent):
210 parts.append(block.text)
211 else:
212 parts.append(str(block))
213 return "\n".join(parts) or "(no output)"
214
215 return "(MCP tool call failed)" # Unreachable, but satisfies type checkers
216
217
218 class MCPResourceWrapper(Tool):
219 """Wraps an MCP resource URI as a read-only nanobot Tool."""
220
221 def __init__(self, session, server_name: str, resource_def, resource_timeout: int = 30):
222 self._session = session
223 self._uri = resource_def.uri
224 self._name = f"mcp_{server_name}_resource_{resource_def.name}"
225 desc = resource_def.description or resource_def.name
226 self._description = f"[MCP Resource] {desc}\nURI: {self._uri}"
227 self._parameters: dict[str, Any] = {
228 "type": "object",
229 "properties": {},
230 "required": [],
231 }
232 self._resource_timeout = resource_timeout
233
234 @property
235 def name(self) -> str:
236 return self._name
237
238 @property
239 def description(self) -> str:
240 return self._description
241
242 @property
243 def parameters(self) -> dict[str, Any]:
244 return self._parameters
245
246 @property
247 def read_only(self) -> bool:
248 return True
249
250 async def execute(self, **kwargs: Any) -> str:
251 from mcp import types
252
253 for attempt in range(2):
254 try:
255 result = await asyncio.wait_for(
256 self._session.read_resource(self._uri),
257 timeout=self._resource_timeout,
258 )
259 except asyncio.TimeoutError:
260 logger.warning(
261 "MCP resource '{}' timed out after {}s", self._name, self._resource_timeout
262 )
263 return f"(MCP resource read timed out after {self._resource_timeout}s)"
264 except asyncio.CancelledError:
265 task = asyncio.current_task()
266 if task is not None and task.cancelling() > 0:
267 raise
268 logger.warning("MCP resource '{}' was cancelled by server/SDK", self._name)
269 return "(MCP resource read was cancelled)"
270 except Exception as exc:
271 if _is_transient(exc):
272 if attempt == 0:
273 logger.warning(
274 "MCP resource '{}' hit transient error ({}), retrying once...",
275 self._name,
276 type(exc).__name__,
277 )
278 await asyncio.sleep(1)
279 continue
280 logger.error(
281 "MCP resource '{}' failed after retry: {}: {}",
282 self._name,
283 type(exc).__name__,
284 exc,
285 )
286 return f"(MCP resource read failed after retry: {type(exc).__name__})"
287 logger.exception(
288 "MCP resource '{}' failed: {}: {}",
289 self._name,
290 type(exc).__name__,
291 exc,
292 )
293 return f"(MCP resource read failed: {type(exc).__name__})"
294 else:
295 parts: list[str] = []
296 for block in result.contents:
297 if isinstance(block, types.TextResourceContents):
298 parts.append(block.text)
299 elif isinstance(block, types.BlobResourceContents):
300 parts.append(f"[Binary resource: {len(block.blob)} bytes]")
301 else:
302 parts.append(str(block))
303 return "\n".join(parts) or "(no output)"
304
305 return "(MCP resource read failed)" # Unreachable
306
307
308 class MCPPromptWrapper(Tool):
309 """Wraps an MCP prompt as a read-only nanobot Tool."""
310
311 def __init__(self, session, server_name: str, prompt_def, prompt_timeout: int = 30):
312 self._session = session
313 self._prompt_name = prompt_def.name
314 self._name = f"mcp_{server_name}_prompt_{prompt_def.name}"
315 desc = prompt_def.description or prompt_def.name
316 self._description = (
317 f"[MCP Prompt] {desc}\n"
318 "Returns a filled prompt template that can be used as a workflow guide."
319 )
320 self._prompt_timeout = prompt_timeout
321
322 # Build parameters from prompt arguments
323 properties: dict[str, Any] = {}
324 required: list[str] = []
325 for arg in prompt_def.arguments or []:
326 prop: dict[str, Any] = {"type": "string"}
327 if getattr(arg, "description", None):
328 prop["description"] = arg.description
329 properties[arg.name] = prop
330 if arg.required:
331 required.append(arg.name)
332 self._parameters: dict[str, Any] = {
333 "type": "object",
334 "properties": properties,
335 "required": required,
336 }
337
338 @property
339 def name(self) -> str:
340 return self._name
341
342 @property
343 def description(self) -> str:
344 return self._description
345
346 @property
347 def parameters(self) -> dict[str, Any]:
348 return self._parameters
349
350 @property
351 def read_only(self) -> bool:
352 return True
353
354 async def execute(self, **kwargs: Any) -> str:
355 from mcp import types
356 from mcp.shared.exceptions import McpError
357
358 for attempt in range(2):
359 try:
360 result = await asyncio.wait_for(
361 self._session.get_prompt(self._prompt_name, arguments=kwargs),
362 timeout=self._prompt_timeout,
363 )
364 except asyncio.TimeoutError:
365 logger.warning(
366 "MCP prompt '{}' timed out after {}s", self._name, self._prompt_timeout
367 )
368 return f"(MCP prompt call timed out after {self._prompt_timeout}s)"
369 except asyncio.CancelledError:
370 task = asyncio.current_task()
371 if task is not None and task.cancelling() > 0:
372 raise
373 logger.warning("MCP prompt '{}' was cancelled by server/SDK", self._name)
374 return "(MCP prompt call was cancelled)"
375 except McpError as exc:
376 logger.error(
377 "MCP prompt '{}' failed: code={} message={}",
378 self._name,
379 exc.error.code,
380 exc.error.message,
381 )
382 return f"(MCP prompt call failed: {exc.error.message} [code {exc.error.code}])"
383 except Exception as exc:
384 if _is_transient(exc):
385 if attempt == 0:
386 logger.warning(
387 "MCP prompt '{}' hit transient error ({}), retrying once...",
388 self._name,
389 type(exc).__name__,
390 )
391 await asyncio.sleep(1)
392 continue
393 logger.error(
394 "MCP prompt '{}' failed after retry: {}: {}",
395 self._name,
396 type(exc).__name__,
397 exc,
398 )
399 return f"(MCP prompt call failed after retry: {type(exc).__name__})"
400 logger.exception(
401 "MCP prompt '{}' failed: {}: {}",
402 self._name,
403 type(exc).__name__,
404 exc,
405 )
406 return f"(MCP prompt call failed: {type(exc).__name__})"
407 else:
408 parts: list[str] = []
409 for message in result.messages:
410 content = message.content
411 if isinstance(content, types.TextContent):
412 parts.append(content.text)
413 elif isinstance(content, list):
414 for block in content:
415 if isinstance(block, types.TextContent):
416 parts.append(block.text)
417 else:
418 parts.append(str(block))
419 else:
420 parts.append(str(content))
421 return "\n".join(parts) or "(no output)"
422
423 return "(MCP prompt call failed)" # Unreachable
424
425
426 async def connect_mcp_servers(
427 mcp_servers: dict, registry: ToolRegistry
428 ) -> dict[str, AsyncExitStack]:
429 """Connect to configured MCP servers and register their tools, resources, prompts.
430
431 Returns a dict mapping server name -> its dedicated AsyncExitStack.
432 Each server gets its own stack and runs in its own task to prevent
433 cancel scope conflicts when multiple MCP servers are configured.
434 """
435 from mcp import ClientSession, StdioServerParameters
436 from mcp.client.sse import sse_client
437 from mcp.client.stdio import stdio_client
438 from mcp.client.streamable_http import streamable_http_client
439
440 async def connect_single_server(name: str, cfg) -> tuple[str, AsyncExitStack | None]:
441 server_stack = AsyncExitStack()
442 await server_stack.__aenter__()
443
444 try:
445 transport_type = cfg.type
446 if not transport_type:
447 if cfg.command:
448 transport_type = "stdio"
449 elif cfg.url:
450 transport_type = (
451 "sse" if cfg.url.rstrip("/").endswith("/sse") else "streamableHttp"
452 )
453 else:
454 logger.warning("MCP server '{}': no command or url configured, skipping", name)
455 await server_stack.aclose()
456 return name, None
457
458 if transport_type == "stdio":
459 command, args, env = _normalize_windows_stdio_command(
460 cfg.command,
461 cfg.args,
462 cfg.env or None,
463 )
464 params = StdioServerParameters(
465 command=command,
466 args=args,
467 env=env,
468 )
469 read, write = await server_stack.enter_async_context(stdio_client(params))
470 elif transport_type == "sse":
471
472 def httpx_client_factory(
473 headers: dict[str, str] | None = None,
474 timeout: httpx.Timeout | None = None,
475 auth: httpx.Auth | None = None,
476 ) -> httpx.AsyncClient:
477 merged_headers = {
478 "Accept": "application/json, text/event-stream",
479 **(cfg.headers or {}),
480 **(headers or {}),
481 }
482 return httpx.AsyncClient(
483 headers=merged_headers or None,
484 follow_redirects=True,
485 timeout=timeout,
486 auth=auth,
487 )
488
489 read, write = await server_stack.enter_async_context(
490 sse_client(cfg.url, httpx_client_factory=httpx_client_factory)
491 )
492 elif transport_type == "streamableHttp":
493 http_client = await server_stack.enter_async_context(
494 httpx.AsyncClient(
495 headers=cfg.headers or None,
496 follow_redirects=True,
497 timeout=None,
498 )
499 )
500 read, write, _ = await server_stack.enter_async_context(
501 streamable_http_client(cfg.url, http_client=http_client)
502 )
503 else:
504 logger.warning("MCP server '{}': unknown transport type '{}'", name, transport_type)
505 await server_stack.aclose()
506 return name, None
507
508 session = await server_stack.enter_async_context(ClientSession(read, write))
509 await session.initialize()
510
511 tools = await session.list_tools()
512 enabled_tools = set(cfg.enabled_tools)
513 allow_all_tools = "*" in enabled_tools
514 registered_count = 0
515 matched_enabled_tools: set[str] = set()
516 available_raw_names = [tool_def.name for tool_def in tools.tools]
517 available_wrapped_names = [f"mcp_{name}_{tool_def.name}" for tool_def in tools.tools]
518 for tool_def in tools.tools:
519 wrapped_name = f"mcp_{name}_{tool_def.name}"
520 if (
521 not allow_all_tools
522 and tool_def.name not in enabled_tools
523 and wrapped_name not in enabled_tools
524 ):
525 logger.debug(
526 "MCP: skipping tool '{}' from server '{}' (not in enabledTools)",
527 wrapped_name,
528 name,
529 )
530 continue
531 wrapper = MCPToolWrapper(session, name, tool_def, tool_timeout=cfg.tool_timeout)
532 registry.register(wrapper)
533 logger.debug("MCP: registered tool '{}' from server '{}'", wrapper.name, name)
534 registered_count += 1
535 if enabled_tools:
536 if tool_def.name in enabled_tools:
537 matched_enabled_tools.add(tool_def.name)
538 if wrapped_name in enabled_tools:
539 matched_enabled_tools.add(wrapped_name)
540
541 if enabled_tools and not allow_all_tools:
542 unmatched_enabled_tools = sorted(enabled_tools - matched_enabled_tools)
543 if unmatched_enabled_tools:
544 logger.warning(
545 "MCP server '{}': enabledTools entries not found: {}. Available raw names: {}. "
546 "Available wrapped names: {}",
547 name,
548 ", ".join(unmatched_enabled_tools),
549 ", ".join(available_raw_names) or "(none)",
550 ", ".join(available_wrapped_names) or "(none)",
551 )
552
553 try:
554 resources_result = await session.list_resources()
555 for resource in resources_result.resources:
556 wrapper = MCPResourceWrapper(
557 session, name, resource, resource_timeout=cfg.tool_timeout
558 )
559 registry.register(wrapper)
560 registered_count += 1
561 logger.debug(
562 "MCP: registered resource '{}' from server '{}'", wrapper.name, name
563 )
564 except Exception as e:
565 logger.debug("MCP server '{}': resources not supported or failed: {}", name, e)
566
567 try:
568 prompts_result = await session.list_prompts()
569 for prompt in prompts_result.prompts:
570 wrapper = MCPPromptWrapper(
571 session, name, prompt, prompt_timeout=cfg.tool_timeout
572 )
573 registry.register(wrapper)
574 registered_count += 1
575 logger.debug("MCP: registered prompt '{}' from server '{}'", wrapper.name, name)
576 except Exception as e:
577 logger.debug("MCP server '{}': prompts not supported or failed: {}", name, e)
578
579 logger.info(
580 "MCP server '{}': connected, {} capabilities registered", name, registered_count
581 )
582 return name, server_stack
583
584 except Exception as e:
585 hint = ""
586 text = str(e).lower()
587 if any(
588 marker in text
589 for marker in (
590 "parse error",
591 "invalid json",
592 "unexpected token",
593 "jsonrpc",
594 "content-length",
595 )
596 ):
597 hint = (
598 " Hint: this looks like stdio protocol pollution. Make sure the MCP server writes "
599 "only JSON-RPC to stdout and sends logs/debug output to stderr instead."
600 )
601 logger.error("MCP server '{}': failed to connect: {}{}", name, e, hint)
602 try:
603 await server_stack.aclose()
604 except Exception:
605 pass
606 return name, None
607
608 server_stacks: dict[str, AsyncExitStack] = {}
609
610 tasks: list[asyncio.Task] = []
611 for name, cfg in mcp_servers.items():
612 task = asyncio.create_task(connect_single_server(name, cfg))
613 tasks.append(task)
614
615 results = await asyncio.gather(*tasks, return_exceptions=True)
616
617 for i, result in enumerate(results):
618 name = list(mcp_servers.keys())[i]
619 if isinstance(result, BaseException):
620 if not isinstance(result, asyncio.CancelledError):
621 logger.error("MCP server '{}' connection task failed: {}", name, result)
622 elif result is not None and result[1] is not None:
623 server_stacks[result[0]] = result[1]
624
625 return server_stacks
626
626 lines PYTHON