api/main.py

"""FastAPI entry point. Lifespan owns the browser session."""
from __future__ import annotations

import asyncio
import logging
import sys
from contextlib import asynccontextmanager
from pathlib import Path

import structlog
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from api.routes_auth import router as auth_router  # noqa: E402
from api.routes_chat import router as chat_router  # noqa: E402
from api.routes_control import router as control_router  # noqa: E402
from api.routes_images import router as images_router  # noqa: E402
from api.routes_models import router as models_router  # noqa: E402
from api.routes_sessions import router as sessions_router  # noqa: E402
from chatgpt_client.browser import SessionDeadError, session  # noqa: E402
from chatgpt_client.models import registry  # noqa: E402
from config import settings  # noqa: E402


def _setup_logging() -> None:
    logging.basicConfig(level=getattr(logging, settings.log_level.upper(), logging.INFO))
    structlog.configure(
        processors=[
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.processors.add_log_level,
            structlog.processors.StackInfoRenderer(),
            structlog.dev.ConsoleRenderer(),
        ]
    )


_setup_logging()
log = structlog.get_logger("chatgpt_bridge")


async def _models_refresh_loop() -> None:
    while True:
        await asyncio.sleep(settings.models_refresh_interval_s)
        try:
            await registry.refresh(session)
        except Exception as e:
            log.warning("models.refresh.loop_error", error=str(e))


@asynccontextmanager
async def lifespan(app: FastAPI):
    try:
        await session.start()
    except SessionDeadError as e:
        log.error("startup.session_dead", error=str(e))
        yield
        return
    except Exception as e:
        log.exception("startup.failed", error=str(e))
        yield
        return

    await registry.refresh(session)
    refresh_task = asyncio.create_task(_models_refresh_loop())
    log.info("api.ready", host=settings.host, port=settings.port, models=registry.slugs)

    try:
        yield
    finally:
        refresh_task.cancel()
        try:
            await refresh_task
        except (asyncio.CancelledError, Exception):
            pass
        await session.stop()


app = FastAPI(title="ChatGPT Bridge", version="0.1.0", lifespan=lifespan)


@app.middleware("http")
async def auth_middleware(request: Request, call_next):
    if settings.api_key and request.url.path.startswith("/v1/"):
        auth = request.headers.get("authorization", "")
        token = auth.removeprefix("Bearer ").strip()
        if token != settings.api_key:
            return JSONResponse(status_code=401, content={"error": {"message": "Invalid API key"}})
    return await call_next(request)


@app.get("/healthz")
async def healthz():
    return {
        "status": "ok" if session.alive else "session_dead",
        "models": registry.slugs,
    }


app.include_router(chat_router)
app.include_router(models_router)
app.include_router(sessions_router)
app.include_router(images_router)
app.include_router(auth_router)
app.include_router(control_router)