Blame for tiny_computer_agent/accessibility.py

"""Elements from the operating system's accessibility tree.

OCR reads pixels and *guesses* which marks are widgets. The accessibility tree
simply knows: every application that publishes one advertises each button, entry,
menu item and label with a role, a name, and its real screen extents. That is
faster than OCR, exact on small text the pixels lose, and - the part that matters
for a tiny model - it produces the same short, labelled menu the rest of the
harness already expects.

The catch is honest and load-bearing: not every application publishes a tree.
GTK/Qt/Electron and most native Windows apps do; Tk barely does, and a raw canvas
does not at all. So this is a detector, not a replacement - when the tree is
empty, `AutoDetector` falls through to OCR, and then to the numbered grid.

Backends:

* **AT-SPI** (Linux) - via PyGObject's ``Atspi`` typelib, implemented here and
  exercised against a real GTK window.
* **UI Automation** (Windows) - *not implemented*: :func:`accessibility_available`
  reports it plainly rather than pretending, so an agent on Windows keeps using
  OCR instead of silently indexing nothing.

The coordinates AT-SPI reports are already absolute screen pixels, which is why
:class:`AtspiDetector` sets ``absolute_coordinates``: :class:`IndexedScreen` must
not re-map them through a downscale, or every click would land in the wrong place.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Protocol, Sequence

from tiny_computer_agent.observe import Element
from tiny_computer_agent.perception import DetectorUnavailable

# Roles that are worth clicking. Deliberately not "everything with a name": a
# frame, a panel or a scrollbar is named too, and offering them to a model is how
# you get an agent that clicks the window furniture.
INTERACTIVE_ROLES: frozenset[str] = frozenset({
    "button", "push button", "toggle button", "check box", "radio button",
    "menu item", "check menu item", "radio menu item", "page tab", "list item",
    "combo box", "entry", "text", "password text", "spin button", "slider",
    "link", "icon", "table cell", "tree item", "tool bar", "menu",
})

# Named things that are not clickable, but are worth *reading* (a verifier, or a
# VLM, needs the words on screen).
TEXT_ROLES: frozenset[str] = frozenset({
    "label", "paragraph", "heading", "static", "status bar", "text",
})


class AccessibleNode:
    """One node of an accessibility tree, reduced to what a decision needs."""

    __slots__ = ("role", "name", "x", "y", "width", "height", "depth")

    def __init__(self, role: str, name: str, x: int, y: int, width: int, height: int,
                 depth: int = 0) -> None:
        self.role = role
        self.name = name
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.depth = depth

    def __repr__(self) -> str:                                    # pragma: no cover
        return (f"AccessibleNode({self.role!r}, {self.name!r}, "
                f"{self.x},{self.y} {self.width}x{self.height})")

    def __eq__(self, other: object) -> bool:
        return isinstance(other, AccessibleNode) and self.__slots__ and all(
            getattr(self, slot) == getattr(other, slot) for slot in self.__slots__
        )

    def __hash__(self) -> int:
        return hash((self.role, self.name, self.x, self.y, self.width, self.height))


class TreeReader(Protocol):
    """Reads the accessibility tree of the whole desktop."""

    name: str

    def read(self) -> tuple[AccessibleNode, ...]: ...


class AtspiReader:
    """Reads the live AT-SPI desktop through PyGObject.

    Everything here is guarded: a machine without the typelib, without a session
    bus, or without the accessibility bridge raises :class:`DetectorUnavailable`
    with the fix, rather than returning an empty index that would look like "the
    screen is empty" when it is really "the bridge is off".
    """

    name = "atspi"

    def __init__(self, *, desktop: int = 0, max_depth: int = 10,
                 max_nodes: int = 600, timeout_ms: int = 500) -> None:
        self.desktop = desktop
        self.max_depth = max_depth
        self.max_nodes = max_nodes
        self.timeout_ms = timeout_ms

    @staticmethod
    def _atspi():
        try:
            import gi
        except ImportError as exc:                       # pragma: no cover - env
            raise _unavailable(
                "the AT-SPI reader needs PyGObject: pip install PyGObject, and "
                "apt install gir1.2-atspi-2.0 at-spi2-core"
            ) from exc
        try:
            gi.require_version("Atspi", "2.0")
            from gi.repository import Atspi
        except Exception as exc:                          # pragma: no cover - env
            raise _unavailable(
                "the Atspi typelib is missing: apt install gir1.2-atspi-2.0 "
                "at-spi2-core (and run inside a session bus)"
            ) from exc
        return Atspi

    def read(self) -> tuple[AccessibleNode, ...]:
        Atspi = self._atspi()
        try:
            desktop = Atspi.get_desktop(self.desktop)
        except Exception as exc:                          # pragma: no cover - env
            raise _unavailable(
                f"could not reach the accessibility bus: {type(exc).__name__}: {exc}. "
                "Start a session bus (dbus-run-session) and leave NO_AT_BRIDGE unset."
            ) from exc

        nodes: list[AccessibleNode] = []
        for index in range(desktop.get_child_count()):
            self._walk(Atspi, desktop.get_child_at_index(index), 0, nodes)
            if len(nodes) >= self.max_nodes:
                break
        return tuple(nodes)

    def _walk(self, Atspi, node: Any, depth: int, out: list[AccessibleNode]) -> None:
        if depth > self.max_depth or len(out) >= self.max_nodes:
            return
        try:
            role = (node.get_role_name() or "").strip()
            name = " ".join((node.get_name() or "").split())
        except Exception:
            return
        extents = self._extents(Atspi, node)
        if extents is not None:
            if not name and role.lower() in INTERACTIVE_ROLES:
                # A GTK list row is an unnamed container whose only child is the
                # text label: without this, the row is invisible to the index and
                # nothing in a list can be clicked at all.
                name = self._borrowed_name(Atspi, node)
            if name:
                x, y, width, height = extents
                out.append(AccessibleNode(role, name, x, y, width, height, depth))
        try:
            children = node.get_child_count()
        except Exception:
            children = 0
        for index in range(children):
            try:
                child = node.get_child_at_index(index)
            except Exception:
                continue
            self._walk(Atspi, child, depth + 1, out)

    def _borrowed_name(self, Atspi, node: Any, depth: int = 0) -> str:
        """The first text a child of this container offers, for an unnamed row."""
        if depth > 2:
            return ""
        try:
            children = node.get_child_count()
        except Exception:
            return ""
        for index in range(min(children, 8)):
            try:
                child = node.get_child_at_index(index)
                child_name = " ".join((child.get_name() or "").split())
            except Exception:
                continue
            if child_name:
                return child_name
            deeper = self._borrowed_name(Atspi, child, depth + 1)
            if deeper:
                return deeper
        return ""

    def _extents(self, Atspi, node: Any) -> tuple[int, int, int, int] | None:
        try:
            component = node.get_component_iface()
            if component is None:
                return None
            box = component.get_extents(Atspi.CoordType.SCREEN)
        except Exception:
            return None
        if box is None:
            return None
        return int(box.x), int(box.y), int(box.width), int(box.height)


# UI Automation control types -> the roles the rest of the harness cares about.
_UIA_ROLES: dict[str, str] = {
    "ButtonControl": "button",
    "SplitButtonControl": "button",
    "CheckBoxControl": "check box",
    "RadioButtonControl": "radio button",
    "HyperlinkControl": "link",
    "MenuItemControl": "menu item",
    "TabItemControl": "page tab",
    "ListItemControl": "list item",
    "EditControl": "entry",
    "DocumentControl": "entry",
    "ComboBoxControl": "combo box",
    "SliderControl": "slider",
    "SpinnerControl": "spin button",
    "TreeItemControl": "tree item",
    "DataItemControl": "table cell",
    "MenuControl": "menu",
    "ToolBarControl": "tool bar",
    "TextControl": "label",
    "ImageControl": "icon",
}


class UiaReader:
    """Reads the Windows UI Automation tree.

    **Status: implemented, not verified here.** There is no Windows in this
    repository's sandbox, so the walking and the mapping are exercised against a
    synthetic control tree in the tests, and the real backend is meant to be
    validated by the ``windows-latest`` CI job. :func:`accessibility_available`
    reports the truth on whatever machine it is asked on, so a non-Windows machine
    says "uia is Windows-only" rather than returning an empty index.

    ``provider`` exists for exactly that reason: the tests inject a fake root and
    assert the mapping, the filtering and the geometry without a Windows box.
    """

    name = "uia"

    def __init__(self, *, max_depth: int = 12, max_nodes: int = 600,
                 provider: Any = None) -> None:
        self.max_depth = max_depth
        self.max_nodes = max_nodes
        self._provider = provider

    @staticmethod
    def _automation():
        import platform

        if platform.system() != "Windows":
            raise _unavailable(
                f"UI Automation is Windows-only (this is {platform.system()}); "
                "use the AT-SPI detector on Linux"
            )
        try:
            import uiautomation  # noqa: F401
        except ImportError as exc:
            raise _unavailable(
                "UI Automation needs the uiautomation package: pip install uiautomation"
            ) from exc
        return uiautomation

    def _root(self) -> Any:
        if self._provider is not None:
            return self._provider()
        return self._automation().GetRootControl()

    def read(self) -> tuple[AccessibleNode, ...]:
        nodes: list[AccessibleNode] = []
        try:
            root = self._root()
        except Exception as exc:
            if isinstance(exc, DetectorUnavailable):
                raise
            raise _unavailable(f"could not reach UI Automation: {type(exc).__name__}: {exc}")
        self._walk(root, 0, nodes)
        return tuple(nodes)

    def _walk(self, control: Any, depth: int, out: list[AccessibleNode]) -> None:
        if depth > self.max_depth or len(out) >= self.max_nodes:
            return
        name = " ".join(str(getattr(control, "Name", "") or "").split())
        control_type = str(getattr(control, "ControlTypeName", "") or "")
        role = _UIA_ROLES.get(control_type, control_type.replace("Control", "").lower())
        box = getattr(control, "BoundingRectangle", None)
        if box is not None and name:
            left, top, right, bottom = (int(box.left), int(box.top),
                                        int(box.right), int(box.bottom))
            out.append(AccessibleNode(role, name, left, top, right - left, bottom - top, depth))
        try:
            children = list(control.GetChildren()) if hasattr(control, "GetChildren") else []
        except Exception:
            children = []
        for child in children:
            self._walk(child, depth + 1, out)


class UiaDetector:
    """The tree detector over Windows UI Automation.

    Thin on purpose: it is the same filtering as :class:`AtspiDetector`, with a
    ``uia`` reader. Kept as its own name so the CLI can say ``--detector uia``.
    """

    name = "uia"
    absolute_coordinates = True

    def __init__(self, reader: TreeReader | None = None, **kwargs: Any) -> None:
        self._detector = AtspiDetector(reader=reader or UiaReader(), name="uia", **kwargs)

    def detect(self, image: Any = None) -> tuple[Element, ...]:
        return self._detector.detect(image)

    def coordinates_of_last_detect(self) -> str:
        return "screen"

    def text(self) -> str:
        return self._detector.text()


def _unavailable(detail: str) -> DetectorUnavailable:
    return DetectorUnavailable(detail)


@dataclass
class AtspiDetector:
    """Turns the accessibility tree into the harness's element index.

    ``kind`` is the AT-SPI role (``button``, ``check box``, ...), so a policy can
    distinguish a button from a text field, and ``absolute_coordinates`` tells
    :class:`IndexedScreen` these are already screen pixels.
    """

    name: str = "atspi"
    absolute_coordinates: bool = True
    reader: TreeReader | None = field(default=None)
    interactive_roles: frozenset[str] = INTERACTIVE_ROLES
    include_text: bool = False
    min_label_chars: int = 1
    max_label_chars: int = 64
    min_size: int = 4
    max_nodes: int = 120
    _last_roles: tuple[str, ...] = field(default=(), init=False, repr=False)

    def coordinates_of_last_detect(self) -> str:
        """Elements come back in absolute screen pixels, always."""
        return "screen"

    def detect(self, image: Any = None) -> tuple[Element, ...]:
        """Index the live tree. The image is accepted and deliberately ignored."""
        reader = self.reader if self.reader is not None else AtspiReader()
        nodes = reader.read()
        wanted = self.interactive_roles | (TEXT_ROLES if self.include_text else frozenset())
        elements: list[Element] = []
        seen: set[tuple[str, int, int]] = set()
        for node in nodes:
            role = node.role.lower()
            if role not in wanted:
                continue
            label = node.name
            if not label or not (self.min_label_chars <= len(label) <= self.max_label_chars):
                continue
            if node.width < self.min_size or node.height < self.min_size:
                continue
            if node.x < 0 or node.y < 0:                  # off screen
                continue
            key = (label.lower(), node.x, node.y)
            if key in seen:
                continue
            seen.add(key)
            elements.append(Element(label=label, kind=role, x=node.x, y=node.y,
                                    width=node.width, height=node.height))
            if len(elements) >= self.max_nodes:
                break
        self._last_roles = tuple(element.kind for element in elements)
        return tuple(elements)

    def text(self) -> str:
        """The words on screen, for a verifier - straight from the tree."""
        reader = self.reader if self.reader is not None else AtspiReader()
        return " ".join(node.name for node in reader.read()
                        if node.name and node.role.lower() in (TEXT_ROLES | self.interactive_roles))


def accessibility_available(backend: str = "atspi") -> tuple[bool, str]:
    """(available, detail) for preflight and the doctor, honestly reported."""
    if backend == "uia":
        try:
            automation = UiaReader._automation()
        except Exception as exc:
            return False, str(exc)
        try:
            root = automation.GetRootControl()
            children = root.GetChildren()
        except Exception as exc:
            return False, f"uiautomation is installed but the tree is unreachable: {exc}"
        return True, f"UI Automation ({len(list(children))} top-level window(s))"
    if backend != "atspi":
        return False, f"unknown accessibility backend {backend!r}"
    try:
        Atspi = AtspiReader._atspi()
    except Exception as exc:
        return False, str(exc)
    try:
        desktop = Atspi.get_desktop(0)
        applications = desktop.get_child_count()
    except Exception as exc:
        return False, f"Atspi is installed but the a11y bus is unreachable: {exc}"
    return True, f"AT-SPI via PyGObject ({applications} application(s) on the desktop)"


__all__ = [
    "INTERACTIVE_ROLES",
    "TEXT_ROLES",
    "AccessibleNode",
    "AtspiDetector",
    "AtspiReader",
    "TreeReader",
    "UiaDetector",
    "UiaReader",
    "accessibility_available",
]