From 8c99dc015a465bb39fd53ff37ad26463455e2269 Mon Sep 17 00:00:00 2001 From: showtan001 <240788545@qq.com> Date: Mon, 31 Aug 2026 10:43:37 +0800 Subject: [PATCH] feat: add on-demand device screenshots --- README.md | 5 ++ apps/device-host-agent/host_agent/web/app.py | 38 ++++++++++ .../host_agent/web/templates/devices.html | 60 +++++++++++++++- .../tests/host_agent/web/test_templates.py | 1 + apps/device-host-agent/tests/test_web_app.py | 72 +++++++++++++++++++ device/manager.py | 8 ++- driver/android_driver.py | 7 ++ driver/base.py | 9 +++ driver/wda_driver.py | 7 ++ tests/test_device_manager.py | 26 +++++++ 10 files changed, 231 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 62026e5..dbbdd83 100644 --- a/README.md +++ b/README.md @@ -125,6 +125,11 @@ and final replies in a local SQLite database. View them at `http://127.0.0.1:8765/conversations`; image bytes are excluded. Set `HOST_AGENT_CONVERSATION_LOG_PATH` to change the database path. +The authenticated `Devices` page has an on-demand `Get screenshot` button for +each connected device. Screenshots are captured only after the operator clicks +the button; the page does not auto-refresh or capture screenshots as part of +heartbeat synchronization. + In local mode, Appium supervision is enabled by default. Host Agent probes `/status`, adopts a healthy existing Appium instance, starts Appium when no listener exists, restarts only processes it started if they crash, and stops diff --git a/apps/device-host-agent/host_agent/web/app.py b/apps/device-host-agent/host_agent/web/app.py index 9c99baa..a787f8f 100644 --- a/apps/device-host-agent/host_agent/web/app.py +++ b/apps/device-host-agent/host_agent/web/app.py @@ -11,6 +11,7 @@ import jinja2 from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response +from core.errors import DeviceNotFoundError, DeviceOfflineError, DeviceRuntimeError from device.manager import DeviceManager from host_agent.assignment import AssignmentExecutor from host_agent.client import ( @@ -488,6 +489,43 @@ def create_console_app( error=None, ) + @app.post("/api/devices/{device_id}/screenshot") + async def api_device_screenshot( + device_id: str, + session: SessionState = Depends(require_csrf), + ) -> Response: + """Capture one on-demand screenshot for a connected local device.""" + try: + screenshot = await asyncio.to_thread( + lambda: manager.active_driver(device_id).screenshot() + ) + except DeviceNotFoundError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + except DeviceOfflineError as exc: + raise HTTPException(status_code=503, detail=str(exc)) from exc + except DeviceRuntimeError as exc: + raise HTTPException(status_code=502, detail=str(exc)) from exc + except Exception as exc: + raise HTTPException( + status_code=502, + detail=str(exc) or "failed to capture device screenshot", + ) from exc + + if not isinstance(screenshot, bytes) or not screenshot: + raise HTTPException( + status_code=502, + detail="device returned an empty screenshot", + ) + return Response( + content=screenshot, + media_type="image/png", + headers={ + "Cache-Control": "no-store", + "Pragma": "no-cache", + "X-Content-Type-Options": "nosniff", + }, + ) + @app.post("/devices/save") async def devices_save( request: Request, diff --git a/apps/device-host-agent/host_agent/web/templates/devices.html b/apps/device-host-agent/host_agent/web/templates/devices.html index 10b8929..6f68e74 100644 --- a/apps/device-host-agent/host_agent/web/templates/devices.html +++ b/apps/device-host-agent/host_agent/web/templates/devices.html @@ -5,13 +5,20 @@

{{ error }}

{% endif %} - + {% for device in devices %} +
IDNameDriverCloud ID
IDNameDriverCloud IDScreenshot
{{ device["device_id"] }} {{ device["name"] or "" }} {{ device["driver_type"] }} {{ device["cloud_device_id"] or "" }} + + + Edit
@@ -34,4 +41,55 @@
+ + {% endblock %} diff --git a/apps/device-host-agent/tests/host_agent/web/test_templates.py b/apps/device-host-agent/tests/host_agent/web/test_templates.py index a7f38d4..08a0cee 100644 --- a/apps/device-host-agent/tests/host_agent/web/test_templates.py +++ b/apps/device-host-agent/tests/host_agent/web/test_templates.py @@ -47,6 +47,7 @@ def test_devices_renders(env, sample_session) -> None: **make_devices_context(sample_session) ) assert '
' in html + assert 'class="screenshot-button"' in html def test_account_renders(env, sample_session) -> None: diff --git a/apps/device-host-agent/tests/test_web_app.py b/apps/device-host-agent/tests/test_web_app.py index a73fea6..434c51c 100644 --- a/apps/device-host-agent/tests/test_web_app.py +++ b/apps/device-host-agent/tests/test_web_app.py @@ -242,6 +242,78 @@ def test_add_device_appears_in_devices_page_and_manager(tmp_path) -> None: assert [device.id for device in context["manager"].list_devices()] == ["device-a"] +def test_devices_page_captures_screenshot_only_when_button_endpoint_is_called( + tmp_path, +) -> None: + class ScreenshotDriver: + def __init__(self) -> None: + self.capture_count = 0 + + def connect(self) -> None: + pass + + def disconnect(self) -> None: + pass + + def screenshot(self) -> bytes: + self.capture_count += 1 + return b"fake-png" + + driver = ScreenshotDriver() + client, context = _build_client(tmp_path) + context["config_store"].add( + device_id="device-a", + name="Lab iPhone", + driver_type="wda", + connection_info={}, + ) + context["manager"].register_device("device-a", lambda: driver) + context["manager"].connect("device-a") + csrf_token = _login(client) + + page = client.get("/devices") + assert page.status_code == 200 + assert 'class="screenshot-button"' in page.text + assert 'data-device-id="device-a"' in page.text + assert driver.capture_count == 0 + + response = client.post( + "/api/devices/device-a/screenshot", + headers={"X-CSRF-Token": csrf_token}, + ) + + assert response.status_code == 200 + assert response.content == b"fake-png" + assert response.headers["content-type"] == "image/png" + assert response.headers["cache-control"] == "no-store" + assert driver.capture_count == 1 + + +def test_device_screenshot_requires_csrf_and_connected_device(tmp_path) -> None: + class ScreenshotDriver: + def connect(self) -> None: + pass + + def disconnect(self) -> None: + pass + + def screenshot(self) -> bytes: + return b"fake-png" + + client, context = _build_client(tmp_path) + context["manager"].register_device("device-a", ScreenshotDriver) + csrf_token = _login(client) + + missing_csrf = client.post("/api/devices/device-a/screenshot") + assert missing_csrf.status_code == 403 + + offline = client.post( + "/api/devices/device-a/screenshot", + headers={"X-CSRF-Token": csrf_token}, + ) + assert offline.status_code == 503 + + def test_remove_device_unregisters_from_manager(tmp_path) -> None: client, context = _build_client(tmp_path) csrf_token = _login(client) diff --git a/device/manager.py b/device/manager.py index 709a886..6db28a8 100644 --- a/device/manager.py +++ b/device/manager.py @@ -130,7 +130,13 @@ class DeviceManager: if driver is None: return False try: - driver.screenshot() + health_check = getattr(driver, "health_check", None) + if callable(health_check): + health_check() + else: + # Compatibility for drivers implemented before health_check + # existed. Built-in drivers use the non-screen health check. + driver.screenshot() except Exception: self.mark_error(device_id, offline=True) return False diff --git a/driver/android_driver.py b/driver/android_driver.py index 160530d..bf639d5 100644 --- a/driver/android_driver.py +++ b/driver/android_driver.py @@ -81,6 +81,13 @@ class AndroidDriver(Driver): except Exception as exc: raise DriverError("screenshot failed") from exc + def health_check(self) -> None: + client = self._require_client() + try: + client.get_status() + except Exception as exc: + raise DriverError("health check failed") from exc + def tap(self, x: float, y: float) -> None: client = self._require_client() try: diff --git a/driver/base.py b/driver/base.py index f4cc0e7..c330e5c 100644 --- a/driver/base.py +++ b/driver/base.py @@ -25,6 +25,15 @@ class Driver(ABC): def screenshot(self) -> bytes: """Return the current screen as image bytes.""" + def health_check(self) -> None: + """Verify the live session without reading the device screen. + + Drivers with a transport-level status endpoint should override this + method. The default is a no-op for legacy drivers that do not expose + a separate health check. + """ + return None + @abstractmethod def tap(self, x: float, y: float) -> None: """Tap the screen at the given coordinates.""" diff --git a/driver/wda_driver.py b/driver/wda_driver.py index 169daca..9dd0c4e 100644 --- a/driver/wda_driver.py +++ b/driver/wda_driver.py @@ -71,6 +71,13 @@ class WDADriver(Driver): except Exception as exc: raise DriverError("screenshot failed") from exc + def health_check(self) -> None: + client = self._require_client() + try: + client.get_status() + except Exception as exc: + raise DriverError("health check failed") from exc + def tap(self, x: float, y: float) -> None: client = self._require_client() try: diff --git a/tests/test_device_manager.py b/tests/test_device_manager.py index 871708f..9667a3f 100644 --- a/tests/test_device_manager.py +++ b/tests/test_device_manager.py @@ -46,3 +46,29 @@ def test_probe_marks_connected_device_offline_when_driver_is_unreachable() -> No assert manager.probe("iphone-1") is False assert manager.status("iphone-1") == "offline" + + +def test_probe_uses_health_check_without_capturing_screen() -> None: + class HealthCheckedDriver: + def __init__(self) -> None: + self.health_checks = 0 + self.screenshots = 0 + + def connect(self) -> None: + return None + + def screenshot(self) -> bytes: + self.screenshots += 1 + return b"screen" + + def health_check(self) -> None: + self.health_checks += 1 + + driver = HealthCheckedDriver() + manager = DeviceManager() + manager.register_device("iphone-1", lambda: driver) # type: ignore[arg-type] + manager.connect("iphone-1") + + assert manager.probe("iphone-1") is True + assert driver.health_checks == 1 + assert driver.screenshots == 0