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

app/modules/proxy/sticky_repository.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from sqlalchemy import delete, select
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql import Insert, func
+
+from app.db.models import StickySession
+
+
+class StickySessionsRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def get_account_id(self, key: str) -> str | None:
+        if not key:
+            return None
+        result = await self._session.execute(select(StickySession.account_id).where(StickySession.key == key))
+        return result.scalar_one_or_none()
+
+    async def upsert(self, key: str, account_id: str) -> StickySession:
+        statement = self._build_upsert_statement(key, account_id)
+        await self._session.execute(statement)
+        await self._session.commit()
+        row = await self._session.get(StickySession, key)
+        if row is None:
+            raise RuntimeError(f"StickySession upsert failed for key={key!r}")
+        await self._session.refresh(row)
+        return row
+
+    async def delete(self, key: str) -> bool:
+        if not key:
+            return False
+        result = await self._session.execute(
+            delete(StickySession).where(StickySession.key == key).returning(StickySession.key)
+        )
+        await self._session.commit()
+        return result.scalar_one_or_none() is not None
+
+    def _build_upsert_statement(self, key: str, account_id: str) -> Insert:
+        dialect = self._session.get_bind().dialect.name
+        if dialect == "postgresql":
+            insert_fn = pg_insert
+        elif dialect == "sqlite":
+            insert_fn = sqlite_insert
+        else:
+            raise RuntimeError(f"StickySession upsert unsupported for dialect={dialect!r}")
+        statement = insert_fn(StickySession).values(key=key, account_id=account_id)
+        return statement.on_conflict_do_update(
+            index_elements=[StickySession.key],
+            set_={
+                "account_id": account_id,
+                "updated_at": func.now(),
+            },
+        )

app/modules/request_logs/repository.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-import asyncio
 from datetime import datetime
 
+import anyio
 from sqlalchemy import and_, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -33,6 +33,7 @@ class RequestLogsRepository:
         requested_at: datetime | None = None,
         cached_input_tokens: int | None = None,
         reasoning_tokens: int | None = None,
+        reasoning_effort: str | None = None,
     ) -> RequestLog:
         resolved_request_id = ensure_request_id(request_id)
         log = RequestLog(
@@ -43,6 +44,7 @@ class RequestLogsRepository:
             output_tokens=output_tokens,
             cached_input_tokens=cached_input_tokens,
             reasoning_tokens=reasoning_tokens,
+            reasoning_effort=reasoning_effort,
             latency_ms=latency_ms,
             status=status,
             error_code=error_code,
@@ -95,6 +97,7 @@ async def _safe_rollback(session: AsyncSession) -> None:
     if not session.in_transaction():
         return
     try:
-        await asyncio.shield(session.rollback())
-    except Exception:
+        with anyio.CancelScope(shield=True):
+            await session.rollback()
+    except BaseException:
         return

app/modules/request_logs/schemas.py

@@ -17,6 +17,8 @@ class RequestLogEntry(DashboardModel):
     error_code: str | None = None
     error_message: str | None = None
     tokens: int | None = None
+    cached_input_tokens: int | None = None
+    reasoning_effort: str | None = None
     cost_usd: float | None = None
     latency_ms: int | None = None
 

app/modules/request_logs/service.py

@@ -3,7 +3,12 @@ from __future__ import annotations
 from datetime import datetime
 from typing import cast
 
-from app.core.usage.logs import RequestLogLike, cost_from_log, total_tokens_from_log
+from app.core.usage.logs import (
+    RequestLogLike,
+    cached_input_tokens_from_log,
+    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
@@ -70,10 +75,12 @@ def _to_entry(log: RequestLog) -> RequestLogEntry:
         account_id=log.account_id,
         request_id=log.request_id,
         model=log.model,
+        reasoning_effort=log.reasoning_effort,
         status=_log_status(log),
         error_code=log.error_code,
         error_message=log.error_message,
         tokens=total_tokens_from_log(log_like),
+        cached_input_tokens=cached_input_tokens_from_log(log_like),
         cost_usd=cost_from_log(log_like, precision=6),
         latency_ms=log.latency_ms,
     )

app/modules/settings/__init__.py

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

app/modules/settings/api.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Body, Depends
+
+from app.dependencies import SettingsContext, get_settings_context
+from app.modules.settings.schemas import DashboardSettingsResponse, DashboardSettingsUpdateRequest
+from app.modules.settings.service import DashboardSettingsData
+
+router = APIRouter(prefix="/api/settings", tags=["dashboard"])
+
+
+@router.get("", response_model=DashboardSettingsResponse)
+async def get_settings(
+    context: SettingsContext = Depends(get_settings_context),
+) -> DashboardSettingsResponse:
+    settings = await context.service.get_settings()
+    return DashboardSettingsResponse(
+        sticky_threads_enabled=settings.sticky_threads_enabled,
+        prefer_earlier_reset_accounts=settings.prefer_earlier_reset_accounts,
+    )
+
+
+@router.put("", response_model=DashboardSettingsResponse)
+async def update_settings(
+    payload: DashboardSettingsUpdateRequest = Body(...),
+    context: SettingsContext = Depends(get_settings_context),
+) -> DashboardSettingsResponse:
+    updated = await context.service.update_settings(
+        DashboardSettingsData(
+            sticky_threads_enabled=payload.sticky_threads_enabled,
+            prefer_earlier_reset_accounts=payload.prefer_earlier_reset_accounts,
+        )
+    )
+    return DashboardSettingsResponse(
+        sticky_threads_enabled=updated.sticky_threads_enabled,
+        prefer_earlier_reset_accounts=updated.prefer_earlier_reset_accounts,
+    )

app/modules/settings/repository.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db.models import DashboardSettings
+
+_SETTINGS_ID = 1
+
+
+class SettingsRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def get_or_create(self) -> DashboardSettings:
+        existing = await self._session.get(DashboardSettings, _SETTINGS_ID)
+        if existing is not None:
+            return existing
+
+        row = DashboardSettings(
+            id=_SETTINGS_ID,
+            sticky_threads_enabled=False,
+            prefer_earlier_reset_accounts=False,
+        )
+        self._session.add(row)
+        await self._session.commit()
+        await self._session.refresh(row)
+        return row
+
+    async def update(
+        self,
+        *,
+        sticky_threads_enabled: bool,
+        prefer_earlier_reset_accounts: bool,
+    ) -> DashboardSettings:
+        settings = await self.get_or_create()
+        settings.sticky_threads_enabled = sticky_threads_enabled
+        settings.prefer_earlier_reset_accounts = prefer_earlier_reset_accounts
+        await self._session.commit()
+        await self._session.refresh(settings)
+        return settings

app/modules/settings/schemas.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from app.modules.shared.schemas import DashboardModel
+
+
+class DashboardSettingsResponse(DashboardModel):
+    sticky_threads_enabled: bool
+    prefer_earlier_reset_accounts: bool
+
+
+class DashboardSettingsUpdateRequest(DashboardModel):
+    sticky_threads_enabled: bool
+    prefer_earlier_reset_accounts: bool

app/modules/settings/service.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from app.modules.settings.repository import SettingsRepository
+
+
+@dataclass(frozen=True, slots=True)
+class DashboardSettingsData:
+    sticky_threads_enabled: bool
+    prefer_earlier_reset_accounts: bool
+
+
+class SettingsService:
+    def __init__(self, repository: SettingsRepository) -> None:
+        self._repository = repository
+
+    async def get_settings(self) -> DashboardSettingsData:
+        row = await self._repository.get_or_create()
+        return DashboardSettingsData(
+            sticky_threads_enabled=row.sticky_threads_enabled,
+            prefer_earlier_reset_accounts=row.prefer_earlier_reset_accounts,
+        )
+
+    async def update_settings(self, payload: DashboardSettingsData) -> DashboardSettingsData:
+        row = await self._repository.update(
+            sticky_threads_enabled=payload.sticky_threads_enabled,
+            prefer_earlier_reset_accounts=payload.prefer_earlier_reset_accounts,
+        )
+        return DashboardSettingsData(
+            sticky_threads_enabled=row.sticky_threads_enabled,
+            prefer_earlier_reset_accounts=row.prefer_earlier_reset_accounts,
+        )

app/modules/shared/schemas.py

@@ -1,8 +1,22 @@
 from __future__ import annotations
 
-from pydantic import BaseModel, ConfigDict
+from datetime import datetime
+
+from pydantic import BaseModel, ConfigDict, field_serializer
 from pydantic.alias_generators import to_camel
 
 
 class DashboardModel(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        ser_json_timedelta="iso8601",
+    )
+
+    @field_serializer("*", when_used="json")
+    def serialize_datetime_as_utc(value, _info):
+        if isinstance(value, datetime):
+            if value.tzinfo is None:
+                return value.isoformat() + "Z"
+            return value.isoformat().replace("+00:00", "Z")
+        return value

app/modules/usage/schemas.py

@@ -30,6 +30,7 @@ class UsageCost(DashboardModel):
 class UsageMetrics(DashboardModel):
     requests_7d: int | None = Field(default=None, alias="requests7d")
     tokens_secondary_window: int | None = None
+    cached_tokens_secondary_window: int | None = None
     error_rate_7d: float | None = Field(default=None, alias="errorRate7d")
     top_error: str | None = None
 

app/modules/usage/service.py

@@ -4,7 +4,13 @@ from datetime import timedelta
 from typing import cast
 
 from app.core import usage as usage_core
-from app.core.usage.logs import RequestLogLike, cost_from_log, total_tokens_from_log, usage_tokens_from_log
+from app.core.usage.logs import (
+    RequestLogLike,
+    cached_input_tokens_from_log,
+    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,
@@ -181,9 +187,11 @@ def _usage_metrics(logs_secondary: list[RequestLog]) -> UsageMetricsSummary:
         error_rate = len(error_logs) / total_requests
     top_error = _top_error_code(error_logs)
     tokens_secondary = _sum_tokens(logs_secondary)
+    cached_tokens_secondary = _sum_cached_input_tokens(logs_secondary)
     return UsageMetricsSummary(
         requests_7d=total_requests,
         tokens_secondary_window=tokens_secondary,
+        cached_tokens_secondary_window=cached_tokens_secondary,
         error_rate_7d=error_rate,
         top_error=top_error,
     )
@@ -196,6 +204,13 @@ def _sum_tokens(logs: list[RequestLog]) -> int:
     return total
 
 
+def _sum_cached_input_tokens(logs: list[RequestLog]) -> int:
+    total = 0
+    for log in logs:
+        total += cached_input_tokens_from_log(cast(RequestLogLike, log)) or 0
+    return total
+
+
 def _top_error_code(logs: list[RequestLog]) -> str | None:
     counts: dict[str, int] = {}
     for log in logs:
@@ -242,6 +257,7 @@ def _metrics_summary_to_model(metrics: UsageMetricsSummary) -> UsageMetrics:
     return UsageMetrics(
         requests_7d=metrics.requests_7d,
         tokens_secondary_window=metrics.tokens_secondary_window,
+        cached_tokens_secondary_window=metrics.cached_tokens_secondary_window,
         error_rate_7d=metrics.error_rate_7d,
         top_error=metrics.top_error,
     )

app/modules/usage/updater.py

@@ -2,7 +2,9 @@ from __future__ import annotations
 
 import logging
 import math
-from typing import Mapping
+from collections import Counter
+from datetime import datetime
+from typing import Mapping, Protocol
 
 from app.core.auth.refresh import RefreshError
 from app.core.clients.usage import UsageFetchError, fetch_usage
@@ -14,15 +16,31 @@ 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.usage.repository import UsageRepository
 
 logger = logging.getLogger(__name__)
 
 
+class UsageRepositoryPort(Protocol):
+    async def add_entry(
+        self,
+        account_id: str,
+        used_percent: float,
+        input_tokens: int | None = None,
+        output_tokens: int | None = None,
+        recorded_at: datetime | None = None,
+        window: str | None = None,
+        reset_at: int | None = None,
+        window_minutes: int | None = None,
+        credits_has: bool | None = None,
+        credits_unlimited: bool | None = None,
+        credits_balance: float | None = None,
+    ) -> UsageHistory | None: ...
+
+
 class UsageUpdater:
     def __init__(
         self,
-        usage_repo: UsageRepository,
+        usage_repo: UsageRepositoryPort,
         accounts_repo: AccountsRepository | None = None,
     ) -> None:
         self._usage_repo = usage_repo
@@ -38,6 +56,7 @@ class UsageUpdater:
         if not settings.usage_refresh_enabled:
             return
 
+        shared_chatgpt_account_ids = _shared_chatgpt_account_ids(accounts)
         now = utcnow()
         interval = settings.usage_refresh_interval_seconds
         for account in accounts:
@@ -46,11 +65,16 @@ class UsageUpdater:
             latest = latest_usage.get(account.id)
             if latest and (now - latest.recorded_at).total_seconds() < interval:
                 continue
+            usage_account_id = (
+                None
+                if account.chatgpt_account_id and account.chatgpt_account_id in shared_chatgpt_account_ids
+                else account.chatgpt_account_id
+            )
             # NOTE: AsyncSession is not safe for concurrent use. Run sequentially
             # within the request-scoped session to avoid PK collisions and
             # flush-time warnings (SAWarning: Session.add during flush).
             try:
-                await self._refresh_account(account)
+                await self._refresh_account(account, usage_account_id=usage_account_id)
             except Exception as exc:
                 logger.warning(
                     "Usage refresh failed account_id=%s request_id=%s error=%s",
@@ -62,12 +86,12 @@ class UsageUpdater:
                 # swallow per-account failures so the whole refresh loop keeps going
                 continue
 
-    async def _refresh_account(self, account: Account) -> None:
+    async def _refresh_account(self, account: Account, *, usage_account_id: str | None) -> None:
         access_token = self._encryptor.decrypt(account.access_token_encrypted)
         try:
             payload = await fetch_usage(
                 access_token=access_token,
-                account_id=account.id,
+                account_id=usage_account_id,
             )
         except UsageFetchError as exc:
             if exc.status_code != 401 or not self._auth_manager:
@@ -80,7 +104,7 @@ class UsageUpdater:
             try:
                 payload = await fetch_usage(
                     access_token=access_token,
-                    account_id=account.id,
+                    account_id=usage_account_id,
                 )
             except UsageFetchError:
                 return
@@ -145,3 +169,8 @@ def _window_minutes(limit_seconds: int | None) -> int | None:
     if not limit_seconds or limit_seconds <= 0:
         return None
     return max(1, math.ceil(limit_seconds / 60))
+
+
+def _shared_chatgpt_account_ids(accounts: list[Account]) -> set[str]:
+    counts = Counter(account.chatgpt_account_id for account in accounts if account.chatgpt_account_id)
+    return {account_id for account_id, count in counts.items() if count > 1}