feat(local): scale and encode desktop screenshots at the source

pyproject.toml

@@ -38,7 +38,7 @@ cli = [
   "typer>=0.12,<1",
 ]
 desktop = [
-  "hai-drivers[desktop]>=0.1.0",
+  "hai-drivers[desktop]>=0.1.2",
 ]
 browser = [
   "hai-drivers[web]>=0.1.0",

src/hai_agents_cli/app.py

@@ -21,6 +21,7 @@
 from hai_agents_common import credentials
 from hai_agents_common.credentials import absolute_share_url, make_client
 from hai_agents_common.jsonable import to_jsonable
+from hai_agents_local import desktop as desktop_defaults
 
 from . import auth, mcp_hosts
 
@@ -586,11 +587,31 @@ def local_browser(
 def local_desktop(
     ctx: typer.Context,
     session_id: str | None = typer.Option(None, "--session-id", help="Session id to serve. Generated when omitted."),
+    max_width: int = typer.Option(
+        desktop_defaults.DEFAULT_MAX_WIDTH, "--max-width", help="Cap screenshot width in pixels. 0 disables."
+    ),
+    max_height: int = typer.Option(0, "--max-height", help="Cap screenshot height in pixels. 0 disables."),
+    image_format: str = typer.Option(
+        desktop_defaults.DEFAULT_IMAGE_FORMAT, "--image-format", help="Screenshot encoding: png, jpeg, or webp."
+    ),
+    quality: int = typer.Option(
+        desktop_defaults.DEFAULT_QUALITY, "--quality", help="Encoding quality (1-100) for jpeg/webp."
+    ),
 ) -> None:
     """Serve desktop commands on this machine's mouse, keyboard, and screen."""
     from hai_agents_local import PyautoguiDesktopBridge
 
-    _run_bridge(_state(ctx), PyautoguiDesktopBridge, session_id)
+    if image_format not in ("png", "jpeg", "webp"):
+        _raise_cli_error(ValueError(f"--image-format must be png, jpeg, or webp; got {image_format!r}"))
+    _run_bridge(
+        _state(ctx),
+        PyautoguiDesktopBridge,
+        session_id,
+        max_width=max_width or None,
+        max_height=max_height or None,
+        image_format=image_format,
+        quality=quality,
+    )
 
 
 def _run_bridge(state: AppState, bridge_type: type[LocalBridge], session_id: str | None, **options: Any) -> None:

src/hai_agents_local/desktop.py

@@ -1,27 +1,65 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
-from .bridge import LocalBridge
+from .bridge import LocalBridge, TokenSource
 
 if TYPE_CHECKING:
-    from hai_drivers.desktop.local import LocalDesktopDriver
+    from hai_drivers.desktop.interface import DesktopDriverInterface
 
+ImageFormat = Literal["png", "jpeg", "webp"]
 
-class PyautoguiDesktopBridge(LocalBridge["LocalDesktopDriver"]):
-    """Serves desktop environments (mouse, keyboard, screen, files, shell) on this machine via pyautogui."""
+DEFAULT_MAX_WIDTH = 1920
+DEFAULT_IMAGE_FORMAT: ImageFormat = "jpeg"
+DEFAULT_QUALITY = 85
+
+
+class PyautoguiDesktopBridge(LocalBridge["DesktopDriverInterface"]):
+    """Serves desktop environments (mouse, keyboard, screen, files, shell) on this machine via pyautogui.
+
+    Screenshots are downscaled and encoded here, before they cross the network; set every
+    knob to None to serve raw native-resolution captures.
+    """
 
     environment_kind = "desktop"
 
-    def create_driver(self) -> LocalDesktopDriver:
+    def __init__(
+        self,
+        environment_id: str | None = None,
+        *,
+        api_key: TokenSource,
+        base_url: str | None = None,
+        session_id: str | None = None,
+        max_width: int | None = DEFAULT_MAX_WIDTH,
+        max_height: int | None = None,
+        image_format: ImageFormat | None = DEFAULT_IMAGE_FORMAT,
+        quality: int = DEFAULT_QUALITY,
+    ) -> None:
+        super().__init__(environment_id, api_key=api_key, base_url=base_url, session_id=session_id)
+        self.max_width = max_width
+        self.max_height = max_height
+        self.image_format = image_format
+        self.quality = quality
+
+    def create_driver(self) -> DesktopDriverInterface:
         # Runtime import: hai-drivers is absent unless installed with hai-agents[desktop].
         try:
             from hai_drivers.desktop.local import LocalDesktopDriver
+            from hai_drivers.desktop.scaled import ScaledDesktopDriver
         except ImportError as exc:
             raise ImportError(
-                "Local desktop control requires extra deps. Install with: pip install 'hai-agents[desktop]'"
+                "Local desktop control requires hai-drivers>=0.1.2. Install with: pip install 'hai-agents[desktop]'"
             ) from exc
-        return LocalDesktopDriver()
+        driver = LocalDesktopDriver()
+        if self.max_width is None and self.max_height is None and self.image_format is None:
+            return driver
+        return ScaledDesktopDriver(
+            driver,
+            max_width=self.max_width,
+            max_height=self.max_height,
+            image_format=self.image_format,
+            quality=self.quality,
+        )
 
     def driver_interface(self) -> type:
         # Runtime import: hai-drivers is absent unless installed with hai-agents[desktop].

tests/test_local.py

@@ -119,6 +119,28 @@ def test_bridge_rejects_non_uuid_session_id(self):
         with pytest.raises(ValueError, match="must be a UUID"):
             PyautoguiDesktopBridge(api_key=API_KEY, session_id="my-laptop-1")
 
+    def test_desktop_screenshots_are_scaled_at_the_source(self, monkeypatch):
+        from hai_drivers.desktop import local as local_module
+        from hai_drivers.desktop.scaled import ScaledDesktopDriver
+
+        monkeypatch.setattr(local_module, "LocalDesktopDriver", _FakeLocalDriver)
+        bridge = PyautoguiDesktopBridge(api_key=API_KEY, max_width=1280, image_format="webp", quality=70)
+        driver = bridge.create_driver()
+        assert isinstance(driver, ScaledDesktopDriver)
+        assert (driver._max_width, driver._image_format, driver._quality) == (1280, "webp", 70)
+        assert isinstance(driver._driver, _FakeLocalDriver)
+
+    def test_desktop_bridge_with_all_knobs_off_serves_the_raw_driver(self, monkeypatch):
+        from hai_drivers.desktop import local as local_module
+
+        monkeypatch.setattr(local_module, "LocalDesktopDriver", _FakeLocalDriver)
+        bridge = PyautoguiDesktopBridge(api_key=API_KEY, max_width=None, image_format=None)
+        assert isinstance(bridge.create_driver(), _FakeLocalDriver)
+
+
+class _FakeLocalDriver:
+    """Stands in for LocalDesktopDriver, whose constructor needs a display."""
+
 
 class TestAutoStart:
     def test_create_session_starts_bridges_and_stamps_matching_session_ids(self, monkeypatch):