refactor(computer_use): unify cua parse image sniff with backend.image_dimensions_from_bytes; inline _first_nonempty_list; compact prose

This commit is contained in:
Teknium
2026-09-02 16:23:58 -07:00
parent 3fefdd879b
commit 49cef0bae3
+76 -187
View File
@@ -1,8 +1,6 @@
"""Pure parsing helpers for the cua-driver backend: MCP result flattening,
``list_windows`` / ``get_window_state`` payload normalisation, key combos.
No I/O, no module state — everything here is a function of its inputs, which
is what makes it safe to share between the MCP and CLI transports.
"""Pure parsing helpers for the cua-driver backend: MCP result flattening, ``list_windows``
/ ``get_window_state`` payload normalisation, key combos. No I/O, no module state — every
function depends only on its inputs, so the MCP and CLI transports share them safely.
"""
from __future__ import annotations
@@ -11,19 +9,13 @@ import json
import re
from typing import Any, Dict, List, Optional, Tuple
from tools.computer_use.backend import ActionResult, UIElement
from tools.computer_use.backend import ActionResult, UIElement, image_dimensions_from_bytes
_MISSING = object()
# Linux/X11 can surface GNOME Shell / desktop backdrop windows before real app
# windows with no useful z-order. They are targetable X11 windows but capture
# as empty through get_window_state, so default app capture must skip them.
_NON_APP_WINDOW_TITLE_PREFIXES = (
"@!", # GNOME Shell background/monitor helper windows
"Desktop",
"gnome-shell",
"GNOME Shell",
)
# Linux/X11 surfaces GNOME Shell / desktop backdrop windows ahead of real app windows with
# no useful z-order; they are targetable but capture as empty, so default capture skips them.
_NON_APP_WINDOW_TITLE_PREFIXES = ("@!", "Desktop", "gnome-shell", "GNOME Shell") # "@!" = GNOME helpers
_ELEMENT_LINE_RE = re.compile(
r'^\s*(?:-\s+)?\[(\d+)\]\s+(\w+)'
@@ -35,25 +27,18 @@ _ELEMENT_LINE_RE = re.compile(
r'(?:\s+(?:\(\d+\)\s+)?id=([^\s\[\]]+))?', # optional id=value (after an optional (order))
re.MULTILINE,
)
"""Element line of the get_window_state AX-tree markdown.
cua-driver renders each actionable node as ``[N] AXRole`` followed by a label
in one of four forms — ``= "value"``, ``"quoted"``, ``(paren)``, ``id=Label``
(optionally after an ``(order)`` number). A parenthesised pure-digit group is
an ORDER index, not a label, and is excluded so the id= label wins. Group 1 is
the index, group 2 the role, groups 3-6 the label in whichever form matched.
"""
"""get_window_state AX-tree markdown line: ``[N] AXRole`` + label in one of four forms
(``= "value"``, ``"quoted"``, ``(paren)``, ``id=Label`` optionally after an ``(order)`` number).
A parenthesised pure-digit group is an ORDER index, not a label, and is excluded so the id=
label wins. Group 1 index, group 2 role, groups 3-6 the label in whichever form matched."""
def _mcp_field(obj, snake: str, camel: str, default=None):
"""Read an MCP model field across the 1.x -> 2.x rename.
mcp 2.0 exposes snake_case attributes and keeps camelCase only as a
serialization alias, so ``getattr(result, "isError", False)`` reads False
for every result on 2.x and a denied call would look like a success.
Deliberately duplicated from ``tools.mcp_tool.mcp_field`` so computer_use
never loads the much larger config-driven MCP client module.
"""
"""Read an MCP model field across the 1.x -> 2.x rename: mcp 2.0 exposes snake_case
attributes and keeps camelCase only as a serialization alias, so ``getattr(result,
"isError", False)`` is False for every 2.x result and a denied call looks like success.
Deliberately duplicated from ``tools.mcp_tool.mcp_field`` so computer_use never loads
the much larger config-driven MCP client module."""
value = getattr(obj, snake, _MISSING)
if value is not _MISSING:
return value
@@ -61,21 +46,11 @@ def _mcp_field(obj, snake: str, camel: str, default=None):
return default if value is _MISSING else value
def _action_result_from(
name: str,
ok: bool,
message: str,
meta: Dict[str, Any],
structured: Dict[str, Any],
*,
requested_delivery: Optional[str] = None,
) -> ActionResult:
"""Build an ActionResult, lifting cua-driver's structured verdict.
structuredContent is canonical, the flattened ``meta`` copy the fallback.
Every structured field is additive: a driver that omits one leaves the
attribute ``None`` so old drivers see unchanged behavior.
"""
def _action_result_from(name: str, ok: bool, message: str, meta: Dict[str, Any],
structured: Dict[str, Any], *, requested_delivery: Optional[str] = None) -> ActionResult:
"""Build an ActionResult, lifting cua-driver's structured verdict. structuredContent is
canonical, the flattened ``meta`` copy the fallback. Every structured field is additive:
a driver that omits one leaves the attribute ``None`` so old drivers see unchanged behavior."""
sc = structured if isinstance(structured, dict) else {}
def _pick(key: str) -> Any:
@@ -85,10 +60,7 @@ def _action_result_from(
return value if isinstance(value, typ) else None
return ActionResult(
ok=ok,
action=name,
message=message,
meta=meta,
ok=ok, action=name, message=message, meta=meta,
verified=_typed(_pick("verified"), bool),
effect=_typed(_pick("effect"), str),
escalation=_typed(_pick("escalation"), dict),
@@ -103,24 +75,15 @@ def _action_result_from(
def _z_index_uninformative(windows: List[Dict[str, Any]]) -> bool:
"""True when every window shares the same z_index (common on Linux/X11)."""
if not windows:
return True
return len({w.get("z_index", 0) for w in windows}) <= 1
def _parse_xprop_net_active_window(stdout: str) -> Optional[int]:
"""Parse ``xprop -root _NET_ACTIVE_WINDOW`` stdout into a window id.
Accepts the ``window id # 0x...`` form, falling back to the first hex token.
"""
"""Parse ``xprop -root _NET_ACTIVE_WINDOW`` stdout into a window id: the ``window id #
0x...`` form, falling back to the first hex token."""
text = stdout or ""
match = re.search(r"window id # (0x[0-9a-fA-F]+)", text) or re.search(r"(0x[0-9a-fA-F]+)", text)
if not match:
return None
try:
return int(match.group(1), 16)
except ValueError:
return None
return int(match.group(1), 16) if match else None
def _is_real_app_window(w: Dict[str, Any]) -> bool:
@@ -133,12 +96,9 @@ def _is_real_app_window(w: Dict[str, Any]) -> bool:
def _parse_elements_from_tree(markdown: str) -> List[UIElement]:
"""Parse UIElements from get_window_state AX-tree markdown.
Last-resort fallback for drivers without ``structuredContent.elements``.
Bounds always come back ``(0, 0, 0, 0)`` — the markdown carries none —
which is fine for element-index clicks (the driver resolves the frame).
"""
"""Parse UIElements from get_window_state AX-tree markdown — last-resort fallback for
drivers without ``structuredContent.elements``. Bounds are always ``(0, 0, 0, 0)`` (the
markdown carries none), fine for element-index clicks since the driver resolves the frame."""
return [
UIElement(
index=int(m.group(1)),
@@ -152,79 +112,36 @@ def _parse_elements_from_tree(markdown: str) -> List[UIElement]:
def _parse_elements_from_structured(raw_elements: List[Dict[str, Any]]) -> List[UIElement]:
"""Read the canonical ``structuredContent.elements`` array.
Each entry has ``element_index``, ``role``, ``label`` and, when the AT-SPI /
AXFrame call returned usable bounds, ``frame`` ``{x, y, w, h}`` — so real
pixel bounds survive (the markdown path loses them). Malformed entries are
skipped rather than failing the whole walk.
"""
"""Read the canonical ``structuredContent.elements`` array: ``element_index``, ``role``,
``label`` and, when AT-SPI / AXFrame returned usable bounds, ``frame`` ``{x, y, w, h}`` —
so real pixel bounds survive (the markdown path loses them). Malformed entries are skipped."""
elements: List[UIElement] = []
for raw in raw_elements:
if not isinstance(raw, dict):
continue
idx = raw.get("element_index")
idx = raw.get("element_index") if isinstance(raw, dict) else None
if not isinstance(idx, int):
continue
role = raw.get("role") if isinstance(raw.get("role"), str) else ""
label = raw.get("label") if isinstance(raw.get("label"), str) else ""
frame = raw.get("frame") if isinstance(raw.get("frame"), dict) else None
role, label, frame, token = (raw.get(k) for k in ("role", "label", "frame", "element_token"))
bounds: Tuple[int, int, int, int] = (0, 0, 0, 0)
if frame:
if isinstance(frame, dict) and frame:
try:
bounds = tuple(int(frame.get(k, 0)) for k in ("x", "y", "w", "h")) # type: ignore[assignment]
except (TypeError, ValueError):
bounds = (0, 0, 0, 0)
# Opaque element_token (`s{snapshot_hex}:{index}`) — the driver owns
# the parse + LRU semantics; we treat it as a black-box string.
raw_token = raw.get("element_token")
elements.append(UIElement(
index=idx,
role=role,
label=label,
role=role if isinstance(role, str) else "",
label=label if isinstance(label, str) else "",
bounds=bounds,
element_token=raw_token if isinstance(raw_token, str) and raw_token else None,
# Opaque `s{snapshot_hex}:{index}` token — the driver owns parse + LRU semantics.
element_token=token if isinstance(token, str) and token else None,
))
return elements
def _image_dimensions_from_bytes(raw: bytes) -> Tuple[int, int]:
"""Best-effort PNG/JPEG dimension sniffing without extra dependencies."""
if raw.startswith(b"\x89PNG\r\n\x1a\n") and len(raw) >= 24:
width = int.from_bytes(raw[16:20], "big")
height = int.from_bytes(raw[20:24], "big")
if width > 0 and height > 0:
return width, height
if raw.startswith(b"\xff\xd8"):
i = 2
n = len(raw)
while i + 9 < n:
if raw[i] != 0xFF:
i += 1
continue
marker = raw[i + 1]
i += 2
if marker in {0xD8, 0xD9} or 0xD0 <= marker <= 0xD7:
continue
if i + 2 > n:
break
segment_len = int.from_bytes(raw[i:i + 2], "big")
if segment_len < 2 or i + segment_len > n:
break
if marker in {
0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7,
0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF,
}:
if segment_len >= 7:
height = int.from_bytes(raw[i + 3:i + 5], "big")
width = int.from_bytes(raw[i + 5:i + 7], "big")
if width > 0 and height > 0:
return width, height
break
i += segment_len
return 0, 0
"""Best-effort PNG/JPEG dimension sniff; ``(0, 0)`` when unreadable or non-positive."""
dims = image_dimensions_from_bytes(raw)
return dims if dims and dims[0] > 0 and dims[1] > 0 else (0, 0)
def _split_tree_text(full_text: str) -> Tuple[str, str]:
@@ -251,31 +168,21 @@ def _parse_key_combo(keys: str) -> Tuple[Optional[str], List[str]]:
def _extract_tool_result(mcp_result: Any) -> Dict[str, Any]:
"""Flatten an mcp CallToolResult into
``{data, images, image_mime_types, structuredContent, isError}``.
``data`` is the joined text parts (parsed as JSON when it looks like JSON);
``image_mime_types`` is parallel to ``images`` with ``""`` where the part
carried no mimeType (older drivers — callers then sniff the base64 prefix).
"""
"""Flatten an mcp CallToolResult into ``{data, images, image_mime_types, structuredContent,
isError}``. ``data`` is the joined text parts (parsed as JSON when it looks like JSON);
``image_mime_types`` is parallel to ``images`` with ``""`` where the part carried no
mimeType (older drivers — callers then sniff the base64 prefix)."""
data: Any = None
images: List[str] = []
image_mime_types: List[str] = []
# Identity, not truthiness: mocks/proxies synthesize truthy attributes.
is_error = _mcp_field(mcp_result, "is_error", "isError", False) is True
structured: Optional[Dict] = (
_mcp_field(mcp_result, "structured_content", "structuredContent") or None
)
text_chunks: List[str] = []
for part in getattr(mcp_result, "content", []) or []:
ptype = getattr(part, "type", None)
if ptype == "text":
text_chunks.append(getattr(part, "text", "") or "")
elif ptype == "image":
b64 = getattr(part, "data", None)
if b64:
images.append(b64)
image_mime_types.append(_mcp_field(part, "mime_type", "mimeType") or "")
elif ptype == "image" and getattr(part, "data", None):
images.append(part.data)
image_mime_types.append(_mcp_field(part, "mime_type", "mimeType") or "")
if text_chunks:
joined = "\n".join(t for t in text_chunks if t)
try:
@@ -286,19 +193,16 @@ def _extract_tool_result(mcp_result: Any) -> Dict[str, Any]:
"data": data,
"images": images,
"image_mime_types": image_mime_types,
"structuredContent": structured,
"isError": is_error,
"structuredContent": _mcp_field(mcp_result, "structured_content", "structuredContent") or None,
# Identity, not truthiness: mocks/proxies synthesize truthy attributes.
"isError": _mcp_field(mcp_result, "is_error", "isError", False) is True,
}
def _image_from_tool_result(out: Dict[str, Any]) -> tuple[Optional[str], Optional[str]]:
"""Pull ``(b64, mime_type)`` out of a flattened tool result.
cua-driver delivers screenshots either as an MCP ``image`` part
(``out["images"]``) or as ``screenshot_png_b64`` inside structuredContent
(newer builds, and the CLI transport); checking both keeps capture()
robust when the driver moves the image between the two.
"""
"""Pull ``(b64, mime_type)`` out of a flattened tool result. cua-driver delivers screenshots
as an MCP ``image`` part (``out["images"]``) or as ``screenshot_png_b64`` in structuredContent
(newer builds, CLI transport); checking both keeps capture() robust to the driver moving it."""
images = out.get("images") or []
if images and images[0]:
mimes = out.get("image_mime_types") or []
@@ -322,12 +226,9 @@ def _positive_int(value: Any) -> Optional[int]:
def _is_placeholder_id(value: Any) -> bool:
"""True when *value* is a schema-filler id (``0`` / negative) rather than a target.
Some providers emit every optional integer zero-filled; treating that as a
targeting request would drop the caller's ``app=``. Non-numeric values are
NOT placeholders — they still reach the validation error.
"""
"""True when *value* is a schema-filler id (``0`` / negative) rather than a target: some
providers zero-fill every optional integer, and treating that as targeting would drop the
caller's ``app=``. Non-numeric values are NOT placeholders — they still reach validation."""
if isinstance(value, bool) or not isinstance(value, (int, str)):
return False
try:
@@ -337,25 +238,19 @@ def _is_placeholder_id(value: Any) -> bool:
def _ingest_windows(raw_windows: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Normalise cua-driver ``list_windows`` entries, dropping unusable ones.
Every downstream call needs an integer ``pid`` and ``window_id``. On X11
the PID comes from the optional ``_NET_WM_PID`` property, so root/panel/
popup windows report ``pid: null`` — skip those instead of aborting the
whole enumeration. ``z_index``: higher = closer to front; Wayland's null
(undefined stacking) sorts lowest so real windows stay above the desktop.
"""
"""Normalise cua-driver ``list_windows`` entries, dropping unusable ones. Every downstream
call needs integer ``pid`` and ``window_id``; on X11 the PID comes from the optional
``_NET_WM_PID`` property, so root/panel/popup windows report ``pid: null`` — skip those
instead of aborting the enumeration. ``z_index``: higher = closer to front; Wayland's null
(undefined stacking) sorts lowest so real windows stay above the desktop."""
windows: List[Dict[str, Any]] = []
for w in raw_windows:
if not isinstance(w, dict): # untrusted compatibility envelopes
continue
pid_int = _positive_int(w.get("pid"))
window_id_int = _positive_int(w.get("window_id"))
pid_int, window_id_int = _positive_int(w.get("pid")), _positive_int(w.get("window_id"))
if pid_int is None or window_id_int is None:
continue
z_raw = w.get("z_index")
app_name = w.get("app_name", "")
title = w.get("title", "")
z_raw, app_name, title = w.get("z_index"), w.get("app_name", ""), w.get("title", "")
windows.append({
"app_name": app_name if isinstance(app_name, str) else "",
"pid": pid_int,
@@ -368,25 +263,19 @@ def _ingest_windows(raw_windows: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
return windows
def _first_nonempty_list(*containers: Any, keys: Tuple[str, ...]) -> List[Any]:
for container in containers:
if not isinstance(container, dict):
continue
for key in keys:
value = container.get(key)
if isinstance(value, list) and value:
return value
return []
def _windows_from_tool_result(out: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Return list_windows payloads across cua-driver result shapes."""
structured = out.get("structuredContent")
if isinstance(structured, dict):
windows = structured.get("windows")
if isinstance(windows, list) and windows:
return windows
return _first_nonempty_list(out.get("data"), out, keys=("windows", "_legacy_windows"))
"""Return list_windows payloads across cua-driver result shapes: structuredContent.windows,
then ``windows`` / ``_legacy_windows`` in the text payload, then on the envelope itself."""
candidates = ((out.get("structuredContent"), ("windows",)),
(out.get("data"), ("windows", "_legacy_windows")),
(out, ("windows", "_legacy_windows")))
for container, keys in candidates:
if isinstance(container, dict):
for key in keys:
value = container.get(key)
if isinstance(value, list) and value:
return value
return []
def _apps_from_windows(windows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: