fix(mcp): defer stale stack cleanup during reconnect

This commit is contained in:
flyzstu
2026-07-11 11:45:20 +08:00
committed by Xubin Ren
parent c111aaa7ee
commit f3d1b9ca2d
3 changed files with 56 additions and 10 deletions
+8 -1
View File
@@ -361,6 +361,7 @@ class AgentLoop:
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_connecting = False
self._active_tasks: dict[str, list[asyncio.Task]] = {} # session_key -> tasks
self._background_tasks: list[asyncio.Task] = []
@@ -1177,9 +1178,15 @@ class AgentLoop:
if self._background_tasks:
await asyncio.gather(*self._background_tasks, return_exceptions=True)
self._background_tasks.clear()
for name, stack in self._mcp_stacks.items():
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()
+13 -7
View File
@@ -1177,7 +1177,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)
await _close_server(state, name)
_retire_server_stack(state, name)
state._mcp_servers = next_servers
retry_missing = sorted(
@@ -1341,7 +1341,7 @@ async def _refresh_terminated_server(
logger.warning("MCP server '{}' session terminated; refreshing connection", server_name)
_unregister_server_tools(state, registry, server_name)
await _close_server(state, server_name)
_retire_server_stack(state, server_name)
connected = await connect_mcp_servers({server_name: cfg}, registry)
state._mcp_stacks.update(connected)
@@ -1378,11 +1378,17 @@ def _unregister_server_tools(state: Any, registry: ToolRegistry, server_name: st
return removed
async def _close_server(state: Any, server_name: str) -> None:
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.
"""
stack = state._mcp_stacks.pop(server_name, None)
if stack is None:
return
try:
await stack.aclose()
except (RuntimeError, BaseExceptionGroup):
logger.debug("MCP server '{}' cleanup error (can be ignored)", server_name)
retired = getattr(state, "_mcp_retired_stacks", None)
if retired is not None:
retired.append((server_name, stack))