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

app/core/auth/__init__.py

@@ -90,7 +90,17 @@ def claims_from_auth(auth: AuthFile) -> AccountClaims:
     )
 
 
+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:
@@ -24,6 +27,8 @@ class AccountState:
     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
@@ -36,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)
@@ -95,18 +105,35 @@ def select_account(states: Iterable[AccountState], now: float | None = None) ->
             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)
 
 
 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:

app/core/config/settings.py

@@ -40,6 +40,8 @@ class Settings(BaseSettings):
     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

app/core/usage/__init__.py

@@ -17,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,
 }
 

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

@@ -30,8 +30,11 @@ def apply_usage_quota(
                 reset_at = secondary_reset
             return status, used_percent, reset_at
         if status == AccountStatus.QUOTA_EXCEEDED:
-            status = AccountStatus.ACTIVE
-            reset_at = None
+            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
 
@@ -45,8 +48,11 @@ def apply_usage_quota(
                 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
+            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
 

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/db/migrations/__init__.py

@@ -8,7 +8,13 @@ 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
+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 (
@@ -21,6 +27,7 @@ _INSERT_MIGRATION = """
 INSERT INTO schema_migrations (name, applied_at)
 VALUES (:name, :applied_at)
 ON CONFLICT(name) DO NOTHING
+RETURNING name
 """
 
 
@@ -32,6 +39,10 @@ class Migration:
 
 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),
 )
 
 
@@ -54,8 +65,8 @@ async def _apply_migration(session: AsyncSession, migration: Migration) -> bool:
                 "applied_at": _utcnow_iso(),
             },
         )
-        rowcount = getattr(result, "rowcount", 0) or 0
-        if not rowcount:
+        inserted = result.scalar_one_or_none()
+        if inserted is None:
             return False
         await migration.run(session)
     return True

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/models.py

@@ -24,6 +24,7 @@ class Account(Base):
     __tablename__ = "accounts"
 
     id: Mapped[str] = mapped_column(String, primary_key=True)
+    chatgpt_account_id: Mapped[str | None] = mapped_column(String, nullable=True)
     email: Mapped[str] = mapped_column(String, unique=True, nullable=False)
     plan_type: Mapped[str] = mapped_column(String, nullable=False)
 
@@ -40,6 +41,7 @@ class Account(Base):
         nullable=False,
     )
     deactivation_reason: Mapped[str | None] = mapped_column(Text, nullable=True)
+    reset_at: Mapped[int | None] = mapped_column(Integer, nullable=True)
 
 
 class UsageHistory(Base):
@@ -71,12 +73,43 @@ class RequestLog(Base):
     output_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
     cached_input_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
     reasoning_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    reasoning_effort: Mapped[str | None] = mapped_column(String, nullable=True)
     latency_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
     status: Mapped[str] = mapped_column(String, nullable=False)
     error_code: Mapped[str | None] = mapped_column(String, nullable=True)
     error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
 
 
+class StickySession(Base):
+    __tablename__ = "sticky_sessions"
+
+    key: Mapped[str] = mapped_column(String, primary_key=True)
+    account_id: Mapped[str] = mapped_column(String, ForeignKey("accounts.id"), nullable=False)
+    created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), nullable=False)
+    updated_at: Mapped[datetime] = mapped_column(
+        DateTime,
+        server_default=func.now(),
+        onupdate=func.now(),
+        nullable=False,
+    )
+
+
+class DashboardSettings(Base):
+    __tablename__ = "dashboard_settings"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=False)
+    sticky_threads_enabled: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
+    prefer_earlier_reset_accounts: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), nullable=False)
+    updated_at: Mapped[datetime] = mapped_column(
+        DateTime,
+        server_default=func.now(),
+        onupdate=func.now(),
+        nullable=False,
+    )
+
+
 Index("idx_usage_recorded_at", UsageHistory.recorded_at)
 Index("idx_usage_account_time", UsageHistory.account_id, UsageHistory.recorded_at)
 Index("idx_logs_account_time", RequestLog.account_id, RequestLog.requested_at)
+Index("idx_sticky_account", StickySession.account_id)

app/db/session.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
-import asyncio
 import logging
+import sqlite3
 from pathlib import Path
-from typing import AsyncIterator
+from typing import AsyncIterator, Awaitable, TypeVar
 
+import anyio
+from sqlalchemy import event
+from sqlalchemy.engine import Engine
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
 
 from app.core.config.settings import get_settings
@@ -14,33 +17,86 @@ DATABASE_URL = get_settings().database_url
 
 logger = logging.getLogger(__name__)
 
-engine = create_async_engine(DATABASE_URL, echo=False)
+_SQLITE_BUSY_TIMEOUT_MS = 5_000
+_SQLITE_BUSY_TIMEOUT_SECONDS = _SQLITE_BUSY_TIMEOUT_MS / 1000
+
+
+def _is_sqlite_url(url: str) -> bool:
+    return url.startswith("sqlite+aiosqlite:///") or url.startswith("sqlite:///")
+
+
+def _is_sqlite_memory_url(url: str) -> bool:
+    return _is_sqlite_url(url) and ":memory:" in url
+
+
+def _configure_sqlite_engine(engine: Engine, *, enable_wal: bool) -> None:
+    @event.listens_for(engine, "connect")
+    def _set_sqlite_pragmas(dbapi_connection: sqlite3.Connection, _: object) -> None:
+        cursor: sqlite3.Cursor = dbapi_connection.cursor()
+        try:
+            if enable_wal:
+                cursor.execute("PRAGMA journal_mode=WAL")
+            cursor.execute("PRAGMA synchronous=NORMAL")
+            cursor.execute("PRAGMA foreign_keys=ON")
+            cursor.execute(f"PRAGMA busy_timeout={_SQLITE_BUSY_TIMEOUT_MS}")
+        finally:
+            cursor.close()
+
+
+if _is_sqlite_url(DATABASE_URL):
+    engine = create_async_engine(
+        DATABASE_URL,
+        echo=False,
+        connect_args={"timeout": _SQLITE_BUSY_TIMEOUT_SECONDS},
+    )
+    _configure_sqlite_engine(engine.sync_engine, enable_wal=not _is_sqlite_memory_url(DATABASE_URL))
+else:
+    engine = create_async_engine(DATABASE_URL, echo=False)
+
 SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
 
+_T = TypeVar("_T")
+
 
 def _ensure_sqlite_dir(url: str) -> None:
-    prefix = "sqlite+aiosqlite:///"
-    if not url.startswith(prefix):
+    if not (url.startswith("sqlite+aiosqlite:") or url.startswith("sqlite:")):
+        return
+
+    marker = ":///"
+    marker_index = url.find(marker)
+    if marker_index < 0:
         return
-    path = url[len(prefix) :]
-    if path == ":memory:":
+
+    # Works for both relative (sqlite+aiosqlite:///./db.sqlite) and absolute
+    # paths (sqlite+aiosqlite:////var/lib/app/db.sqlite).
+    path = url[marker_index + len(marker) :]
+    path = path.partition("?")[0]
+    path = path.partition("#")[0]
+
+    if not path or path == ":memory:":
         return
+
     Path(path).expanduser().parent.mkdir(parents=True, exist_ok=True)
 
 
+async def _shielded(awaitable: Awaitable[_T]) -> _T:
+    with anyio.CancelScope(shield=True):
+        return await awaitable
+
+
 async def _safe_rollback(session: AsyncSession) -> None:
     if not session.in_transaction():
         return
     try:
-        await asyncio.shield(session.rollback())
-    except Exception:
+        await _shielded(session.rollback())
+    except BaseException:
         return
 
 
 async def _safe_close(session: AsyncSession) -> None:
     try:
-        await asyncio.shield(session.close())
-    except Exception:
+        await _shielded(session.close())
+    except BaseException:
         return
 
 
@@ -74,3 +130,7 @@ async def init_db() -> None:
             logger.exception("Failed to apply database migrations")
             if get_settings().database_migrations_fail_fast:
                 raise
+
+
+async def close_db() -> None:
+    await engine.dispose()

app/dependencies.py

@@ -12,8 +12,11 @@ from app.modules.accounts.repository import AccountsRepository
 from app.modules.accounts.service import AccountsService
 from app.modules.oauth.service import OauthService
 from app.modules.proxy.service import ProxyService
+from app.modules.proxy.sticky_repository import StickySessionsRepository
 from app.modules.request_logs.repository import RequestLogsRepository
 from app.modules.request_logs.service import RequestLogsService
+from app.modules.settings.repository import SettingsRepository
+from app.modules.settings.service import SettingsService
 from app.modules.usage.repository import UsageRepository
 from app.modules.usage.service import UsageService
 
@@ -49,6 +52,13 @@ class RequestLogsContext:
     service: RequestLogsService
 
 
+@dataclass(slots=True)
+class SettingsContext:
+    session: AsyncSession
+    repository: SettingsRepository
+    service: SettingsService
+
+
 def get_accounts_context(
     session: AsyncSession = Depends(get_session),
 ) -> AccountsContext:
@@ -104,7 +114,15 @@ def get_proxy_context(
     accounts_repository = AccountsRepository(session)
     usage_repository = UsageRepository(session)
     request_logs_repository = RequestLogsRepository(session)
-    service = ProxyService(accounts_repository, usage_repository, request_logs_repository)
+    sticky_repository = StickySessionsRepository(session)
+    settings_repository = SettingsRepository(session)
+    service = ProxyService(
+        accounts_repository,
+        usage_repository,
+        request_logs_repository,
+        sticky_repository,
+        settings_repository,
+    )
     return ProxyContext(service=service)
 
 
@@ -114,3 +132,11 @@ def get_request_logs_context(
     repository = RequestLogsRepository(session)
     service = RequestLogsService(repository)
     return RequestLogsContext(session=session, repository=repository, service=service)
+
+
+def get_settings_context(
+    session: AsyncSession = Depends(get_session),
+) -> SettingsContext:
+    repository = SettingsRepository(session)
+    service = SettingsService(repository)
+    return SettingsContext(session=session, repository=repository, service=service)

app/main.py

@@ -18,12 +18,13 @@ from starlette.exceptions import HTTPException as StarletteHTTPException
 from app.core.clients.http import close_http_client, init_http_client
 from app.core.errors import dashboard_error
 from app.core.utils.request_id import get_request_id, reset_request_id, set_request_id
-from app.db.session import init_db
+from app.db.session import close_db, init_db
 from app.modules.accounts import api as accounts_api
 from app.modules.health import api as health_api
 from app.modules.oauth import api as oauth_api
 from app.modules.proxy import api as proxy_api
 from app.modules.request_logs import api as request_logs_api
+from app.modules.settings import api as settings_api
 from app.modules.usage import api as usage_api
 
 logger = logging.getLogger(__name__)
@@ -37,7 +38,10 @@ async def lifespan(_: FastAPI):
     try:
         yield
     finally:
-        await close_http_client()
+        try:
+            await close_http_client()
+        finally:
+            await close_db()
 
 
 def create_app() -> FastAPI:
@@ -103,6 +107,7 @@ def create_app() -> FastAPI:
     app.include_router(usage_api.router)
     app.include_router(request_logs_api.router)
     app.include_router(oauth_api.router)
+    app.include_router(settings_api.router)
     app.include_router(health_api.router)
 
     static_dir = Path(__file__).parent / "static"
@@ -116,6 +121,10 @@ def create_app() -> FastAPI:
     async def spa_accounts():
         return FileResponse(index_html, media_type="text/html")
 
+    @app.get("/settings", include_in_schema=False)
+    async def spa_settings():
+        return FileResponse(index_html, media_type="text/html")
+
     app.mount("/dashboard", StaticFiles(directory=static_dir, html=True), name="dashboard")
 
     return app