feat: checkpoint device agent runtime milestones

device/__init__.py

@@ -0,0 +1,3 @@
+from device.manager import DEFAULT_MANAGER, DeviceManager, DriverFactory
+
+__all__ = ["DEFAULT_MANAGER", "DeviceManager", "DriverFactory"]

device/manager.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import replace
+from threading import RLock
+from time import sleep
+
+from core.errors import (
+    DeviceBusyError,
+    DeviceNotFoundError,
+    DeviceOfflineError,
+    DriverError,
+)
+from core.models import Device, DeviceStatus
+from driver.base import Driver
+
+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()