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

fix(mcp): reconnect terminated sessions

Browse Source
This commit is contained in:
chengyongru
2026-06-04 10:43:09 +08:00
committed by Xubin Ren
parent 7c3808327f
commit e9145b7acd
3 changed files with 393 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/agent/test_mcp_connection.py
+130
View File
@@ -4,14 +4,19 @@ from __future__ import annotations
import asyncio
from contextlib import AsyncExitStack
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock
import pytest
from mcp import types as mcp_types
from mcp.shared.exceptions import McpError
from mcp.types import ErrorData
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools import mcp as mcp_runtime
from nanobot.agent.tools.base import Tool
from nanobot.agent.tools.mcp import MCPResourceWrapper, MCPToolWrapper
from nanobot.bus.queue import MessageBus
from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import MCPServerConfig
@@ -218,3 +223,128 @@ async def test_reload_mcp_servers_retries_configured_server_without_live_stack(
assert result["retried"] == ["browserbase"]
assert loop.tools.has("mcp_browserbase_navigate")
await loop.close_mcp()
@pytest.mark.asyncio
async def test_mcp_tool_reconnects_after_session_terminated(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
):
loop = _make_loop(tmp_path, mcp_servers={"remote": object()})
closed: list[str] = []
sessions: list[Any] = []
connect_count = 0
async def _mark_closed(name: str) -> None:
closed.append(name)
class _FakeSession:
def __init__(self, index: int) -> None:
self.index = index
self.call_count = 0
async def call_tool(self, _name: str, arguments: dict[str, Any]) -> Any:
self.call_count += 1
assert arguments == {"symbol": "AAPL"}
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
session = _FakeSession(connect_count)
sessions.append(session)
tool_def = SimpleNamespace(
name="quote",
description="quote tool",
inputSchema={"type": "object", "properties": {}},
)
registry.register(MCPToolWrapper(session, name, tool_def, tool_timeout=5))
stack = AsyncExitStack()
await stack.__aenter__()
stack.push_async_callback(_mark_closed, name)
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(symbol="AAPL")
assert output == "recovered"
assert connect_count == 2
assert closed == ["remote"]
assert sessions[0].call_count == 1
assert sessions[1].call_count == 1
assert "remote" in loop._mcp_stacks
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,
monkeypatch: pytest.MonkeyPatch,
):
loop = _make_loop(tmp_path, mcp_servers={"remote": object()})
closed: list[str] = []
connect_count = 0
async def _mark_closed(name: str) -> None:
closed.append(name)
class _DeadSession:
async def read_resource(self, _uri: str) -> Any:
raise McpError(ErrorData(code=-32000, message="Session terminated"))
class _LiveSession:
async def read_resource(self, uri: str) -> Any:
await asyncio.sleep(0)
return SimpleNamespace(
contents=[
mcp_types.TextResourceContents(
uri=uri,
text=f"fresh:{uri.rsplit('/', maxsplit=1)[-1]}",
)
]
)
async def _fake_connect(servers, registry):
nonlocal connect_count
stacks = {}
for name in servers:
connect_count += 1
session = _DeadSession() if connect_count == 1 else _LiveSession()
for resource_name in ("alpha", "beta"):
resource_def = SimpleNamespace(
name=resource_name,
uri=f"file:///{resource_name}",
description=f"{resource_name} resource",
)
registry.register(MCPResourceWrapper(session, name, resource_def))
stack = AsyncExitStack()
await stack.__aenter__()
stack.push_async_callback(_mark_closed, name)
stacks[name] = stack
return stacks
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
await loop._connect_mcp()
old_alpha = loop.tools.get("mcp_remote_resource_alpha")
old_beta = loop.tools.get("mcp_remote_resource_beta")
assert isinstance(old_alpha, MCPResourceWrapper)
assert isinstance(old_beta, MCPResourceWrapper)
outputs = await asyncio.gather(old_alpha.execute(), old_beta.execute())
assert outputs == ["fresh:alpha", "fresh:beta"]
assert connect_count == 2
assert closed == ["remote"]
tests/agent/test_mcp_transient_retry.py
+98
View File
@@ -13,6 +13,7 @@ from nanobot.agent.tools.mcp import (
MCPPromptWrapper,
MCPResourceWrapper,
MCPToolWrapper,
_is_session_terminated,
_is_transient,
)
@@ -35,6 +36,14 @@ class _FakeEndOfStreamError(Exception):
_FakeEndOfStreamError.__name__ = "EndOfStream"
def _session_terminated_error() -> McpError:
return McpError(ErrorData(code=-32000, message="Session terminated"))
def _connection_closed_error() -> McpError:
return McpError(ErrorData(code=-32000, message="Connection closed"))
def test_is_transient_recognizes_closed_resource():
assert _is_transient(_FakeClosedResourceError("gone"))
@@ -67,6 +76,14 @@ def test_is_transient_rejects_timeout():
assert not _is_transient(TimeoutError("timeout"))
def test_is_session_terminated_recognizes_mcp_error():
assert _is_session_terminated(_session_terminated_error())
def test_is_session_terminated_recognizes_connection_closed_mcp_error():
assert _is_session_terminated(_connection_closed_error())
# ---------------------------------------------------------------------------
# MCPToolWrapper retry behaviour
# ---------------------------------------------------------------------------
@@ -219,6 +236,35 @@ async def test_tool_retry_on_end_of_stream():
assert session.call_tool.call_count == 2
@pytest.mark.asyncio
async def test_tool_reconnects_when_transient_retry_reveals_terminated_session():
"""Tool should reconnect if a stale session reports termination after transient retry."""
old_session = AsyncMock()
old_session.call_tool = AsyncMock(
side_effect=[_FakeClosedResourceError("closed"), _session_terminated_error()]
)
new_session = AsyncMock()
new_session.call_tool = AsyncMock(return_value=_make_tool_result("fresh"))
wrapper = MCPToolWrapper(old_session, "test_server", _make_tool_def(), tool_timeout=5)
replacement = MCPToolWrapper(new_session, "test_server", _make_tool_def(), tool_timeout=5)
async def reconnect(server_name: str, tool_name: str, stale_tool):
assert server_name == "test_server"
assert tool_name == "mcp_test_server_test_tool"
assert stale_tool is wrapper
return replacement
wrapper.set_reconnect_handler(reconnect)
with patch("nanobot.agent.tools.mcp.asyncio.sleep", new_callable=AsyncMock):
output = await wrapper.execute(foo="bar")
assert output == "fresh"
assert old_session.call_tool.call_count == 2
assert new_session.call_tool.call_count == 1
# ---------------------------------------------------------------------------
# MCPResourceWrapper retry behaviour
# ---------------------------------------------------------------------------
@@ -284,6 +330,32 @@ async def test_resource_no_retry_on_non_transient():
assert session.read_resource.call_count == 1
@pytest.mark.asyncio
async def test_resource_reconnects_on_session_terminated():
"""Resource should reconnect once when the MCP SDK reports a dead session."""
old_session = AsyncMock()
old_session.read_resource = AsyncMock(side_effect=_session_terminated_error())
new_session = AsyncMock()
new_session.read_resource = AsyncMock(return_value=_make_resource_result("fresh"))
wrapper = MCPResourceWrapper(old_session, "test_server", _make_resource_def())
replacement = MCPResourceWrapper(new_session, "test_server", _make_resource_def())
async def reconnect(server_name: str, tool_name: str, stale_tool):
assert server_name == "test_server"
assert tool_name == "mcp_test_server_resource_test_resource"
assert stale_tool is wrapper
return replacement
wrapper.set_reconnect_handler(reconnect)
output = await wrapper.execute()
assert output == "fresh"
assert old_session.read_resource.call_count == 1
assert new_session.read_resource.call_count == 1
# ---------------------------------------------------------------------------
# MCPPromptWrapper retry behaviour
# ---------------------------------------------------------------------------
@@ -366,3 +438,29 @@ async def test_prompt_no_retry_on_non_transient():
assert "RuntimeError" in output
assert session.get_prompt.call_count == 1
@pytest.mark.asyncio
async def test_prompt_reconnects_on_session_terminated():
"""Prompt should reconnect once before falling back to McpError handling."""
old_session = AsyncMock()
old_session.get_prompt = AsyncMock(side_effect=_session_terminated_error())
new_session = AsyncMock()
new_session.get_prompt = AsyncMock(return_value=_make_prompt_result("fresh prompt"))
wrapper = MCPPromptWrapper(old_session, "test_server", _make_prompt_def())
replacement = MCPPromptWrapper(new_session, "test_server", _make_prompt_def())
async def reconnect(server_name: str, tool_name: str, stale_tool):
assert server_name == "test_server"
assert tool_name == "mcp_test_server_prompt_test_prompt"
assert stale_tool is wrapper
return replacement
wrapper.set_reconnect_handler(reconnect)
output = await wrapper.execute()
assert output == "fresh prompt"
assert old_session.get_prompt.call_count == 1
assert new_session.get_prompt.call_count == 1