codexLB-up 1.25.0b10__py3-none-any.whl

app/core/openai/model_refresh_scheduler.py

@@ -0,0 +1,423 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from dataclasses import dataclass, field
+
+from app.core.auth.refresh import RefreshError
+from app.core.cache.invalidation import NAMESPACE_MODEL_REGISTRY, get_cache_invalidation_poller
+from app.core.clients.codex_version import get_codex_version_cache
+from app.core.clients.http import refresh_http_client
+from app.core.clients.model_fetcher import ModelFetchError, fetch_models_for_plan
+from app.core.config.settings import get_settings
+from app.core.crypto import TokenEncryptor
+from app.core.openai.model_registry import (
+    UpstreamModel,
+    _merge_service_tier_metadata,
+    get_model_registry,
+)
+from app.core.openai.model_registry_store import (
+    encode_registry_export,
+    persist_registry_snapshot,
+    reconcile_model_registry_from_store,
+)
+from app.core.scheduling.leader_election_handle import get_leader_election as _get_leader_election
+from app.core.upstream_proxy import ResolvedUpstreamRoute, resolve_upstream_route
+from app.db.models import Account, AccountStatus
+from app.db.session import detach_session_objects, get_background_session
+from app.modules.accounts.auth_manager import AuthManager
+from app.modules.accounts.background_repository import BackgroundAccountsRepository
+from app.modules.accounts.repository import AccountsRepository
+from app.modules.proxy.account_cache import get_account_selection_cache
+
+logger = logging.getLogger(__name__)
+
+# Registry refresh cadence (fixed; issue #1340 / PRINCIPLES.md P2). The
+# scheduler keeps ``interval_seconds`` as a constructor field so tests can
+# exercise the loop with a short interval.
+_REFRESH_INTERVAL_SECONDS = 300
+
+
+@dataclass(slots=True)
+class _TransportRecoveryState:
+    attempted: bool = False
+
+
+@dataclass(slots=True)
+class _FetchResult:
+    models: list[UpstreamModel]
+    account_models: dict[str, tuple[str, list[UpstreamModel]]]
+
+
+async def _warm_codex_version_cache() -> None:
+    """Refresh the in-process Codex client version on every replica.
+
+    The version is presented as the outbound fingerprint of non-native
+    requests (``codex_cli_rs/<version>``); upstream gates newer models on it.
+    It used to be fetched only inside the leader's model refresh, so a
+    non-leader replica -- for instance the live color of a blue/green pair
+    whose standby still holds the scheduler lease -- served the configured
+    fallback version indefinitely and had its non-native ``gpt-6-astra``
+    requests rejected with "requires a newer version of Codex". The fetch is
+    a public GitHub/npm lookup (no account token) cached for an hour, so
+    every replica may perform it, and it must never fail the tick.
+    """
+    try:
+        await get_codex_version_cache().get_version()
+    except Exception:  # pragma: no cover - the cache itself already logs and falls back
+        logger.warning("Codex client version warm-up failed; keeping the cached or default version", exc_info=True)
+
+
+@dataclass(slots=True)
+class ModelRefreshScheduler:
+    interval_seconds: int
+    enabled: bool
+    _task: asyncio.Task[None] | None = None
+    _stop: asyncio.Event = field(default_factory=asyncio.Event)
+
+    async def start(self) -> None:
+        if not self.enabled:
+            return
+        if self._task and not self._task.done():
+            return
+        self._stop.clear()
+        self._task = asyncio.create_task(self._run_loop())
+
+    async def stop(self) -> None:
+        if not self._task:
+            return
+        self._stop.set()
+        self._task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self._task
+        self._task = None
+
+    async def _run_loop(self) -> None:
+        while not self._stop.is_set():
+            await _warm_codex_version_cache()
+            await self._refresh_once()
+            try:
+                await asyncio.wait_for(self._stop.wait(), timeout=self.interval_seconds)
+            except asyncio.TimeoutError:
+                continue
+
+    async def _refresh_once(self) -> None:
+        ran_as_leader = await _get_leader_election().run_if_leader(self._refresh_as_leader)
+        if not ran_as_leader:
+            # Never fetch upstream on a non-leader; reconcile from the persisted
+            # snapshot instead. This is the TTL backstop for a lost invalidation
+            # bump — the 0.5s cache-invalidation poller is the fast path. Also
+            # covers losing the lease mid-refresh, where run_if_leader returns
+            # None after cancelling the body.
+            await reconcile_model_registry_from_store()
+
+    async def _refresh_as_leader(self) -> bool:
+        try:
+            async with get_background_session() as session:
+                accounts_repo = AccountsRepository(session)
+                accounts = await accounts_repo.list_accounts()
+                detach_session_objects(session)
+            grouped = _group_by_plan(accounts)
+            if not grouped:
+                await get_model_registry().clear()
+                get_account_selection_cache().invalidate()
+                logger.info("Model registry cleared because no active accounts remain")
+                await _persist_registry_state_and_bump()
+                return True
+
+            encryptor = TokenEncryptor()
+            per_plan_results: dict[str, list[UpstreamModel]] = {}
+            per_account_results: dict[str, tuple[str, list[UpstreamModel]]] = {}
+            active_account_plans: dict[str, str] = {}
+
+            for plan_type, candidates in grouped.items():
+                for account in candidates:
+                    active_account_plans[account.id] = plan_type
+                result = await _fetch_with_failover(
+                    candidates,
+                    encryptor,
+                )
+                if result is not None:
+                    per_plan_results[plan_type] = result.models
+                    per_account_results.update(result.account_models)
+
+            if per_plan_results:
+                registry = get_model_registry()
+                await registry.update(
+                    per_plan_results,
+                    per_account_results=per_account_results,
+                    active_account_plans=active_account_plans,
+                )
+                snapshot = registry.get_snapshot()
+                total_models = len(snapshot.models) if snapshot else 0
+                logger.info(
+                    "Model registry refreshed plans=%d total_models=%d",
+                    len(per_plan_results),
+                    total_models,
+                )
+                get_account_selection_cache().invalidate()
+                await _persist_registry_state_and_bump()
+            else:
+                logger.warning("Model registry refresh failed for all plans")
+                # Every upstream fetch failed, so the leader made no change and
+                # never advances the persisted ``refreshed_at``. Followers drop
+                # to the bootstrap floor once the store row ages past
+                # ``model_registry_snapshot_max_age_seconds``; reconcile here so
+                # the leader applies the same expiry instead of serving its now
+                # stale in-memory catalog indefinitely under a prolonged
+                # upstream outage. On a still-fresh row this is a no-op because
+                # the leader's applied content hash already matches the store.
+                await reconcile_model_registry_from_store()
+        except Exception:
+            logger.exception("Model registry refresh loop failed")
+        # Ran as leader (even on internal failure): signal completion so the
+        # caller does not additionally reconcile from the persisted snapshot.
+        return True
+
+
+async def _persist_registry_state_and_bump() -> None:
+    """Persist the leader's registry state, then bump the bus (write-then-bump).
+
+    A persist failure degrades to leader-local refresh behavior: the in-memory
+    registry already holds the refreshed catalog and persistence is retried on
+    the next cycle. The applied-hash marker is reset on failure because the
+    in-memory state now diverges from the persisted row; leaving the old hash
+    in place would make a later reconcile (e.g. after losing leadership) treat
+    the store's row as already applied and never converge back to it.
+    """
+    registry = get_model_registry()
+    try:
+        export = await registry.export_state()
+        encoded = encode_registry_export(export)
+        async with get_background_session() as session:
+            changed = await persist_registry_snapshot(
+                session,
+                encoded=encoded,
+                leader_id=get_settings().http_responses_session_bridge_instance_id,
+            )
+        registry.note_applied_content_hash(encoded.content_hash)
+    except Exception:
+        registry.note_applied_content_hash(None)
+        logger.warning(
+            "Model registry snapshot persist failed; serving leader-local refresh until next cycle",
+            exc_info=True,
+        )
+        return
+    if changed:
+        poller = get_cache_invalidation_poller()
+        if poller is not None:
+            await poller.bump(NAMESPACE_MODEL_REGISTRY)
+
+
+def _group_by_plan(accounts: list[Account]) -> dict[str, list[Account]]:
+    grouped: dict[str, list[Account]] = {}
+    for account in accounts:
+        if account.status != AccountStatus.ACTIVE:
+            continue
+        plan_type = account.plan_type
+        if not plan_type:
+            continue
+        grouped.setdefault(plan_type, []).append(account)
+    return grouped
+
+
+def _error_summary(exc: BaseException) -> str:
+    if isinstance(exc, ModelFetchError):
+        summary = f"status={exc.status_code} transport={exc.transport_error}"
+        if exc.message:
+            summary = f"{summary} message={_compact_error_message(exc.message)}"
+        return summary
+    if isinstance(exc, RefreshError):
+        summary = f"code={exc.code} permanent={exc.is_permanent} transport={exc.transport_error}"
+        if exc.message:
+            summary = f"{summary} message={_compact_error_message(exc.message)}"
+        return summary
+
+    message = _compact_error_message(str(exc))
+    if message:
+        return f"{exc.__class__.__name__}: {message}"
+    return exc.__class__.__name__
+
+
+def _compact_error_message(message: str) -> str:
+    return " ".join(message.split())
+
+
+async def _fetch_with_failover(
+    candidates: list[Account],
+    encryptor: TokenEncryptor,
+    accounts_repo: AccountsRepository | None = None,
+) -> _FetchResult | None:
+    transport_recovery = _TransportRecoveryState()
+    successful_results: list[list[UpstreamModel]] = []
+    account_models: dict[str, tuple[str, list[UpstreamModel]]] = {}
+    auth_accounts_repo = accounts_repo or BackgroundAccountsRepository()
+    auth_manager = AuthManager(auth_accounts_repo)
+
+    for account in candidates:
+        try:
+            account = await _ensure_fresh_with_transport_recovery(
+                auth_manager,
+                account,
+                transport_recovery=transport_recovery,
+            )
+            models = await _fetch_models_with_transport_recovery(
+                account,
+                encryptor,
+                transport_recovery=transport_recovery,
+            )
+            successful_results.append(models)
+            account_models[account.id] = (account.plan_type, models)
+        except ModelFetchError as exc:
+            if exc.status_code == 401:
+                try:
+                    account = await _ensure_fresh_with_transport_recovery(
+                        auth_manager,
+                        account,
+                        force=True,
+                        transport_recovery=transport_recovery,
+                    )
+                    models = await _fetch_models_with_transport_recovery(
+                        account,
+                        encryptor,
+                        transport_recovery=transport_recovery,
+                    )
+                    successful_results.append(models)
+                    account_models[account.id] = (account.plan_type, models)
+                    continue
+                except (ModelFetchError, RefreshError) as retry_exc:
+                    logger.warning(
+                        "Model fetch auth retry failed account=%s plan=%s initial_error=%s retry_error=%s",
+                        account.id,
+                        account.plan_type,
+                        _error_summary(exc),
+                        _error_summary(retry_exc),
+                    )
+                    continue
+            logger.warning(
+                "Model fetch failed account=%s plan=%s error=%s",
+                account.id,
+                account.plan_type,
+                _error_summary(exc),
+            )
+            continue
+        except RefreshError as exc:
+            logger.warning(
+                "Token refresh failed for model fetch account=%s plan=%s error=%s",
+                account.id,
+                account.plan_type,
+                _error_summary(exc),
+            )
+            continue
+        except Exception as exc:
+            logger.warning(
+                "Unexpected error during model fetch account=%s plan=%s error=%s",
+                account.id,
+                account.plan_type,
+                _error_summary(exc),
+                exc_info=True,
+            )
+            continue
+    merged_models = _merge_same_plan_model_results(successful_results)
+    if not successful_results:
+        return None
+    return _FetchResult(models=merged_models, account_models=account_models)
+
+
+def _merge_same_plan_model_results(successful_results: list[list[UpstreamModel]]) -> list[UpstreamModel]:
+    if not successful_results:
+        return []
+
+    merged_by_slug: dict[str, UpstreamModel] = {}
+    for models in successful_results:
+        for model in models:
+            existing = merged_by_slug.get(model.slug)
+            merged_by_slug[model.slug] = model if existing is None else _merge_service_tier_metadata(existing, model)
+    return list(merged_by_slug.values())
+
+
+async def _ensure_fresh_with_transport_recovery(
+    auth_manager: AuthManager,
+    account: Account,
+    *,
+    transport_recovery: _TransportRecoveryState,
+    force: bool = False,
+) -> Account:
+    try:
+        return await auth_manager.ensure_fresh(account, force=force)
+    except RefreshError as exc:
+        if not exc.transport_error or transport_recovery.attempted:
+            raise
+
+        await _refresh_http_client_after_transport_error(account, exc)
+        transport_recovery.attempted = True
+        return await auth_manager.ensure_fresh(account, force=force)
+
+
+async def _fetch_models_with_transport_recovery(
+    account: Account,
+    encryptor: TokenEncryptor,
+    *,
+    transport_recovery: _TransportRecoveryState,
+) -> list[UpstreamModel]:
+    access_token = encryptor.decrypt(account.access_token_encrypted)
+    account_id = account.chatgpt_account_id
+    route = await _resolve_upstream_route_for_account(account, operation="model_discovery")
+
+    try:
+        return await fetch_models_for_plan(
+            access_token,
+            account_id,
+            route=route,
+            allow_direct_egress=route is None,
+        )
+    except ModelFetchError as exc:
+        if not exc.transport_error or transport_recovery.attempted:
+            raise
+
+        await _refresh_http_client_after_transport_error(account, exc)
+        transport_recovery.attempted = True
+        access_token = encryptor.decrypt(account.access_token_encrypted)
+        account_id = account.chatgpt_account_id
+        route = await _resolve_upstream_route_for_account(account, operation="model_discovery")
+        return await fetch_models_for_plan(
+            access_token,
+            account_id,
+            route=route,
+            allow_direct_egress=route is None,
+        )
+
+
+async def _resolve_upstream_route_for_account(account: Account, *, operation: str) -> ResolvedUpstreamRoute | None:
+    async with get_background_session() as session:
+        return await resolve_upstream_route(
+            session,
+            account_id=account.id,
+            operation=operation,
+            scope="account",
+        )
+
+
+async def _refresh_http_client_after_transport_error(account: Account, transport_exc: BaseException) -> None:
+    try:
+        await refresh_http_client()
+    except Exception as refresh_exc:
+        logger.warning(
+            "Model fetch transport recovery failed account=%s plan=%s transport_error=%s refresh_error=%s",
+            account.id,
+            account.plan_type,
+            _error_summary(transport_exc),
+            _error_summary(refresh_exc),
+        )
+        raise
+    logger.info(
+        "Refreshed shared HTTP client after model fetch transport error; retrying account=%s plan=%s error=%s",
+        account.id,
+        account.plan_type,
+        _error_summary(transport_exc),
+    )
+
+
+def build_model_refresh_scheduler() -> ModelRefreshScheduler:
+    return ModelRefreshScheduler(interval_seconds=_REFRESH_INTERVAL_SECONDS, enabled=True)