creator-hub/browser_gateway/browser/cdp.py

"""Shared restricted CDP transport used by platform browsers."""

from __future__ import annotations

import json
import logging
import threading
import time
from collections import deque
from collections.abc import Callable
from contextlib import suppress

import websocket

LOG = logging.getLogger("creatorhub.browser.cdp")

CONTROL_TIMEOUT = 15.0
RESOLVE_TIMEOUT = 90.0


class BrowserError(RuntimeError):
    def __init__(self, message: str):
        super().__init__(message)
        self.uncertain = False


class CDPConnection:
    def __init__(self, socket: websocket.WebSocket) -> None:
        self.socket = socket
        self._lock = threading.RLock()
        self._next_id = 0
        self._pending: deque[dict] = deque()

    def notify(self, method: str, params: dict | None = None) -> None:
        with self._lock:
            self._next_id += 1
            self.socket.send(
                json.dumps({"id": self._next_id, "method": method, "params": params or {}})
            )

    def command(self, method: str, params: dict | None = None) -> dict:
        with self._lock:
            self._next_id += 1
            command_id = self._next_id
            self.socket.send(
                json.dumps({"id": command_id, "method": method, "params": params or {}})
            )
            deadline = time.monotonic() + CONTROL_TIMEOUT
            while True:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    self._terminate_evaluation(method)
                    raise BrowserError(f"CDP command timed out: {method}")
                message = self._take_pending(command_id)
                if message is None:
                    remaining = deadline - time.monotonic()
                    if remaining <= 0:
                        self._terminate_evaluation(method)
                        raise BrowserError(f"CDP command timed out: {method}")
                    try:
                        message = self._receive(remaining, method)
                    except BrowserError:
                        self._terminate_evaluation(method)
                        raise
                if not isinstance(message, dict):
                    raise BrowserError(f"CDP command returned invalid message: {method}")
                if message.get("id") != command_id:
                    self._pending.append(message)
                    continue
                if message.get("error"):
                    raise BrowserError(f"CDP command rejected: {method}")
                result = message.get("result")
                if not isinstance(result, dict):
                    raise BrowserError(f"CDP command returned invalid result: {method}")
                return result

    def wait_event(
        self,
        method: str,
        predicate: Callable[[dict], bool],
        timeout: float = CONTROL_TIMEOUT,
    ) -> dict:
        with self._lock:
            deadline = time.monotonic() + timeout
            while True:
                for index, message in enumerate(self._pending):
                    params = message.get("params", {})
                    if (
                        message.get("method") == method
                        and isinstance(params, dict)
                        and predicate(params)
                    ):
                        del self._pending[index]
                        return message
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise BrowserError(f"CDP event timed out: {method}")
                message = self._receive(remaining, method)
                if isinstance(message, dict):
                    params = message.get("params", {})
                    if (
                        message.get("method") == method
                        and isinstance(params, dict)
                        and predicate(params)
                    ):
                        return message
                    self._pending.append(message)

    def _terminate_evaluation(self, method: str) -> None:
        if method != "Runtime.evaluate":
            return
        try:
            self._next_id += 1
            command_id = self._next_id
            self.socket.settimeout(1.0)
            self.socket.send(
                json.dumps(
                    {
                        "id": command_id,
                        "method": "Runtime.terminateExecution",
                        "params": {},
                    }
                )
            )
            deadline = time.monotonic() + 1.0
            while time.monotonic() < deadline:
                message = json.loads(self.socket.recv())
                if isinstance(message, dict) and message.get("id") == command_id:
                    return
                if isinstance(message, dict):
                    self._pending.append(message)
        except (
            OSError,
            TypeError,
            websocket.WebSocketException,
            json.JSONDecodeError,
        ) as exc:
            LOG.debug("failed to terminate timed-out CDP evaluation", exc_info=exc)

    def _take_pending(self, command_id: int) -> dict | None:
        for index, message in enumerate(self._pending):
            if message.get("id") == command_id:
                del self._pending[index]
                return message
        return None

    def _receive(self, timeout: float, operation: str) -> object:
        self.socket.settimeout(timeout)
        try:
            return json.loads(self.socket.recv())
        except (TimeoutError, websocket.WebSocketTimeoutException) as exc:
            raise BrowserError(f"CDP command timed out during {operation}") from exc
        except (
            OSError,
            TypeError,
            websocket.WebSocketException,
            json.JSONDecodeError,
        ) as exc:
            raise BrowserError(f"CDP message failed during {operation}") from exc

    def evaluate(self, expression: str) -> object:
        result = self.command(
            "Runtime.evaluate",
            {
                "expression": expression,
                "awaitPromise": True,
                "returnByValue": True,
                "userGesture": False,
            },
        )
        if result.get("exceptionDetails"):
            raise BrowserError("page evaluation failed")
        value = result.get("result", {}).get("value")
        if "value" not in result.get("result", {}):
            raise BrowserError("page evaluation returned no value")
        return value

    def close(self) -> None:
        with suppress(OSError, websocket.WebSocketException):
            self.socket.close()