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

app/modules/accounts/auth_manager.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+import logging
 from datetime import datetime
 from typing import Protocol
 
-from app.core.auth import DEFAULT_PLAN
+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
@@ -29,9 +30,13 @@ class AccountsRepositoryPort(Protocol):
         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
@@ -39,8 +44,8 @@ class AuthManager:
 
     async def ensure_fresh(self, account: Account, *, force: bool = False) -> Account:
         if force or should_refresh(account.last_refresh):
-            return await self.refresh_account(account)
-        return account
+            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)
@@ -58,6 +63,8 @@ class AuthManager:
         account.refresh_token_encrypted = self._encryptor.encrypt(result.refresh_token)
         account.id_token_encrypted = self._encryptor.encrypt(result.id_token)
         account.last_refresh = utcnow()
+        if result.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,
@@ -76,5 +83,39 @@ class AuthManager:
             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(getattr(result, "rowcount", 0) or 0)
+        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(getattr(result, "rowcount", 0) or 0)
+        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(getattr(result, "rowcount", 0) or 0)
+        return result.scalar_one_or_none() is not None

app/modules/accounts/service.py

@@ -9,7 +9,7 @@ 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
@@ -66,12 +66,14 @@ class AccountsService:
         claims = claims_from_auth(auth)
 
         email = claims.email or DEFAULT_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)
-        account_id = claims.account_id or fallback_account_id(email)
         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),

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,
@@ -294,16 +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
+        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_id = account_id or fallback_account_id(email)
 
         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),

app/modules/proxy/load_balancer.py

@@ -6,6 +6,7 @@ from typing import Iterable
 
 from app.core.balancer import (
     AccountState,
+    SelectionResult,
     handle_permanent_failure,
     handle_quota_exceeded,
     handle_rate_limit,
@@ -15,6 +16,7 @@ from app.core.balancer.types import UpstreamError
 from app.core.usage.quota import apply_usage_quota
 from app.db.models import Account, UsageHistory
 from app.modules.accounts.repository import AccountsRepository
+from app.modules.proxy.sticky_repository import StickySessionsRepository
 from app.modules.usage.repository import UsageRepository
 from app.modules.usage.updater import UsageUpdater
 
@@ -35,13 +37,25 @@ class AccountSelection:
 
 
 class LoadBalancer:
-    def __init__(self, accounts_repo: AccountsRepository, usage_repo: UsageRepository) -> None:
+    def __init__(
+        self,
+        accounts_repo: AccountsRepository,
+        usage_repo: UsageRepository,
+        sticky_repo: StickySessionsRepository | None = None,
+    ) -> None:
         self._accounts_repo = accounts_repo
         self._usage_repo = usage_repo
         self._usage_updater = UsageUpdater(usage_repo, accounts_repo)
+        self._sticky_repo = sticky_repo
         self._runtime: dict[str, RuntimeState] = {}
 
-    async def select_account(self) -> AccountSelection:
+    async def select_account(
+        self,
+        sticky_key: str | None = None,
+        *,
+        reallocate_sticky: bool = False,
+        prefer_earlier_reset_accounts: bool = False,
+    ) -> AccountSelection:
         accounts = await self._accounts_repo.list_accounts()
         latest_primary = await self._usage_repo.latest_by_account()
         await self._usage_updater.refresh_accounts(accounts, latest_primary)
@@ -55,7 +69,13 @@ class LoadBalancer:
             runtime=self._runtime,
         )
 
-        result = select_account(states)
+        result = await self._select_with_stickiness(
+            states=states,
+            account_map=account_map,
+            sticky_key=sticky_key,
+            reallocate_sticky=reallocate_sticky,
+            prefer_earlier_reset_accounts=prefer_earlier_reset_accounts,
+        )
         for state in states:
             account = account_map.get(state.account_id)
             if account:
@@ -74,6 +94,39 @@ class LoadBalancer:
             return AccountSelection(account=None, error_message=result.error_message)
         return AccountSelection(account=selected, error_message=None)
 
+    async def _select_with_stickiness(
+        self,
+        *,
+        states: list[AccountState],
+        account_map: dict[str, Account],
+        sticky_key: str | None,
+        reallocate_sticky: bool,
+        prefer_earlier_reset_accounts: bool,
+    ) -> SelectionResult:
+        if not sticky_key or not self._sticky_repo:
+            return select_account(states, prefer_earlier_reset=prefer_earlier_reset_accounts)
+
+        if reallocate_sticky:
+            chosen = select_account(states, prefer_earlier_reset=prefer_earlier_reset_accounts)
+            if chosen.account is not None and chosen.account.account_id in account_map:
+                await self._sticky_repo.upsert(sticky_key, chosen.account.account_id)
+            return chosen
+
+        existing = await self._sticky_repo.get_account_id(sticky_key)
+        if existing:
+            pinned = next((state for state in states if state.account_id == existing), None)
+            if pinned is None:
+                await self._sticky_repo.delete(sticky_key)
+            else:
+                pinned_result = select_account([pinned], prefer_earlier_reset=prefer_earlier_reset_accounts)
+                if pinned_result.account is not None:
+                    return pinned_result
+
+        chosen = select_account(states, prefer_earlier_reset=prefer_earlier_reset_accounts)
+        if chosen.account is not None and chosen.account.account_id in account_map:
+            await self._sticky_repo.upsert(sticky_key, chosen.account.account_id)
+        return chosen
+
     async def mark_rate_limit(self, account: Account, error: UpstreamError) -> None:
         state = self._state_for(account)
         handle_rate_limit(state, error)
@@ -103,6 +156,8 @@ class LoadBalancer:
             used_percent=None,
             reset_at=runtime.reset_at,
             cooldown_until=runtime.cooldown_until,
+            secondary_used_percent=None,
+            secondary_reset_at=None,
             last_error_at=runtime.last_error_at,
             last_selected_at=runtime.last_selected_at,
             error_count=runtime.error_count,
@@ -116,14 +171,21 @@ class LoadBalancer:
         runtime.last_error_at = state.last_error_at
         runtime.error_count = state.error_count
 
-        if account.status != state.status or account.deactivation_reason != state.deactivation_reason:
+        reset_at_int = int(state.reset_at) if state.reset_at else None
+        status_changed = account.status != state.status
+        reason_changed = account.deactivation_reason != state.deactivation_reason
+        reset_changed = account.reset_at != reset_at_int
+
+        if status_changed or reason_changed or reset_changed:
             await self._accounts_repo.update_status(
                 account.id,
                 state.status,
                 state.deactivation_reason,
+                reset_at_int,
             )
             account.status = state.status
             account.deactivation_reason = state.deactivation_reason
+            account.reset_at = reset_at_int
 
 
 def _build_states(
@@ -161,12 +223,17 @@ def _state_from_account(
     secondary_used = secondary_entry.used_percent if secondary_entry else None
     secondary_reset = secondary_entry.reset_at if secondary_entry else None
 
+    # Use account.reset_at from DB as the authoritative source for runtime reset
+    # This survives across requests since LoadBalancer is instantiated per-request
+    db_reset_at = float(account.reset_at) if account.reset_at else None
+    effective_runtime_reset = db_reset_at or runtime.reset_at
+
     status, used_percent, reset_at = apply_usage_quota(
         status=account.status,
         primary_used=primary_used,
         primary_reset=primary_reset,
         primary_window_minutes=primary_window_minutes,
-        runtime_reset=runtime.reset_at,
+        runtime_reset=effective_runtime_reset,
         secondary_used=secondary_used,
         secondary_reset=secondary_reset,
     )
@@ -177,6 +244,8 @@ def _state_from_account(
         used_percent=used_percent,
         reset_at=reset_at,
         cooldown_until=runtime.cooldown_until,
+        secondary_used_percent=secondary_used,
+        secondary_reset_at=secondary_reset,
         last_error_at=runtime.last_error_at,
         last_selected_at=runtime.last_selected_at,
         error_count=runtime.error_count,

app/modules/proxy/service.py

@@ -2,9 +2,13 @@ from __future__ import annotations
 
 import logging
 import time
+from collections.abc import Sequence
 from datetime import timedelta
+from hashlib import sha256
 from typing import AsyncIterator, Mapping
 
+import anyio
+
 from app.core import usage as usage_core
 from app.core.auth.refresh import RefreshError
 from app.core.balancer import PERMANENT_FAILURE_CODES
@@ -12,13 +16,14 @@ 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 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
-from app.core.utils.request_id import ensure_request_id
+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, UsageHistory
@@ -40,8 +45,10 @@ from app.modules.proxy.helpers import (
     _window_snapshot,
 )
 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
 
@@ -54,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(
@@ -70,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,
@@ -82,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(
@@ -91,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)
@@ -189,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(
@@ -289,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
@@ -370,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")
@@ -436,12 +464,9 @@ class ProxyService:
         await self._handle_stream_error(account, _upstream_error_from_openai(error), code)
 
     async def _handle_stream_error(self, account: Account, error: UpstreamError, code: str) -> None:
-        if code == "rate_limit_exceeded":
+        if code in {"rate_limit_exceeded", "usage_limit_reached"}:
             await self._load_balancer.mark_rate_limit(account, error)
             return
-        if code == "usage_limit_reached":
-            await self._load_balancer.mark_quota_exceeded(account, error)
-            return
         if code in {"insufficient_quota", "usage_not_included", "quota_exceeded"}:
             await self._load_balancer.mark_quota_exceeded(account, error)
             return
@@ -456,3 +481,102 @@ class _RetryableStreamError(Exception):
         super().__init__(code)
         self.code = code
         self.error = error
+
+
+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
+
+    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
+    )
+
+    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 _hash_identifier(value: str) -> str:
+    digest = sha256(value.encode("utf-8")).hexdigest()
+    return f"sha256:{digest[:12]}"
+
+
+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 _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 _sticky_key_from_payload(payload: ResponsesRequest) -> str | None:
+    value = payload.prompt_cache_key
+    if not value:
+        return None
+    stripped = value.strip()
+    return stripped or None
+
+
+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