codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/proxy/_service/websocket/protocol.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+from app.core.clock import Clock, Scheduler
+
+
+class _WebSocketServiceProtocol(Protocol):
+    _capability_router: Any
+    _background_cleanup_tasks: Any
+    _clock: Clock
+    _remaining_budget_seconds: Any
+    _acquire_account_response_create_lease_or_overload: Any
+    _acquire_request_state_response_create_admission: Any
+    _cancel_request_state_api_key_reservation_heartbeat: Any
+    _connect_proxy_websocket: Any
+    _decide_websocket_failover_action: Any
+    _downstream_websocket_is_idle: Any
+    _drain_deferred_account_error_backoffs: Any
+    _drain_deferred_keyed_stream_health: Any
+    _emit_pending_websocket_keepalive: Any
+    _emit_websocket_connect_failure: Any
+    _emit_websocket_connect_timeout: Any
+    _emit_websocket_proxy_request_timeout: Any
+    _emit_websocket_security_work_missing_pool: Any
+    _emit_websocket_terminal_error: Any
+    _encryptor: Any
+    _ensure_fresh_with_budget: Any
+    _fail_expired_pending_websocket_requests: Any
+    _fail_pending_websocket_requests: Any
+    _finalize_websocket_request_state: Any
+    _get_work_admission: Any
+    _handle_precreated_websocket_auth_failure: Any
+    _handle_stream_error: Any
+    _handle_websocket_connect_error: Any
+    _load_balancer: Any
+    _maybe_touch_request_state_api_key_reservation: Any
+    _next_websocket_receive_timeout: Any
+    _open_upstream_websocket: Any
+    _open_upstream_websocket_with_budget: Any
+    _prepare_response_bridge_request_state: Any
+    _prepare_websocket_response_create_request: Any
+    _process_upstream_websocket_text: Any
+    _raise_for_unsupported_input_image_references: Any
+    _refresh_websocket_api_key_policy: Any
+    _revalidate_open_websocket_account: Any
+    _relay_upstream_websocket_messages: Any
+    _release_websocket_request_state_reservation: Any
+    _release_websocket_reservation: Any
+    _release_request_state_account_response_create_lease: Any
+    _remember_websocket_previous_response_owner: Any
+    _remember_websocket_previous_response_owner_miss: Any
+    _repo_factory: Any
+    _reserve_websocket_api_key_usage: Any
+    _resolve_file_account_for_responses: Any
+    _resolve_compact_turn_state_owner: Any
+    _resolve_upstream_route_for_account: Any
+    _resolve_websocket_previous_response_owner: Any
+    _retry_websocket_connect_after_401: Any
+    _scheduler: Scheduler
+    _select_account_with_budget_compatible: Any
+    _select_websocket_connect_account: Any
+    _send_downstream_websocket_bytes: Any
+    _send_downstream_websocket_text: Any
+    _settle_stream_api_key_usage: Any
+    _start_request_state_api_key_reservation_heartbeat: Any
+    _try_open_websocket_connect_attempt: Any
+    _touch_active_websocket_thread_affinity: Any
+    _websocket_continuity_index: Any
+    _websocket_continuity_state_for_request: Any
+    _websocket_previous_response_account_index: Any
+    _write_request_log: Any
+    _write_websocket_connect_failure: Any
+    proxy_responses_websocket: Any

app/modules/proxy/account_cache.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import anyio
+from sqlalchemy import select
+
+from app.core.cache.invalidation import (
+    NAMESPACE_ACCOUNT_ROUTING,
+    NAMESPACE_ACCOUNT_SELECTION,
+    get_cache_invalidation_poller,
+)
+from app.db.models import Account, AccountStatus
+from app.db.session import SessionLocal, close_session
+
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncSession
+
+    from app.modules.proxy.load_balancer import SelectionInputs
+
+_AssignedAccountsKey = tuple[str, ...] | None
+_CacheKey = tuple[str | None, str | None, str | None, str, _AssignedAccountsKey]
+
+
+@dataclass(slots=True)
+class _CachedSelectionInputs:
+    data: SelectionInputs
+    expires_at: float
+
+
+class AccountSelectionCache:
+    def __init__(self, ttl_seconds: int | None = None) -> None:
+        if ttl_seconds is None:
+            import sys
+
+            ttl_seconds = 0 if "pytest" in sys.modules else 5
+        if ttl_seconds < 0:
+            raise ValueError("ttl_seconds must be non-negative")
+        self._ttl_seconds = ttl_seconds
+        self._cache: dict[_CacheKey, _CachedSelectionInputs] = {}
+        self._lock = anyio.Lock()
+        self._generation: int = 0
+
+    @property
+    def generation(self) -> int:
+        return self._generation
+
+    async def get(self, key: _CacheKey = (None, None, None, "", None)) -> SelectionInputs | None:
+        if self._ttl_seconds == 0:
+            return None
+        entry = self._cache.get(key)
+        if entry is None:
+            return None
+        if time.monotonic() >= entry.expires_at:
+            return None
+        return entry.data
+
+    async def set(
+        self,
+        data: SelectionInputs,
+        key: _CacheKey = (None, None, None, "", None),
+        *,
+        generation: int | None = None,
+    ) -> None:
+        async with self._lock:
+            if generation is not None and generation != self._generation:
+                return
+            self._cache[key] = _CachedSelectionInputs(
+                data=data,
+                expires_at=time.monotonic() + self._ttl_seconds,
+            )
+
+    def invalidate(self, *, propagate: bool = True) -> None:
+        """Invalidate the local cache and, unless ``propagate`` is False, enqueue a
+        coalesced cross-replica ``account_selection`` bump.
+
+        The cache-invalidation poller callback registers ``propagate=False`` so a
+        remote bump never re-bumps (feedback-loop prevention).
+        """
+        self._generation += 1
+        self._cache.clear()
+        if propagate:
+            poller = get_cache_invalidation_poller()
+            if poller is not None:
+                poller.request_bump(NAMESPACE_ACCOUNT_SELECTION)
+
+
+_ROUTING_UNAVAILABLE_STATUSES = frozenset(
+    {
+        AccountStatus.PAUSED,
+        AccountStatus.DEACTIVATED,
+    }
+)
+
+
+class RoutingAvailabilityCache:
+    """Cluster-coherent view of which accounts are unavailable for routing.
+
+    The cache keeps a snapshot of committed account statuses (``{account_id: status}``)
+    seeded at poller start and rebuilt on every ``account_routing`` bump. An account is
+    routing-unavailable when its committed status is PAUSED / DEACTIVATED, or the id
+    is absent from the snapshot (deleted), or a local mark
+    overlay entry exists (covering the same-replica window between a mark and the
+    snapshot rebuild). RATE_LIMITED and QUOTA_EXCEEDED deliberately do NOT map to
+    unavailable, preserving cooldown-state bridge-session reuse.
+
+    When the snapshot is unseeded (unit tests, poller not running) the cache degrades
+    to the historical process-local set semantics.
+    """
+
+    def __init__(self, session_factory: Callable[[], AsyncSession] | None = None) -> None:
+        self._session_factory = session_factory
+        self._snapshot: dict[str, AccountStatus] | None = None
+        self._local_marks: set[str] = set()
+
+    @property
+    def seeded(self) -> bool:
+        return self._snapshot is not None
+
+    def mark_unavailable(self, account_id: str) -> None:
+        self._local_marks.add(account_id)
+        _request_account_routing_bump()
+
+    def clear_unavailable(self, account_id: str) -> None:
+        self._local_marks.discard(account_id)
+        if self._snapshot is not None:
+            self._snapshot[account_id] = AccountStatus.ACTIVE
+        _request_account_routing_bump()
+
+    def is_unavailable(self, account_id: str) -> bool:
+        if account_id in self._local_marks:
+            return True
+        snapshot = self._snapshot
+        if snapshot is None:
+            return False
+        status = snapshot.get(account_id)
+        return status is None or status in _ROUTING_UNAVAILABLE_STATUSES
+
+    async def refresh_from_db(self) -> None:
+        """Rebuild the snapshot from committed account statuses.
+
+        Local overlay marks whose committed status became routable again are dropped —
+        this is what lets a reactivation or re-authentication served by another replica
+        clear this replica's marker without a restart. Only marks that already existed
+        when this refresh started are eligible to be dropped: a mark added while the
+        SELECT is in flight may not be reflected in the rows it read (the status commit
+        can land after the read), so filtering it against that snapshot would silently
+        lose the mark. Such marks are preserved and re-evaluated by the next refresh,
+        which the mark's own queued ``account_routing`` bump guarantees.
+
+        Database errors propagate to the caller: when invoked as an
+        ``account_routing`` invalidation callback the poller then leaves the
+        namespace version unacknowledged and retries on the next poll cycle, so
+        a transient failure cannot make a replica permanently miss a pause,
+        deletion, or deactivation.
+        """
+        marks_before_refresh = frozenset(self._local_marks)
+        factory = self._session_factory or SessionLocal
+        session = factory()
+        try:
+            result = await session.execute(select(Account.id, Account.status))
+            snapshot: dict[str, AccountStatus] = {account_id: status for account_id, status in result.all()}
+        finally:
+            await close_session(session)
+        self._snapshot = snapshot
+        self._local_marks = {
+            account_id
+            for account_id in self._local_marks
+            if account_id not in marks_before_refresh
+            or (status := snapshot.get(account_id)) is None
+            or status in _ROUTING_UNAVAILABLE_STATUSES
+        }
+
+    def reset(self) -> None:
+        """Drop all state (snapshot back to unseeded). Test isolation helper."""
+        self._snapshot = None
+        self._local_marks.clear()
+
+
+_account_selection_cache = AccountSelectionCache()
+_routing_availability_cache = RoutingAvailabilityCache()
+
+
+def get_account_selection_cache() -> AccountSelectionCache:
+    return _account_selection_cache
+
+
+def get_routing_availability_cache() -> RoutingAvailabilityCache:
+    return _routing_availability_cache
+
+
+def _request_account_routing_bump() -> None:
+    poller = get_cache_invalidation_poller()
+    if poller is not None:
+        poller.request_bump(NAMESPACE_ACCOUNT_ROUTING)
+
+
+def mark_account_routing_unavailable(account_id: str) -> None:
+    _routing_availability_cache.mark_unavailable(account_id)
+
+
+def clear_account_routing_unavailable(account_id: str) -> None:
+    _routing_availability_cache.clear_unavailable(account_id)
+
+
+def clear_all_account_routing_unavailable() -> None:
+    _routing_availability_cache.reset()
+
+
+def is_account_routing_unavailable(account_id: str) -> bool:
+    return _routing_availability_cache.is_unavailable(account_id)
+
+
+async def propagate_account_routing_change() -> bool:
+    """Durably bump the ``account_routing`` namespace before returning.
+
+    Used by API-endpoint mutation paths (pause/reactivate/delete, OAuth re-auth,
+    proxy-binding reactivation) so the cross-replica signal is written before the
+    HTTP response returns. Returns False when no poller is wired or the bump failed
+    after retries; the coalesced bump enqueued by ``mark_``/``clear_`` remains the
+    fallback path.
+    """
+    poller = get_cache_invalidation_poller()
+    if poller is None:
+        return False
+    return await poller.bump(NAMESPACE_ACCOUNT_ROUTING)

app/modules/proxy/account_eligibility.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import time
+from collections.abc import Collection
+
+from app.core.auth import token_expiry_epoch_ms
+from app.core.crypto import TokenEncryptor
+from app.db.models import Account, AccountStatus
+
+ROUTABLE_STATUSES = (AccountStatus.ACTIVE, AccountStatus.REAUTH_REQUIRED)
+"""Statuses whose sessions keep serving requests (reauth stays routable until expiry)."""
+
+HARD_OWNER_UNAVAILABLE_STATUSES = frozenset(
+    {
+        AccountStatus.PAUSED,
+        AccountStatus.DEACTIVATED,
+        AccountStatus.RATE_LIMITED,
+        AccountStatus.QUOTA_EXCEEDED,
+        AccountStatus.REAUTH_REQUIRED,
+    }
+)
+"""Statuses a hard continuity owner can sit in while it cannot serve its thread.
+
+Usable directly in SQL, unlike :func:`reauth_access_token_is_expired`, which
+needs a decrypted token. ``REAUTH_REQUIRED`` belongs here only because every
+consumer pairs the set with a long inactivity grace: an owner that has been
+warning about reauthentication for hours without serving its thread is not
+coming back on its own, whatever its stored token claims. Anything deciding
+routability *now* must use :data:`ROUTABLE_STATUSES` plus the expiry check
+instead.
+"""
+
+
+def stored_access_token_expires_at(
+    encrypted_access_token: bytes | None,
+    encryptor: TokenEncryptor,
+) -> float | None:
+    """Return the stored access token's JWT expiry in epoch seconds when known."""
+    if not encrypted_access_token:
+        return None
+    try:
+        access_token = encryptor.decrypt(encrypted_access_token)
+    except Exception:
+        return None
+    expires_ms = token_expiry_epoch_ms(access_token)
+    return None if expires_ms is None else expires_ms / 1000.0
+
+
+def account_access_token_expires_at(account: Account, encryptor: TokenEncryptor) -> float | None:
+    """Return an account snapshot's known access-token expiry."""
+    encrypted_access_token = getattr(account, "access_token_encrypted", None)
+    return stored_access_token_expires_at(encrypted_access_token, encryptor)
+
+
+def reauth_access_token_is_expired(
+    status: AccountStatus,
+    access_token_expires_at: float | None,
+    *,
+    now: float | None = None,
+) -> bool:
+    """Return whether a reauthentication-warning account is known unusable."""
+    return (
+        status == AccountStatus.REAUTH_REQUIRED
+        and access_token_expires_at is not None
+        and access_token_expires_at <= (time.time() if now is None else now)
+    )
+
+
+def all_accounts_require_reauthentication(
+    accounts: Collection[Account],
+    encryptor: TokenEncryptor,
+    *,
+    now: float,
+) -> bool:
+    """Return whether every candidate is reauthentication-blocked by known expiry.
+
+    ``now`` is the caller's epoch clock sample (the balancer passes
+    ``self._clock.time()``) so selection never mixes clock domains.
+    """
+    return bool(accounts) and all(
+        reauth_access_token_is_expired(
+            account.status,
+            account_access_token_expires_at(account, encryptor),
+            now=now,
+        )
+        for account in accounts
+    )

app/modules/proxy/additional_model_limits.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from app.modules.usage.additional_quota_keys import (
+    get_additional_display_label_for_quota_key,
+    get_additional_quota_definition_for_model,
+    get_additional_quota_key_for_model,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class AdditionalModelLimit:
+    model: str
+    quota_key: str
+    display_label: str
+
+
+def get_additional_model_limit(model: str | None) -> AdditionalModelLimit | None:
+    resolved = get_additional_quota_definition_for_model(model)
+    if resolved is None:
+        return None
+    normalized_model = model.strip().lower() if model is not None else None
+    if normalized_model is None:
+        return None
+    return AdditionalModelLimit(
+        model=normalized_model,
+        quota_key=resolved.quota_key,
+        display_label=resolved.display_label,
+    )
+
+
+def get_additional_quota_key_for_model_id(model: str | None) -> str | None:
+    resolved = get_additional_model_limit(model)
+    return resolved.quota_key if resolved is not None else None
+
+
+def get_additional_display_label_for_model(model: str | None) -> str | None:
+    quota_key = get_additional_quota_key_for_model(model)
+    return get_additional_display_label_for_quota_key(quota_key)