codex-lb 0.1.5__tar.gz → 0.2.0__tar.gz

codex_lb-0.2.0/.github/release-please-manifest.json

@@ -0,0 +1,3 @@
+{
+  ".": "0.2.0"
+}

codex_lb-0.2.0/.pre-commit-config.yaml

@@ -0,0 +1,14 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.14.10
+    hooks:
+      - id: ruff
+        args: ["--fix"]
+      - id: ruff-format
+  - repo: local
+    hooks:
+      - id: ty
+        name: ty
+        entry: .venv/bin/python -m ty check
+        language: system
+        pass_filenames: false

{codex_lb-0.1.5 → codex_lb-0.2.0}/AGENTS.md

@@ -33,3 +33,20 @@
 - Services should be instantiated inside context providers and receive repositories via constructor injection.
 - Background tasks or standalone scripts must create and manage their own session; do not reuse request contexts.
 - When adding a new module, define `api.py` endpoints that depend on a module-specific context provider.
+
+## Git Workflow & Contribution
+1. **Important**: Create branches, commits, or PRs **only upon explicit user request**. Implicit actions are not allowed.
+2. **Branch Naming**: Use prefixes like `feature/`, `fix/`, `chore/` (e.g., `feature/add-login`).
+3. **Commit Messages**: Follow [Conventional Commits](https://www.conventionalcommits.org/).
+   - Format: `<type>(<scope>): <description>`
+   - Types: `feat`, `fix`, `docs`, `refactor`, `chore`, `test`
+   - Example: `feat(api): add auth endpoint`
+4. **Workflow**:
+   ```bash
+   git checkout -b feature/add-login
+   git commit -m "feat(api): add auth endpoint"
+   # Only on explicit request:
+   git push -u origin feature/add-login
+   gh pr create --title "feat(api): add auth" --body "..."
+   ```
+5. **Best Practices**: Commit often in small units. Do not commit directly to `main`. Always check `git diff` before pushing.

{codex_lb-0.1.5 → codex_lb-0.2.0}/CHANGELOG.md

@@ -1,5 +1,26 @@
 # Changelog
 
+## [0.2.0](https://github.com/Soju06/codex-lb/compare/v0.1.5...v0.2.0) (2026-01-19)
+
+
+### Features
+
+* add ty type checking and pre-commit hook
+* add health response schema and typed context cleanup
+
+
+### Bug Fixes
+
+* normalize stored plan types (pro/team/business/enterprise/edu) so accounts no longer show as unknown
+* prevent rate-limit status when usage is below 100% by using cooldown/backoff and primary-window quota checks
+* surface per-account quota reset times by applying primary/secondary reset windows with fallbacks
+
+
+### Refactor
+
+* move auth/usage helpers into module boundaries and extract proxy helpers
+* tighten typing across services and tests
+
 ## [0.1.5](https://github.com/Soju06/codex-lb/compare/v0.1.4...v0.1.5) (2026-01-14)
 
 

{codex_lb-0.1.5 → codex_lb-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: codex-lb
-Version: 0.1.5
+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>

{codex_lb-0.1.5 → codex_lb-0.2.0}/app/__init__.py

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

{codex_lb-0.1.5 → codex_lb-0.2.0}/app/core/auth/__init__.py

@@ -82,10 +82,11 @@ 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,
     )
 
 

{codex_lb-0.1.5 → codex_lb-0.2.0}/app/core/balancer/logic.py

@@ -23,6 +23,7 @@ class AccountState:
     status: AccountStatus
     used_percent: float | None = None
     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
@@ -59,6 +60,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,6 +89,10 @@ 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]:
@@ -94,14 +105,13 @@ def select_account(states: Iterable[AccountState], now: float | None = 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()
     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:

{codex_lb-0.1.5 → codex_lb-0.2.0}/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"
 
 

{codex_lb-0.1.5 → codex_lb-0.2.0}/app/core/config/settings.py

@@ -39,6 +39,7 @@ 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
 
     @field_validator("database_url")
     @classmethod
@@ -61,7 +62,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)

codex_lb-0.2.0/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

codex_lb-0.2.0/app/core/types.py

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

{codex_lb-0.1.5 → codex_lb-0.2.0}/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,
@@ -134,9 +135,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)

codex_lb-0.2.0/app/core/usage/quota.py

@@ -0,0 +1,58 @@
+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:
+            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:
+            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

{codex_lb-0.1.5 → codex_lb-0.2.0}/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:

codex_lb-0.2.0/app/db/migrations/__init__.py

@@ -0,0 +1,80 @@
+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 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
+"""
+
+
+@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),
+)
+
+
+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(),
+            },
+        )
+        rowcount = getattr(result, "rowcount", 0) or 0
+        if not rowcount:
+            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()

codex_lb-0.2.0/app/db/migrations/versions/__init__.py

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

codex_lb-0.2.0/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

{codex_lb-0.1.5 → codex_lb-0.2.0}/app/db/session.py

@@ -1,15 +1,19 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 from pathlib import Path
 from typing import AsyncIterator
 
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
 
 from app.core.config.settings import get_settings
+from app.db.migrations import run_migrations
 
 DATABASE_URL = get_settings().database_url
 
+logger = logging.getLogger(__name__)
+
 engine = create_async_engine(DATABASE_URL, echo=False)
 SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
 
@@ -60,3 +64,13 @@ async def init_db() -> None:
 
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.create_all)
+
+    async with SessionLocal() as session:
+        try:
+            updated = await run_migrations(session)
+            if updated:
+                logger.info("Applied database migrations count=%s", updated)
+        except Exception:
+            logger.exception("Failed to apply database migrations")
+            if get_settings().database_migrations_fail_fast:
+                raise

{codex_lb-0.1.5 → codex_lb-0.2.0}/app/dependencies.py

@@ -22,8 +22,6 @@ from app.modules.usage.service import UsageService
 class AccountsContext:
     session: AsyncSession
     repository: AccountsRepository
-    usage_repository: UsageRepository
-    request_logs_repository: RequestLogsRepository
     service: AccountsService
 
 
@@ -31,8 +29,6 @@ class AccountsContext:
 class UsageContext:
     session: AsyncSession
     usage_repository: UsageRepository
-    request_logs_repository: RequestLogsRepository
-    accounts_repository: AccountsRepository
     service: UsageService
 
 
@@ -63,8 +59,6 @@ def get_accounts_context(
     return AccountsContext(
         session=session,
         repository=repository,
-        usage_repository=usage_repository,
-        request_logs_repository=request_logs_repository,
         service=service,
     )
 
@@ -79,8 +73,6 @@ def get_usage_context(
     return UsageContext(
         session=session,
         usage_repository=usage_repository,
-        request_logs_repository=request_logs_repository,
-        accounts_repository=accounts_repository,
         service=service,
     )
 

{codex_lb-0.1.5 → codex_lb-0.2.0}/app/main.py

@@ -11,7 +11,7 @@ from fastapi.exception_handlers import (
     request_validation_exception_handler,
 )
 from fastapi.exceptions import RequestValidationError
-from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
+from fastapi.responses import FileResponse, JSONResponse, RedirectResponse, Response
 from fastapi.staticfiles import StaticFiles
 from starlette.exceptions import HTTPException as StarletteHTTPException
 
@@ -57,7 +57,7 @@ def create_app() -> FastAPI:
         return response
 
     @app.middleware("http")
-    async def api_unhandled_error_middleware(request: Request, call_next) -> JSONResponse:
+    async def api_unhandled_error_middleware(request: Request, call_next) -> Response:
         try:
             return await call_next(request)
         except Exception:
@@ -76,7 +76,7 @@ def create_app() -> FastAPI:
     async def _validation_error_handler(
         request: Request,
         exc: RequestValidationError,
-    ) -> JSONResponse:
+    ) -> Response:
         if request.url.path.startswith("/api/"):
             return JSONResponse(
                 status_code=422,
@@ -88,7 +88,7 @@ def create_app() -> FastAPI:
     async def _http_error_handler(
         request: Request,
         exc: StarletteHTTPException,
-    ) -> JSONResponse:
+    ) -> Response:
         if request.url.path.startswith("/api/"):
             detail = exc.detail if isinstance(exc.detail, str) else "Request failed"
             return JSONResponse(

{codex_lb-0.1.5/app/modules/proxy → codex_lb-0.2.0/app/modules/accounts}/auth_manager.py

@@ -1,15 +1,39 @@
 from __future__ import annotations
 
+from datetime import datetime
+from typing import Protocol
+
+from app.core.auth import DEFAULT_PLAN
 from app.core.auth.refresh import RefreshError, refresh_access_token, should_refresh
 from app.core.balancer import PERMANENT_FAILURE_CODES
 from app.core.crypto import TokenEncryptor
+from app.core.plan_types import coerce_account_plan_type
 from app.core.utils.time import utcnow
 from app.db.models import Account, AccountStatus
-from app.modules.accounts.repository import AccountsRepository
+
+
+class AccountsRepositoryPort(Protocol):
+    async def update_status(
+        self,
+        account_id: str,
+        status: AccountStatus,
+        deactivation_reason: str | None = None,
+    ) -> bool: ...
+
+    async def update_tokens(
+        self,
+        account_id: str,
+        access_token_encrypted: bytes,
+        refresh_token_encrypted: bytes,
+        id_token_encrypted: bytes,
+        last_refresh: datetime,
+        plan_type: str | None = None,
+        email: str | None = None,
+    ) -> bool: ...
 
 
 class AuthManager:
-    def __init__(self, repo: AccountsRepository) -> None:
+    def __init__(self, repo: AccountsRepositoryPort) -> None:
         self._repo = repo
         self._encryptor = TokenEncryptor()
 
@@ -34,8 +58,13 @@ class AuthManager:
         account.refresh_token_encrypted = self._encryptor.encrypt(result.refresh_token)
         account.id_token_encrypted = self._encryptor.encrypt(result.id_token)
         account.last_refresh = utcnow()
-        if result.plan_type:
-            account.plan_type = result.plan_type
+        if result.plan_type is not None:
+            account.plan_type = coerce_account_plan_type(
+                result.plan_type,
+                account.plan_type or DEFAULT_PLAN,
+            )
+        elif not account.plan_type:
+            account.plan_type = DEFAULT_PLAN
         if result.email:
             account.email = result.email
 

{codex_lb-0.1.5 → codex_lb-0.2.0}/app/modules/accounts/repository.py

@@ -48,12 +48,12 @@ class AccountsRepository:
             .values(status=status, deactivation_reason=deactivation_reason)
         )
         await self._session.commit()
-        return bool(result.rowcount)
+        return bool(getattr(result, "rowcount", 0) or 0)
 
     async def delete(self, account_id: str) -> bool:
         result = await self._session.execute(delete(Account).where(Account.id == account_id))
         await self._session.commit()
-        return bool(result.rowcount)
+        return bool(getattr(result, "rowcount", 0) or 0)
 
     async def update_tokens(
         self,
@@ -77,4 +77,4 @@ class AccountsRepository:
             values["email"] = email
         result = await self._session.execute(update(Account).where(Account.id == account_id).values(**values))
         await self._session.commit()
-        return bool(result.rowcount)
+        return bool(getattr(result, "rowcount", 0) or 0)