codex-lb 0.1.5__py3-none-any.whl → 0.3.0__py3-none-any.whl

app/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.1"
+__version__ = "0.2.0"
 
 from app.main import app as app
 

app/core/auth/__init__.py

@@ -82,14 +82,25 @@ def extract_id_token_claims(id_token: str) -> IdTokenClaims:
 def claims_from_auth(auth: AuthFile) -> AccountClaims:
     claims = extract_id_token_claims(auth.tokens.id_token)
     auth_claims = claims.auth or OpenAIAuthClaims()
+    plan_type = auth_claims.chatgpt_plan_type or claims.chatgpt_plan_type
     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,
+        plan_type=plan_type,
     )
 
 
+def generate_unique_account_id(account_id: str | None, email: str | None) -> str:
+    if account_id and email and email != DEFAULT_EMAIL:
+        email_hash = hashlib.sha256(email.encode()).hexdigest()[:8]
+        return f"{account_id}_{email_hash}"
+    if account_id:
+        return account_id
+    return fallback_account_id(email)
+
+
 def fallback_account_id(email: str | None) -> str:
+    """Generate a fallback account ID when no OpenAI account ID is available."""
     if email and email != DEFAULT_EMAIL:
         digest = hashlib.sha256(email.encode()).hexdigest()[:12]
         return f"email_{digest}"

app/core/balancer/logic.py

@@ -16,6 +16,9 @@ PERMANENT_FAILURE_CODES = {
     "account_deleted": "Account has been deleted",
 }
 
+SECONDS_PER_DAY = 60 * 60 * 24
+UNKNOWN_RESET_BUCKET_DAYS = 10_000
+
 
 @dataclass
 class AccountState:
@@ -23,6 +26,9 @@ class AccountState:
     status: AccountStatus
     used_percent: float | None = None
     reset_at: float | None = None
+    cooldown_until: float | None = None
+    secondary_used_percent: float | None = None
+    secondary_reset_at: int | None = None
     last_error_at: float | None = None
     last_selected_at: float | None = None
     error_count: int = 0
@@ -35,7 +41,12 @@ class SelectionResult:
     error_message: str | None
 
 
-def select_account(states: Iterable[AccountState], now: float | None = None) -> SelectionResult:
+def select_account(
+    states: Iterable[AccountState],
+    now: float | None = None,
+    *,
+    prefer_earlier_reset: bool = False,
+) -> SelectionResult:
     current = now or time.time()
     available: list[AccountState] = []
     all_states = list(states)
@@ -59,6 +70,12 @@ def select_account(states: Iterable[AccountState], now: float | None = None) ->
                 state.reset_at = None
             else:
                 continue
+        if state.cooldown_until and current >= state.cooldown_until:
+            state.cooldown_until = None
+            state.last_error_at = None
+            state.error_count = 0
+        if state.cooldown_until and current < state.cooldown_until:
+            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:
@@ -82,14 +99,29 @@ def select_account(states: Iterable[AccountState], now: float | None = None) ->
             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")
+        cooldowns = [s.cooldown_until for s in all_states if s.cooldown_until and s.cooldown_until > current]
+        if cooldowns:
+            wait_seconds = max(0.0, min(cooldowns) - 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
+    def _usage_sort_key(state: AccountState) -> tuple[float, float, float, str]:
+        primary_used = state.used_percent if state.used_percent is not None else 0.0
+        secondary_used = state.secondary_used_percent if state.secondary_used_percent is not None else primary_used
         last_selected = state.last_selected_at or 0.0
-        return used, last_selected, state.account_id
-
-    selected = min(available, key=_sort_key)
+        return secondary_used, primary_used, last_selected, state.account_id
+
+    def _reset_first_sort_key(state: AccountState) -> tuple[int, float, float, float, str]:
+        reset_bucket_days = UNKNOWN_RESET_BUCKET_DAYS
+        if state.secondary_reset_at is not None:
+            reset_bucket_days = max(
+                0,
+                int((state.secondary_reset_at - current) // SECONDS_PER_DAY),
+            )
+        secondary_used, primary_used, last_selected, account_id = _usage_sort_key(state)
+        return reset_bucket_days, secondary_used, primary_used, last_selected, account_id
+
+    selected = min(available, key=_reset_first_sort_key if prefer_earlier_reset else _usage_sort_key)
     return SelectionResult(selected, None)
 
 
@@ -97,11 +129,16 @@ 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
+
     message = error.get("message")
     delay = parse_retry_after(message) if message else None
     if delay is None:
         delay = backoff_seconds(state.error_count)
-    state.reset_at = time.time() + delay
+    state.cooldown_until = time.time() + delay
 
 
 def handle_quota_exceeded(state: AccountState, error: UpstreamError) -> None:

app/core/clients/proxy.py

@@ -18,7 +18,6 @@ IGNORE_INBOUND_HEADERS = {"authorization", "chatgpt-account-id", "content-length
 
 _ERROR_TYPE_CODE_MAP = {
     "rate_limit_exceeded": "rate_limit_exceeded",
-    "usage_limit_reached": "rate_limit_exceeded",
     "usage_not_included": "usage_not_included",
     "insufficient_quota": "insufficient_quota",
     "quota_exceeded": "quota_exceeded",
@@ -64,12 +63,11 @@ def _normalize_error_code(code: str | None, error_type: str | None) -> str:
     if code:
         normalized_code = code.lower()
         mapped = _ERROR_TYPE_CODE_MAP.get(normalized_code)
-        return mapped or code
+        return mapped or normalized_code
     normalized_type = error_type.lower() if error_type else None
     if normalized_type:
         mapped = _ERROR_TYPE_CODE_MAP.get(normalized_type)
-        if mapped:
-            return mapped
+        return mapped or normalized_type
     return "upstream_error"
 
 

app/core/config/settings.py

@@ -39,6 +39,9 @@ class Settings(BaseSettings):
     usage_refresh_enabled: bool = True
     usage_refresh_interval_seconds: int = 60
     encryption_key_file: Path = DEFAULT_ENCRYPTION_KEY_FILE
+    database_migrations_fail_fast: bool = True
+    log_proxy_request_shape: bool = False
+    log_proxy_request_shape_raw_cache_key: bool = False
 
     @field_validator("database_url")
     @classmethod
@@ -61,7 +64,7 @@ class Settings(BaseSettings):
             return value.expanduser()
         if isinstance(value, str):
             return Path(value).expanduser()
-        return value
+        raise TypeError("encryption_key_file must be a path")
 
 
 @lru_cache(maxsize=1)

app/core/plan_types.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from typing import Final
+
+ACCOUNT_PLAN_TYPES: Final[set[str]] = {
+    "free",
+    "plus",
+    "pro",
+    "team",
+    "business",
+    "enterprise",
+    "edu",
+}
+
+RATE_LIMIT_PLAN_TYPES: Final[set[str]] = {
+    *ACCOUNT_PLAN_TYPES,
+    "guest",
+    "go",
+    "free_workspace",
+    "education",
+    "quorum",
+    "k12",
+}
+
+
+def _clean_plan_type(value: str | None) -> str | None:
+    if value is None:
+        return None
+    cleaned = value.strip()
+    return cleaned or None
+
+
+def normalize_account_plan_type(value: str | None) -> str | None:
+    cleaned = _clean_plan_type(value)
+    if not cleaned:
+        return None
+    normalized = cleaned.lower()
+    return normalized if normalized in ACCOUNT_PLAN_TYPES else None
+
+
+def canonicalize_account_plan_type(value: str | None) -> str | None:
+    cleaned = _clean_plan_type(value)
+    if not cleaned:
+        return None
+    normalized = cleaned.lower()
+    if normalized in ACCOUNT_PLAN_TYPES:
+        return normalized
+    return cleaned
+
+
+def coerce_account_plan_type(value: str | None, default: str) -> str:
+    cleaned = _clean_plan_type(value)
+    if cleaned is None:
+        return default
+    canonical = canonicalize_account_plan_type(cleaned)
+    return canonical if canonical is not None else default
+
+
+def normalize_rate_limit_plan_type(value: str | None) -> str | None:
+    cleaned = _clean_plan_type(value)
+    if not cleaned:
+        return None
+    normalized = cleaned.lower()
+    return normalized if normalized in RATE_LIMIT_PLAN_TYPES else None

app/core/types.py

@@ -1,4 +1,6 @@
 from __future__ import annotations
 
-type JsonValue = bool | int | float | str | None | list[JsonValue] | dict[str, JsonValue]
-type JsonObject = dict[str, JsonValue]
+from collections.abc import Mapping
+
+type JsonValue = bool | int | float | str | None | list[JsonValue] | Mapping[str, JsonValue]
+type JsonObject = Mapping[str, JsonValue]

app/core/usage/__init__.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import Iterable, Mapping
 
+from app.core.plan_types import normalize_account_plan_type
 from app.core.usage.types import (
     UsageCostSummary,
     UsageHistoryPayload,
@@ -16,12 +17,14 @@ from app.db.models import Account
 PLAN_CAPACITY_CREDITS_PRIMARY = {
     "plus": 225.0,
     "business": 225.0,
+    "team": 225.0,
     "pro": 1500.0,
 }
 
 PLAN_CAPACITY_CREDITS_SECONDARY = {
     "plus": 7560.0,
     "business": 7560.0,
+    "team": 7560.0,
     "pro": 50400.0,
 }
 
@@ -134,9 +137,9 @@ def summarize_usage_window(
 
 
 def capacity_for_plan(plan_type: str | None, window: str) -> float | None:
-    if not plan_type:
+    normalized = normalize_account_plan_type(plan_type)
+    if not normalized:
         return None
-    normalized = plan_type.lower()
     window_key = _normalize_window_key(window)
     if window_key == "primary":
         return PLAN_CAPACITY_CREDITS_PRIMARY.get(normalized)

app/core/usage/logs.py

@@ -13,6 +13,17 @@ class RequestLogLike(Protocol):
     reasoning_tokens: int | None
 
 
+def cached_input_tokens_from_log(log: RequestLogLike) -> int | None:
+    cached_tokens = log.cached_input_tokens
+    if cached_tokens is None:
+        return None
+    cached_tokens = max(0, int(cached_tokens))
+    input_tokens = log.input_tokens
+    if input_tokens is not None:
+        cached_tokens = min(cached_tokens, int(input_tokens))
+    return cached_tokens
+
+
 def usage_tokens_from_log(log: RequestLogLike) -> UsageTokens | None:
     input_tokens = log.input_tokens
     if input_tokens is None:
@@ -20,8 +31,7 @@ def usage_tokens_from_log(log: RequestLogLike) -> UsageTokens | None:
     output_tokens = log.output_tokens if log.output_tokens is not None else log.reasoning_tokens
     if output_tokens is None:
         return None
-    cached_tokens = log.cached_input_tokens or 0
-    cached_tokens = max(0, min(cached_tokens, input_tokens))
+    cached_tokens = cached_input_tokens_from_log(log) or 0
     return UsageTokens(
         input_tokens=float(input_tokens),
         output_tokens=float(output_tokens),

app/core/usage/quota.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import time
+
+from app.core import usage as usage_core
+from app.db.models import AccountStatus
+
+
+def apply_usage_quota(
+    *,
+    status: AccountStatus,
+    primary_used: float | None,
+    primary_reset: int | None,
+    primary_window_minutes: int | None,
+    runtime_reset: float | None,
+    secondary_used: float | None,
+    secondary_reset: int | None,
+) -> tuple[AccountStatus, float | None, float | 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 not None:
+        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:
+            if runtime_reset and runtime_reset > time.time():
+                reset_at = runtime_reset
+            else:
+                status = AccountStatus.ACTIVE
+                reset_at = None
+    elif status == AccountStatus.QUOTA_EXCEEDED and secondary_reset is not None:
+        reset_at = secondary_reset
+
+    if primary_used is not None:
+        if primary_used >= 100.0:
+            status = AccountStatus.RATE_LIMITED
+            used_percent = 100.0
+            if primary_reset is not None:
+                reset_at = primary_reset
+            else:
+                reset_at = _fallback_primary_reset(primary_window_minutes) or reset_at
+            return status, used_percent, reset_at
+        if status == AccountStatus.RATE_LIMITED:
+            if runtime_reset and runtime_reset > time.time():
+                reset_at = runtime_reset
+            else:
+                status = AccountStatus.ACTIVE
+                reset_at = None
+
+    return status, used_percent, reset_at
+
+
+def _fallback_primary_reset(primary_window_minutes: int | None) -> float | None:
+    window_minutes = primary_window_minutes or usage_core.default_window_minutes("primary")
+    if not window_minutes:
+        return None
+    return time.time() + float(window_minutes) * 60.0

app/core/usage/types.py

@@ -67,8 +67,9 @@ class UsageCostSummary:
 class UsageMetricsSummary:
     requests_7d: int | None
     tokens_secondary_window: int | None
-    error_rate_7d: float | None
-    top_error: str | None
+    cached_tokens_secondary_window: int | None = None
+    error_rate_7d: float | None = None
+    top_error: str | None = None
 
 
 @dataclass(frozen=True)

app/core/utils/sse.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Mapping
 
-from app.core.types import JsonObject
+from app.core.errors import ResponseFailedEvent
+from app.core.types import JsonValue
 
+type JsonPayload = Mapping[str, JsonValue] | ResponseFailedEvent
 
-def format_sse_event(payload: JsonObject) -> str:
+
+def format_sse_event(payload: JsonPayload) -> str:
     data = json.dumps(payload, ensure_ascii=True, separators=(",", ":"))
     event_type = payload.get("type")
     if isinstance(event_type, str) and event_type:

app/db/migrations/__init__.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Awaitable, Callable, Final
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db.migrations.versions import (
+    add_accounts_chatgpt_account_id,
+    add_accounts_reset_at,
+    add_dashboard_settings,
+    add_request_logs_reasoning_effort,
+    normalize_account_plan_types,
+)
+
+_CREATE_MIGRATIONS_TABLE = """
+CREATE TABLE IF NOT EXISTS schema_migrations (
+    name TEXT PRIMARY KEY,
+    applied_at TEXT NOT NULL
+)
+"""
+
+_INSERT_MIGRATION = """
+INSERT INTO schema_migrations (name, applied_at)
+VALUES (:name, :applied_at)
+ON CONFLICT(name) DO NOTHING
+RETURNING name
+"""
+
+
+@dataclass(frozen=True)
+class Migration:
+    name: str
+    run: Callable[[AsyncSession], Awaitable[None]]
+
+
+MIGRATIONS: Final[tuple[Migration, ...]] = (
+    Migration("001_normalize_account_plan_types", normalize_account_plan_types.run),
+    Migration("002_add_request_logs_reasoning_effort", add_request_logs_reasoning_effort.run),
+    Migration("003_add_accounts_reset_at", add_accounts_reset_at.run),
+    Migration("004_add_accounts_chatgpt_account_id", add_accounts_chatgpt_account_id.run),
+    Migration("005_add_dashboard_settings", add_dashboard_settings.run),
+)
+
+
+async def run_migrations(session: AsyncSession) -> int:
+    await _ensure_schema_migrations(session)
+    applied_count = 0
+    for migration in MIGRATIONS:
+        applied_now = await _apply_migration(session, migration)
+        if applied_now:
+            applied_count += 1
+    return applied_count
+
+
+async def _apply_migration(session: AsyncSession, migration: Migration) -> bool:
+    async with _migration_transaction(session):
+        result = await session.execute(
+            text(_INSERT_MIGRATION),
+            {
+                "name": migration.name,
+                "applied_at": _utcnow_iso(),
+            },
+        )
+        inserted = result.scalar_one_or_none()
+        if inserted is None:
+            return False
+        await migration.run(session)
+    return True
+
+
+async def _ensure_schema_migrations(session: AsyncSession) -> None:
+    async with _migration_transaction(session):
+        await session.execute(text(_CREATE_MIGRATIONS_TABLE))
+
+
+@asynccontextmanager
+async def _migration_transaction(session: AsyncSession):
+    if session.in_transaction():
+        async with session.begin_nested():
+            yield
+    else:
+        async with session.begin():
+            yield
+
+
+def _utcnow_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()

app/db/migrations/versions/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

app/db/migrations/versions/add_accounts_chatgpt_account_id.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+
+async def run(session: AsyncSession) -> None:
+    bind = session.get_bind()
+    dialect = getattr(getattr(bind, "dialect", None), "name", None)
+    if dialect == "sqlite":
+        await _sqlite_add_column_if_missing(session, "accounts", "chatgpt_account_id", "VARCHAR")
+    elif dialect == "postgresql":
+        await session.execute(
+            text("ALTER TABLE accounts ADD COLUMN IF NOT EXISTS chatgpt_account_id VARCHAR"),
+        )
+
+
+async def _sqlite_add_column_if_missing(
+    session: AsyncSession,
+    table: str,
+    column: str,
+    column_type: str,
+) -> None:
+    result = await session.execute(text(f"PRAGMA table_info({table})"))
+    rows = result.fetchall()
+    existing = {row[1] for row in rows if len(row) > 1}
+    if column in existing:
+        return
+    await session.execute(text(f"ALTER TABLE {table} ADD COLUMN {column} {column_type}"))

app/db/migrations/versions/add_accounts_reset_at.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+
+async def run(session: AsyncSession) -> None:
+    bind = session.get_bind()
+    dialect = getattr(getattr(bind, "dialect", None), "name", None)
+    if dialect == "sqlite":
+        await _sqlite_add_column_if_missing(session, "accounts", "reset_at", "INTEGER")
+    elif dialect == "postgresql":
+        await session.execute(
+            text("ALTER TABLE accounts ADD COLUMN IF NOT EXISTS reset_at INTEGER"),
+        )
+
+
+async def _sqlite_add_column_if_missing(
+    session: AsyncSession,
+    table: str,
+    column: str,
+    column_type: str,
+) -> None:
+    result = await session.execute(text(f"PRAGMA table_info({table})"))
+    rows = result.fetchall()
+    existing = {row[1] for row in rows if len(row) > 1}
+    if column in existing:
+        return
+    await session.execute(text(f"ALTER TABLE {table} ADD COLUMN {column} {column_type}"))

app/db/migrations/versions/add_dashboard_settings.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from sqlalchemy import inspect
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import Session
+
+from app.db.models import DashboardSettings
+
+
+def _settings_table_exists(session: Session) -> bool:
+    inspector = inspect(session.connection())
+    return inspector.has_table("dashboard_settings")
+
+
+async def run(session: AsyncSession) -> None:
+    exists = await session.run_sync(_settings_table_exists)
+    if not exists:
+        return
+
+    row = await session.get(DashboardSettings, 1)
+    if row is not None:
+        return
+
+    session.add(
+        DashboardSettings(
+            id=1,
+            sticky_threads_enabled=False,
+            prefer_earlier_reset_accounts=False,
+        )
+    )
+    await session.flush()

app/db/migrations/versions/add_request_logs_reasoning_effort.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from sqlalchemy import inspect, text
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import Session
+
+
+def _request_logs_column_state(session: Session) -> tuple[bool, bool]:
+    conn = session.connection()
+    inspector = inspect(conn)
+    if not inspector.has_table("request_logs"):
+        return False, False
+    columns = {column["name"] for column in inspector.get_columns("request_logs")}
+    return True, "reasoning_effort" in columns
+
+
+async def run(session: AsyncSession) -> None:
+    has_table, has_column = await session.run_sync(_request_logs_column_state)
+    if not has_table or has_column:
+        return
+    await session.execute(text("ALTER TABLE request_logs ADD COLUMN reasoning_effort VARCHAR"))

app/db/migrations/versions/normalize_account_plan_types.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.auth import DEFAULT_PLAN
+from app.core.plan_types import coerce_account_plan_type
+from app.db.models import Account
+
+
+async def run(session: AsyncSession) -> None:
+    result = await session.execute(select(Account))
+    accounts = list(result.scalars().all())
+    for account in accounts:
+        coerced = coerce_account_plan_type(account.plan_type, DEFAULT_PLAN)
+        if account.plan_type != coerced:
+            account.plan_type = coerced