SKY-10723: Add worker-side CDP livestream publishing (#6895)

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:

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(

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,

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"

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:

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:

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

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)
 

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

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

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

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

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)

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 == []