From 8eebb745bdaf88d8ea2292a356c1d362b3763d6e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:34:18 -0800 Subject: [PATCH] fix(proxy): keep FedRAMP and non-api.openai.com keys home, survive catalog redirect loops, and keep fetched Codex releases cached --- .../common_utils/codex_catalog_source.py | 25 +++++----- .../proxy/common_utils/codex_model_catalog.py | 21 ++++++-- .../common_utils/test_codex_catalog_source.py | 43 +++++++++++++++- .../common_utils/test_codex_model_catalog.py | 49 ++++++++++++++++++- 4 files changed, 119 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/common_utils/codex_catalog_source.py b/litellm/proxy/common_utils/codex_catalog_source.py index 0a10ad24649..5bf95c540b0 100644 --- a/litellm/proxy/common_utils/codex_catalog_source.py +++ b/litellm/proxy/common_utils/codex_catalog_source.py @@ -40,7 +40,6 @@ import httpx from pydantic import ConfigDict, 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, @@ -49,6 +48,7 @@ from litellm.proxy.common_utils.codex_model_catalog import ( CodexStockModel, bundled_codex_models, codex_base_instructions, + uses_vendored_codex_catalog, ) from litellm.types.llms.base_model import LiteLLMBaseModel @@ -127,17 +127,12 @@ def _default_client() -> _AsyncGetClient: return get_async_httpx_client(llm_provider=httpxSpecialProvider.CodexModelCatalog) +_VENDORED_SOURCE: Final = f"vendored {CODEX_BUNDLED_VERSION}" + + def vendored_codex_stock() -> CodexStock: """The pair vendored with LiteLLM, served when the requesting release's own cannot be.""" - return CodexStock( - models=bundled_codex_models(), - instructions=codex_base_instructions(), - source=f"vendored {CODEX_BUNDLED_VERSION}", - ) - - -def uses_vendored_codex_catalog() -> bool: - return is_offline_mode() or os.getenv("LITELLM_LOCAL_CODEX_CATALOG", "").strip().lower() == "true" + return CodexStock(models=bundled_codex_models(), instructions=codex_base_instructions(), source=_VENDORED_SOURCE) def codex_release_tag(client_version: str | None) -> str | None: @@ -271,9 +266,15 @@ class CodexCatalogSource: self._cache[slot] = entry self._cache.move_to_end(slot) while len(self._cache) > self._cached_versions: - self._cache.popitem(last=False) + del self._cache[self._eviction_victim()] return entry.stock + def _eviction_victim(self) -> _CacheSlot: + """The least recently asked release still on the vendored pair, else the least recently asked one, so + releases no catalog answered push each other out before a fetched release goes.""" + vendored: Final = (slot for slot, cached in self._cache.items() if cached.stock.source == _VENDORED_SOURCE) + return next(vendored, next(iter(self._cache))) + async def _fetch( self, tag: str, version: str, openai_api_key: str | None, *, github_fallback: bool ) -> _Fetched | _FetchFailed: @@ -345,7 +346,7 @@ class CodexCatalogSource: response: Final = await client.get( url, headers=headers, timeout=self._timeout, max_response_bytes=CODEX_CATALOG_MAX_BYTES ) - except (httpx.HTTPError, httpx.InvalidURL, HTTPResponseLimitError) as error: + except (httpx.HTTPError, httpx.InvalidURL, HTTPResponseLimitError, ValueError) 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 13de829fbd5..03b84e690b5 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -37,6 +37,7 @@ from packaging.version import Version from pydantic import ConfigDict, Field, StringConstraints, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.litellm_core_utils.reasoning_effort_strength import nearest_declared_reasoning_effort from litellm.types.llms.base_model import LiteLLMBaseModel @@ -583,7 +584,7 @@ def _catalog_row( OPENAI_API_HOST: Final = "api.openai.com" -OPENAI_HOST_SUFFIX: Final = "openai.com" +FEDRAMP_KEY_PREFIXES: Final = ("sk-proj-fedramp-", "sk-svcacct-fedramp-") _DEPLOYMENT_BASE_PARAMS: Final = ("api_base", "base_url") @@ -591,7 +592,7 @@ _DEPLOYMENT_BASE_PARAMS: Final = ("api_base", "base_url") def _is_openai_host(url: object) -> bool: host: Final = urlparse(url).hostname if isinstance(url, str) else None - return host is not None and (host == OPENAI_API_HOST or host.endswith(f".{OPENAI_HOST_SUFFIX}")) + return host == OPENAI_API_HOST def _is_openai_origin(bases: Sequence[object]) -> bool: @@ -646,9 +647,13 @@ def _deployment_openai_key(deployment: DeploymentTypedDict) -> str | None: *(params.get(key) for key in _DEPLOYMENT_BASE_PARAMS), *(credential.get(key) for key in _DEPLOYMENT_BASE_PARAMS), ) - api_key: Final = effective.get("api_key") if _deployment_provider(effective) != "openai" or not _is_openai_origin(bases): return None + key: Final = _resolved_key(effective.get("api_key")) + return None if key is None or key.startswith(FEDRAMP_KEY_PREFIXES) else key + + +def _resolved_key(api_key: object) -> str | None: if not isinstance(api_key, str) or not api_key: return os.getenv("OPENAI_API_KEY") or None from litellm.secret_managers.main import ( @@ -658,6 +663,10 @@ def _deployment_openai_key(deployment: DeploymentTypedDict) -> str | None: return (get_secret_str(api_key) if api_key.startswith("os.environ/") else api_key) or None +def uses_vendored_codex_catalog() -> bool: + return is_offline_mode() or os.getenv("LITELLM_LOCAL_CODEX_CATALOG", "").strip().lower() == "true" + + def codex_catalog_openai_key( llm_router: Router | None, entries: Sequence[tuple[str, str]], @@ -666,7 +675,11 @@ def codex_catalog_openai_key( """The OpenAI API key the catalog fetch authenticates with: the first one among the listed entries' deployments that routes to OpenAI with every `api_base` and `base_url` it or its credential sets being OpenAI's own, ``OPENAI_API_KEY`` from the environment for one that sets no key. A key for an - OpenAI-compatible server never leaves for OpenAI, and a listing without an OpenAI deployment sends none.""" + OpenAI-compatible server or a FedRAMP key (Codex sends those to ``gov.chatgpt.com``) never leaves for + OpenAI's commercial catalog, a listing without an OpenAI deployment sends none, and no key is looked up + while the catalog stays local.""" + if uses_vendored_codex_catalog(): + return None registry: Final = llm_router.registry() if llm_router is not None else None deployments: Final = ( tuple( 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 a359e01f472..ba209dd9251 100644 --- a/tests/unit/proxy/common_utils/test_codex_catalog_source.py +++ b/tests/unit/proxy/common_utils/test_codex_catalog_source.py @@ -8,7 +8,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter -from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPResponseLimitError from litellm.proxy.common_utils.codex_catalog_source import ( CODEX_CATALOG_BASE_URL, CODEX_CATALOG_MAX_BYTES, @@ -282,6 +282,47 @@ def test_the_cache_keeps_only_the_newest_releases() -> None: assert tree.calls.count(newest_models_url) == 1, tree.calls +def test_releases_no_catalog_answered_leave_the_cache_before_a_fetched_one() -> None: + tree: Final = _RawTree(_tree(RELEASE)) + clock: Final = _Clock() + source: Final = _source(tree, clock, retry_seconds=300, cached_versions=2) + + async def walk() -> None: + await source.stock(RELEASE) + for minor in range(201, 205): + clock.now += 301 + await source.stock(f"0.{minor}.0") + clock.now += 301 + await source.stock(RELEASE) + + asyncio.run(walk()) + models_url: Final = next(url for url in _tree(RELEASE) if url.endswith("models.json")) + assert tree.calls.count(models_url) == 1, tree.calls + + +def test_a_catalog_url_that_redirects_to_itself_serves_the_vendored_pair_and_waits_out_the_retry_window() -> None: + hops: Final[list[str]] = [] + + def redirect_to_itself(request: httpx.Request) -> httpx.Response: + hops.append(str(request.url)) + return httpx.Response(302, headers={"location": str(request.url)}) + + source: Final = CodexCatalogSource( + client=AsyncHTTPHandler(transport=httpx.MockTransport(redirect_to_itself)), clock=_Clock() + ) + + async def ask_twice() -> tuple[str, int, str]: + first: Final = await source.stock(RELEASE) + hops_for_first: Final = len(hops) + second: Final = await source.stock(RELEASE) + return first.source, hops_for_first, second.source + + first_source, hops_for_first, second_source = asyncio.run(ask_twice()) + assert first_source == second_source == vendored_codex_stock().source + assert hops_for_first > 0 + assert len(hops) == hops_for_first + + OPENAI_KEY: Final = "sk-unit-openai-key" OPENAI_ENTRY: Final = {**FRESH_ENTRY, "shell_type": "unified_exec", "base_instructions": "# OpenAI live prompt"} OPENAI_CATALOG: Final = json.dumps({"models": [OPENAI_ENTRY]}).encode() 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 161d13c0486..a12bd76b5fa 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -688,7 +688,9 @@ def test_the_catalog_is_cut_at_the_byte_limit_only_for_a_codex_that_enforces_one listing: Final = tuple( {"id": name, "object": "model", "created": 0, "owned_by": "x", "mode": "chat"} for name in names ) - body: Final = codex_model_list_body(listing, tuple((name, name) for name in names), None, client_version=client_version) + body: Final = codex_model_list_body( + listing, tuple((name, name) for name in names), None, client_version=client_version + ) slugs: Final = [entry["slug"] for entry in json.loads(body)["models"]] assert (len(body.encode()) <= CODEX_CATALOG_BYTE_LIMIT) is capped, len(body.encode()) assert (slugs == list(names)) is not capped, slugs @@ -968,6 +970,12 @@ def test_codex_tool_search_is_offered_off_the_route_and_model_info_and_only_when assert models[0]["context_window"] == 1000000 +@pytest.fixture +def catalog_fetching(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("LITELLM_LOCAL_CODEX_CATALOG", raising=False) + monkeypatch.delenv("LITELLM_OFFLINE_MODE", raising=False) + + def _router_with(*litellm_params: Mapping[str, str]) -> litellm.Router: return litellm.Router( model_list=[ @@ -977,6 +985,7 @@ def _router_with(*litellm_params: Mapping[str, str]) -> litellm.Router: ) +@pytest.mark.usefixtures("catalog_fetching") def test_the_catalog_key_is_the_listed_openai_deployments_own_and_the_environments_only_for_a_keyless_one( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -994,6 +1003,7 @@ def test_the_catalog_key_is_the_listed_openai_deployments_own_and_the_environmen assert codex_catalog_openai_key(router, entries[3:]) == "sk-from-env" +@pytest.mark.usefixtures("catalog_fetching") def test_a_bare_openai_model_name_supplies_its_key_and_a_bare_alias_of_another_provider_does_not( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1009,6 +1019,7 @@ def test_a_bare_openai_model_name_supplies_its_key_and_a_bare_alias_of_another_p assert codex_catalog_openai_key(router, entries[2:]) is None +@pytest.mark.usefixtures("catalog_fetching") def test_the_catalog_key_is_absent_without_an_openai_deployment_or_environment_key( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1018,6 +1029,7 @@ def test_the_catalog_key_is_absent_without_an_openai_deployment_or_environment_k assert codex_catalog_openai_key(None, [("catalog-model-0", "catalog-model-0")]) is None +@pytest.mark.usefixtures("catalog_fetching") def test_a_key_for_an_openai_compatible_server_or_a_redirected_environment_never_goes_to_openai( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1036,7 +1048,7 @@ def test_a_key_for_an_openai_compatible_server_or_a_redirected_environment_never entries: Final = [(f"catalog-model-{index}", f"catalog-model-{index}") for index in range(4)] assert codex_catalog_openai_key(router, entries[:2]) is None assert codex_catalog_openai_key(router, entries) == "sk-real-openai" - assert codex_catalog_openai_key(router, entries[3:]) == "sk-eu-openai" + assert codex_catalog_openai_key(router, entries[3:]) is None monkeypatch.delenv("OPENAI_BASE_URL") monkeypatch.setattr(litellm, "api_base", "https://llm-gateway.internal.example/v1") assert codex_catalog_openai_key(router, entries[:2]) is None @@ -1044,6 +1056,39 @@ def test_a_key_for_an_openai_compatible_server_or_a_redirected_environment_never assert codex_catalog_openai_key(router, entries[:2]) == "sk-through-the-gateway" +@pytest.mark.usefixtures("catalog_fetching") +def test_a_fedramp_key_or_a_gov_openai_base_never_goes_to_openais_commercial_catalog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-proj-fedramp-from-env") + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + router: Final = _router_with( + {"model": "openai/gpt-5.5", "api_key": "sk-gov-base", "api_base": "https://gov.api.openai.com/v1"}, + {"model": "openai/gpt-5.5", "api_key": "sk-proj-fedramp-own"}, + {"model": "openai/gpt-5.5", "api_key": "sk-svcacct-fedramp-own"}, + {"model": "openai/gpt-5.5"}, + {"model": "openai/gpt-5.5", "api_key": "sk-commercial"}, + ) + entries: Final = [(f"catalog-model-{index}", f"catalog-model-{index}") for index in range(5)] + assert codex_catalog_openai_key(router, entries[:4]) is None + assert codex_catalog_openai_key(router, entries) == "sk-commercial" + + +@pytest.mark.parametrize("setting", ["LITELLM_LOCAL_CODEX_CATALOG", "LITELLM_OFFLINE_MODE"]) +@pytest.mark.usefixtures("catalog_fetching") +def test_no_catalog_key_is_looked_up_while_the_catalog_stays_local( + monkeypatch: pytest.MonkeyPatch, setting: str +) -> None: + monkeypatch.setenv("CATALOG_OPENAI_KEY", "sk-from-secret-ref") + router: Final = _router_with({"model": "openai/gpt-5.5", "api_key": "os.environ/CATALOG_OPENAI_KEY"}) + entries: Final = [("catalog-model-0", "catalog-model-0")] + assert codex_catalog_openai_key(router, entries) == "sk-from-secret-ref" + monkeypatch.setenv(setting, "True") + assert codex_catalog_openai_key(router, entries) is None + + +@pytest.mark.usefixtures("catalog_fetching") def test_a_base_url_or_a_credential_naming_another_server_keeps_the_key_from_openai( monkeypatch: pytest.MonkeyPatch, ) -> None: