Nexus/browser-worker/main.py

"""Browser Worker — Playwright Chromium + WebSocket stream."""

from __future__ import annotations

import asyncio
import json
import logging
from contextlib import asynccontextmanager
from typing import Any

from fastapi import FastAPI, Header, HTTPException, WebSocket, WebSocketDisconnect
from pydantic import BaseModel, Field

from config import settings
from session_manager import manager

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("browser_worker")


def _check_key(header: str | None) -> None:
    if not settings.WORKER_SECRET:
        raise HTTPException(status_code=503, detail="Worker secret not configured")
    if not header or header != settings.WORKER_SECRET:
        raise HTTPException(status_code=401, detail="Invalid worker key")


class CreateSessionBody(BaseModel):
    session_id: str = Field(min_length=8, max_length=64)
    viewport_w: int = Field(default=1280, ge=320, le=1920)
    viewport_h: int = Field(default=720, ge=240, le=1080)


@asynccontextmanager
async def lifespan(_app: FastAPI):
    await manager.start()
    logger.info("Playwright Chromium ready")
    yield
    await manager.stop()


app = FastAPI(title="Nexus Browser Worker", lifespan=lifespan)


@app.get("/health")
async def health():
    return {
        "status": "ok",
        "sessions": manager.session_count(),
        "max_sessions": settings.WORKER_MAX_SESSIONS,
    }


@app.post("/v1/sessions")
async def create_session(
    body: CreateSessionBody,
    x_nexus_worker_key: str | None = Header(default=None, alias="X-Nexus-Worker-Key"),
):
    _check_key(x_nexus_worker_key)
    try:
        await manager.reserve_session(body.session_id, body.viewport_w, body.viewport_h)
    except ValueError as exc:
        code = str(exc)
        if code == "max_sessions":
            raise HTTPException(status_code=503, detail="Max sessions reached") from exc
        if code == "session_exists":
            raise HTTPException(status_code=409, detail="Session already exists") from exc
        raise HTTPException(status_code=400, detail=code) from exc
    return {"session_id": body.session_id, "viewport_w": body.viewport_w, "viewport_h": body.viewport_h}


@app.delete("/v1/sessions/{session_id}")
async def delete_session(
    session_id: str,
    x_nexus_worker_key: str | None = Header(default=None, alias="X-Nexus-Worker-Key"),
):
    _check_key(x_nexus_worker_key)
    if not await manager.destroy_session(session_id):
        raise HTTPException(status_code=404, detail="Session not found")
    return {"ok": True}


@app.websocket("/v1/sessions/{session_id}/stream")
async def session_stream(
    websocket: WebSocket,
    session_id: str,
    x_nexus_worker_key: str | None = Header(default=None, alias="X-Nexus-Worker-Key"),
):
    key = x_nexus_worker_key or websocket.headers.get("x-nexus-worker-key")
    if not settings.WORKER_SECRET or key != settings.WORKER_SECRET:
        await websocket.close(code=4401, reason="Invalid worker key")
        return

    if not manager.get_session(session_id):
        await websocket.close(code=4404, reason="Session not found")
        return

    await websocket.accept()

    outbound: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

    async def send(msg: dict[str, Any]) -> None:
        await outbound.put(msg)

    try:
        await manager.attach_stream(session_id, send)
    except ValueError:
        await websocket.close(code=4404, reason="Session not found")
        return
    except RuntimeError:
        await manager.destroy_session(session_id)
        await websocket.close(code=4503, reason="Browser not ready")
        return

    async def pump_out() -> None:
        while True:
            msg = await outbound.get()
            await websocket.send_json(msg)

    pump_task = asyncio.create_task(pump_out(), name=f"ws_out_{session_id}")
    try:
        while True:
            raw = await websocket.receive_text()
            try:
                msg = json.loads(raw)
            except json.JSONDecodeError:
                await send({"type": "ERROR", "message": "Invalid JSON", "code": "BAD_JSON"})
                continue
            if not isinstance(msg, dict):
                continue
            await manager.handle_message(session_id, msg)
    except WebSocketDisconnect:
        pass
    finally:
        pump_task.cancel()
        try:
            await pump_task
        except asyncio.CancelledError:
            pass
        await manager.destroy_session(session_id)


if __name__ == "__main__":
    import uvicorn

    host, _, port = settings.WORKER_BIND.partition(":")
    uvicorn.run("main:app", host=host or "0.0.0.0", port=int(port or 8443), log_level="info")