diff --git a/cdp_use/client.py b/cdp_use/client.py index 893215d..46a57e5 100644 --- a/cdp_use/client.py +++ b/cdp_use/client.py @@ -4,6 +4,7 @@ import re import time from typing import TYPE_CHECKING, Any, Dict, Optional +from cdp_use.recorder import Recorder import websockets @@ -240,6 +241,7 @@ def __init__( self.msg_id: int = 0 self.pending_requests: Dict[int, asyncio.Future] = {} self._message_handler_task = None + self._recorder: Optional[Recorder] = None # Initialize the type-safe CDP library from cdp_use.cdp.library import CDPLibrary @@ -279,6 +281,12 @@ async def start(self): async def stop(self): """Stop the message handler and close the WebSocket connection""" + # Stop active recording before closing the WebSocket so screencastFrameAck + # calls in the recorder can still complete cleanly + if self._recorder is not None: + await self._recorder.stop() + self._recorder = None + # Cancel the message handler task if self._message_handler_task: self._message_handler_task.cancel() @@ -400,3 +408,18 @@ async def emit_event( produced by your application rather than the browser. """ return await self._event_registry.handle_event(method, params or {}, session_id) + + async def start_recording(self, output_dir: str) -> "Recorder": + """ + Start recording browser session frames. + + :param output_dir: Directory to save frames + :return: Recorder instance + :raises RuntimeError: If recording is already in progress + """ + if self._recorder is not None: + raise RuntimeError("Recording already in progress. Call recorder.stop() first.") + recorder = Recorder(self, output_dir) + await recorder.start() + self._recorder = recorder + return recorder diff --git a/cdp_use/recorder.py b/cdp_use/recorder.py new file mode 100644 index 0000000..0903e60 --- /dev/null +++ b/cdp_use/recorder.py @@ -0,0 +1,90 @@ +import os +import base64 +import asyncio +from typing import Any + + +class Recorder: + """ + Records browser frames using CDP screencast and saves them as JPEG images. + """ + + def __init__(self, client: Any, output_dir: str): + self.client = client + self.output_dir = output_dir + self.frame_count = 0 + self._running = False + self._queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + self._worker_task: asyncio.Task | None = None + + async def start(self) -> None: + """ + Start recording frames. + """ + os.makedirs(self.output_dir, exist_ok=True) + self._running = True + + async def worker(): + while self._running or not self._queue.empty(): + event = await self._queue.get() + + try: + self.frame_count += 1 + + # Decode base64 frame into binary image + image_data = base64.b64decode(event["data"]) + + filename = os.path.join( + self.output_dir, f"frame_{self.frame_count:04d}.jpg" + ) + + with open(filename, "wb") as f: + f.write(image_data) + + print(f"Saved {filename}") + + # Acknowledge frame so Chrome continues sending frames + await self.client.send.Page.screencastFrameAck({ + "sessionId": event["sessionId"] + }) + finally: + self._queue.task_done() + + def on_frame(event: dict, session_id: str) -> None: + if self._running: + self._queue.put_nowait(event) + + self.client.register.Page.screencastFrame(on_frame) + + self._worker_task = asyncio.create_task(worker()) + + await self.client.send.Page.enable() + + await self.client.send.Page.startScreencast({ + "format": "jpeg", + "quality": 50, + "everyNthFrame": 1 + }) + + async def stop(self) -> None: + """ + Stop recording and finalize frame saving. + """ + self._running = False + + # Clear the client's back-reference so start_recording() can be called again + if self.client._recorder is self: + self.client._recorder = None + + await self.client.send.Page.stopScreencast() + + await self._queue.join() + + if self._worker_task: + self._worker_task.cancel() + try: + await self._worker_task + except asyncio.CancelledError: + pass + + print(f"Recording saved to {self.output_dir}") \ No newline at end of file diff --git a/examples/record.py b/examples/record.py new file mode 100644 index 0000000..b7541fa --- /dev/null +++ b/examples/record.py @@ -0,0 +1,22 @@ +import asyncio +from cdp_use.client import CDPClient + + +async def main(): + ws_url = "ws://localhost:9222/devtools/page/XXXX" + + async with CDPClient(ws_url) as client: + # Start recording browser session + recorder = await client.start_recording("recording_output") + + try: + # Perform actions while recording + await client.send.Page.navigate({"url": "https://youtube.com"}) + await asyncio.sleep(5) + finally: + # Ensure recording is properly stopped + await recorder.stop() + + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/tests/test_recorder.py b/tests/test_recorder.py new file mode 100644 index 0000000..dc1cd2d --- /dev/null +++ b/tests/test_recorder.py @@ -0,0 +1,92 @@ +# ABOUTME: Tests for three specific fixes in recorder.py and client.py. +# ABOUTME: Covers try/finally task_done, double-start guard, and _recorder back-ref clearing. + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from cdp_use.recorder import Recorder + + +def make_fake_client(): + """Return a minimal fake client that stubs the CDP calls Recorder needs.""" + client = MagicMock() + client._recorder = None + client.send.Page.enable = AsyncMock() + client.send.Page.startScreencast = AsyncMock() + client.send.Page.stopScreencast = AsyncMock() + client.send.Page.screencastFrameAck = AsyncMock() + # register.Page.screencastFrame just stores the callback; capture it for tests + client.register.Page.screencastFrame = MagicMock() + return client + + +# --------------------------------------------------------------------------- +# Fix 1: try/finally ensures task_done() is called even when frame processing +# raises, so queue.join() doesn't hang. +# --------------------------------------------------------------------------- + +def test_task_done_called_after_frame_processing_exception(): + async def run(): + client = make_fake_client() + recorder = Recorder(client, "/tmp/test_frames") + + await recorder.start() + + # Grab the on_frame callback that was registered + on_frame = client.register.Page.screencastFrame.call_args[0][0] + + # Push a malformed event — missing "data" key — so base64.b64decode raises + on_frame({"sessionId": "abc"}, "abc") + + # queue.join() must complete; if task_done() wasn't called it would hang + await asyncio.wait_for(recorder._queue.join(), timeout=2.0) + + recorder._running = False + recorder._worker_task.cancel() + try: + await recorder._worker_task + except (asyncio.CancelledError, KeyError): + pass + + asyncio.run(run()) + + +# --------------------------------------------------------------------------- +# Fix 2: Calling start_recording() twice raises RuntimeError. +# --------------------------------------------------------------------------- + +def test_start_recording_twice_raises(): + async def run(): + from cdp_use.client import CDPClient + + client = CDPClient("ws://fake") + # Simulate a recorder already being active + client._recorder = object() + + with pytest.raises(RuntimeError, match="Recording already in progress"): + await client.start_recording("/tmp/out") + + asyncio.run(run()) + + +# --------------------------------------------------------------------------- +# Fix 3: After recorder.stop(), client._recorder is cleared to None. +# --------------------------------------------------------------------------- + +def test_recorder_stop_clears_client_back_ref(): + async def run(): + client = make_fake_client() + recorder = Recorder(client, "/tmp/test_frames") + + await recorder.start() + + # Point the client back-ref at this recorder (as start_recording() does) + client._recorder = recorder + + await recorder.stop() + + assert client._recorder is None + + asyncio.run(run())