codex-lb 0.1.2__py3-none-any.whl

app/db/models.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from datetime import datetime
+from enum import Enum
+
+from sqlalchemy import Boolean, DateTime, Float, ForeignKey, Index, Integer, LargeBinary, String, Text, func
+from sqlalchemy import Enum as SqlEnum
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class AccountStatus(str, Enum):
+    ACTIVE = "active"
+    RATE_LIMITED = "rate_limited"
+    QUOTA_EXCEEDED = "quota_exceeded"
+    PAUSED = "paused"
+    DEACTIVATED = "deactivated"
+
+
+class Account(Base):
+    __tablename__ = "accounts"
+
+    id: Mapped[str] = mapped_column(String, primary_key=True)
+    email: Mapped[str] = mapped_column(String, unique=True, nullable=False)
+    plan_type: Mapped[str] = mapped_column(String, nullable=False)
+
+    access_token_encrypted: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
+    refresh_token_encrypted: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
+    id_token_encrypted: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
+
+    last_refresh: Mapped[datetime] = mapped_column(DateTime, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), nullable=False)
+
+    status: Mapped[AccountStatus] = mapped_column(
+        SqlEnum(AccountStatus, name="account_status", validate_strings=True),
+        default=AccountStatus.ACTIVE,
+        nullable=False,
+    )
+    deactivation_reason: Mapped[str | None] = mapped_column(Text, nullable=True)
+
+
+class UsageHistory(Base):
+    __tablename__ = "usage_history"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+    account_id: Mapped[str] = mapped_column(String, ForeignKey("accounts.id"), nullable=False)
+    recorded_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), nullable=False)
+    window: Mapped[str | None] = mapped_column(String, nullable=True)
+    used_percent: Mapped[float] = mapped_column(Float, nullable=False)
+    input_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    output_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    reset_at: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    window_minutes: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    credits_has: Mapped[bool | None] = mapped_column(Boolean, nullable=True)
+    credits_unlimited: Mapped[bool | None] = mapped_column(Boolean, nullable=True)
+    credits_balance: Mapped[float | None] = mapped_column(Float, nullable=True)
+
+
+class RequestLog(Base):
+    __tablename__ = "request_logs"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+    account_id: Mapped[str] = mapped_column(String, ForeignKey("accounts.id"), nullable=False)
+    request_id: Mapped[str] = mapped_column(String, nullable=False)
+    requested_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), nullable=False)
+    model: Mapped[str] = mapped_column(String, nullable=False)
+    input_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    output_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    cached_input_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    reasoning_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    latency_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    status: Mapped[str] = mapped_column(String, nullable=False)
+    error_code: Mapped[str | None] = mapped_column(String, nullable=True)
+    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+
+
+Index("idx_usage_recorded_at", UsageHistory.recorded_at)
+Index("idx_usage_account_time", UsageHistory.account_id, UsageHistory.recorded_at)
+Index("idx_logs_account_time", RequestLog.account_id, RequestLog.requested_at)

app/db/session.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import AsyncIterator
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+
+from app.core.config.settings import get_settings
+
+DATABASE_URL = get_settings().database_url
+
+engine = create_async_engine(DATABASE_URL, echo=False)
+SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
+
+
+def _ensure_sqlite_dir(url: str) -> None:
+    prefix = "sqlite+aiosqlite:///"
+    if not url.startswith(prefix):
+        return
+    path = url[len(prefix) :]
+    if path == ":memory:":
+        return
+    Path(path).expanduser().parent.mkdir(parents=True, exist_ok=True)
+
+
+async def get_session() -> AsyncIterator[AsyncSession]:
+    async with SessionLocal() as session:
+        try:
+            yield session
+        except Exception:
+            await session.rollback()
+            raise
+        finally:
+            if session.in_transaction():
+                await session.rollback()
+
+
+async def init_db() -> None:
+    from app.db.models import Base
+
+    _ensure_sqlite_dir(DATABASE_URL)
+
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)

app/dependencies.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+
+from fastapi import Depends
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db.session import SessionLocal, get_session
+from app.modules.accounts.repository import AccountsRepository
+from app.modules.accounts.service import AccountsService
+from app.modules.oauth.service import OauthService
+from app.modules.proxy.service import ProxyService
+from app.modules.request_logs.repository import RequestLogsRepository
+from app.modules.request_logs.service import RequestLogsService
+from app.modules.usage.repository import UsageRepository
+from app.modules.usage.service import UsageService
+
+
+@dataclass(slots=True)
+class AccountsContext:
+    session: AsyncSession
+    repository: AccountsRepository
+    usage_repository: UsageRepository
+    request_logs_repository: RequestLogsRepository
+    service: AccountsService
+
+
+@dataclass(slots=True)
+class UsageContext:
+    session: AsyncSession
+    usage_repository: UsageRepository
+    request_logs_repository: RequestLogsRepository
+    accounts_repository: AccountsRepository
+    service: UsageService
+
+
+@dataclass(slots=True)
+class OauthContext:
+    service: OauthService
+
+
+@dataclass(slots=True)
+class ProxyContext:
+    service: ProxyService
+
+
+@dataclass(slots=True)
+class RequestLogsContext:
+    session: AsyncSession
+    repository: RequestLogsRepository
+    service: RequestLogsService
+
+
+def get_accounts_context(
+    session: AsyncSession = Depends(get_session),
+) -> AccountsContext:
+    repository = AccountsRepository(session)
+    usage_repository = UsageRepository(session)
+    request_logs_repository = RequestLogsRepository(session)
+    service = AccountsService(repository, usage_repository, request_logs_repository)
+    return AccountsContext(
+        session=session,
+        repository=repository,
+        usage_repository=usage_repository,
+        request_logs_repository=request_logs_repository,
+        service=service,
+    )
+
+
+def get_usage_context(
+    session: AsyncSession = Depends(get_session),
+) -> UsageContext:
+    usage_repository = UsageRepository(session)
+    request_logs_repository = RequestLogsRepository(session)
+    accounts_repository = AccountsRepository(session)
+    service = UsageService(usage_repository, request_logs_repository, accounts_repository)
+    return UsageContext(
+        session=session,
+        usage_repository=usage_repository,
+        request_logs_repository=request_logs_repository,
+        accounts_repository=accounts_repository,
+        service=service,
+    )
+
+
+@asynccontextmanager
+async def _accounts_repo_context() -> AsyncIterator[AccountsRepository]:
+    async with SessionLocal() as session:
+        try:
+            yield AccountsRepository(session)
+        except Exception:
+            await session.rollback()
+            raise
+        finally:
+            if session.in_transaction():
+                await session.rollback()
+
+
+def get_oauth_context(
+    session: AsyncSession = Depends(get_session),
+) -> OauthContext:
+    accounts_repository = AccountsRepository(session)
+    return OauthContext(service=OauthService(accounts_repository, repo_factory=_accounts_repo_context))
+
+
+def get_proxy_context(
+    session: AsyncSession = Depends(get_session),
+) -> ProxyContext:
+    accounts_repository = AccountsRepository(session)
+    usage_repository = UsageRepository(session)
+    request_logs_repository = RequestLogsRepository(session)
+    service = ProxyService(accounts_repository, usage_repository, request_logs_repository)
+    return ProxyContext(service=service)
+
+
+def get_request_logs_context(
+    session: AsyncSession = Depends(get_session),
+) -> RequestLogsContext:
+    repository = RequestLogsRepository(session)
+    service = RequestLogsService(repository)
+    return RequestLogsContext(session=session, repository=repository, service=service)

app/main.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import logging
+from contextlib import asynccontextmanager
+from pathlib import Path
+from uuid import uuid4
+
+from fastapi import FastAPI, Request
+from fastapi.exception_handlers import (
+    http_exception_handler,
+    request_validation_exception_handler,
+)
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
+from fastapi.staticfiles import StaticFiles
+from starlette.exceptions import HTTPException as StarletteHTTPException
+
+from app.core.clients.http import close_http_client, init_http_client
+from app.core.errors import dashboard_error
+from app.core.utils.request_id import get_request_id, reset_request_id, set_request_id
+from app.db.session import init_db
+from app.modules.accounts import api as accounts_api
+from app.modules.health import api as health_api
+from app.modules.oauth import api as oauth_api
+from app.modules.proxy import api as proxy_api
+from app.modules.request_logs import api as request_logs_api
+from app.modules.usage import api as usage_api
+
+logger = logging.getLogger(__name__)
+
+
+@asynccontextmanager
+async def lifespan(_: FastAPI):
+    await init_db()
+    await init_http_client()
+
+    try:
+        yield
+    finally:
+        await close_http_client()
+
+
+def create_app() -> FastAPI:
+    app = FastAPI(title="codex-lb", version="0.1.0", lifespan=lifespan)
+
+    @app.middleware("http")
+    async def request_id_middleware(request: Request, call_next) -> JSONResponse:
+        inbound_request_id = request.headers.get("x-request-id") or request.headers.get("request-id")
+        request_id = inbound_request_id or str(uuid4())
+        token = set_request_id(request_id)
+        try:
+            response = await call_next(request)
+        except Exception:
+            reset_request_id(token)
+            raise
+        response.headers.setdefault("x-request-id", request_id)
+        return response
+
+    @app.middleware("http")
+    async def api_unhandled_error_middleware(request: Request, call_next) -> JSONResponse:
+        try:
+            return await call_next(request)
+        except Exception:
+            if request.url.path.startswith("/api/"):
+                logger.exception(
+                    "Unhandled API error request_id=%s",
+                    get_request_id(),
+                )
+                return JSONResponse(
+                    status_code=500,
+                    content=dashboard_error("internal_error", "Unexpected error"),
+                )
+            raise
+
+    @app.exception_handler(RequestValidationError)
+    async def _validation_error_handler(
+        request: Request,
+        exc: RequestValidationError,
+    ) -> JSONResponse:
+        if request.url.path.startswith("/api/"):
+            return JSONResponse(
+                status_code=422,
+                content=dashboard_error("validation_error", "Invalid request payload"),
+            )
+        return await request_validation_exception_handler(request, exc)
+
+    @app.exception_handler(StarletteHTTPException)
+    async def _http_error_handler(
+        request: Request,
+        exc: StarletteHTTPException,
+    ) -> JSONResponse:
+        if request.url.path.startswith("/api/"):
+            detail = exc.detail if isinstance(exc.detail, str) else "Request failed"
+            return JSONResponse(
+                status_code=exc.status_code,
+                content=dashboard_error(f"http_{exc.status_code}", detail),
+            )
+        return await http_exception_handler(request, exc)
+
+    app.include_router(proxy_api.router)
+    app.include_router(proxy_api.usage_router)
+    app.include_router(accounts_api.router)
+    app.include_router(usage_api.router)
+    app.include_router(request_logs_api.router)
+    app.include_router(oauth_api.router)
+    app.include_router(health_api.router)
+
+    static_dir = Path(__file__).parent / "static"
+    index_html = static_dir / "index.html"
+
+    @app.get("/", include_in_schema=False)
+    async def root_redirect():
+        return RedirectResponse(url="/dashboard", status_code=302)
+
+    @app.get("/accounts", include_in_schema=False)
+    async def spa_accounts():
+        return FileResponse(index_html, media_type="text/html")
+
+    app.mount("/dashboard", StaticFiles(directory=static_dir, html=True), name="dashboard")
+
+    return app
+
+
+app = create_app()

app/modules/accounts/api.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Depends, File, UploadFile
+from fastapi.responses import JSONResponse
+
+from app.core.errors import dashboard_error
+from app.dependencies import AccountsContext, get_accounts_context
+from app.modules.accounts.schemas import (
+    AccountDeleteResponse,
+    AccountImportResponse,
+    AccountPauseResponse,
+    AccountReactivateResponse,
+    AccountsResponse,
+)
+
+router = APIRouter(prefix="/api/accounts", tags=["dashboard"])
+
+
+@router.get("", response_model=AccountsResponse)
+async def list_accounts(
+    context: AccountsContext = Depends(get_accounts_context),
+) -> AccountsResponse:
+    accounts = await context.service.list_accounts()
+    return AccountsResponse(accounts=accounts)
+
+
+@router.post("/import", response_model=AccountImportResponse)
+async def import_account(
+    auth_json: UploadFile = File(...),
+    context: AccountsContext = Depends(get_accounts_context),
+) -> AccountImportResponse | JSONResponse:
+    raw = await auth_json.read()
+    try:
+        return await context.service.import_account(raw)
+    except Exception:
+        return JSONResponse(
+            status_code=400,
+            content=dashboard_error("invalid_auth_json", "Invalid auth.json payload"),
+        )
+
+
+@router.post("/{account_id}/reactivate", response_model=AccountReactivateResponse)
+async def reactivate_account(
+    account_id: str,
+    context: AccountsContext = Depends(get_accounts_context),
+) -> AccountReactivateResponse | JSONResponse:
+    success = await context.service.reactivate_account(account_id)
+    if not success:
+        return JSONResponse(
+            status_code=404,
+            content=dashboard_error("account_not_found", "Account not found"),
+        )
+    return AccountReactivateResponse(status="reactivated")
+
+
+@router.post("/{account_id}/pause", response_model=AccountPauseResponse)
+async def pause_account(
+    account_id: str,
+    context: AccountsContext = Depends(get_accounts_context),
+) -> AccountPauseResponse | JSONResponse:
+    success = await context.service.pause_account(account_id)
+    if not success:
+        return JSONResponse(
+            status_code=404,
+            content=dashboard_error("account_not_found", "Account not found"),
+        )
+    return AccountPauseResponse(status="paused")
+
+
+@router.delete("/{account_id}", response_model=AccountDeleteResponse)
+async def delete_account(
+    account_id: str,
+    context: AccountsContext = Depends(get_accounts_context),
+) -> AccountDeleteResponse | JSONResponse:
+    success = await context.service.delete_account(account_id)
+    if not success:
+        return JSONResponse(
+            status_code=404,
+            content=dashboard_error("account_not_found", "Account not found"),
+        )
+    return AccountDeleteResponse(status="deleted")

app/modules/accounts/repository.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import delete, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db.models import Account, AccountStatus
+
+
+class AccountsRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def list_accounts(self) -> list[Account]:
+        result = await self._session.execute(select(Account).order_by(Account.email))
+        return list(result.scalars().all())
+
+    async def upsert(self, account: Account) -> Account:
+        existing = await self._session.get(Account, account.id)
+        if existing:
+            existing.email = account.email
+            existing.plan_type = account.plan_type
+            existing.access_token_encrypted = account.access_token_encrypted
+            existing.refresh_token_encrypted = account.refresh_token_encrypted
+            existing.id_token_encrypted = account.id_token_encrypted
+            existing.last_refresh = account.last_refresh
+            existing.status = account.status
+            existing.deactivation_reason = account.deactivation_reason
+            await self._session.commit()
+            await self._session.refresh(existing)
+            return existing
+
+        self._session.add(account)
+        await self._session.commit()
+        await self._session.refresh(account)
+        return account
+
+    async def update_status(
+        self,
+        account_id: str,
+        status: AccountStatus,
+        deactivation_reason: str | None = None,
+    ) -> bool:
+        result = await self._session.execute(
+            update(Account)
+            .where(Account.id == account_id)
+            .values(status=status, deactivation_reason=deactivation_reason)
+        )
+        await self._session.commit()
+        return bool(result.rowcount)
+
+    async def delete(self, account_id: str) -> bool:
+        result = await self._session.execute(delete(Account).where(Account.id == account_id))
+        await self._session.commit()
+        return bool(result.rowcount)
+
+    async def update_tokens(
+        self,
+        account_id: str,
+        access_token_encrypted: bytes,
+        refresh_token_encrypted: bytes,
+        id_token_encrypted: bytes,
+        last_refresh: datetime,
+        plan_type: str | None = None,
+        email: str | None = None,
+    ) -> bool:
+        values = {
+            "access_token_encrypted": access_token_encrypted,
+            "refresh_token_encrypted": refresh_token_encrypted,
+            "id_token_encrypted": id_token_encrypted,
+            "last_refresh": last_refresh,
+        }
+        if plan_type is not None:
+            values["plan_type"] = plan_type
+        if email is not None:
+            values["email"] = email
+        result = await self._session.execute(update(Account).where(Account.id == account_id).values(**values))
+        await self._session.commit()
+        return bool(result.rowcount)

app/modules/accounts/schemas.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import List
+
+from pydantic import Field
+
+from app.modules.shared.schemas import DashboardModel
+
+
+class AccountUsage(DashboardModel):
+    primary_remaining_percent: float | None = None
+    secondary_remaining_percent: float | None = None
+
+
+class AccountTokenStatus(DashboardModel):
+    expires_at: datetime | None = None
+    state: str | None = None
+
+
+class AccountAuthStatus(DashboardModel):
+    access: AccountTokenStatus | None = None
+    refresh: AccountTokenStatus | None = None
+    id_token: AccountTokenStatus | None = None
+
+
+class AccountSummary(DashboardModel):
+    account_id: str
+    email: str
+    display_name: str
+    plan_type: str
+    status: str
+    usage: AccountUsage | None = None
+    reset_at_primary: datetime | None = None
+    reset_at_secondary: datetime | None = None
+    last_refresh_at: datetime | None = None
+    capacity_credits_primary: float | None = None
+    remaining_credits_primary: float | None = None
+    capacity_credits_secondary: float | None = None
+    remaining_credits_secondary: float | None = None
+    cost_usd_24h: float | None = None
+    deactivation_reason: str | None = None
+    auth: AccountAuthStatus | None = None
+
+
+class AccountsResponse(DashboardModel):
+    accounts: List[AccountSummary] = Field(default_factory=list)
+
+
+class AccountImportResponse(DashboardModel):
+    account_id: str
+    email: str
+    plan_type: str
+    status: str
+
+
+class AccountPauseResponse(DashboardModel):
+    status: str
+
+
+class AccountReactivateResponse(DashboardModel):
+    status: str
+
+
+class AccountDeleteResponse(DashboardModel):
+    status: str