diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index 2aea1927..ab95b446 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -10,6 +10,25 @@ from loguru import logger from nanobot.agent.tools.base import Tool from nanobot.agent.tools.registry import ToolRegistry +# Transient connection errors that warrant a single retry. +# These typically happen when an MCP server restarts or a network +# connection is interrupted between calls. +_TRANSIENT_EXC_NAMES: frozenset[str] = frozenset(( + "ClosedResourceError", + "BrokenResourceError", + "EndOfStream", + "BrokenPipeError", + "ConnectionResetError", + "ConnectionRefusedError", + "ConnectionAbortedError", + "ConnectionError", +)) + + +def _is_transient(exc: BaseException) -> bool: + """Check if an exception looks like a transient connection error.""" + return type(exc).__name__ in _TRANSIENT_EXC_NAMES + def _extract_nullable_branch(options: Any) -> tuple[dict[str, Any], bool] | None: """Return the single non-null branch for nullable unions.""" @@ -99,38 +118,61 @@ class MCPToolWrapper(Tool): async def execute(self, **kwargs: Any) -> str: from mcp import types - try: - result = await asyncio.wait_for( - self._session.call_tool(self._original_name, arguments=kwargs), - timeout=self._tool_timeout, - ) - except asyncio.TimeoutError: - logger.warning("MCP tool '{}' timed out after {}s", self._name, self._tool_timeout) - return f"(MCP tool call timed out after {self._tool_timeout}s)" - except asyncio.CancelledError: - # MCP SDK's anyio cancel scopes can leak CancelledError on timeout/failure. - # Re-raise only if our task was externally cancelled (e.g. /stop). - task = asyncio.current_task() - if task is not None and task.cancelling() > 0: - raise - logger.warning("MCP tool '{}' was cancelled by server/SDK", self._name) - return "(MCP tool call was cancelled)" - except Exception as exc: - logger.exception( - "MCP tool '{}' failed: {}: {}", - self._name, - type(exc).__name__, - exc, - ) - return f"(MCP tool call failed: {type(exc).__name__})" - - parts = [] - for block in result.content: - if isinstance(block, types.TextContent): - parts.append(block.text) + for attempt in range(2): # At most 1 retry + try: + result = await asyncio.wait_for( + self._session.call_tool(self._original_name, arguments=kwargs), + timeout=self._tool_timeout, + ) + except asyncio.TimeoutError: + logger.warning( + "MCP tool '{}' timed out after {}s", self._name, self._tool_timeout + ) + return f"(MCP tool call timed out after {self._tool_timeout}s)" + except asyncio.CancelledError: + # MCP SDK's anyio cancel scopes can leak CancelledError on timeout/failure. + # Re-raise only if our task was externally cancelled (e.g. /stop). + task = asyncio.current_task() + if task is not None and task.cancelling() > 0: + raise + logger.warning("MCP tool '{}' was cancelled by server/SDK", self._name) + return "(MCP tool call was cancelled)" + except Exception as exc: + if _is_transient(exc): + if attempt == 0: + logger.warning( + "MCP tool '{}' hit transient error ({}), retrying once...", + self._name, + type(exc).__name__, + ) + await asyncio.sleep(1) # Brief backoff before retry + continue + # Second transient failure — give up with retry-specific message + logger.error( + "MCP tool '{}' failed after retry: {}: {}", + self._name, + type(exc).__name__, + exc, + ) + return f"(MCP tool call failed after retry: {type(exc).__name__})" + logger.exception( + "MCP tool '{}' failed: {}: {}", + self._name, + type(exc).__name__, + exc, + ) + return f"(MCP tool call failed: {type(exc).__name__})" else: - parts.append(str(block)) - return "\n".join(parts) or "(no output)" + # Success — extract result + parts = [] + for block in result.content: + if isinstance(block, types.TextContent): + parts.append(block.text) + else: + parts.append(str(block)) + return "\n".join(parts) or "(no output)" + + return "(MCP tool call failed)" # Unreachable, but satisfies type checkers class MCPResourceWrapper(Tool): @@ -168,40 +210,59 @@ class MCPResourceWrapper(Tool): async def execute(self, **kwargs: Any) -> str: from mcp import types - try: - result = await asyncio.wait_for( - self._session.read_resource(self._uri), - timeout=self._resource_timeout, - ) - except asyncio.TimeoutError: - logger.warning( - "MCP resource '{}' timed out after {}s", self._name, self._resource_timeout - ) - return f"(MCP resource read timed out after {self._resource_timeout}s)" - except asyncio.CancelledError: - task = asyncio.current_task() - if task is not None and task.cancelling() > 0: - raise - logger.warning("MCP resource '{}' was cancelled by server/SDK", self._name) - return "(MCP resource read was cancelled)" - except Exception as exc: - logger.exception( - "MCP resource '{}' failed: {}: {}", - self._name, - type(exc).__name__, - exc, - ) - return f"(MCP resource read failed: {type(exc).__name__})" - - parts: list[str] = [] - for block in result.contents: - if isinstance(block, types.TextResourceContents): - parts.append(block.text) - elif isinstance(block, types.BlobResourceContents): - parts.append(f"[Binary resource: {len(block.blob)} bytes]") + for attempt in range(2): + try: + result = await asyncio.wait_for( + self._session.read_resource(self._uri), + timeout=self._resource_timeout, + ) + except asyncio.TimeoutError: + logger.warning( + "MCP resource '{}' timed out after {}s", self._name, self._resource_timeout + ) + return f"(MCP resource read timed out after {self._resource_timeout}s)" + except asyncio.CancelledError: + task = asyncio.current_task() + if task is not None and task.cancelling() > 0: + raise + logger.warning("MCP resource '{}' was cancelled by server/SDK", self._name) + return "(MCP resource read was cancelled)" + except Exception as exc: + if _is_transient(exc): + if attempt == 0: + logger.warning( + "MCP resource '{}' hit transient error ({}), retrying once...", + self._name, + type(exc).__name__, + ) + await asyncio.sleep(1) + continue + logger.error( + "MCP resource '{}' failed after retry: {}: {}", + self._name, + type(exc).__name__, + exc, + ) + return f"(MCP resource read failed after retry: {type(exc).__name__})" + logger.exception( + "MCP resource '{}' failed: {}: {}", + self._name, + type(exc).__name__, + exc, + ) + return f"(MCP resource read failed: {type(exc).__name__})" else: - parts.append(str(block)) - return "\n".join(parts) or "(no output)" + parts: list[str] = [] + for block in result.contents: + if isinstance(block, types.TextResourceContents): + parts.append(block.text) + elif isinstance(block, types.BlobResourceContents): + parts.append(f"[Binary resource: {len(block.blob)} bytes]") + else: + parts.append(str(block)) + return "\n".join(parts) or "(no output)" + + return "(MCP resource read failed)" # Unreachable class MCPPromptWrapper(Tool): @@ -254,52 +315,72 @@ class MCPPromptWrapper(Tool): from mcp import types from mcp.shared.exceptions import McpError - try: - result = await asyncio.wait_for( - self._session.get_prompt(self._prompt_name, arguments=kwargs), - timeout=self._prompt_timeout, - ) - except asyncio.TimeoutError: - 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() - if task is not None and task.cancelling() > 0: - raise - logger.warning("MCP prompt '{}' was cancelled by server/SDK", self._name) - return "(MCP prompt call was cancelled)" - except McpError as exc: - logger.error( - "MCP prompt '{}' failed: code={} 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, - ) - return f"(MCP prompt call failed: {type(exc).__name__})" - - parts: list[str] = [] - for message in result.messages: - content = message.content - # content is a single ContentBlock (not a list) in MCP SDK >= 1.x - if isinstance(content, types.TextContent): - parts.append(content.text) - elif isinstance(content, list): - for block in content: - if isinstance(block, types.TextContent): - parts.append(block.text) - else: - parts.append(str(block)) + for attempt in range(2): + try: + result = await asyncio.wait_for( + self._session.get_prompt(self._prompt_name, arguments=kwargs), + timeout=self._prompt_timeout, + ) + except asyncio.TimeoutError: + 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() + if task is not None and task.cancelling() > 0: + raise + logger.warning("MCP prompt '{}' was cancelled by server/SDK", self._name) + return "(MCP prompt call was cancelled)" + except McpError as exc: + logger.error( + "MCP prompt '{}' failed: code={} 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: + if _is_transient(exc): + if attempt == 0: + logger.warning( + "MCP prompt '{}' hit transient error ({}), retrying once...", + self._name, + type(exc).__name__, + ) + await asyncio.sleep(1) + continue + logger.error( + "MCP prompt '{}' failed after retry: {}: {}", + self._name, + type(exc).__name__, + exc, + ) + return f"(MCP prompt call failed after retry: {type(exc).__name__})" + logger.exception( + "MCP prompt '{}' failed: {}: {}", + self._name, + type(exc).__name__, + exc, + ) + return f"(MCP prompt call failed: {type(exc).__name__})" else: - parts.append(str(content)) - return "\n".join(parts) or "(no output)" + parts: list[str] = [] + for message in result.messages: + content = message.content + if isinstance(content, types.TextContent): + parts.append(content.text) + elif isinstance(content, list): + for block in content: + if isinstance(block, types.TextContent): + parts.append(block.text) + else: + parts.append(str(block)) + else: + parts.append(str(content)) + return "\n".join(parts) or "(no output)" + + return "(MCP prompt call failed)" # Unreachable async def connect_mcp_servers( diff --git a/tests/agent/test_mcp_transient_retry.py b/tests/agent/test_mcp_transient_retry.py new file mode 100644 index 00000000..823bb1e8 --- /dev/null +++ b/tests/agent/test_mcp_transient_retry.py @@ -0,0 +1,344 @@ +"""Tests for MCP tool/resource/prompt transient error retry.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest +from mcp import types as mcp_types +from mcp.shared.exceptions import McpError +from mcp.types import ErrorData + +from nanobot.agent.tools.mcp import ( + MCPPromptWrapper, + MCPResourceWrapper, + MCPToolWrapper, + _is_transient, +) + +# --------------------------------------------------------------------------- +# _is_transient helper +# --------------------------------------------------------------------------- + + +class _FakeClosedResourceError(Exception): + pass + + +_FakeClosedResourceError.__name__ = "ClosedResourceError" + + +class _FakeEndOfStreamError(Exception): + pass + + +_FakeEndOfStreamError.__name__ = "EndOfStream" + + +def test_is_transient_recognizes_closed_resource(): + assert _is_transient(_FakeClosedResourceError("gone")) + + +def test_is_transient_recognizes_broken_pipe(): + assert _is_transient(BrokenPipeError("pipe")) + + +def test_is_transient_recognizes_connection_reset(): + assert _is_transient(ConnectionResetError("reset")) + + +def test_is_transient_recognizes_connection_refused(): + assert _is_transient(ConnectionRefusedError("refused")) + + +def test_is_transient_recognizes_end_of_stream(): + assert _is_transient(_FakeEndOfStreamError("eof")) + + +def test_is_transient_rejects_value_error(): + assert not _is_transient(ValueError("nope")) + + +def test_is_transient_rejects_runtime_error(): + assert not _is_transient(RuntimeError("nope")) + + +def test_is_transient_rejects_timeout(): + assert not _is_transient(TimeoutError("timeout")) + + +# --------------------------------------------------------------------------- +# MCPToolWrapper retry behaviour +# --------------------------------------------------------------------------- + + +def _make_tool_def(name="test_tool"): + return SimpleNamespace( + name=name, + description="A test tool", + inputSchema={"type": "object", "properties": {}}, + ) + + +def _make_tool_result(text): + """Build a mock tool result with proper MCP TextContent.""" + return SimpleNamespace(content=[mcp_types.TextContent(type="text", text=text)]) + + +@pytest.mark.asyncio +async def test_tool_retries_on_transient_error(): + """Tool should retry once when a transient error occurs, then succeed.""" + session = AsyncMock() + result = _make_tool_result("ok") + exc = _FakeClosedResourceError("connection lost") + session.call_tool = AsyncMock(side_effect=[exc, result]) + + wrapper = MCPToolWrapper(session, "test_server", _make_tool_def(), tool_timeout=5) + + with patch("nanobot.agent.tools.mcp.asyncio.sleep", new_callable=AsyncMock): + output = await wrapper.execute(foo="bar") + + assert output == "ok" + assert session.call_tool.call_count == 2 + + +@pytest.mark.asyncio +async def test_tool_fails_after_retry_exhausted(): + """Tool should fail with retry message when both attempts hit transient errors.""" + session = AsyncMock() + exc1 = _FakeClosedResourceError("still dead") + exc2 = _FakeClosedResourceError("still dead again") + session.call_tool = AsyncMock(side_effect=[exc1, exc2]) + + wrapper = MCPToolWrapper(session, "test_server", _make_tool_def(), tool_timeout=5) + + with patch("nanobot.agent.tools.mcp.asyncio.sleep", new_callable=AsyncMock): + output = await wrapper.execute() + + assert "failed after retry" in output + assert "ClosedResourceError" in output + assert session.call_tool.call_count == 2 + + +@pytest.mark.asyncio +async def test_tool_no_retry_on_non_transient_error(): + """Tool should NOT retry on non-transient errors like ValueError.""" + session = AsyncMock() + session.call_tool = AsyncMock(side_effect=ValueError("bad input")) + + wrapper = MCPToolWrapper(session, "test_server", _make_tool_def(), tool_timeout=5) + output = await wrapper.execute() + + assert "ValueError" in output + assert "retry" not in output + assert session.call_tool.call_count == 1 + + +@pytest.mark.asyncio +async def test_tool_no_retry_on_timeout(): + """Timeouts should not trigger retry (they have their own handling).""" + session = AsyncMock() + session.call_tool = AsyncMock(side_effect=asyncio.TimeoutError()) + + wrapper = MCPToolWrapper(session, "test_server", _make_tool_def(), tool_timeout=5) + output = await wrapper.execute() + + assert "timed out" in output + assert session.call_tool.call_count == 1 + + +@pytest.mark.asyncio +async def test_tool_success_on_first_try_no_retry(): + """Normal success path — no retry logic involved.""" + session = AsyncMock() + result = _make_tool_result("hello") + session.call_tool = AsyncMock(return_value=result) + + wrapper = MCPToolWrapper(session, "test_server", _make_tool_def(), tool_timeout=5) + output = await wrapper.execute() + + assert output == "hello" + assert session.call_tool.call_count == 1 + + +@pytest.mark.asyncio +async def test_tool_retry_on_connection_reset(): + """ConnectionResetError (a stdlib exception) should also trigger retry.""" + session = AsyncMock() + result = _make_tool_result("recovered") + session.call_tool = AsyncMock( + side_effect=[ConnectionResetError("reset by peer"), result] + ) + + wrapper = MCPToolWrapper(session, "test_server", _make_tool_def(), tool_timeout=5) + + with patch("nanobot.agent.tools.mcp.asyncio.sleep", new_callable=AsyncMock): + output = await wrapper.execute() + + assert output == "recovered" + assert session.call_tool.call_count == 2 + + +@pytest.mark.asyncio +async def test_tool_retry_on_end_of_stream(): + """EndOfStream (anyio) should trigger retry.""" + session = AsyncMock() + result = _make_tool_result("back") + session.call_tool = AsyncMock(side_effect=[_FakeEndOfStreamError("eof"), result]) + + wrapper = MCPToolWrapper(session, "test_server", _make_tool_def(), tool_timeout=5) + + with patch("nanobot.agent.tools.mcp.asyncio.sleep", new_callable=AsyncMock): + output = await wrapper.execute() + + assert output == "back" + assert session.call_tool.call_count == 2 + + +# --------------------------------------------------------------------------- +# MCPResourceWrapper retry behaviour +# --------------------------------------------------------------------------- + + +def _make_resource_def(name="test_resource"): + return SimpleNamespace( + name=name, + uri="file:///test", + description="A test resource", + ) + + +def _make_resource_result(text): + return SimpleNamespace( + contents=[mcp_types.TextResourceContents(uri="file:///test", text=text)] + ) + + +@pytest.mark.asyncio +async def test_resource_retries_on_transient_error(): + """Resource should retry once on transient connection error.""" + session = AsyncMock() + result = _make_resource_result("data") + exc = _FakeClosedResourceError("gone") + session.read_resource = AsyncMock(side_effect=[exc, result]) + + wrapper = MCPResourceWrapper(session, "test_server", _make_resource_def()) + + with patch("nanobot.agent.tools.mcp.asyncio.sleep", new_callable=AsyncMock): + output = await wrapper.execute() + + assert output == "data" + assert session.read_resource.call_count == 2 + + +@pytest.mark.asyncio +async def test_resource_fails_after_retry_exhausted(): + """Resource should fail with retry message when both attempts fail.""" + session = AsyncMock() + exc = _FakeClosedResourceError("dead") + session.read_resource = AsyncMock(side_effect=[exc, exc]) + + wrapper = MCPResourceWrapper(session, "test_server", _make_resource_def()) + + with patch("nanobot.agent.tools.mcp.asyncio.sleep", new_callable=AsyncMock): + output = await wrapper.execute() + + assert "failed after retry" in output + assert session.read_resource.call_count == 2 + + +@pytest.mark.asyncio +async def test_resource_no_retry_on_non_transient(): + """Resource should not retry on non-transient errors.""" + session = AsyncMock() + session.read_resource = AsyncMock(side_effect=RuntimeError("bad")) + + wrapper = MCPResourceWrapper(session, "test_server", _make_resource_def()) + output = await wrapper.execute() + + assert "RuntimeError" in output + assert session.read_resource.call_count == 1 + + +# --------------------------------------------------------------------------- +# MCPPromptWrapper retry behaviour +# --------------------------------------------------------------------------- + + +def _make_prompt_def(name="test_prompt"): + return SimpleNamespace( + name=name, + description="A test prompt", + arguments=[], + ) + + +def _make_prompt_result(text): + return SimpleNamespace( + messages=[ + SimpleNamespace( + content=mcp_types.TextContent(type="text", text=text), + ) + ] + ) + + +@pytest.mark.asyncio +async def test_prompt_retries_on_transient_error(): + """Prompt should retry once on transient connection error.""" + session = AsyncMock() + result = _make_prompt_result("prompt text") + exc = _FakeClosedResourceError("gone") + session.get_prompt = AsyncMock(side_effect=[exc, result]) + + wrapper = MCPPromptWrapper(session, "test_server", _make_prompt_def()) + + with patch("nanobot.agent.tools.mcp.asyncio.sleep", new_callable=AsyncMock): + output = await wrapper.execute() + + assert output == "prompt text" + assert session.get_prompt.call_count == 2 + + +@pytest.mark.asyncio +async def test_prompt_fails_after_retry_exhausted(): + """Prompt should fail with retry message when both attempts fail.""" + session = AsyncMock() + exc = _FakeClosedResourceError("dead") + session.get_prompt = AsyncMock(side_effect=[exc, exc]) + + wrapper = MCPPromptWrapper(session, "test_server", _make_prompt_def()) + + with patch("nanobot.agent.tools.mcp.asyncio.sleep", new_callable=AsyncMock): + output = await wrapper.execute() + + assert "failed after retry" in output + assert session.get_prompt.call_count == 2 + + +@pytest.mark.asyncio +async def test_prompt_no_retry_on_mcp_error(): + """McpError (application-level) should NOT trigger retry.""" + session = AsyncMock() + session.get_prompt = AsyncMock( + side_effect=McpError(ErrorData(code=-1, message="not found")) + ) + + wrapper = MCPPromptWrapper(session, "test_server", _make_prompt_def()) + output = await wrapper.execute() + + assert "not found" in output + assert session.get_prompt.call_count == 1 + + +@pytest.mark.asyncio +async def test_prompt_no_retry_on_non_transient(): + """Non-transient errors should not trigger retry for prompts.""" + session = AsyncMock() + session.get_prompt = AsyncMock(side_effect=RuntimeError("bad")) + + wrapper = MCPPromptWrapper(session, "test_server", _make_prompt_def()) + output = await wrapper.execute() + + assert "RuntimeError" in output + assert session.get_prompt.call_count == 1