codex-lb 0.1.2__py3-none-any.whl

app/__init__.py

@@ -0,0 +1,5 @@
+__version__ = "0.1.1"
+
+from app.main import app as app
+
+__all__ = ["app", "__version__"]

app/cli.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import argparse
+import os
+
+import uvicorn
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run the codex-lb API server.")
+    parser.add_argument("--host", default=os.getenv("HOST", "127.0.0.1"))
+    parser.add_argument("--port", type=int, default=int(os.getenv("PORT", "2455")))
+
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+
+    uvicorn.run("app.main:app", host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

app/core/auth/__init__.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import base64
+import hashlib
+import json
+from dataclasses import dataclass
+from datetime import datetime
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field
+
+DEFAULT_EMAIL = "unknown@example.com"
+DEFAULT_PLAN = "unknown"
+
+
+class AuthTokens(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, extra="ignore")
+
+    id_token: str = Field(alias="idToken")
+    access_token: str = Field(alias="accessToken")
+    refresh_token: str = Field(alias="refreshToken")
+    account_id: str | None = Field(default=None, alias="accountId")
+
+
+class AuthFile(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, extra="ignore")
+
+    openai_api_key: str | None = Field(default=None, alias="OPENAI_API_KEY")
+    tokens: AuthTokens
+    last_refresh_at: datetime | None = Field(default=None, alias="lastRefreshAt")
+
+
+class OpenAIAuthClaims(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    chatgpt_account_id: str | None = None
+    chatgpt_plan_type: str | None = None
+
+
+class IdTokenClaims(BaseModel):
+    model_config = ConfigDict(extra="ignore", populate_by_name=True)
+
+    email: str | None = None
+    chatgpt_account_id: str | None = None
+    chatgpt_plan_type: str | None = None
+    exp: int | float | str | None = None
+    auth: OpenAIAuthClaims | None = Field(
+        default=None,
+        alias="https://api.openai.com/auth",
+    )
+
+
+@dataclass
+class AccountClaims:
+    account_id: str | None
+    email: str | None
+    plan_type: str | None
+
+
+def parse_auth_json(raw: bytes) -> AuthFile:
+    data = json.loads(raw)
+    model = AuthFile.model_validate(data)
+    return model
+
+
+def extract_id_token_claims(id_token: str) -> IdTokenClaims:
+    try:
+        parts = id_token.split(".")
+        if len(parts) < 2:
+            return IdTokenClaims()
+        payload = parts[1]
+        padding = "=" * (-len(payload) % 4)
+        decoded = base64.urlsafe_b64decode(payload + padding)
+        data = json.loads(decoded)
+        if not isinstance(data, dict):
+            return IdTokenClaims()
+        return IdTokenClaims.model_validate(data)
+    except Exception:
+        return IdTokenClaims()
+
+
+def claims_from_auth(auth: AuthFile) -> AccountClaims:
+    claims = extract_id_token_claims(auth.tokens.id_token)
+    auth_claims = claims.auth or OpenAIAuthClaims()
+    return AccountClaims(
+        account_id=auth.tokens.account_id or auth_claims.chatgpt_account_id or claims.chatgpt_account_id,
+        email=claims.email,
+        plan_type=auth_claims.chatgpt_plan_type or claims.chatgpt_plan_type,
+    )
+
+
+def fallback_account_id(email: str | None) -> str:
+    if email and email != DEFAULT_EMAIL:
+        digest = hashlib.sha256(email.encode()).hexdigest()[:12]
+        return f"email_{digest}"
+    return f"local_{uuid4().hex[:12]}"

app/core/auth/models.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from pydantic import AliasChoices, BaseModel, ConfigDict, Field, StrictStr, field_validator
+
+from app.core.types import JsonObject
+
+
+class OAuthTokenPayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    access_token: StrictStr | None = None
+    refresh_token: StrictStr | None = None
+    id_token: StrictStr | None = None
+    authorization_code: StrictStr | None = None
+    code_verifier: StrictStr | None = None
+    error: JsonObject | StrictStr | None = None
+    error_description: StrictStr | None = None
+    message: StrictStr | None = None
+    error_code: StrictStr | None = None
+    code: StrictStr | None = None
+    status: StrictStr | None = None
+
+
+class DeviceCodePayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    device_auth_id: StrictStr | None = None
+    user_code: StrictStr | None = Field(
+        default=None,
+        validation_alias=AliasChoices("user_code", "usercode"),
+    )
+    interval: int | None = None
+    expires_in: int | None = None
+    expires_at: StrictStr | None = None
+
+    @field_validator("interval", mode="before")
+    @classmethod
+    def _parse_interval(cls, value: object) -> int | None:
+        if value is None:
+            return None
+        if isinstance(value, int):
+            return value
+        if isinstance(value, str):
+            stripped = value.strip()
+            if not stripped:
+                return None
+            if stripped.isdigit():
+                return int(stripped)
+        raise ValueError("Invalid interval")

app/core/auth/refresh.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+
+import aiohttp
+from pydantic import ValidationError
+
+from app.core.auth import OpenAIAuthClaims, extract_id_token_claims
+from app.core.auth.models import OAuthTokenPayload
+from app.core.balancer import PERMANENT_FAILURE_CODES
+from app.core.clients.http import get_http_client
+from app.core.config.settings import get_settings
+from app.core.types import JsonObject
+from app.core.utils.request_id import get_request_id
+from app.core.utils.time import to_utc_naive, utcnow
+
+TOKEN_REFRESH_INTERVAL_DAYS = 8
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class TokenRefreshResult:
+    access_token: str
+    refresh_token: str
+    id_token: str
+    account_id: str | None
+    plan_type: str | None
+    email: str | None
+
+
+class RefreshError(Exception):
+    def __init__(self, code: str, message: str, is_permanent: bool) -> None:
+        super().__init__(message)
+        self.code = code
+        self.message = message
+        self.is_permanent = is_permanent
+
+
+def should_refresh(last_refresh: datetime, now: datetime | None = None) -> bool:
+    current = to_utc_naive(now) if now is not None else utcnow()
+    last = to_utc_naive(last_refresh)
+    interval_days = get_settings().token_refresh_interval_days or TOKEN_REFRESH_INTERVAL_DAYS
+    return current - last > timedelta(days=interval_days)
+
+
+def classify_refresh_error(code: str | None) -> bool:
+    if not code:
+        return False
+    return code in PERMANENT_FAILURE_CODES
+
+
+async def refresh_access_token(
+    refresh_token: str,
+    *,
+    session: aiohttp.ClientSession | None = None,
+) -> TokenRefreshResult:
+    settings = get_settings()
+    url = f"{settings.auth_base_url.rstrip('/')}/oauth/token"
+    payload = {
+        "grant_type": "refresh_token",
+        "client_id": settings.oauth_client_id,
+        "refresh_token": refresh_token,
+        "scope": settings.oauth_scope,
+    }
+    timeout = aiohttp.ClientTimeout(total=settings.token_refresh_timeout_seconds)
+
+    client_session = session or get_http_client().session
+    headers: dict[str, str] = {}
+    request_id = get_request_id()
+    if request_id:
+        headers["x-request-id"] = request_id
+    async with client_session.post(url, json=payload, headers=headers, timeout=timeout) as resp:
+        data = await _safe_json(resp)
+        try:
+            payload_data = OAuthTokenPayload.model_validate(data)
+        except ValidationError as exc:
+            logger.warning(
+                "Token refresh response invalid request_id=%s",
+                get_request_id(),
+            )
+            raise RefreshError("invalid_response", "Refresh response invalid", False) from exc
+        if resp.status >= 400:
+            logger.warning(
+                "Token refresh failed request_id=%s status=%s",
+                get_request_id(),
+                resp.status,
+            )
+            raise _refresh_error_from_payload(payload_data, resp.status)
+
+    if not payload_data.access_token or not payload_data.refresh_token or not payload_data.id_token:
+        raise RefreshError("invalid_response", "Refresh response missing tokens", False)
+
+    claims = extract_id_token_claims(payload_data.id_token)
+    auth_claims = claims.auth or OpenAIAuthClaims()
+    account_id = auth_claims.chatgpt_account_id or claims.chatgpt_account_id
+    plan_type = auth_claims.chatgpt_plan_type or claims.chatgpt_plan_type
+    email = claims.email
+
+    return TokenRefreshResult(
+        access_token=payload_data.access_token,
+        refresh_token=payload_data.refresh_token,
+        id_token=payload_data.id_token,
+        account_id=account_id,
+        plan_type=plan_type,
+        email=email,
+    )
+
+
+async def _safe_json(resp: aiohttp.ClientResponse) -> JsonObject:
+    try:
+        data = await resp.json(content_type=None)
+    except Exception:
+        text = await resp.text()
+        return {"error": {"message": text.strip()}}
+    return data if isinstance(data, dict) else {"error": {"message": str(data)}}
+
+
+def _refresh_error_from_payload(payload: OAuthTokenPayload, status_code: int) -> RefreshError:
+    code = _extract_error_code(payload) or f"http_{status_code}"
+    message = _extract_error_message(payload) or f"Token refresh failed ({status_code})"
+    return RefreshError(code, message, classify_refresh_error(code))
+
+
+def _extract_error_code(payload: OAuthTokenPayload) -> str | None:
+    error = payload.error
+    if isinstance(error, dict):
+        code = error.get("code") or error.get("error")
+        return code if isinstance(code, str) else None
+    if isinstance(error, str):
+        return error
+    return payload.error_code or payload.code
+
+
+def _extract_error_message(payload: OAuthTokenPayload) -> str | None:
+    error = payload.error
+    if isinstance(error, dict):
+        message = error.get("message") or error.get("error_description")
+        return message if isinstance(message, str) else None
+    if isinstance(error, str):
+        return payload.error_description or error
+    return payload.message

app/core/balancer/__init__.py

@@ -0,0 +1,19 @@
+from app.core.balancer.logic import (
+    PERMANENT_FAILURE_CODES,
+    AccountState,
+    SelectionResult,
+    handle_permanent_failure,
+    handle_quota_exceeded,
+    handle_rate_limit,
+    select_account,
+)
+
+__all__ = [
+    "PERMANENT_FAILURE_CODES",
+    "AccountState",
+    "SelectionResult",
+    "handle_permanent_failure",
+    "handle_quota_exceeded",
+    "handle_rate_limit",
+    "select_account",
+]

app/core/balancer/logic.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Iterable
+
+from app.core.balancer.types import UpstreamError
+from app.core.utils.retry import parse_retry_after
+from app.db.models import AccountStatus
+
+PERMANENT_FAILURE_CODES = {
+    "refresh_token_expired": "Refresh token expired - re-login required",
+    "refresh_token_reused": "Refresh token was reused - re-login required",
+    "refresh_token_invalidated": "Refresh token was revoked - re-login required",
+    "account_suspended": "Account has been suspended",
+    "account_deleted": "Account has been deleted",
+}
+
+
+@dataclass
+class AccountState:
+    account_id: str
+    status: AccountStatus
+    used_percent: float | None = None
+    reset_at: int | None = None
+    last_error_at: float | None = None
+    last_selected_at: float | None = None
+    error_count: int = 0
+    deactivation_reason: str | None = None
+
+
+@dataclass
+class SelectionResult:
+    account: AccountState | None
+    error_message: str | None
+
+
+def select_account(states: Iterable[AccountState], now: float | None = None) -> SelectionResult:
+    current = now or time.time()
+    available: list[AccountState] = []
+    all_states = list(states)
+
+    for state in all_states:
+        if state.status == AccountStatus.DEACTIVATED:
+            continue
+        if state.status == AccountStatus.PAUSED:
+            continue
+        if state.status == AccountStatus.RATE_LIMITED:
+            if state.reset_at and current >= state.reset_at:
+                state.status = AccountStatus.ACTIVE
+                state.error_count = 0
+                state.reset_at = None
+            else:
+                continue
+        if state.status == AccountStatus.QUOTA_EXCEEDED:
+            if state.reset_at and current >= state.reset_at:
+                state.status = AccountStatus.ACTIVE
+                state.used_percent = 0.0
+                state.reset_at = None
+            else:
+                continue
+        if state.error_count >= 3:
+            backoff = min(300, 30 * (2 ** (state.error_count - 3)))
+            if state.last_error_at and current - state.last_error_at < backoff:
+                continue
+        available.append(state)
+
+    if not available:
+        deactivated = [s for s in all_states if s.status == AccountStatus.DEACTIVATED]
+        paused = [s for s in all_states if s.status == AccountStatus.PAUSED]
+        rate_limited = [s for s in all_states if s.status == AccountStatus.RATE_LIMITED]
+        quota_exceeded = [s for s in all_states if s.status == AccountStatus.QUOTA_EXCEEDED]
+
+        if paused and deactivated and not rate_limited and not quota_exceeded:
+            return SelectionResult(None, "All accounts are paused or require re-authentication")
+        if paused and not rate_limited and not quota_exceeded:
+            return SelectionResult(None, "All accounts are paused")
+        if deactivated and not rate_limited and not quota_exceeded:
+            return SelectionResult(None, "All accounts require re-authentication")
+        if quota_exceeded:
+            reset_candidates = [s.reset_at for s in quota_exceeded if s.reset_at]
+            if reset_candidates:
+                wait_seconds = max(0, min(reset_candidates) - int(current))
+                return SelectionResult(None, f"Rate limit exceeded. Try again in {wait_seconds:.0f}s")
+        return SelectionResult(None, "No available accounts")
+
+    def _sort_key(state: AccountState) -> tuple[float, float, str]:
+        used = state.used_percent if state.used_percent is not None else 0.0
+        last_selected = state.last_selected_at or 0.0
+        return used, last_selected, state.account_id
+
+    selected = min(available, key=_sort_key)
+    return SelectionResult(selected, None)
+
+
+def handle_rate_limit(state: AccountState, error: UpstreamError) -> None:
+    state.status = AccountStatus.RATE_LIMITED
+    state.error_count += 1
+    state.last_error_at = time.time()
+
+    reset_at = _extract_reset_at(error)
+    if reset_at is not None:
+        state.reset_at = reset_at
+        return
+
+    message = error.get("message")
+    delay = parse_retry_after(message) if message else None
+    if delay:
+        state.reset_at = int(time.time() + delay)
+    else:
+        state.reset_at = int(time.time() + 300)
+
+
+def handle_quota_exceeded(state: AccountState, error: UpstreamError) -> None:
+    state.status = AccountStatus.QUOTA_EXCEEDED
+    state.used_percent = 100.0
+
+    reset_at = _extract_reset_at(error)
+    if reset_at is not None:
+        state.reset_at = reset_at
+    else:
+        state.reset_at = int(time.time() + 3600)
+
+
+def handle_permanent_failure(state: AccountState, error_code: str) -> None:
+    state.status = AccountStatus.DEACTIVATED
+    state.deactivation_reason = PERMANENT_FAILURE_CODES.get(
+        error_code,
+        f"Authentication failed: {error_code}",
+    )
+
+
+def _extract_reset_at(error: UpstreamError) -> int | None:
+    reset_at = error.get("resets_at")
+    if reset_at is not None:
+        return int(reset_at)
+    reset_in = error.get("resets_in_seconds")
+    if reset_in is not None:
+        return int(time.time() + float(reset_in))
+    return None

app/core/balancer/types.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from typing import TypedDict
+
+
+class UpstreamError(TypedDict, total=False):
+    message: str
+    resets_at: int | float
+    resets_in_seconds: int | float

app/core/clients/http.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import aiohttp
+from aiohttp_retry import RetryClient
+
+
+@dataclass(slots=True)
+class HttpClient:
+    session: aiohttp.ClientSession
+    retry_client: RetryClient
+
+
+_http_client: HttpClient | None = None
+
+
+async def init_http_client() -> HttpClient:
+    global _http_client
+    if _http_client is not None:
+        return _http_client
+    session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=None))
+    retry_client = RetryClient(client_session=session, raise_for_status=False)
+    _http_client = HttpClient(session=session, retry_client=retry_client)
+    return _http_client
+
+
+async def close_http_client() -> None:
+    global _http_client
+    if _http_client is None:
+        return
+    await _http_client.retry_client.close()
+    _http_client = None
+
+
+def get_http_client() -> HttpClient:
+    if _http_client is None:
+        raise RuntimeError("HTTP client not initialized")
+    return _http_client