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

fix: preserve real cancellation in MCP paths

Browse Source
This commit is contained in:
adabarbulescu
2026-07-18 17:36:35 +08:00
committed by Xubin Ren
parent 07ad0bafa8
commit d4f5abe004
5 changed files with 80 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/tools/mcp.py
+8 -8
View File
@@ -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):