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

app/modules/proxy/service.py

@@ -2,10 +2,12 @@ from __future__ import annotations
 
 import logging
 import time
+from collections.abc import Sequence
 from datetime import timedelta
-from typing import AsyncIterator, Iterable, Mapping
+from hashlib import sha256
+from typing import AsyncIterator, Mapping
 
-from pydantic import ValidationError
+import anyio
 
 from app.core import usage as usage_core
 from app.core.auth.refresh import RefreshError
@@ -14,28 +16,41 @@ from app.core.balancer.types import UpstreamError
 from app.core.clients.proxy import ProxyResponseError, filter_inbound_headers
 from app.core.clients.proxy import compact_responses as core_compact_responses
 from app.core.clients.proxy import stream_responses as core_stream_responses
+from app.core.config.settings import get_settings
 from app.core.crypto import TokenEncryptor
-from app.core.errors import OpenAIErrorDetail, OpenAIErrorEnvelope, openai_error, response_failed_event
-from app.core.openai.models import OpenAIError, OpenAIResponsePayload
+from app.core.errors import openai_error, response_failed_event
+from app.core.openai.models import OpenAIResponsePayload
 from app.core.openai.parsing import parse_sse_event
 from app.core.openai.requests import ResponsesCompactRequest, ResponsesRequest
-from app.core.usage.types import UsageWindowRow, UsageWindowSummary
-from app.core.utils.request_id import ensure_request_id
+from app.core.usage.types import UsageWindowRow
+from app.core.utils.request_id import ensure_request_id, get_request_id
 from app.core.utils.sse import format_sse_event
 from app.core.utils.time import utcnow
-from app.db.models import Account, AccountStatus, UsageHistory
+from app.db.models import Account, UsageHistory
+from app.modules.accounts.auth_manager import AuthManager
 from app.modules.accounts.repository import AccountsRepository
-from app.modules.proxy.auth_manager import AuthManager
-from app.modules.proxy.load_balancer import LoadBalancer
-from app.modules.proxy.types import (
-    CreditStatusDetailsData,
-    RateLimitStatusDetailsData,
-    RateLimitStatusPayloadData,
-    RateLimitWindowSnapshotData,
+from app.modules.proxy.helpers import (
+    _apply_error_metadata,
+    _credits_headers,
+    _credits_snapshot,
+    _header_account_id,
+    _normalize_error_code,
+    _parse_openai_error,
+    _plan_type_for_accounts,
+    _rate_limit_details,
+    _rate_limit_headers,
+    _select_accounts_for_limits,
+    _summarize_window,
+    _upstream_error_from_openai,
+    _window_snapshot,
 )
-from app.modules.proxy.usage_updater import UsageUpdater
+from app.modules.proxy.load_balancer import LoadBalancer
+from app.modules.proxy.sticky_repository import StickySessionsRepository
+from app.modules.proxy.types import RateLimitStatusPayloadData
 from app.modules.request_logs.repository import RequestLogsRepository
+from app.modules.settings.repository import SettingsRepository
 from app.modules.usage.repository import UsageRepository
+from app.modules.usage.updater import UsageUpdater
 
 logger = logging.getLogger(__name__)
 
@@ -46,13 +61,16 @@ class ProxyService:
         accounts_repo: AccountsRepository,
         usage_repo: UsageRepository,
         logs_repo: RequestLogsRepository,
+        sticky_repo: StickySessionsRepository,
+        settings_repo: SettingsRepository,
     ) -> None:
         self._accounts_repo = accounts_repo
         self._usage_repo = usage_repo
         self._logs_repo = logs_repo
+        self._settings_repo = settings_repo
         self._encryptor = TokenEncryptor()
         self._auth_manager = AuthManager(accounts_repo)
-        self._load_balancer = LoadBalancer(accounts_repo, usage_repo)
+        self._load_balancer = LoadBalancer(accounts_repo, usage_repo, sticky_repo)
         self._usage_updater = UsageUpdater(usage_repo, accounts_repo)
 
     def stream_responses(
@@ -62,6 +80,7 @@ class ProxyService:
         *,
         propagate_http_errors: bool = False,
     ) -> AsyncIterator[str]:
+        _maybe_log_proxy_request_shape("stream", payload, headers)
         filtered = filter_inbound_headers(headers)
         return self._stream_with_retry(
             payload,
@@ -74,8 +93,16 @@ class ProxyService:
         payload: ResponsesCompactRequest,
         headers: Mapping[str, str],
     ) -> OpenAIResponsePayload:
+        _maybe_log_proxy_request_shape("compact", payload, headers)
         filtered = filter_inbound_headers(headers)
-        selection = await self._load_balancer.select_account()
+        settings = await self._settings_repo.get_or_create()
+        prefer_earlier_reset = settings.prefer_earlier_reset_accounts
+        sticky_key = _sticky_key_from_compact_payload(payload) if settings.sticky_threads_enabled else None
+        selection = await self._load_balancer.select_account(
+            sticky_key=sticky_key,
+            reallocate_sticky=sticky_key is not None,
+            prefer_earlier_reset_accounts=prefer_earlier_reset,
+        )
         account = selection.account
         if not account:
             raise ProxyResponseError(
@@ -83,7 +110,7 @@ class ProxyService:
                 openai_error("no_accounts", selection.error_message or "No active accounts available"),
             )
         account = await self._ensure_fresh(account)
-        account_id = _header_account_id(account.id)
+        account_id = _header_account_id(account.chatgpt_account_id)
 
         async def _call_compact(target: Account) -> OpenAIResponsePayload:
             access_token = self._encryptor.decrypt(target.access_token_encrypted)
@@ -181,9 +208,15 @@ class ProxyService:
         propagate_http_errors: bool,
     ) -> AsyncIterator[str]:
         request_id = ensure_request_id()
+        settings = await self._settings_repo.get_or_create()
+        prefer_earlier_reset = settings.prefer_earlier_reset_accounts
+        sticky_key = _sticky_key_from_payload(payload) if settings.sticky_threads_enabled else None
         max_attempts = 3
         for attempt in range(max_attempts):
-            selection = await self._load_balancer.select_account()
+            selection = await self._load_balancer.select_account(
+                sticky_key=sticky_key,
+                prefer_earlier_reset_accounts=prefer_earlier_reset,
+            )
             account = selection.account
             if not account:
                 event = response_failed_event(
@@ -281,8 +314,9 @@ class ProxyService:
     ) -> AsyncIterator[str]:
         account_id_value = account.id
         access_token = self._encryptor.decrypt(account.access_token_encrypted)
-        account_id = _header_account_id(account_id_value)
+        account_id = _header_account_id(account.chatgpt_account_id)
         model = payload.model
+        reasoning_effort = payload.reasoning.effort if payload.reasoning else None
         start = time.monotonic()
         status = "success"
         error_code = None
@@ -304,7 +338,11 @@ class ProxyService:
                 return
             event = parse_sse_event(first)
             if event and event.type in ("response.failed", "error"):
-                error = event.response.error if event.type == "response.failed" else event.error
+                if event.type == "response.failed":
+                    response = event.response
+                    error = response.error if response else None
+                else:
+                    error = event.error
                 code = _normalize_error_code(
                     error.code if error else None,
                     error.type if error else None,
@@ -326,7 +364,11 @@ class ProxyService:
                     event_type = event.type
                     if event_type in ("response.failed", "error"):
                         status = "error"
-                        error = event.response.error if event_type == "response.failed" else event.error
+                        if event_type == "response.failed":
+                            response = event.response
+                            error = response.error if response else None
+                        else:
+                            error = event.error
                         error_code = _normalize_error_code(
                             error.code if error else None,
                             error.type if error else None,
@@ -354,27 +396,29 @@ class ProxyService:
             reasoning_tokens = (
                 usage.output_tokens_details.reasoning_tokens if usage and usage.output_tokens_details else None
             )
-            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,
-                )
+            with anyio.CancelScope(shield=True):
+                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,
+                        reasoning_effort=reasoning_effort,
+                        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")
@@ -432,297 +476,107 @@ class ProxyService:
         await self._load_balancer.record_error(account)
 
 
-def _header_account_id(account_id: str | None) -> str | None:
-    if not account_id:
-        return None
-    if account_id.startswith(("email_", "local_")):
-        return None
-    return account_id
-
-
-KNOWN_PLAN_TYPES = {
-    "guest",
-    "free",
-    "go",
-    "plus",
-    "pro",
-    "free_workspace",
-    "team",
-    "business",
-    "education",
-    "quorum",
-    "k12",
-    "enterprise",
-    "edu",
-}
-
-PLAN_TYPE_PRIORITY = (
-    "enterprise",
-    "business",
-    "team",
-    "pro",
-    "plus",
-    "education",
-    "edu",
-    "free_workspace",
-    "free",
-    "go",
-    "guest",
-    "quorum",
-    "k12",
-)
-
-
-def _select_accounts_for_limits(accounts: Iterable[Account]) -> list[Account]:
-    return [account for account in accounts if account.status not in (AccountStatus.DEACTIVATED, AccountStatus.PAUSED)]
-
-
-def _summarize_window(
-    rows: list[UsageWindowRow],
-    account_map: dict[str, Account],
-    window: str,
-) -> UsageWindowSummary | None:
-    if not rows:
-        return None
-    return usage_core.summarize_usage_window(rows, account_map, window)
-
-
-def _window_snapshot(
-    summary: UsageWindowSummary | None,
-    rows: list[UsageWindowRow],
-    window: str,
-    now_epoch: int,
-) -> RateLimitWindowSnapshotData | None:
-    if summary is None:
-        return None
-
-    used_percent = _normalize_used_percent(summary.used_percent, rows)
-    if used_percent is None:
-        return None
-
-    reset_at = summary.reset_at
-    if reset_at is None:
-        return None
+class _RetryableStreamError(Exception):
+    def __init__(self, code: str, error: UpstreamError) -> None:
+        super().__init__(code)
+        self.code = code
+        self.error = error
 
-    window_minutes = summary.window_minutes or usage_core.default_window_minutes(window)
-    if not window_minutes:
-        return None
 
-    limit_window_seconds = int(window_minutes * 60)
-    reset_after_seconds = max(0, int(reset_at) - now_epoch)
+def _maybe_log_proxy_request_shape(
+    kind: str,
+    payload: ResponsesRequest | ResponsesCompactRequest,
+    headers: Mapping[str, str],
+) -> None:
+    settings = get_settings()
+    if not settings.log_proxy_request_shape:
+        return
 
-    return RateLimitWindowSnapshotData(
-        used_percent=_percent_to_int(used_percent),
-        limit_window_seconds=limit_window_seconds,
-        reset_after_seconds=reset_after_seconds,
-        reset_at=int(reset_at),
+    request_id = get_request_id()
+    prompt_cache_key = getattr(payload, "prompt_cache_key", None)
+    if prompt_cache_key is None and payload.model_extra:
+        extra_value = payload.model_extra.get("prompt_cache_key")
+        if isinstance(extra_value, str):
+            prompt_cache_key = extra_value
+    prompt_cache_key_hash = _hash_identifier(prompt_cache_key) if isinstance(prompt_cache_key, str) else None
+    prompt_cache_key_raw = (
+        _truncate_identifier(prompt_cache_key)
+        if settings.log_proxy_request_shape_raw_cache_key and isinstance(prompt_cache_key, str)
+        else None
     )
 
-
-def _normalize_used_percent(
-    value: float | None,
-    rows: Iterable[UsageWindowRow],
-) -> float | None:
-    if value is not None:
-        return value
-    values = [row.used_percent for row in rows if row.used_percent is not None]
-    if not values:
-        return None
-    return sum(values) / len(values)
-
-
-def _percent_to_int(value: float) -> int:
-    bounded = max(0.0, min(100.0, value))
-    return int(bounded)
-
-
-def _rate_limit_details(
-    primary: RateLimitWindowSnapshotData | None,
-    secondary: RateLimitWindowSnapshotData | None,
-) -> RateLimitStatusDetailsData | None:
-    if not primary and not secondary:
-        return None
-    used_percents = [window.used_percent for window in (primary, secondary) if window]
-    limit_reached = any(used >= 100 for used in used_percents)
-    return RateLimitStatusDetailsData(
-        allowed=not limit_reached,
-        limit_reached=limit_reached,
-        primary_window=primary,
-        secondary_window=secondary,
+    extra_keys = sorted(payload.model_extra.keys()) if payload.model_extra else []
+    fields_set = sorted(payload.model_fields_set)
+    input_summary = _summarize_input(payload.input)
+    header_keys = _interesting_header_keys(headers)
+
+    logger.warning(
+        "proxy_request_shape request_id=%s kind=%s model=%s stream=%s input=%s "
+        "prompt_cache_key=%s prompt_cache_key_raw=%s fields=%s extra=%s headers=%s",
+        request_id,
+        kind,
+        payload.model,
+        getattr(payload, "stream", None),
+        input_summary,
+        prompt_cache_key_hash,
+        prompt_cache_key_raw,
+        fields_set,
+        extra_keys,
+        header_keys,
     )
 
 
-def _aggregate_credits(entries: Iterable[UsageHistory]) -> tuple[bool, bool, float] | None:
-    has_data = False
-    has_credits = False
-    unlimited = False
-    balance_total = 0.0
-
-    for entry in entries:
-        credits_has = entry.credits_has
-        credits_unlimited = entry.credits_unlimited
-        credits_balance = entry.credits_balance
-        if credits_has is None and credits_unlimited is None and credits_balance is None:
-            continue
-        has_data = True
-        if credits_has is True:
-            has_credits = True
-        if credits_unlimited is True:
-            unlimited = True
-        if credits_balance is not None and not credits_unlimited:
-            try:
-                balance_total += float(credits_balance)
-            except (TypeError, ValueError):
-                continue
-
-    if not has_data:
-        return None
-    if unlimited:
-        has_credits = True
-    return has_credits, unlimited, balance_total
-
+def _hash_identifier(value: str) -> str:
+    digest = sha256(value.encode("utf-8")).hexdigest()
+    return f"sha256:{digest[:12]}"
 
-def _credits_snapshot(entries: Iterable[UsageHistory]) -> CreditStatusDetailsData | None:
-    aggregate = _aggregate_credits(entries)
-    if aggregate is None:
-        return None
-    has_credits, unlimited, balance_total = aggregate
-    balance_value = str(round(balance_total, 2))
-    return CreditStatusDetailsData(
-        has_credits=has_credits,
-        unlimited=unlimited,
-        balance=balance_value,
-        approx_local_messages=None,
-        approx_cloud_messages=None,
-    )
 
+def _summarize_input(items: Sequence[object]) -> str:
+    if not items:
+        return "0"
+    type_counts: dict[str, int] = {}
+    for item in items:
+        type_name = type(item).__name__
+        type_counts[type_name] = type_counts.get(type_name, 0) + 1
+    summary = ",".join(f"{key}={type_counts[key]}" for key in sorted(type_counts))
+    return f"{len(items)}({summary})"
 
-def _plan_type_for_accounts(accounts: Iterable[Account]) -> str:
-    normalized = [_normalize_plan_type(account.plan_type) for account in accounts]
-    filtered = [plan for plan in normalized if plan is not None]
-    if not filtered:
-        return "guest"
-    unique = set(filtered)
-    if len(unique) == 1:
-        return filtered[0]
-    for plan in PLAN_TYPE_PRIORITY:
-        if plan in unique:
-            return plan
-    return "guest"
 
-
-def _normalize_plan_type(value: str | None) -> str | None:
-    if not value:
-        return None
-    normalized = value.strip().lower()
-    if normalized not in KNOWN_PLAN_TYPES:
-        return None
-    return normalized
-
-
-def _rate_limit_headers(
-    window_label: str,
-    summary: UsageWindowSummary,
-) -> dict[str, str]:
-    used_percent = summary.used_percent
-    window_minutes = summary.window_minutes
-    if used_percent is None or window_minutes is None:
-        return {}
-    headers = {
-        f"x-codex-{window_label}-used-percent": str(float(used_percent)),
-        f"x-codex-{window_label}-window-minutes": str(int(window_minutes)),
-    }
-    reset_at = summary.reset_at
-    if reset_at is not None:
-        headers[f"x-codex-{window_label}-reset-at"] = str(int(reset_at))
-    return headers
-
-
-def _credits_headers(entries: Iterable[UsageHistory]) -> dict[str, str]:
-    aggregate = _aggregate_credits(entries)
-    if aggregate is None:
-        return {}
-    has_credits, unlimited, balance_total = aggregate
-    balance_value = f"{balance_total:.2f}"
-    return {
-        "x-codex-credits-has-credits": "true" if has_credits else "false",
-        "x-codex-credits-unlimited": "true" if unlimited else "false",
-        "x-codex-credits-balance": balance_value,
+def _truncate_identifier(value: str, *, max_length: int = 96) -> str:
+    if len(value) <= max_length:
+        return value
+    return f"{value[:48]}...{value[-16:]}"
+
+
+def _interesting_header_keys(headers: Mapping[str, str]) -> list[str]:
+    allowlist = {
+        "user-agent",
+        "x-request-id",
+        "request-id",
+        "x-openai-client-id",
+        "x-openai-client-version",
+        "x-openai-client-arch",
+        "x-openai-client-os",
+        "x-openai-client-user-agent",
+        "x-codex-session-id",
+        "x-codex-conversation-id",
     }
+    return sorted({key.lower() for key in headers.keys() if key.lower() in allowlist})
 
 
-def _normalize_error_code(code: str | None, error_type: str | None) -> str:
-    value = code or error_type
+def _sticky_key_from_payload(payload: ResponsesRequest) -> str | None:
+    value = payload.prompt_cache_key
     if not value:
-        return "upstream_error"
-    return value.lower()
-
-
-def _parse_openai_error(payload: OpenAIErrorEnvelope) -> OpenAIError | None:
-    error = payload.get("error")
-    if not error:
         return None
-    try:
-        return OpenAIError.model_validate(error)
-    except ValidationError:
-        if not isinstance(error, dict):
-            return None
-        return OpenAIError(
-            message=_coerce_str(error.get("message")),
-            type=_coerce_str(error.get("type")),
-            code=_coerce_str(error.get("code")),
-            param=_coerce_str(error.get("param")),
-            plan_type=_coerce_str(error.get("plan_type")),
-            resets_at=_coerce_number(error.get("resets_at")),
-            resets_in_seconds=_coerce_number(error.get("resets_in_seconds")),
-        )
-
-
-def _coerce_str(value: object) -> str | None:
-    return value if isinstance(value, str) else None
-
-
-def _coerce_number(value: object) -> int | float | None:
-    if isinstance(value, (int, float)):
-        return value
-    if isinstance(value, str):
-        try:
-            return float(value.strip())
-        except ValueError:
-            return None
-    return None
-
-
-def _apply_error_metadata(target: OpenAIErrorDetail, error: OpenAIError | None) -> None:
-    if not error:
-        return
-    if error.plan_type is not None:
-        target["plan_type"] = error.plan_type
-    if error.resets_at is not None:
-        target["resets_at"] = error.resets_at
-    if error.resets_in_seconds is not None:
-        target["resets_in_seconds"] = error.resets_in_seconds
-
-
-class _RetryableStreamError(Exception):
-    def __init__(self, code: str, error: UpstreamError) -> None:
-        super().__init__(code)
-        self.code = code
-        self.error = error
+    stripped = value.strip()
+    return stripped or None
 
 
-def _upstream_error_from_openai(error: OpenAIError | None) -> UpstreamError:
-    if not error:
-        return {}
-    data = error.model_dump(exclude_none=True)
-    payload: UpstreamError = {}
-    message = data.get("message")
-    if isinstance(message, str):
-        payload["message"] = message
-    resets_at = data.get("resets_at")
-    if isinstance(resets_at, (int, float)):
-        payload["resets_at"] = resets_at
-    resets_in_seconds = data.get("resets_in_seconds")
-    if isinstance(resets_in_seconds, (int, float)):
-        payload["resets_in_seconds"] = resets_in_seconds
-    return payload
+def _sticky_key_from_compact_payload(payload: ResponsesCompactRequest) -> str | None:
+    if not payload.model_extra:
+        return None
+    value = payload.model_extra.get("prompt_cache_key")
+    if not isinstance(value, str):
+        return None
+    stripped = value.strip()
+    return stripped or None

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(),
+            },
+        )