"""Playwright Chromium session manager with CDP screencast.""" from __future__ import annotations import asyncio import logging import time from dataclasses import dataclass, field from typing import Any, Callable, Awaitable from playwright.async_api import Browser, BrowserContext, Page, async_playwright from browser_stealth import ( CHROMIUM_ARGS, CONTEXT_OPTIONS, IGNORE_DEFAULT_ARGS, STEALTH_INIT_SCRIPT, ) from browser_url_safe import validate_navigation_url from config import settings def _headless() -> bool: return settings.WORKER_HEADLESS.lower() in ("1", "true", "yes") logger = logging.getLogger("browser_worker.session") SendFn = Callable[[dict[str, Any]], Awaitable[None]] @dataclass class BrowserSession: session_id: str viewport_w: int viewport_h: int send: SendFn ready: bool = False context: BrowserContext | None = None page: Page | None = None cdp: Any = None last_activity: float = field(default_factory=time.monotonic) reserved_at: float = field(default_factory=time.monotonic) _screencast_task: asyncio.Task | None = None _title_task: asyncio.Task | None = None async def touch(self) -> None: self.last_activity = time.monotonic() async def close(self) -> None: for task in (self._screencast_task, self._title_task): if task and not task.done(): task.cancel() try: await task except asyncio.CancelledError: pass self._screencast_task = None self._title_task = None if self.cdp: try: await self.cdp.send("Page.stopScreencast") except Exception: pass self.cdp = None if self.context: try: await self.context.close() except Exception as exc: logger.warning("context close failed session=%s: %s", self.session_id, exc) self.context = None self.page = None self.ready = False class SessionManager: def __init__(self) -> None: self._browser: Browser | None = None self._playwright = None self._sessions: dict[str, BrowserSession] = {} self._lock = asyncio.Lock() self._idle_task: asyncio.Task | None = None async def start(self) -> None: self._playwright = await async_playwright().start() self._browser = await self._playwright.chromium.launch( headless=_headless(), args=CHROMIUM_ARGS, ignore_default_args=IGNORE_DEFAULT_ARGS, ) self._idle_task = asyncio.create_task(self._idle_loop(), name="browser_idle") async def stop(self) -> None: if self._idle_task: self._idle_task.cancel() try: await self._idle_task except asyncio.CancelledError: pass for sid in list(self._sessions): await self.destroy_session(sid) if self._browser: await self._browser.close() self._browser = None if self._playwright: await self._playwright.stop() self._playwright = None async def _idle_loop(self) -> None: while True: await asyncio.sleep(60) now = time.monotonic() stale: list[str] = [] for sid, sess in self._sessions.items(): if sess.ready: if now - sess.last_activity > settings.WORKER_IDLE_TIMEOUT_SEC: stale.append(sid) elif now - sess.reserved_at > settings.WORKER_RESERVE_TIMEOUT_SEC: stale.append(sid) for sid in stale: logger.info("idle timeout session=%s", sid) await self.destroy_session(sid) def session_count(self) -> int: return len(self._sessions) async def reserve_session( self, session_id: str, viewport_w: int, viewport_h: int, ) -> None: async with self._lock: if session_id in self._sessions: raise ValueError("session_exists") if len(self._sessions) >= settings.WORKER_MAX_SESSIONS: raise ValueError("max_sessions") self._sessions[session_id] = BrowserSession( session_id=session_id, viewport_w=viewport_w, viewport_h=viewport_h, send=self._noop_send, ) async def attach_stream(self, session_id: str, send: SendFn) -> BrowserSession: async with self._lock: session = self._sessions.get(session_id) if not session: raise ValueError("session_not_found") if session.ready: session.send = send return session if not self._browser: raise RuntimeError("browser_not_ready") session.send = send context = await self._browser.new_context( viewport={"width": session.viewport_w, "height": session.viewport_h}, ignore_https_errors=False, **CONTEXT_OPTIONS, ) await context.add_init_script(STEALTH_INIT_SCRIPT) page = await context.new_page() session.context = context session.page = page cdp = await context.new_cdp_session(page) session.cdp = cdp await cdp.send("Page.startScreencast", { "format": "jpeg", "quality": max(60, min(settings.WORKER_SCREENCAST_QUALITY, 100)), "everyNthFrame": 1, }) session._screencast_task = asyncio.create_task( self._screencast_loop(session, cdp), name=f"screencast_{session_id}", ) session._title_task = asyncio.create_task( self._title_loop(session), name=f"title_{session_id}", ) session.ready = True await session.send({ "type": "BROWSER_INIT", "session_id": session_id, "viewport_w": session.viewport_w, "viewport_h": session.viewport_h, }) return session async def _noop_send(self, _msg: dict[str, Any]) -> None: return async def destroy_session(self, session_id: str) -> bool: async with self._lock: session = self._sessions.pop(session_id, None) if not session: return False await session.close() return True def get_session(self, session_id: str) -> BrowserSession | None: return self._sessions.get(session_id) async def _screencast_loop(self, session: BrowserSession, cdp: Any) -> None: """CDPSession has no wait_for_event — use on() + asyncio.Queue.""" frame_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() def _on_screencast_frame(params: dict[str, Any]) -> None: try: frame_queue.put_nowait(params) except asyncio.QueueFull: pass cdp.on("Page.screencastFrame", _on_screencast_frame) try: while session.session_id in self._sessions and session.ready: try: msg = await asyncio.wait_for(frame_queue.get(), timeout=30.0) except asyncio.TimeoutError: continue await session.touch() await session.send({ "type": "SCREENCAST_FRAME", "data_b64": msg.get("data", ""), "metadata": msg.get("metadata") or {}, }) await cdp.send("Page.screencastFrameAck", {"sessionId": msg["sessionId"]}) except asyncio.CancelledError: raise except Exception as exc: logger.warning("screencast ended session=%s: %s", session.session_id, exc) async def _title_loop(self, session: BrowserSession) -> None: page = session.page if not page: return last_url = "" last_title = "" try: while session.session_id in self._sessions and session.ready: await asyncio.sleep(1) await session.touch() try: url = page.url title = await page.title() except Exception: continue if url != last_url or title != last_title: last_url = url last_title = title await session.send({ "type": "PAGE_INFO", "url": url, "title": title, }) except asyncio.CancelledError: raise async def handle_message(self, session_id: str, msg: dict[str, Any]) -> None: session = self.get_session(session_id) if not session or not session.page or not session.ready: return await session.touch() page = session.page mtype = msg.get("type") if mtype == "PING": await session.send({"type": "PONG"}) return if mtype == "NAVIGATE": url = msg.get("url", "") normalized, err = validate_navigation_url(str(url)) if err: await session.send({"type": "ERROR", "message": err, "code": "URL_BLOCKED"}) return try: await page.goto(normalized, wait_until="domcontentloaded", timeout=30_000) await session.send({ "type": "PAGE_INFO", "url": page.url, "title": await page.title(), }) except Exception as exc: await session.send({ "type": "ERROR", "message": str(exc)[:500], "code": "NAVIGATE_FAILED", }) return if mtype == "RELOAD": await page.reload(wait_until="domcontentloaded", timeout=30_000) return if mtype == "GO_BACK": await page.go_back(wait_until="domcontentloaded", timeout=30_000) return if mtype == "GO_FORWARD": await page.go_forward(wait_until="domcontentloaded", timeout=30_000) return if mtype == "VIEWPORT": w = int(msg.get("width") or session.viewport_w) h = int(msg.get("height") or session.viewport_h) w = max(320, min(w, 1920)) h = max(240, min(h, 1080)) session.viewport_w = w session.viewport_h = h await page.set_viewport_size({"width": w, "height": h}) return if mtype == "MOUSE": event = msg.get("event") x = float(msg.get("x", 0)) y = float(msg.get("y", 0)) button = msg.get("button", "left") steps = int(msg.get("steps") or 1) steps = max(1, min(steps, 25)) if event == "move": await page.mouse.move(x, y, steps=steps) elif event == "down": await page.mouse.move(x, y) await page.mouse.down(button=button) elif event == "up": await page.mouse.move(x, y) await page.mouse.up(button=button) elif event == "click": await page.mouse.click(x, y, button=button, delay=50) elif event == "wheel": delta_y = float(msg.get("delta_y", 0)) await page.mouse.wheel(0, delta_y) return if mtype == "KEY": event = msg.get("event") key = msg.get("key", "") text = msg.get("text", "") if event == "down" and key: await page.keyboard.down(key) elif event == "up" and key: await page.keyboard.up(key) elif event == "type" and text: await page.keyboard.type(text) return manager = SessionManager()