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

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

@@ -6,21 +6,25 @@ from typing import Iterable
 
 from app.core.balancer import (
     AccountState,
+    SelectionResult,
     handle_permanent_failure,
     handle_quota_exceeded,
     handle_rate_limit,
     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.proxy.sticky_repository import StickySessionsRepository
 from app.modules.usage.repository import UsageRepository
+from app.modules.usage.updater import UsageUpdater
 
 
 @dataclass
 class RuntimeState:
     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
@@ -33,13 +37,25 @@ class AccountSelection:
 
 
 class LoadBalancer:
-    def __init__(self, accounts_repo: AccountsRepository, usage_repo: UsageRepository) -> None:
+    def __init__(
+        self,
+        accounts_repo: AccountsRepository,
+        usage_repo: UsageRepository,
+        sticky_repo: StickySessionsRepository | None = None,
+    ) -> None:
         self._accounts_repo = accounts_repo
         self._usage_repo = usage_repo
         self._usage_updater = UsageUpdater(usage_repo, accounts_repo)
+        self._sticky_repo = sticky_repo
         self._runtime: dict[str, RuntimeState] = {}
 
-    async def select_account(self) -> AccountSelection:
+    async def select_account(
+        self,
+        sticky_key: str | None = None,
+        *,
+        reallocate_sticky: bool = False,
+        prefer_earlier_reset_accounts: bool = False,
+    ) -> AccountSelection:
         accounts = await self._accounts_repo.list_accounts()
         latest_primary = await self._usage_repo.latest_by_account()
         await self._usage_updater.refresh_accounts(accounts, latest_primary)
@@ -53,7 +69,13 @@ class LoadBalancer:
             runtime=self._runtime,
         )
 
-        result = select_account(states)
+        result = await self._select_with_stickiness(
+            states=states,
+            account_map=account_map,
+            sticky_key=sticky_key,
+            reallocate_sticky=reallocate_sticky,
+            prefer_earlier_reset_accounts=prefer_earlier_reset_accounts,
+        )
         for state in states:
             account = account_map.get(state.account_id)
             if account:
@@ -72,6 +94,39 @@ class LoadBalancer:
             return AccountSelection(account=None, error_message=result.error_message)
         return AccountSelection(account=selected, error_message=None)
 
+    async def _select_with_stickiness(
+        self,
+        *,
+        states: list[AccountState],
+        account_map: dict[str, Account],
+        sticky_key: str | None,
+        reallocate_sticky: bool,
+        prefer_earlier_reset_accounts: bool,
+    ) -> SelectionResult:
+        if not sticky_key or not self._sticky_repo:
+            return select_account(states, prefer_earlier_reset=prefer_earlier_reset_accounts)
+
+        if reallocate_sticky:
+            chosen = select_account(states, prefer_earlier_reset=prefer_earlier_reset_accounts)
+            if chosen.account is not None and chosen.account.account_id in account_map:
+                await self._sticky_repo.upsert(sticky_key, chosen.account.account_id)
+            return chosen
+
+        existing = await self._sticky_repo.get_account_id(sticky_key)
+        if existing:
+            pinned = next((state for state in states if state.account_id == existing), None)
+            if pinned is None:
+                await self._sticky_repo.delete(sticky_key)
+            else:
+                pinned_result = select_account([pinned], prefer_earlier_reset=prefer_earlier_reset_accounts)
+                if pinned_result.account is not None:
+                    return pinned_result
+
+        chosen = select_account(states, prefer_earlier_reset=prefer_earlier_reset_accounts)
+        if chosen.account is not None and chosen.account.account_id in account_map:
+            await self._sticky_repo.upsert(sticky_key, chosen.account.account_id)
+        return chosen
+
     async def mark_rate_limit(self, account: Account, error: UpstreamError) -> None:
         state = self._state_for(account)
         handle_rate_limit(state, error)
@@ -100,6 +155,9 @@ class LoadBalancer:
             status=account.status,
             used_percent=None,
             reset_at=runtime.reset_at,
+            cooldown_until=runtime.cooldown_until,
+            secondary_used_percent=None,
+            secondary_reset_at=None,
             last_error_at=runtime.last_error_at,
             last_selected_at=runtime.last_selected_at,
             error_count=runtime.error_count,
@@ -109,17 +167,25 @@ 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
 
-        if account.status != state.status or account.deactivation_reason != state.deactivation_reason:
+        reset_at_int = int(state.reset_at) if state.reset_at else None
+        status_changed = account.status != state.status
+        reason_changed = account.deactivation_reason != state.deactivation_reason
+        reset_changed = account.reset_at != reset_at_int
+
+        if status_changed or reason_changed or reset_changed:
             await self._accounts_repo.update_status(
                 account.id,
                 state.status,
                 state.deactivation_reason,
+                reset_at_int,
             )
             account.status = state.status
             account.deactivation_reason = state.deactivation_reason
+            account.reset_at = reset_at_int
 
 
 def _build_states(
@@ -152,13 +218,22 @@ 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(
+    # Use account.reset_at from DB as the authoritative source for runtime reset
+    # This survives across requests since LoadBalancer is instantiated per-request
+    db_reset_at = float(account.reset_at) if account.reset_at else None
+    effective_runtime_reset = db_reset_at or runtime.reset_at
+
+    status, used_percent, reset_at = apply_usage_quota(
         status=account.status,
         primary_used=primary_used,
-        runtime_reset=runtime.reset_at,
+        primary_reset=primary_reset,
+        primary_window_minutes=primary_window_minutes,
+        runtime_reset=effective_runtime_reset,
         secondary_used=secondary_used,
         secondary_reset=secondary_reset,
     )
@@ -168,41 +243,11 @@ def _state_from_account(
         status=status,
         used_percent=used_percent,
         reset_at=reset_at,
+        cooldown_until=runtime.cooldown_until,
+        secondary_used_percent=secondary_used,
+        secondary_reset_at=secondary_reset,
         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: 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 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