diff --git a/litellm/proxy/common_utils/codex_catalog_source.py b/litellm/proxy/common_utils/codex_catalog_source.py index 75526454d66..935ad66246c 100644 --- a/litellm/proxy/common_utils/codex_catalog_source.py +++ b/litellm/proxy/common_utils/codex_catalog_source.py @@ -15,6 +15,7 @@ fetch at a mirror of the openai/codex raw tree. from __future__ import annotations import asyncio +import hashlib import os import re import time @@ -30,6 +31,7 @@ from pydantic import ValidationError from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.offline_mode import is_offline_mode +from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError from litellm.proxy.common_utils.codex_model_catalog import ( CODEX_BUNDLED_VERSION, CodexStock, @@ -80,6 +82,13 @@ def _default_client() -> _AsyncGetClient: return get_async_httpx_client(llm_provider=httpxSpecialProvider.CodexModelCatalog) +def _cache_slot(tag: str, openai_api_key: str | None) -> str: + """Catalogs fetched with different OpenAI keys are kept apart, since OpenAI may answer per account.""" + if openai_api_key is None: + return tag + return f"{tag}#{hashlib.sha256(openai_api_key.encode()).hexdigest()[:16]}" + + def vendored_codex_stock() -> CodexStock: """The pair vendored with LiteLLM, served when the requesting release's own cannot be.""" return CodexStock( @@ -139,11 +148,12 @@ class CodexCatalogSource: tag: Final = codex_release_tag(client_version) if tag is None or client_version is None or uses_vendored_codex_catalog(): return vendored_codex_stock() - cached: Final = self._cache.get(tag) + slot: Final = _cache_slot(tag, openai_api_key) + cached: Final = self._cache.get(slot) if cached is not None and (cached.retry_at is None or cached.retry_at > self._clock()): - self._cache.move_to_end(tag) + self._cache.move_to_end(slot) return cached.stock - pending: Final = self._pending.get(tag) + pending: Final = self._pending.get(slot) if pending is not None: return await asyncio.shield(pending) if len(self._pending) >= self._cached_versions: @@ -154,12 +164,12 @@ class CodexCatalogSource: vendored_codex_stock().source, ) return vendored_codex_stock() - task: Final = asyncio.ensure_future(self._resolve(tag, client_version, openai_api_key)) - self._pending[tag] = task - task.add_done_callback(lambda _: self._pending.pop(tag, None)) + task: Final = asyncio.ensure_future(self._resolve(slot, tag, client_version, openai_api_key)) + self._pending[slot] = task + task.add_done_callback(lambda _: self._pending.pop(slot, None)) return await asyncio.shield(task) - async def _resolve(self, tag: str, version: str, openai_api_key: str | None) -> CodexStock: + async def _resolve(self, slot: str, tag: str, version: str, openai_api_key: str | None) -> CodexStock: fetched: Final = await self._fetch(tag, version, openai_api_key) if isinstance(fetched, _FetchFailed): verbose_proxy_logger.warning( @@ -178,8 +188,8 @@ class CodexCatalogSource: if isinstance(fetched, _FetchFailed) else _Cached(stock=fetched, retry_at=None) ) - self._cache[tag] = entry - self._cache.move_to_end(tag) + self._cache[slot] = entry + self._cache.move_to_end(slot) while len(self._cache) > self._cached_versions: self._cache.popitem(last=False) return entry.stock @@ -243,7 +253,7 @@ class CodexCatalogSource: response: Final = await client.get( url, headers=headers, timeout=self._timeout, max_response_bytes=CODEX_CATALOG_MAX_BYTES ) - except httpx.HTTPError as error: + except (httpx.HTTPError, HTTPResponseLimitError) as error: return _FetchFailed(reason=f"{type(error).__name__} fetching {url}") if response.status_code != 200: return _FetchFailed(reason=f"HTTP {response.status_code} from {url}") diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 90a47e12757..496110d36c0 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -31,6 +31,7 @@ from itertools import accumulate, pairwise from pathlib import Path from types import MappingProxyType, NoneType from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias +from urllib.parse import urlparse from packaging.version import Version from pydantic import ConfigDict, Field, StringConstraints, TypeAdapter, ValidationError @@ -569,23 +570,46 @@ def _catalog_row( ) +OPENAI_API_HOST: Final = "api.openai.com" + + +def _is_openai_origin(api_base: object) -> bool: + """Whether a deployment's api_base is OpenAI itself, so its key may go to OpenAI's catalog endpoint.""" + if not api_base: + return True + return isinstance(api_base, str) and urlparse(api_base).hostname == OPENAI_API_HOST + + def _deployment_openai_key(deployment: DeploymentTypedDict) -> str | None: params: Final = MappingProxyType[str, object](deployment["litellm_params"]) model: Final = params.get("model") provider: Final = params.get("custom_llm_provider") or (model.partition("/")[0] if isinstance(model, str) else "") api_key: Final = params.get("api_key") - if provider != "openai" or not isinstance(api_key, str) or not api_key: + if ( + provider != "openai" + or not _is_openai_origin(params.get("api_base")) + or not isinstance(api_key, str) + or not api_key + ): return None return (get_secret_str(api_key) if api_key.startswith("os.environ/") else api_key) or None +def _environment_openai_key() -> str | None: + if not _is_openai_origin(os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE")): + return None + return os.getenv("OPENAI_API_KEY") or None + + def codex_catalog_openai_key( llm_router: Router | None, entries: Sequence[tuple[str, str]], team_id: str | None = None, ) -> str | None: - """The OpenAI API key the catalog fetch authenticates with: the first one among the - listed entries' OpenAI deployments, else ``OPENAI_API_KEY`` from the environment.""" + """The OpenAI API key the catalog fetch authenticates with: the first one among the listed entries' + deployments that is `openai/` with no api_base or OpenAI's own, else ``OPENAI_API_KEY`` from the + environment unless ``OPENAI_BASE_URL`` points elsewhere. A key for an OpenAI-compatible server never + leaves for OpenAI.""" registry: Final = llm_router.registry() if llm_router is not None else None deployments: Final = ( tuple( @@ -598,7 +622,7 @@ def codex_catalog_openai_key( if registry is not None else () ) - return next(filter(None, map(_deployment_openai_key, deployments)), None) or os.getenv("OPENAI_API_KEY") or None + return next(filter(None, map(_deployment_openai_key, deployments)), None) or _environment_openai_key() def codex_catalog_rows( diff --git a/tests/unit/proxy/common_utils/test_codex_catalog_source.py b/tests/unit/proxy/common_utils/test_codex_catalog_source.py index 966d0b00a10..8c93f9aa947 100644 --- a/tests/unit/proxy/common_utils/test_codex_catalog_source.py +++ b/tests/unit/proxy/common_utils/test_codex_catalog_source.py @@ -7,6 +7,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError from litellm.proxy.common_utils.codex_catalog_source import ( CODEX_CATALOG_BASE_URL, CODEX_CATALOG_MAX_BYTES, @@ -300,3 +301,39 @@ def test_a_release_asked_for_while_the_in_flight_fetches_are_at_the_cap_gets_the source_label, calls = asyncio.run(scenario()) assert source_label == vendored_codex_stock().source and not any("0.152.0" in url for url in calls), calls + + +def test_a_download_the_client_cuts_at_the_byte_limit_falls_through_to_the_next_source() -> None: + class _CuttingTree(_RawTree): + async def get( + self, + url: str, + *, + headers: dict[str, str] | None = None, + timeout: float | None = None, + max_response_bytes: int | None = None, + ) -> httpx.Response: + if url == _openai_url(RELEASE): + self.calls.append(url) + raise HTTPResponseLimitError(f"{url} exceeded {max_response_bytes} bytes") + return await super().get(url, headers=headers, timeout=timeout, max_response_bytes=max_response_bytes) + + tree = _CuttingTree(_tree(RELEASE)) + stock = asyncio.run(_source(tree).stock(RELEASE, openai_api_key=OPENAI_KEY)) + assert stock.source == f"github rust-v{RELEASE}" and tree.calls[0] == _openai_url(RELEASE), tree.calls + + +def test_catalogs_fetched_with_different_openai_keys_are_cached_apart() -> None: + tree = _RawTree({**_tree(RELEASE), _openai_url(RELEASE): (200, OPENAI_CATALOG)}) + + async def scenario() -> None: + source = _source(tree) + await source.stock(RELEASE, openai_api_key=OPENAI_KEY) + await source.stock(RELEASE, openai_api_key=OPENAI_KEY) + await source.stock(RELEASE, openai_api_key="sk-unit-other-key") + + asyncio.run(scenario()) + assert tree.headers == [ + {"Authorization": f"Bearer {OPENAI_KEY}"}, + {"Authorization": "Bearer sk-unit-other-key"}, + ], tree.headers diff --git a/tests/unit/proxy/common_utils/test_codex_model_catalog.py b/tests/unit/proxy/common_utils/test_codex_model_catalog.py index 45861aea67c..16f507cc6fb 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -978,3 +978,21 @@ def test_the_catalog_key_is_absent_without_an_openai_deployment_or_environment_k router: Final = _router_with({"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-ant-not-openai"}) assert codex_catalog_openai_key(router, [("catalog-model-0", "catalog-model-0")]) is None assert codex_catalog_openai_key(None, [("catalog-model-0", "catalog-model-0")]) is None + + +def test_a_key_for_an_openai_compatible_server_or_a_redirected_environment_never_goes_to_openai( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + monkeypatch.setenv("OPENAI_BASE_URL", "https://llm-gateway.internal.example/v1") + router: Final = _router_with( + { + "model": "openai/local-llama", + "api_key": "sk-internal-server", + "api_base": "https://inference.internal.example/v1", + }, + {"model": "openai/gpt-5.5", "api_key": "sk-real-openai", "api_base": "https://api.openai.com/v1"}, + ) + entries: Final = [(f"catalog-model-{index}", f"catalog-model-{index}") for index in range(2)] + assert codex_catalog_openai_key(router, entries[:1]) is None + assert codex_catalog_openai_key(router, entries) == "sk-real-openai"