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

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
+36 -72
View File
@@ -271,15 +271,11 @@ async def test_connect_mcp_servers_enabled_tools_supports_raw_names(
) -> None:
fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
registry = ToolRegistry()
stack = AsyncExitStack()
await stack.__aenter__()
try:
await connect_mcp_servers(
{"test": MCPServerConfig(command="fake", enabled_tools=["demo"])},
registry,
stack,
)
finally:
stacks = await connect_mcp_servers(
{"test": MCPServerConfig(command="fake", enabled_tools=["demo"])},
registry,
)
for stack in stacks.values():
await stack.aclose()
assert registry.tool_names == ["mcp_test_demo"]
@@ -291,15 +287,11 @@ async def test_connect_mcp_servers_enabled_tools_defaults_to_all(
) -> None:
fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
registry = ToolRegistry()
stack = AsyncExitStack()
await stack.__aenter__()
try:
await connect_mcp_servers(
{"test": MCPServerConfig(command="fake")},
registry,
stack,
)
finally:
stacks = await connect_mcp_servers(
{"test": MCPServerConfig(command="fake")},
registry,
)
for stack in stacks.values():
await stack.aclose()
assert registry.tool_names == ["mcp_test_demo", "mcp_test_other"]
@@ -311,15 +303,11 @@ async def test_connect_mcp_servers_enabled_tools_supports_wrapped_names(
) -> None:
fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
registry = ToolRegistry()
stack = AsyncExitStack()
await stack.__aenter__()
try:
await connect_mcp_servers(
{"test": MCPServerConfig(command="fake", enabled_tools=["mcp_test_demo"])},
registry,
stack,
)
finally:
stacks = await connect_mcp_servers(
{"test": MCPServerConfig(command="fake", enabled_tools=["mcp_test_demo"])},
registry,
)
for stack in stacks.values():
await stack.aclose()
assert registry.tool_names == ["mcp_test_demo"]
@@ -331,15 +319,11 @@ async def test_connect_mcp_servers_enabled_tools_empty_list_registers_none(
) -> None:
fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
registry = ToolRegistry()
stack = AsyncExitStack()
await stack.__aenter__()
try:
await connect_mcp_servers(
{"test": MCPServerConfig(command="fake", enabled_tools=[])},
registry,
stack,
)
finally:
stacks = await connect_mcp_servers(
{"test": MCPServerConfig(command="fake", enabled_tools=[])},
registry,
)
for stack in stacks.values():
await stack.aclose()
assert registry.tool_names == []
@@ -358,15 +342,11 @@ async def test_connect_mcp_servers_enabled_tools_warns_on_unknown_entries(
monkeypatch.setattr("nanobot.agent.tools.mcp.logger.warning", _warning)
stack = AsyncExitStack()
await stack.__aenter__()
try:
await connect_mcp_servers(
{"test": MCPServerConfig(command="fake", enabled_tools=["unknown"])},
registry,
stack,
)
finally:
stacks = await connect_mcp_servers(
{"test": MCPServerConfig(command="fake", enabled_tools=["unknown"])},
registry,
)
for stack in stacks.values():
await stack.aclose()
assert registry.tool_names == []
@@ -389,9 +369,7 @@ def _make_resource_def(
return SimpleNamespace(name=name, uri=uri, description=description)
def _make_resource_wrapper(
session: object, *, timeout: float = 0.1
) -> MCPResourceWrapper:
def _make_resource_wrapper(session: object, *, timeout: float = 0.1) -> MCPResourceWrapper:
return MCPResourceWrapper(session, "srv", _make_resource_def(), resource_timeout=timeout)
@@ -434,9 +412,7 @@ async def test_resource_wrapper_execute_handles_timeout() -> None:
await asyncio.sleep(1)
return SimpleNamespace(contents=[])
wrapper = _make_resource_wrapper(
SimpleNamespace(read_resource=read_resource), timeout=0.01
)
wrapper = _make_resource_wrapper(SimpleNamespace(read_resource=read_resource), timeout=0.01)
result = await wrapper.execute()
assert result == "(MCP resource read timed out after 0.01s)"
@@ -464,20 +440,14 @@ def _make_prompt_def(
return SimpleNamespace(name=name, description=description, arguments=arguments)
def _make_prompt_wrapper(
session: object, *, timeout: float = 0.1
) -> MCPPromptWrapper:
return MCPPromptWrapper(
session, "srv", _make_prompt_def(), prompt_timeout=timeout
)
def _make_prompt_wrapper(session: object, *, timeout: float = 0.1) -> MCPPromptWrapper:
return MCPPromptWrapper(session, "srv", _make_prompt_def(), prompt_timeout=timeout)
def test_prompt_wrapper_properties() -> None:
arg1 = SimpleNamespace(name="topic", required=True)
arg2 = SimpleNamespace(name="style", required=False)
wrapper = MCPPromptWrapper(
None, "myserver", _make_prompt_def(arguments=[arg1, arg2])
)
wrapper = MCPPromptWrapper(None, "myserver", _make_prompt_def(arguments=[arg1, arg2]))
assert wrapper.name == "mcp_myserver_prompt_myprompt"
assert "[MCP Prompt]" in wrapper.description
assert "A test prompt" in wrapper.description
@@ -528,9 +498,7 @@ async def test_prompt_wrapper_execute_handles_timeout() -> None:
await asyncio.sleep(1)
return SimpleNamespace(messages=[])
wrapper = _make_prompt_wrapper(
SimpleNamespace(get_prompt=get_prompt), timeout=0.01
)
wrapper = _make_prompt_wrapper(SimpleNamespace(get_prompt=get_prompt), timeout=0.01)
result = await wrapper.execute()
assert result == "(MCP prompt call timed out after 0.01s)"
@@ -616,15 +584,11 @@ async def test_connect_registers_resources_and_prompts(
prompt_names=["prompt_c"],
)
registry = ToolRegistry()
stack = AsyncExitStack()
await stack.__aenter__()
try:
await connect_mcp_servers(
{"test": MCPServerConfig(command="fake")},
registry,
stack,
)
finally:
stacks = await connect_mcp_servers(
{"test": MCPServerConfig(command="fake")},
registry,
)
for stack in stacks.values():
await stack.aclose()
assert "mcp_test_tool_a" in registry.tool_names