From 25f8e28734056d662f48403c7e585c3f2583ab46 Mon Sep 17 00:00:00 2001 From: LawyZheng Date: Tue, 30 Jun 2026 02:19:41 +0800 Subject: [PATCH] SKY-10723: Add worker-side CDP livestream publishing (#6895) --- skyvern/webeye/browser_artifacts.py | 4 + skyvern/webeye/browser_factory.py | 3 + skyvern/webeye/browser_state.py | 4 +- skyvern/webeye/cdp_frame_publisher.py | 294 +++++++++++ skyvern/webeye/real_browser_manager.py | 171 ++++++- skyvern/webeye/real_browser_state.py | 22 + .../test_browser_context_factory_marker.py | 70 +++ tests/unit/test_real_browser_manager.py | 8 +- ...st_real_browser_manager_frame_publisher.py | 365 ++++++++++++++ ...r_manager_publisher_self_stops_on_close.py | 107 ++++ ...al_browser_manager_publisher_start_race.py | 149 ++++++ .../test_real_browser_state_on_close_hook.py | 79 +++ ...test_stream_ref_dec_no_publisher_access.py | 83 ++++ tests/unit/test_worker_cdp_frame_publisher.py | 465 ++++++++++++++++++ 14 files changed, 1815 insertions(+), 9 deletions(-) create mode 100644 skyvern/webeye/cdp_frame_publisher.py create mode 100644 tests/unit/test_browser_context_factory_marker.py create mode 100644 tests/unit/test_real_browser_manager_frame_publisher.py create mode 100644 tests/unit/test_real_browser_manager_publisher_self_stops_on_close.py create mode 100644 tests/unit/test_real_browser_manager_publisher_start_race.py create mode 100644 tests/unit/test_real_browser_state_on_close_hook.py create mode 100644 tests/unit/test_stream_ref_dec_no_publisher_access.py create mode 100644 tests/unit/test_worker_cdp_frame_publisher.py diff --git a/skyvern/webeye/browser_artifacts.py b/skyvern/webeye/browser_artifacts.py index a49e71232..05bfd95e8 100644 --- a/skyvern/webeye/browser_artifacts.py +++ b/skyvern/webeye/browser_artifacts.py @@ -22,6 +22,10 @@ class BrowserArtifacts(BaseModel): traces_dir: str | None = None browser_session_dir: str | None = None browser_console_log_path: str | None = None + # Set by remote-CDP creators so RealBrowserManager attaches a CDP frame + # publisher. Local Playwright contexts leave it False. Type must stay + # ``bool`` — the manager guards with ``is True`` identity, not truthiness. + needs_cdp_frame_publisher: bool = False _browser_console_log_lock: asyncio.Lock = PrivateAttr(default_factory=asyncio.Lock) async def append_browser_console_log(self, msg: str) -> int: diff --git a/skyvern/webeye/browser_factory.py b/skyvern/webeye/browser_factory.py index bcab43eeb..1c3c31bdb 100644 --- a/skyvern/webeye/browser_factory.py +++ b/skyvern/webeye/browser_factory.py @@ -1008,6 +1008,9 @@ async def _connect_to_cdp_browser( browser_artifacts = BrowserContextFactory.build_browser_artifacts( har_path=browser_args["record_har_path"], ) + # Single chokepoint for OSS remote-CDP creation; stamp the marker so + # RealBrowserManager attaches the CDP frame publisher. + browser_artifacts.needs_cdp_frame_publisher = True LOG.info("Connecting browser CDP connection", remote_browser_url=remote_browser_url) cdp_headers = merge_cdp_connect_headers( diff --git a/skyvern/webeye/browser_state.py b/skyvern/webeye/browser_state.py index 01e0d8842..cbf053a9c 100644 --- a/skyvern/webeye/browser_state.py +++ b/skyvern/webeye/browser_state.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Literal, Protocol +from typing import Awaitable, Callable, Literal, Protocol from playwright.async_api import BrowserContext, Page, Playwright @@ -18,6 +18,8 @@ class BrowserState(Protocol): browser_cleanup: BrowserCleanupFunc pw: Playwright + def add_on_close(self, callback: Callable[[], Awaitable[None]]) -> None: ... + async def check_and_fix_state( self, url: str | None = None, diff --git a/skyvern/webeye/cdp_frame_publisher.py b/skyvern/webeye/cdp_frame_publisher.py new file mode 100644 index 000000000..10d986813 --- /dev/null +++ b/skyvern/webeye/cdp_frame_publisher.py @@ -0,0 +1,294 @@ +"""Worker-side CDP frame publisher for reused-CDP / remote browser contexts the in-process screencast cannot reach.""" + +from __future__ import annotations + +import asyncio +import base64 +import binascii +import hashlib +import os +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING + +import structlog + +from skyvern.forge import app +from skyvern.forge.sdk.api.files import get_skyvern_temp_dir + +if TYPE_CHECKING: + from playwright.async_api import CDPSession, Page + + from skyvern.webeye.browser_state import BrowserState + +LOG = structlog.get_logger() + + +DEFAULT_CAPTURE_INTERVAL_SECONDS: float = 1.0 + + +def _write_frame_atomically(temp_dir: Path, stream_key: str, data: bytes) -> None: + """Atomic tempfile+``os.replace`` write; intended to run on a worker thread.""" + temp_dir.mkdir(parents=True, exist_ok=True) + target = temp_dir / stream_key + tmp_fd, tmp_path = tempfile.mkstemp(dir=str(temp_dir), prefix=f".{stream_key}.", suffix=".tmp") + fp_taken = False + try: + with os.fdopen(tmp_fd, "wb") as fp: + fp_taken = True + fp.write(data) + os.replace(tmp_path, target) + tmp_path = None # type: ignore[assignment] + finally: + if not fp_taken: + try: + os.close(tmp_fd) + except OSError: + pass + if tmp_path is not None: + try: + os.unlink(tmp_path) + except OSError: + pass + + +class CDPFramePublisher: + """Periodically publishes the active page's PNG to the streaming storage key. + + ``stream_key`` is the bare key the API-side WebSocket polls via + ``StorageBase.get_streaming_file`` (e.g. ``"wr_123.png"``). CDP failures + are tolerated; the loop continues. ``RealBrowserManager`` calls + :meth:`start` after a working page exists and :meth:`stop` before closing + the browser context. + """ + + def __init__( + self, + *, + browser_state: BrowserState, + stream_key: str, + organization_id: str, + capture_interval_seconds: float = DEFAULT_CAPTURE_INTERVAL_SECONDS, + ) -> None: + self._browser_state = browser_state + self._stream_key = stream_key + self._organization_id = organization_id + self._capture_interval_seconds = max(capture_interval_seconds, 0.1) + + self._task: asyncio.Task[None] | None = None + self._stopped = asyncio.Event() + self._cdp_session: CDPSession | None = None + self._attached_page: Page | None = None + # Digest of the last frame whose write + upload both succeeded. Lets us + # dedupe identical frames without losing a retry on transient upload + # failure. + self._last_published_digest: bytes | None = None + + @property + def stream_key(self) -> str: + return self._stream_key + + @property + def is_running(self) -> bool: + return self._task is not None and not self._task.done() + + async def start(self) -> None: + """Spawn the background publish loop. Idempotent.""" + if self.is_running: + return + self._stopped.clear() + self._task = asyncio.create_task(self._run(), name=f"cdp-frame-publisher:{self._stream_key}") + LOG.info( + "CDP frame publisher started", + stream_key=self._stream_key, + organization_id=self._organization_id, + interval_seconds=self._capture_interval_seconds, + ) + + async def stop(self) -> None: + """Cancel the loop, detach CDP session, and reset state. Idempotent.""" + self._stopped.set() + task = self._task + self._task = None + if task is not None and not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + await self._detach_cdp_session() + self._last_published_digest = None + LOG.info( + "CDP frame publisher stopped", + stream_key=self._stream_key, + organization_id=self._organization_id, + ) + + async def _run(self) -> None: + try: + while not self._stopped.is_set(): + # Self-terminate if the browser context is gone (persistent-session + # handoff with ``close_browser_on_completion=False``, crash, or any + # other path that drops the context without firing on-close). + if not self._browser_state_is_connected(): + LOG.info( + "CDP frame publisher self-terminating: browser context disconnected", + stream_key=self._stream_key, + organization_id=self._organization_id, + ) + self._stopped.set() + break + + try: + await self._publish_one_frame() + except asyncio.CancelledError: + raise + except Exception: + LOG.warning( + "CDP frame publish iteration failed", + stream_key=self._stream_key, + organization_id=self._organization_id, + exc_info=True, + ) + await self._detach_cdp_session() + + try: + await asyncio.wait_for(self._stopped.wait(), timeout=self._capture_interval_seconds) + except asyncio.TimeoutError: + continue + except asyncio.CancelledError: + raise + except asyncio.CancelledError: + return + + def _browser_state_is_connected(self) -> bool: + """Cheap, never-raises wrapper around ``BrowserState.is_connected``. + + A stale ``BrowserState`` whose underlying Playwright driver has gone + away will report ``False`` here; we use it as the loop's keep-going + signal so the publisher does not spin forever after teardown. + """ + try: + return bool(self._browser_state.is_connected()) + except Exception: + return False + + async def _publish_one_frame(self) -> None: + page = await self._browser_state.get_working_page() + if page is None: + return + + if page is not self._attached_page or self._cdp_session is None: + await self._detach_cdp_session() + try: + self._cdp_session = await page.context.new_cdp_session(page) + except Exception: + LOG.warning( + "Could not open CDP session for frame publishing", + stream_key=self._stream_key, + organization_id=self._organization_id, + exc_info=True, + ) + self._cdp_session = None + self._attached_page = None + return + self._attached_page = page + self._last_published_digest = None + LOG.info( + "CDP frame publisher attached to page", + stream_key=self._stream_key, + organization_id=self._organization_id, + page_url=getattr(page, "url", ""), + ) + + try: + result = await self._cdp_session.send( + "Page.captureScreenshot", + { + "format": "png", + "captureBeyondViewport": False, + }, + ) + except Exception: + # Transient failure (target detached, navigation in progress, etc.). + # Reattach on the next tick; logged at debug so a misbehaving remote + # CDP doesn't flood the warning stream at 1 FPS. + LOG.debug( + "Page.captureScreenshot failed; will reattach next tick", + stream_key=self._stream_key, + organization_id=self._organization_id, + exc_info=True, + ) + await self._detach_cdp_session() + return + + encoded = result.get("data", "") if isinstance(result, dict) else "" + if not encoded: + return + try: + data = base64.b64decode(encoded, validate=False) + except (binascii.Error, ValueError): + return + if not data: + return + + # Content-addressable dedupe hash (not a security boundary); SHA-256 + # to satisfy security scanners that block SHA-1. + digest = hashlib.sha256(data).digest() + if digest == self._last_published_digest: + return + + write_ok = await self._write_frame(data) + if write_ok: + # Dedupe only after both local write and upload succeed, so a + # transient upload failure retries instead of getting deduped away. + self._last_published_digest = digest + + async def _write_frame(self, data: bytes) -> bool: + """Persist one frame; True iff both the local write and the upload succeeded.""" + temp_dir = Path(get_skyvern_temp_dir()) / self._organization_id + try: + # Blocking I/O runs on a worker thread so a large flush does not + # stall the event loop shared with other publishers / agent work. + await asyncio.to_thread(_write_frame_atomically, temp_dir, self._stream_key, data) + except OSError: + LOG.warning( + "Failed to write streaming frame to disk", + stream_key=self._stream_key, + organization_id=self._organization_id, + exc_info=True, + ) + return False + + # Local-disk storage reads the temp file directly; remote object-storage + # backends need an explicit upload. + try: + await app.STORAGE.save_streaming_file(self._organization_id, self._stream_key) + except Exception: + LOG.debug( + "save_streaming_file failed; will retry on next iteration", + stream_key=self._stream_key, + organization_id=self._organization_id, + exc_info=True, + ) + return False + return True + + async def _detach_cdp_session(self) -> None: + session = self._cdp_session + self._cdp_session = None + self._attached_page = None + if session is None: + return + try: + await session.detach() + except Exception: + pass + + +def stream_key_for_workflow_run(workflow_run_id: str) -> str: + return f"{workflow_run_id}.png" + + +def stream_key_for_task(task_id: str) -> str: + return f"{task_id}.png" diff --git a/skyvern/webeye/real_browser_manager.py b/skyvern/webeye/real_browser_manager.py index c82bc424b..b3ef3cee3 100644 --- a/skyvern/webeye/real_browser_manager.py +++ b/skyvern/webeye/real_browser_manager.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import os import structlog @@ -9,6 +10,7 @@ from skyvern.exceptions import MissingBrowserState from skyvern.forge import app from skyvern.forge.sdk.api.files import resolve_run_download_id from skyvern.forge.sdk.core import skyvern_context +from skyvern.forge.sdk.routes.streaming.registries import set_deferred_close_params, stream_ref_active from skyvern.forge.sdk.schemas.tasks import Task from skyvern.forge.sdk.workflow.models.workflow import WorkflowRun from skyvern.schemas.runs import ProxyLocation, ProxyLocationInput @@ -16,6 +18,11 @@ from skyvern.webeye.browser_artifacts import VideoArtifact from skyvern.webeye.browser_factory import BrowserContextFactory, rebind_download_dir from skyvern.webeye.browser_manager import BrowserManager from skyvern.webeye.browser_state import BrowserState +from skyvern.webeye.cdp_frame_publisher import ( + CDPFramePublisher, + stream_key_for_task, + stream_key_for_workflow_run, +) from skyvern.webeye.real_browser_state import RealBrowserState from skyvern.webeye.session_cookies import persist_session_cookies from skyvern.webeye.video_utils import finalize_webm @@ -32,9 +39,116 @@ def _merge_proxy_session_headers( return app.AGENT_FUNCTION.merge_proxy_session_extra_http_headers(extra_http_headers, proxy_session_id) +def _resolve_stream_key(*, workflow_run_id: str | None, task_id: str | None) -> str | None: + """Pick the stream key that the API-side WebSocket polls for this entity. + + Workflow-run streams always read ``{workflow_run_id}.png``; task streams use + that same key when the task belongs to a workflow run, and fall back to + ``{task_id}.png`` only for standalone tasks. See + ``skyvern/forge/sdk/routes/streaming/screenshot.py``. + """ + if workflow_run_id: + return stream_key_for_workflow_run(workflow_run_id) + if task_id: + return stream_key_for_task(task_id) + return None + + class RealBrowserManager(BrowserManager): def __init__(self) -> None: self.pages: dict[str, BrowserState] = {} + # CDP frame publishers keyed by stream key (``{wr}.png`` / ``{task}.png``). + self._frame_publishers: dict[str, CDPFramePublisher] = {} + # Serializes the check/create/start/store/register sequence in + # ``_start_frame_publisher`` so concurrent attaches for one stream key + # cannot orphan a publisher loop. + self._publisher_lock = asyncio.Lock() + + async def _start_frame_publisher( + self, + *, + browser_state: BrowserState, + workflow_run_id: str | None = None, + task_id: str | None = None, + organization_id: str | None = None, + ) -> None: + """Best-effort start a CDP frame publisher for this entity. + + Gated on ``browser_state.browser_artifacts.needs_cdp_frame_publisher``, + which remote-CDP creators stamp. Local Playwright contexts leave it + False and skip publishing. Never raises. + """ + # Strict equality; MagicMock attributes are truthy by default. + if browser_state.browser_artifacts.needs_cdp_frame_publisher is not True: + return + stream_key = _resolve_stream_key(workflow_run_id=workflow_run_id, task_id=task_id) + if not stream_key or not organization_id: + return + async with self._publisher_lock: + if stream_key in self._frame_publishers: + return + try: + publisher = CDPFramePublisher( + browser_state=browser_state, + stream_key=stream_key, + organization_id=organization_id, + ) + await publisher.start() + self._frame_publishers[stream_key] = publisher + except Exception: + LOG.warning( + "Failed to start CDP frame publisher; livestream may be unavailable", + stream_key=stream_key, + organization_id=organization_id, + exc_info=True, + ) + return + # Tie publisher lifetime to BrowserState.close() so any close path + # stops it without needing to know about the publisher registry. + captured_stream_key = stream_key + + async def _on_browser_state_close() -> None: + # Pop under the same lock that guards ``_start_frame_publisher`` + # so a concurrent restart cannot slip past the registry check + # and orphan a second publisher. ``pub.stop()`` runs outside + # the lock — it awaits the task's exit and must not block + # other publishers from starting. + async with self._publisher_lock: + pub = self._frame_publishers.pop(captured_stream_key, None) + if pub is None: + return + try: + await pub.stop() + except Exception: + LOG.debug( + "CDP frame publisher stop raised during browser-state close; ignored", + stream_key=captured_stream_key, + exc_info=True, + ) + + browser_state.add_on_close(_on_browser_state_close) + + async def _stop_frame_publisher( + self, + *, + workflow_run_id: str | None = None, + task_id: str | None = None, + ) -> None: + """Best-effort: stop the publisher matching this entity. Idempotent.""" + stream_key = _resolve_stream_key(workflow_run_id=workflow_run_id, task_id=task_id) + if not stream_key: + return + publisher = self._frame_publishers.pop(stream_key, None) + if publisher is None: + return + try: + await publisher.stop() + except Exception: + LOG.debug( + "CDP frame publisher stop raised; ignored", + stream_key=stream_key, + exc_info=True, + ) @staticmethod async def _create_browser_state( @@ -172,6 +286,12 @@ class RealBrowserManager(BrowserManager): cdp_connect_headers=task.cdp_connect_headers, browser_address=task.browser_address, ) + await self._start_frame_publisher( + browser_state=browser_state, + workflow_run_id=task.workflow_run_id, + task_id=task.task_id, + organization_id=task.organization_id, + ) return browser_state async def get_or_create_for_workflow_run( @@ -209,6 +329,14 @@ class RealBrowserManager(BrowserManager): self.pages[workflow_run_id] = browser_state if parent_workflow_run_id: self.pages[parent_workflow_run_id] = browser_state + # The workflow-run streaming endpoint reads ``{workflow_run_id}.png``, so the + # child needs its own publisher even when reusing the parent's browser state — + # the parent's publisher writes a different key. + await self._start_frame_publisher( + browser_state=browser_state, + workflow_run_id=workflow_run_id, + organization_id=workflow_run.organization_id, + ) return browser_state if browser_session_id: @@ -305,6 +433,11 @@ class RealBrowserManager(BrowserManager): browser_address=workflow_run.browser_address, browser_profile_id=browser_profile_id, ) + await self._start_frame_publisher( + browser_state=browser_state, + workflow_run_id=workflow_run.workflow_run_id, + organization_id=workflow_run.organization_id, + ) return browser_state def get_for_workflow_run( @@ -411,6 +544,21 @@ class RealBrowserManager(BrowserManager): async def close(self) -> None: LOG.info("Closing BrowserManager") + # Stop all streaming frame publishers before closing browsers so CDP + # sessions detach cleanly. Cancellation here is best-effort and must + # not block manager shutdown. + for stream_key in list(self._frame_publishers.keys()): + publisher = self._frame_publishers.pop(stream_key, None) + if publisher is None: + continue + try: + await publisher.stop() + except Exception: + LOG.debug( + "CDP frame publisher stop raised during manager close; ignored", + stream_key=stream_key, + exc_info=True, + ) for browser_state in self.pages.values(): await browser_state.close() self.pages = dict() @@ -435,6 +583,12 @@ class RealBrowserManager(BrowserManager): trace_path = f"{browser_state_to_close.browser_artifacts.traces_dir}/{task_id}.zip" await browser_state_to_close.browser_context.tracing.stop(path=trace_path) LOG.info("Stopped tracing", trace_path=trace_path) + # Standalone-task only: a workflow-owned task's publisher is keyed + # by ``workflow_run_id`` (see ``_resolve_stream_key``) and is stopped + # by ``cleanup_for_workflow_run``. Passing ``task_id`` here is the + # honest signal — it hits ``{task_id}.png`` for standalone tasks + # and is a deliberate no-op for workflow tasks. + await self._stop_frame_publisher(task_id=task_id) await browser_state_to_close.close(close_browser_on_completion=close_browser_on_completion) LOG.info("Task is cleaned up") @@ -467,9 +621,15 @@ class RealBrowserManager(BrowserManager): if child_workflow_run_ids: for child_id in child_workflow_run_ids: self.pages.pop(child_id, None) + # Child workflows skip their own cleanup, so the publishers + # started for inherited child runs would otherwise leak until + # process shutdown. Stop them here. + await self._stop_frame_publisher(workflow_run_id=child_id) - from skyvern.forge.sdk.routes.streaming.registries import set_deferred_close_params, stream_ref_active - + # Dual-stop is intentional and safe: both the explicit + # ``_stop_frame_publisher`` above and the ``add_on_close`` callback + # registered in ``_start_frame_publisher`` may fire for the same + # stream key. ``dict.pop(key, None)`` makes the second pop a no-op. streams_active = stream_ref_active(workflow_run_id) if browser_state_to_close: @@ -509,7 +669,14 @@ class RealBrowserManager(BrowserManager): workflow_run_id=workflow_run_id, ) set_deferred_close_params(workflow_run_id, close_browser_on_completion) + # Keep the publisher running while streams are attached. The + # eventual ``close(True)`` fires the on-close callback that + # stops it; ``close(False)`` is covered by the publisher's + # own disconnect-driven self-termination. else: + # Detach the publisher's CDP session before the Playwright context + # closes; otherwise the stale session can race the teardown. + await self._stop_frame_publisher(workflow_run_id=workflow_run_id) await browser_state_to_close.close(close_browser_on_completion=effective_close) if not streams_active: diff --git a/skyvern/webeye/real_browser_state.py b/skyvern/webeye/real_browser_state.py index 41f980c4f..96b145644 100644 --- a/skyvern/webeye/real_browser_state.py +++ b/skyvern/webeye/real_browser_state.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import random import time +from collections.abc import Awaitable, Callable from typing import Literal from urllib.parse import urlparse @@ -70,6 +71,21 @@ class RealBrowserState(BrowserState): self.browser_context = browser_context self.browser_artifacts = browser_artifacts self.browser_cleanup = browser_cleanup + # One-shot callbacks fired first inside ``close()``. Cleared after + # firing so re-entry into ``close()`` is safe. + self._on_close_callbacks: list[Callable[[], Awaitable[None]]] = [] + + def add_on_close(self, callback: Callable[[], Awaitable[None]]) -> None: + self._on_close_callbacks.append(callback) + + async def _run_on_close_callbacks(self) -> None: + callbacks = self._on_close_callbacks + self._on_close_callbacks = [] + for callback in callbacks: + try: + await callback() + except Exception: + LOG.debug("on-close callback raised; ignored", exc_info=True) async def __assert_page(self) -> Page: page = await self.get_working_page() @@ -536,6 +552,12 @@ class RealBrowserState(BrowserState): async def close(self, close_browser_on_completion: bool = True) -> None: LOG.info("Closing browser state", sampling=True) + # Only fire on-close observers on a real teardown. Shared / parent-child + # close calls pass ``close_browser_on_completion=False`` to leave the + # browser alive for another run; firing callbacks then would stop the + # surviving run's publisher and freeze its livestream. + if close_browser_on_completion: + await self._run_on_close_callbacks() try: async with asyncio.timeout(BROWSER_CLOSE_TIMEOUT): if self.browser_context and close_browser_on_completion: diff --git a/tests/unit/test_browser_context_factory_marker.py b/tests/unit/test_browser_context_factory_marker.py new file mode 100644 index 000000000..287d4f570 --- /dev/null +++ b/tests/unit/test_browser_context_factory_marker.py @@ -0,0 +1,70 @@ +"""``_connect_to_cdp_browser`` stamps ``needs_cdp_frame_publisher``. + +It is the single chokepoint for remote-CDP creation here — ``cdp-connect`` +always, plus ``chromium-headless`` / ``chromium-headful`` when +``browser_address`` is set — so one stamp there covers every remote-CDP path. +Ordinary local creators leave the marker False; the factory does not +auto-stamp. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from skyvern.webeye import browser_factory as factory_module + + +@pytest.mark.asyncio +async def test_connect_to_cdp_browser_stamps_marker(monkeypatch: pytest.MonkeyPatch) -> None: + """The chokepoint stamps; every remote-CDP path inherits the marker.""" + fake_context = MagicMock() + fake_browser = MagicMock() + fake_browser.contexts = [fake_context] + + monkeypatch.setattr( + factory_module, + "_connect_over_cdp_with_diagnostics", + AsyncMock(return_value=fake_browser), + ) + + _, browser_artifacts, _ = await factory_module._connect_to_cdp_browser( + playwright=MagicMock(), + remote_browser_url="ws://remote.example/cdp", + ) + + assert browser_artifacts.needs_cdp_frame_publisher is True + + +@pytest.mark.asyncio +async def test_ordinary_local_creator_leaves_marker_false(monkeypatch: pytest.MonkeyPatch) -> None: + """The factory does not auto-stamp; a local creator's marker stays False.""" + from skyvern.webeye.browser_artifacts import BrowserArtifacts + from skyvern.webeye.browser_factory import BrowserContextFactory + + async def _local_creator(playwright: Any, **kwargs: Any) -> tuple[Any, BrowserArtifacts, None]: + return object(), BrowserArtifacts(), None + + monkeypatch.setattr(factory_module, "restore_session_cookies", AsyncMock()) + monkeypatch.setattr(factory_module, "set_browser_console_log", lambda **_: None) + monkeypatch.setattr(factory_module, "set_popup_video_listener", lambda **_: None) + monkeypatch.setattr(factory_module, "set_download_file_listener", lambda **_: None) + monkeypatch.setattr(factory_module, "set_dialog_handler", lambda **_: None) + + class _FakeAgentFunction: + async def setup_browser_context_extensions(self, **_: Any) -> None: + return None + + class _FakeApp: + AGENT_FUNCTION = _FakeAgentFunction() + + monkeypatch.setattr(factory_module, "app", _FakeApp()) + + BrowserContextFactory.register_type("test-local", _local_creator) + monkeypatch.setattr(factory_module.settings, "BROWSER_TYPE", "test-local") + + _, artifacts, _ = await BrowserContextFactory.create_browser_context(playwright=object()) + + assert artifacts.needs_cdp_frame_publisher is False diff --git a/tests/unit/test_real_browser_manager.py b/tests/unit/test_real_browser_manager.py index 886cb31ef..998d9ca68 100644 --- a/tests/unit/test_real_browser_manager.py +++ b/tests/unit/test_real_browser_manager.py @@ -577,12 +577,8 @@ async def test_cleanup_persists_session_cookies_when_close_deferred_for_streams( persist_mock = AsyncMock() monkeypatch.setattr("skyvern.webeye.real_browser_manager.persist_session_cookies", persist_mock) - # streaming is a namespace package; patch the module object so resolution does not depend on the - # parent binding the submodule, which varies with import order across the suite. - from skyvern.forge.sdk.routes.streaming import registries as streaming_registries - - monkeypatch.setattr(streaming_registries, "stream_ref_active", lambda wrid: True) - monkeypatch.setattr(streaming_registries, "set_deferred_close_params", lambda *a, **k: None) + monkeypatch.setattr("skyvern.webeye.real_browser_manager.stream_ref_active", lambda wrid: True) + monkeypatch.setattr("skyvern.webeye.real_browser_manager.set_deferred_close_params", lambda *a, **k: None) await manager.cleanup_for_workflow_run("wfr_streamed", task_ids=[], close_browser_on_completion=True) diff --git a/tests/unit/test_real_browser_manager_frame_publisher.py b/tests/unit/test_real_browser_manager_frame_publisher.py new file mode 100644 index 000000000..76a6ebfe4 --- /dev/null +++ b/tests/unit/test_real_browser_manager_frame_publisher.py @@ -0,0 +1,365 @@ +"""Wire-up tests for ``RealBrowserManager``'s CDP frame publisher helpers. + +Exercises ``_start_frame_publisher`` / ``_stop_frame_publisher`` and their +cleanup interaction with ``self._frame_publishers``. The publisher loop itself +is covered by ``test_worker_cdp_frame_publisher.py``. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from skyvern.webeye import real_browser_manager as manager_module +from skyvern.webeye.real_browser_manager import RealBrowserManager, _resolve_stream_key + + +def test_resolve_stream_key_prefers_workflow_run() -> None: + assert _resolve_stream_key(workflow_run_id="wr_1", task_id="tsk_2") == "wr_1.png" + + +def test_resolve_stream_key_falls_back_to_task() -> None: + assert _resolve_stream_key(workflow_run_id=None, task_id="tsk_2") == "tsk_2.png" + + +def test_resolve_stream_key_returns_none_when_no_identifier() -> None: + assert _resolve_stream_key(workflow_run_id=None, task_id=None) is None + + +class _RecordingPublisher: + """Stand-in for ``CDPFramePublisher`` that tracks lifecycle calls.""" + + def __init__( + self, + *, + browser_state: object, + stream_key: str, + organization_id: str, + ) -> None: + self.browser_state = browser_state + self.stream_key = stream_key + self.organization_id = organization_id + self.start_called = 0 + self.stop_called = 0 + + async def start(self) -> None: + self.start_called += 1 + + async def stop(self) -> None: + self.stop_called += 1 + + +def _marked_browser_state(*, needs_publisher: bool = True) -> SimpleNamespace: + """Fake BrowserState carrying the marker that the factory stamps.""" + return SimpleNamespace( + browser_artifacts=SimpleNamespace(needs_cdp_frame_publisher=needs_publisher), + add_on_close=lambda _cb: None, + ) + + +@pytest.fixture +def patched_publisher(monkeypatch: pytest.MonkeyPatch) -> list[_RecordingPublisher]: + created: list[_RecordingPublisher] = [] + + def _factory(**kwargs: object) -> _RecordingPublisher: + pub = _RecordingPublisher(**kwargs) # type: ignore[arg-type] + created.append(pub) + return pub + + monkeypatch.setattr(manager_module, "CDPFramePublisher", _factory) + return created + + +@pytest.mark.asyncio +async def test_start_frame_publisher_gated_off_when_state_unmarked( + patched_publisher: list[_RecordingPublisher], +) -> None: + """The marker is the gate: an unmarked BrowserState produces no publisher + even when stream_key + organization_id are valid.""" + manager = RealBrowserManager() + await manager._start_frame_publisher( + browser_state=_marked_browser_state(needs_publisher=False), + workflow_run_id="wr_42", + organization_id="o_1", + ) + assert patched_publisher == [] + assert manager._frame_publishers == {} + + +@pytest.mark.asyncio +async def test_start_frame_publisher_creates_publisher_with_workflow_key( + patched_publisher: list[_RecordingPublisher], +) -> None: + manager = RealBrowserManager() + browser_state = _marked_browser_state() + + await manager._start_frame_publisher( + browser_state=browser_state, + workflow_run_id="wr_42", + task_id="tsk_99", + organization_id="o_1", + ) + + assert len(patched_publisher) == 1 + pub = patched_publisher[0] + assert pub.stream_key == "wr_42.png" + assert pub.organization_id == "o_1" + assert pub.start_called == 1 + assert manager._frame_publishers["wr_42.png"] is pub + + +@pytest.mark.asyncio +async def test_start_frame_publisher_uses_task_key_when_no_workflow( + patched_publisher: list[_RecordingPublisher], +) -> None: + manager = RealBrowserManager() + await manager._start_frame_publisher( + browser_state=_marked_browser_state(), + task_id="tsk_99", + organization_id="o_1", + ) + + assert patched_publisher[0].stream_key == "tsk_99.png" + + +@pytest.mark.asyncio +async def test_start_frame_publisher_is_noop_when_already_running( + patched_publisher: list[_RecordingPublisher], +) -> None: + manager = RealBrowserManager() + + for _ in range(3): + await manager._start_frame_publisher( + browser_state=_marked_browser_state(), + workflow_run_id="wr_42", + organization_id="o_1", + ) + + assert len(patched_publisher) == 1 + assert patched_publisher[0].start_called == 1 + + +@pytest.mark.asyncio +async def test_start_frame_publisher_requires_organization_id( + patched_publisher: list[_RecordingPublisher], +) -> None: + manager = RealBrowserManager() + await manager._start_frame_publisher( + browser_state=_marked_browser_state(), + workflow_run_id="wr_42", + organization_id=None, + ) + assert patched_publisher == [] + assert manager._frame_publishers == {} + + +@pytest.mark.asyncio +async def test_start_frame_publisher_swallows_start_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + created: list[_RecordingPublisher] = [] + + class _ExplodingPublisher(_RecordingPublisher): + async def start(self) -> None: # noqa: D401 + raise RuntimeError("kaboom") + + def _factory(**kwargs: object) -> _ExplodingPublisher: + pub = _ExplodingPublisher(**kwargs) # type: ignore[arg-type] + created.append(pub) + return pub + + monkeypatch.setattr(manager_module, "CDPFramePublisher", _factory) + + manager = RealBrowserManager() + # Must not raise: livestream is best-effort. + await manager._start_frame_publisher( + browser_state=_marked_browser_state(), + workflow_run_id="wr_42", + organization_id="o_1", + ) + # The factory must have been invoked — otherwise this test would silently + # exit at the feature gate without exercising the exception-handling path. + assert len(created) == 1 + assert manager._frame_publishers == {} + + +@pytest.mark.asyncio +async def test_stop_frame_publisher_calls_stop_and_pops( + patched_publisher: list[_RecordingPublisher], +) -> None: + manager = RealBrowserManager() + await manager._start_frame_publisher( + browser_state=_marked_browser_state(), + workflow_run_id="wr_42", + organization_id="o_1", + ) + + await manager._stop_frame_publisher(workflow_run_id="wr_42") + + assert patched_publisher[0].stop_called == 1 + assert "wr_42.png" not in manager._frame_publishers + + +@pytest.mark.asyncio +async def test_stop_frame_publisher_is_noop_when_missing( + patched_publisher: list[_RecordingPublisher], +) -> None: + manager = RealBrowserManager() + # Must not raise on stop without start. + await manager._stop_frame_publisher(workflow_run_id="wr_does_not_exist") + assert patched_publisher == [] + + +@pytest.mark.asyncio +async def test_stop_frame_publisher_swallows_stop_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _ExplodingStopPublisher(_RecordingPublisher): + async def stop(self) -> None: + raise RuntimeError("kaboom-stop") + + def _factory(**kwargs: object) -> _ExplodingStopPublisher: + return _ExplodingStopPublisher(**kwargs) # type: ignore[arg-type] + + monkeypatch.setattr(manager_module, "CDPFramePublisher", _factory) + + manager = RealBrowserManager() + await manager._start_frame_publisher( + browser_state=_marked_browser_state(), + workflow_run_id="wr_42", + organization_id="o_1", + ) + + # Must not raise. + await manager._stop_frame_publisher(workflow_run_id="wr_42") + assert manager._frame_publishers == {} + + +@pytest.mark.asyncio +async def test_inherited_child_workflow_starts_publisher_for_child_key( + patched_publisher: list[_RecordingPublisher], +) -> None: + """Child workflow runs that reuse the parent's browser must still publish to the + child's own ``{child_workflow_run_id}.png`` stream key. The workflow-run streaming + endpoint reads that key, so a missing publisher leaves the child livestream blank. + """ + manager = RealBrowserManager() + parent_state = MagicMock() + # Parent was created via a remote-CDP creator; the marker is set to True. + # Bool literal (not a Mock auto-attribute) so the gate accepts it. + parent_state.browser_artifacts.needs_cdp_frame_publisher = True + manager.pages["wfr_parent"] = parent_state + + workflow_run = MagicMock() + workflow_run.workflow_run_id = "wfr_child" + workflow_run.parent_workflow_run_id = "wfr_parent" + workflow_run.organization_id = "o_1" + workflow_run.browser_profile_id = None + workflow_run.proxy_location = None + workflow_run.extra_http_headers = None + workflow_run.cdp_connect_headers = None + workflow_run.browser_address = None + + result = await manager.get_or_create_for_workflow_run( + workflow_run=workflow_run, + url=None, + browser_session_id=None, + ) + + assert result is parent_state + assert len(patched_publisher) == 1 + pub = patched_publisher[0] + assert pub.stream_key == "wfr_child.png" + assert pub.organization_id == "o_1" + assert pub.start_called == 1 + assert manager._frame_publishers["wfr_child.png"] is pub + + +@pytest.mark.asyncio +async def test_close_stops_all_publishers(patched_publisher: list[_RecordingPublisher]) -> None: + manager = RealBrowserManager() + await manager._start_frame_publisher( + browser_state=_marked_browser_state(), + workflow_run_id="wr_1", + organization_id="o_1", + ) + await manager._start_frame_publisher( + browser_state=_marked_browser_state(), + workflow_run_id="wr_2", + organization_id="o_1", + ) + + # Replace browser_state .close() so we don't need real BrowserState plumbing. + closed_states: list[object] = [] + + async def _close_state() -> None: + closed_states.append(object()) + + manager.pages = { + "wr_1": SimpleNamespace(close=AsyncMock(side_effect=_close_state)), + "wr_2": SimpleNamespace(close=AsyncMock(side_effect=_close_state)), + } + + await manager.close() + + assert all(pub.stop_called == 1 for pub in patched_publisher) + assert manager._frame_publishers == {} + assert manager.pages == {} + + +@pytest.mark.asyncio +async def test_cleanup_for_workflow_run_stops_child_publishers( + patched_publisher: list[_RecordingPublisher], +) -> None: + """Parent cleanup must stop publishers for every child workflow run id it + pops. Child workflows skip their own cleanup, so otherwise those publishers + leak until process shutdown. + """ + manager = RealBrowserManager() + + # Parent + two inherited child runs all share the same browser state. + shared_state = SimpleNamespace( + browser_context=None, + browser_artifacts=SimpleNamespace(traces_dir=None, needs_cdp_frame_publisher=True), + close=AsyncMock(), + add_on_close=lambda _cb: None, + ) + manager.pages = { + "wfr_parent": shared_state, + "wfr_child_a": shared_state, + "wfr_child_b": shared_state, + } + + # Start a publisher for each entity, matching the keys the streaming + # endpoint reads. + for wr_id in ("wfr_parent", "wfr_child_a", "wfr_child_b"): + await manager._start_frame_publisher( + browser_state=shared_state, + workflow_run_id=wr_id, + organization_id="o_1", + ) + assert set(manager._frame_publishers.keys()) == { + "wfr_parent.png", + "wfr_child_a.png", + "wfr_child_b.png", + } + + await manager.cleanup_for_workflow_run( + workflow_run_id="wfr_parent", + task_ids=[], + close_browser_on_completion=True, + child_workflow_run_ids=["wfr_child_a", "wfr_child_b"], + ) + + # Every publisher (parent + both children) is stopped and removed from + # the registry. Without the fix the child publishers would still be in + # the dict because their stop was never called. + assert manager._frame_publishers == {} + child_pubs = [pub for pub in patched_publisher if pub.stream_key.startswith("wfr_child_")] + assert len(child_pubs) == 2 + assert all(pub.stop_called == 1 for pub in child_pubs) + # Child entries are popped from the pages map regardless of stream state. + assert "wfr_child_a" not in manager.pages + assert "wfr_child_b" not in manager.pages diff --git a/tests/unit/test_real_browser_manager_publisher_self_stops_on_close.py b/tests/unit/test_real_browser_manager_publisher_self_stops_on_close.py new file mode 100644 index 000000000..dcd8785bf --- /dev/null +++ b/tests/unit/test_real_browser_manager_publisher_self_stops_on_close.py @@ -0,0 +1,107 @@ +"""``_start_frame_publisher`` ties the publisher to ``BrowserState.close()``. + +The publisher follows the browser state regardless of who closes it, so the +API-side ``stream_ref_dec`` does not need to reach into publisher internals. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from types import SimpleNamespace +from typing import Any + +import pytest + +from skyvern.webeye import real_browser_manager as manager_module +from skyvern.webeye.real_browser_manager import RealBrowserManager + + +class _RecordingPublisher: + def __init__(self, *, browser_state: Any, stream_key: str, organization_id: str) -> None: + self.browser_state = browser_state + self.stream_key = stream_key + self.organization_id = organization_id + self.stop_called = 0 + + async def start(self) -> None: + return None + + async def stop(self) -> None: + self.stop_called += 1 + + +class _FakeBrowserState: + """Stand-in for ``RealBrowserState`` with the marker stamped on.""" + + def __init__(self) -> None: + self.browser_artifacts = SimpleNamespace(needs_cdp_frame_publisher=True) + self._on_close_callbacks: list[Callable[[], Awaitable[None]]] = [] + + def add_on_close(self, callback: Callable[[], Awaitable[None]]) -> None: + self._on_close_callbacks.append(callback) + + async def fire_close(self) -> None: + for cb in self._on_close_callbacks: + await cb() + self._on_close_callbacks.clear() + + +@pytest.fixture +def patched_publisher_factory(monkeypatch: pytest.MonkeyPatch) -> list[_RecordingPublisher]: + created: list[_RecordingPublisher] = [] + + def _factory(**kwargs: Any) -> _RecordingPublisher: + pub = _RecordingPublisher(**kwargs) + created.append(pub) + return pub + + monkeypatch.setattr(manager_module, "CDPFramePublisher", _factory) + return created + + +@pytest.mark.asyncio +async def test_publisher_stops_when_browser_state_closes( + patched_publisher_factory: list[_RecordingPublisher], +) -> None: + manager = RealBrowserManager() + browser_state = _FakeBrowserState() + + await manager._start_frame_publisher( + browser_state=browser_state, + workflow_run_id="wr_self_stop", + organization_id="o_1", + ) + assert manager._frame_publishers["wr_self_stop.png"] is patched_publisher_factory[0] + + # Simulate BrowserState.close() invoking its registered callbacks. + await browser_state.fire_close() + + assert patched_publisher_factory[0].stop_called == 1 + # Registry is cleared so a re-start later would not see a stale entry. + assert "wr_self_stop.png" not in manager._frame_publishers + + +@pytest.mark.asyncio +async def test_double_close_callback_safe( + patched_publisher_factory: list[_RecordingPublisher], +) -> None: + """Worker cleanup may stop the publisher explicitly AND BrowserState.close + may fire the registered callback. The second stop must be a no-op.""" + manager = RealBrowserManager() + browser_state = _FakeBrowserState() + + await manager._start_frame_publisher( + browser_state=browser_state, + workflow_run_id="wr_double", + organization_id="o_1", + ) + + # Explicit worker-side stop (e.g. cleanup_for_workflow_run synchronous branch) + await manager._stop_frame_publisher(workflow_run_id="wr_double") + assert patched_publisher_factory[0].stop_called == 1 + assert "wr_double.png" not in manager._frame_publishers + + # Then the registered on-close callback runs (e.g. via browser_state.close() + # later). The pop returns None, so no second stop is dispatched. + await browser_state.fire_close() + assert patched_publisher_factory[0].stop_called == 1 diff --git a/tests/unit/test_real_browser_manager_publisher_start_race.py b/tests/unit/test_real_browser_manager_publisher_start_race.py new file mode 100644 index 000000000..9d58ca837 --- /dev/null +++ b/tests/unit/test_real_browser_manager_publisher_start_race.py @@ -0,0 +1,149 @@ +"""Regression: concurrent ``_start_frame_publisher`` for one key must not +orphan a publisher loop. The lock around check/create/start/store serializes +starts so exactly one publisher is created per stream key. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any + +import pytest + +from skyvern.webeye import real_browser_manager as manager_module +from skyvern.webeye.real_browser_manager import RealBrowserManager + + +def _marked_state() -> SimpleNamespace: + return SimpleNamespace( + browser_artifacts=SimpleNamespace(needs_cdp_frame_publisher=True), + add_on_close=lambda _cb: None, + ) + + +@pytest.mark.asyncio +async def test_start_frame_publisher_serializes_concurrent_starts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances: list[Any] = [] + barrier = asyncio.Event() + first_start_entered = asyncio.Event() + + class _BlockingPublisher: + def __init__( + self, + *, + browser_state: Any, + stream_key: str, + organization_id: str, + ) -> None: + self.stream_key = stream_key + self.organization_id = organization_id + self.start_calls = 0 + self.stop_calls = 0 + instances.append(self) + + async def start(self) -> None: + self.start_calls += 1 + first_start_entered.set() + # Block here so a second concurrent call gets a chance to race the + # registry check before this publisher is stored. + await barrier.wait() + + async def stop(self) -> None: + self.stop_calls += 1 + + monkeypatch.setattr(manager_module, "CDPFramePublisher", _BlockingPublisher) + + manager = RealBrowserManager() + + task_a = asyncio.create_task( + manager._start_frame_publisher( + browser_state=_marked_state(), + workflow_run_id="wr_race", + organization_id="o_1", + ) + ) + # Wait until publisher A is suspended inside start(); without the lock fix + # this is exactly the window where publisher B can slip past the registry + # check. + await first_start_entered.wait() + + task_b = asyncio.create_task( + manager._start_frame_publisher( + browser_state=_marked_state(), + workflow_run_id="wr_race", + organization_id="o_1", + ) + ) + # Give task B a turn so it advances as far as it can (either to the lock + # under the fix, or into the factory under the bug). + for _ in range(5): + await asyncio.sleep(0) + + barrier.set() + await asyncio.gather(task_a, task_b) + + assert len(instances) == 1, ( + f"expected exactly one publisher under concurrent start, got {len(instances)} — " + "the check-then-await-then-store race lets a second publisher slip past the " + "registry check before the first is stored, then orphans the loser." + ) + assert instances[0].start_calls == 1 + assert manager._frame_publishers["wr_race.png"] is instances[0] + # No orphan loop: the single publisher we created can be stopped through + # the registry, with no second instance lingering. + await manager._stop_frame_publisher(workflow_run_id="wr_race") + assert instances[0].stop_calls == 1 + assert manager._frame_publishers == {} + + +@pytest.mark.asyncio +async def test_on_close_callback_pops_under_publisher_lock( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The ``add_on_close`` callback registered in ``_start_frame_publisher`` + must pop ``_frame_publishers`` under ``_publisher_lock`` so a concurrent + re-start cannot read a stale registry between the pop and the publisher's + actual stop.""" + callbacks: list[Any] = [] + + class _FakePublisher: + def __init__(self, *, browser_state: Any, stream_key: str, organization_id: str) -> None: + self.stop_calls = 0 + + async def start(self) -> None: + return None + + async def stop(self) -> None: + self.stop_calls += 1 + + monkeypatch.setattr(manager_module, "CDPFramePublisher", _FakePublisher) + + state = SimpleNamespace( + browser_artifacts=SimpleNamespace(needs_cdp_frame_publisher=True), + add_on_close=lambda cb: callbacks.append(cb), + ) + + manager = RealBrowserManager() + await manager._start_frame_publisher( + browser_state=state, + workflow_run_id="wr_lock", + organization_id="o_1", + ) + assert "wr_lock.png" in manager._frame_publishers + assert len(callbacks) == 1 + + # Hold the lock from outside, then fire the on-close callback. It must wait + # on the lock before popping rather than racing the start path. + async with manager._publisher_lock: + cb_task = asyncio.create_task(callbacks[0]()) + for _ in range(5): + await asyncio.sleep(0) + # Lock still held by us, so the callback should not have popped yet. + assert "wr_lock.png" in manager._frame_publishers + assert not cb_task.done() + + await cb_task + assert "wr_lock.png" not in manager._frame_publishers diff --git a/tests/unit/test_real_browser_state_on_close_hook.py b/tests/unit/test_real_browser_state_on_close_hook.py new file mode 100644 index 000000000..63aed8e96 --- /dev/null +++ b/tests/unit/test_real_browser_state_on_close_hook.py @@ -0,0 +1,79 @@ +"""``RealBrowserState.close`` runs registered on-close callbacks first. + +Lets ``_start_frame_publisher`` register ``publisher.stop`` on the browser +state so any caller of ``close()`` stops the publisher implicitly. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock + +import pytest + +from skyvern.webeye.real_browser_state import RealBrowserState + + +def _bare_state() -> RealBrowserState: + """RealBrowserState with no live browser context; ``close`` should still + invoke registered callbacks.""" + # ``pw`` is set to a no-op stub so ``close`` can short-circuit out of the + # playwright stop branch without raising. + pw_stub: Any = type("_PW", (), {"stop": AsyncMock()})() + return RealBrowserState(pw=pw_stub, browser_context=None) + + +@pytest.mark.asyncio +async def test_close_runs_registered_callbacks_in_order() -> None: + state = _bare_state() + calls: list[str] = [] + + async def cb_one() -> None: + calls.append("one") + + async def cb_two() -> None: + calls.append("two") + + state.add_on_close(cb_one) + state.add_on_close(cb_two) + + await state.close() + + assert calls == ["one", "two"] + + +@pytest.mark.asyncio +async def test_close_callbacks_are_one_shot() -> None: + """Subsequent ``close()`` calls must not re-fire previous callbacks.""" + state = _bare_state() + invocations = 0 + + async def cb() -> None: + nonlocal invocations + invocations += 1 + + state.add_on_close(cb) + await state.close() + await state.close() + assert invocations == 1 + + +@pytest.mark.asyncio +async def test_close_swallows_callback_errors() -> None: + """A misbehaving on-close callback must not block the rest of close.""" + state = _bare_state() + later_ran = False + + async def cb_explodes() -> None: + raise RuntimeError("boom") + + async def cb_later() -> None: + nonlocal later_ran + later_ran = True + + state.add_on_close(cb_explodes) + state.add_on_close(cb_later) + + # Must not raise — close() is the universal teardown path. + await state.close() + assert later_ran is True diff --git a/tests/unit/test_stream_ref_dec_no_publisher_access.py b/tests/unit/test_stream_ref_dec_no_publisher_access.py new file mode 100644 index 000000000..7e3c205f1 --- /dev/null +++ b/tests/unit/test_stream_ref_dec_no_publisher_access.py @@ -0,0 +1,83 @@ +"""Regression: ``stream_ref_dec`` must not touch publisher state on the +browser manager. Publisher lifecycle is driven through ``BrowserState.close()``; +``stream_ref_dec`` only closes/evicts the browser state. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from skyvern.forge.sdk.routes.streaming import registries + + +class _ExplodingPublisherStop: + """Sentinel — any access to ``_stop_frame_publisher`` should fail the test.""" + + def __get__(self, instance: object, owner: type | None = None) -> object: + raise AssertionError( + "stream_ref_dec must not read _stop_frame_publisher from app.BROWSER_MANAGER. " + "Worker-side publisher lifecycle is driven by BrowserState.close(), not the " + "API process." + ) + + +@pytest.mark.asyncio +async def test_stream_ref_dec_does_not_touch_publisher_on_deferred_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_run_id = "wr_no_api_publisher_access" + registries.stream_ref_inc(workflow_run_id) + registries.set_deferred_close_params(workflow_run_id, True) + + close_mock = AsyncMock() + fake_state = SimpleNamespace(close=close_mock) + + class _ManagerWithExplodingStop: + # Reading this attribute MUST fail — the API code path is not allowed + # to ask the worker manager about publishers. + _stop_frame_publisher = _ExplodingPublisherStop() + + def __init__(self) -> None: + self.pages: dict[str, object] = {workflow_run_id: fake_state} + self.evict_page = Mock() + + fake_manager = _ManagerWithExplodingStop() + fake_app = SimpleNamespace(BROWSER_MANAGER=fake_manager) + + import skyvern.forge as forge_module + + monkeypatch.setattr(forge_module, "app", fake_app) + + # Pre-fix the descriptor raises on read; the post-fix code path doesn't + # read it at all, so this must succeed. + await registries.stream_ref_dec(workflow_run_id) + + close_mock.assert_awaited_once_with(close_browser_on_completion=True) + fake_manager.evict_page.assert_called_once_with(workflow_run_id) + + +@pytest.mark.asyncio +async def test_stream_ref_dec_handles_missing_browser_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cross-process reality: API's ``BROWSER_MANAGER.pages`` is empty.""" + workflow_run_id = "wr_cross_process_missing_state" + registries.stream_ref_inc(workflow_run_id) + registries.set_deferred_close_params(workflow_run_id, True) + + evict_mock = Mock() + # No _stop_frame_publisher attribute at all — the API code path must not + # care whether the worker BrowserManager exposes one. + fake_manager = SimpleNamespace(pages={}, evict_page=evict_mock) + fake_app = SimpleNamespace(BROWSER_MANAGER=fake_manager) + + import skyvern.forge as forge_module + + monkeypatch.setattr(forge_module, "app", fake_app) + + # Must not raise even though pages is empty (cross-process case). + await registries.stream_ref_dec(workflow_run_id) + evict_mock.assert_called_once_with(workflow_run_id) diff --git a/tests/unit/test_worker_cdp_frame_publisher.py b/tests/unit/test_worker_cdp_frame_publisher.py new file mode 100644 index 000000000..be7dd0305 --- /dev/null +++ b/tests/unit/test_worker_cdp_frame_publisher.py @@ -0,0 +1,465 @@ +from __future__ import annotations + +import asyncio +import base64 +import os +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from skyvern.webeye import cdp_frame_publisher as publisher_module +from skyvern.webeye.cdp_frame_publisher import ( + CDPFramePublisher, + _write_frame_atomically, + stream_key_for_task, + stream_key_for_workflow_run, +) + +ORG_ID = "o_test_123" +STREAM_KEY = "wr_test_456.png" + + +def _make_cdp_session(frame_bytes_seq: list[bytes]) -> MagicMock: + """CDP session whose ``send`` returns successive base64-encoded PNGs.""" + encoded_seq = [base64.b64encode(b).decode("ascii") for b in frame_bytes_seq] + call_state = {"i": 0, "screenshot_params": []} + + async def _send(method: str, params: dict | None = None) -> dict: + if method != "Page.captureScreenshot": + return {} + call_state["screenshot_params"].append(params or {}) + idx = min(call_state["i"], len(encoded_seq) - 1) + call_state["i"] += 1 + return {"data": encoded_seq[idx]} + + session = MagicMock() + session.send = AsyncMock(side_effect=_send) + session.detach = AsyncMock() + session._screenshot_calls = call_state["screenshot_params"] # for assertions + return session + + +def _make_page_with_session(session: MagicMock) -> MagicMock: + page = MagicMock() + page.context = MagicMock() + page.context.new_cdp_session = AsyncMock(return_value=session) + return page + + +def _make_browser_state(page: MagicMock | None, *, connected: bool = True) -> SimpleNamespace: + return SimpleNamespace( + get_working_page=AsyncMock(return_value=page), + is_connected=lambda: connected, + ) + + +@pytest.fixture +def streaming_temp_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + monkeypatch.setattr(publisher_module, "get_skyvern_temp_dir", lambda: str(tmp_path)) + return tmp_path + + +@pytest.fixture +def fake_storage(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + save_mock = AsyncMock() + fake_app = SimpleNamespace(STORAGE=SimpleNamespace(save_streaming_file=save_mock)) + monkeypatch.setattr(publisher_module, "app", fake_app) + return save_mock + + +async def _drive_publish_once(pub: CDPFramePublisher) -> None: + """Call the internal one-frame routine without starting the long-running loop.""" + await pub._publish_one_frame() + + +def _stream_path(temp_dir: Path) -> Path: + return temp_dir / ORG_ID / STREAM_KEY + + +def test_stream_key_helpers() -> None: + assert stream_key_for_workflow_run("wr_42") == "wr_42.png" + assert stream_key_for_task("tsk_7") == "tsk_7.png" + + +@pytest.mark.asyncio +async def test_publishes_initial_frame_to_streaming_storage(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + frame_bytes = b"\xff\xd8jpegbytes1" + session = _make_cdp_session([frame_bytes]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + await _drive_publish_once(pub) + + target = _stream_path(streaming_temp_dir) + assert target.exists() + assert target.read_bytes() == frame_bytes + fake_storage.assert_awaited_once_with(ORG_ID, STREAM_KEY) + page.context.new_cdp_session.assert_awaited_once_with(page) + + +@pytest.mark.asyncio +async def test_skips_unchanged_frame(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + frame = b"\xff\xd8same" + session = _make_cdp_session([frame, frame, frame]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + await _drive_publish_once(pub) + await _drive_publish_once(pub) + await _drive_publish_once(pub) + + # captureScreenshot is called every tick, but only the first upload happens. + assert session.send.await_count == 3 + fake_storage.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_publishes_again_when_frame_changes(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + session = _make_cdp_session([b"frame-a", b"frame-b"]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + await _drive_publish_once(pub) + await _drive_publish_once(pub) + + assert fake_storage.await_count == 2 + assert _stream_path(streaming_temp_dir).read_bytes() == b"frame-b" + + +@pytest.mark.asyncio +async def test_reattaches_on_page_switch(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + session_a = _make_cdp_session([b"frame-a"]) + page_a = _make_page_with_session(session_a) + session_b = _make_cdp_session([b"frame-b"]) + page_b = _make_page_with_session(session_b) + + current = {"page": page_a} + + async def _get_page() -> MagicMock: + return current["page"] + + state = SimpleNamespace(get_working_page=_get_page) + pub = CDPFramePublisher(browser_state=state, stream_key=STREAM_KEY, organization_id=ORG_ID) + + await _drive_publish_once(pub) + current["page"] = page_b + await _drive_publish_once(pub) + + page_a.context.new_cdp_session.assert_awaited_once_with(page_a) + page_b.context.new_cdp_session.assert_awaited_once_with(page_b) + # Old session must be detached when the page changes. + session_a.detach.assert_awaited() + assert _stream_path(streaming_temp_dir).read_bytes() == b"frame-b" + + +@pytest.mark.asyncio +async def test_no_working_page_is_noop(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + pub = CDPFramePublisher( + browser_state=_make_browser_state(None), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + await _drive_publish_once(pub) + + assert not _stream_path(streaming_temp_dir).exists() + fake_storage.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_capture_screenshot_failure_is_non_fatal(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + session = MagicMock() + session.send = AsyncMock(side_effect=RuntimeError("CDP boom")) + session.detach = AsyncMock() + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + # Must not raise. + await _drive_publish_once(pub) + + assert not _stream_path(streaming_temp_dir).exists() + fake_storage.assert_not_awaited() + # Failed session should be detached so the next tick reattaches. + session.detach.assert_awaited() + + +@pytest.mark.asyncio +async def test_save_streaming_file_failure_is_non_fatal( + streaming_temp_dir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + save_mock = AsyncMock(side_effect=RuntimeError("remote-storage boom")) + fake_app = SimpleNamespace(STORAGE=SimpleNamespace(save_streaming_file=save_mock)) + monkeypatch.setattr(publisher_module, "app", fake_app) + + session = _make_cdp_session([b"frame-x"]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + # Must not raise even though the remote upload raises. + await _drive_publish_once(pub) + + # Local file is still written, so local-disk storage can serve it. + assert _stream_path(streaming_temp_dir).read_bytes() == b"frame-x" + save_mock.assert_awaited_once_with(ORG_ID, STREAM_KEY) + + +@pytest.mark.asyncio +async def test_invalid_base64_is_silently_skipped(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + session = MagicMock() + # Returning a payload without "data" should not crash; b64 decode error + # on garbage should also not crash. + session.send = AsyncMock(return_value={"data": "$$$ not base64 $$$"}) + session.detach = AsyncMock() + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + await _drive_publish_once(pub) + + fake_storage.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_start_is_idempotent(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + session = _make_cdp_session([b"x"]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + capture_interval_seconds=10.0, + ) + + await pub.start() + task_after_first = pub._task + await pub.start() + assert pub._task is task_after_first + await pub.stop() + + +@pytest.mark.asyncio +async def test_stop_cancels_task_and_detaches_session(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + session = _make_cdp_session([b"x"]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + capture_interval_seconds=10.0, + ) + + await pub.start() + # Yield so the loop has a chance to publish at least one frame. + await asyncio.sleep(0) + await asyncio.sleep(0) + await pub.stop() + + assert pub._task is None + assert pub._cdp_session is None + assert pub._attached_page is None + # Stop is idempotent. + await pub.stop() + + +@pytest.mark.asyncio +async def test_run_self_terminates_when_browser_state_disconnects( + streaming_temp_dir: Path, fake_storage: AsyncMock +) -> None: + """``close_browser_on_completion=False`` returns the browser to the pool + without firing on-close. The publisher must still self-terminate when its + underlying ``BrowserState`` reports disconnected so the loop does not + spin forever after the run ends.""" + session = _make_cdp_session([b"png-frame"]) + page = _make_page_with_session(session) + connected = {"flag": True} + browser_state = SimpleNamespace( + get_working_page=AsyncMock(return_value=page), + is_connected=lambda: connected["flag"], + ) + pub = CDPFramePublisher( + browser_state=browser_state, + stream_key=STREAM_KEY, + organization_id=ORG_ID, + capture_interval_seconds=0.1, + ) + + await pub.start() + # Let at least one tick run while connected. + await asyncio.sleep(0.15) + assert pub.is_running + + connected["flag"] = False + # Next iteration should observe disconnect and exit the loop. + await asyncio.sleep(0.25) + assert not pub.is_running + assert pub._stopped.is_set() + + # ``stop()`` is still safe to call after self-termination. + await pub.stop() + + +@pytest.mark.asyncio +async def test_temp_dir_oserror_is_non_fatal( + streaming_temp_dir: Path, fake_storage: AsyncMock, monkeypatch: pytest.MonkeyPatch +) -> None: + # Simulate a temp dir that cannot be created by patching mkdir. + real_mkdir = Path.mkdir + + def _boom(self: Path, *args: object, **kwargs: object) -> None: + if str(self).endswith(ORG_ID): + raise OSError("disk gone") + return real_mkdir(self, *args, **kwargs) # type: ignore[no-any-return] + + monkeypatch.setattr(Path, "mkdir", _boom) + + session = _make_cdp_session([b"frame"]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + await _drive_publish_once(pub) + + fake_storage.assert_not_awaited() + # Org directory was never successfully created, so the stream file shouldn't exist. + assert not (streaming_temp_dir / ORG_ID).exists() or not _stream_path(streaming_temp_dir).exists() + + +@pytest.mark.asyncio +async def test_retries_after_upload_failure_on_same_bytes( + streaming_temp_dir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Upload failure on tick N must not silently dedupe the same frame on tick N+1.""" + call_state = {"n": 0} + + async def _flaky_save(org: str, key: str) -> None: + call_state["n"] += 1 + if call_state["n"] == 1: + raise RuntimeError("transient remote-storage failure") + # Subsequent calls succeed. + + fake_app = SimpleNamespace(STORAGE=SimpleNamespace(save_streaming_file=AsyncMock(side_effect=_flaky_save))) + monkeypatch.setattr(publisher_module, "app", fake_app) + + same_frame = b"jpeg-same-bytes" + session = _make_cdp_session([same_frame, same_frame, same_frame]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + await _drive_publish_once(pub) # save fails; do not mark published + await _drive_publish_once(pub) # retry: save succeeds; mark published + await _drive_publish_once(pub) # identical bytes already marked => skip + + assert call_state["n"] == 2 # exactly the retry attempted + assert _stream_path(streaming_temp_dir).read_bytes() == same_frame + + +@pytest.mark.asyncio +async def test_atomic_write_leaves_no_tmp_files(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + session = _make_cdp_session([b"frame-a", b"frame-b"]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + await _drive_publish_once(pub) + await _drive_publish_once(pub) + + org_dir = streaming_temp_dir / ORG_ID + leftovers = [p for p in os.listdir(org_dir) if p != STREAM_KEY] + assert leftovers == [] + + +@pytest.mark.asyncio +async def test_capture_uses_png_format(streaming_temp_dir: Path, fake_storage: AsyncMock) -> None: + """The frontend WebSocket consumes screenshots as PNG; capture must request PNG. + + Frame bytes here start with the PNG magic (``\\x89PNG\\r\\n\\x1a\\n``) + so the round-trip assertion below also confirms PNG bytes hit disk. + """ + png_bytes = b"\x89PNG\r\n\x1a\nfake-png-payload" + session = _make_cdp_session([png_bytes]) + page = _make_page_with_session(session) + pub = CDPFramePublisher( + browser_state=_make_browser_state(page), + stream_key=STREAM_KEY, + organization_id=ORG_ID, + ) + + await _drive_publish_once(pub) + + screenshot_calls = session._screenshot_calls + assert len(screenshot_calls) == 1 + params = screenshot_calls[0] + assert params.get("format") == "png" + # Quality is JPEG-only; PNG capture must omit it. + assert "quality" not in params + + written = _stream_path(streaming_temp_dir).read_bytes() + assert written.startswith(b"\x89PNG") + + +def test_write_frame_atomically_writes_bytes_via_tempfile(tmp_path: Path) -> None: + """The sync helper used by ``_write_frame`` writes via tempfile+replace and + leaves no ``.tmp`` siblings behind on success.""" + target_dir = tmp_path / "o_org_x" + _write_frame_atomically(target_dir, "wr_x.png", b"payload-bytes") + + assert (target_dir / "wr_x.png").read_bytes() == b"payload-bytes" + # No leftover temp files in the org dir. + leftovers = [p for p in target_dir.iterdir() if p.name.startswith(".") and p.suffix == ".tmp"] + assert leftovers == [] + + +def test_write_frame_atomically_cleans_tempfile_on_replace_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """If ``os.replace`` fails, the tempfile must be unlinked so a failing + publisher cannot accumulate ``.tmp`` debris in the streaming dir.""" + target_dir = tmp_path / "o_org_x" + + def _boom(src: str, dst: str) -> None: # type: ignore[unused-argument] + raise OSError("replace failed") + + monkeypatch.setattr(publisher_module.os, "replace", _boom) + + with pytest.raises(OSError, match="replace failed"): + _write_frame_atomically(target_dir, "wr_x.png", b"payload-bytes") + + # Target was never written, and the temp file was cleaned up. + assert not (target_dir / "wr_x.png").exists() + leftovers = [p for p in target_dir.iterdir() if p.name.startswith(".") and p.suffix == ".tmp"] + assert leftovers == []