Implement Apex Agent MVP scaffold
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user