fix: cover MCP reconnect edge cases

maintainer edit: handle prompt sessions that report Connection closed outside McpError, and match reconnect registration prefixes with the same sanitization used by MCP wrapper names.
This commit is contained in:
chengyongru
2026-06-04 10:43:09 +08:00
committed by Xubin Ren
parent e9145b7acd
commit d0eba7cd9d
3 changed files with 83 additions and 4 deletions
+49
View File
@@ -288,6 +288,55 @@ async def test_mcp_tool_reconnects_after_session_terminated(
assert loop.tools.get("mcp_remote_quote") is not old_tool
@pytest.mark.asyncio
async def test_mcp_reconnect_handler_uses_sanitized_server_prefix(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
):
loop = _make_loop(tmp_path, mcp_servers={"remote_": object()})
connect_count = 0
class _FakeSession:
def __init__(self, index: int) -> None:
self.index = index
async def call_tool(self, _name: str, arguments: dict[str, Any]) -> Any:
assert arguments == {}
if self.index == 1:
raise McpError(ErrorData(code=-32000, message="Session terminated"))
return SimpleNamespace(
content=[mcp_types.TextContent(type="text", text="recovered")]
)
async def _fake_connect(servers, registry):
nonlocal connect_count
stacks = {}
for name in servers:
connect_count += 1
tool_def = SimpleNamespace(
name="quote",
description="quote tool",
inputSchema={"type": "object", "properties": {}},
)
registry.register(MCPToolWrapper(_FakeSession(connect_count), name, tool_def))
stack = AsyncExitStack()
await stack.__aenter__()
stacks[name] = stack
return stacks
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
await loop._connect_mcp()
old_tool = loop.tools.get("mcp_remote_quote")
assert isinstance(old_tool, MCPToolWrapper)
output = await old_tool.execute()
assert output == "recovered"
assert connect_count == 2
assert loop.tools.get("mcp_remote_quote") is not old_tool
@pytest.mark.asyncio
async def test_concurrent_mcp_reconnect_reuses_fresh_session(
tmp_path,