feat: add on-demand device screenshots

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,

apps/device-host-agent/host_agent/web/templates/devices.html

@@ -5,13 +5,20 @@
     <p class="error">{{ error }}</p>
     {% endif %}
     <table>
-      <thead><tr><th>ID</th><th>Name</th><th>Driver</th><th>Cloud ID</th><th></th></tr></thead>
+      <thead><tr><th>ID</th><th>Name</th><th>Driver</th><th>Cloud ID</th><th>Screenshot</th><th></th></tr></thead>
       <tbody>{% for device in devices %}
         <tr>
           <td>{{ device["device_id"] }}</td>
           <td>{{ device["name"] or "" }}</td>
           <td>{{ device["driver_type"] }}</td>
           <td>{{ device["cloud_device_id"] or "" }}</td>
+          <td>
+            <button type="button" class="screenshot-button" data-device-id="{{ device["device_id"] }}">Get screenshot</button>
+            <div class="screenshot-preview" data-screenshot-preview hidden>
+              <p class="screenshot-status" data-screenshot-status></p>
+              <img alt="Current screen for {{ device["device_id"] }}" data-screenshot-image hidden>
+            </div>
+          </td>
           <td>
             <a href="/devices?edit={{ device["device_id"] }}">Edit</a>
             <form class="inline" method="post" action="/devices/remove">
@@ -34,4 +41,55 @@
       </label><br>
       <button type="submit">Save</button>
     </form>
+    <style>
+      .screenshot-preview { margin-top: 0.5rem; max-width: 260px; }
+      .screenshot-preview img { display: block; width: 100%; height: auto; border: 1px solid #c8d0d6; }
+      .screenshot-status { margin: 0 0 0.35rem; color: #5e6b73; }
+      .screenshot-status.error { color: #b00020; }
+    </style>
+    <script>
+      (() => {
+        const csrfInput = document.querySelector('input[name="csrf_token"]');
+        const csrfToken = csrfInput ? csrfInput.value : "";
+        document.querySelectorAll(".screenshot-button").forEach((button) => {
+          button.addEventListener("click", async () => {
+            const deviceId = button.dataset.deviceId;
+            const preview = button.parentElement.querySelector("[data-screenshot-preview]");
+            const status = preview.querySelector("[data-screenshot-status]");
+            const image = preview.querySelector("[data-screenshot-image]");
+            const previousUrl = image.dataset.objectUrl;
+            if (previousUrl) URL.revokeObjectURL(previousUrl);
+            button.disabled = true;
+            preview.hidden = false;
+            image.hidden = true;
+            status.classList.remove("error");
+            status.textContent = "Capturing...";
+            try {
+              const response = await fetch(
+                "/api/devices/" + encodeURIComponent(deviceId) + "/screenshot",
+                { method: "POST", headers: { "X-CSRF-Token": csrfToken } }
+              );
+              if (!response.ok) {
+                let detail = "Screenshot failed.";
+                try {
+                  const payload = await response.json();
+                  if (payload.detail) detail = payload.detail;
+                } catch (_) {}
+                throw new Error(detail);
+              }
+              const objectUrl = URL.createObjectURL(await response.blob());
+              image.src = objectUrl;
+              image.dataset.objectUrl = objectUrl;
+              image.hidden = false;
+              status.textContent = "Captured.";
+            } catch (error) {
+              status.classList.add("error");
+              status.textContent = error.message || "Screenshot failed.";
+            } finally {
+              button.disabled = false;
+            }
+          });
+        });
+      })();
+    </script>
 {% endblock %}

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 '<form method="post" action="/devices/save">' in html
+    assert 'class="screenshot-button"' in html
 
 
 def test_account_renders(env, sample_session) -> None:

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)