codex-lb 0.1.4__py3-none-any.whl → 0.2.0__py3-none-any.whl

app/modules/accounts/service.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta, timezone
+from typing import cast
 
 from app.core import usage as usage_core
 from app.core.auth import (
@@ -12,7 +13,8 @@ from app.core.auth import (
     parse_auth_json,
 )
 from app.core.crypto import TokenEncryptor
-from app.core.usage.logs import cost_from_log
+from app.core.plan_types import coerce_account_plan_type
+from app.core.usage.logs import RequestLogLike, cost_from_log
 from app.core.utils.time import from_epoch_seconds, to_utc_naive, utcnow
 from app.db.models import Account, AccountStatus, UsageHistory
 from app.modules.accounts.repository import AccountsRepository
@@ -23,9 +25,9 @@ from app.modules.accounts.schemas import (
     AccountTokenStatus,
     AccountUsage,
 )
-from app.modules.proxy.usage_updater import UsageUpdater
 from app.modules.request_logs.repository import RequestLogsRepository
 from app.modules.usage.repository import UsageRepository
+from app.modules.usage.updater import UsageUpdater
 
 
 class AccountsService:
@@ -64,7 +66,7 @@ class AccountsService:
         claims = claims_from_auth(auth)
 
         email = claims.email or DEFAULT_EMAIL
-        plan_type = claims.plan_type or DEFAULT_PLAN
+        plan_type = coerce_account_plan_type(claims.plan_type, DEFAULT_PLAN)
         account_id = claims.account_id or fallback_account_id(email)
         last_refresh = to_utc_naive(auth.last_refresh_at) if auth.last_refresh_at else utcnow()
 
@@ -107,6 +109,7 @@ class AccountsService:
         secondary_usage: UsageHistory | None,
         cost_usd_24h: float | None,
     ) -> AccountSummary:
+        plan_type = coerce_account_plan_type(account.plan_type, DEFAULT_PLAN)
         auth_status = self._build_auth_status(account)
         primary_used_percent = _normalize_used_percent(primary_usage) or 0.0
         secondary_used_percent = _normalize_used_percent(secondary_usage) or 0.0
@@ -114,8 +117,8 @@ class AccountsService:
         secondary_remaining_percent = usage_core.remaining_percent_from_used(secondary_used_percent) or 0.0
         reset_at_primary = from_epoch_seconds(primary_usage.reset_at) if primary_usage is not None else None
         reset_at_secondary = from_epoch_seconds(secondary_usage.reset_at) if secondary_usage is not None else None
-        capacity_primary = usage_core.capacity_for_plan(account.plan_type, "primary")
-        capacity_secondary = usage_core.capacity_for_plan(account.plan_type, "secondary")
+        capacity_primary = usage_core.capacity_for_plan(plan_type, "primary")
+        capacity_secondary = usage_core.capacity_for_plan(plan_type, "secondary")
         remaining_credits_primary = usage_core.remaining_credits_from_percent(
             primary_used_percent,
             capacity_primary,
@@ -128,7 +131,7 @@ class AccountsService:
             account_id=account.id,
             email=account.email,
             display_name=account.email,
-            plan_type=account.plan_type,
+            plan_type=plan_type,
             status=account.status.value,
             usage=AccountUsage(
                 primary_remaining_percent=primary_remaining_percent,
@@ -186,7 +189,7 @@ class AccountsService:
         logs = await self._logs_repo.list_since(since)
         totals: dict[str, float] = {}
         for log in logs:
-            cost = cost_from_log(log)
+            cost = cost_from_log(cast(RequestLogLike, log))
             if cost is None:
                 continue
             totals[log.account_id] = totals.get(log.account_id, 0.0) + cost

app/modules/health/api.py

@@ -2,9 +2,11 @@ from __future__ import annotations
 
 from fastapi import APIRouter
 
+from app.modules.health.schemas import HealthResponse
+
 router = APIRouter(tags=["health"])
 
 
-@router.get("/health")
-async def health_check() -> dict:
-    return {"status": "ok"}
+@router.get("/health", response_model=HealthResponse)
+async def health_check() -> HealthResponse:
+    return HealthResponse(status="ok")

app/modules/health/schemas.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict
+
+
+class HealthResponse(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    status: str

app/modules/oauth/service.py

@@ -28,6 +28,7 @@ from app.core.clients.oauth import (
 )
 from app.core.config.settings import get_settings
 from app.core.crypto import TokenEncryptor
+from app.core.plan_types import coerce_account_plan_type
 from app.core.utils.time import utcnow
 from app.db.models import Account, AccountStatus
 from app.modules.accounts.repository import AccountsRepository
@@ -295,7 +296,10 @@ class OauthService:
         auth_claims = claims.auth or OpenAIAuthClaims()
         account_id = auth_claims.chatgpt_account_id or claims.chatgpt_account_id
         email = claims.email or DEFAULT_EMAIL
-        plan_type = auth_claims.chatgpt_plan_type or claims.chatgpt_plan_type or DEFAULT_PLAN
+        plan_type = coerce_account_plan_type(
+            auth_claims.chatgpt_plan_type or claims.chatgpt_plan_type,
+            DEFAULT_PLAN,
+        )
         account_id = account_id or fallback_account_id(email)
 
         account = Account(

app/modules/proxy/helpers.py

@@ -0,0 +1,285 @@
+from __future__ import annotations
+
+from typing import Iterable
+
+from pydantic import ValidationError
+
+from app.core import usage as usage_core
+from app.core.balancer.types import UpstreamError
+from app.core.errors import OpenAIErrorDetail, OpenAIErrorEnvelope
+from app.core.openai.models import OpenAIError
+from app.core.plan_types import normalize_rate_limit_plan_type
+from app.core.usage.types import UsageWindowRow, UsageWindowSummary
+from app.db.models import Account, AccountStatus, UsageHistory
+from app.modules.proxy.types import (
+    CreditStatusDetailsData,
+    RateLimitStatusDetailsData,
+    RateLimitWindowSnapshotData,
+)
+
+PLAN_TYPE_PRIORITY = (
+    "enterprise",
+    "business",
+    "team",
+    "pro",
+    "plus",
+    "education",
+    "edu",
+    "free_workspace",
+    "free",
+    "go",
+    "guest",
+    "quorum",
+    "k12",
+)
+
+
+def _header_account_id(account_id: str | None) -> str | None:
+    if not account_id:
+        return None
+    if account_id.startswith(("email_", "local_")):
+        return None
+    return account_id
+
+
+def _select_accounts_for_limits(accounts: Iterable[Account]) -> list[Account]:
+    return [account for account in accounts if account.status not in (AccountStatus.DEACTIVATED, AccountStatus.PAUSED)]
+
+
+def _summarize_window(
+    rows: list[UsageWindowRow],
+    account_map: dict[str, Account],
+    window: str,
+) -> UsageWindowSummary | None:
+    if not rows:
+        return None
+    return usage_core.summarize_usage_window(rows, account_map, window)
+
+
+def _window_snapshot(
+    summary: UsageWindowSummary | None,
+    rows: list[UsageWindowRow],
+    window: str,
+    now_epoch: int,
+) -> RateLimitWindowSnapshotData | None:
+    if summary is None:
+        return None
+
+    used_percent = _normalize_used_percent(summary.used_percent, rows)
+    if used_percent is None:
+        return None
+
+    reset_at = summary.reset_at
+    if reset_at is None:
+        return None
+
+    window_minutes = summary.window_minutes or usage_core.default_window_minutes(window)
+    if not window_minutes:
+        return None
+
+    limit_window_seconds = int(window_minutes * 60)
+    reset_after_seconds = max(0, int(reset_at) - now_epoch)
+
+    return RateLimitWindowSnapshotData(
+        used_percent=_percent_to_int(used_percent),
+        limit_window_seconds=limit_window_seconds,
+        reset_after_seconds=reset_after_seconds,
+        reset_at=int(reset_at),
+    )
+
+
+def _normalize_used_percent(
+    value: float | None,
+    rows: Iterable[UsageWindowRow],
+) -> float | None:
+    if value is not None:
+        return value
+    values = [row.used_percent for row in rows if row.used_percent is not None]
+    if not values:
+        return None
+    return sum(values) / len(values)
+
+
+def _percent_to_int(value: float) -> int:
+    bounded = max(0.0, min(100.0, value))
+    return int(bounded)
+
+
+def _rate_limit_details(
+    primary: RateLimitWindowSnapshotData | None,
+    secondary: RateLimitWindowSnapshotData | None,
+) -> RateLimitStatusDetailsData | None:
+    if not primary and not secondary:
+        return None
+    used_percents = [window.used_percent for window in (primary, secondary) if window]
+    limit_reached = any(used >= 100 for used in used_percents)
+    return RateLimitStatusDetailsData(
+        allowed=not limit_reached,
+        limit_reached=limit_reached,
+        primary_window=primary,
+        secondary_window=secondary,
+    )
+
+
+def _aggregate_credits(entries: Iterable[UsageHistory]) -> tuple[bool, bool, float] | None:
+    has_data = False
+    has_credits = False
+    unlimited = False
+    balance_total = 0.0
+
+    for entry in entries:
+        credits_has = entry.credits_has
+        credits_unlimited = entry.credits_unlimited
+        credits_balance = entry.credits_balance
+        if credits_has is None and credits_unlimited is None and credits_balance is None:
+            continue
+        has_data = True
+        if credits_has is True:
+            has_credits = True
+        if credits_unlimited is True:
+            unlimited = True
+        if credits_balance is not None and not credits_unlimited:
+            try:
+                balance_total += float(credits_balance)
+            except (TypeError, ValueError):
+                continue
+
+    if not has_data:
+        return None
+    if unlimited:
+        has_credits = True
+    return has_credits, unlimited, balance_total
+
+
+def _credits_snapshot(entries: Iterable[UsageHistory]) -> CreditStatusDetailsData | None:
+    aggregate = _aggregate_credits(entries)
+    if aggregate is None:
+        return None
+    has_credits, unlimited, balance_total = aggregate
+    balance_value = str(round(balance_total, 2))
+    return CreditStatusDetailsData(
+        has_credits=has_credits,
+        unlimited=unlimited,
+        balance=balance_value,
+        approx_local_messages=None,
+        approx_cloud_messages=None,
+    )
+
+
+def _plan_type_for_accounts(accounts: Iterable[Account]) -> str:
+    normalized = [_normalize_plan_type(account.plan_type) for account in accounts]
+    filtered = [plan for plan in normalized if plan is not None]
+    if not filtered:
+        return "guest"
+    unique = set(filtered)
+    if len(unique) == 1:
+        return filtered[0]
+    for plan in PLAN_TYPE_PRIORITY:
+        if plan in unique:
+            return plan
+    return "guest"
+
+
+def _normalize_plan_type(value: str | None) -> str | None:
+    return normalize_rate_limit_plan_type(value)
+
+
+def _rate_limit_headers(
+    window_label: str,
+    summary: UsageWindowSummary,
+) -> dict[str, str]:
+    used_percent = summary.used_percent
+    window_minutes = summary.window_minutes
+    if used_percent is None or window_minutes is None:
+        return {}
+    headers = {
+        f"x-codex-{window_label}-used-percent": str(float(used_percent)),
+        f"x-codex-{window_label}-window-minutes": str(int(window_minutes)),
+    }
+    reset_at = summary.reset_at
+    if reset_at is not None:
+        headers[f"x-codex-{window_label}-reset-at"] = str(int(reset_at))
+    return headers
+
+
+def _credits_headers(entries: Iterable[UsageHistory]) -> dict[str, str]:
+    aggregate = _aggregate_credits(entries)
+    if aggregate is None:
+        return {}
+    has_credits, unlimited, balance_total = aggregate
+    balance_value = f"{balance_total:.2f}"
+    return {
+        "x-codex-credits-has-credits": "true" if has_credits else "false",
+        "x-codex-credits-unlimited": "true" if unlimited else "false",
+        "x-codex-credits-balance": balance_value,
+    }
+
+
+def _normalize_error_code(code: str | None, error_type: str | None) -> str:
+    value = code or error_type
+    if not value:
+        return "upstream_error"
+    return value.lower()
+
+
+def _parse_openai_error(payload: OpenAIErrorEnvelope) -> OpenAIError | None:
+    error = payload.get("error")
+    if not error:
+        return None
+    try:
+        return OpenAIError.model_validate(error)
+    except ValidationError:
+        if not isinstance(error, dict):
+            return None
+        return OpenAIError(
+            message=_coerce_str(error.get("message")),
+            type=_coerce_str(error.get("type")),
+            code=_coerce_str(error.get("code")),
+            param=_coerce_str(error.get("param")),
+            plan_type=_coerce_str(error.get("plan_type")),
+            resets_at=_coerce_number(error.get("resets_at")),
+            resets_in_seconds=_coerce_number(error.get("resets_in_seconds")),
+        )
+
+
+def _coerce_str(value: object) -> str | None:
+    return value if isinstance(value, str) else None
+
+
+def _coerce_number(value: object) -> int | float | None:
+    if isinstance(value, (int, float)):
+        return value
+    if isinstance(value, str):
+        try:
+            return float(value.strip())
+        except ValueError:
+            return None
+    return None
+
+
+def _apply_error_metadata(target: OpenAIErrorDetail, error: OpenAIError | None) -> None:
+    if not error:
+        return
+    if error.plan_type is not None:
+        target["plan_type"] = error.plan_type
+    if error.resets_at is not None:
+        target["resets_at"] = error.resets_at
+    if error.resets_in_seconds is not None:
+        target["resets_in_seconds"] = error.resets_in_seconds
+
+
+def _upstream_error_from_openai(error: OpenAIError | None) -> UpstreamError:
+    if not error:
+        return {}
+    data = error.model_dump(exclude_none=True)
+    payload: UpstreamError = {}
+    message = data.get("message")
+    if isinstance(message, str):
+        payload["message"] = message
+    resets_at = data.get("resets_at")
+    if isinstance(resets_at, (int, float)):
+        payload["resets_at"] = resets_at
+    resets_in_seconds = data.get("resets_in_seconds")
+    if isinstance(resets_in_seconds, (int, float)):
+        payload["resets_in_seconds"] = resets_in_seconds
+    return payload

app/modules/proxy/load_balancer.py

@@ -12,15 +12,17 @@ from app.core.balancer import (
     select_account,
 )
 from app.core.balancer.types import UpstreamError
-from app.db.models import Account, AccountStatus, UsageHistory
+from app.core.usage.quota import apply_usage_quota
+from app.db.models import Account, UsageHistory
 from app.modules.accounts.repository import AccountsRepository
-from app.modules.proxy.usage_updater import UsageUpdater
 from app.modules.usage.repository import UsageRepository
+from app.modules.usage.updater import UsageUpdater
 
 
 @dataclass
 class RuntimeState:
-    reset_at: int | None = None
+    reset_at: float | None = None
+    cooldown_until: float | None = None
     last_error_at: float | None = None
     last_selected_at: float | None = None
     error_count: int = 0
@@ -100,6 +102,7 @@ class LoadBalancer:
             status=account.status,
             used_percent=None,
             reset_at=runtime.reset_at,
+            cooldown_until=runtime.cooldown_until,
             last_error_at=runtime.last_error_at,
             last_selected_at=runtime.last_selected_at,
             error_count=runtime.error_count,
@@ -109,6 +112,7 @@ class LoadBalancer:
     async def _sync_state(self, account: Account, state: AccountState) -> None:
         runtime = self._runtime.setdefault(account.id, RuntimeState())
         runtime.reset_at = state.reset_at
+        runtime.cooldown_until = state.cooldown_until
         runtime.last_error_at = state.last_error_at
         runtime.error_count = state.error_count
 
@@ -152,12 +156,16 @@ def _state_from_account(
     runtime: RuntimeState,
 ) -> AccountState:
     primary_used = primary_entry.used_percent if primary_entry else None
+    primary_reset = primary_entry.reset_at if primary_entry else None
+    primary_window_minutes = primary_entry.window_minutes if primary_entry else None
     secondary_used = secondary_entry.used_percent if secondary_entry else None
     secondary_reset = secondary_entry.reset_at if secondary_entry else None
 
-    status, used_percent, reset_at = _apply_secondary_quota(
+    status, used_percent, reset_at = apply_usage_quota(
         status=account.status,
         primary_used=primary_used,
+        primary_reset=primary_reset,
+        primary_window_minutes=primary_window_minutes,
         runtime_reset=runtime.reset_at,
         secondary_used=secondary_used,
         secondary_reset=secondary_reset,
@@ -168,41 +176,9 @@ def _state_from_account(
         status=status,
         used_percent=used_percent,
         reset_at=reset_at,
+        cooldown_until=runtime.cooldown_until,
         last_error_at=runtime.last_error_at,
         last_selected_at=runtime.last_selected_at,
         error_count=runtime.error_count,
         deactivation_reason=account.deactivation_reason,
     )
-
-
-def _apply_secondary_quota(
-    *,
-    status: AccountStatus,
-    primary_used: float | None,
-    runtime_reset: int | None,
-    secondary_used: float | None,
-    secondary_reset: int | None,
-) -> tuple[AccountStatus, float | None, int | None]:
-    used_percent = primary_used
-    reset_at = runtime_reset
-
-    if status in (AccountStatus.DEACTIVATED, AccountStatus.PAUSED):
-        return status, used_percent, reset_at
-
-    if secondary_used is None:
-        if status == AccountStatus.QUOTA_EXCEEDED and secondary_reset is not None:
-            reset_at = secondary_reset
-        return status, used_percent, reset_at
-
-    if secondary_used >= 100.0:
-        status = AccountStatus.QUOTA_EXCEEDED
-        used_percent = 100.0
-        if secondary_reset is not None:
-            reset_at = secondary_reset
-        return status, used_percent, reset_at
-
-    if status == AccountStatus.QUOTA_EXCEEDED:
-        status = AccountStatus.ACTIVE
-        reset_at = None
-
-    return status, used_percent, reset_at