fix: preserve real cancellation in MCP paths
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user