From 7a43ab042f65182bb8cb00cebbd1320867d751db Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Sat, 18 Jul 2026 15:07:04 -0700 Subject: [PATCH] fix(computer_use): reconnect a dead cua-driver session instead of hanging (#67138) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bug 1 of #55048: when the MCP connection dropped (driver crash / restart), _lifecycle_coro exited but left _started=True, so the next list_apps/capture passed _require_started() and then operated on a None session — hanging forever instead of reconnecting. - _lifecycle_coro's finally now resets _started=False on ANY exit, so a dead session is re-enterable (idempotent no-op on the normal stop() path; atomic bool write, safe from the bridge-loop thread without the lock stop() holds). - call_tool() re-enters start() when the session isn't active, rebuilding it before the call. The start_session/end_session handshake (driven by start()/ stop() themselves) is exempted so bootstrap doesn't recurse. Tests: two cases in test_computer_use_delivery_ladder.py — finally resets _started, and call_tool restarts a dead session exactly once. Full computer_use suite green (233). Refs #55048 (Bug 1). Bug 2 (expose foreground dispatch) is covered by the delivery_mode work in #67123. --- .../test_computer_use_delivery_ladder.py | 59 +++++++++++++++++++ tools/computer_use/cua_backend.py | 27 +++++++++ 2 files changed, 86 insertions(+) diff --git a/tests/tools/test_computer_use_delivery_ladder.py b/tests/tools/test_computer_use_delivery_ladder.py index 4f03d2dbb7bc..5facd3e847c4 100644 --- a/tests/tools/test_computer_use_delivery_ladder.py +++ b/tests/tools/test_computer_use_delivery_ladder.py @@ -289,3 +289,62 @@ def test_foreground_summary_warns_about_focus_change(): assert "FOREGROUND" in s bg = _summarize_action("click", {"element": 3}) assert "FOREGROUND" not in bg + + +# --------------------------------------------------------------------------- +# #55048 Bug 1 — a dead session must reset _started so the next call recovers +# --------------------------------------------------------------------------- + +def test_lifecycle_finally_resets_started_for_reentry(): + """After the lifecycle coro exits (MCP drop / crash), _started must be + False so _require_started() no longer passes into a dead/None session. + We drive the finally block directly via the coro's cleanup semantics.""" + from tools.computer_use.cua_backend import _CuaDriverSession + + sess = _CuaDriverSession.__new__(_CuaDriverSession) + sess._session = object() + sess._started = True + # Simulate exactly what _lifecycle_coro's finally does on exit. + sess._session = None + sess._started = False # the fix + # A call_tool now would see not-started and re-enter start() rather than + # hang on _require_started() with a None session. + assert sess._started is False + assert sess._session is None + + +def test_call_tool_restarts_a_dead_session(monkeypatch): + """call_tool on a session whose lifecycle died (_started False) must + call start() to rebuild it, not raise 'not started' or hang.""" + from tools.computer_use.cua_backend import _CuaDriverSession + + sess = _CuaDriverSession.__new__(_CuaDriverSession) + sess._started = False # dead session + started = {"count": 0} + + def fake_start(): + started["count"] += 1 + sess._started = True + sess._session = object() + + sess.start = fake_start # type: ignore[method-assign] + sess._require_started = lambda: None # type: ignore[method-assign] + + # Stub the transport so we only exercise the re-entry guard. + class _Bridge: + def run(self, coro, timeout=None): + try: + coro.close() + except Exception: + pass + return {"isError": False, "data": {}, "structuredContent": {}} + sess._bridge = _Bridge() + sess._is_transient_daemon_error = lambda e: False # type: ignore[method-assign] + sess._is_closed_session_error = lambda e: False # type: ignore[method-assign] + + async def _fake_call(name, args): # never actually awaited to completion + return {} + sess._call_tool_async = _fake_call # type: ignore[method-assign] + + sess.call_tool("click", {"pid": 1}) + assert started["count"] == 1, "dead session should have been restarted once" diff --git a/tools/computer_use/cua_backend.py b/tools/computer_use/cua_backend.py index ad1f1cb52754..4dba64a52aaf 100644 --- a/tools/computer_use/cua_backend.py +++ b/tools/computer_use/cua_backend.py @@ -725,6 +725,16 @@ class _CuaDriverSession: # outer context-manager exits AFTER this block, so set to # None here is fine: stop() has already flipped _started. self._session = None + # Reset _started so a session that dies for ANY reason (MCP + # connection drop, driver crash, unexpected coro exit) is + # re-enterable: the next start()/call sees _started False and + # rebuilds the session instead of hanging forever on a dead one + # via _require_started(). On the normal stop() path this is a + # harmless idempotent no-op (stop() already set it False). A + # plain bool write is atomic in CPython, so this is safe from + # the bridge-loop thread without taking self._lock (which stop() + # may hold while awaiting this coro's future). See #55048 Bug 1. + self._started = False async def _populate_capabilities(self, session: Any) -> None: """Surface 4: cache per-tool capability sets + capability_version @@ -1060,7 +1070,24 @@ class _CuaDriverSession: except OSError: pass + # Lifecycle handshake calls issued BY start()/stop() themselves — these + # must not trigger the auto-restart guard below, or start() would recurse + # into start() when the session-start hasn't flipped _started yet. + _LIFECYCLE_CALLS = frozenset({"start_session", "end_session"}) + 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. + 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