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

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,42 +1,102 @@
 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
+from app.db.migrations import run_migrations
 
 DATABASE_URL = get_settings().database_url
 
-engine = create_async_engine(DATABASE_URL, echo=False)
+logger = logging.getLogger(__name__)
+
+_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
-    path = url[len(prefix) :]
-    if path == ":memory:":
+
+    marker = ":///"
+    marker_index = url.find(marker)
+    if marker_index < 0:
         return
+
+    # 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
 
 
@@ -60,3 +120,17 @@ 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
+
+
+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
 
@@ -22,8 +25,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 +32,6 @@ class AccountsContext:
 class UsageContext:
     session: AsyncSession
     usage_repository: UsageRepository
-    request_logs_repository: RequestLogsRepository
-    accounts_repository: AccountsRepository
     service: UsageService
 
 
@@ -53,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:
@@ -63,8 +69,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 +83,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,
     )
 
@@ -112,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)
 
 
@@ -122,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

@@ -11,19 +11,20 @@ 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
 
 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:
@@ -57,7 +61,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 +80,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 +92,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(
@@ -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

app/modules/accounts/auth_manager.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+from typing import Protocol
+
+from app.core.auth import DEFAULT_PLAN, OpenAIAuthClaims, extract_id_token_claims
+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
+
+
+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,
+        chatgpt_account_id: str | None = None,
+    ) -> bool: ...
+
+
+logger = logging.getLogger(__name__)
+
+
+class AuthManager:
+    def __init__(self, repo: AccountsRepositoryPort) -> None:
+        self._repo = repo
+        self._encryptor = TokenEncryptor()
+
+    async def ensure_fresh(self, account: Account, *, force: bool = False) -> Account:
+        if force or should_refresh(account.last_refresh):
+            account = await self.refresh_account(account)
+        return await self._ensure_chatgpt_account_id(account)
+
+    async def refresh_account(self, account: Account) -> Account:
+        refresh_token = self._encryptor.decrypt(account.refresh_token_encrypted)
+        try:
+            result = await refresh_access_token(refresh_token)
+        except RefreshError as exc:
+            if exc.is_permanent:
+                reason = PERMANENT_FAILURE_CODES.get(exc.code, exc.message)
+                await self._repo.update_status(account.id, AccountStatus.DEACTIVATED, reason)
+                account.status = AccountStatus.DEACTIVATED
+                account.deactivation_reason = reason
+            raise
+
+        account.access_token_encrypted = self._encryptor.encrypt(result.access_token)
+        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.account_id:
+            account.chatgpt_account_id = result.account_id
+        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
+
+        await self._repo.update_tokens(
+            account.id,
+            access_token_encrypted=account.access_token_encrypted,
+            refresh_token_encrypted=account.refresh_token_encrypted,
+            id_token_encrypted=account.id_token_encrypted,
+            last_refresh=account.last_refresh,
+            plan_type=account.plan_type,
+            email=account.email,
+            chatgpt_account_id=account.chatgpt_account_id,
+        )
+        return account
+
+    async def _ensure_chatgpt_account_id(self, account: Account) -> Account:
+        if account.chatgpt_account_id:
+            return account
+        try:
+            id_token = self._encryptor.decrypt(account.id_token_encrypted)
+        except Exception:
+            return account
+        raw_account_id = _chatgpt_account_id_from_id_token(id_token)
+        if not raw_account_id:
+            return account
+
+        account.chatgpt_account_id = raw_account_id
+        try:
+            await self._repo.update_tokens(
+                account.id,
+                access_token_encrypted=account.access_token_encrypted,
+                refresh_token_encrypted=account.refresh_token_encrypted,
+                id_token_encrypted=account.id_token_encrypted,
+                last_refresh=account.last_refresh,
+                plan_type=account.plan_type,
+                email=account.email,
+                chatgpt_account_id=raw_account_id,
+            )
+        except Exception:
+            logger.warning("Failed to persist chatgpt_account_id account_id=%s", account.id, exc_info=True)
+        return account
+
+
+def _chatgpt_account_id_from_id_token(id_token: str) -> str | None:
+    claims = extract_id_token_claims(id_token)
+    auth_claims = claims.auth or OpenAIAuthClaims()
+    return auth_claims.chatgpt_account_id or claims.chatgpt_account_id

app/modules/accounts/repository.py

@@ -19,6 +19,7 @@ class AccountsRepository:
     async def upsert(self, account: Account) -> Account:
         existing = await self._session.get(Account, account.id)
         if existing:
+            existing.chatgpt_account_id = account.chatgpt_account_id
             existing.email = account.email
             existing.plan_type = account.plan_type
             existing.access_token_encrypted = account.access_token_encrypted
@@ -41,19 +42,21 @@ class AccountsRepository:
         account_id: str,
         status: AccountStatus,
         deactivation_reason: str | None = None,
+        reset_at: int | None = None,
     ) -> bool:
         result = await self._session.execute(
             update(Account)
             .where(Account.id == account_id)
-            .values(status=status, deactivation_reason=deactivation_reason)
+            .values(status=status, deactivation_reason=deactivation_reason, reset_at=reset_at)
+            .returning(Account.id)
         )
         await self._session.commit()
-        return bool(result.rowcount)
+        return result.scalar_one_or_none() is not None
 
     async def delete(self, account_id: str) -> bool:
-        result = await self._session.execute(delete(Account).where(Account.id == account_id))
+        result = await self._session.execute(delete(Account).where(Account.id == account_id).returning(Account.id))
         await self._session.commit()
-        return bool(result.rowcount)
+        return result.scalar_one_or_none() is not None
 
     async def update_tokens(
         self,
@@ -64,6 +67,7 @@ class AccountsRepository:
         last_refresh: datetime,
         plan_type: str | None = None,
         email: str | None = None,
+        chatgpt_account_id: str | None = None,
     ) -> bool:
         values = {
             "access_token_encrypted": access_token_encrypted,
@@ -75,6 +79,10 @@ class AccountsRepository:
             values["plan_type"] = plan_type
         if email is not None:
             values["email"] = email
-        result = await self._session.execute(update(Account).where(Account.id == account_id).values(**values))
+        if chatgpt_account_id is not None:
+            values["chatgpt_account_id"] = chatgpt_account_id
+        result = await self._session.execute(
+            update(Account).where(Account.id == account_id).values(**values).returning(Account.id)
+        )
         await self._session.commit()
-        return bool(result.rowcount)
+        return result.scalar_one_or_none() is not None

app/modules/accounts/service.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta, timezone
+from typing import cast
 
 from app.core import usage as usage_core
 from app.core.auth import (
@@ -8,11 +9,12 @@ from app.core.auth import (
     DEFAULT_PLAN,
     claims_from_auth,
     extract_id_token_claims,
-    fallback_account_id,
+    generate_unique_account_id,
     parse_auth_json,
 )
 from app.core.crypto import TokenEncryptor
-from app.core.usage.logs import cost_from_log
+from app.core.plan_types import coerce_account_plan_type
+from app.core.usage.logs import RequestLogLike, cost_from_log
 from app.core.utils.time import from_epoch_seconds, to_utc_naive, utcnow
 from app.db.models import Account, AccountStatus, UsageHistory
 from app.modules.accounts.repository import AccountsRepository
@@ -23,9 +25,9 @@ from app.modules.accounts.schemas import (
     AccountTokenStatus,
     AccountUsage,
 )
-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.updater import UsageUpdater
 
 
 class AccountsService:
@@ -64,12 +66,14 @@ class AccountsService:
         claims = claims_from_auth(auth)
 
         email = claims.email or DEFAULT_EMAIL
-        plan_type = claims.plan_type or DEFAULT_PLAN
-        account_id = claims.account_id or fallback_account_id(email)
+        raw_account_id = claims.account_id
+        account_id = generate_unique_account_id(raw_account_id, email)
+        plan_type = coerce_account_plan_type(claims.plan_type, DEFAULT_PLAN)
         last_refresh = to_utc_naive(auth.last_refresh_at) if auth.last_refresh_at else utcnow()
 
         account = Account(
             id=account_id,
+            chatgpt_account_id=raw_account_id,
             email=email,
             plan_type=plan_type,
             access_token_encrypted=self._encryptor.encrypt(auth.tokens.access_token),
@@ -107,6 +111,7 @@ class AccountsService:
         secondary_usage: UsageHistory | None,
         cost_usd_24h: float | None,
     ) -> AccountSummary:
+        plan_type = coerce_account_plan_type(account.plan_type, DEFAULT_PLAN)
         auth_status = self._build_auth_status(account)
         primary_used_percent = _normalize_used_percent(primary_usage) or 0.0
         secondary_used_percent = _normalize_used_percent(secondary_usage) or 0.0
@@ -114,8 +119,8 @@ class AccountsService:
         secondary_remaining_percent = usage_core.remaining_percent_from_used(secondary_used_percent) or 0.0
         reset_at_primary = from_epoch_seconds(primary_usage.reset_at) if primary_usage is not None else None
         reset_at_secondary = from_epoch_seconds(secondary_usage.reset_at) if secondary_usage is not None else None
-        capacity_primary = usage_core.capacity_for_plan(account.plan_type, "primary")
-        capacity_secondary = usage_core.capacity_for_plan(account.plan_type, "secondary")
+        capacity_primary = usage_core.capacity_for_plan(plan_type, "primary")
+        capacity_secondary = usage_core.capacity_for_plan(plan_type, "secondary")
         remaining_credits_primary = usage_core.remaining_credits_from_percent(
             primary_used_percent,
             capacity_primary,
@@ -128,7 +133,7 @@ class AccountsService:
             account_id=account.id,
             email=account.email,
             display_name=account.email,
-            plan_type=account.plan_type,
+            plan_type=plan_type,
             status=account.status.value,
             usage=AccountUsage(
                 primary_remaining_percent=primary_remaining_percent,
@@ -186,7 +191,7 @@ class AccountsService:
         logs = await self._logs_repo.list_since(since)
         totals: dict[str, float] = {}
         for log in logs:
-            cost = cost_from_log(log)
+            cost = cost_from_log(cast(RequestLogLike, log))
             if cost is None:
                 continue
             totals[log.account_id] = totals.get(log.account_id, 0.0) + cost

app/modules/health/api.py

@@ -2,9 +2,11 @@ from __future__ import annotations
 
 from fastapi import APIRouter
 
+from app.modules.health.schemas import HealthResponse
+
 router = APIRouter(tags=["health"])
 
 
-@router.get("/health")
-async def health_check() -> dict:
-    return {"status": "ok"}
+@router.get("/health", response_model=HealthResponse)
+async def health_check() -> HealthResponse:
+    return HealthResponse(status="ok")

app/modules/health/schemas.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict
+
+
+class HealthResponse(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    status: str

app/modules/oauth/service.py

@@ -15,7 +15,7 @@ from app.core.auth import (
     DEFAULT_PLAN,
     OpenAIAuthClaims,
     extract_id_token_claims,
-    fallback_account_id,
+    generate_unique_account_id,
 )
 from app.core.clients.oauth import (
     OAuthError,
@@ -28,6 +28,7 @@ from app.core.clients.oauth import (
 )
 from app.core.config.settings import get_settings
 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
@@ -293,13 +294,17 @@ class OauthService:
     async def _persist_tokens(self, tokens: OAuthTokens) -> None:
         claims = extract_id_token_claims(tokens.id_token)
         auth_claims = claims.auth or OpenAIAuthClaims()
-        account_id = auth_claims.chatgpt_account_id or claims.chatgpt_account_id
+        raw_account_id = auth_claims.chatgpt_account_id or claims.chatgpt_account_id
         email = claims.email or DEFAULT_EMAIL
-        plan_type = auth_claims.chatgpt_plan_type or claims.chatgpt_plan_type or DEFAULT_PLAN
-        account_id = account_id or fallback_account_id(email)
+        account_id = generate_unique_account_id(raw_account_id, email)
+        plan_type = coerce_account_plan_type(
+            auth_claims.chatgpt_plan_type or claims.chatgpt_plan_type,
+            DEFAULT_PLAN,
+        )
 
         account = Account(
             id=account_id,
+            chatgpt_account_id=raw_account_id,
             email=email,
             plan_type=plan_type,
             access_token_encrypted=self._encryptor.encrypt(tokens.access_token),