Nexus/browser-worker/session_manager.py

"""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()