feat: capture-chunk reassembly + screencast API (v2.36.2)

ceki_sdk/__init__.py

@@ -21,7 +21,7 @@
 from ._profile import BrowserProfile
 from .humanize import HumanProfile
 
-__version__ = "2.36.1"
+__version__ = "2.36.2"
 __all__ = [
     "connect",
     "ConnectOptions",

ceki_sdk/_browser.py

@@ -40,6 +40,7 @@
 TabOpenedCallback = Callable[[str], Awaitable[None]]
 SimpleCallback = Callable[[], Awaitable[None]]
 UserEventCallback = Callable[[list[dict[str, Any]]], Awaitable[None]]
+CaptureFrameCallback = Callable[[dict[str, Any]], Awaitable[None]]
 
 _ERROR_TERMINAL = {-1011, -1012, -1015, -1018}
 
@@ -98,6 +99,7 @@ def __init__(self, client: "Client", match: Match, *, human="natural") -> None:
         self._provider_disconnected_callbacks: list[SimpleCallback] = []
         self._provider_reconnected_callbacks: list[SimpleCallback] = []
         self._user_event_callbacks: list[UserEventCallback] = []
+        self._capture_frame_callbacks: list[CaptureFrameCallback] = []
         self._ended = asyncio.Event()
         self._ended_reason: str | None = None
 
@@ -254,6 +256,39 @@ def on_provider_reconnected(self, callback: SimpleCallback) -> None:
     def on_user_event(self, callback: UserEventCallback) -> None:
         self._user_event_callbacks.append(callback)
 
+    def on_capture_frame(self, callback: CaptureFrameCallback) -> None:
+        """Register a callback that receives screencast video frames.
+
+        Frames arrive on the P2P ``ceki-capture`` data channel — the extension
+        intercepts ``Page.startScreencast`` and streams frames there via its
+        capture bridge instead of emitting CDP ``Page.screencastFrame`` events.
+        Each callback is invoked with the raw capture frame dict::
+
+            {"type": "video-frame", "data": "<base64 jpeg>",
+             "width": ..., "height": ..., "timestamp": ...}
+
+        Frames that exceed the chunk threshold arrive as ``capture-chunk``
+        fragments and are reassembled transparently before delivery.
+        """
+        self._capture_frame_callbacks.append(callback)
+
+    async def start_screencast(self, **params: Any) -> dict[str, Any]:
+        """Start streaming video frames to :meth:`on_capture_frame` callbacks.
+
+        Sends ``Page.startScreencast`` (intercepted by the extension and served
+        by its capture bridge). Supported params mirror CDP::
+
+            maxWidth, maxHeight, quality, everyNthFrame, maxFrameRate
+
+        Frames arrive asynchronously on the capture data channel and are
+        delivered to every registered callback.
+        """
+        return await self.send({"method": "Page.startScreencast", "params": params})
+
+    async def stop_screencast(self) -> dict[str, Any]:
+        """Stop the screencast stream (sends ``Page.stopScreencast``)."""
+        return await self.send({"method": "Page.stopScreencast"})
+
     async def switch_tab(self) -> None:
         await self._client._ws_send({"type": "switch_tab", "session_id": self.session_id})
 
@@ -944,6 +979,18 @@ async def _on_cdp_event(self, msg: dict[str, Any]) -> None:
         for cb in self._event_callbacks:
             asyncio.create_task(cast(Coroutine, cb(method, params)))
 
+    async def _on_capture_data(self, msg: dict[str, Any]) -> None:
+        """Dispatch a capture-DC message to :meth:`on_capture_frame` callbacks.
+
+        Only ``video-frame`` messages are forwarded — ``video-stopped`` and
+        screenshot messages that also travel on the capture DC are not frames
+        and are ignored by the screencast API.
+        """
+        if msg.get("type") != "video-frame":
+            return
+        for cb in self._capture_frame_callbacks:
+            asyncio.create_task(cast(Coroutine, cb(msg)))
+
     async def _on_tab_opened(self, msg: dict[str, Any]) -> None:
         url = msg.get("url", "")
         for cb in self._tab_opened_callbacks:

ceki_sdk/_client.py

@@ -684,6 +684,17 @@ async def _on_cdp(msg: dict[str, Any]) -> None:
 
             transport.on_cdp_message = _on_cdp
 
+            # Wire capture-data callback → route capture frames to the active
+            # browser. video-frame messages arrive on the ceki-capture DC (the
+            # extension intercepts Page.startScreencast and streams frames via
+            # capture-bridge) rather than as CDP screencastFrame events.
+            async def _on_capture(msg: dict[str, Any]) -> None:
+                browser = self._active_browsers.get(session_id)
+                if browser:
+                    await browser._on_capture_data(msg)
+
+            transport.on_capture_data = _on_capture
+
             # Wire connection state callback for lifecycle monitoring
             async def _on_conn_state(state: str) -> None:
                 log.info("p2p: connection state -> %s", state)
@@ -793,6 +804,15 @@ async def _on_cdp(msg_inner: dict[str, Any]) -> None:
 
             transport.on_cdp_message = _on_cdp
 
+            # Wire capture-data callback → route capture frames to the active
+            # browser (same as _init_p2p — video-frame arrives on ceki-capture DC).
+            async def _on_capture(msg_inner: dict[str, Any]) -> None:
+                browser = self._active_browsers.get(session_id)
+                if browser:
+                    await browser._on_capture_data(msg_inner)
+
+            transport.on_capture_data = _on_capture
+
             # Wire connection state callback
             async def _on_conn_state(state: str) -> None:
                 log.info("p2p: connection state -> %s", state)

ceki_sdk/_webrtc.py

@@ -22,6 +22,7 @@
 import logging
 import os
 import re
+import time
 from typing import Any, Callable, Coroutine
 
 log = logging.getLogger(__name__)
@@ -126,6 +127,7 @@ def __init__(
     ) -> None:
         self._pc: Any = None  # aiortc.RTCPeerConnection
         self._cmd_dc: Any = None  # aiortc.RTCDataChannel
+        self._capture_dc: Any = None  # aiortc.RTCDataChannel (ceki-capture)
 
         # ICE servers: constructor arg → CEKI_TURN_SERVERS env → default STUN
         env_servers_raw = os.environ.get("CEKI_TURN_SERVERS")
@@ -175,6 +177,9 @@ def __init__(
         # Callbacks — set by consumer (_client.py)
         self.on_ice_candidate: Callable[[dict[str, Any]], Coroutine[Any, Any, None] | None] | None = None
         self.on_cdp_message: Callable[[dict[str, Any]], Coroutine[Any, Any, None] | None] | None = None
+        self.on_capture_data: (
+            Callable[[dict[str, Any]], Coroutine[Any, Any, None] | None] | None
+        ) = None
         self.on_connection_state: Callable[[str], Coroutine[Any, Any, None] | None] | None = None
         self.on_data_channel_state: Callable[[str], Coroutine[Any, Any, None] | None] | None = None
 
@@ -188,6 +193,15 @@ def __init__(
         # a lost chunk surfaces as an SDK-side timeout, existing mechanism.
         self._pending_chunks: dict[str, dict[str, Any]] = {}
 
+        # Capture-chunk reassembly buffer for large video-frame/screenshot data
+        # sent over the ceki-capture DC. Keyed by frameId →
+        # {chunks:[slice,...], received, total, received_at}. The capture DC is
+        # created with ordered:false, so fragments can arrive out of order and
+        # can be dropped — incomplete frames are pruned after
+        # ``_capture_stale_ms`` (mirrors extension CaptureChunkReassembler).
+        self._pending_capture_frames: dict[str, dict[str, Any]] = {}
+        self._capture_stale_ms = 5000
+
     async def _ensure_pc(self) -> Any:
         """Lazy-create the RTCPeerConnection on first use."""
         if self._pc is not None:
@@ -240,8 +254,8 @@ def _on_dc(channel: Any) -> None:
                 self._cmd_dc = channel
                 self._wire_cmd_dc(channel)
             elif channel.label == "ceki-capture":
-                # Agent doesn't process capture frames, but log it
-                log.info("webrtc: ceki-capture channel opened (no-op for agent)")
+                self._capture_dc = channel
+                self._wire_capture_dc(channel)
 
         return self._pc
 
@@ -283,6 +297,104 @@ async def _on_message(message: str | bytes) -> None:
             if self.on_cdp_message:
                 await self.on_cdp_message(data)
 
+    def _wire_capture_dc(self, channel: Any) -> None:
+        """Set up message/close handlers on the ceki-capture data channel.
+
+        Mirror of ``_wire_cmd_dc`` for the capture channel: small frames
+        (single ``video-frame`` / screenshot messages) are forwarded to
+        ``on_capture_data`` unchanged, while ``capture-chunk`` fragments are
+        reassembled transparently before delivery.
+        """
+
+        @channel.on("open")
+        async def _on_open() -> None:
+            log.info("webrtc: ceki-capture DC opened")
+
+        @channel.on("close")
+        async def _on_close() -> None:
+            log.info("webrtc: ceki-capture DC closed")
+
+        @channel.on("message")
+        async def _on_message(message: str | bytes) -> None:
+            try:
+                data = json.loads(message if isinstance(message, str) else message.decode())
+            except (json.JSONDecodeError, UnicodeDecodeError) as exc:
+                log.warning("webrtc: failed to parse capture DC message: %s", exc)
+                return
+
+            # Chunked capture frames (large messages from the extension) are
+            # reassembled transparently here — individual chunks are never
+            # forwarded to on_capture_data.
+            if data.get("type") == "capture-chunk":
+                restored = self._buffer_capture_chunk(data)
+                if restored is None:
+                    return  # not complete yet (or malformed)
+                if self.on_capture_data:
+                    await self.on_capture_data(restored)
+                return
+
+            if self.on_capture_data:
+                await self.on_capture_data(data)
+
+    def _buffer_capture_chunk(self, chunk: dict[str, Any]) -> dict[str, Any] | None:
+        """Buffer one capture-chunk fragment and return the reassembled frame.
+
+        Mirrors the extension's ``CaptureChunkReassembler.handle``: fragments
+        are buffered per ``frameId`` until all ``total`` have arrived (in any
+        order — the capture DC is ordered:false), then the concatenated
+        payload is parsed back into the original frame and returned.
+        Incomplete frames are pruned after ``_capture_stale_ms`` so a dropped
+        fragment cannot leak memory forever.
+        """
+        frame_id = chunk.get("frameId")
+        seq = chunk.get("seq")
+        total = chunk.get("total")
+        payload = chunk.get("payload")
+        if (
+            not isinstance(frame_id, str)
+            or not isinstance(seq, int)
+            or not isinstance(total, int)
+            or seq < 0
+            or total <= 0
+            or seq >= total
+            or not isinstance(payload, str)
+        ):
+            log.warning("webrtc: malformed capture-chunk, dropping")
+            return None
+
+        now = time.monotonic()
+        for fid in [
+            fid
+            for fid, entry in self._pending_capture_frames.items()
+            if now - entry["received_at"] > self._capture_stale_ms
+        ]:
+            log.debug("webrtc: pruning stale capture-chunk frame %s", fid)
+            del self._pending_capture_frames[fid]
+
+        entry = self._pending_capture_frames.get(frame_id)
+        if entry is None:
+            entry = {
+                "chunks": [""] * total,
+                "received": 0,
+                "total": total,
+                "received_at": now,
+            }
+            self._pending_capture_frames[frame_id] = entry
+
+        if entry["chunks"][seq] == "":
+            entry["chunks"][seq] = payload
+            entry["received"] += 1
+
+        if entry["received"] != entry["total"]:
+            return None
+
+        del self._pending_capture_frames[frame_id]
+        try:
+            return json.loads("".join(entry["chunks"]))
+        except (json.JSONDecodeError, UnicodeDecodeError) as exc:
+            log.warning("webrtc: failed to reassemble capture-chunk frame %s: %s", frame_id, exc)
+            return None
+
     def _buffer_chunk(self, chunk: dict[str, Any]) -> dict[str, Any] | None:
         """Buffer one CDP chunk fragment and return the reassembled message.
 
@@ -520,6 +632,12 @@ async def close(self) -> None:
             except Exception:
                 pass
             self._cmd_dc = None
+        if self._capture_dc is not None:
+            try:
+                self._capture_dc.close()
+            except Exception:
+                pass
+            self._capture_dc = None
         if self._pc is not None:
             try:
                 await self._pc.close()
@@ -529,4 +647,6 @@ async def close(self) -> None:
         self._dc_open_event.clear()
         self._local_fingerprint = None
         self._pending_remote_candidates.clear()
+        self._pending_chunks.clear()
+        self._pending_capture_frames.clear()
         log.info("webrtc: transport closed")