fix(cua-driver): reconnect MCP stdio session once on ClosedResourceError after daemon restart (#40570)

Salvaged from #40282; cleaned up, re-verified against main, tests added.

Co-authored-by: jeeves-assistant <jeeves-assistant@users.noreply.github.com>
This commit is contained in:
Teknium
2026-06-06 18:35:12 -07:00
committed by GitHub
co-authored by jeeves-assistant
parent 97524344ad
commit 365437e4aa
2 changed files with 106 additions and 1 deletions
+34 -1
View File
@@ -277,9 +277,42 @@ class _CuaDriverSession:
result = await self._session.call_tool(name, args)
return _extract_tool_result(result)
@staticmethod
def _is_closed_session_error(exc: Exception) -> bool:
"""Return True for MCP/stdio failures that are recoverable by reconnecting."""
name = exc.__class__.__name__
module = getattr(exc.__class__, "__module__", "")
return (
name in {"ClosedResourceError", "BrokenResourceError", "EndOfStream"}
or (module.startswith("anyio") and "Resource" in name)
or isinstance(exc, (BrokenPipeError, EOFError))
)
def _restart_session_locked(self) -> None:
"""Recreate the MCP session after the daemon/stdin transport was closed."""
try:
if self._started:
self._bridge.run(self._aexit(), timeout=5.0)
except Exception as e:
logger.debug("cua-driver session cleanup before reconnect failed: %s", e)
self._started = False
self._bridge.run(self._aenter(), timeout=15.0)
self._started = True
def call_tool(self, name: str, args: Dict[str, Any], timeout: float = 30.0) -> Dict[str, Any]:
self._require_started()
return self._bridge.run(self._call_tool_async(name, args), timeout=timeout)
try:
return self._bridge.run(self._call_tool_async(name, args), timeout=timeout)
except Exception as e:
if not self._is_closed_session_error(e):
raise
# Daemon restart closes the cached stdio channel. Reconnect once and
# retry exactly one more time — never loop, to avoid hammering a
# genuinely dead daemon.
logger.warning("cua-driver MCP session closed during %s; reconnecting once", name)
with self._lock:
self._restart_session_locked()
return self._bridge.run(self._call_tool_async(name, args), timeout=timeout)
def _extract_tool_result(mcp_result: Any) -> Dict[str, Any]: