From 4ec6a35b3f6e36764593cbfbc3c59920b73b910d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 08:29:46 -0800 Subject: [PATCH] fix(proxy): apply JWT route policy before native Codex routes serve LiteLLM models A standard JWT caller now gets a LiteLLM model on /backend-api/codex/models and /backend-api/codex/responses only when the JWT auth flow would admit a direct /v1/responses call for that model, so team_allowed_routes, admin_allowed_routes and RBAC route limits hold on the handoff too. The Fireworks Responses transform also leaves additional_tools input items in place when the tools param is not a list, instead of raising while hoisting --- .../fireworks_ai/responses/transformation.py | 18 ++-- .../codex_passthrough_endpoints.py | 89 +++++++++++++++--- ...t_fireworks_ai_responses_transformation.py | 15 +++ .../test_codex_passthrough_endpoints.py | 91 ++++++++++++++++++- 4 files changed, 193 insertions(+), 20 deletions(-) diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py index c811c06236e..e841cc87c49 100644 --- a/litellm/llms/fireworks_ai/responses/transformation.py +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -21,7 +21,7 @@ from litellm.llms.fireworks_ai.common_utils import ( ) from litellm.llms.openai.responses.instruction_items import developer_items_as_system, instruction_items_folded from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig -from litellm.responses.additional_tools import hoist_additional_tools +from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam from litellm.types.responses.main import DeleteResponseResult @@ -100,11 +100,17 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): headers: dict[str, str], # mutable-ok: overrides the base class signature ) -> dict[str, object]: # mutable-ok: overrides the base class signature validated_input: Final = self._validate_input_param(input) - tools_param: Final = cast( # cast-ok: the base signature leaves the params bag untyped - "Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None", response_api_optional_request_params.get("tools") - ) - hoisted: Final = hoist_additional_tools( - validated_input, tools_param, drop_params=drops_unsupported_params(litellm_params.drop_params) + tools_param: Final = response_api_optional_request_params.get("tools") + hoisted: Final = ( + hoist_additional_tools( + validated_input, + cast( # cast-ok: the base signature leaves the params bag untyped + "Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None", tools_param + ), + drop_params=drops_unsupported_params(litellm_params.drop_params), + ) + if tools_param is None or isinstance(tools_param, list) + else HoistedAdditionalTools(input=validated_input, tools=(), hoisted=()) ) hoisted_params: Final = ( {**response_api_optional_request_params, "tools": list(hoisted.tools)} diff --git a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py index 1f65d3f092b..d321f5b5861 100644 --- a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py @@ -4,6 +4,7 @@ import hmac import json import re from collections.abc import Iterator, Mapping +from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol @@ -25,6 +26,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import check_model_access_helper, check_tools_allowlist +from litellm.proxy.auth.handle_jwt import JWTAuthManager from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.codex_model_catalog import ( @@ -175,7 +177,7 @@ def codex_shared_cache() -> _SharedCache: return proxy_logging_obj.internal_usage_cache.dual_cache -def _may_call_litellm_responses(user_api_key_dict: UserAPIKeyAuth) -> bool: +def _key_may_call_litellm_responses(user_api_key_dict: UserAPIKeyAuth) -> bool: if not RouteChecks.is_allowed_on_proxy(_LITELLM_RESPONSES_PATH): return False try: @@ -187,6 +189,59 @@ def _may_call_litellm_responses(user_api_key_dict: UserAPIKeyAuth) -> bool: return True +async def _jwt_policy_admits_litellm_responses( + gateway_token: str, model: str, request_headers: Mapping[str, str] +) -> bool: + from litellm.proxy.proxy_server import ( + general_settings_view, + jwt_handler, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + try: + _ = await JWTAuthManager.auth_builder( + api_key=gateway_token, + jwt_handler=jwt_handler, + request_data={"model": model}, + general_settings=dict(general_settings_view()), + route=_LITELLM_RESPONSES_PATH, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + request_headers=request_headers, + request_method="POST", + ) + except Exception: # noqa: BLE001 # any JWT policy refusal keeps the caller on ChatGPT's models + return False + return True + + +@dataclass(frozen=True, slots=True) +class _LitellmResponsesAccess: + key_admits: bool + jwt_token: str | None + request_headers: Mapping[str, str] + + async def admits(self, model: str) -> bool: + if not self.key_admits: + return False + if self.jwt_token is None: + return True + return await _jwt_policy_admits_litellm_responses(self.jwt_token, model, self.request_headers) + + +def _litellm_responses_access(user_api_key_dict: UserAPIKeyAuth, request: Request) -> _LitellmResponsesAccess: + standard_jwt: Final = bool(user_api_key_dict.jwt_claims) and not user_api_key_dict.token + return _LitellmResponsesAccess( + key_admits=_key_may_call_litellm_responses(user_api_key_dict), + jwt_token=_gateway_token(request) if standard_jwt else None, + request_headers=request.headers, + ) + + def _account_slugs_cache_key(headers: Mapping[str, str]) -> str | None: account_id: Final = headers.get("chatgpt-account-id") if not account_id: @@ -264,13 +319,14 @@ async def _serves_through_litellm( model: object, llm_router: "Router | None", user_api_key_dict: UserAPIKeyAuth, + access: _LitellmResponsesAccess, caller: str, headers: Mapping[str, str], cache: _SharedCache, ) -> bool: - if not routes_to_litellm(model, llm_router, user_api_key_dict.team_id) or not _may_call_litellm_responses( - user_api_key_dict - ): + if not isinstance(model, str) or not routes_to_litellm(model, llm_router, user_api_key_dict.team_id): + return False + if not await access.admits(model): return False account_slugs: Final = await _account_slugs(headers, caller, cache) return account_slugs is not None and model not in account_slugs @@ -305,6 +361,7 @@ async def native_turn_serves_through_litellm( model, llm_router, user_api_key_dict, + _litellm_responses_access(user_api_key_dict, request), _models_etag_caller(user_api_key_dict, gateway_token), _forwarded_headers(request), codex_shared_cache(), @@ -523,11 +580,11 @@ def _chatgpt_body(body: Mapping[str, object], client_model: str | None) -> Mappi async def _litellm_catalog_rows( - user_api_key_dict: UserAPIKeyAuth, llm_router: "Router | None" + user_api_key_dict: UserAPIKeyAuth, access: _LitellmResponsesAccess, llm_router: "Router | None" ) -> tuple[CodexCatalogRow, ...]: from litellm.proxy.proxy_server import codex_catalog_rows_for_key - if llm_router is None or not _may_call_litellm_responses(user_api_key_dict): + if llm_router is None or not access.key_admits: return () try: rows: Final = await codex_catalog_rows_for_key(user_api_key_dict) @@ -537,7 +594,9 @@ async def _litellm_catalog_rows( type(error).__name__, ) return () - return tuple(row for row in rows if routes_to_litellm(row.id, llm_router, user_api_key_dict.team_id)) + routed: Final = tuple(row for row in rows if routes_to_litellm(row.id, llm_router, user_api_key_dict.team_id)) + admitted: Final = await asyncio.gather(*(access.admits(row.id) for row in routed)) + return tuple(row for row, row_admitted in zip(routed, admitted, strict=True) if row_admitted) def _models_etag_caller(user_api_key_dict: UserAPIKeyAuth, gateway_token: str) -> str: @@ -570,6 +629,7 @@ async def _cached_models_etag(caller: str, upstream_etag: str, cache: _SharedCac async def _key_models_etag( upstream_etag: str, user_api_key_dict: UserAPIKeyAuth, + access: _LitellmResponsesAccess, caller: str, llm_router: "Router | None", cache: _SharedCache, @@ -577,7 +637,7 @@ async def _key_models_etag( cached: Final = await _cached_models_etag(caller, upstream_etag, cache) if cached is not None: return cached - rows: Final = await _litellm_catalog_rows(user_api_key_dict, llm_router) + rows: Final = await _litellm_catalog_rows(user_api_key_dict, access, llm_router) return await _remembered_models_etag(caller, upstream_etag, rows, cache) @@ -596,6 +656,7 @@ def _with_etag(response: Response, content: bytes | str, etag: str | None) -> Re async def _catalog( request: Request, user_api_key_dict: UserAPIKeyAuth, + access: _LitellmResponsesAccess, caller: str, headers: Mapping[str, str], cache: _SharedCache, @@ -619,7 +680,7 @@ async def _catalog( account_slugs_key: Final = _account_slugs_cache_key(headers) if account_slugs_key is not None: await _remember_account_slugs(account_slugs_key, slugs, cache) - rows: Final = await _litellm_catalog_rows(user_api_key_dict, llm_router) + rows: Final = await _litellm_catalog_rows(user_api_key_dict, access, llm_router) upstream_etag: Final = upstream.headers.get(_ETAG_HEADER) etag: Final = None if upstream_etag is None else await _remembered_models_etag(caller, upstream_etag, rows, cache) if not rows: @@ -634,6 +695,7 @@ async def _responses( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, + access: _LitellmResponsesAccess, caller: str, headers: Mapping[str, str], cache: _SharedCache, @@ -641,7 +703,7 @@ async def _responses( from litellm.proxy.proxy_server import llm_router body: Final = _JSON_OBJECT.validate_python(await read_request_body(request)) - if await _serves_through_litellm(body.get("model"), llm_router, user_api_key_dict, caller, headers, cache): + if await _serves_through_litellm(body.get("model"), llm_router, user_api_key_dict, access, caller, headers, cache): litellm_body: Final = _litellm_body(body, llm_router, user_api_key_dict.team_id, headers) await check_tools_allowlist( request_body=dict(litellm_body), @@ -667,7 +729,7 @@ async def _responses( upstream_models_etag: Final = upstream.headers.get(_MODELS_ETAG_HEADER) if upstream_models_etag is not None: upstream.headers[_MODELS_ETAG_HEADER] = await _key_models_etag( - upstream_models_etag, user_api_key_dict, caller, llm_router, cache + upstream_models_etag, user_api_key_dict, access, caller, llm_router, cache ) return upstream @@ -718,6 +780,7 @@ async def codex_passthrough_route( headers: Final = _forwarded_headers(request) caller: Final = _models_etag_caller(user_api_key_dict, gateway_token) + access: Final = _litellm_responses_access(user_api_key_dict, request) if endpoint == "models": - return await _catalog(request, user_api_key_dict, caller, headers, cache) - return await _responses(request, fastapi_response, user_api_key_dict, caller, headers, cache) + return await _catalog(request, user_api_key_dict, access, caller, headers, cache) + return await _responses(request, fastapi_response, user_api_key_dict, access, caller, headers, cache) diff --git a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py index a786565fd74..b5aad23c4f0 100644 --- a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py +++ b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -539,6 +539,21 @@ def test_transform_request_merges_codex_additional_tools_into_existing_tools_onl assert tuple(request["input"]) == (user_item,) +@pytest.mark.parametrize("drop_params", [True, None], ids=["drop_params", "no-drop_params"]) +def test_transform_request_leaves_codex_additional_tools_in_input_when_tools_is_not_a_list( + drop_params: bool | None, +) -> None: + user_item: Final = {"role": "user", "content": "Sleep for two seconds."} + request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request( + model="accounts/fireworks/models/glm-5p3", + input=cast(ResponseInputParam, [CODEX_ADDITIONAL_TOOLS_ITEM, user_item]), + response_api_optional_request_params={"tools": {"type": "function", "name": "get_weather"}}, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + assert tuple(request["input"]) == (CODEX_ADDITIONAL_TOOLS_ITEM, user_item) + + def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None: client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) pydantic_input: Final = cast( diff --git a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py index 9cbbd936889..1a4751b4598 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py @@ -1,14 +1,18 @@ import asyncio import base64 import json +import time from collections.abc import Callable, Iterator, Sequence from datetime import datetime from typing import Final import httpx +import jwt import pytest import respx +from cryptography.hazmat.primitives.asymmetric import rsa from fastapi.testclient import TestClient +from jwt.algorithms import RSAAlgorithm import litellm from litellm import Router @@ -16,7 +20,8 @@ from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.llms.chatgpt.common_utils import OPENAI_AUTH_CLAIM_KEY from litellm.proxy import proxy_server -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_TeamTableCachedObj, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.codex_model_catalog import bundled_codex_models from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import ( @@ -986,6 +991,90 @@ def test_proxy_allowed_routes_keep_litellm_models_off_native_routes_unless_v1_re assert chatgpt.called is not through_litellm +_JWT_ISSUER: Final = "https://idp.codex-test.example" +_JWT_JWKS_URL: Final = f"{_JWT_ISSUER}/keys" +_JWT_AUDIENCE: Final = "litellm-proxy" +_JWT_KID: Final = "codex-test-key" +_JWT_TEAM: Final = "codex-team" + + +def _team_jwt(monkeypatch: pytest.MonkeyPatch, team_allowed_routes: list[str]) -> str: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_jwk: Final = { + **json.loads(RSAAlgorithm.to_jwk(private_key.public_key())), + "kid": _JWT_KID, + "alg": "RS256", + "use": "sig", + } + keys_cache: Final = DualCache() + keys_cache.set_cache(key=f"litellm_jwt_auth_keys_{_JWT_JWKS_URL}", value=[public_jwk]) + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=keys_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + issuers=[{"issuer": _JWT_ISSUER, "jwks_url": _JWT_JWKS_URL, "audience": _JWT_AUDIENCE}], + team_ids_jwt_field="groups", + team_allowed_routes=team_allowed_routes, + ), + ) + monkeypatch.setattr(proxy_server, "jwt_handler", handler) + monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True}) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + proxy_server.user_api_key_cache.set_cache( + key=f"team_id:{_JWT_TEAM}", + value=LiteLLM_TeamTableCachedObj(team_id=_JWT_TEAM, models=["chatgpt-native", "open-weights"]), + ) + issued_at: Final = int(time.time()) + return jwt.encode( + { + "sub": "codex-user", + "iss": _JWT_ISSUER, + "aud": _JWT_AUDIENCE, + "iat": issued_at, + "exp": issued_at + 300, + "groups": [_JWT_TEAM], + }, + private_key, + algorithm="RS256", + headers={"kid": _JWT_KID}, + ) + + +@pytest.mark.parametrize( + ("team_allowed_routes", "through_litellm"), + [ + pytest.param(["/backend-api/codex/*"], False, id="team allowed only the native routes"), + pytest.param(["/backend-api/codex/*", "openai_routes"], True, id="team allowed openai routes too"), + ], +) +def test_jwt_team_routes_keep_litellm_models_off_native_routes_unless_v1_responses_is_allowed( + client: TestClient, + respx_mock: respx.MockRouter, + mixed_router: Router, + monkeypatch: pytest.MonkeyPatch, + team_allowed_routes: list[str], + through_litellm: bool, +) -> None: + token: Final = _team_jwt(monkeypatch, team_allowed_routes) + _account_catalog(respx_mock) + provider: Final = _open_weights_provider(respx_mock) + chatgpt: Final = _chatgpt_responses(respx_mock) + catalog: Final = client.get( + "/backend-api/codex/models?client_version=1.0.0", headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": token} + ) + response: Final = client.post( + "/backend-api/codex/responses", + json={"model": "open-weights", "input": "Hello", "stream": True}, + headers={**_ACCOUNT_HEADERS, "x-litellm-api-key": token}, + ) + assert (catalog.status_code, response.status_code) == (200, 200), response.text + assert ("open-weights" in [model["slug"] for model in catalog.json()["models"]]) is through_litellm + assert provider.called is through_litellm + assert chatgpt.called is not through_litellm + + _USER_TURN: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Hello"}]} _ASSISTANT_TURN: Final = {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Hi"}]} _FIREWORKS_REASONING: Final = {