fix(mcp): reconnect terminated sessions

tests/agent/test_mcp_connection.py

@@ -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"]