Nexus/server/api/browser_session.py

"""Remote browser session API — bridge to Browser Worker."""

from __future__ import annotations

import asyncio
import json
import logging
import uuid
from typing import Any, Optional

import websockets
from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect, Query
from pydantic import BaseModel, Field

from server.api.auth_jwt import get_current_admin
from server.api.websocket import _verify_ws_token
from server.config import settings
from server.domain.models import Admin
from server.infrastructure.browser import session_registry
from server.infrastructure.browser.worker_client import (
    BrowserWorkerError,
    browser_enabled,
    create_worker_session,
    delete_worker_session,
    worker_health,
    worker_stream_url,
)
from server.utils.browser_url_safe import validate_navigation_url

logger = logging.getLogger("nexus.browser.session")

router = APIRouter(prefix="/api/browser", tags=["browser"])
ws_router = APIRouter(tags=["browser"])


class CreateBrowserSessionBody(BaseModel):
    viewport_w: int = Field(default=1280, ge=320, le=1920)
    viewport_h: int = Field(default=720, ge=240, le=1080)


@router.get("/status")
async def browser_status(admin: Admin = Depends(get_current_admin)):
    del admin
    enabled = browser_enabled()
    health = await worker_health() if enabled else None
    return {
        "enabled": enabled,
        "worker_ok": health is not None,
        "worker_sessions": health.get("sessions") if health else None,
        "max_sessions": settings.BROWSER_MAX_SESSIONS,
    }


@router.post("/sessions")
async def create_browser_session(
    body: CreateBrowserSessionBody,
    admin: Admin = Depends(get_current_admin),
):
    if not browser_enabled():
        raise HTTPException(status_code=503, detail="服务端浏览器未启用")
    # 清理断线残留（避免 409 导致前端无法重连）
    for stale_id in session_registry.session_ids_for_admin(admin.id):
        session_registry.remove(stale_id)
        await delete_worker_session(stale_id)
    if session_registry.total_count() >= settings.BROWSER_MAX_SESSIONS:
        raise HTTPException(status_code=503, detail="浏览器会话已达上限")

    session_id = str(uuid.uuid4())
    try:
        await create_worker_session(session_id, body.viewport_w, body.viewport_h)
    except BrowserWorkerError as exc:
        raise HTTPException(status_code=exc.status_code, detail=str(exc)) from exc

    session_registry.register(session_id, admin.id, body.viewport_w, body.viewport_h)
    return {
        "session_id": session_id,
        "viewport_w": body.viewport_w,
        "viewport_h": body.viewport_h,
    }


@router.delete("/sessions/{session_id}")
async def delete_browser_session(
    session_id: str,
    admin: Admin = Depends(get_current_admin),
):
    sess = session_registry.get(session_id)
    if not sess or sess.admin_id != admin.id:
        raise HTTPException(status_code=404, detail="Session not found")
    session_registry.remove(session_id)
    await delete_worker_session(session_id)
    return {"ok": True}


def _validate_client_navigate(msg: dict[str, Any]) -> dict[str, Any] | None:
    if msg.get("type") != "NAVIGATE":
        return msg
    url = msg.get("url", "")
    normalized, err = validate_navigation_url(str(url))
    if err:
        return {"type": "ERROR", "message": err, "code": "URL_BLOCKED"}
    out = dict(msg)
    out["url"] = normalized
    return out


@ws_router.websocket("/ws/browser/{session_id}")
async def browser_ws(
    websocket: WebSocket,
    session_id: str,
    token: Optional[str] = Query(None),
):
    if not browser_enabled():
        await websocket.close(code=4503, reason="Browser not enabled")
        return
    if not token:
        await websocket.close(code=4001, reason="Missing JWT token")
        return

    admin = await _verify_ws_token(token)
    if not admin:
        await websocket.close(code=4401, reason="Invalid or expired JWT token")
        return

    sess = session_registry.get(session_id)
    if not sess or sess.admin_id != admin.id:
        await websocket.close(code=4403, reason="Session not authorized")
        return

    await websocket.accept()
    worker_url = worker_stream_url(session_id)
    headers = {"X-Nexus-Worker-Key": settings.BROWSER_WORKER_SECRET}

    try:
        async with websockets.connect(
            worker_url,
            additional_headers=headers,
            open_timeout=15,
            max_size=8 * 1024 * 1024,
        ) as worker_ws:

            async def client_to_worker() -> None:
                try:
                    while True:
                        raw = await websocket.receive_text()
                        session_registry.touch(session_id)
                        try:
                            msg = json.loads(raw)
                        except json.JSONDecodeError:
                            await websocket.send_json({
                                "type": "ERROR",
                                "message": "Invalid JSON",
                                "code": "BAD_JSON",
                            })
                            continue
                        if not isinstance(msg, dict):
                            continue
                        validated = _validate_client_navigate(msg)
                        if validated:
                            await worker_ws.send(json.dumps(validated))
                except WebSocketDisconnect:
                    pass

            async def worker_to_client() -> None:
                try:
                    async for raw in worker_ws:
                        session_registry.touch(session_id)
                        if isinstance(raw, bytes):
                            raw = raw.decode("utf-8", errors="replace")
                        await websocket.send_text(raw)
                except websockets.ConnectionClosed:
                    pass

            done, pending = await asyncio.wait(
                [asyncio.create_task(client_to_worker()), asyncio.create_task(worker_to_client())],
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
            for task in done:
                try:
                    await task
                except Exception:
                    pass
    except Exception as exc:
        logger.warning("browser bridge error session=%s: %s", session_id, exc)
        try:
            await websocket.send_json({
                "type": "ERROR",
                "message": "无法连接 Browser Worker",
                "code": "WORKER_UNREACHABLE",
            })
        except Exception:
            pass
    finally:
        session_registry.remove(session_id)
        await delete_worker_session(session_id)