diff --git a/plugins/_browser/api/ws_browser.py b/plugins/_browser/api/ws_browser.py index 566ec10b8..f1c28c42d 100644 --- a/plugins/_browser/api/ws_browser.py +++ b/plugins/_browser/api/ws_browser.py @@ -477,7 +477,10 @@ class WsBrowser(WsHandler): def frame_consumer(frame: dict[str, Any]): return asyncio.run_coroutine_threadsafe(emit_frame(frame), server_loop) - await runtime.call("attach_screencast_consumer", stream_id, frame_consumer) + def stop_consumer() -> None: + server_loop.call_soon_threadsafe(stop_event.set) + + await runtime.call("attach_screencast_consumer", stream_id, frame_consumer, stop_consumer) while True: if stop_event.is_set(): diff --git a/plugins/_browser/helpers/runtime.py b/plugins/_browser/helpers/runtime.py index 6031a7f5f..8f27cdc73 100644 --- a/plugins/_browser/helpers/runtime.py +++ b/plugins/_browser/helpers/runtime.py @@ -318,8 +318,10 @@ class _BrowserScreencast: self.session = session self.mime = mime self.frame_consumer: Any | None = None + self.stop_callback: Any | None = None self.queue = asyncio.Queue(maxsize=1) self.stopped = False + self._closed = False self._ack_tasks: set[asyncio.Task] = set() self._expected_width = 0 self._expected_height = 0 @@ -399,15 +401,17 @@ class _BrowserScreencast: raise RuntimeError("Browser screencast stopped.") return frame - async def attach_consumer(self, frame_consumer: Any) -> None: + async def attach_consumer(self, frame_consumer: Any, stop_callback: Any | None = None) -> None: self.frame_consumer = frame_consumer + self.stop_callback = stop_callback frame = await self.pop_frame() if frame: await self._deliver_frame(frame) async def stop(self) -> None: - if self.stopped: + if self._closed: return + self._closed = True self.stopped = True self._drop_queued_frames() with contextlib.suppress(asyncio.QueueFull): @@ -431,6 +435,7 @@ class _BrowserScreencast: async def _handle_frame(self, params: dict[str, Any]) -> None: stop_after_ack = False + notify_stop = False try: data = params.get("data") or "" if data: @@ -453,6 +458,7 @@ class _BrowserScreencast: except Exception: if self.frame_consumer: stop_after_ack = True + notify_stop = True else: raise finally: @@ -465,6 +471,14 @@ class _BrowserScreencast: ) if stop_after_ack: self.stopped = True + if notify_stop: + self._notify_stopped() + + def _notify_stopped(self) -> None: + if not self.stop_callback: + return + with contextlib.suppress(Exception): + self.stop_callback() async def _deliver_frame(self, frame: dict[str, Any]) -> None: if not self.frame_consumer: @@ -1694,11 +1708,16 @@ class _BrowserRuntimeCore: raise KeyError("Browser screencast is not active.") return await screencast.pop_frame() - async def attach_screencast_consumer(self, stream_id: str, frame_consumer: Any) -> None: + async def attach_screencast_consumer( + self, + stream_id: str, + frame_consumer: Any, + stop_callback: Any | None = None, + ) -> None: screencast = self.screencasts.get(str(stream_id or "")) if not screencast: raise KeyError("Browser screencast is not active.") - await screencast.attach_consumer(frame_consumer) + await screencast.attach_consumer(frame_consumer, stop_callback) async def stop_screencast(self, stream_id: str) -> None: screencast = self.screencasts.pop(str(stream_id or ""), None) diff --git a/plugins/_browser/webui/browser-store.js b/plugins/_browser/webui/browser-store.js index 1b8bb9e26..6e9a779a8 100644 --- a/plugins/_browser/webui/browser-store.js +++ b/plugins/_browser/webui/browser-store.js @@ -1423,8 +1423,8 @@ const model = { paintFrameBitmap(bitmap) { const canvas = this._frameCanvas; if (!canvas || !bitmap?.width || !bitmap?.height) return false; - canvas.width = bitmap.width; - canvas.height = bitmap.height; + if (canvas.width !== bitmap.width) canvas.width = bitmap.width; + if (canvas.height !== bitmap.height) canvas.height = bitmap.height; const context = canvas.getContext("2d"); if (!context) return false; context.drawImage(bitmap, 0, 0); diff --git a/tests/test_browser_agent_regressions.py b/tests/test_browser_agent_regressions.py index c11ab4da3..67ee8c9ba 100644 --- a/tests/test_browser_agent_regressions.py +++ b/tests/test_browser_agent_regressions.py @@ -1180,6 +1180,9 @@ def test_browser_viewer_defaults_to_live_screencast_with_snapshot_fallback(): assert '"Page.screencastFrameAck"' in runtime assert "async def attach_screencast_consumer" in runtime assert "async def attach_consumer" in runtime + assert "stop_callback: Any | None = None" in runtime + assert "def _notify_stopped(self) -> None:" in runtime + assert "server_loop.call_soon_threadsafe(stop_event.set)" in ws_browser assert "await self._deliver_frame(" in runtime assert "await asyncio.wrap_future(future)" in runtime assert '"Page.stopScreencast"' in runtime @@ -1213,6 +1216,7 @@ def test_browser_viewer_defaults_to_live_screencast_with_snapshot_fallback(): assert "let dimensions = options?.dimensions || null" in browser_store assert "dimensions ||= await loadFrameDimensions(frameSrc)" in browser_store assert "paintFrameBitmap(bitmap)" in browser_store + assert "if (canvas.width !== bitmap.width) canvas.width = bitmap.width;" in browser_store assert "freezeCanvasFrameToImage()" in browser_store assert 'this._frameCanvas.toDataURL("image/jpeg", 0.86)' in browser_store assert "this.frameSrc = frameSrc;" in browser_store @@ -1697,6 +1701,60 @@ async def test_browser_screencast_acks_after_consumer_settles(): await screencast.stop() +@pytest.mark.anyio +async def test_browser_screencast_notifies_consumer_when_frame_task_stops_before_delivery(): + class FakeSession: + def __init__(self): + self.handlers = {} + self.sent = [] + self.detached = False + + def on(self, event, handler): + self.handlers[event] = handler + + async def send(self, method, params=None): + self.sent.append((method, params or {})) + + async def detach(self): + self.detached = True + + session = FakeSession() + delivered = [] + stopped = [] + screencast = _BrowserScreencast( + stream_id="stream", + browser_id=7, + session=session, + mime="image/jpeg", + ) + + await screencast.start(quality=92, every_nth_frame=1, viewport={"width": 640, "height": 480}) + await screencast.attach_consumer( + lambda frame: delivered.append(frame), + lambda: stopped.append(True), + ) + + def fail_jpeg_probe(_data): + raise RuntimeError("jpeg probe failed") + + screencast._jpeg_size = fail_jpeg_probe + session.handlers["Page.screencastFrame"]( + {"data": SMALL_JPEG_10X10, "metadata": {}, "sessionId": 29} + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert delivered == [] + assert stopped == [True] + assert screencast.stopped is True + assert ("Page.screencastFrameAck", {"sessionId": 29}) in session.sent + + await screencast.stop() + + assert ("Page.stopScreencast", {}) in session.sent + assert session.detached is True + + @pytest.mark.anyio async def test_browser_screencast_attach_consumer_flushes_queued_frame(): class FakeSession: