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

fix(mcp): keep transport cleanup in owner tasks

Browse Source
This commit is contained in:
Xubin Ren
2026-07-11 11:45:20 +08:00
parent bd0dd85f44
commit edf78e7054
5 changed files with 187 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/context.py
+4
View File
@@ -44,6 +44,10 @@ async def connect_mcp(state: Any, tools: ToolRegistry) -> None:
await mcp_tools.connect_missing_servers(state, tools)
async def close_mcp(state: Any) -> None:
await mcp_tools.close_mcp_servers(state)
async def handle_runtime_control(state: Any, msg: InboundMessage, tools: ToolRegistry) -> bool:
return await mcp_tools.handle_runtime_control(state, msg, tools)
nanobot/agent/loop.py
+4 -15
View File
@@ -7,7 +7,7 @@ import dataclasses
import os
import time
from collections.abc import Mapping
from contextlib import AsyncExitStack, nullcontext, suppress
from contextlib import nullcontext, suppress
from dataclasses import dataclass, field
from enum import Enum, auto
from functools import partial
@@ -82,6 +82,7 @@ from nanobot.utils.runtime import (
)
if TYPE_CHECKING:
from nanobot.agent.tools.mcp import MCPConnection
from nanobot.config.schema import (
ChannelsConfig,
ProviderConfig,
@@ -360,8 +361,7 @@ class AgentLoop:
self._unified_session = unified_session
self._running = False
self._mcp_servers = mcp_servers or {}
self._mcp_stacks: dict[str, AsyncExitStack] = {}
self._mcp_retired_stacks: list[tuple[str, AsyncExitStack]] = []
self._mcp_stacks: dict[str, MCPConnection] = {}
self._mcp_connecting = False
self._active_tasks: dict[str, list[asyncio.Task]] = {} # session_key -> tasks
self._background_tasks: list[asyncio.Task] = []
@@ -1178,18 +1178,7 @@ class AgentLoop:
if self._background_tasks:
await asyncio.gather(*self._background_tasks, return_exceptions=True)
self._background_tasks.clear()
stacks = [*self._mcp_retired_stacks, *self._mcp_stacks.items()]
self._mcp_retired_stacks.clear()
for name, stack in stacks:
try:
await stack.aclose()
except asyncio.CancelledError as exc:
if not str(exc).startswith("Cancelled via cancel scope"):
raise
logger.debug("MCP server '{}' cleanup cancelled by SDK (can be ignored)", name)
except (RuntimeError, BaseExceptionGroup):
logger.debug("MCP server '{}' cleanup error (can be ignored)", name)
self._mcp_stacks.clear()
await agent_context.close_mcp(self)
def _schedule_background(self, coro) -> None:
"""Schedule a coroutine as a tracked background task (drained on shutdown)."""
nanobot/agent/tools/mcp.py
+131 -41
View File
@@ -9,7 +9,7 @@ import shutil
import urllib.parse
from collections.abc import Awaitable, Callable
from contextlib import AsyncExitStack, suppress
from typing import Any, Mapping
from typing import Any, Mapping, Protocol
from weakref import WeakKeyDictionary
import httpx
@@ -54,6 +54,26 @@ _RELOAD_LOCKS: WeakKeyDictionary[Any, asyncio.Lock] = WeakKeyDictionary()
_ReconnectCallback = Callable[[str, str, Tool], Awaitable[Tool | None]]
class MCPConnection(Protocol):
async def aclose(self) -> None: ...
class _OwnedMCPConnection:
"""Close an MCP transport from the task that originally opened it."""
def __init__(self, owner: asyncio.Task[None], close_requested: asyncio.Event) -> None:
self._owner = owner
self._close_requested = close_requested
async def aclose(self) -> None:
self._close_requested.set()
try:
await asyncio.shield(self._owner)
except asyncio.CancelledError:
if not self._owner.cancelled():
raise
def _is_malformed_mcp_progress_notification(message: Any) -> bool:
payload = _mcp_jsonrpc_payload(message)
if _payload_value(payload, "method") != "notifications/progress":
@@ -814,19 +834,19 @@ class MCPPromptWrapper(_MCPWrapperBase):
async def connect_mcp_servers(
mcp_servers: dict, registry: ToolRegistry
) -> dict[str, AsyncExitStack]:
) -> dict[str, MCPConnection]:
"""Connect to configured MCP servers and register their tools, resources, prompts.
Returns a dict mapping server name -> its dedicated AsyncExitStack.
Each server gets its own stack to prevent cancel scope conflicts
when multiple MCP servers are configured.
Returns one connection handle per server. Each handle keeps the task that
entered the MCP SDK contexts alive so reconnect and shutdown can close
AnyIO cancel scopes from their owning task.
"""
from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamable_http_client
async def connect_single_server(name: str, cfg) -> tuple[str, AsyncExitStack | None]:
async def open_single_server(name: str, cfg) -> tuple[str, AsyncExitStack | None]:
server_stack = AsyncExitStack()
await server_stack.__aenter__()
@@ -1045,7 +1065,43 @@ async def connect_mcp_servers(
await server_stack.aclose()
return name, None
server_stacks: dict[str, AsyncExitStack] = {}
async def connect_single_server(name: str, cfg) -> tuple[str, MCPConnection | None]:
loop = asyncio.get_running_loop()
ready: asyncio.Future[bool] = loop.create_future()
close_requested = asyncio.Event()
async def own_connection() -> None:
stack: AsyncExitStack | None = None
try:
_, stack = await open_single_server(name, cfg)
if not ready.done():
ready.set_result(stack is not None)
if stack is not None:
await close_requested.wait()
except BaseException as exc:
if not ready.done():
ready.set_exception(exc)
raise
finally:
if stack is not None:
await stack.aclose()
owner = asyncio.create_task(own_connection(), name=f"mcp:{name}")
connection = _OwnedMCPConnection(owner, close_requested)
try:
connected = await ready
except BaseException:
close_requested.set()
owner.cancel()
with suppress(BaseException):
await asyncio.shield(owner)
raise
if not connected:
await connection.aclose()
return name, None
return name, connection
server_stacks: dict[str, MCPConnection] = {}
for name, cfg in mcp_servers.items():
try:
@@ -1124,31 +1180,44 @@ def runtime_lines(
async def connect_missing_servers(state: Any, registry: ToolRegistry) -> None:
"""Connect configured MCP servers that are not currently live."""
missing_servers = {
name: cfg for name, cfg in state._mcp_servers.items() if name not in state._mcp_stacks
}
if state._mcp_connecting or not missing_servers:
return
state._mcp_connecting = True
try:
connected = await connect_mcp_servers(missing_servers, registry)
state._mcp_stacks.update(connected)
_attach_reconnect_handlers(state, registry, connected)
if connected:
logger.info("MCP connected servers: {}", sorted(connected))
else:
logger.warning("No MCP servers connected successfully (will retry next message)")
except asyncio.CancelledError:
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)
finally:
state._mcp_connecting = False
async with _reload_lock(state):
if getattr(state, "_mcp_closing", False):
return
missing_servers = {
name: cfg for name, cfg in state._mcp_servers.items() if name not in state._mcp_stacks
}
if state._mcp_connecting or not missing_servers:
return
state._mcp_connecting = True
try:
connected = await connect_mcp_servers(missing_servers, registry)
if getattr(state, "_mcp_closing", False):
for connection in connected.values():
await connection.aclose()
return
state._mcp_stacks.update(connected)
_attach_reconnect_handlers(state, registry, connected)
if connected:
logger.info("MCP connected servers: {}", sorted(connected))
else:
logger.warning("No MCP servers connected successfully (will retry next message)")
except asyncio.CancelledError:
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)
finally:
state._mcp_connecting = False
async def reload_servers(state: Any, registry: ToolRegistry) -> dict[str, Any]:
"""Reconcile live MCP connections with the current config file."""
async with _reload_lock(state):
if getattr(state, "_mcp_closing", False):
return {
"ok": False,
"message": "MCP connections are shutting down.",
"requires_restart": True,
}
try:
from nanobot.config.loader import load_config, resolve_config_env_vars
@@ -1177,7 +1246,7 @@ async def reload_servers(state: Any, registry: ToolRegistry) -> dict[str, Any]:
tools_removed = 0
for name in [*removed, *changed]:
tools_removed += _unregister_server_tools(state, registry, name)
_retire_server_stack(state, name)
await _close_server(state, name)
state._mcp_servers = next_servers
retry_missing = sorted(
@@ -1187,9 +1256,17 @@ async def reload_servers(state: Any, registry: ToolRegistry) -> dict[str, Any]:
)
to_connect_names = sorted(set(added) | set(changed) | set(retry_missing))
to_connect = {name: next_servers[name] for name in to_connect_names}
connected: dict[str, AsyncExitStack] = {}
connected: dict[str, MCPConnection] = {}
if to_connect:
connected = await connect_mcp_servers(to_connect, registry)
if getattr(state, "_mcp_closing", False):
for connection in connected.values():
await connection.aclose()
return {
"ok": False,
"message": "MCP connections are shutting down.",
"requires_restart": True,
}
state._mcp_stacks.update(connected)
_attach_reconnect_handlers(state, registry, connected)
@@ -1323,6 +1400,8 @@ async def _refresh_terminated_server(
stale_tool: Tool,
) -> Tool | None:
async with _reload_lock(state):
if getattr(state, "_mcp_closing", False):
return None
cfg = state._mcp_servers.get(server_name)
if cfg is None:
logger.warning(
@@ -1341,9 +1420,13 @@ async def _refresh_terminated_server(
logger.warning("MCP server '{}' session terminated; refreshing connection", server_name)
_unregister_server_tools(state, registry, server_name)
_retire_server_stack(state, server_name)
await _close_server(state, server_name)
connected = await connect_mcp_servers({server_name: cfg}, registry)
if getattr(state, "_mcp_closing", False):
for connection in connected.values():
await connection.aclose()
return None
state._mcp_stacks.update(connected)
_attach_reconnect_handlers(state, registry, connected)
if server_name not in connected:
@@ -1378,17 +1461,24 @@ def _unregister_server_tools(state: Any, registry: ToolRegistry, server_name: st
return removed
def _retire_server_stack(state: Any, server_name: str) -> None:
"""Remove a stale MCP stack from active use without closing it mid-turn.
MCP stream transports use AnyIO cancel scopes. Closing a stack from the
reconnecting dispatch task can inject ``CancelledError`` into the task that
originally opened it (often ``AgentLoop.run``), which crashes the gateway.
Retired stacks are closed later by ``AgentLoop.close_mcp`` during shutdown.
"""
async def _close_server(state: Any, server_name: str) -> None:
stack = state._mcp_stacks.pop(server_name, None)
if stack is None:
return
retired = getattr(state, "_mcp_retired_stacks", None)
if retired is not None:
retired.append((server_name, stack))
try:
await stack.aclose()
except (RuntimeError, BaseExceptionGroup):
logger.debug("MCP server '{}' cleanup error (can be ignored)", server_name)
async def close_mcp_servers(state: Any) -> None:
"""Close every MCP connection while excluding reconnect and hot reload."""
state._mcp_closing = True
async with _reload_lock(state):
connections = list(state._mcp_stacks.items())
state._mcp_stacks.clear()
for name, connection in connections:
try:
await connection.aclose()
except (RuntimeError, BaseExceptionGroup):
logger.debug("MCP server '{}' cleanup error (can be ignored)", name)