Skyvern/skyvern/webeye/real_browser_manager.py

from __future__ import annotations

import asyncio
import os

import structlog
from playwright.async_api import async_playwright

from skyvern.exceptions import FailedToNavigateToUrl, 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
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

LOG = structlog.get_logger()

# Only driver/transport-level CDP drops trigger the cached-PBS evict + reconnect path.
# Playwright also surfaces page/context-only closes ("Target page, context or browser
# has been closed") with text that overlaps a transport drop; treating those as cached
# CDP drops would tear down a healthy PBS over a recoverable page-level state.
_CACHED_CDP_DROP_ERROR_SUBSTRINGS = ("Connection closed while reading from the driver",)


def _is_cached_cdp_drop_error(exc: FailedToNavigateToUrl) -> bool:
    message = exc.error_message or ""
    return any(needle in message for needle in _CACHED_CDP_DROP_ERROR_SUBSTRINGS)


async def _rebind_pbs_download_dir(
    browser_state: BrowserState,
    workflow_run: WorkflowRun,
    browser_session_id: str,
) -> None:
    browser_context = browser_state.browser_context
    adopted_browser = browser_context.browser if browser_context else None
    if adopted_browser is None:
        return
    try:
        rebind_run_id = resolve_run_download_id(skyvern_context.current(), fallback_run_id=workflow_run.workflow_run_id)
        await rebind_download_dir(adopted_browser, run_id=rebind_run_id)
    except Exception:
        LOG.warning(
            "Failed to rebind download dir on adopted browser session",
            browser_session_id=browser_session_id,
            workflow_run_id=workflow_run.workflow_run_id,
            exc_info=True,
        )


def _merge_proxy_session_headers(
    extra_http_headers: dict[str, str] | None,
    proxy_session_id: str | None,
) -> dict[str, str] | None:
    if not proxy_session_id:
        return extra_http_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(
        proxy_location: ProxyLocationInput = None,
        url: str | None = None,
        task_id: str | None = None,
        workflow_run_id: str | None = None,
        workflow_permanent_id: str | None = None,
        script_id: str | None = None,
        organization_id: str | None = None,
        extra_http_headers: dict[str, str] | None = None,
        cdp_connect_headers: dict[str, str] | None = None,
        browser_address: str | None = None,
        browser_profile_id: str | None = None,
    ) -> BrowserState:
        pw = await async_playwright().start()
        (
            browser_context,
            browser_artifacts,
            browser_cleanup,
        ) = await BrowserContextFactory.create_browser_context(
            pw,
            proxy_location=proxy_location,
            url=url,
            task_id=task_id,
            workflow_run_id=workflow_run_id,
            workflow_permanent_id=workflow_permanent_id,
            script_id=script_id,
            organization_id=organization_id,
            extra_http_headers=extra_http_headers,
            cdp_connect_headers=cdp_connect_headers,
            browser_address=browser_address,
            browser_profile_id=browser_profile_id,
        )
        return RealBrowserState(
            pw=pw,
            browser_context=browser_context,
            page=None,
            browser_artifacts=browser_artifacts,
            browser_cleanup=browser_cleanup,
            release_driver_on_close=browser_address is not None,
        )

    def evict_page(self, page_id: str) -> None:
        self.pages.pop(page_id, None)

    def get_for_task(self, task_id: str, workflow_run_id: str | None = None) -> BrowserState | None:
        if task_id in self.pages:
            return self.pages[task_id]

        if workflow_run_id and workflow_run_id in self.pages:
            LOG.info(
                "Browser state for task not found. Using browser state for workflow run",
                sampling=True,
                task_id=task_id,
                workflow_run_id=workflow_run_id,
            )
            self.pages[task_id] = self.pages[workflow_run_id]
            return self.pages[task_id]

        return None

    async def get_or_create_for_task(
        self,
        task: Task,
        browser_session_id: str | None = None,
    ) -> BrowserState:
        browser_state = self.get_for_task(task_id=task.task_id, workflow_run_id=task.workflow_run_id)
        if browser_state is not None:
            return browser_state

        if browser_session_id:
            LOG.info(
                "Getting browser state for task from persistent sessions manager",
                browser_session_id=browser_session_id,
            )
            browser_state = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_state(
                browser_session_id, organization_id=task.organization_id
            )
            if browser_state is None:
                LOG.warning(
                    "Browser state not found in persistent sessions manager",
                    browser_session_id=browser_session_id,
                )
            else:
                if task.organization_id:
                    LOG.info("User to occupy browser session here", browser_session_id=browser_session_id)
                else:
                    LOG.warning("Organization ID is not set for task", task_id=task.task_id)
                page = await browser_state.get_working_page()
                if page:
                    await browser_state.navigate_to_url(page=page, url=task.url)
                else:
                    LOG.warning("Browser state has no page", workflow_run_id=task.workflow_run_id)

        proxy_location = task.proxy_location
        extra_http_headers = task.extra_http_headers
        if browser_state is None:
            LOG.info("Creating browser state for task", task_id=task.task_id)
            if browser_session_id and task.organization_id:
                session = await app.PERSISTENT_SESSIONS_MANAGER.get_session(browser_session_id, task.organization_id)
                if session:
                    if session.proxy_location is not None:
                        proxy_location = session.proxy_location
                    extra_http_headers = _merge_proxy_session_headers(extra_http_headers, session.proxy_session_id)
            browser_state = await self._create_browser_state(
                proxy_location=proxy_location,
                url=task.url,
                task_id=task.task_id,
                workflow_permanent_id=task.workflow_permanent_id,
                organization_id=task.organization_id,
                extra_http_headers=extra_http_headers,
                cdp_connect_headers=task.cdp_connect_headers,
                browser_address=task.browser_address,
            )

            if browser_session_id:
                await app.PERSISTENT_SESSIONS_MANAGER.set_browser_state(
                    browser_session_id,
                    browser_state,
                    organization_id=task.organization_id,
                )

        self.pages[task.task_id] = browser_state
        if task.workflow_run_id:
            self.pages[task.workflow_run_id] = browser_state

        # The URL here is only used when creating a new page, and not when using an existing page.
        # This will make sure browser_state.page is not None.
        await browser_state.get_or_create_page(
            url=task.url,
            proxy_location=proxy_location,
            task_id=task.task_id,
            workflow_permanent_id=task.workflow_permanent_id,
            organization_id=task.organization_id,
            extra_http_headers=extra_http_headers,
            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(
        self,
        workflow_run: WorkflowRun,
        url: str | None = None,
        browser_session_id: str | None = None,
        browser_profile_id: str | None = None,
        navigate: bool = True,
    ) -> BrowserState:
        parent_workflow_run_id = workflow_run.parent_workflow_run_id
        workflow_run_id = workflow_run.workflow_run_id
        if browser_profile_id is None:
            browser_profile_id = workflow_run.browser_profile_id

        # Check own cache entry first so navigate_to_url is only called on the first step.
        # Don't pass parent_workflow_run_id here — that lookup is deferred to the block
        # below so PBS runs don't accidentally inherit the parent's browser.
        browser_state = self.get_for_workflow_run(workflow_run_id=workflow_run_id)
        if browser_state:
            LOG.debug("Returning cached browser state for workflow run", workflow_run_id=workflow_run_id)
            return browser_state

        # When an explicit browser_session_id is provided (e.g. from a workflow
        # trigger block), skip the parent workflow lookup so the child uses the
        # specified persistent session instead of inheriting the parent's browser.
        # Note: at this point workflow_run_id is guaranteed not in self.pages (caught above),
        # so the call below can only match via parent_workflow_run_id.
        if not browser_session_id:
            browser_state = self.get_for_workflow_run(
                workflow_run_id=workflow_run_id, parent_workflow_run_id=parent_workflow_run_id
            )
            if browser_state:
                # always keep the browser state for the workflow run and the parent workflow run synced
                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:
            LOG.info(
                "Getting browser state for workflow run from persistent sessions manager",
                browser_session_id=browser_session_id,
            )
            browser_state = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_state(
                browser_session_id, organization_id=workflow_run.organization_id
            )
            if browser_state is None:
                LOG.warning(
                    "Browser state not found in persistent sessions manager", browser_session_id=browser_session_id
                )
            else:
                LOG.info("Used to occupy browser session here", browser_session_id=browser_session_id)
                await _rebind_pbs_download_dir(browser_state, workflow_run, browser_session_id)
                page = await browser_state.get_working_page()
                if page:
                    if url and navigate:
                        try:
                            await browser_state.navigate_to_url(page=page, url=url)
                        except FailedToNavigateToUrl as nav_exc:
                            if not _is_cached_cdp_drop_error(nav_exc):
                                raise
                            if not app.PERSISTENT_SESSIONS_MANAGER.supports_evict_and_reconnect():
                                # Default OSS impl: ``get_browser_state`` is an in-memory
                                # dict lookup, so an evict would tear down the only cached
                                # BrowserState without any way to reconnect — and would
                                # break profile/video cleanup at ``close_session`` later.
                                # Re-raise the original navigation error untouched.
                                raise
                            LOG.warning(
                                "Cached browser CDP appears dead at first goto — evicting and reconnecting once",
                                browser_session_id=browser_session_id,
                                workflow_run_id=workflow_run.workflow_run_id,
                                error_message=nav_exc.error_message,
                            )
                            await app.PERSISTENT_SESSIONS_MANAGER.evict_cached_browser_state(
                                browser_session_id,
                                organization_id=workflow_run.organization_id,
                                expected=browser_state,
                            )
                            browser_state = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_state(
                                browser_session_id,
                                organization_id=workflow_run.organization_id,
                            )
                            if browser_state is None:
                                raise
                            await _rebind_pbs_download_dir(browser_state, workflow_run, browser_session_id)
                            page = await browser_state.get_working_page()
                            if page is not None:
                                await browser_state.navigate_to_url(page=page, url=url)
                            else:
                                # The fresh CDP connection has no working page (e.g. the
                                # prior context closed its last tab during the dead-CDP
                                # window). The outer ``get_or_create_page`` below mirrors
                                # the normal-path behavior and will produce a page +
                                # navigate to ``url``, so don't fail a recoverable
                                # session here — fall through.
                                LOG.info(
                                    "Recovered PBS reconnect has no working page — deferring to get_or_create_page",
                                    browser_session_id=browser_session_id,
                                    workflow_run_id=workflow_run.workflow_run_id,
                                )
                else:
                    LOG.warning("Browser state has no page", workflow_run_id=workflow_run.workflow_run_id)

        proxy_location = workflow_run.proxy_location
        extra_http_headers = workflow_run.extra_http_headers
        if browser_state is None:
            LOG.info(
                "Creating browser state for workflow run",
                sampling=True,
                workflow_run_id=workflow_run.workflow_run_id,
            )
            if browser_session_id and workflow_run.organization_id:
                session = await app.PERSISTENT_SESSIONS_MANAGER.get_session(
                    browser_session_id, workflow_run.organization_id
                )
                if session:
                    if session.proxy_location is not None:
                        proxy_location = session.proxy_location
                    extra_http_headers = _merge_proxy_session_headers(extra_http_headers, session.proxy_session_id)
            browser_state = await self._create_browser_state(
                proxy_location=proxy_location,
                url=url,
                workflow_run_id=workflow_run.workflow_run_id,
                workflow_permanent_id=workflow_run.workflow_permanent_id,
                organization_id=workflow_run.organization_id,
                extra_http_headers=extra_http_headers,
                cdp_connect_headers=workflow_run.cdp_connect_headers,
                browser_address=workflow_run.browser_address,
                browser_profile_id=browser_profile_id,
            )

            if browser_session_id:
                await app.PERSISTENT_SESSIONS_MANAGER.set_browser_state(
                    browser_session_id,
                    browser_state,
                    organization_id=workflow_run.organization_id,
                )

        self.pages[workflow_run_id] = browser_state
        # Only sync the parent's entry when the child is sharing the parent's
        # browser.  When an explicit browser_session_id is provided the child
        # has its own browser, and overwriting the parent's entry would break
        # subsequent parent blocks.
        if parent_workflow_run_id and not browser_session_id:
            self.pages[parent_workflow_run_id] = browser_state

        # The URL here is only used when creating a new page, and not when using an existing page.
        # This will make sure browser_state.page is not None.
        # When navigate is False, the URL has already been used for proxy selection in
        # _create_browser_state above; we skip navigation so the caller (e.g. a generated
        # script) performs the first goto itself, avoiding a redundant page load.
        await browser_state.get_or_create_page(
            url=url if navigate else None,
            proxy_location=proxy_location,
            workflow_run_id=workflow_run.workflow_run_id,
            workflow_permanent_id=workflow_run.workflow_permanent_id,
            organization_id=workflow_run.organization_id,
            extra_http_headers=extra_http_headers,
            cdp_connect_headers=workflow_run.cdp_connect_headers,
            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(
        self, workflow_run_id: str, parent_workflow_run_id: str | None = None
    ) -> BrowserState | None:
        # Priority: parent first, then own entry.
        # Callers that need to avoid parent inheritance must omit parent_workflow_run_id.
        # See get_or_create_for_workflow_run() for the two-phase lookup pattern.
        if parent_workflow_run_id and parent_workflow_run_id in self.pages:
            return self.pages[parent_workflow_run_id]

        if workflow_run_id in self.pages:
            return self.pages[workflow_run_id]

        return None

    def set_video_artifact_for_task(self, task: Task, artifacts: list[VideoArtifact]) -> None:
        if task.workflow_run_id and task.workflow_run_id in self.pages:
            self.pages[task.workflow_run_id].browser_artifacts.video_artifacts = artifacts
            return
        if task.task_id in self.pages:
            self.pages[task.task_id].browser_artifacts.video_artifacts = artifacts
            return

        raise MissingBrowserState(task_id=task.task_id)

    async def get_video_artifacts(
        self,
        browser_state: BrowserState,
        task_id: str = "",
        workflow_id: str = "",
        workflow_run_id: str = "",
        finalize: bool = True,
    ) -> list[VideoArtifact]:
        if len(browser_state.browser_artifacts.video_artifacts) == 0:
            LOG.warning(
                "Video data not found for task",
                task_id=task_id,
                workflow_id=workflow_id,
                workflow_run_id=workflow_run_id,
            )
            return []

        for i, video_artifact in enumerate(browser_state.browser_artifacts.video_artifacts):
            path = video_artifact.video_path
            if path and os.path.exists(path=path):
                # Only the local Playwright-launched recording path produces WebM
                # that needs the remux fix-up. Other producers (e.g. fully formed
                # MP4 downloaded from a remote source) are already container-valid
                # and would be corrupted by ``finalize_webm`` — read those raw.
                is_webm = path.lower().endswith(".webm")
                if finalize and is_webm:
                    # Remux via ffmpeg so the WebM container has a valid Duration + Cues,
                    # even when browser_context.close() was killed mid-finalization.
                    browser_state.browser_artifacts.video_artifacts[i].video_data = await finalize_webm(path)
                else:
                    # Per-step snapshot while recording is still open — skip ffmpeg: the file is
                    # partial, so remux would either fail or be thrown away by the final pass.
                    with open(path, "rb") as f:
                        browser_state.browser_artifacts.video_artifacts[i].video_data = f.read()
            else:
                LOG.debug(
                    "Video path not found",
                    task_id=task_id,
                    workflow_id=workflow_id,
                    workflow_run_id=workflow_run_id,
                    video_path=path,
                )

        return browser_state.browser_artifacts.video_artifacts

    async def get_har_data(
        self,
        browser_state: BrowserState,
        task_id: str = "",
        workflow_id: str = "",
        workflow_run_id: str = "",
    ) -> bytes:
        if browser_state:
            path = browser_state.browser_artifacts.har_path
            if path and os.path.exists(path=path):
                with open(path, "rb") as f:
                    return f.read()
        LOG.warning(
            "HAR data not found for task",
            task_id=task_id,
            workflow_id=workflow_id,
            workflow_run_id=workflow_run_id,
        )
        return b""

    async def get_browser_console_log(
        self,
        browser_state: BrowserState,
        task_id: str = "",
        workflow_id: str = "",
        workflow_run_id: str = "",
    ) -> bytes:
        if browser_state.browser_artifacts.browser_console_log_path is None:
            LOG.warning(
                "browser console log not found for task",
                task_id=task_id,
                workflow_id=workflow_id,
                workflow_run_id=workflow_run_id,
            )
            return b""

        return await browser_state.browser_artifacts.read_browser_console_log()

    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()
        LOG.info("BrowserManger is closed")

    async def cleanup_for_task(
        self,
        task_id: str,
        close_browser_on_completion: bool = True,
        browser_session_id: str | None = None,
        organization_id: str | None = None,
    ) -> BrowserState | None:
        """
        Developer notes: handle errors here. Do not raise error from this function.
        If error occurs, log it and address the cleanup error.
        """
        LOG.info("Cleaning up for task")
        browser_state_to_close = self.pages.pop(task_id, None)
        if browser_state_to_close:
            # Stop tracing before closing the browser if tracing is enabled
            if browser_state_to_close.browser_context and browser_state_to_close.browser_artifacts.traces_dir:
                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)
            # A state backing a persistent session stays cached in the sessions
            # manager for reuse; its driver is released when the session closes.
            await browser_state_to_close.close(
                close_browser_on_completion=close_browser_on_completion,
                release_driver=False if browser_session_id else None,
            )
        LOG.info("Task is cleaned up")

        if browser_session_id:
            if organization_id:
                await app.PERSISTENT_SESSIONS_MANAGER.release_browser_session(
                    browser_session_id, organization_id=organization_id
                )
                LOG.info("Released browser session", browser_session_id=browser_session_id)
            else:
                LOG.warning("Organization ID not specified, cannot release browser session", task_id=task_id)

        return browser_state_to_close

    async def cleanup_for_workflow_run(
        self,
        workflow_run_id: str,
        task_ids: list[str],
        close_browser_on_completion: bool = True,
        browser_session_id: str | None = None,
        organization_id: str | None = None,
        child_workflow_run_ids: list[str] | None = None,
    ) -> BrowserState | None:
        LOG.info("Cleaning up for workflow run", sampling=True)
        browser_state_to_close = self.pages.get(workflow_run_id)

        # Pop child workflow_run entries first — these are orphaned because child
        # workflows skip clean_up_workflow. Must happen before the shared check
        # so the task loop can correctly detect when the browser is no longer shared.
        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)

        # 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:
            # If another workflow run still references this browser state (e.g. a
            # parent whose in-memory browser was shared via use_parent_browser_session),
            # skip closing the browser so the parent can continue using it.
            shared = any(bs is browser_state_to_close for bs in self.pages.values())
            effective_close = close_browser_on_completion and not shared
            if shared:
                LOG.info(
                    "Browser state is shared with another workflow run, skipping browser close",
                    sampling=True,
                    workflow_run_id=workflow_run_id,
                )

            # Stop tracing before closing the browser if tracing is enabled.
            # Skip when the browser is shared — Playwright supports only one active
            # tracing session per context, so stopping here would kill the parent's trace.
            if (
                browser_state_to_close.browser_context
                and browser_state_to_close.browser_artifacts.traces_dir
                and not shared
            ):
                trace_path = f"{browser_state_to_close.browser_artifacts.traces_dir}/{workflow_run_id}.zip"
                await browser_state_to_close.browser_context.tracing.stop(path=trace_path)
                LOG.info("Stopped tracing", trace_path=trace_path)

            if streams_active:
                # Defer close until the last stream disconnects. Persist session cookies first: the
                # deferred close() runs after store_browser_session archives the dir, too late for it.
                await persist_session_cookies(
                    browser_state_to_close.browser_context,
                    browser_state_to_close.browser_artifacts.browser_session_dir,
                )
                LOG.info(
                    "Deferring browser close — active CDP streams",
                    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,
                    release_driver=False if (shared or browser_session_id) else None,
                )

        if not streams_active:
            self.pages.pop(workflow_run_id, None)
        for task_id in task_ids:
            task_browser_state = self.pages.pop(task_id, None)
            if task_browser_state is None or streams_active:
                continue
            # Same shared-state check for task-level entries
            shared = any(bs is task_browser_state for bs in self.pages.values())
            effective_close = close_browser_on_completion and not shared
            if shared:
                LOG.info(
                    "Browser state is shared with another workflow run, skipping browser close",
                    sampling=True,
                    task_id=task_id,
                    workflow_run_id=workflow_run_id,
                )
            try:
                await task_browser_state.close(
                    close_browser_on_completion=effective_close,
                    release_driver=False if (shared or browser_session_id) else None,
                )
            except Exception:
                LOG.info(
                    "Failed to close the browser state from the task block, might because it's already closed.",
                    exc_info=True,
                    task_id=task_id,
                    workflow_run_id=workflow_run_id,
                )
        LOG.info("Workflow run is cleaned up", sampling=True)

        if browser_session_id:
            if organization_id:
                await app.PERSISTENT_SESSIONS_MANAGER.release_browser_session(
                    browser_session_id, organization_id=organization_id
                )
                LOG.info("Released browser session", browser_session_id=browser_session_id)
            else:
                LOG.warning(
                    "Organization ID not specified, cannot release browser session", workflow_run_id=workflow_run_id
                )

        return browser_state_to_close

    async def get_or_create_for_script(
        self,
        script_id: str | None = None,
        browser_session_id: str | None = None,
    ) -> BrowserState:
        browser_state = self.get_for_script(script_id=script_id)
        if browser_state:
            return browser_state

        if browser_session_id:
            LOG.info(
                "Getting browser state for script",
                browser_session_id=browser_session_id,
            )
            browser_state = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_state(
                browser_session_id, organization_id=script_id
            )
            if browser_state is None:
                LOG.warning(
                    "Browser state not found in persistent sessions manager",
                    browser_session_id=browser_session_id,
                )
            else:
                page = await browser_state.get_working_page()
                if not page:
                    LOG.warning("Browser state has no page to run the script", script_id=script_id)
        proxy_location = ProxyLocation.RESIDENTIAL
        if not browser_state:
            browser_state = await self._create_browser_state(
                proxy_location=proxy_location,
                script_id=script_id,
            )

        if script_id:
            self.pages[script_id] = browser_state
        await browser_state.get_or_create_page(
            proxy_location=proxy_location,
            script_id=script_id,
        )

        return browser_state

    def get_for_script(self, script_id: str | None = None) -> BrowserState | None:
        if script_id and script_id in self.pages:
            return self.pages[script_id]
        return None