diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index fbbadfcd..a7f621b7 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -81,6 +81,7 @@ from nanobot.session.manager import ( replay_max_messages_for_context, ) from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator +from nanobot.utils.cancellation import task_is_cancelling from nanobot.utils.document import extract_documents, reference_non_image_attachments from nanobot.utils.helpers import image_placeholder_text from nanobot.utils.helpers import truncate_text as truncate_text_fn @@ -998,8 +999,11 @@ class AgentLoop: except asyncio.CancelledError: # Preserve real task cancellation so shutdown can complete cleanly. # Only ignore non-task CancelledError signals that may leak from integrations. - if not self._running or asyncio.current_task().cancelling(): + if not self._running or task_is_cancelling(): raise + logger.warning( + "Ignoring leaked CancelledError while consuming inbound messages" + ) continue except Exception as e: logger.warning("Error consuming inbound message: {}, continuing...", e) diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index 2d17be59..3cf1147c 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -30,6 +30,7 @@ from nanobot.security.network import ( resolve_url_target, validate_url_target, ) +from nanobot.utils.cancellation import task_is_cancelling # Transient connection errors that warrant a single retry. # These typically happen when an MCP server restarts or a network @@ -487,8 +488,7 @@ class MCPToolWrapper(_MCPWrapperBase): 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: + if task_is_cancelling(): raise logger.warning("MCP tool '{}' was cancelled by server/SDK", self._name) return ToolResult.error("(MCP tool call was cancelled)") @@ -650,8 +650,7 @@ class MCPResourceWrapper(_MCPWrapperBase): ) 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: + if task_is_cancelling(): raise logger.warning("MCP resource '{}' was cancelled by server/SDK", self._name) return "(MCP resource read was cancelled)" @@ -764,8 +763,7 @@ class MCPPromptWrapper(_MCPWrapperBase): ) 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: + if task_is_cancelling(): raise logger.warning("MCP prompt '{}' was cancelled by server/SDK", self._name) return "(MCP prompt call was cancelled)" @@ -1145,6 +1143,8 @@ async def connect_missing_servers(state: Any, registry: ToolRegistry) -> None: else: logger.warning("No MCP servers connected successfully (will retry next message)") except asyncio.CancelledError: + if task_is_cancelling(): + raise logger.warning("MCP connection cancelled (will retry next message)") except BaseException as e: logger.warning("Failed to connect MCP servers (will retry next message): {}", e) @@ -1411,7 +1411,7 @@ async def _close_server(state: Any, server_name: str) -> None: try: await stack.aclose() except asyncio.CancelledError: - if asyncio.current_task().cancelling() > 0: + if task_is_cancelling(): raise logger.debug("MCP server '{}' cleanup error (can be ignored)", server_name) except (RuntimeError, BaseExceptionGroup): @@ -1428,7 +1428,7 @@ async def close_mcp_servers(state: Any) -> None: try: await connection.aclose() except asyncio.CancelledError: - if asyncio.current_task().cancelling() > 0: + if task_is_cancelling(): raise logger.debug("MCP server '{}' cleanup error (can be ignored)", name) except (RuntimeError, BaseExceptionGroup): diff --git a/nanobot/utils/cancellation.py b/nanobot/utils/cancellation.py new file mode 100644 index 00000000..29f68948 --- /dev/null +++ b/nanobot/utils/cancellation.py @@ -0,0 +1,10 @@ +"""Async cancellation helpers.""" + +from __future__ import annotations + +import asyncio + + +def task_is_cancelling() -> bool: + task = asyncio.current_task() + return task is not None and task.cancelling() > 0 diff --git a/tests/agent/test_task_cancel.py b/tests/agent/test_task_cancel.py index 6f5ff982..1b3fe97b 100644 --- a/tests/agent/test_task_cancel.py +++ b/tests/agent/test_task_cancel.py @@ -111,6 +111,34 @@ class TestHandleStop: class TestDispatch: + @pytest.mark.asyncio + async def test_run_logs_and_continues_after_leaked_cancelled_error(self, monkeypatch): + loop, bus = _make_loop() + loop._connect_mcp = AsyncMock() + loop.close_mcp = AsyncMock() + loop.auto_compact.check_expired = MagicMock() + warnings: list[str] = [] + calls = 0 + + async def consume_once_then_stop(): + nonlocal calls + calls += 1 + if calls == 1: + raise asyncio.CancelledError() + loop.stop() + raise asyncio.TimeoutError() + + monkeypatch.setattr(bus, "consume_inbound", consume_once_then_stop) + monkeypatch.setattr( + "nanobot.agent.loop.logger.warning", + lambda message, *args, **kwargs: warnings.append(message), + ) + + await loop.run() + + assert calls == 2 + assert any("Ignoring leaked CancelledError" in warning for warning in warnings) + def test_exec_tool_not_registered_when_disabled(self): from nanobot.agent.tools.shell import ExecToolConfig from nanobot.config.schema import ToolsConfig diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index ccc2abeb..8bf932a7 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -144,6 +144,35 @@ def _make_wrapper(session: object, *, timeout: float = 0.1) -> MCPToolWrapper: return MCPToolWrapper(session, "test", tool_def, tool_timeout=timeout) +@pytest.mark.asyncio +async def test_connect_missing_servers_propagates_external_cancellation(monkeypatch) -> None: + started = asyncio.Event() + + async def connect_mcp_servers(_servers: dict, _registry: ToolRegistry) -> dict: + started.set() + await asyncio.sleep(60) + return {} + + class State: + pass + + state = State() + state._mcp_closing = False + state._mcp_servers = {"test": MCPServerConfig(command="fake")} + state._mcp_stacks = {} + state._mcp_connecting = False + monkeypatch.setattr(mcp_mod, "connect_mcp_servers", connect_mcp_servers) + + task = asyncio.create_task(mcp_mod.connect_missing_servers(state, ToolRegistry())) + await asyncio.wait_for(started.wait(), timeout=1.0) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert state._mcp_connecting is False + + def test_wrapper_preserves_non_nullable_unions() -> None: tool_def = SimpleNamespace( name="demo",