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

app/core/balancer/logic.py

@@ -5,7 +5,7 @@ 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.core.utils.retry import backoff_seconds, parse_retry_after
 from app.db.models import AccountStatus
 
 PERMANENT_FAILURE_CODES = {
@@ -22,7 +22,7 @@ class AccountState:
     account_id: str
     status: AccountStatus
     used_percent: float | None = None
-    reset_at: int | None = None
+    reset_at: float | None = None
     last_error_at: float | None = None
     last_selected_at: float | None = None
     error_count: int = 0
@@ -97,18 +97,11 @@ 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)
+    if delay is None:
+        delay = backoff_seconds(state.error_count)
+    state.reset_at = time.time() + delay
 
 
 def handle_quota_exceeded(state: AccountState, error: UpstreamError) -> None:

app/core/utils/retry.py

@@ -1,8 +1,13 @@
 from __future__ import annotations
 
+import random
 import re
 
 _RETRY_PATTERN = re.compile(r"(?i)try again in\s*(\d+(?:\.\d+)?)\s*(s|ms|seconds?)")
+_BACKOFF_INITIAL_DELAY_MS = 200
+_BACKOFF_FACTOR = 2.0
+_BACKOFF_JITTER_MIN = 0.9
+_BACKOFF_JITTER_MAX = 1.1
 
 
 def parse_retry_after(message: str) -> float | None:
@@ -14,3 +19,12 @@ def parse_retry_after(message: str) -> float | None:
     if unit == "ms":
         return value / 1000
     return value
+
+
+def backoff_seconds(attempt: int) -> float:
+    if attempt < 1:
+        attempt = 1
+    exponent = _BACKOFF_FACTOR ** (attempt - 1)
+    base_ms = _BACKOFF_INITIAL_DELAY_MS * exponent
+    jitter = random.uniform(_BACKOFF_JITTER_MIN, _BACKOFF_JITTER_MAX)
+    return (base_ms * jitter) / 1000.0

app/db/session.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from pathlib import Path
 from typing import AsyncIterator
 
@@ -23,16 +24,33 @@ def _ensure_sqlite_dir(url: str) -> None:
     Path(path).expanduser().parent.mkdir(parents=True, exist_ok=True)
 
 
+async def _safe_rollback(session: AsyncSession) -> None:
+    if not session.in_transaction():
+        return
+    try:
+        await asyncio.shield(session.rollback())
+    except Exception:
+        return
+
+
+async def _safe_close(session: AsyncSession) -> None:
+    try:
+        await asyncio.shield(session.close())
+    except Exception:
+        return
+
+
 async def get_session() -> AsyncIterator[AsyncSession]:
-    async with SessionLocal() as session:
-        try:
-            yield session
-        except Exception:
-            await session.rollback()
-            raise
-        finally:
-            if session.in_transaction():
-                await session.rollback()
+    session = SessionLocal()
+    try:
+        yield session
+    except BaseException:
+        await _safe_rollback(session)
+        raise
+    finally:
+        if session.in_transaction():
+            await _safe_rollback(session)
+        await _safe_close(session)
 
 
 async def init_db() -> None:

app/dependencies.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass
 from fastapi import Depends
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.db.session import SessionLocal, get_session
+from app.db.session import SessionLocal, _safe_close, _safe_rollback, get_session
 from app.modules.accounts.repository import AccountsRepository
 from app.modules.accounts.service import AccountsService
 from app.modules.oauth.service import OauthService
@@ -87,15 +87,16 @@ def get_usage_context(
 
 @asynccontextmanager
 async def _accounts_repo_context() -> AsyncIterator[AccountsRepository]:
-    async with SessionLocal() as session:
-        try:
-            yield AccountsRepository(session)
-        except Exception:
-            await session.rollback()
-            raise
-        finally:
-            if session.in_transaction():
-                await session.rollback()
+    session = SessionLocal()
+    try:
+        yield AccountsRepository(session)
+    except BaseException:
+        await _safe_rollback(session)
+        raise
+    finally:
+        if session.in_transaction():
+            await _safe_rollback(session)
+        await _safe_close(session)
 
 
 def get_oauth_context(

app/modules/proxy/load_balancer.py

@@ -20,7 +20,7 @@ from app.modules.usage.repository import UsageRepository
 
 @dataclass
 class RuntimeState:
-    reset_at: int | None = None
+    reset_at: float | None = None
     last_error_at: float | None = None
     last_selected_at: float | None = None
     error_count: int = 0
@@ -179,10 +179,10 @@ def _apply_secondary_quota(
     *,
     status: AccountStatus,
     primary_used: float | None,
-    runtime_reset: int | None,
+    runtime_reset: float | None,
     secondary_used: float | None,
     secondary_reset: int | None,
-) -> tuple[AccountStatus, float | None, int | None]:
+) -> tuple[AccountStatus, float | None, float | None]:
     used_percent = primary_used
     reset_at = runtime_reset
 

app/modules/proxy/service.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import time
 from datetime import timedelta
 from typing import AsyncIterator, Iterable, Mapping
@@ -36,6 +37,8 @@ from app.modules.proxy.usage_updater import UsageUpdater
 from app.modules.request_logs.repository import RequestLogsRepository
 from app.modules.usage.repository import UsageRepository
 
+logger = logging.getLogger(__name__)
+
 
 class ProxyService:
     def __init__(
@@ -191,6 +194,7 @@ class ProxyService:
                 yield format_sse_event(event)
                 return
 
+            account_id_value = account.id
             try:
                 account = await self._ensure_fresh(account)
                 async for line in self._stream_once(
@@ -243,7 +247,15 @@ class ProxyService:
                     await self._load_balancer.mark_permanent_failure(account, exc.code)
                 continue
             except Exception:
-                await self._load_balancer.record_error(account)
+                try:
+                    await self._load_balancer.record_error(account)
+                except Exception:
+                    logger.warning(
+                        "Failed to record proxy error account_id=%s request_id=%s",
+                        account_id_value,
+                        request_id,
+                        exc_info=True,
+                    )
                 if attempt == max_attempts - 1:
                     event = response_failed_event(
                         "upstream_error",
@@ -267,8 +279,9 @@ class ProxyService:
         request_id: str,
         allow_retry: bool,
     ) -> AsyncIterator[str]:
+        account_id_value = account.id
         access_token = self._encryptor.decrypt(account.access_token_encrypted)
-        account_id = _header_account_id(account.id)
+        account_id = _header_account_id(account_id_value)
         model = payload.model
         start = time.monotonic()
         status = "success"
@@ -341,19 +354,27 @@ class ProxyService:
             reasoning_tokens = (
                 usage.output_tokens_details.reasoning_tokens if usage and usage.output_tokens_details else None
             )
-            await self._logs_repo.add_log(
-                account_id=account.id,
-                request_id=request_id,
-                model=model,
-                input_tokens=input_tokens,
-                output_tokens=output_tokens,
-                cached_input_tokens=cached_input_tokens,
-                reasoning_tokens=reasoning_tokens,
-                latency_ms=latency_ms,
-                status=status,
-                error_code=error_code,
-                error_message=error_message,
-            )
+            try:
+                await self._logs_repo.add_log(
+                    account_id=account_id_value,
+                    request_id=request_id,
+                    model=model,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                    cached_input_tokens=cached_input_tokens,
+                    reasoning_tokens=reasoning_tokens,
+                    latency_ms=latency_ms,
+                    status=status,
+                    error_code=error_code,
+                    error_message=error_message,
+                )
+            except Exception:
+                logger.warning(
+                    "Failed to persist request log account_id=%s request_id=%s",
+                    account_id_value,
+                    request_id,
+                    exc_info=True,
+                )
 
     async def _refresh_usage(self, accounts: list[Account]) -> None:
         latest_usage = await self._usage_repo.latest_by_account(window="primary")

app/modules/request_logs/repository.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from datetime import datetime
 
 from sqlalchemy import and_, select
@@ -49,9 +50,13 @@ class RequestLogsRepository:
             requested_at=requested_at or utcnow(),
         )
         self._session.add(log)
-        await self._session.commit()
-        await self._session.refresh(log)
-        return log
+        try:
+            await self._session.commit()
+            await self._session.refresh(log)
+            return log
+        except BaseException:
+            await _safe_rollback(self._session)
+            raise
 
     async def list_recent(
         self,
@@ -84,3 +89,12 @@ class RequestLogsRepository:
             stmt = stmt.limit(limit)
         result = await self._session.execute(stmt)
         return list(result.scalars().all())
+
+
+async def _safe_rollback(session: AsyncSession) -> None:
+    if not session.in_transaction():
+        return
+    try:
+        await asyncio.shield(session.rollback())
+    except Exception:
+        return

app/static/index.js

@@ -569,16 +569,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 +591,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 +1184,6 @@
 						accountsResult.value,
 						primaryUsage,
 						secondaryUsage,
-						summary,
 					);
 					this.applyData(
 						{

{codex_lb-0.1.3.dist-info → codex_lb-0.1.5.dist-info}/METADATA

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

{codex_lb-0.1.3.dist-info → codex_lb-0.1.5.dist-info}/RECORD

@@ -1,6 +1,6 @@
 app/__init__.py,sha256=X_4Y93UbuflRBRZd65MtjsVxlZW8d9ZUyi0n83Q-Lrg,89
 app/cli.py,sha256=gkIAkYOT9SbQjUDnVmwhVKZeKjL3YJCMrOjFINwBx54,544
-app/dependencies.py,sha256=43XzEFcOJbfw3Wb1O-fGyLemfgu9RmYPrULUUTPst54,3933
+app/dependencies.py,sha256=9VvBKrUctmuab793Yq2MsHmAh7upXt0viTeofdop8lA,3966
 app/main.py,sha256=JMvTfBdS8wTDNqjQ5QSdqKjjhE0CkpfkbFW6fWNMpy0,4351
 app/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 app/core/crypto.py,sha256=zUz2GVqXigzXB5zX2Diq2cR41Kl37bBqxJiZmWGiIcY,1145
@@ -10,7 +10,7 @@ app/core/auth/__init__.py,sha256=wR52dbGd66W1haAOUXn5RMzn3bj-EeXm4vm7zqI3kHE,281
 app/core/auth/models.py,sha256=iyygu0uHK7fu8txyIyCkXmlgqYPbqEZu3IFysD1t-gk,1557
 app/core/auth/refresh.py,sha256=jhYxT2mQFX4CCvuBQfYNcI53iOtWa0vSkTuZAfbpd9k,5105
 app/core/balancer/__init__.py,sha256=bGy7gITRPObt6P9NdG3pg8t3qWoSW0XWoADqfNYrhdg,405
-app/core/balancer/logic.py,sha256=JFQjzYFzTFRl2Ps24Me-KFZ8jVLMe3Jj5Xzm3aJzFs4,5168
+app/core/balancer/logic.py,sha256=aBuluReYRaYRmXNBntyUhOy1oPsK6Ol4Yt6WSkBLzn4,5059
 app/core/balancer/types.py,sha256=gDgjlTy-NH3YhHYl2-YYpIabnckN9Q8-4cRy6S1u0K4,191
 app/core/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 app/core/clients/http.py,sha256=yfFwsaIZNbSmNtyV02WnMKn00JdoNNSlSIx5xB3QY4o,987
@@ -30,12 +30,12 @@ app/core/usage/pricing.py,sha256=6p8rJ26Gk61mz2t_h9sa0T7NiPDUTiNpzoDewMzT6E0,546
 app/core/usage/types.py,sha256=CbFF6JYSLvALa2P0qYuj9J9b_w53Fgjg0JOu4RzMWbI,2104
 app/core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 app/core/utils/request_id.py,sha256=gxafI-Se8dRQTir-HNGelMzC9S2gwYZntTkVzEqzp_I,705
-app/core/utils/retry.py,sha256=E1s397Bse072sdeEZ2uri5sPAOlTIdyccB-F9jh_E64,408
+app/core/utils/retry.py,sha256=UmBap1Wh-CBT7r4fHzVb_PI9-LR9-HjUtDzRnhRjP2U,822
 app/core/utils/sse.py,sha256=ER3_7mm50BUC1CPIKoVOmQikz-vlEWKNSRDEVqzI_KA,387
 app/core/utils/time.py,sha256=B6FfSe43Eq_puE6eourly1X3gajyihK2VOAwJ8M3wyI,497
 app/db/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 app/db/models.py,sha256=JCknQCuzjHfgyuSzjHqBmeIE-0XisIAQGEhEWqmzabs,3841
-app/db/session.py,sha256=9UBO4fZKWbN3ihbhNxlkFDDXBMRnlV4M0INvatAeyn0,1211
+app/db/session.py,sha256=te9-diH008J-YUAnmE7-OCBUII4LB2jgLlV_ipgoej4,1584
 app/modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 app/modules/accounts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 app/modules/accounts/api.py,sha256=rSkYrg3_p_JE1kHZ4xXa5lBFTdT40v1oNAmlTDGOguY,2807
@@ -52,14 +52,14 @@ app/modules/oauth/templates/oauth_success.html,sha256=YNSGUIozcZEJQjpFtM2sgF4n8j
 app/modules/proxy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 app/modules/proxy/api.py,sha256=BR_qNlNZg2Ft5GYQ-7AzlLlJFf6RVJmlzJzFr_KHUIc,2921
 app/modules/proxy/auth_manager.py,sha256=sFcbStd-OK0g8S56z7XcSkFJ22kEirBU0fYdz2drS0o,2185
-app/modules/proxy/load_balancer.py,sha256=k53RXn9v9fQPXEJ3zWbkhh_Jd82PK1tUef3ZvCdSzog,7441
+app/modules/proxy/load_balancer.py,sha256=L1a7duA4mcR61huXKU052QNz4JAGgUwhMuOPpv5BV4w,7447
 app/modules/proxy/schemas.py,sha256=55pXtUCl2R_93kAPOJJ7Ji4Jn3qVu10vq2KSCCkNdp4,2748
-app/modules/proxy/service.py,sha256=P5TyqZIo63Yz70LGAB5Vb5Um9Wy7Toj9qEyHJ39Qero,26018
+app/modules/proxy/service.py,sha256=XBpL1OIZuzss8nvsYajC5KsaRdGztjci5RgnWORZ1ss,26816
 app/modules/proxy/types.py,sha256=iqEyoO8vGr8N5oEzUSvVWCai7UZbJAU62IvO7JNS9qs,927
 app/modules/proxy/usage_updater.py,sha256=t5wpn3SmPnTXceiQFF0tw-9V5oPPhP4C0anwbI7T5v0,5588
 app/modules/request_logs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 app/modules/request_logs/api.py,sha256=6nV2Uv_hnK7WI3gNpKrgTx4MyUQIXk1QxKp40nij0Xo,1037
-app/modules/request_logs/repository.py,sha256=S86lCDNVAb4N4SvZs5Cnj96ypGzV39r3syS649iOoQ8,2991
+app/modules/request_logs/repository.py,sha256=SnwDwD84wgslmdIgYFGD8ReUYmIyhxT6MOjFAolLBC8,3333
 app/modules/request_logs/schemas.py,sha256=GSCi4TEWMmQ-THsQl2irRrA_msU8jzsqKSBDFn7hiJU,592
 app/modules/request_logs/service.py,sha256=SuJeeFJy0qhMe6cUgQ3-zPrENoPlcNnDmTvG7aYsLmE,2386
 app/modules/shared/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -72,9 +72,9 @@ app/modules/usage/service.py,sha256=_LQA3-L1FxIpzPoIQ-rxz3n3KjiHpTSxRimZapzwx3o,
 app/static/7.css,sha256=9EHW2Ouff2fRXgcQbYuCuglxTFgQZGNLLTXKTS6S5aI,80687
 app/static/index.css,sha256=ct1FfBg0PY7c6KxSS6A7JM_WDwdeJqhQa4pXTgyJ19Q,8786
 app/static/index.html,sha256=FmWQ0l40rTBsHIR9i0ttGTzG17ePASsO2VDv6Q48CNw,24477
-app/static/index.js,sha256=L6Gie0iSAHyrPLvO_QtzEis4rxhVqNe2m-g4h3lB6GU,52419
-codex_lb-0.1.3.dist-info/METADATA,sha256=f7L6g4cETQ9XBxXs-pzJKK1hdnuF9FKIhe6kZTLDnM8,3828
-codex_lb-0.1.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-codex_lb-0.1.3.dist-info/entry_points.txt,sha256=SEa5T6Uz2Fhy574No6Y0XyGmYi3PXLrhu2xStJTqyI8,42
-codex_lb-0.1.3.dist-info/licenses/LICENSE,sha256=cHPibxiL0TXwrUX_kNY6ym544EX1UCzKhxdaca5cFuk,1062
-codex_lb-0.1.3.dist-info/RECORD,,
+app/static/index.js,sha256=2gbVWp_Ot00NLdPpLE2Cc-YWk38H04x8TZw3OFO2SKk,52209
+codex_lb-0.1.5.dist-info/METADATA,sha256=Cjw3NQBgM5hFgPAQUZ81PFBOI2mAq8SiwViALHOkxxA,3828
+codex_lb-0.1.5.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+codex_lb-0.1.5.dist-info/entry_points.txt,sha256=SEa5T6Uz2Fhy574No6Y0XyGmYi3PXLrhu2xStJTqyI8,42
+codex_lb-0.1.5.dist-info/licenses/LICENSE,sha256=cHPibxiL0TXwrUX_kNY6ym544EX1UCzKhxdaca5cFuk,1062
+codex_lb-0.1.5.dist-info/RECORD,,