refactor(computer_use): unify session redeclare path, inline single-use session helpers

This commit is contained in:
Teknium
2026-09-02 16:07:05 -07:00
parent 5629d0fe3d
commit eb837d2e36

View File

@@ -9,6 +9,7 @@ from __future__ import annotations
import asyncio
import base64
import concurrent.futures
import contextlib
import json
import logging
import os
@@ -40,10 +41,8 @@ class _AsyncBridge:
try:
self._loop.run_forever()
finally:
try:
with contextlib.suppress(Exception):
self._loop.close()
except Exception:
pass
self._thread = threading.Thread(target=_run, daemon=True, name="cua-driver-loop")
self._thread.start()
@@ -126,32 +125,26 @@ def _cli_run_json(cmd: List[str], env: Dict[str, str], name: str, timeout: float
def _cli_result(parsed: Any, shot_file: Optional[str]) -> Dict[str, Any]:
"""Remap a ``cua-driver call`` JSON body into the ``_extract_tool_result`` shape."""
images: List[str] = []
if not isinstance(parsed, dict):
return {"data": None, "images": [], "structuredContent": None, "isError": False}
# Logical failures may be reported in-band even when the subprocess exits 0 — fail closed.
is_error = parsed.get("isError") is True or parsed.get("is_error") is True
shot = parsed.get("screenshot_png_b64")
if not shot:
# Screenshot was routed to a file (ours or the daemon's choice).
fpath = parsed.get("screenshot_file_path") or shot_file
if fpath and os.path.exists(fpath):
try:
with open(fpath, "rb") as fh:
shot = base64.b64encode(fh.read()).decode("ascii")
except Exception as e:
logger.debug("cua-driver CLI fallback: failed reading %s: %s", fpath, e)
data: Any = None
is_error = False
if isinstance(parsed, dict):
# Logical failures may be reported in-band even when the subprocess
# exits 0 — preserve the bit so callers fail closed.
is_error = parsed.get("isError") is True or parsed.get("is_error") is True
shot = parsed.get("screenshot_png_b64")
if not shot:
# Screenshot was routed to a file (ours or the daemon's choice).
fpath = parsed.get("screenshot_file_path") or shot_file
if fpath and os.path.exists(fpath):
try:
with open(fpath, "rb") as fh:
shot = base64.b64encode(fh.read()).decode("ascii")
except Exception as e:
logger.debug("cua-driver CLI fallback: failed reading %s: %s", fpath, e)
if shot:
images.append(shot)
tree = parsed.get("tree_markdown")
if tree is not None:
ec = parsed.get("element_count")
summary = f"{ec} elements" if ec is not None else ""
data = f"{summary}\n{tree}" if summary else tree
structured = parsed if isinstance(parsed, dict) else None
return {"data": data, "images": images, "structuredContent": structured, "isError": is_error}
tree = parsed.get("tree_markdown")
if tree is not None:
ec = parsed.get("element_count")
data = f"{ec} elements\n{tree}" if ec is not None else tree
return {"data": data, "images": [shot] if shot else [], "structuredContent": parsed, "isError": is_error}
class _CuaDriverSession:
@@ -348,7 +341,6 @@ class _CuaDriverSession:
logger.debug("cua-driver transport reset callback failed: %s", exc)
def _stop_lifecycle_locked(self) -> None:
"""Signal shutdown and wait (5s) for the lifecycle coroutine to unwind."""
self._signal_shutdown_locked()
fut = self._lifecycle_future
if fut is None:
@@ -367,10 +359,8 @@ class _CuaDriverSession:
loop = self._bridge._loop
event = self._shutdown_event
if loop is not None and event is not None and loop.is_running():
try:
with contextlib.suppress(RuntimeError): # loop closed — nothing to signal
loop.call_soon_threadsafe(event.set)
except RuntimeError: # loop closed — nothing to signal
pass
async def _call_tool_async(self, name: str, args: Dict[str, Any]) -> Dict[str, Any]:
result = await self._session.call_tool(name, args)
@@ -451,10 +441,6 @@ class _CuaDriverSession:
return any(needle in msg for needle in ("Resource temporarily unavailable", "os error 35",
"daemon transport error", "daemon proxy"))
@classmethod
def _transport_replay_is_safe(cls, name: str) -> bool:
return name in cls._TRANSPORT_REPLAY_SAFE_TOOLS
@staticmethod
def _unknown_transport_outcome(name: str, exc: Exception) -> Dict[str, Any]:
return _outcome_unknown(
@@ -487,22 +473,23 @@ class _CuaDriverSession:
return first_result
logger.warning("cua-driver session %s ended during %s; reviving and retrying once",
session_id, name)
revive_result = self._run_call("start_session", {"session": session_id}, 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))
if not self._redeclare_session(timeout, "cua-driver session %s could not be revived: %s"):
return first_result
return self._run_call(name, args, timeout)
def _restore_declared_session_after_transport_reset(self, timeout: float) -> None:
"""Re-attach the public label inside a replacement private lifecycle."""
session_id = getattr(self, "_declared_session_id", None)
if not session_id:
return
def _redeclare_session(self, timeout: float, failure_msg: str) -> bool:
"""start_session with the declared id; log *failure_msg* and return False on rejection."""
session_id = self._declared_session_id
result = self._run_call("start_session", {"session": session_id}, timeout)
if result.get("isError") is True:
logger.warning("cua-driver public session label %s could not be restored: %s",
session_id, self._logical_error_text(result))
logger.warning(failure_msg, session_id, self._logical_error_text(result))
return False
return True
def _restore_declared_session_after_transport_reset(self, timeout: float) -> None:
"""Re-attach the public label inside a replacement private lifecycle."""
if getattr(self, "_declared_session_id", None):
self._redeclare_session(timeout, "cua-driver public session label %s could not be restored: %s")
def _restart_session_locked(self) -> None:
"""Recreate the MCP session after the transport closed. Caller holds self._lock."""
@@ -519,11 +506,6 @@ class _CuaDriverSession:
self._start_lifecycle_locked()
self._started = True
def _restart_and_restore(self, timeout: float) -> None:
with self._lock:
self._restart_session_locked()
self._restore_declared_session_after_transport_reset(timeout)
def _cli_command(self, name: str, args: Dict[str, Any]) -> Tuple[List[str], Dict[str, str], Optional[str]]:
"""Build ``(cmd, child_env, shot_file)`` for the CLI fallback. ``get_window_state``
routes its screenshot to a temp file (``screenshot_out_file``) so the daemon returns
@@ -564,10 +546,8 @@ class _CuaDriverSession:
return _cli_result(parsed, shot_file)
finally:
if shot_file and os.path.exists(shot_file):
try:
with contextlib.suppress(OSError):
os.remove(shot_file)
except OSError:
pass
def call_tool(self, name: str, args: Dict[str, Any], timeout: float = 30.0) -> Dict[str, Any]:
# A prior MCP timeout marks the session suspect (possibly wedged): recreate it
@@ -597,7 +577,7 @@ class _CuaDriverSession:
"for recreation before the next call", name)
return self._timeout_outcome(name, e)
if self._is_transient_daemon_error(e):
if not self._transport_replay_is_safe(name):
if name not in self._TRANSPORT_REPLAY_SAFE_TOOLS:
self._notify_transport_reset()
return self._unknown_transport_outcome(name, e)
logger.warning("cua-driver MCP transport failed on %s (%s); "
@@ -606,8 +586,10 @@ class _CuaDriverSession:
if not self._is_closed_session_error(e):
raise
logger.warning("cua-driver MCP session closed during %s; reconnecting once", name)
self._restart_and_restore(timeout)
if not self._transport_replay_is_safe(name):
with self._lock:
self._restart_session_locked()
self._restore_declared_session_after_transport_reset(timeout)
if name not in self._TRANSPORT_REPLAY_SAFE_TOOLS:
return self._unknown_transport_outcome(name, e)
result = self._run_call(name, args, timeout)