From 02ef4d23f2ba1bbf87f30db9e987a9eae131ba7b Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Mon, 6 Jul 2026 13:22:15 +0800 Subject: [PATCH] Implement Apex Agent MVP scaffold --- api/__init__.py | 2 + api/errors.py | 31 ++++ api/mcp.py | 162 +++++++++++++++++++ api/rest.py | 109 +++++++++++++ core/__init__.py | 2 + core/device_manager.py | 155 +++++++++++++++++++ core/driver.py | 68 ++++++++ core/errors.py | 30 ++++ core/models.py | 180 ++++++++++++++++++++++ core/wda_driver.py | 155 +++++++++++++++++++ openspec/changes/apex-agent-mvp/tasks.md | 58 +++++++ pyproject.toml | 31 ++++ runtime/__init__.py | 2 + runtime/context.py | 30 ++++ runtime/executor.py | 148 ++++++++++++++++++ runtime/planner.py | 40 +++++ runtime/task.py | 151 ++++++++++++++++++ storage/__init__.py | 2 + storage/artifact_store.py | 71 +++++++++ storage/task_metadata.py | 97 ++++++++++++ storage/timeline.py | 62 ++++++++ tests/__init__.py | 2 + tests/fakes.py | 85 ++++++++++ tests/test_device_manager.py | 33 ++++ tests/test_executor.py | 26 ++++ tests/test_find_helpers.py | 46 ++++++ tests/test_mcp.py | 38 +++++ tests/test_rest_api.py | 44 ++++++ tests/test_scene_builder.py | 45 ++++++ tests/test_smoke.py | 9 ++ tests/test_task_loop.py | 78 ++++++++++ tests/test_timeline.py | 33 ++++ tests/test_wda_integration.py | 28 ++++ tools/__init__.py | 2 + tools/_device.py | 13 ++ tools/describe_screen.py | 27 ++++ tools/find_icon.py | 19 +++ tools/find_text.py | 59 +++++++ tools/input_text.py | 15 ++ tools/launch_app.py | 25 +++ tools/screenshot.py | 13 ++ tools/swipe.py | 31 ++++ tools/tap.py | 16 ++ tools/ui_tree.py | 26 ++++ vision/__init__.py | 2 + vision/icon_detector.py | 36 +++++ vision/ocr.py | 188 +++++++++++++++++++++++ vision/scene_builder.py | 91 +++++++++++ vision/ui_parser.py | 151 ++++++++++++++++++ 49 files changed, 2767 insertions(+) create mode 100644 api/__init__.py create mode 100644 api/errors.py create mode 100644 api/mcp.py create mode 100644 api/rest.py create mode 100644 core/__init__.py create mode 100644 core/device_manager.py create mode 100644 core/driver.py create mode 100644 core/errors.py create mode 100644 core/models.py create mode 100644 core/wda_driver.py create mode 100644 openspec/changes/apex-agent-mvp/tasks.md create mode 100644 pyproject.toml create mode 100644 runtime/__init__.py create mode 100644 runtime/context.py create mode 100644 runtime/executor.py create mode 100644 runtime/planner.py create mode 100644 runtime/task.py create mode 100644 storage/__init__.py create mode 100644 storage/artifact_store.py create mode 100644 storage/task_metadata.py create mode 100644 storage/timeline.py create mode 100644 tests/__init__.py create mode 100644 tests/fakes.py create mode 100644 tests/test_device_manager.py create mode 100644 tests/test_executor.py create mode 100644 tests/test_find_helpers.py create mode 100644 tests/test_mcp.py create mode 100644 tests/test_rest_api.py create mode 100644 tests/test_scene_builder.py create mode 100644 tests/test_smoke.py create mode 100644 tests/test_task_loop.py create mode 100644 tests/test_timeline.py create mode 100644 tests/test_wda_integration.py create mode 100644 tools/__init__.py create mode 100644 tools/_device.py create mode 100644 tools/describe_screen.py create mode 100644 tools/find_icon.py create mode 100644 tools/find_text.py create mode 100644 tools/input_text.py create mode 100644 tools/launch_app.py create mode 100644 tools/screenshot.py create mode 100644 tools/swipe.py create mode 100644 tools/tap.py create mode 100644 tools/ui_tree.py create mode 100644 vision/__init__.py create mode 100644 vision/icon_detector.py create mode 100644 vision/ocr.py create mode 100644 vision/scene_builder.py create mode 100644 vision/ui_parser.py diff --git a/api/__init__.py b/api/__init__.py new file mode 100644 index 0000000..50675b5 --- /dev/null +++ b/api/__init__.py @@ -0,0 +1,2 @@ +"""MCP and REST API surfaces for Apex Agent.""" + diff --git a/api/errors.py b/api/errors.py new file mode 100644 index 0000000..0b1d6c9 --- /dev/null +++ b/api/errors.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +from core.errors import ( + DeviceNotFoundError, + DeviceOfflineError, + ElementNotFoundError, + DriverError, +) + + +def semantic_error(exc: Exception) -> str: + if isinstance(exc, DeviceOfflineError): + return "device offline" + if isinstance(exc, DeviceNotFoundError): + return "device not found" + if isinstance(exc, ElementNotFoundError): + return "element not found" + if isinstance(exc, DriverError): + return "driver error" + return "operation failed" + + +def call_with_semantic_errors(func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + try: + return func(*args, **kwargs) + except Exception as exc: + return {"ok": False, "error": semantic_error(exc)} + diff --git a/api/mcp.py b/api/mcp.py new file mode 100644 index 0000000..39005d6 --- /dev/null +++ b/api/mcp.py @@ -0,0 +1,162 @@ +import base64 +from collections.abc import Callable +from typing import Any + +from api.errors import call_with_semantic_errors +from core.device_manager import DEFAULT_MANAGER, DeviceManager +from tools.describe_screen import describe_screen +from tools.find_icon import find_icon_on_screen +from tools.find_text import find_text_on_screen +from tools.input_text import input_text +from tools.launch_app import launch_app +from tools.screenshot import take_screenshot +from tools.swipe import swipe +from tools.tap import tap +from tools.ui_tree import get_ui_tree + + +def tool_handlers( + *, + manager: DeviceManager | None = None, +) -> dict[str, Callable[..., Any]]: + device_manager = manager or DEFAULT_MANAGER + + def _screenshot(device_id: str | None = None) -> dict[str, Any]: + image = take_screenshot(device_id, manager=device_manager) + return { + "ok": True, + "image_base64": base64.b64encode(image).decode("ascii"), + "mime_type": "image/png", + } + + return { + "take_screenshot": lambda device_id=None: call_with_semantic_errors( + _screenshot, + device_id, + ), + "tap": lambda x, y, device_id=None: call_with_semantic_errors( + tap, + x, + y, + device_id=device_id, + manager=device_manager, + ), + "swipe": lambda start_x, start_y, end_x, end_y, duration_ms=500, device_id=None: call_with_semantic_errors( + swipe, + start_x, + start_y, + end_x, + end_y, + duration_ms=duration_ms, + device_id=device_id, + manager=device_manager, + ), + "input_text": lambda text, device_id=None: call_with_semantic_errors( + input_text, + text, + device_id=device_id, + manager=device_manager, + ), + "launch_app": lambda app_id, device_id=None: call_with_semantic_errors( + launch_app, + app_id, + device_id=device_id, + manager=device_manager, + ), + "find_text": lambda query, device_id=None: call_with_semantic_errors( + find_text_on_screen, + query, + device_id=device_id, + manager=device_manager, + ), + "find_icon": lambda name, device_id=None: call_with_semantic_errors( + find_icon_on_screen, + name, + device_id=device_id, + manager=device_manager, + ), + "get_ui_tree": lambda device_id=None: call_with_semantic_errors( + get_ui_tree, + device_id, + manager=device_manager, + ), + "describe_screen": lambda device_id=None: call_with_semantic_errors( + lambda: describe_screen(device_id, manager=device_manager).to_dict() + ), + "list_devices": lambda: [ + device.to_dict() for device in device_manager.list_devices() + ], + "device_status": lambda device_id: call_with_semantic_errors( + lambda: {"device_id": device_id, "status": device_manager.status(device_id)} + ), + } + + +def create_mcp_server(*, manager: DeviceManager | None = None) -> Any: + try: + from mcp.server.fastmcp import FastMCP + except ImportError as exc: + raise RuntimeError("mcp SDK is not installed") from exc + + handlers = tool_handlers(manager=manager) + server = FastMCP("apex-agent") + + @server.tool(name="take_screenshot") + def _take_screenshot(device_id: str | None = None) -> dict[str, Any]: + return handlers["take_screenshot"](device_id=device_id) + + @server.tool(name="tap") + def _tap(x: float, y: float, device_id: str | None = None) -> dict[str, Any]: + return handlers["tap"](x=x, y=y, device_id=device_id) + + @server.tool(name="swipe") + def _swipe( + start_x: float, + start_y: float, + end_x: float, + end_y: float, + duration_ms: int = 500, + device_id: str | None = None, + ) -> dict[str, Any]: + return handlers["swipe"]( + start_x=start_x, + start_y=start_y, + end_x=end_x, + end_y=end_y, + duration_ms=duration_ms, + device_id=device_id, + ) + + @server.tool(name="input_text") + def _input_text(text: str, device_id: str | None = None) -> dict[str, Any]: + return handlers["input_text"](text=text, device_id=device_id) + + @server.tool(name="launch_app") + def _launch_app(app_id: str, device_id: str | None = None) -> dict[str, Any]: + return handlers["launch_app"](app_id=app_id, device_id=device_id) + + @server.tool(name="find_text") + def _find_text(query: str, device_id: str | None = None) -> dict[str, Any]: + return handlers["find_text"](query=query, device_id=device_id) + + @server.tool(name="find_icon") + def _find_icon(name: str, device_id: str | None = None) -> dict[str, Any]: + return handlers["find_icon"](name=name, device_id=device_id) + + @server.tool(name="get_ui_tree") + def _get_ui_tree(device_id: str | None = None) -> Any: + return handlers["get_ui_tree"](device_id=device_id) + + @server.tool(name="describe_screen") + def _describe_screen(device_id: str | None = None) -> dict[str, Any]: + return handlers["describe_screen"](device_id=device_id) + + @server.tool(name="list_devices") + def _list_devices() -> list[dict[str, Any]]: + return handlers["list_devices"]() + + @server.tool(name="device_status") + def _device_status(device_id: str) -> dict[str, Any]: + return handlers["device_status"](device_id=device_id) + + return server diff --git a/api/rest.py b/api/rest.py new file mode 100644 index 0000000..8dcacbc --- /dev/null +++ b/api/rest.py @@ -0,0 +1,109 @@ +from typing import Any + +from api.errors import semantic_error +from core.device_manager import DEFAULT_MANAGER, DeviceManager +from core.models import Task +from runtime.executor import Executor, default_tool_registry +from runtime.task import TaskRunner +from storage.task_metadata import TaskMetadataStore +from tools.launch_app import launch_app +from tools.screenshot import take_screenshot +from tools.tap import tap + + +def create_app( + *, + manager: DeviceManager | None = None, + task_runner: TaskRunner | None = None, + metadata_store: TaskMetadataStore | None = None, +) -> Any: + from fastapi import BackgroundTasks, FastAPI, HTTPException + from pydantic import BaseModel + + device_manager = manager or DEFAULT_MANAGER + store = metadata_store or TaskMetadataStore() + runner = task_runner or TaskRunner( + metadata_store=store, + executor=Executor(tools=default_tool_registry(manager=device_manager)), + ) + app = FastAPI(title="Apex Agent API") + + class TapRequest(BaseModel): + x: float + y: float + + class ScreenshotResponse(BaseModel): + image_base64: str + mime_type: str = "image/png" + + class LaunchRequest(BaseModel): + app_id: str + + class AgentTaskRequest(BaseModel): + goal: str + device_id: str + + @app.get("/devices") + def devices() -> list[dict[str, Any]]: + return [device.to_dict() for device in device_manager.list_devices()] + + @app.post("/devices/{device_id}/tap") + def tap_device(device_id: str, request: TapRequest) -> dict[str, Any]: + return _raise_semantic( + lambda: tap( + request.x, + request.y, + device_id=device_id, + manager=device_manager, + ) + ) + + @app.post("/devices/{device_id}/screenshot") + def screenshot_device(device_id: str) -> dict[str, str]: + import base64 + + image = _raise_semantic( + lambda: take_screenshot(device_id, manager=device_manager) + ) + response = ScreenshotResponse( + image_base64=base64.b64encode(image).decode("ascii") + ) + return response.model_dump() + + @app.post("/devices/{device_id}/launch") + def launch_device(device_id: str, request: LaunchRequest) -> dict[str, Any]: + return _raise_semantic( + lambda: launch_app( + request.app_id, + device_id=device_id, + manager=device_manager, + ) + ) + + @app.post("/agent/task") + def start_task( + request: AgentTaskRequest, + background_tasks: BackgroundTasks, + ) -> dict[str, str]: + task = Task(goal=request.goal, device_id=request.device_id) + store.create_task(task) + background_tasks.add_task(runner.run, task) + return {"task_id": task.id, "status": task.status} + + @app.get("/task/{task_id}") + def get_task(task_id: str) -> dict[str, Any]: + task = store.get_task(task_id) + if task is None: + raise HTTPException(status_code=404, detail="task not found") + return task + + return app + + +def _raise_semantic(func: Any) -> Any: + from fastapi import HTTPException + + try: + return func() + except Exception as exc: + raise HTTPException(status_code=400, detail=semantic_error(exc)) from exc diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000..e24d38b --- /dev/null +++ b/core/__init__.py @@ -0,0 +1,2 @@ +"""Core device and domain abstractions for Apex Agent.""" + diff --git a/core/device_manager.py b/core/device_manager.py new file mode 100644 index 0000000..3f76be7 --- /dev/null +++ b/core/device_manager.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import replace +from threading import RLock +from time import sleep + +from core.driver import Driver +from core.errors import ( + DeviceBusyError, + DeviceNotFoundError, + DeviceOfflineError, + DriverError, +) +from core.models import Device, DeviceStatus + +DriverFactory = Callable[[], Driver] + + +class DeviceManager: + def __init__(self) -> None: + self._devices: dict[str, Device] = {} + self._factories: dict[str, DriverFactory] = {} + self._drivers: dict[str, Driver] = {} + self._lock = RLock() + + def register_device( + self, + device_id: str, + driver_factory: DriverFactory, + *, + name: str | None = None, + driver_type: str = "wda", + connection_info: dict[str, object] | None = None, + status: DeviceStatus = "idle", + ) -> Device: + with self._lock: + device = Device( + id=device_id, + name=name, + status=status, + driver_type=driver_type, + connection_info=dict(connection_info or {}), + ) + self._devices[device_id] = device + self._factories[device_id] = driver_factory + self._drivers.pop(device_id, None) + return replace(device) + + def unregister_device(self, device_id: str) -> None: + with self._lock: + driver = self._drivers.pop(device_id, None) + if driver: + driver.disconnect() + self._devices.pop(device_id, None) + self._factories.pop(device_id, None) + + def list_devices(self) -> list[Device]: + with self._lock: + return [replace(device) for device in self._devices.values()] + + def status(self, device_id: str) -> DeviceStatus: + with self._lock: + return self._device(device_id).status + + def connect( + self, + device_id: str, + *, + max_retries: int = 2, + retry_backoff_seconds: float = 0.25, + ) -> Driver: + if max_retries < 1: + raise ValueError("max_retries must be at least 1") + + last_error: Exception | None = None + for attempt in range(1, max_retries + 1): + with self._lock: + device = self._device(device_id) + if device.status == "busy" and device_id in self._drivers: + return self._drivers[device_id] + if device.status == "busy": + raise DeviceBusyError(f"device {device_id} is busy") + driver = self._factories[device_id]() + + try: + driver.connect() + except Exception as exc: + last_error = exc + with self._lock: + status: DeviceStatus = ( + "offline" if attempt == max_retries else "error" + ) + self._set_status(device_id, status) + if attempt < max_retries: + sleep(retry_backoff_seconds) + continue + + with self._lock: + self._drivers[device_id] = driver + self._set_status(device_id, "busy") + return driver + + raise DeviceOfflineError( + f"device {device_id} is offline" + ) from last_error + + def disconnect(self, device_id: str) -> None: + with self._lock: + self._device(device_id) + driver = self._drivers.pop(device_id, None) + if driver: + driver.disconnect() + with self._lock: + self._set_status(device_id, "idle") + + def mark_error(self, device_id: str, *, offline: bool = False) -> None: + with self._lock: + self._device(device_id) + self._drivers.pop(device_id, None) + self._set_status(device_id, "offline" if offline else "error") + + def active_driver(self, device_id: str | None = None) -> Driver: + with self._lock: + if device_id is None: + busy_ids = [ + known_id + for known_id, device in self._devices.items() + if device.status == "busy" and known_id in self._drivers + ] + if len(busy_ids) != 1: + raise DriverError( + "exactly one connected device is required when no device_id is given" + ) + device_id = busy_ids[0] + + self._device(device_id) + driver = self._drivers.get(device_id) + if not driver: + raise DeviceOfflineError(f"device {device_id} is not connected") + return driver + + def _device(self, device_id: str) -> Device: + try: + return self._devices[device_id] + except KeyError as exc: + raise DeviceNotFoundError(f"unknown device {device_id}") from exc + + def _set_status(self, device_id: str, status: DeviceStatus) -> None: + device = self._devices[device_id] + self._devices[device_id] = replace(device, status=status) + + +DEFAULT_MANAGER = DeviceManager() + diff --git a/core/driver.py b/core/driver.py new file mode 100644 index 0000000..a6ec952 --- /dev/null +++ b/core/driver.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class Driver(ABC): + """Driver-independent device capability interface. + + Implementations may hold a live connection handle, but no task or business + state belongs here. + """ + + @abstractmethod + def connect(self) -> None: + """Open the device connection.""" + + @abstractmethod + def disconnect(self) -> None: + """Close the device connection.""" + + @abstractmethod + def screenshot(self) -> bytes: + """Return the current screen as image bytes.""" + + @abstractmethod + def tap(self, x: float, y: float) -> None: + """Tap the screen at the given coordinates.""" + + @abstractmethod + def swipe( + self, + start_x: float, + start_y: float, + end_x: float, + end_y: float, + duration_ms: int = 500, + ) -> None: + """Swipe between two screen coordinates.""" + + @abstractmethod + def input(self, text: str) -> None: + """Input text into the current focused field.""" + + @abstractmethod + def launch(self, app_id: str) -> None: + """Launch an app by bundle id or driver-supported app identifier.""" + + @abstractmethod + def terminate(self, app_id: str) -> None: + """Terminate an app by bundle id or driver-supported app identifier.""" + + @abstractmethod + def tree(self) -> Any: + """Return the raw UI tree from the device driver.""" + + @abstractmethod + def home(self) -> None: + """Press the device home button.""" + + @abstractmethod + def lock(self) -> None: + """Lock the device.""" + + @abstractmethod + def unlock(self) -> None: + """Unlock the device.""" + diff --git a/core/errors.py b/core/errors.py new file mode 100644 index 0000000..b3fd5da --- /dev/null +++ b/core/errors.py @@ -0,0 +1,30 @@ +from __future__ import annotations + + +class ApexAgentError(Exception): + """Base error for semantic errors surfaced above the driver layer.""" + + +class DriverError(ApexAgentError): + """A driver operation failed.""" + + +class DeviceNotFoundError(ApexAgentError): + """A requested device id is unknown.""" + + +class DeviceOfflineError(ApexAgentError): + """A device is unavailable or lost its connection.""" + + +class DeviceBusyError(ApexAgentError): + """A device is already in use.""" + + +class ElementNotFoundError(ApexAgentError): + """A requested screen element was not found.""" + + +class TaskFailedError(ApexAgentError): + """A task failed before reaching its goal.""" + diff --git a/core/models.py b/core/models.py new file mode 100644 index 0000000..3e3710c --- /dev/null +++ b/core/models.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Any, Literal +from uuid import uuid4 + +DeviceStatus = Literal["idle", "busy", "offline", "error"] +TaskStatus = Literal["created", "running", "completed", "failed", "cancelled"] +StepStatus = Literal["pending", "running", "completed", "failed"] + + +def utc_now() -> datetime: + return datetime.now(UTC) + + +@dataclass(frozen=True) +class Bounds: + x: float + y: float + width: float + height: float + + @property + def right(self) -> float: + return self.x + self.width + + @property + def bottom(self) -> float: + return self.y + self.height + + @property + def center(self) -> tuple[float, float]: + return (self.x + self.width / 2, self.y + self.height / 2) + + def to_dict(self) -> dict[str, float]: + return { + "x": self.x, + "y": self.y, + "width": self.width, + "height": self.height, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "Bounds": + return cls( + x=float(data["x"]), + y=float(data["y"]), + width=float(data["width"]), + height=float(data["height"]), + ) + + +@dataclass +class Device: + id: str + status: DeviceStatus = "idle" + name: str | None = None + driver_type: str = "wda" + connection_info: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "name": self.name, + "status": self.status, + "driver_type": self.driver_type, + "connection_info": dict(self.connection_info), + } + + +@dataclass +class SceneElement: + id: str + type: str + bounds: Bounds + text: str | None = None + confidence: float | None = None + source: str | None = None + + @property + def center(self) -> tuple[float, float]: + return self.bounds.center + + def to_dict(self) -> dict[str, Any]: + data: dict[str, Any] = { + "id": self.id, + "type": self.type, + "text": self.text, + "bounds": self.bounds.to_dict(), + "confidence": self.confidence, + } + if self.source: + data["source"] = self.source + return data + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "SceneElement": + return cls( + id=str(data["id"]), + type=str(data.get("type") or "unknown"), + text=data.get("text"), + bounds=Bounds.from_dict(data["bounds"]), + confidence=data.get("confidence"), + source=data.get("source"), + ) + + +@dataclass +class Scene: + width: int + height: int + elements: list[SceneElement] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return { + "screen": {"width": self.width, "height": self.height}, + "elements": [element.to_dict() for element in self.elements], + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "Scene": + screen = data.get("screen") or {} + return cls( + width=int(screen.get("width") or data.get("width") or 0), + height=int(screen.get("height") or data.get("height") or 0), + elements=[ + SceneElement.from_dict(element) + for element in data.get("elements", []) + ], + ) + + +@dataclass +class Task: + goal: str + device_id: str + id: str = field(default_factory=lambda: uuid4().hex) + status: TaskStatus = "created" + created_at: datetime = field(default_factory=utc_now) + updated_at: datetime = field(default_factory=utc_now) + completed_at: datetime | None = None + failure_reason: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "goal": self.goal, + "device_id": self.device_id, + "status": self.status, + "created_at": self.created_at.isoformat(), + "updated_at": self.updated_at.isoformat(), + "completed_at": self.completed_at.isoformat() + if self.completed_at + else None, + "failure_reason": self.failure_reason, + } + + +@dataclass +class Step: + description: str + action: str + args: dict[str, Any] = field(default_factory=dict) + id: str = field(default_factory=lambda: uuid4().hex) + status: StepStatus = "pending" + result: dict[str, Any] | None = None + error: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "description": self.description, + "action": self.action, + "args": dict(self.args), + "status": self.status, + "result": self.result, + "error": self.error, + } + diff --git a/core/wda_driver.py b/core/wda_driver.py new file mode 100644 index 0000000..c3c4d27 --- /dev/null +++ b/core/wda_driver.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from core.driver import Driver +from core.errors import DeviceOfflineError, DriverError + + +@dataclass(frozen=True) +class WDADriverConfig: + server_url: str = "http://127.0.0.1:4723" + platform_name: str = "iOS" + automation_name: str = "XCUITest" + device_name: str | None = None + udid: str | None = None + wda_local_port: int | None = None + no_reset: bool = True + extra_capabilities: dict[str, Any] = field(default_factory=dict) + + +class WDADriver(Driver): + def __init__(self, config: WDADriverConfig | None = None) -> None: + self.config = config or WDADriverConfig() + self._client: Any | None = None + + def connect(self) -> None: + try: + from appium import webdriver + from appium.options.ios import XCUITestOptions + from appium.webdriver.client_config import AppiumClientConfig + except ImportError as exc: + raise DriverError("Appium Python client is not installed") from exc + + capabilities: dict[str, Any] = { + "platformName": self.config.platform_name, + "automationName": self.config.automation_name, + "noReset": self.config.no_reset, + **self.config.extra_capabilities, + } + if self.config.device_name: + capabilities["deviceName"] = self.config.device_name + if self.config.udid: + capabilities["udid"] = self.config.udid + if self.config.wda_local_port: + capabilities["wdaLocalPort"] = self.config.wda_local_port + + options = XCUITestOptions().load_capabilities(capabilities) + client_config = AppiumClientConfig(remote_server_addr=self.config.server_url) + try: + self._client = webdriver.Remote( + options=options, + client_config=client_config, + ) + except Exception as exc: + self._client = None + raise DeviceOfflineError("device offline") from exc + + def disconnect(self) -> None: + client = self._require_client() + try: + client.quit() + finally: + self._client = None + + def screenshot(self) -> bytes: + client = self._require_client() + try: + return client.get_screenshot_as_png() + except Exception as exc: + raise DriverError("screenshot failed") from exc + + def tap(self, x: float, y: float) -> None: + client = self._require_client() + try: + client.execute_script("mobile: tap", {"x": x, "y": y}) + except Exception as exc: + raise DriverError("tap failed") from exc + + def swipe( + self, + start_x: float, + start_y: float, + end_x: float, + end_y: float, + duration_ms: int = 500, + ) -> None: + client = self._require_client() + try: + client.execute_script( + "mobile: dragFromToForDuration", + { + "fromX": start_x, + "fromY": start_y, + "toX": end_x, + "toY": end_y, + "duration": duration_ms / 1000, + }, + ) + except Exception as exc: + raise DriverError("swipe failed") from exc + + def input(self, text: str) -> None: + client = self._require_client() + try: + client.switch_to.active_element.send_keys(text) + except Exception as exc: + raise DriverError("text input failed") from exc + + def launch(self, app_id: str) -> None: + client = self._require_client() + try: + client.activate_app(app_id) + except Exception as exc: + raise DriverError("app launch failed") from exc + + def terminate(self, app_id: str) -> None: + client = self._require_client() + try: + client.terminate_app(app_id) + except Exception as exc: + raise DriverError("app terminate failed") from exc + + def tree(self) -> str: + client = self._require_client() + try: + return client.page_source + except Exception as exc: + raise DriverError("ui tree retrieval failed") from exc + + def home(self) -> None: + client = self._require_client() + try: + client.execute_script("mobile: pressButton", {"name": "home"}) + except Exception as exc: + raise DriverError("home failed") from exc + + def lock(self) -> None: + client = self._require_client() + try: + client.lock() + except Exception as exc: + raise DriverError("lock failed") from exc + + def unlock(self) -> None: + client = self._require_client() + try: + client.unlock() + except Exception as exc: + raise DriverError("unlock failed") from exc + + def _require_client(self) -> Any: + if self._client is None: + raise DeviceOfflineError("device offline") + return self._client diff --git a/openspec/changes/apex-agent-mvp/tasks.md b/openspec/changes/apex-agent-mvp/tasks.md new file mode 100644 index 0000000..8092580 --- /dev/null +++ b/openspec/changes/apex-agent-mvp/tasks.md @@ -0,0 +1,58 @@ +## 1. Project scaffolding + +- [x] 1.1 Create the package layout: `core/`, `tools/`, `vision/`, `runtime/`, `api/`, `storage/`, `tests/` (each with `__init__.py`) +- [x] 1.2 Add MVP dependencies to `pyproject.toml` (Appium Python client, an OCR engine e.g. PaddleOCR, FastAPI, an MCP server SDK, SQLite driver, test framework) +- [x] 1.3 Add `core/models.py` with shared data classes: `Device`, `Scene`, `SceneElement`, `Task`, `Step` +- [x] 1.4 Set up a basic test runner config and a smoke test that imports every new package + +## 2. Device management (capability: device-management) + +- [x] 2.1 Define the abstract `Driver` interface in `core/driver.py` (`connect`, `disconnect`, `screenshot`, `tap`, `swipe`, `input`, `launch`, `terminate`, `tree`, `home`, `lock`, `unlock`) with no stored business state +- [x] 2.2 Implement `core/device_manager.py`: device discovery/registration, status tracking (`idle`/`busy`/`offline`/`error`), `list_devices()`, `connect()`, `disconnect()`, `status()` +- [x] 2.3 Implement `core/wda_driver.py` (`WDADriver`) using the Appium Python client against WebDriverAgent, covering all `Driver` methods +- [x] 2.4 Add connection-loss handling: mark device `offline`/`error` and bound retries on `connect()` instead of hanging +- [x] 2.5 Write unit tests for `DeviceManager` state transitions using a fake/mock `Driver` +- [x] 2.6 Write an integration test (skippable without hardware) that connects to a real/simulated device via `WDADriver` and takes a screenshot + +## 3. Perception pipeline (capability: scene-perception) + +- [x] 3.1 Implement `vision/ocr.py`: run OCR over a screenshot and return text boxes with bounds/confidence +- [x] 3.2 Implement `vision/ui_parser.py`: parse the raw UI tree returned by `Driver.tree()` into a normalized element list (type, text, bounds) +- [x] 3.3 Implement `vision/scene_builder.py`: fuse OCR output + parsed UI tree into a single `Scene`, deduping overlapping elements by bounding-box IoU (prefer UI-tree bounds/type on overlap) +- [x] 3.4 Implement `vision/icon_detector.py` as a minimal stub (template match or "not found") satisfying the `find_icon` contract +- [x] 3.5 Implement `tools/screenshot.py`, `tools/tap.py`, `tools/swipe.py`, `tools/input_text.py`, `tools/launch_app.py`, `tools/ui_tree.py` as thin wrappers over the active device's `Driver` +- [x] 3.6 Implement `tools/describe_screen.py` (returns the fused `Scene`) and `find_text`/`find_icon` helpers that search a `Scene` and return coordinates or "not found" +- [x] 3.7 Write unit tests for `scene_builder` dedup logic using fixture screenshots/tree/OCR data +- [x] 3.8 Write unit tests for `find_text`/`find_icon` against fixture `Scene` objects (found and not-found cases) + +## 4. Agent runtime (capability: agent-runtime) + +- [x] 4.1 Implement `runtime/context.py`: per-task in-memory context holding Scene history and executed step results +- [x] 4.2 Implement `runtime/planner.py`: given a goal + current `Scene` (+ context), produce an ordered list of intended next steps +- [x] 4.3 Implement `runtime/executor.py`: turn a planned step into `tools/` calls, with retry/backoff and wait-for-element handling, bounded by configurable max-retries and max-steps +- [x] 4.4 Implement `runtime/task.py`: drive the Observe→Think→Act→Observe loop for a task from start to completion/failure, wiring Planner + Executor + Context together +- [x] 4.5 Add max-step/max-retry ceiling handling that fails the task with a clear reason instead of looping indefinitely +- [x] 4.6 Write unit tests for `Executor` retry/backoff behavior using a fake tool that fails N times then succeeds +- [x] 4.7 Write an end-to-end test of the loop against a mocked `Driver`/`Scene` sequence simulating "open app → find search → tap → type" + +## 5. Task memory / timeline (capability: task-memory) + +- [x] 5.1 Implement `storage/artifact_store.py`: write screenshots and per-step JSON (Scene, prompt, tool call, result) to `tasks/history//NNN.{png,json}` +- [x] 5.2 Implement `storage/timeline.py`: append-and-read API for a task's ordered step records, backed by `artifact_store` +- [x] 5.3 Add SQLite-backed task metadata storage (task id, device id, status, start/end timestamps) with create/update/query functions +- [x] 5.4 Wire `runtime/task.py` (from section 4) to write a timeline record after every Executor step and update task metadata on completion/failure +- [x] 5.5 Write unit tests verifying timeline records are written in order and are readable after a simulated process restart (re-opening the same task directory) + +## 6. MCP tool server & REST API (capability: mcp-tool-server) + +- [x] 6.1 Implement `api/mcp.py`: register `take_screenshot`, `tap`, `swipe`, `input_text`, `launch_app`, `find_text`, `find_icon`, `get_ui_tree`, `describe_screen`, `list_devices`, `device_status` as MCP tools calling directly into `tools/` +- [x] 6.2 Add a semantic error-translation layer so driver/framework errors (WDA connection errors, element-not-found) surface as clear MCP tool errors (e.g. "device offline", "element not found") +- [x] 6.3 Implement `api/rest.py` (FastAPI) with `GET /devices`, `POST /devices/{id}/tap`, `POST /devices/{id}/screenshot`, `POST /devices/{id}/launch`, `POST /agent/task`, `GET /task/{id}`, calling the same `tools/`/`runtime/` functions as the MCP tools +- [x] 6.4 Write a test that starts a task via `POST /agent/task` and polls `GET /task/{id}` until it completes (using a mocked device) +- [x] 6.5 Write a test that calls each MCP tool against a mocked device and asserts no Appium/WDA-specific type or string leaks into the response + +## 7. End-to-end validation + +- [ ] 7.1 Run the full Observe→Think→Act→Observe loop against a real (or simulator) iPhone for a simple goal (e.g. "open an app and search for a term"), verifying tap/input/screenshot all work through the real `WDADriver` +- [ ] 7.2 Confirm the resulting task's timeline directory contains a complete, ordered set of per-step screenshots and JSON records +- [ ] 7.3 Confirm an MCP client (or a manual MCP tool call harness) can drive the same real-device flow end-to-end through `api/mcp.py` diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4cb675e --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,31 @@ +[project] +name = "agentic-modile-control" +version = "0.1.0" +requires-python = ">=3.14" +dependencies = [ + "Appium-Python-Client>=5.1.1", + "fastapi>=0.115.0", + "mcp>=1.27,<2", + "paddleocr>=3.0.0", + "uvicorn[standard]>=0.30.0", +] + +[build-system] +requires = ["setuptools>=69"] +build-backend = "setuptools.build_meta" + +[dependency-groups] +dev = [ + "httpx>=0.27.0", + "pytest>=8.3.0", +] + +[tool.setuptools.packages.find] +include = ["api*", "core*", "runtime*", "storage*", "tools*", "vision*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-ra" +markers = [ + "integration: tests requiring external hardware or services", +] diff --git a/runtime/__init__.py b/runtime/__init__.py new file mode 100644 index 0000000..1a8b987 --- /dev/null +++ b/runtime/__init__.py @@ -0,0 +1,2 @@ +"""Agent runtime loop, planner, executor, and in-run context.""" + diff --git a/runtime/context.py b/runtime/context.py new file mode 100644 index 0000000..6be4916 --- /dev/null +++ b/runtime/context.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from core.models import Scene + +if TYPE_CHECKING: + from runtime.executor import StepResult + + +@dataclass +class TaskContext: + task_id: str + goal: str + scenes: list[Scene] = field(default_factory=list) + step_results: list["StepResult"] = field(default_factory=list) + + def add_scene(self, scene: Scene) -> None: + self.scenes.append(scene) + + def add_step_result(self, result: "StepResult") -> None: + self.step_results.append(result) + + @property + def latest_scene(self) -> Scene | None: + if not self.scenes: + return None + return self.scenes[-1] + diff --git a/runtime/executor.py b/runtime/executor.py new file mode 100644 index 0000000..c6280ca --- /dev/null +++ b/runtime/executor.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from time import sleep +from typing import Any + +from core.device_manager import DeviceManager +from core.errors import ElementNotFoundError +from runtime.context import TaskContext +from runtime.planner import PlannedStep + +ToolCallable = Callable[..., Any] + + +@dataclass(frozen=True) +class StepResult: + step: PlannedStep + success: bool + attempts: int + result: Any = None + error: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "step": { + "action": self.step.action, + "description": self.step.description, + "args": dict(self.step.args), + "expected_text": self.step.expected_text, + }, + "success": self.success, + "attempts": self.attempts, + "result": self.result, + "error": self.error, + } + + +@dataclass +class ExecutorConfig: + max_retries: int = 3 + backoff_seconds: float = 0.25 + + +class Executor: + def __init__( + self, + *, + tools: dict[str, ToolCallable] | None = None, + config: ExecutorConfig | None = None, + ) -> None: + self.tools = tools or default_tool_registry() + self.config = config or ExecutorConfig() + + def execute( + self, + step: PlannedStep, + *, + context: TaskContext | None = None, + ) -> StepResult: + if self.config.max_retries < 1: + raise ValueError("max_retries must be at least 1") + + last_error: Exception | None = None + for attempt in range(1, self.config.max_retries + 1): + try: + result = self._execute_once(step, context=context) + return StepResult( + step=step, + success=True, + attempts=attempt, + result=result, + ) + except Exception as exc: + last_error = exc + if attempt < self.config.max_retries: + sleep(self.config.backoff_seconds) + + return StepResult( + step=step, + success=False, + attempts=self.config.max_retries, + error=str(last_error) if last_error else "step failed", + ) + + def _execute_once( + self, + step: PlannedStep, + *, + context: TaskContext | None, + ) -> Any: + if step.action == "wait_for_text": + query = step.args["query"] + scene = context.latest_scene if context else None + if scene is None: + raise ElementNotFoundError("element not found") + result = self.tools["find_text"](scene=scene, query=query) + if not result.get("found"): + raise ElementNotFoundError("element not found") + return result + + tool = self.tools.get(step.action) + if tool is None: + raise KeyError(f"unknown tool action {step.action}") + return tool(**step.args) + + +def default_tool_registry( + *, + manager: DeviceManager | None = None, +) -> dict[str, ToolCallable]: + from tools.describe_screen import describe_screen + from tools.find_icon import find_icon, find_icon_on_screen + from tools.find_text import find_text, find_text_on_screen + from tools.input_text import input_text + from tools.launch_app import launch_app, terminate_app + from tools.screenshot import take_screenshot + from tools.swipe import swipe + from tools.tap import tap + from tools.ui_tree import get_ui_tree + + return { + "take_screenshot": _bind_manager(take_screenshot, manager), + "screenshot": _bind_manager(take_screenshot, manager), + "tap": _bind_manager(tap, manager), + "swipe": _bind_manager(swipe, manager), + "input_text": _bind_manager(input_text, manager), + "launch_app": _bind_manager(launch_app, manager), + "terminate_app": _bind_manager(terminate_app, manager), + "get_ui_tree": _bind_manager(get_ui_tree, manager), + "ui_tree": _bind_manager(get_ui_tree, manager), + "describe_screen": _bind_manager(describe_screen, manager), + "find_text": find_text, + "find_text_on_screen": _bind_manager(find_text_on_screen, manager), + "find_icon": find_icon, + "find_icon_on_screen": _bind_manager(find_icon_on_screen, manager), + } + + +def _bind_manager(func: ToolCallable, manager: DeviceManager | None) -> ToolCallable: + if manager is None: + return func + + def wrapper(*args: Any, **kwargs: Any) -> Any: + kwargs.setdefault("manager", manager) + return func(*args, **kwargs) + + return wrapper diff --git a/runtime/planner.py b/runtime/planner.py new file mode 100644 index 0000000..f147c0f --- /dev/null +++ b/runtime/planner.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from core.models import Scene +from runtime.context import TaskContext + + +@dataclass(frozen=True) +class PlannedStep: + action: str + description: str + args: dict[str, Any] = field(default_factory=dict) + expected_text: str | None = None + + +class Planner: + def plan( + self, + *, + goal: str, + scene: Scene, + context: TaskContext, + ) -> list[PlannedStep]: + if context.step_results: + return [] + return [ + PlannedStep( + action="describe_screen", + description=f"Observe current screen for goal: {goal}", + args={}, + ) + ] + + def goal_reached(self, *, goal: str, scene: Scene, context: TaskContext) -> bool: + return bool(context.step_results) and all( + result.success for result in context.step_results + ) + diff --git a/runtime/task.py b/runtime/task.py new file mode 100644 index 0000000..3d047d1 --- /dev/null +++ b/runtime/task.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace + +from core.models import Scene, Task, utc_now +from runtime.context import TaskContext +from runtime.executor import Executor +from runtime.planner import PlannedStep, Planner +from storage.task_metadata import TaskMetadataStore +from storage.timeline import Timeline +from tools.describe_screen import describe_screen +from tools.screenshot import take_screenshot + + +@dataclass +class TaskRunnerConfig: + max_steps: int = 20 + + +Observer = Callable[[str], Scene] +ScreenshotProvider = Callable[[str], bytes] + + +class TaskRunner: + def __init__( + self, + *, + planner: Planner | None = None, + executor: Executor | None = None, + timeline: Timeline | None = None, + metadata_store: TaskMetadataStore | None = None, + config: TaskRunnerConfig | None = None, + observer: Observer | None = None, + screenshot_provider: ScreenshotProvider | None = None, + ) -> None: + self.planner = planner or Planner() + self.executor = executor or Executor() + self.timeline = timeline + self.metadata_store = metadata_store + self.config = config or TaskRunnerConfig() + self.observer = observer or (lambda device_id: describe_screen(device_id)) + self.screenshot_provider = screenshot_provider or ( + lambda device_id: take_screenshot(device_id) + ) + + def run(self, task: Task) -> Task: + context = TaskContext(task_id=task.id, goal=task.goal) + self._update_task(task, status="running") + + for _ in range(self.config.max_steps): + scene = self.observer(task.device_id) + context.add_scene(scene) + steps = self.planner.plan(goal=task.goal, scene=scene, context=context) + if not steps or self.planner.goal_reached( + goal=task.goal, + scene=scene, + context=context, + ): + self._update_task(task, status="completed", completed=True) + return task + + for step in steps: + result = self.executor.execute( + self._step_for_device(step, task.device_id), + context=context, + ) + context.add_step_result(result) + self._append_timeline(task, scene, step, result) + if not result.success: + self._update_task( + task, + status="failed", + completed=True, + failure_reason=result.error or "step failed", + ) + return task + + self._update_task( + task, + status="failed", + completed=True, + failure_reason=f"max steps exceeded: {self.config.max_steps}", + ) + return task + + def _append_timeline( + self, + task: Task, + scene: Scene, + step: PlannedStep, + result: object, + ) -> None: + if not self.timeline: + return + try: + screenshot = self.screenshot_provider(task.device_id) + except Exception: + screenshot = None + self.timeline.append( + task_id=task.id, + scene=scene.to_dict(), + prompt=task.goal, + tool_call={ + "action": step.action, + "description": step.description, + "args": step.args, + }, + result=result.to_dict() if hasattr(result, "to_dict") else {"result": result}, + screenshot=screenshot, + ) + + def _step_for_device(self, step: PlannedStep, device_id: str) -> PlannedStep: + device_scoped_actions = { + "take_screenshot", + "screenshot", + "tap", + "swipe", + "input_text", + "launch_app", + "terminate_app", + "get_ui_tree", + "ui_tree", + "describe_screen", + "find_text_on_screen", + "find_icon_on_screen", + } + if step.action not in device_scoped_actions or "device_id" in step.args: + return step + return replace(step, args={**step.args, "device_id": device_id}) + + def _update_task( + self, + task: Task, + *, + status: str, + completed: bool = False, + failure_reason: str | None = None, + ) -> None: + task.status = status # type: ignore[assignment] + task.updated_at = utc_now() + if completed: + task.completed_at = utc_now() + task.failure_reason = failure_reason + if self.metadata_store: + self.metadata_store.update_task( + task.id, + status=task.status, + completed=completed, + failure_reason=failure_reason, + ) diff --git a/storage/__init__.py b/storage/__init__.py new file mode 100644 index 0000000..091aa5f --- /dev/null +++ b/storage/__init__.py @@ -0,0 +1,2 @@ +"""Local artifact and task metadata storage.""" + diff --git a/storage/artifact_store.py b/storage/artifact_store.py new file mode 100644 index 0000000..eb677e3 --- /dev/null +++ b/storage/artifact_store.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import json +from dataclasses import asdict, is_dataclass +from datetime import date, datetime +from pathlib import Path +from typing import Any + + +class ArtifactStore: + def __init__(self, root: str | Path = "tasks/history") -> None: + self.root = Path(root) + + def task_dir(self, task_id: str) -> Path: + return self.root / task_id + + def write_step( + self, + *, + task_id: str, + index: int, + screenshot: bytes | None, + record: dict[str, Any], + ) -> dict[str, str | None]: + task_dir = self.task_dir(task_id) + task_dir.mkdir(parents=True, exist_ok=True) + stem = f"{index:03d}" + screenshot_path: Path | None = None + if screenshot is not None: + screenshot_path = task_dir / f"{stem}.png" + screenshot_path.write_bytes(screenshot) + + json_path = task_dir / f"{stem}.json" + payload = { + **record, + "screenshot_path": str(screenshot_path) if screenshot_path else None, + } + json_path.write_text( + json.dumps(_jsonable(payload), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + return { + "json_path": str(json_path), + "screenshot_path": str(screenshot_path) if screenshot_path else None, + } + + def read_steps(self, task_id: str) -> list[dict[str, Any]]: + task_dir = self.task_dir(task_id) + if not task_dir.exists(): + return [] + steps = [] + for path in sorted(task_dir.glob("*.json")): + steps.append(json.loads(path.read_text(encoding="utf-8"))) + return steps + + +def _jsonable(value: Any) -> Any: + if hasattr(value, "to_dict"): + return value.to_dict() + if is_dataclass(value): + return asdict(value) + if isinstance(value, dict): + return {key: _jsonable(inner) for key, inner in value.items()} + if isinstance(value, list): + return [_jsonable(inner) for inner in value] + if isinstance(value, tuple): + return [_jsonable(inner) for inner in value] + if isinstance(value, (datetime, date)): + return value.isoformat() + return value + diff --git a/storage/task_metadata.py b/storage/task_metadata.py new file mode 100644 index 0000000..7d86ef1 --- /dev/null +++ b/storage/task_metadata.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import sqlite3 +from pathlib import Path +from typing import Any + +from core.models import Task, TaskStatus, utc_now + + +class TaskMetadataStore: + def __init__(self, db_path: str | Path = "tasks/tasks.sqlite3") -> None: + self.db_path = Path(db_path) + self.db_path.parent.mkdir(parents=True, exist_ok=True) + self._ensure_schema() + + def create_task(self, task: Task) -> None: + with self._connect() as connection: + connection.execute( + """ + insert into tasks ( + id, goal, device_id, status, created_at, updated_at, + completed_at, failure_reason + ) values (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + task.id, + task.goal, + task.device_id, + task.status, + task.created_at.isoformat(), + task.updated_at.isoformat(), + task.completed_at.isoformat() if task.completed_at else None, + task.failure_reason, + ), + ) + + def update_task( + self, + task_id: str, + *, + status: TaskStatus | None = None, + failure_reason: str | None = None, + completed: bool = False, + ) -> None: + updates: dict[str, Any] = {"updated_at": utc_now().isoformat()} + if status: + updates["status"] = status + if failure_reason is not None: + updates["failure_reason"] = failure_reason + if completed: + updates["completed_at"] = utc_now().isoformat() + + assignments = ", ".join(f"{key} = ?" for key in updates) + values = [*updates.values(), task_id] + with self._connect() as connection: + connection.execute( + f"update tasks set {assignments} where id = ?", + values, + ) + + def get_task(self, task_id: str) -> dict[str, Any] | None: + with self._connect() as connection: + row = connection.execute( + "select * from tasks where id = ?", + (task_id,), + ).fetchone() + return dict(row) if row else None + + def list_tasks(self) -> list[dict[str, Any]]: + with self._connect() as connection: + rows = connection.execute( + "select * from tasks order by created_at desc" + ).fetchall() + return [dict(row) for row in rows] + + def _ensure_schema(self) -> None: + with self._connect() as connection: + connection.execute( + """ + create table if not exists tasks ( + id text primary key, + goal text not null, + device_id text not null, + status text not null, + created_at text not null, + updated_at text not null, + completed_at text, + failure_reason text + ) + """ + ) + + def _connect(self) -> sqlite3.Connection: + connection = sqlite3.connect(self.db_path) + connection.row_factory = sqlite3.Row + return connection + diff --git a/storage/timeline.py b/storage/timeline.py new file mode 100644 index 0000000..30b3c9d --- /dev/null +++ b/storage/timeline.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime +from typing import Any + +from storage.artifact_store import ArtifactStore + + +@dataclass(frozen=True) +class TimelineRecord: + index: int + scene: dict[str, Any] + prompt: str + tool_call: dict[str, Any] + result: dict[str, Any] + timestamp: str + screenshot_path: str | None = None + + +class Timeline: + def __init__(self, artifact_store: ArtifactStore | None = None) -> None: + self.artifact_store = artifact_store or ArtifactStore() + + def append( + self, + *, + task_id: str, + scene: Any, + prompt: str, + tool_call: dict[str, Any], + result: dict[str, Any], + screenshot: bytes | None = None, + ) -> TimelineRecord: + index = len(self.read(task_id)) + 1 + record = { + "index": index, + "scene": scene, + "prompt": prompt, + "tool_call": tool_call, + "result": result, + "timestamp": datetime.now().astimezone().isoformat(), + } + paths = self.artifact_store.write_step( + task_id=task_id, + index=index, + screenshot=screenshot, + record=record, + ) + return TimelineRecord( + index=index, + scene=record["scene"], + prompt=prompt, + tool_call=tool_call, + result=result, + timestamp=record["timestamp"], + screenshot_path=paths["screenshot_path"], + ) + + def read(self, task_id: str) -> list[dict[str, Any]]: + return self.artifact_store.read_steps(task_id) + diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..39239e5 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,2 @@ +"""Test package for Apex Agent.""" + diff --git a/tests/fakes.py b/tests/fakes.py new file mode 100644 index 0000000..5b83c2e --- /dev/null +++ b/tests/fakes.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from typing import Any + +from core.driver import Driver + +PNG_10X20 = ( + b"\x89PNG\r\n\x1a\n" + b"\x00\x00\x00\r" + b"IHDR" + b"\x00\x00\x00\x0a" + b"\x00\x00\x00\x14" +) + +TREE_XML = """ + + + + +""" + + +class FakeDriver(Driver): + def __init__( + self, + *, + fail_connect: bool = False, + tree: Any = TREE_XML, + screenshot: bytes = PNG_10X20, + ) -> None: + self.fail_connect = fail_connect + self.connected = False + self.calls: list[tuple[str, tuple[Any, ...]]] = [] + self._tree = tree + self._screenshot = screenshot + + def connect(self) -> None: + self.calls.append(("connect", ())) + if self.fail_connect: + raise RuntimeError("connection refused") + self.connected = True + + def disconnect(self) -> None: + self.calls.append(("disconnect", ())) + self.connected = False + + def screenshot(self) -> bytes: + self.calls.append(("screenshot", ())) + return self._screenshot + + def tap(self, x: float, y: float) -> None: + self.calls.append(("tap", (x, y))) + + def swipe( + self, + start_x: float, + start_y: float, + end_x: float, + end_y: float, + duration_ms: int = 500, + ) -> None: + self.calls.append(("swipe", (start_x, start_y, end_x, end_y, duration_ms))) + + def input(self, text: str) -> None: + self.calls.append(("input", (text,))) + + def launch(self, app_id: str) -> None: + self.calls.append(("launch", (app_id,))) + + def terminate(self, app_id: str) -> None: + self.calls.append(("terminate", (app_id,))) + + def tree(self) -> Any: + self.calls.append(("tree", ())) + return self._tree + + def home(self) -> None: + self.calls.append(("home", ())) + + def lock(self) -> None: + self.calls.append(("lock", ())) + + def unlock(self) -> None: + self.calls.append(("unlock", ())) + diff --git a/tests/test_device_manager.py b/tests/test_device_manager.py new file mode 100644 index 0000000..647d715 --- /dev/null +++ b/tests/test_device_manager.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +import pytest + +from core.device_manager import DeviceManager +from core.errors import DeviceOfflineError +from tests.fakes import FakeDriver + + +def test_device_manager_connect_disconnect_transitions() -> None: + driver = FakeDriver() + manager = DeviceManager() + manager.register_device("iphone-1", lambda: driver, connection_info={"wda_port": 8100}) + + assert manager.list_devices()[0].status == "idle" + assert manager.connect("iphone-1", max_retries=1) is driver + assert manager.status("iphone-1") == "busy" + + manager.disconnect("iphone-1") + assert manager.status("iphone-1") == "idle" + assert driver.calls[0] == ("connect", ()) + assert driver.calls[-1] == ("disconnect", ()) + + +def test_device_manager_marks_unreachable_device_offline() -> None: + manager = DeviceManager() + manager.register_device("iphone-1", lambda: FakeDriver(fail_connect=True)) + + with pytest.raises(DeviceOfflineError): + manager.connect("iphone-1", max_retries=2, retry_backoff_seconds=0) + + assert manager.status("iphone-1") == "offline" + diff --git a/tests/test_executor.py b/tests/test_executor.py new file mode 100644 index 0000000..a7d325a --- /dev/null +++ b/tests/test_executor.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from runtime.executor import Executor, ExecutorConfig +from runtime.planner import PlannedStep + + +def test_executor_retries_until_transient_tool_succeeds() -> None: + calls = {"count": 0} + + def flaky_tool() -> dict[str, bool]: + calls["count"] += 1 + if calls["count"] < 3: + raise RuntimeError("transient") + return {"ok": True} + + executor = Executor( + tools={"flaky": flaky_tool}, + config=ExecutorConfig(max_retries=3, backoff_seconds=0), + ) + + result = executor.execute(PlannedStep(action="flaky", description="retry")) + + assert result.success is True + assert result.attempts == 3 + assert result.result == {"ok": True} + diff --git a/tests/test_find_helpers.py b/tests/test_find_helpers.py new file mode 100644 index 0000000..3c86dc8 --- /dev/null +++ b/tests/test_find_helpers.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from core.models import Bounds, Scene, SceneElement +from tools.find_icon import find_icon +from tools.find_text import find_text + + +def test_find_text_found_and_not_found() -> None: + scene = Scene( + width=100, + height=100, + elements=[ + SceneElement( + id="button", + type="button", + text="Search", + bounds=Bounds(10, 20, 40, 20), + ) + ], + ) + + assert find_text(scene, "sea")["found"] is True + assert find_text(scene, "missing") == { + "found": False, + "query": "missing", + "reason": "not found", + } + + +def test_find_icon_found_and_not_found() -> None: + scene = Scene( + width=100, + height=100, + elements=[ + SceneElement( + id="settings-icon", + type="image", + text="Settings", + bounds=Bounds(50, 50, 20, 20), + ) + ], + ) + + assert find_icon(scene, "settings")["found"] is True + assert find_icon(scene, "profile")["reason"] == "not found" + diff --git a/tests/test_mcp.py b/tests/test_mcp.py new file mode 100644 index 0000000..67cb7e3 --- /dev/null +++ b/tests/test_mcp.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from api.mcp import tool_handlers +from core.device_manager import DeviceManager +from tests.fakes import FakeDriver + + +def test_mcp_handlers_are_semantic_and_hide_driver_terms() -> None: + manager = DeviceManager() + manager.register_device("iphone-1", lambda: FakeDriver()) + manager.connect("iphone-1", max_retries=1) + handlers = tool_handlers(manager=manager) + + responses = [ + handlers["take_screenshot"](device_id="iphone-1"), + handlers["tap"](x=1, y=2, device_id="iphone-1"), + handlers["swipe"]( + start_x=1, + start_y=2, + end_x=3, + end_y=4, + device_id="iphone-1", + ), + handlers["input_text"](text="hello", device_id="iphone-1"), + handlers["launch_app"](app_id="com.example.app", device_id="iphone-1"), + handlers["find_text"](query="Search", device_id="iphone-1"), + handlers["find_icon"](name="Settings", device_id="iphone-1"), + handlers["get_ui_tree"](device_id="iphone-1"), + handlers["describe_screen"](device_id="iphone-1"), + handlers["list_devices"](), + handlers["device_status"](device_id="iphone-1"), + ] + + serialized = repr(responses) + assert "WDA" not in serialized + assert "Appium" not in serialized + assert "XCUI" not in serialized + assert all(response is not None for response in responses) diff --git a/tests/test_rest_api.py b/tests/test_rest_api.py new file mode 100644 index 0000000..c7d6a31 --- /dev/null +++ b/tests/test_rest_api.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import pytest + +from core.device_manager import DeviceManager +from core.models import Task +from storage.task_metadata import TaskMetadataStore +from tests.fakes import FakeDriver + + +def test_rest_start_task_and_poll_until_complete(tmp_path) -> None: + pytest.importorskip("fastapi") + from fastapi.testclient import TestClient + + from api.rest import create_app + + driver = FakeDriver() + manager = DeviceManager() + manager.register_device("iphone-1", lambda: driver) + manager.connect("iphone-1", max_retries=1) + store = TaskMetadataStore(tmp_path / "tasks.sqlite3") + + class CompletingRunner: + def run(self, task: Task) -> None: + store.update_task(task.id, status="completed", completed=True) + + app = create_app( + manager=manager, + metadata_store=store, + task_runner=CompletingRunner(), + ) + client = TestClient(app) + + response = client.post( + "/agent/task", + json={"goal": "search", "device_id": "iphone-1"}, + ) + assert response.status_code == 200 + task_id = response.json()["task_id"] + + status_response = client.get(f"/task/{task_id}") + assert status_response.status_code == 200 + assert status_response.json()["status"] == "completed" + diff --git a/tests/test_scene_builder.py b/tests/test_scene_builder.py new file mode 100644 index 0000000..bf42d1d --- /dev/null +++ b/tests/test_scene_builder.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from core.models import Bounds, SceneElement +from vision.scene_builder import bbox_iou, build_scene + + +def test_scene_builder_merges_overlapping_ocr_into_ui_element() -> None: + ui_button = SceneElement( + id="ui-button", + type="button", + text=None, + bounds=Bounds(10, 10, 100, 40), + confidence=1.0, + source="ui", + ) + ocr_label = SceneElement( + id="ocr-label", + type="text", + text="Search", + bounds=Bounds(12, 12, 96, 36), + confidence=0.92, + source="ocr", + ) + ocr_only = SceneElement( + id="ocr-only", + type="text", + text="Footer", + bounds=Bounds(0, 100, 40, 20), + confidence=0.8, + source="ocr", + ) + + scene = build_scene( + screen_width=120, + screen_height=140, + ui_elements=[ui_button], + ocr_elements=[ocr_label, ocr_only], + ) + + assert len(scene.elements) == 2 + assert scene.elements[0].type == "button" + assert scene.elements[0].text == "Search" + assert scene.elements[1].text == "Footer" + assert bbox_iou(ui_button.bounds, ocr_label.bounds) > 0.5 + diff --git a/tests/test_smoke.py b/tests/test_smoke.py new file mode 100644 index 0000000..a20eba4 --- /dev/null +++ b/tests/test_smoke.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +import importlib + + +def test_imports_new_packages() -> None: + for package in ("core", "tools", "vision", "runtime", "api", "storage"): + importlib.import_module(package) + diff --git a/tests/test_task_loop.py b/tests/test_task_loop.py new file mode 100644 index 0000000..756e14a --- /dev/null +++ b/tests/test_task_loop.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from core.models import Bounds, Scene, SceneElement, Task +from runtime.executor import Executor, ExecutorConfig +from runtime.planner import PlannedStep, Planner +from runtime.task import TaskRunner, TaskRunnerConfig +from storage.artifact_store import ArtifactStore +from storage.task_metadata import TaskMetadataStore +from storage.timeline import Timeline +from tests.fakes import PNG_10X20 + + +class ScriptedPlanner(Planner): + def __init__(self, steps: list[PlannedStep]) -> None: + self.steps = steps + + def plan(self, *, goal, scene, context): + if len(context.step_results) >= len(self.steps): + return [] + return [self.steps[len(context.step_results)]] + + def goal_reached(self, *, goal, scene, context): + return len(context.step_results) >= len(self.steps) and all( + result.success for result in context.step_results + ) + + +def test_task_runner_executes_loop_and_writes_timeline(tmp_path) -> None: + scene = Scene( + width=10, + height=20, + elements=[ + SceneElement( + id="search", + type="input", + text="Search", + bounds=Bounds(1, 2, 4, 4), + ) + ], + ) + planner = ScriptedPlanner( + [ + PlannedStep(action="tap", description="tap search", args={"x": 3, "y": 4}), + PlannedStep( + action="input_text", + description="type query", + args={"text": "Mac mini"}, + ), + ] + ) + executor = Executor( + tools={ + "tap": lambda **kwargs: {"ok": True, **kwargs}, + "input_text": lambda **kwargs: {"ok": True, **kwargs}, + }, + config=ExecutorConfig(max_retries=1, backoff_seconds=0), + ) + metadata = TaskMetadataStore(tmp_path / "tasks.sqlite3") + timeline = Timeline(ArtifactStore(tmp_path / "history")) + task = Task(goal="open app and search", device_id="iphone-1") + metadata.create_task(task) + + runner = TaskRunner( + planner=planner, + executor=executor, + metadata_store=metadata, + timeline=timeline, + config=TaskRunnerConfig(max_steps=5), + observer=lambda device_id: scene, + screenshot_provider=lambda device_id: PNG_10X20, + ) + + result = runner.run(task) + + assert result.status == "completed" + assert len(timeline.read(task.id)) == 2 + assert metadata.get_task(task.id)["status"] == "completed" + diff --git a/tests/test_timeline.py b/tests/test_timeline.py new file mode 100644 index 0000000..3455e16 --- /dev/null +++ b/tests/test_timeline.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from storage.artifact_store import ArtifactStore +from storage.timeline import Timeline +from tests.fakes import PNG_10X20 + + +def test_timeline_records_survive_reopening_store(tmp_path) -> None: + store = ArtifactStore(tmp_path / "history") + timeline = Timeline(store) + timeline.append( + task_id="task-1", + scene={"screen": {"width": 1, "height": 1}, "elements": []}, + prompt="goal", + tool_call={"action": "tap"}, + result={"ok": True}, + screenshot=PNG_10X20, + ) + timeline.append( + task_id="task-1", + scene={"screen": {"width": 1, "height": 1}, "elements": []}, + prompt="goal", + tool_call={"action": "input_text"}, + result={"ok": True}, + screenshot=PNG_10X20, + ) + + reopened = Timeline(ArtifactStore(tmp_path / "history")) + records = reopened.read("task-1") + + assert [record["index"] for record in records] == [1, 2] + assert records[0]["screenshot_path"].endswith("001.png") + diff --git a/tests/test_wda_integration.py b/tests/test_wda_integration.py new file mode 100644 index 0000000..50235be --- /dev/null +++ b/tests/test_wda_integration.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import os + +import pytest + +from core.wda_driver import WDADriver, WDADriverConfig + + +@pytest.mark.integration +def test_wda_driver_screenshot_against_real_device() -> None: + server_url = os.getenv("APEX_WDA_SERVER_URL") + if not server_url: + pytest.skip("set APEX_WDA_SERVER_URL to run WDA hardware integration") + + driver = WDADriver( + WDADriverConfig( + server_url=server_url, + udid=os.getenv("APEX_WDA_UDID") or None, + device_name=os.getenv("APEX_WDA_DEVICE_NAME") or None, + ) + ) + driver.connect() + try: + assert driver.screenshot() + finally: + driver.disconnect() + diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 0000000..36337fc --- /dev/null +++ b/tools/__init__.py @@ -0,0 +1,2 @@ +"""Semantic capability wrappers over the active device driver.""" + diff --git a/tools/_device.py b/tools/_device.py new file mode 100644 index 0000000..85043e1 --- /dev/null +++ b/tools/_device.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from core.device_manager import DEFAULT_MANAGER, DeviceManager +from core.driver import Driver + + +def get_driver( + device_id: str | None = None, + *, + manager: DeviceManager | None = None, +) -> Driver: + return (manager or DEFAULT_MANAGER).active_driver(device_id) + diff --git a/tools/describe_screen.py b/tools/describe_screen.py new file mode 100644 index 0000000..97c36aa --- /dev/null +++ b/tools/describe_screen.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from core.device_manager import DeviceManager +from core.models import Scene +from tools._device import get_driver +from vision.ocr import PaddleOCREngine, run_ocr +from vision.scene_builder import build_scene, infer_png_size +from vision.ui_parser import parse_ui_tree + + +def describe_screen( + device_id: str | None = None, + *, + manager: DeviceManager | None = None, + ocr_engine: PaddleOCREngine | None = None, +) -> Scene: + driver = get_driver(device_id, manager=manager) + screenshot = driver.screenshot() + raw_tree = driver.tree() + width, height = infer_png_size(screenshot) + return build_scene( + screen_width=width, + screen_height=height, + ui_elements=parse_ui_tree(raw_tree), + ocr_elements=run_ocr(screenshot, engine=ocr_engine), + ) + diff --git a/tools/find_icon.py b/tools/find_icon.py new file mode 100644 index 0000000..0ed23d7 --- /dev/null +++ b/tools/find_icon.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from core.device_manager import DeviceManager +from core.models import Scene +from tools.describe_screen import describe_screen +from vision.icon_detector import find_icon as find_icon_in_scene + + +def find_icon(scene: Scene, name: str) -> dict[str, object]: + return find_icon_in_scene(scene, name).to_dict() + + +def find_icon_on_screen( + name: str, + *, + device_id: str | None = None, + manager: DeviceManager | None = None, +) -> dict[str, object]: + return find_icon(describe_screen(device_id, manager=manager), name) diff --git a/tools/find_text.py b/tools/find_text.py new file mode 100644 index 0000000..f79f6e6 --- /dev/null +++ b/tools/find_text.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from core.device_manager import DeviceManager +from core.models import Scene, SceneElement +from tools.describe_screen import describe_screen + + +def find_text( + scene: Scene, + query: str, + *, + exact: bool = False, + case_sensitive: bool = False, +) -> dict[str, object]: + match = _find_element(scene, query, exact=exact, case_sensitive=case_sensitive) + if not match: + return {"found": False, "query": query, "reason": "not found"} + x, y = match.center + return { + "found": True, + "query": query, + "x": x, + "y": y, + "element": match.to_dict(), + } + + +def find_text_on_screen( + query: str, + *, + device_id: str | None = None, + manager: DeviceManager | None = None, + exact: bool = False, + case_sensitive: bool = False, +) -> dict[str, object]: + return find_text( + describe_screen(device_id, manager=manager), + query, + exact=exact, + case_sensitive=case_sensitive, + ) + + +def _find_element( + scene: Scene, + query: str, + *, + exact: bool, + case_sensitive: bool, +) -> SceneElement | None: + needle = query if case_sensitive else query.casefold() + for element in scene.elements: + if not element.text: + continue + haystack = element.text if case_sensitive else element.text.casefold() + matched = haystack == needle if exact else needle in haystack + if matched: + return element + return None diff --git a/tools/input_text.py b/tools/input_text.py new file mode 100644 index 0000000..178c701 --- /dev/null +++ b/tools/input_text.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +from core.device_manager import DeviceManager +from tools._device import get_driver + + +def input_text( + text: str, + *, + device_id: str | None = None, + manager: DeviceManager | None = None, +) -> dict[str, object]: + get_driver(device_id, manager=manager).input(text) + return {"ok": True, "action": "input_text", "text": text} + diff --git a/tools/launch_app.py b/tools/launch_app.py new file mode 100644 index 0000000..f2e1890 --- /dev/null +++ b/tools/launch_app.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from core.device_manager import DeviceManager +from tools._device import get_driver + + +def launch_app( + app_id: str, + *, + device_id: str | None = None, + manager: DeviceManager | None = None, +) -> dict[str, object]: + get_driver(device_id, manager=manager).launch(app_id) + return {"ok": True, "action": "launch_app", "app_id": app_id} + + +def terminate_app( + app_id: str, + *, + device_id: str | None = None, + manager: DeviceManager | None = None, +) -> dict[str, object]: + get_driver(device_id, manager=manager).terminate(app_id) + return {"ok": True, "action": "terminate_app", "app_id": app_id} + diff --git a/tools/screenshot.py b/tools/screenshot.py new file mode 100644 index 0000000..5a3cee0 --- /dev/null +++ b/tools/screenshot.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from core.device_manager import DeviceManager +from tools._device import get_driver + + +def take_screenshot( + device_id: str | None = None, + *, + manager: DeviceManager | None = None, +) -> bytes: + return get_driver(device_id, manager=manager).screenshot() + diff --git a/tools/swipe.py b/tools/swipe.py new file mode 100644 index 0000000..b7b22df --- /dev/null +++ b/tools/swipe.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from core.device_manager import DeviceManager +from tools._device import get_driver + + +def swipe( + start_x: float, + start_y: float, + end_x: float, + end_y: float, + *, + duration_ms: int = 500, + device_id: str | None = None, + manager: DeviceManager | None = None, +) -> dict[str, object]: + get_driver(device_id, manager=manager).swipe( + start_x, + start_y, + end_x, + end_y, + duration_ms, + ) + return { + "ok": True, + "action": "swipe", + "start": {"x": start_x, "y": start_y}, + "end": {"x": end_x, "y": end_y}, + "duration_ms": duration_ms, + } + diff --git a/tools/tap.py b/tools/tap.py new file mode 100644 index 0000000..abee2fc --- /dev/null +++ b/tools/tap.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +from core.device_manager import DeviceManager +from tools._device import get_driver + + +def tap( + x: float, + y: float, + *, + device_id: str | None = None, + manager: DeviceManager | None = None, +) -> dict[str, object]: + get_driver(device_id, manager=manager).tap(x, y) + return {"ok": True, "action": "tap", "x": x, "y": y} + diff --git a/tools/ui_tree.py b/tools/ui_tree.py new file mode 100644 index 0000000..7fb240c --- /dev/null +++ b/tools/ui_tree.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from typing import Any + +from core.device_manager import DeviceManager +from tools._device import get_driver +from vision.ui_parser import parse_ui_tree + + +def get_raw_ui_tree( + device_id: str | None = None, + *, + manager: DeviceManager | None = None, +) -> Any: + return get_driver(device_id, manager=manager).tree() + + +def get_ui_tree( + device_id: str | None = None, + *, + manager: DeviceManager | None = None, +) -> list[dict[str, object]]: + return [ + element.to_dict() + for element in parse_ui_tree(get_raw_ui_tree(device_id, manager=manager)) + ] diff --git a/vision/__init__.py b/vision/__init__.py new file mode 100644 index 0000000..fd27b2d --- /dev/null +++ b/vision/__init__.py @@ -0,0 +1,2 @@ +"""Screen perception helpers for OCR, UI tree parsing, and Scene fusion.""" + diff --git a/vision/icon_detector.py b/vision/icon_detector.py new file mode 100644 index 0000000..13559d8 --- /dev/null +++ b/vision/icon_detector.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from core.models import Scene + + +@dataclass(frozen=True) +class IconSearchResult: + found: bool + name: str + x: float | None = None + y: float | None = None + reason: str | None = None + + def to_dict(self) -> dict[str, object]: + return { + "found": self.found, + "name": self.name, + "x": self.x, + "y": self.y, + "reason": self.reason, + } + + +def find_icon(scene: Scene, name: str) -> IconSearchResult: + query = name.casefold() + for element in scene.elements: + if element.type not in {"image", "icon", "button"}: + continue + label = (element.text or element.id).casefold() + if query in label: + x, y = element.center + return IconSearchResult(found=True, name=name, x=x, y=y) + return IconSearchResult(found=False, name=name, reason="not found") + diff --git a/vision/ocr.py b/vision/ocr.py new file mode 100644 index 0000000..03fa214 --- /dev/null +++ b/vision/ocr.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +import os +import tempfile +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from core.models import Bounds, SceneElement + + +@dataclass(frozen=True) +class OCRBox: + text: str + bounds: Bounds + confidence: float | None = None + + def to_scene_element(self, element_id: str) -> SceneElement: + return SceneElement( + id=element_id, + type="text", + text=self.text, + bounds=self.bounds, + confidence=self.confidence, + source="ocr", + ) + + +class PaddleOCREngine: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + self._engine: Any | None = None + + def extract(self, image: bytes | str | Path) -> list[OCRBox]: + engine = self._load() + image_input, temp_path = _image_input(image) + try: + if hasattr(engine, "predict"): + raw = engine.predict(input=image_input) + else: + raw = engine.ocr(image_input, cls=True) + return parse_paddle_result(raw) + finally: + if temp_path: + temp_path.unlink(missing_ok=True) + + def _load(self) -> Any: + if self._engine is None: + from paddleocr import PaddleOCR + + self._engine = PaddleOCR(**self.kwargs) + return self._engine + + +def run_ocr( + image: bytes | str | Path, + *, + engine: PaddleOCREngine | None = None, + strict: bool = False, +) -> list[SceneElement]: + try: + boxes = (engine or PaddleOCREngine()).extract(image) + except ImportError: + if strict: + raise + return [] + return [box.to_scene_element(f"ocr-{index:03d}") for index, box in enumerate(boxes)] + + +def parse_paddle_result(raw: Any) -> list[OCRBox]: + boxes: list[OCRBox] = [] + for item in _flatten_pages(raw): + parsed = _parse_line(item) + if parsed: + boxes.append(parsed) + return boxes + + +def _image_input(image: bytes | str | Path) -> tuple[str, Path | None]: + if isinstance(image, bytes): + handle = tempfile.NamedTemporaryFile(delete=False, suffix=".png") + try: + handle.write(image) + finally: + handle.close() + return handle.name, Path(handle.name) + return os.fspath(image), None + + +def _flatten_pages(raw: Any) -> Iterable[Any]: + if raw is None: + return [] + if isinstance(raw, dict): + return _dict_lines(raw) + if isinstance(raw, list): + flattened: list[Any] = [] + for page in raw: + if isinstance(page, dict): + flattened.extend(_dict_lines(page)) + elif _looks_like_ocr_line(page): + flattened.append(page) + elif isinstance(page, list): + flattened.extend(page) + return flattened + json_attr = getattr(raw, "json", None) + if callable(json_attr): + return _flatten_pages(json_attr) + return [] + + +def _dict_lines(data: dict[str, Any]) -> list[Any]: + payload = data.get("res") if isinstance(data.get("res"), dict) else data + texts = payload.get("rec_texts") or payload.get("texts") or [] + scores = payload.get("rec_scores") or payload.get("scores") or [] + boxes = ( + payload.get("rec_boxes") + or payload.get("rec_polys") + or payload.get("dt_polys") + or [] + ) + return [ + { + "text": text, + "confidence": scores[index] if index < len(scores) else None, + "points": boxes[index] if index < len(boxes) else None, + } + for index, text in enumerate(texts) + ] + + +def _looks_like_ocr_line(value: Any) -> bool: + return isinstance(value, (list, tuple)) and len(value) >= 2 + + +def _parse_line(line: Any) -> OCRBox | None: + if isinstance(line, dict): + text = line.get("text") + points = line.get("points") or line.get("box") or line.get("bounds") + confidence = line.get("confidence") + if not text or not points: + return None + return OCRBox(str(text), _bounds_from_points(points), _float_or_none(confidence)) + + if not _looks_like_ocr_line(line): + return None + + points = line[0] + text_payload = line[1] + if isinstance(text_payload, (list, tuple)) and text_payload: + text = text_payload[0] + confidence = text_payload[1] if len(text_payload) > 1 else None + else: + text = text_payload + confidence = None + if not text: + return None + return OCRBox(str(text), _bounds_from_points(points), _float_or_none(confidence)) + + +def _bounds_from_points(points: Any) -> Bounds: + if isinstance(points, dict): + return Bounds.from_dict(points) + if ( + isinstance(points, (list, tuple)) + and len(points) == 4 + and all(isinstance(value, (int, float)) for value in points) + ): + x1, y1, x2, y2 = [float(value) for value in points] + return Bounds(x1, y1, x2 - x1, y2 - y1) + + xs: list[float] = [] + ys: list[float] = [] + for point in points: + if isinstance(point, dict): + xs.append(float(point["x"])) + ys.append(float(point["y"])) + else: + xs.append(float(point[0])) + ys.append(float(point[1])) + return Bounds(min(xs), min(ys), max(xs) - min(xs), max(ys) - min(ys)) + + +def _float_or_none(value: Any) -> float | None: + if value is None: + return None + return float(value) + diff --git a/vision/scene_builder.py b/vision/scene_builder.py new file mode 100644 index 0000000..dcb34f8 --- /dev/null +++ b/vision/scene_builder.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from struct import unpack + +from core.models import Bounds, Scene, SceneElement + + +def build_scene( + *, + screen_width: int, + screen_height: int, + ui_elements: list[SceneElement] | None = None, + ocr_elements: list[SceneElement] | None = None, + iou_threshold: float = 0.5, +) -> Scene: + ui_elements = ui_elements or [] + ocr_elements = ocr_elements or [] + merged: list[SceneElement] = [] + used_ocr: set[int] = set() + + for ui_index, ui_element in enumerate(ui_elements): + best_index: int | None = None + best_iou = 0.0 + for ocr_index, ocr_element in enumerate(ocr_elements): + if ocr_index in used_ocr: + continue + score = bbox_iou(ui_element.bounds, ocr_element.bounds) + if score > best_iou: + best_iou = score + best_index = ocr_index + + element = _with_id(ui_element, f"ui-{ui_index:03d}") + if best_index is not None and best_iou >= iou_threshold: + used_ocr.add(best_index) + ocr_element = ocr_elements[best_index] + element = replace( + element, + text=element.text or ocr_element.text, + confidence=_best_confidence(element.confidence, ocr_element.confidence), + ) + merged.append(element) + + for ocr_index, ocr_element in enumerate(ocr_elements): + if ocr_index in used_ocr: + continue + merged.append(_with_id(ocr_element, f"ocr-{ocr_index:03d}")) + + return Scene(width=screen_width, height=screen_height, elements=merged) + + +def bbox_iou(first: Bounds, second: Bounds) -> float: + x_left = max(first.x, second.x) + y_top = max(first.y, second.y) + x_right = min(first.right, second.right) + y_bottom = min(first.bottom, second.bottom) + if x_right <= x_left or y_bottom <= y_top: + return 0.0 + + intersection = (x_right - x_left) * (y_bottom - y_top) + first_area = first.width * first.height + second_area = second.width * second.height + union = first_area + second_area - intersection + if union <= 0: + return 0.0 + return intersection / union + + +def infer_png_size(image: bytes | str | Path | None) -> tuple[int, int]: + if image is None: + return (0, 0) + data = Path(image).read_bytes() if not isinstance(image, bytes) else image + if len(data) >= 24 and data[:8] == b"\x89PNG\r\n\x1a\n": + width, height = unpack(">II", data[16:24]) + return (int(width), int(height)) + return (0, 0) + + +def _with_id(element: SceneElement, fallback_id: str) -> SceneElement: + if element.id: + return element + return replace(element, id=fallback_id) + + +def _best_confidence(first: float | None, second: float | None) -> float | None: + values = [value for value in (first, second) if value is not None] + if not values: + return None + return max(values) + diff --git a/vision/ui_parser.py b/vision/ui_parser.py new file mode 100644 index 0000000..2c92788 --- /dev/null +++ b/vision/ui_parser.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import re +import xml.etree.ElementTree as ET +from typing import Any + +from core.models import Bounds, SceneElement + +_ANDROID_BOUNDS = re.compile(r"\[(?P-?\d+),(?P-?\d+)\]\[(?P-?\d+),(?P-?\d+)\]") + + +def parse_ui_tree(raw_tree: Any) -> list[SceneElement]: + if raw_tree is None: + return [] + if isinstance(raw_tree, str): + return _parse_xml(raw_tree) + if isinstance(raw_tree, dict): + elements: list[SceneElement] = [] + _parse_dict_node(raw_tree, elements) + return _with_stable_ids(elements, "ui") + if isinstance(raw_tree, list): + elements = [] + for node in raw_tree: + if isinstance(node, dict): + _parse_dict_node(node, elements) + return _with_stable_ids(elements, "ui") + return [] + + +def _parse_xml(raw_tree: str) -> list[SceneElement]: + if not raw_tree.strip(): + return [] + root = ET.fromstring(raw_tree) + elements = [] + for node in root.iter(): + bounds = _bounds_from_attrs(node.attrib) + if bounds is None or bounds.width <= 0 or bounds.height <= 0: + continue + elements.append( + SceneElement( + id="", + type=_normalize_type(_local_name(node.tag), node.attrib), + text=_text_from_attrs(node.attrib), + bounds=bounds, + confidence=1.0, + source="ui", + ) + ) + return _with_stable_ids(elements, "ui") + + +def _parse_dict_node(node: dict[str, Any], elements: list[SceneElement]) -> None: + attrs = dict(node) + children = attrs.pop("children", None) or attrs.pop("nodes", None) or [] + bounds = _bounds_from_attrs(attrs) + if bounds and bounds.width > 0 and bounds.height > 0: + elements.append( + SceneElement( + id="", + type=_normalize_type( + str(attrs.get("type") or attrs.get("class") or "unknown"), + attrs, + ), + text=_text_from_attrs(attrs), + bounds=bounds, + confidence=1.0, + source="ui", + ) + ) + for child in children: + if isinstance(child, dict): + _parse_dict_node(child, elements) + + +def _with_stable_ids(elements: list[SceneElement], prefix: str) -> list[SceneElement]: + return [ + SceneElement( + id=element.id or f"{prefix}-{index:03d}", + type=element.type, + text=element.text, + bounds=element.bounds, + confidence=element.confidence, + source=element.source, + ) + for index, element in enumerate(elements) + ] + + +def _bounds_from_attrs(attrs: dict[str, Any]) -> Bounds | None: + if all(key in attrs for key in ("x", "y", "width", "height")): + return Bounds( + float(attrs["x"]), + float(attrs["y"]), + float(attrs["width"]), + float(attrs["height"]), + ) + if all(key in attrs for key in ("left", "top", "right", "bottom")): + left = float(attrs["left"]) + top = float(attrs["top"]) + right = float(attrs["right"]) + bottom = float(attrs["bottom"]) + return Bounds(left, top, right - left, bottom - top) + raw_bounds = attrs.get("bounds") or attrs.get("rect") + if isinstance(raw_bounds, dict): + return _bounds_from_attrs(raw_bounds) + if isinstance(raw_bounds, str): + match = _ANDROID_BOUNDS.fullmatch(raw_bounds.strip()) + if match: + x1 = float(match.group("x1")) + y1 = float(match.group("y1")) + x2 = float(match.group("x2")) + y2 = float(match.group("y2")) + return Bounds(x1, y1, x2 - x1, y2 - y1) + return None + + +def _text_from_attrs(attrs: dict[str, Any]) -> str | None: + for key in ("label", "name", "text", "value", "placeholder"): + value = attrs.get(key) + if value not in (None, ""): + return str(value) + return None + + +def _normalize_type(raw_type: str, attrs: dict[str, Any]) -> str: + candidate = ( + attrs.get("role") + or attrs.get("type") + or attrs.get("class") + or attrs.get("className") + or raw_type + ) + normalized = str(candidate).split(".")[-1].replace("XCUIElementType", "").lower() + if "button" in normalized: + return "button" + if "textfield" in normalized or "textarea" in normalized or "input" in normalized: + return "input" + if "statictext" in normalized or normalized in {"text", "label"}: + return "text" + if "image" in normalized or "icon" in normalized: + return "image" + if "cell" in normalized or "row" in normalized: + return "cell" + if "window" in normalized: + return "window" + return normalized or "unknown" + + +def _local_name(tag: str) -> str: + return tag.rsplit("}", 1)[-1] +