Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(mcp): support multiple MCP servers by connecting each in isolated task

Browse Source 
Each MCP server now connects in its own asyncio.Task to isolate anyio
cancel scopes and prevent 'exit cancel scope in different task' errors
when multiple servers (especially mixed transport types) are configured.

Changes:
- connect_mcp_servers() returns dict[str, AsyncExitStack] instead of None
- Each server runs in separate task via asyncio.gather()
- AgentLoop uses _mcp_stacks dict to track per-server stacks
- Tests updated to handle new API
This commit is contained in:
worenidewen
2026-04-10 23:51:50 +08:00
parent 9bccfa63d2
commit a167959027
3 changed files with 247 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/tools/mcp.py
+59 -33
View File
@@ -57,9 +57,7 @@ def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]:
if "properties" in normalized and isinstance(normalized["properties"], dict):
normalized["properties"] = {
name: _normalize_schema_for_openai(prop)
if isinstance(prop, dict)
else prop
name: _normalize_schema_for_openai(prop) if isinstance(prop, dict) else prop
for name, prop in normalized["properties"].items()
}
@@ -138,9 +136,7 @@ class MCPToolWrapper(Tool):
class MCPResourceWrapper(Tool):
"""Wraps an MCP resource URI as a read-only nanobot Tool."""
def __init__(
self, session, server_name: str, resource_def, resource_timeout: int = 30
):
def __init__(self, session, server_name: str, resource_def, resource_timeout: int = 30):
self._session = session
self._uri = resource_def.uri
self._name = f"mcp_{server_name}_resource_{resource_def.name}"
@@ -211,9 +207,7 @@ class MCPResourceWrapper(Tool):
class MCPPromptWrapper(Tool):
"""Wraps an MCP prompt as a read-only nanobot Tool."""
def __init__(
self, session, server_name: str, prompt_def, prompt_timeout: int = 30
):
def __init__(self, session, server_name: str, prompt_def, prompt_timeout: int = 30):
self._session = session
self._prompt_name = prompt_def.name
self._name = f"mcp_{server_name}_prompt_{prompt_def.name}"
@@ -266,9 +260,7 @@ class MCPPromptWrapper(Tool):
timeout=self._prompt_timeout,
)
except asyncio.TimeoutError:
logger.warning(
"MCP prompt '{}' timed out after {}s", self._name, self._prompt_timeout
)
logger.warning("MCP prompt '{}' timed out after {}s", self._name, self._prompt_timeout)
return f"(MCP prompt call timed out after {self._prompt_timeout}s)"
except asyncio.CancelledError:
task = asyncio.current_task()
@@ -279,13 +271,17 @@ class MCPPromptWrapper(Tool):
except McpError as exc:
logger.error(
"MCP prompt '{}' failed: code={} message={}",
self._name, exc.error.code, exc.error.message,
self._name,
exc.error.code,
exc.error.message,
)
return f"(MCP prompt call failed: {exc.error.message} [code {exc.error.code}])"
except Exception as exc:
logger.exception(
"MCP prompt '{}' failed: {}: {}",
self._name, type(exc).__name__, exc,
self._name,
type(exc).__name__,
exc,
)
return f"(MCP prompt call failed: {type(exc).__name__})"
@@ -307,35 +303,44 @@ class MCPPromptWrapper(Tool):
async def connect_mcp_servers(
mcp_servers: dict, registry: ToolRegistry, stack: AsyncExitStack
) -> None:
"""Connect to configured MCP servers and register their tools, resources, and prompts."""
mcp_servers: dict, registry: ToolRegistry
) -> dict[str, AsyncExitStack]:
"""Connect to configured MCP servers and register their tools, resources, prompts.
Returns a dict mapping server name -> its dedicated AsyncExitStack.
Each server gets its own stack and runs in its own task to prevent
cancel scope conflicts when multiple MCP servers are configured.
"""
from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamable_http_client
for name, cfg in mcp_servers.items():
async def connect_single_server(name: str, cfg) -> tuple[str, AsyncExitStack | None]:
server_stack = AsyncExitStack()
await server_stack.__aenter__()
try:
transport_type = cfg.type
if not transport_type:
if cfg.command:
transport_type = "stdio"
elif cfg.url:
# Convention: URLs ending with /sse use SSE transport; others use streamableHttp
transport_type = (
"sse" if cfg.url.rstrip("/").endswith("/sse") else "streamableHttp"
)
else:
logger.warning("MCP server '{}': no command or url configured, skipping", name)
continue
await server_stack.aclose()
return name, None
if transport_type == "stdio":
params = StdioServerParameters(
command=cfg.command, args=cfg.args, env=cfg.env or None
)
read, write = await stack.enter_async_context(stdio_client(params))
read, write = await server_stack.enter_async_context(stdio_client(params))
elif transport_type == "sse":
def httpx_client_factory(
headers: dict[str, str] | None = None,
timeout: httpx.Timeout | None = None,
@@ -353,27 +358,26 @@ async def connect_mcp_servers(
auth=auth,
)
read, write = await stack.enter_async_context(
read, write = await server_stack.enter_async_context(
sse_client(cfg.url, httpx_client_factory=httpx_client_factory)
)
elif transport_type == "streamableHttp":
# Always provide an explicit httpx client so MCP HTTP transport does not
# inherit httpx's default 5s timeout and preempt the higher-level tool timeout.
http_client = await stack.enter_async_context(
http_client = await server_stack.enter_async_context(
httpx.AsyncClient(
headers=cfg.headers or None,
follow_redirects=True,
timeout=None,
)
)
read, write, _ = await stack.enter_async_context(
read, write, _ = await server_stack.enter_async_context(
streamable_http_client(cfg.url, http_client=http_client)
)
else:
logger.warning("MCP server '{}': unknown transport type '{}'", name, transport_type)
continue
await server_stack.aclose()
return name, None
session = await stack.enter_async_context(ClientSession(read, write))
session = await server_stack.enter_async_context(ClientSession(read, write))
await session.initialize()
tools = await session.list_tools()
@@ -418,7 +422,6 @@ async def connect_mcp_servers(
", ".join(available_wrapped_names) or "(none)",
)
# --- Register resources ---
try:
resources_result = await session.list_resources()
for resource in resources_result.resources:
@@ -433,7 +436,6 @@ async def connect_mcp_servers(
except Exception as e:
logger.debug("MCP server '{}': resources not supported or failed: {}", name, e)
# --- Register prompts ---
try:
prompts_result = await session.list_prompts()
for prompt in prompts_result.prompts:
@@ -442,14 +444,38 @@ async def connect_mcp_servers(
)
registry.register(wrapper)
registered_count += 1
logger.debug(
"MCP: registered prompt '{}' from server '{}'", wrapper.name, name
)
logger.debug("MCP: registered prompt '{}' from server '{}'", wrapper.name, name)
except Exception as e:
logger.debug("MCP server '{}': prompts not supported or failed: {}", name, e)
logger.info(
"MCP server '{}': connected, {} capabilities registered", name, registered_count
)
return name, server_stack
except Exception as e:
logger.error("MCP server '{}': failed to connect: {}", name, e)
try:
await server_stack.aclose()
except Exception:
pass
return name, None
server_stacks: dict[str, AsyncExitStack] = {}
tasks: list[asyncio.Task] = []
for name, cfg in mcp_servers.items():
task = asyncio.create_task(connect_single_server(name, cfg))
tasks.append(task)
results = await asyncio.gather(*tasks, return_exceptions=True)
for i, result in enumerate(results):
name = list(mcp_servers.keys())[i]
if isinstance(result, BaseException):
if not isinstance(result, asyncio.CancelledError):
logger.error("MCP server '{}' connection task failed: {}", name, result)
elif result is not None and result[1] is not None:
server_stacks[result[0]] = result[1]
return server_stacks