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

app/modules/proxy/service.py

@@ -3,9 +3,7 @@ from __future__ import annotations
 import logging
 import time
 from datetime import timedelta
-from typing import AsyncIterator, Iterable, Mapping
-
-from pydantic import ValidationError
+from typing import AsyncIterator, Mapping
 
 from app.core import usage as usage_core
 from app.core.auth.refresh import RefreshError
@@ -15,27 +13,37 @@ from app.core.clients.proxy import ProxyResponseError, filter_inbound_headers
 from app.core.clients.proxy import compact_responses as core_compact_responses
 from app.core.clients.proxy import stream_responses as core_stream_responses
 from app.core.crypto import TokenEncryptor
-from app.core.errors import OpenAIErrorDetail, OpenAIErrorEnvelope, openai_error, response_failed_event
-from app.core.openai.models import OpenAIError, OpenAIResponsePayload
+from app.core.errors import openai_error, response_failed_event
+from app.core.openai.models import OpenAIResponsePayload
 from app.core.openai.parsing import parse_sse_event
 from app.core.openai.requests import ResponsesCompactRequest, ResponsesRequest
-from app.core.usage.types import UsageWindowRow, UsageWindowSummary
+from app.core.usage.types import UsageWindowRow
 from app.core.utils.request_id import ensure_request_id
 from app.core.utils.sse import format_sse_event
 from app.core.utils.time import utcnow
-from app.db.models import Account, AccountStatus, UsageHistory
+from app.db.models import Account, UsageHistory
+from app.modules.accounts.auth_manager import AuthManager
 from app.modules.accounts.repository import AccountsRepository
-from app.modules.proxy.auth_manager import AuthManager
-from app.modules.proxy.load_balancer import LoadBalancer
-from app.modules.proxy.types import (
-    CreditStatusDetailsData,
-    RateLimitStatusDetailsData,
-    RateLimitStatusPayloadData,
-    RateLimitWindowSnapshotData,
+from app.modules.proxy.helpers import (
+    _apply_error_metadata,
+    _credits_headers,
+    _credits_snapshot,
+    _header_account_id,
+    _normalize_error_code,
+    _parse_openai_error,
+    _plan_type_for_accounts,
+    _rate_limit_details,
+    _rate_limit_headers,
+    _select_accounts_for_limits,
+    _summarize_window,
+    _upstream_error_from_openai,
+    _window_snapshot,
 )
-from app.modules.proxy.usage_updater import UsageUpdater
+from app.modules.proxy.load_balancer import LoadBalancer
+from app.modules.proxy.types import RateLimitStatusPayloadData
 from app.modules.request_logs.repository import RequestLogsRepository
 from app.modules.usage.repository import UsageRepository
+from app.modules.usage.updater import UsageUpdater
 
 logger = logging.getLogger(__name__)
 
@@ -304,7 +312,11 @@ class ProxyService:
                 return
             event = parse_sse_event(first)
             if event and event.type in ("response.failed", "error"):
-                error = event.response.error if event.type == "response.failed" else event.error
+                if event.type == "response.failed":
+                    response = event.response
+                    error = response.error if response else None
+                else:
+                    error = event.error
                 code = _normalize_error_code(
                     error.code if error else None,
                     error.type if error else None,
@@ -326,7 +338,11 @@ class ProxyService:
                     event_type = event.type
                     if event_type in ("response.failed", "error"):
                         status = "error"
-                        error = event.response.error if event_type == "response.failed" else event.error
+                        if event_type == "response.failed":
+                            response = event.response
+                            error = response.error if response else None
+                        else:
+                            error = event.error
                         error_code = _normalize_error_code(
                             error.code if error else None,
                             error.type if error else None,
@@ -420,9 +436,12 @@ class ProxyService:
         await self._handle_stream_error(account, _upstream_error_from_openai(error), code)
 
     async def _handle_stream_error(self, account: Account, error: UpstreamError, code: str) -> None:
-        if code in {"rate_limit_exceeded", "usage_limit_reached"}:
+        if code == "rate_limit_exceeded":
             await self._load_balancer.mark_rate_limit(account, error)
             return
+        if code == "usage_limit_reached":
+            await self._load_balancer.mark_quota_exceeded(account, error)
+            return
         if code in {"insufficient_quota", "usage_not_included", "quota_exceeded"}:
             await self._load_balancer.mark_quota_exceeded(account, error)
             return
@@ -432,297 +451,8 @@ class ProxyService:
         await self._load_balancer.record_error(account)
 
 
-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
-
-
-KNOWN_PLAN_TYPES = {
-    "guest",
-    "free",
-    "go",
-    "plus",
-    "pro",
-    "free_workspace",
-    "team",
-    "business",
-    "education",
-    "quorum",
-    "k12",
-    "enterprise",
-    "edu",
-}
-
-PLAN_TYPE_PRIORITY = (
-    "enterprise",
-    "business",
-    "team",
-    "pro",
-    "plus",
-    "education",
-    "edu",
-    "free_workspace",
-    "free",
-    "go",
-    "guest",
-    "quorum",
-    "k12",
-)
-
-
-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:
-    if not value:
-        return None
-    normalized = value.strip().lower()
-    if normalized not in KNOWN_PLAN_TYPES:
-        return None
-    return normalized
-
-
-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
-
-
 class _RetryableStreamError(Exception):
     def __init__(self, code: str, error: UpstreamError) -> None:
         super().__init__(code)
         self.code = code
         self.error = error
-
-
-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/request_logs/service.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 from datetime import datetime
+from typing import cast
 
-from app.core.usage.logs import cost_from_log, total_tokens_from_log
+from app.core.usage.logs import RequestLogLike, 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
@@ -63,6 +64,7 @@ def _log_status(log: RequestLog) -> str:
 
 
 def _to_entry(log: RequestLog) -> RequestLogEntry:
+    log_like = cast(RequestLogLike, log)
     return RequestLogEntry(
         requested_at=log.requested_at,
         account_id=log.account_id,
@@ -71,7 +73,7 @@ def _to_entry(log: RequestLog) -> RequestLogEntry:
         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),
+        tokens=total_tokens_from_log(log_like),
+        cost_usd=cost_from_log(log_like, precision=6),
         latency_ms=log.latency_ms,
     )

app/modules/usage/service.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 from datetime import timedelta
+from typing import cast
 
 from app.core import usage as usage_core
-from app.core.usage.logs import cost_from_log, total_tokens_from_log, usage_tokens_from_log
+from app.core.usage.logs import RequestLogLike, cost_from_log, total_tokens_from_log, usage_tokens_from_log
 from app.core.usage.pricing import CostItem, calculate_costs
 from app.core.usage.types import (
     UsageCostSummary,
@@ -15,7 +16,6 @@ from app.core.usage.types import (
 from app.core.utils.time import from_epoch_seconds, utcnow
 from app.db.models import Account, RequestLog
 from app.modules.accounts.repository import AccountsRepository
-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.schemas import (
@@ -28,6 +28,7 @@ from app.modules.usage.schemas import (
     UsageWindow,
     UsageWindowResponse,
 )
+from app.modules.usage.updater import UsageUpdater
 
 
 class UsageService:
@@ -137,7 +138,7 @@ def _build_account_history(
     for log in logs:
         account_id = log.account_id
         counts[account_id] = counts.get(account_id, 0) + 1
-        cost = cost_from_log(log)
+        cost = cost_from_log(cast(RequestLogLike, log))
         if cost is None:
             continue
         costs[account_id] = costs.get(account_id, 0.0) + cost
@@ -166,7 +167,7 @@ def _build_account_history(
 
 def _log_to_cost_item(log: RequestLog) -> CostItem | None:
     model = log.model
-    usage = usage_tokens_from_log(log)
+    usage = usage_tokens_from_log(cast(RequestLogLike, log))
     if not model or not usage:
         return None
     return CostItem(model=model, usage=usage)
@@ -191,7 +192,7 @@ def _usage_metrics(logs_secondary: list[RequestLog]) -> UsageMetricsSummary:
 def _sum_tokens(logs: list[RequestLog]) -> int:
     total = 0
     for log in logs:
-        total += total_tokens_from_log(log) or 0
+        total += total_tokens_from_log(cast(RequestLogLike, log)) or 0
     return total
 
 
@@ -232,7 +233,7 @@ def _window_snapshot_to_model(snapshot: UsageWindowSnapshot) -> UsageWindow:
 def _cost_summary_to_model(cost: UsageCostSummary) -> UsageCost:
     return UsageCost(
         currency=cost.currency,
-        total_usd_7d=cost.total_usd_7d,
+        totalUsd7d=cost.total_usd_7d,
         by_model=[UsageCostByModel(model=item.model, usd=item.usd) for item in cost.by_model],
     )
 

app/modules/{proxy/usage_updater.py → usage/updater.py}

@@ -12,8 +12,8 @@ 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.auth_manager import AuthManager
 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__)

app/static/index.js

@@ -74,11 +74,15 @@
 		error: "deactivated",
 	};
 
-	const PLAN_LABELS = {
-		plus: "Plus",
-		team: "Team",
-		free: "Free",
-	};
+	const KNOWN_PLAN_TYPES = new Set([
+		"free",
+		"plus",
+		"pro",
+		"team",
+		"business",
+		"enterprise",
+		"edu",
+	]);
 
 	const ROUTING_LABELS = {
 		usage_weighted: "usage weighted",
@@ -92,7 +96,7 @@
 		timeout: "timeout",
 		upstream: "upstream",
 		rate_limit_exceeded: "rate limit",
-		usage_limit_reached: "rate limit",
+		usage_limit_reached: "quota",
 		insufficient_quota: "quota",
 		usage_not_included: "quota",
 		quota_exceeded: "quota",
@@ -444,7 +448,19 @@
 		REQUEST_STATUS_LABELS[status] || "Unknown";
 	const requestStatusClass = (status) =>
 		REQUEST_STATUS_CLASSES[status] || "deactivated";
-	const planLabel = (plan) => PLAN_LABELS[plan] || "Unknown";
+	const normalizePlanType = (plan) => {
+		if (plan === null || plan === undefined) {
+			return null;
+		}
+		const value = String(plan).trim().toLowerCase();
+		return KNOWN_PLAN_TYPES.has(value) ? value : null;
+	};
+	const titleCase = (value) =>
+		value ? value.charAt(0).toUpperCase() + value.slice(1).toLowerCase() : "";
+	const planLabel = (plan) => {
+		const normalized = normalizePlanType(plan);
+		return normalized ? titleCase(normalized) : "Unknown";
+	};
 	const routingLabel = (strategy) => ROUTING_LABELS[strategy] || "unknown";
 	const errorLabel = (code) => ERROR_LABELS[code] || "--";
 	const progressClass = (status) => PROGRESS_CLASS_BY_STATUS[status] || "";
@@ -569,16 +585,9 @@
 			return acc;
 		}, {});
 
-	const mergeUsageIntoAccounts = (
-		accounts,
-		primaryUsage,
-		secondaryUsage,
-		summary,
-	) => {
+	const mergeUsageIntoAccounts = (accounts, primaryUsage, secondaryUsage) => {
 		const primaryMap = buildUsageIndex(primaryUsage || []);
 		const secondaryMap = buildUsageIndex(secondaryUsage || []);
-		const resetAtPrimary = summary?.primaryWindow?.resetAt ?? null;
-		const resetAtSecondary = summary?.secondaryWindow?.resetAt ?? null;
 		return accounts.map((account) => {
 			const primaryRow = primaryMap[account.id];
 			const secondaryRow = secondaryMap[account.id];
@@ -598,8 +607,8 @@
 						account.usage?.secondaryRemainingPercent ??
 						0,
 				},
-				resetAtPrimary: resetAtPrimary ?? account.resetAtPrimary ?? null,
-				resetAtSecondary: resetAtSecondary ?? account.resetAtSecondary ?? null,
+				resetAtPrimary: account.resetAtPrimary ?? null,
+				resetAtSecondary: account.resetAtSecondary ?? null,
 			};
 		});
 	};
@@ -1191,7 +1200,6 @@
 						accountsResult.value,
 						primaryUsage,
 						secondaryUsage,
-						summary,
 					);
 					this.applyData(
 						{

{codex_lb-0.1.4.dist-info → codex_lb-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: codex-lb
-Version: 0.1.4
+Version: 0.2.0
 Summary: Codex load balancer and proxy for ChatGPT accounts with usage dashboard
 Author-email: Soju06 <qlskssk@gmail.com>
 Maintainer-email: Soju06 <qlskssk@gmail.com>