From 56bda4529b23ceac398e810aa21d713f778dd9fa Mon Sep 17 00:00:00 2001 From: Atakan Date: Sat, 25 Jul 2026 18:36:52 +0300 Subject: [PATCH] fix(computer_use): revive ended cua-driver sessions once --- tests/tools/test_computer_use.py | 147 ++++++++++++++++++++++++++++++ tools/computer_use/cua_backend.py | 112 +++++++++++++++++++---- 2 files changed, 240 insertions(+), 19 deletions(-) diff --git a/tests/tools/test_computer_use.py b/tests/tools/test_computer_use.py index d0ac4d54b979..2aaadbb37b4b 100644 --- a/tests/tools/test_computer_use.py +++ b/tests/tools/test_computer_use.py @@ -1556,6 +1556,7 @@ class TestCuaDriverSessionReconnect: session._shutdown_event = None session._lifecycle_future = None session._setup_error = None + session._declared_session_id = None session._call_tool_async = lambda name, args: ("call", name, args) # Record what reconnect does — stop then start, in that order. session._reconnect_log = [] @@ -1590,6 +1591,152 @@ class TestCuaDriverSessionReconnect: assert bridge.calls[1][0] == ("call", "list_apps", {}) assert len(bridge.calls) == 2 + def test_reconnect_retry_can_revive_ended_session(self): + """A reconnect result may still require reviving the declared session.""" + from anyio import ClosedResourceError + + ended = { + "data": "session 'hermes-test' has ended; call start_session to revive it", + "images": [], + "structuredContent": None, + "isError": True, + } + revived = {"data": "revived", "isError": False} + windows = { + "data": "", + "structuredContent": {"windows": []}, + "isError": False, + } + + class FakeBridge: + def __init__(self): + self.calls = [] + self.effects = [ + ClosedResourceError(), + ended, + revived, + windows, + ] + + def run(self, value, timeout=None): + self.calls.append((value, timeout)) + effect = self.effects.pop(0) + if isinstance(effect, Exception): + raise effect + return effect + + bridge = FakeBridge() + session = self._make_session(bridge) + session._declared_session_id = "hermes-test" + + result = session.call_tool( + "list_windows", {"session": "hermes-test"} + ) + + assert result is windows + assert session._reconnect_log == ["stop", "start"] + assert [call[0] for call in bridge.calls] == [ + ("call", "list_windows", {"session": "hermes-test"}), + ("call", "list_windows", {"session": "hermes-test"}), + ("call", "start_session", {"session": "hermes-test"}), + ("call", "list_windows", {"session": "hermes-test"}), + ] + + def test_call_tool_revives_ended_session_then_retries_once(self): + """Logical ended-session errors revive the same id before one replay.""" + ended = { + "data": ( + "session 'hermes-test' has ended; tool call 'list_windows' " + "was rejected. Call start_session with this id to revive it." + ), + "images": [], + "structuredContent": None, + "isError": True, + } + ok = {"data": "revived", "images": [], "structuredContent": None, "isError": False} + windows = { + "data": "", + "images": [], + "structuredContent": {"windows": [{"pid": 1, "window_id": 2}]}, + "isError": False, + } + + class FakeBridge: + def __init__(self): + self.calls = [] + self.effects = [ended, ok, windows] + + def run(self, value, timeout=None): + self.calls.append((value, timeout)) + return self.effects.pop(0) + + bridge = FakeBridge() + session = self._make_session(bridge) + session._declared_session_id = "hermes-test" + + result = session.call_tool( + "list_windows", {"on_screen_only": True, "session": "hermes-test"} + ) + + assert result is windows + assert [call[0] for call in bridge.calls] == [ + ("call", "list_windows", {"on_screen_only": True, "session": "hermes-test"}), + ("call", "start_session", {"session": "hermes-test"}), + ("call", "list_windows", {"on_screen_only": True, "session": "hermes-test"}), + ] + + def test_call_tool_does_not_loop_when_retry_is_still_ended(self): + """A repeated ended-session result is surfaced after one revival.""" + ended = { + "data": "session 'hermes-test' has ended; call start_session to revive it", + "images": [], + "structuredContent": None, + "isError": True, + } + + class FakeBridge: + def __init__(self): + self.calls = [] + self.effects = [ended, {"isError": False}, ended] + + def run(self, value, timeout=None): + self.calls.append((value, timeout)) + return self.effects.pop(0) + + bridge = FakeBridge() + session = self._make_session(bridge) + session._declared_session_id = "hermes-test" + + result = session.call_tool("list_windows", {"session": "hermes-test"}) + + assert result is ended + assert len(bridge.calls) == 3 + + def test_lifecycle_call_does_not_try_to_revive_itself(self): + """start_session failures stay single-shot and cannot recurse.""" + ended = { + "data": "session 'hermes-test' has ended; call start_session to revive it", + "images": [], + "structuredContent": None, + "isError": True, + } + + class FakeBridge: + def __init__(self): + self.calls = [] + + def run(self, value, timeout=None): + self.calls.append((value, timeout)) + return ended + + bridge = FakeBridge() + session = self._make_session(bridge) + + result = session.call_tool("start_session", {"session": "hermes-test"}) + + assert result is ended + assert len(bridge.calls) == 1 + def test_call_tool_does_not_retry_on_unrelated_error(self): """Non-transport errors must propagate without a reconnect attempt.""" class FakeBridge: diff --git a/tools/computer_use/cua_backend.py b/tools/computer_use/cua_backend.py index 79912814e468..dc084d563519 100644 --- a/tools/computer_use/cua_backend.py +++ b/tools/computer_use/cua_backend.py @@ -909,6 +909,10 @@ class _CuaDriverSession: self._shutdown_event: Optional[asyncio.Event] = None # created on bridge loop self._lifecycle_future = None # concurrent.futures.Future self._setup_error: Optional[BaseException] = None + # Stable driver-side identity declared through start_session. + # Used to revive a logical ended-session rejection without + # recursive call_tool re-entry or backend-owned state (#71166). + self._declared_session_id: Optional[str] = None def _require_started(self) -> None: if not self._started: @@ -1169,6 +1173,67 @@ class _CuaDriverSession: when the driver predates the field — older builds had no version).""" return self._capability_version + @staticmethod + def _logical_error_text(result: Dict[str, Any]) -> str: + """Flatten a logical MCP error into text for narrow classification.""" + chunks: List[str] = [] + for value in (result.get("data"), result.get("structuredContent")): + if isinstance(value, str): + chunks.append(value) + elif value is not None: + try: + chunks.append(json.dumps(value, sort_keys=True)) + except (TypeError, ValueError): + chunks.append(str(value)) + return "\n".join(chunks) + + @classmethod + def _is_ended_session_result(cls, result: Any) -> bool: + """Recognise cua-driver's explicit recoverable ended-session result.""" + if not isinstance(result, dict) or result.get("isError") is not True: + return False + message = cls._logical_error_text(result).lower() + return ( + "session" in message + and ("has ended" in message or "session ended" in message) + and "start_session" in message + ) + + def _revive_declared_session_once( + self, + name: str, + args: Dict[str, Any], + first_result: Dict[str, Any], + timeout: float, + ) -> Dict[str, Any]: + """Revive the stable session and replay one rejected tool call once.""" + session_id = self._declared_session_id + if not session_id or name in self._LIFECYCLE_CALLS: + return first_result + + logger.warning( + "cua-driver session %s ended during %s; reviving and retrying once", + session_id, + name, + ) + revive_result = self._bridge.run( + self._call_tool_async("start_session", {"session": session_id}), + timeout=timeout, + ) + if revive_result.get("isError") is True: + logger.warning( + "cua-driver session %s could not be revived: %s", + session_id, + self._logical_error_text(revive_result), + ) + return first_result + + # Return the second result as-is. A second rejection is surfaced; no loop. + return self._bridge.run( + self._call_tool_async(name, args), + timeout=timeout, + ) + @staticmethod def _is_closed_session_error(exc: Exception) -> bool: """Return True for MCP/stdio failures that are recoverable by reconnecting.""" @@ -1347,28 +1412,19 @@ class _CuaDriverSession: def call_tool(self, name: str, args: Dict[str, Any], timeout: float = 30.0) -> Dict[str, Any]: # A prior session may have died (MCP drop / driver crash): its - # lifecycle coro reset _started to False in its finally (#55048 - # Bug 1). Re-enter start() so we rebuild the session instead of - # calling _require_started() straight into a "not started" raise or - # a None session. start() is idempotent when already started. Skip - # this for the start_session/end_session handshake, which start()/ - # stop() drive directly while _started is still in flux. + # lifecycle coro reset _started to False in its finally (#55048). if not self._started and name not in self._LIFECYCLE_CALLS: logger.warning( "cua-driver session not active on %s; (re)starting before call", name ) self.start() self._require_started() - # The cua-driver daemon proxy returns POSIX EAGAIN ("Resource - # temporarily unavailable") for heavier calls like get_window_state when - # its non-blocking socket buffer is full. On some machines/builds this - # is persistent for get_window_state over the MCP stdio bridge, while - # the direct CLI transport keeps working. So: try the MCP path ONCE, - # and on the transient/transport error fall straight through to the CLI - # transport (which has its own retry + screenshot-to-file mitigation) - # rather than burning a long backoff chain on a path that won't recover. + try: - return self._bridge.run(self._call_tool_async(name, args), timeout=timeout) + result = self._bridge.run( + self._call_tool_async(name, args), + timeout=timeout, + ) except Exception as e: if self._is_transient_daemon_error(e): logger.warning( @@ -1378,13 +1434,31 @@ class _CuaDriverSession: return self._call_tool_via_cli(name, args, timeout) 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) + result = self._bridge.run( + self._call_tool_async(name, args), + timeout=timeout, + ) + + # Remember only a successfully declared stable identity. Failed + # start_session calls must not leave stale recovery state behind. + if name == "start_session" and result.get("isError") is not True: + declared_id = args.get("session") + if isinstance(declared_id, str) and declared_id: + self._declared_session_id = declared_id + + if self._is_ended_session_result(result): + result = self._revive_declared_session_once(name, args, result, timeout) + + if ( + name == "end_session" + and result.get("isError") is not True + and args.get("session") == self._declared_session_id + ): + self._declared_session_id = None + return result def _extract_tool_result(mcp_result: Any) -> Dict[str, Any]: