codexLB-up 1.25.0b10__py3-none-any.whl

app/core/balancer/rendezvous_hash.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from hashlib import sha256
+
+
+def select_node(key: str, nodes: Sequence[str]) -> str | None:
+    """Rendezvous (HRW) hashing: returns node with highest sha256(key + node) hash.
+
+    Property: adding/removing 1 node remaps only 1/N keys (vs N-1/N for modulo).
+    Time complexity: O(N) where N = number of nodes.
+    """
+    if not nodes:
+        return None
+    if len(nodes) == 1:
+        return nodes[0]
+
+    def _score(node: str) -> bytes:
+        return sha256(f"{key}:{node}".encode()).digest()
+
+    return max(nodes, key=_score)

app/core/balancer/types.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from typing import Literal, TypedDict
+
+
+class UpstreamError(TypedDict, total=False):
+    message: str
+    resets_at: int | float
+    resets_in_seconds: int | float
+
+
+FailureClass = Literal["rate_limit", "quota", "retryable_transient", "account_unavailable", "non_retryable"]
+FailurePhase = Literal["connect", "first_event", "mid_stream"]
+
+
+class ClassifiedFailure(TypedDict):
+    failure_class: FailureClass
+    phase: FailurePhase
+    error_code: str
+    error: UpstreamError
+    http_status: int | None

app/core/bootstrap.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import hashlib
+import logging
+import secrets
+from dataclasses import dataclass
+from secrets import compare_digest
+
+from app.core.auth.dashboard_users_cache import get_dashboard_users_cache
+from app.core.config.settings import get_settings
+from app.core.config.settings_cache import get_settings_cache
+from app.core.crypto import TokenEncryptor
+from app.db.session import SessionLocal
+from app.modules.dashboard_auth.repository import DashboardAuthRepository
+
+logger = logging.getLogger(__name__)
+_encryptor: TokenEncryptor | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class SharedBootstrapState:
+    """The cross-replica facts the bootstrap token depends on, read uncached.
+
+    ``local_password_configured`` is derived from ``dashboard_users`` (an active
+    account holding a password), never from the legacy settings column.
+    """
+
+    local_password_configured: bool
+    bootstrap_token_encrypted: bytes | None
+    bootstrap_token_hash: bytes | None
+
+
+def _get_manual_bootstrap_token() -> str | None:
+    manual = (get_settings().dashboard_bootstrap_token or "").strip()
+    return manual or None
+
+
+def _hash_bootstrap_token(token: str) -> bytes:
+    return hashlib.sha256(token.encode("utf-8")).digest()
+
+
+def _get_encryptor() -> TokenEncryptor:
+    global _encryptor
+    if _encryptor is None:
+        _encryptor = TokenEncryptor()
+    return _encryptor
+
+
+def log_bootstrap_token(logger: logging.Logger, token: str, *, reason: str = "first-run") -> None:
+    # Emit at WARNING so the one-time token is surfaced regardless of the
+    # container/root logger level (which defaults to WARNING in most docker
+    # setups, silently dropping an INFO log and leaving operators unable to
+    # find the token). See #458.
+    logger.warning(
+        "\n"
+        "============================================\n"
+        "  Dashboard bootstrap token (%s):\n"
+        "  %s\n"
+        "\n"
+        "  Use this token for initial remote setup.\n"
+        "  It is shared across replicas and stays\n"
+        "  valid until a password is set.\n"
+        "============================================",
+        reason,
+        token,
+    )
+
+
+async def _get_shared_bootstrap_state() -> SharedBootstrapState:
+    async with SessionLocal() as session:
+        repository = DashboardAuthRepository(session)
+        settings = await repository.get_settings()
+        auth_state = await repository.get_local_auth_state()
+        return SharedBootstrapState(
+            local_password_configured=auth_state.active_local_password_users > 0,
+            bootstrap_token_encrypted=settings.bootstrap_token_encrypted,
+            bootstrap_token_hash=settings.bootstrap_token_hash,
+        )
+
+
+async def has_active_bootstrap_token() -> bool:
+    manual = _get_manual_bootstrap_token()
+    if manual:
+        return True
+
+    state = await _get_shared_bootstrap_state()
+    return not state.local_password_configured and state.bootstrap_token_hash is not None
+
+
+async def validate_bootstrap_token(submitted_token: str) -> bool:
+    manual = _get_manual_bootstrap_token()
+    if manual is not None:
+        return compare_digest(submitted_token.encode("utf-8"), manual.encode("utf-8"))
+
+    state = await _get_shared_bootstrap_state()
+    if state.local_password_configured or state.bootstrap_token_hash is None:
+        return False
+    return compare_digest(_hash_bootstrap_token(submitted_token), state.bootstrap_token_hash)
+
+
+async def get_bootstrap_validation_status(submitted_token: str) -> str:
+    manual = _get_manual_bootstrap_token()
+    if manual is not None:
+        if compare_digest(submitted_token.encode("utf-8"), manual.encode("utf-8")):
+            return "valid"
+        return "invalid"
+
+    state = await _get_shared_bootstrap_state()
+    if state.bootstrap_token_hash is None:
+        return "password_already_configured" if state.local_password_configured else "unavailable"
+    if compare_digest(_hash_bootstrap_token(submitted_token), state.bootstrap_token_hash):
+        return "valid"
+    if state.local_password_configured:
+        return "password_already_configured"
+    return "invalid"
+
+
+async def _invalidate_bootstrap_caches() -> None:
+    await get_settings_cache().invalidate()
+    await get_dashboard_users_cache().invalidate()
+
+
+async def ensure_auto_bootstrap_token() -> str | None:
+    manual = _get_manual_bootstrap_token()
+
+    async with SessionLocal() as session:
+        repository = DashboardAuthRepository(session)
+        settings = await repository.get_settings()
+        local_password_configured = (await repository.get_local_auth_state()).active_local_password_users > 0
+
+        if manual or local_password_configured:
+            if settings.bootstrap_token_hash is not None:
+                await repository.clear_bootstrap_token()
+                await _invalidate_bootstrap_caches()
+            return None
+
+        if settings.bootstrap_token_hash is not None:
+            encrypted = settings.bootstrap_token_encrypted
+            if encrypted is None:
+                logger.warning(
+                    "Stored bootstrap token hash exists without encrypted token; keeping existing hash untouched"
+                )
+                return None
+            try:
+                return _get_encryptor().decrypt(encrypted)
+            except Exception:
+                logger.warning(
+                    "Stored bootstrap token could not be decrypted; leaving existing token valid. "
+                    "Configure CODEX_LB_DASHBOARD_BOOTSTRAP_TOKEN or restore a shared encryption key "
+                    "to recover it.",
+                    exc_info=True,
+                )
+                return None
+
+        token = secrets.token_urlsafe(32)
+        stored = await repository.store_bootstrap_token_if_absent(
+            _get_encryptor().encrypt(token),
+            _hash_bootstrap_token(token),
+        )
+
+    await _invalidate_bootstrap_caches()
+    if stored:
+        return token
+    return None
+
+
+async def clear_auto_generated_token() -> None:
+    async with SessionLocal() as session:
+        repository = DashboardAuthRepository(session)
+        cleared = await repository.clear_bootstrap_token()
+    if cleared:
+        await _invalidate_bootstrap_caches()

app/core/cache/invalidation.py

@@ -0,0 +1,426 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Awaitable, Callable
+from inspect import isawaitable
+
+from sqlalchemy import select, update
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.metrics.prometheus import (
+    PROMETHEUS_AVAILABLE,
+    cache_invalidation_bump_failures_total,
+    cache_invalidation_poll_failures_total,
+)
+from app.db.models import CacheInvalidation
+from app.db.session import close_session
+
+logger = logging.getLogger(__name__)
+
+NAMESPACE_API_KEY = "api_key"
+NAMESPACE_FIREWALL = "firewall"
+NAMESPACE_ACCOUNT_ROUTING = "account_routing"
+NAMESPACE_ACCOUNT_SELECTION = "account_selection"
+NAMESPACE_SETTINGS = "settings"
+NAMESPACE_RESET_CREDITS = "reset_credits"
+NAMESPACE_MODEL_REGISTRY = "model_registry"
+NAMESPACE_UPSTREAM_ROUTE = "upstream_route"
+NAMESPACE_DASHBOARD_USERS = "dashboard_users"
+# Callback return values are ignored; awaitables are awaited for their side
+# effects only, so callbacks may return a status (e.g. bool) for other callers.
+type InvalidationCallback = Callable[[], object | Awaitable[object]]
+
+# Log-safe labels for namespace values. Static analyzers (CodeQL) classify the
+# NAMESPACE_API_KEY constant as credential-like from its name alone, so log
+# statements must not take dataflow from the constants. These literals mirror
+# the namespace values; test_namespace_log_labels_cover_all_namespaces keeps
+# them in sync.
+_NAMESPACE_LOG_LABELS: dict[str, str] = {
+    "api_key": "api_key",
+    "firewall": "firewall",
+    "account_routing": "account_routing",
+    "account_selection": "account_selection",
+    "settings": "settings",
+    "reset_credits": "reset_credits",
+    "model_registry": "model_registry",
+    "upstream_route": "upstream_route",
+    "dashboard_users": "dashboard_users",
+}
+
+_BUMP_RETRY_ATTEMPTS = 3
+_BUMP_RETRY_BASE_SECONDS = 0.05
+_POLL_FAILURES_WARNING_THRESHOLD = 3
+_POLL_FAILURES_ERROR_THRESHOLD = 10
+
+
+class CacheInvalidationPoller:
+    def __init__(
+        self,
+        session_factory: Callable[[], AsyncSession],
+        poll_interval_seconds: float = 0.5,
+    ) -> None:
+        self._session_factory = session_factory
+        self._poll_interval = poll_interval_seconds
+        self._known_versions: dict[str, int] = {}
+        self._callbacks: dict[str, list[InvalidationCallback]] = {}
+        self._pending_bumps: set[str] = set()
+        self._consecutive_poll_failures = 0
+        self._poll_initialized = False
+        self._task: asyncio.Task[None] | None = None
+        self._stop = asyncio.Event()
+
+    def on_invalidation(self, namespace: str, callback: InvalidationCallback) -> None:
+        self._callbacks.setdefault(namespace, []).append(callback)
+
+    async def initialize(self) -> None:
+        """Seed baseline namespace versions before the process serves traffic.
+
+        MUST run at lifespan startup before local caches / routing snapshots are
+        loaded and before the first poll cycle. It records the current version of
+        every ``cache_invalidation`` row as the baseline and marks the poller
+        initialized, so that any bump a peer commits after this point is observed
+        as a change (fires callbacks) rather than acknowledged as pre-existing
+        state. Without this seed the first poll would treat every existing version
+        row as a baseline and silently drop a peer bump that landed between local
+        cache load and the first poll, leaving security/selection/routing state
+        stale until the fallback TTL or a later bump.
+
+        On success ``_poll_initialized`` is set to ``True``. If the read fails the
+        method raises with state unchanged (baseline empty, ``_poll_initialized``
+        still ``False``), so the caller can retry before background polling. If
+        the caller continues, ``start()`` arms conservative callback delivery for
+        the first successfully observed versions instead of accepting them as a
+        callback-less baseline.
+        """
+        session = self._session_factory()
+        try:
+            result = await session.execute(select(CacheInvalidation.namespace, CacheInvalidation.version))
+            rows = result.all()
+        finally:
+            await close_session(session)
+        for namespace, version in rows:
+            self._known_versions[namespace] = version
+        self._poll_initialized = True
+
+    async def prime(self) -> None:
+        """Record the current version baseline without firing callbacks.
+
+        Runs a single poll so that any bump landing *after* this call is
+        delivered as an invalidation callback rather than absorbed as the
+        poller's initial (callback-less) baseline. Call before a one-shot
+        startup reconcile: a leader that persists-and-bumps in the window
+        between that reconcile's read and the poller's first background tick
+        would otherwise be silently missed until the next full backstop. The
+        first poll only records versions (``_poll_initialized`` is still
+        ``False``), so priming never invokes a callback. Unlike ``initialize``
+        this also flushes any pending bumps and seeds baselines for every
+        namespace observed in the ``cache_invalidation`` table.
+
+        Mirrors ``initialize``'s error contract: if the baseline read fails the
+        poller stays uninitialized (``_poll_initialized`` still ``False``) and
+        this method raises so the caller can retry before background polling.
+        Starting without a successful retry makes the first recovered poll
+        reconcile positive versions through their callbacks before acknowledging
+        them.
+        ``_poll_once`` swallows the read error, so a silent success here would
+        let the first *background* poll absorb a peer bump as the initial
+        baseline, voiding the delivery guarantee priming exists to provide.
+        """
+        if not await self._poll_once():
+            raise RuntimeError("cache invalidation prime failed: baseline version read did not complete")
+
+    async def start(self) -> None:
+        if self._task and not self._task.done():
+            return
+        # Callback-less baseline acquisition is safe only before process-local
+        # state can be served. Once background polling starts, a missing baseline
+        # means the observed versions are uncertain: treat each positive first
+        # observation as a change and reconcile it through the normal callback /
+        # acknowledgement path. A successful prime already set this flag after
+        # recording exact versions, so normal startup remains unchanged.
+        self._poll_initialized = True
+        self._stop.clear()
+        self._task = asyncio.create_task(self._run())
+
+    async def stop(self) -> None:
+        if not self._task:
+            return
+        self._stop.set()
+        self._task.cancel()
+        try:
+            await self._task
+        except asyncio.CancelledError:
+            pass
+        self._task = None
+
+    def request_bump(self, namespace: str) -> None:
+        """Enqueue a coalesced bump flushed at the start of the next poll cycle.
+
+        Multiple requests for the same namespace within one cycle collapse into a
+        single version bump. A failed flush keeps the namespace pending so it is
+        retried on subsequent cycles. A request arriving while a flush is already
+        awaiting the bump for this namespace re-queues it for the next cycle.
+        """
+        self._pending_bumps.add(namespace)
+
+    async def bump(self, namespace: str) -> bool:
+        """Publish now, retaining unfinished work for a later poll cycle.
+
+        Consume only markers queued before this attempt. A request arriving
+        during the write re-adds the marker and must survive completion. An
+        interrupted write may already have committed; a redundant retry is
+        safe, while losing its invalidation is not.
+        """
+        self._pending_bumps.discard(namespace)
+        succeeded = False
+        try:
+            for attempt in range(_BUMP_RETRY_ATTEMPTS):
+                try:
+                    await self._bump_once(namespace)
+                    succeeded = True
+                    return True
+                except OperationalError:
+                    if attempt == _BUMP_RETRY_ATTEMPTS - 1:
+                        self._record_bump_failure(namespace)
+                        return False
+                    await asyncio.sleep(_BUMP_RETRY_BASE_SECONDS * (2**attempt))
+                except Exception:
+                    self._record_bump_failure(namespace)
+                    return False
+            return False
+        finally:
+            # Synchronous cleanup also covers cancellation during retry backoff
+            # and session cleanup without swallowing the cancellation.
+            if not succeeded:
+                self.request_bump(namespace)
+
+    async def bump_local(self, namespace: str) -> bool:
+        """Bump a namespace this replica has ALREADY invalidated locally.
+
+        For callers that both mutate their own in-memory state AND want peers
+        to react (e.g. the reset-credit redeem path evicts the affected
+        account's snapshot locally, then bumps so peers clear their stores).
+        After the shared bump succeeds this records the resulting version as
+        already-observed on THIS poller, so the originating replica does NOT
+        re-run its (possibly whole-store) callback for a bump it already
+        accounted for locally. Peer replicas still observe the bump and fire.
+
+        A peer bump that lands between our commit and the acknowledging read is
+        acknowledged here without firing on this replica; that degrades to the
+        per-replica refresh fallback the reset-credits design already documents
+        for a lost bump, and never affects peers. ``_known_versions`` is only
+        advanced (``max``), never rewound, so a concurrent poll cannot be forced
+        to re-fire by this method.
+        """
+        if not await self.bump(namespace):
+            return False
+        session = self._session_factory()
+        try:
+            version = await session.scalar(
+                select(CacheInvalidation.version).where(CacheInvalidation.namespace == namespace)
+            )
+        except Exception:
+            # Failing to acknowledge only risks one redundant self-invalidation
+            # on the next poll; the bump itself already succeeded for peers.
+            logger.debug("cache_invalidation bump_local acknowledge read failed", exc_info=True)
+            return True
+        finally:
+            await close_session(session)
+        if version is not None:
+            self._known_versions[namespace] = max(self._known_versions.get(namespace, 0), version)
+        return True
+
+    async def _bump_once(self, namespace: str) -> None:
+        session = self._session_factory()
+        try:
+            dialect = session.get_bind().dialect.name
+            if dialect == "postgresql":
+                stmt = (
+                    pg_insert(CacheInvalidation)
+                    .values(namespace=namespace, version=1)
+                    .on_conflict_do_update(
+                        index_elements=[CacheInvalidation.namespace],
+                        set_={"version": CacheInvalidation.version + 1},
+                    )
+                )
+                await session.execute(stmt)
+            elif dialect == "sqlite":
+                stmt = (
+                    sqlite_insert(CacheInvalidation)
+                    .values(namespace=namespace, version=1)
+                    .on_conflict_do_update(
+                        index_elements=[CacheInvalidation.namespace],
+                        set_={"version": CacheInvalidation.version + 1},
+                    )
+                )
+                await session.execute(stmt)
+            else:
+                existing = await session.scalar(
+                    select(CacheInvalidation).where(CacheInvalidation.namespace == namespace)
+                )
+                if existing is None:
+                    session.add(CacheInvalidation(namespace=namespace, version=1))
+                else:
+                    await session.execute(
+                        update(CacheInvalidation)
+                        .where(CacheInvalidation.namespace == namespace)
+                        .values(version=CacheInvalidation.version + 1)
+                    )
+            await session.commit()
+        finally:
+            await close_session(session)
+
+    def _record_bump_failure(self, namespace: str) -> None:
+        logger.error(
+            "cache_invalidation bump failed for namespace %s",
+            _NAMESPACE_LOG_LABELS.get(namespace, "unknown"),
+            exc_info=True,
+        )
+        if PROMETHEUS_AVAILABLE and cache_invalidation_bump_failures_total is not None:
+            cache_invalidation_bump_failures_total.labels(namespace=namespace).inc()
+
+    async def _run(self) -> None:
+        while not self._stop.is_set():
+            try:
+                await self._poll_once()
+            except Exception:
+                logger.debug("cache_invalidation poll failed", exc_info=True)
+            try:
+                await asyncio.wait_for(self._stop.wait(), timeout=self._poll_interval)
+            except asyncio.TimeoutError:
+                continue
+
+    async def _flush_pending_bumps(self) -> None:
+        for namespace in sorted(self._pending_bumps):
+            try:
+                await self.bump(namespace)
+            except asyncio.CancelledError:
+                # Defensive if bump() is replaced: cancellation still retains
+                # the current namespace and propagates to the flush owner.
+                self.request_bump(namespace)
+                raise
+            except Exception:
+                # An abnormal raise must not starve later namespaces.
+                self.request_bump(namespace)
+                logger.warning(
+                    "cache_invalidation flush bump raised for namespace %s; kept pending",
+                    _NAMESPACE_LOG_LABELS.get(namespace, "unknown"),
+                    exc_info=True,
+                )
+
+    async def _poll_once(self) -> bool:
+        """Flush pending bumps and reconcile observed versions once.
+
+        Returns ``True`` when the baseline/version read completed (callbacks may
+        still have failed individually and left their namespace unacknowledged)
+        and ``False`` when the read itself failed, so ``prime`` can surface a
+        failed startup seed instead of silently leaving the poller uninitialized.
+        """
+        await self._flush_pending_bumps()
+        session: AsyncSession | None = None
+        try:
+            session = self._session_factory()
+            result = await session.execute(select(CacheInvalidation.namespace, CacheInvalidation.version))
+            rows = result.all()
+        except Exception:
+            self._record_poll_failure()
+            return False
+        finally:
+            if session is not None:
+                await close_session(session)
+        self._consecutive_poll_failures = 0
+
+        for namespace, version in rows:
+            # Re-read the acknowledged version for THIS namespace immediately
+            # before comparing: a concurrent bump_local() may have advanced it
+            # (recording a local self-suppression) after this cycle's snapshot
+            # read but before we reach this row. Using the freshly-read value
+            # keeps the local ack authoritative.
+            prev = self._known_versions.get(namespace)
+            # Only a STRICTLY newer version is a change. A stale snapshot whose
+            # version is <= the acknowledged version (e.g. read before a
+            # concurrent local bump advanced _known_versions) must neither fire
+            # the callback nor rewind the acknowledged version.
+            changed = (prev is not None and version > prev) or (prev is None and self._poll_initialized and version > 0)
+            if changed and not await self._run_callbacks(namespace):
+                # Do not acknowledge the observed version: a failed callback
+                # (e.g. a transient DB error during a snapshot refresh) must be
+                # retried on the next cycle, or a replica would permanently
+                # miss the invalidation.
+                continue
+            # Never lower an acknowledged version: a concurrent bump_local()
+            # advance must survive an older in-flight poll observation.
+            self._known_versions[namespace] = max(self._known_versions.get(namespace, 0), version)
+        self._poll_initialized = True
+        return True
+
+    async def _run_callbacks(self, namespace: str) -> bool:
+        succeeded = True
+        for cb in self._callbacks.get(namespace, []):
+            try:
+                result = cb()
+                if isawaitable(result):
+                    await result
+            except Exception:
+                succeeded = False
+                logger.warning(
+                    "cache_invalidation callback failed for namespace %s; retrying next poll cycle",
+                    _NAMESPACE_LOG_LABELS.get(namespace, "unknown"),
+                    exc_info=True,
+                )
+        return succeeded
+
+    def _record_poll_failure(self) -> None:
+        self._consecutive_poll_failures += 1
+        if PROMETHEUS_AVAILABLE and cache_invalidation_poll_failures_total is not None:
+            cache_invalidation_poll_failures_total.inc()
+        if self._consecutive_poll_failures >= _POLL_FAILURES_ERROR_THRESHOLD:
+            logger.error(
+                "cache_invalidation poll failed %d consecutive times",
+                self._consecutive_poll_failures,
+                exc_info=True,
+            )
+        elif self._consecutive_poll_failures >= _POLL_FAILURES_WARNING_THRESHOLD:
+            logger.warning(
+                "cache_invalidation poll failed %d consecutive times",
+                self._consecutive_poll_failures,
+                exc_info=True,
+            )
+        else:
+            logger.debug("cache_invalidation poll failed", exc_info=True)
+
+
+_poller: CacheInvalidationPoller | None = None
+
+
+def get_cache_invalidation_poller() -> CacheInvalidationPoller | None:
+    return _poller
+
+
+def set_cache_invalidation_poller(poller: CacheInvalidationPoller | None) -> None:
+    global _poller
+    _poller = poller
+
+
+async def bump_cache_invalidation(namespace: str) -> None:
+    """Publish now or queue a retry; a no-op outside the poller's lifetime."""
+    poller = _poller
+    if poller is None:
+        return
+    await poller.bump(namespace)
+
+
+async def bump_cache_invalidation_local(namespace: str) -> None:
+    """Best-effort bump for a namespace already invalidated locally by the caller.
+
+    Notifies peers while suppressing this replica's own re-invalidation for the
+    bump it just issued. A no-op outside the lifespan poller's lifetime.
+    """
+    poller = _poller
+    if poller is None:
+        return
+    await poller.bump_local(namespace)