codex-lb 0.1.2__py3-none-any.whl

app/modules/proxy/types.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from app.core.types import JsonValue
+
+
+@dataclass(frozen=True)
+class RateLimitWindowSnapshotData:
+    used_percent: int
+    limit_window_seconds: int
+    reset_after_seconds: int
+    reset_at: int
+
+
+@dataclass(frozen=True)
+class RateLimitStatusDetailsData:
+    allowed: bool
+    limit_reached: bool
+    primary_window: RateLimitWindowSnapshotData | None = None
+    secondary_window: RateLimitWindowSnapshotData | None = None
+
+
+@dataclass(frozen=True)
+class CreditStatusDetailsData:
+    has_credits: bool
+    unlimited: bool
+    balance: str | None = None
+    approx_local_messages: list[JsonValue] | None = None
+    approx_cloud_messages: list[JsonValue] | None = None
+
+
+@dataclass(frozen=True)
+class RateLimitStatusPayloadData:
+    plan_type: str
+    rate_limit: RateLimitStatusDetailsData | None = None
+    credits: CreditStatusDetailsData | None = None

app/modules/proxy/usage_updater.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import logging
+import math
+from typing import Mapping
+
+from app.core.auth.refresh import RefreshError
+from app.core.clients.usage import UsageFetchError, fetch_usage
+from app.core.config.settings import get_settings
+from app.core.crypto import TokenEncryptor
+from app.core.usage.models import UsagePayload
+from app.core.utils.request_id import get_request_id
+from app.core.utils.time import utcnow
+from app.db.models import Account, AccountStatus, UsageHistory
+from app.modules.accounts.repository import AccountsRepository
+from app.modules.proxy.auth_manager import AuthManager
+from app.modules.usage.repository import UsageRepository
+
+logger = logging.getLogger(__name__)
+
+
+class UsageUpdater:
+    def __init__(
+        self,
+        usage_repo: UsageRepository,
+        accounts_repo: AccountsRepository | None = None,
+    ) -> None:
+        self._usage_repo = usage_repo
+        self._encryptor = TokenEncryptor()
+        self._auth_manager = AuthManager(accounts_repo) if accounts_repo else None
+
+    async def refresh_accounts(
+        self,
+        accounts: list[Account],
+        latest_usage: Mapping[str, UsageHistory],
+    ) -> None:
+        settings = get_settings()
+        if not settings.usage_refresh_enabled:
+            return
+
+        now = utcnow()
+        interval = settings.usage_refresh_interval_seconds
+        for account in accounts:
+            if account.status == AccountStatus.DEACTIVATED:
+                continue
+            latest = latest_usage.get(account.id)
+            if latest and (now - latest.recorded_at).total_seconds() < interval:
+                continue
+            # NOTE: AsyncSession is not safe for concurrent use. Run sequentially
+            # within the request-scoped session to avoid PK collisions and
+            # flush-time warnings (SAWarning: Session.add during flush).
+            try:
+                await self._refresh_account(account)
+            except Exception as exc:
+                logger.warning(
+                    "Usage refresh failed account_id=%s request_id=%s error=%s",
+                    account.id,
+                    get_request_id(),
+                    exc,
+                    exc_info=True,
+                )
+                # swallow per-account failures so the whole refresh loop keeps going
+                continue
+
+    async def _refresh_account(self, account: Account) -> None:
+        access_token = self._encryptor.decrypt(account.access_token_encrypted)
+        try:
+            payload = await fetch_usage(
+                access_token=access_token,
+                account_id=account.id,
+            )
+        except UsageFetchError as exc:
+            if exc.status_code != 401 or not self._auth_manager:
+                return
+            try:
+                account = await self._auth_manager.ensure_fresh(account, force=True)
+            except RefreshError:
+                return
+            access_token = self._encryptor.decrypt(account.access_token_encrypted)
+            try:
+                payload = await fetch_usage(
+                    access_token=access_token,
+                    account_id=account.id,
+                )
+            except UsageFetchError:
+                return
+
+        rate_limit = payload.rate_limit
+        primary = rate_limit.primary_window if rate_limit else None
+        credits_has, credits_unlimited, credits_balance = _credits_snapshot(payload)
+        primary_window_minutes = _window_minutes(primary.limit_window_seconds) if primary else None
+        secondary = rate_limit.secondary_window if rate_limit else None
+        secondary_window_minutes = _window_minutes(secondary.limit_window_seconds) if secondary else None
+
+        if primary and primary.used_percent is not None:
+            await self._usage_repo.add_entry(
+                account_id=account.id,
+                used_percent=primary.used_percent,
+                input_tokens=None,
+                output_tokens=None,
+                window="primary",
+                reset_at=primary.reset_at,
+                window_minutes=primary_window_minutes,
+                credits_has=credits_has,
+                credits_unlimited=credits_unlimited,
+                credits_balance=credits_balance,
+            )
+
+        if secondary and secondary.used_percent is not None:
+            await self._usage_repo.add_entry(
+                account_id=account.id,
+                used_percent=secondary.used_percent,
+                input_tokens=None,
+                output_tokens=None,
+                window="secondary",
+                reset_at=secondary.reset_at,
+                window_minutes=secondary_window_minutes,
+            )
+
+
+def _credits_snapshot(payload: UsagePayload) -> tuple[bool | None, bool | None, float | None]:
+    credits = payload.credits
+    if credits is None:
+        return None, None, None
+    credits_has = credits.has_credits
+    credits_unlimited = credits.unlimited
+    balance_value = credits.balance
+    return credits_has, credits_unlimited, _parse_credits_balance(balance_value)
+
+
+def _parse_credits_balance(value: str | int | float | None) -> float | None:
+    if value is None:
+        return None
+    if isinstance(value, (int, float)):
+        return float(value)
+    if isinstance(value, str):
+        try:
+            return float(value.strip())
+        except ValueError:
+            return None
+    return None
+
+
+def _window_minutes(limit_seconds: int | None) -> int | None:
+    if not limit_seconds or limit_seconds <= 0:
+        return None
+    return max(1, math.ceil(limit_seconds / 60))

app/modules/request_logs/api.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+from fastapi import APIRouter, Depends, Query
+
+from app.dependencies import RequestLogsContext, get_request_logs_context
+from app.modules.request_logs.schemas import RequestLogsResponse
+
+router = APIRouter(prefix="/api/request-logs", tags=["dashboard"])
+
+
+@router.get("", response_model=RequestLogsResponse)
+async def list_request_logs(
+    limit: int = Query(50, ge=1, le=200),
+    account_id: str | None = Query(default=None, alias="accountId"),
+    status: str | None = Query(default=None),
+    model: str | None = Query(default=None),
+    since: datetime | None = Query(default=None),
+    until: datetime | None = Query(default=None),
+    context: RequestLogsContext = Depends(get_request_logs_context),
+) -> RequestLogsResponse:
+    logs = await context.service.list_recent(
+        limit=limit,
+        since=since,
+        until=until,
+        account_id=account_id,
+        model=model,
+        status=status,
+    )
+    return RequestLogsResponse(requests=logs)

app/modules/request_logs/repository.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import and_, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.utils.request_id import ensure_request_id
+from app.core.utils.time import utcnow
+from app.db.models import RequestLog
+
+
+class RequestLogsRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def list_since(self, since: datetime) -> list[RequestLog]:
+        result = await self._session.execute(select(RequestLog).where(RequestLog.requested_at >= since))
+        return list(result.scalars().all())
+
+    async def add_log(
+        self,
+        account_id: str,
+        request_id: str,
+        model: str,
+        input_tokens: int | None,
+        output_tokens: int | None,
+        latency_ms: int | None,
+        status: str,
+        error_code: str | None,
+        error_message: str | None = None,
+        requested_at: datetime | None = None,
+        cached_input_tokens: int | None = None,
+        reasoning_tokens: int | None = None,
+    ) -> RequestLog:
+        resolved_request_id = ensure_request_id(request_id)
+        log = RequestLog(
+            account_id=account_id,
+            request_id=resolved_request_id,
+            model=model,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cached_input_tokens=cached_input_tokens,
+            reasoning_tokens=reasoning_tokens,
+            latency_ms=latency_ms,
+            status=status,
+            error_code=error_code,
+            error_message=error_message,
+            requested_at=requested_at or utcnow(),
+        )
+        self._session.add(log)
+        await self._session.commit()
+        await self._session.refresh(log)
+        return log
+
+    async def list_recent(
+        self,
+        limit: int = 50,
+        since: datetime | None = None,
+        until: datetime | None = None,
+        account_id: str | None = None,
+        model: str | None = None,
+        status: str | None = None,
+        error_codes: list[str] | None = None,
+    ) -> list[RequestLog]:
+        conditions = []
+        if since is not None:
+            conditions.append(RequestLog.requested_at >= since)
+        if until is not None:
+            conditions.append(RequestLog.requested_at <= until)
+        if account_id is not None:
+            conditions.append(RequestLog.account_id == account_id)
+        if model is not None:
+            conditions.append(RequestLog.model == model)
+        if status is not None:
+            conditions.append(RequestLog.status == status)
+        if error_codes:
+            conditions.append(RequestLog.error_code.in_(error_codes))
+
+        stmt = select(RequestLog).order_by(RequestLog.requested_at.desc())
+        if conditions:
+            stmt = stmt.where(and_(*conditions))
+        if limit:
+            stmt = stmt.limit(limit)
+        result = await self._session.execute(stmt)
+        return list(result.scalars().all())

app/modules/request_logs/schemas.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import List
+
+from pydantic import Field
+
+from app.modules.shared.schemas import DashboardModel
+
+
+class RequestLogEntry(DashboardModel):
+    requested_at: datetime
+    account_id: str
+    request_id: str
+    model: str
+    status: str
+    error_code: str | None = None
+    error_message: str | None = None
+    tokens: int | None = None
+    cost_usd: float | None = None
+    latency_ms: int | None = None
+
+
+class RequestLogsResponse(DashboardModel):
+    requests: List[RequestLogEntry] = Field(default_factory=list)

app/modules/request_logs/service.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+from app.core.usage.logs import cost_from_log, total_tokens_from_log
+from app.db.models import RequestLog
+from app.modules.request_logs.repository import RequestLogsRepository
+from app.modules.request_logs.schemas import RequestLogEntry
+
+RATE_LIMIT_CODES = {"rate_limit_exceeded", "usage_limit_reached"}
+QUOTA_CODES = {"insufficient_quota", "usage_not_included", "quota_exceeded"}
+
+
+class RequestLogsService:
+    def __init__(self, repo: RequestLogsRepository) -> None:
+        self._repo = repo
+
+    async def list_recent(
+        self,
+        limit: int = 50,
+        since: datetime | None = None,
+        until: datetime | None = None,
+        account_id: str | None = None,
+        model: str | None = None,
+        status: str | None = None,
+    ) -> list[RequestLogEntry]:
+        status_filter, error_codes = _map_status_filter(status)
+        logs = await self._repo.list_recent(
+            limit=limit,
+            since=since,
+            until=until,
+            account_id=account_id,
+            model=model,
+            status=status_filter,
+            error_codes=error_codes,
+        )
+        return [_to_entry(log) for log in logs]
+
+
+def _map_status_filter(status: str | None) -> tuple[str | None, list[str] | None]:
+    if not status:
+        return None, None
+    normalized = status.lower()
+    if normalized == "ok":
+        return "success", None
+    if normalized == "rate_limit":
+        return "error", sorted(RATE_LIMIT_CODES)
+    if normalized == "quota":
+        return "error", sorted(QUOTA_CODES)
+    if normalized == "error":
+        return "error", None
+    return status, None
+
+
+def _log_status(log: RequestLog) -> str:
+    if log.status == "success":
+        return "ok"
+    if log.error_code in RATE_LIMIT_CODES:
+        return "rate_limit"
+    if log.error_code in QUOTA_CODES:
+        return "quota"
+    return "error"
+
+
+def _to_entry(log: RequestLog) -> RequestLogEntry:
+    return RequestLogEntry(
+        requested_at=log.requested_at,
+        account_id=log.account_id,
+        request_id=log.request_id,
+        model=log.model,
+        status=_log_status(log),
+        error_code=log.error_code,
+        error_message=log.error_message,
+        tokens=total_tokens_from_log(log),
+        cost_usd=cost_from_log(log, precision=6),
+        latency_ms=log.latency_ms,
+    )

app/modules/shared/schemas.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class DashboardModel(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)

app/modules/usage/api.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Depends, Query
+
+from app.dependencies import UsageContext, get_usage_context
+from app.modules.usage.schemas import UsageHistoryResponse, UsageSummaryResponse, UsageWindowResponse
+
+router = APIRouter(prefix="/api/usage", tags=["dashboard"])
+
+
+@router.get("/summary", response_model=UsageSummaryResponse)
+async def get_usage_summary(
+    context: UsageContext = Depends(get_usage_context),
+) -> UsageSummaryResponse:
+    return await context.service.get_usage_summary()
+
+
+@router.get("/history", response_model=UsageHistoryResponse)
+async def get_usage_history(
+    hours: int = Query(24, ge=1, le=168),
+    context: UsageContext = Depends(get_usage_context),
+) -> UsageHistoryResponse:
+    return await context.service.get_usage_history(hours)
+
+
+@router.get("/window", response_model=UsageWindowResponse)
+async def get_usage_window(
+    window: str = Query("primary", pattern="^(primary|secondary)$"),
+    context: UsageContext = Depends(get_usage_context),
+) -> UsageWindowResponse:
+    return await context.service.get_usage_window(window)

app/modules/usage/repository.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import func, or_, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.usage.types import UsageAggregateRow
+from app.core.utils.time import utcnow
+from app.db.models import UsageHistory
+
+
+class UsageRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def add_entry(
+        self,
+        account_id: str,
+        used_percent: float,
+        input_tokens: int | None = None,
+        output_tokens: int | None = None,
+        recorded_at: datetime | None = None,
+        window: str | None = None,
+        reset_at: int | None = None,
+        window_minutes: int | None = None,
+        credits_has: bool | None = None,
+        credits_unlimited: bool | None = None,
+        credits_balance: float | None = None,
+    ) -> UsageHistory:
+        entry = UsageHistory(
+            account_id=account_id,
+            used_percent=used_percent,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            window=window,
+            reset_at=reset_at,
+            window_minutes=window_minutes,
+            credits_has=credits_has,
+            credits_unlimited=credits_unlimited,
+            credits_balance=credits_balance,
+            recorded_at=recorded_at or utcnow(),
+        )
+        self._session.add(entry)
+        await self._session.commit()
+        await self._session.refresh(entry)
+        return entry
+
+    async def aggregate_since(
+        self,
+        since: datetime,
+        window: str | None = None,
+    ) -> list[UsageAggregateRow]:
+        conditions = [UsageHistory.recorded_at >= since]
+        if window:
+            if window == "primary":
+                conditions.append(or_(UsageHistory.window == "primary", UsageHistory.window.is_(None)))
+            else:
+                conditions.append(UsageHistory.window == window)
+        stmt = (
+            select(
+                UsageHistory.account_id,
+                func.avg(UsageHistory.used_percent).label("used_percent_avg"),
+                func.sum(UsageHistory.input_tokens).label("input_tokens_sum"),
+                func.sum(UsageHistory.output_tokens).label("output_tokens_sum"),
+                func.count(UsageHistory.id).label("samples"),
+                func.max(UsageHistory.recorded_at).label("last_recorded_at"),
+                func.max(UsageHistory.reset_at).label("reset_at_max"),
+                func.max(UsageHistory.window_minutes).label("window_minutes_max"),
+            )
+            .where(*conditions)
+            .group_by(UsageHistory.account_id)
+        )
+        result = await self._session.execute(stmt)
+        rows = result.all()
+        return [
+            UsageAggregateRow(
+                account_id=row.account_id,
+                used_percent_avg=float(row.used_percent_avg) if row.used_percent_avg is not None else None,
+                input_tokens_sum=int(row.input_tokens_sum) if row.input_tokens_sum is not None else None,
+                output_tokens_sum=int(row.output_tokens_sum) if row.output_tokens_sum is not None else None,
+                samples=int(row.samples),
+                last_recorded_at=row.last_recorded_at,
+                reset_at_max=int(row.reset_at_max) if row.reset_at_max is not None else None,
+                window_minutes_max=int(row.window_minutes_max) if row.window_minutes_max is not None else None,
+            )
+            for row in rows
+        ]
+
+    async def latest_by_account(self, window: str | None = None) -> dict[str, UsageHistory]:
+        if window:
+            if window == "primary":
+                conditions = or_(UsageHistory.window == "primary", UsageHistory.window.is_(None))
+            else:
+                conditions = UsageHistory.window == window
+        else:
+            conditions = or_(UsageHistory.window == "primary", UsageHistory.window.is_(None))
+        stmt = select(UsageHistory).where(conditions).order_by(UsageHistory.account_id, UsageHistory.recorded_at.desc())
+        result = await self._session.execute(stmt)
+        latest: dict[str, UsageHistory] = {}
+        for entry in result.scalars().all():
+            if entry.account_id not in latest:
+                latest[entry.account_id] = entry
+        return latest
+
+    async def latest_window_minutes(self, window: str) -> int | None:
+        if window == "primary":
+            conditions = or_(UsageHistory.window == "primary", UsageHistory.window.is_(None))
+        else:
+            conditions = UsageHistory.window == window
+        result = await self._session.execute(select(func.max(UsageHistory.window_minutes)).where(conditions))
+        value = result.scalar_one_or_none()
+        return int(value) if value is not None else None

app/modules/usage/schemas.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import List
+
+from pydantic import Field
+
+from app.modules.shared.schemas import DashboardModel
+
+
+class UsageWindow(DashboardModel):
+    remaining_percent: float
+    capacity_credits: float
+    remaining_credits: float
+    reset_at: datetime | None = None
+    window_minutes: int | None = None
+
+
+class UsageCostByModel(DashboardModel):
+    model: str
+    usd: float
+
+
+class UsageCost(DashboardModel):
+    currency: str
+    total_usd_7d: float = Field(alias="totalUsd7d")
+    by_model: List[UsageCostByModel] = Field(default_factory=list)
+
+
+class UsageMetrics(DashboardModel):
+    requests_7d: int | None = Field(default=None, alias="requests7d")
+    tokens_secondary_window: int | None = None
+    error_rate_7d: float | None = Field(default=None, alias="errorRate7d")
+    top_error: str | None = None
+
+
+class UsageSummaryResponse(DashboardModel):
+    primary_window: UsageWindow
+    secondary_window: UsageWindow | None = None
+    cost: UsageCost
+    metrics: UsageMetrics | None = None
+
+
+class UsageHistoryItem(DashboardModel):
+    account_id: str
+    email: str
+    remaining_percent_avg: float
+    capacity_credits: float
+    remaining_credits: float
+    request_count: int
+    cost_usd: float
+
+
+class UsageHistoryResponse(DashboardModel):
+    window_hours: int
+    accounts: List[UsageHistoryItem] = Field(default_factory=list)
+
+
+class UsageWindowResponse(DashboardModel):
+    window_key: str
+    window_minutes: int | None = None
+    accounts: List[UsageHistoryItem] = Field(default_factory=list)