diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b33e2fe7ff6..f83011835fd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2638,6 +2638,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): default=None, description="Set-up pass-through endpoints for provider-specific endpoints. Docs - https://docs.litellm.ai/docs/proxy/pass_through", ) + enable_openai_websocket_passthrough: bool | None = Field( + default=None, + description="Serve the OpenAI pass-through WebSocket route, which relays frames to OpenAI under the proxy's own provider credential without reading them. Off by default.", + ) user_header_name: str | None = Field( None, description="[DEPRECATED] Use 'user_header_mappings' instead. When set, the header value is treated as the end user id unless overridden by user_header_mappings.", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f83f0303deb..dc693317de0 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2352,6 +2352,13 @@ async def _backfill_null_user_email( return updated_row +class UserNotFoundError(ValueError): + """The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer.""" + + def __init__(self, user_id: str) -> None: + super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.") + + @log_db_metrics async def get_user_object( user_id: str | None, @@ -2457,7 +2464,7 @@ async def get_user_object( value=None, last_db_access_time=last_db_access_time, ) - raise Exception + raise UserNotFoundError(user_id=user_id) if response.organization_memberships is not None and len(response.organization_memberships) > 0: # dump each organization membership to type LiteLLM_OrganizationMembershipTable @@ -2493,7 +2500,9 @@ async def get_user_object( ) return _response - except Exception as e: # if user not in db + except UserNotFoundError: + raise + except Exception as e: _log_budget_lookup_failure("user", e) raise ValueError( f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call. Got error - {e}" @@ -4155,6 +4164,79 @@ async def _granted_model_lists( ) +async def _user_object_or_none( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_UserTable | None: + try: + return await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + except UserNotFoundError: + return None + + +async def enforced_model_allowlists( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[Sequence[str], ...]: + """One model allowlist per level that ``common_checks`` enforces on a request from this identity.""" + key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) + if prisma_client is None: + return (key_models, tuple(valid_token.team_models or ())) + team_object: Final = ( + None + if valid_token.team_id is None + else await get_team_object( + team_id=valid_token.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + user_object: Final = ( + None + if team_object is not None + else await _user_object_or_none( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + project_object: Final = ( + None + if valid_token.project_id is None + else await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + return ( + key_models, + team_object.models if team_object is not None else (), + await _team_member_granted_models( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + user_object.models if user_object is not None else (), + project_object.models if project_object is not None else (), + ) + + async def collect_matched_model_access_groups( model: str | Sequence[str] | None, valid_token: UserAPIKeyAuth | None, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6b1d6405a6a..32da2658b99 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -13,14 +13,16 @@ import inspect import json import os import re -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket from fastapi.responses import StreamingResponse from starlette.websockets import WebSocketState +from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider @@ -35,6 +37,7 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import enforced_model_allowlists from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( @@ -1773,7 +1776,7 @@ def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str] def get_vertex_pass_through_handler( - call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here + call_type: Literal["discovery", "aiplatform"], ) -> BaseVertexAIPassThroughHandler: if call_type == "discovery": return VertexAIDiscoveryPassThroughHandler() @@ -2340,9 +2343,102 @@ _OPENAI_WS_ALL_MODEL_ACCESS: Final = frozenset( ) -def _key_has_model_restrictions(user_api_key_dict: UserAPIKeyAuth) -> bool: - scoped_models: Final = (*user_api_key_dict.models, *user_api_key_dict.team_models) - return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for model in scoped_models) +def _has_model_restrictions(model_allowlists: tuple[Sequence[str], ...]) -> bool: + return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for allowlist in model_allowlists for model in allowlist) + + +@dataclass(frozen=True, slots=True) +class _OpenAIWebsocketRefusal: + close_reason: str + message: str + + +class _OpenAIWebsocketErrorDetail(TypedDict): + type: ReadOnly[Literal["invalid_request_error"]] + message: ReadOnly[str] + + +class _OpenAIWebsocketErrorFrame(TypedDict): + type: ReadOnly[Literal["error"]] + error: ReadOnly[_OpenAIWebsocketErrorDetail] + + +_OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( + close_reason="OpenAI websocket passthrough is disabled", + message=( + "OpenAI websocket passthrough is disabled on this gateway. A proxy admin can turn it on by " + "setting general_settings.enable_openai_websocket_passthrough to true." + ), +) + +_OPENAI_WS_MODEL_RESTRICTED_REFUSAL: Final = _OpenAIWebsocketRefusal( + close_reason="Keys with model restrictions cannot use OpenAI websocket passthrough", + message=( + "Keys with model restrictions cannot use OpenAI websocket passthrough, because this route " + "relays frames to the provider without reading which model they ask for." + ), +) + + +def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, object]) -> bool: + setting: Final = general_settings.get("enable_openai_websocket_passthrough") + if isinstance(setting, str): + return str_to_bool(setting) is True + return setting is True + + +class _OpenAIWebsocketModelAllowlists(Protocol): + async def __call__(self, valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: ... + + +async def _openai_websocket_refusal( + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], + model_allowlists: _OpenAIWebsocketModelAllowlists, +) -> _OpenAIWebsocketRefusal | None: + if not _is_openai_websocket_passthrough_enabled(general_settings): + return _OPENAI_WS_DISABLED_REFUSAL + if _has_model_restrictions(await model_allowlists(user_api_key_dict)): + return _OPENAI_WS_MODEL_RESTRICTED_REFUSAL + return None + + +class _OpenAIWebsocketRelay(Protocol): + async def __call__( + self, + *, + websocket: WebSocket, + target: str, + custom_headers: dict[str, str], # mutable-ok: the relay takes a plain dict of upstream headers + user_api_key_dict: UserAPIKeyAuth, + forward_headers: bool, + endpoint: str, + accept_websocket: bool, + ) -> None: ... + + +def _proxy_general_settings() -> Mapping[str, object]: + from litellm.proxy.proxy_server import general_settings + + return general_settings + + +def _openai_websocket_relay() -> _OpenAIWebsocketRelay: + return websocket_passthrough_request + + +def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + async def resolve(valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: + return await enforced_model_allowlists( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + return resolve @router.websocket("/openai_passthrough/{endpoint:path}") @@ -2351,13 +2447,27 @@ async def openai_websocket_proxy_route( websocket: WebSocket, endpoint: str, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], + general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)], + relay: Annotated[_OpenAIWebsocketRelay, Depends(_openai_websocket_relay)], + model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)], ) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" - if _key_has_model_restrictions(user_api_key_dict): - await websocket.close( - code=1008, - reason="Keys with model restrictions cannot use OpenAI websocket passthrough", - ) + requested_subprotocols: Final = tuple( + protocol.strip() + for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") + if protocol.strip() + ) + negotiated_subprotocol: Final = requested_subprotocols[0] if requested_subprotocols else None + + refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists) + if refusal is not None: + await websocket.accept(subprotocol=negotiated_subprotocol) + error_frame: Final[_OpenAIWebsocketErrorFrame] = { + "type": "error", + "error": {"type": "invalid_request_error", "message": refusal.message}, + } + await websocket.send_text(json.dumps(error_frame)) + await websocket.close(code=1008, reason=refusal.close_reason) return base_target_url: Final = os.getenv("OPENAI_API_BASE") or "https://api.openai.com/" @@ -2393,14 +2503,9 @@ async def openai_websocket_proxy_route( "Authorization": f"Bearer {openai_api_key}" } - requested_subprotocols: Final = tuple( - protocol.strip() - for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") - if protocol.strip() - ) - await websocket.accept(subprotocol=requested_subprotocols[0] if requested_subprotocols else None) + await websocket.accept(subprotocol=negotiated_subprotocol) - await websocket_passthrough_request( + await relay( websocket=websocket, target=wss_target, custom_headers=custom_headers, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61bcdea94d4..026d8d2228c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6810,6 +6810,11 @@ class ProxyConfig: else: general_settings["apply_user_budget_to_team_keys"] = db_value if db_value is None else bool(db_value) + if "enable_openai_websocket_passthrough" not in self._yaml_general_settings_keys: + general_settings["enable_openai_websocket_passthrough"] = _general_settings.get( + "enable_openai_websocket_passthrough" + ) + ## STORE MODEL IN DB ## if "store_model_in_db" in _general_settings: value = _general_settings["store_model_in_db"] diff --git a/test-quality-budget.json b/test-quality-budget.json index 7ca563d25af..3c12371f02f 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -3,7 +3,7 @@ "limit": 733 }, "TQ002": { - "limit": 741 + "limit": 737 }, "TQ003": { "limit": 62 @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 11003 + "limit": 10993 } } diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index be83ca57e76..2284a05b2e9 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -126,14 +126,10 @@ def invalid_sso_user_defined_values(): def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_values): """Test generating JWT token with valid user role""" - token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - valid_sso_user_defined_values - ) + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") # Check that decrypted_token is not None before using json.loads assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -159,9 +155,7 @@ def test_get_cli_jwt_auth_token_includes_team_alias(valid_sso_user_defined_value team_alias="test-team", ) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -188,9 +182,7 @@ def test_get_cli_jwt_auth_token_carries_team_grants_not_user_allowlist( team_model_aliases={"team-fast": "gpt-4.1-mini"}, ) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -207,9 +199,7 @@ def test_get_cli_jwt_auth_token_keeps_user_allowlist_when_no_team( """A session token with no team bound still carries the user's own allowlist.""" token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -222,12 +212,8 @@ def test_get_experimental_ui_login_jwt_auth_token_uses_10_min_expiry( valid_sso_user_defined_values, ): """Test that Experimental UI token uses fixed 10-minute expiry (does not use LITELLM_UI_SESSION_DURATION).""" - token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - valid_sso_user_defined_values - ) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) @@ -244,43 +230,33 @@ def test_experimental_ui_token_ignores_litellm_ui_session_duration( Experimental UI intentionally uses fixed 10-min expiry. If this test fails, the constant was incorrectly wired to the experimental flow.""" # Default LITELLM_UI_SESSION_DURATION is "24h" - token must still expire in ~10 min - token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - valid_sso_user_defined_values - ) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) now = get_utc_datetime() # Must be ~10 min, NOT 24h. If LITELLM_UI_SESSION_DURATION were incorrectly used, this would fail. - assert expires <= now + timedelta( - minutes=11 - ), "Experimental UI must use 10-min expiry, not LITELLM_UI_SESSION_DURATION" + assert expires <= now + timedelta(minutes=11), ( + "Experimental UI must use 10-min expiry, not LITELLM_UI_SESSION_DURATION" + ) def test_get_experimental_ui_login_jwt_auth_token_invalid( invalid_sso_user_defined_values, ): """Test generating JWT token with missing user role""" - with pytest.raises(Exception, match='User role is required for experimental UI login') as exc_info: - ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - invalid_sso_user_defined_values - ) + with pytest.raises(Exception, match="User role is required for experimental UI login") as exc_info: + ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(invalid_sso_user_defined_values) assert str(exc_info.value) == "User role is required for experimental UI login" -def test_get_key_object_from_ui_hash_key_valid( - valid_sso_user_defined_values, monkeypatch -): +def test_get_key_object_from_ui_hash_key_valid(valid_sso_user_defined_values, monkeypatch): """Test getting key object from valid UI hash key""" monkeypatch.setenv("EXPERIMENTAL_UI_LOGIN", "True") # Generate a valid token - token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - valid_sso_user_defined_values - ) + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) # Get key object key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) @@ -309,9 +285,7 @@ def test_get_key_object_from_ui_hash_key_invalid(): ("project", ProxyErrorTypes.project_model_access_denied), ], ) -def test_can_object_call_model_denials_return_forbidden( - object_type, expected_error_type -): +def test_can_object_call_model_denials_return_forbidden(object_type, expected_error_type): with pytest.raises(ProxyException) as exc_info: _can_object_call_model( model="restricted-model", @@ -568,9 +542,7 @@ async def test_get_key_object_should_reconnect_once_on_db_connection_error(): @pytest.mark.asyncio async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_error(): mock_prisma_client = MagicMock() - mock_prisma_client.get_data = AsyncMock( - side_effect=httpx.ConnectError("db not reachable after outage") - ) + mock_prisma_client.get_data = AsyncMock(side_effect=httpx.ConnectError("db not reachable after outage")) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) mock_cache = MagicMock() @@ -613,9 +585,7 @@ class TestAuthCacheRedisWritePolicy: @pytest.mark.asyncio async def test_get_key_object_db_load_publishes_to_redis(self): mock_prisma_client = MagicMock() - mock_prisma_client.get_data = AsyncMock( - return_value=UserAPIKeyAuth(token="hashed-token-db") - ) + mock_prisma_client.get_data = AsyncMock(return_value=UserAPIKeyAuth(token="hashed-token-db")) fake_redis = _fake_redis_cache() cache = UserApiKeyCache() @@ -630,8 +600,7 @@ class TestAuthCacheRedisWritePolicy: assert key_obj.token == "hashed-token-db" fake_redis.async_set_cache.assert_awaited_once() assert ( - fake_redis.async_set_cache.await_args.kwargs.get("key") - or fake_redis.async_set_cache.await_args.args[0] + fake_redis.async_set_cache.await_args.kwargs.get("key") or fake_redis.async_set_cache.await_args.args[0] ) == "hashed-token-db" @@ -640,9 +609,7 @@ def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -664,9 +631,7 @@ def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values assert expires >= get_utc_datetime() + timedelta(hours=23, minutes=59) -def test_get_cli_jwt_auth_token_custom_expiration( - valid_sso_user_defined_values, monkeypatch -): +def test_get_cli_jwt_auth_token_custom_expiration(valid_sso_user_defined_values, monkeypatch): """Test generating CLI JWT token with custom expiration via environment variable""" import importlib @@ -681,14 +646,10 @@ def test_get_cli_jwt_auth_token_custom_expiration( # Also reload auth_checks to pick up the new constant value importlib.reload(auth_checks) - token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token( - valid_sso_user_defined_values - ) + token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -706,18 +667,12 @@ def test_get_cli_jwt_auth_token_unique_per_session(valid_sso_user_defined_values from litellm.constants import CLI_SESSION_KEY_PREFIX def _decode(token: str) -> dict: - decrypted = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted is not None return json.loads(decrypted) - first = _decode( - ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) - ) - second = _decode( - ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) - ) + first = _decode(ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values)) + second = _decode(ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values)) assert first["token"].startswith(f"{CLI_SESSION_KEY_PREFIX}-") assert second["token"].startswith(f"{CLI_SESSION_KEY_PREFIX}-") @@ -740,9 +695,7 @@ def test_get_cli_jwt_auth_token_applies_fallback_budget(valid_sso_user_defined_v def test_get_cli_jwt_auth_token_no_fallback_when_budget_provided( valid_sso_user_defined_values, ): - token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - valid_sso_user_defined_values, max_budget=None - ) + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values, max_budget=None) decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted is not None assert json.loads(decrypted).get("max_budget") is None @@ -945,9 +898,7 @@ async def test_get_user_object_upsert_includes_user_email(): mock_prisma_client.db.litellm_usertable.create.assert_called_once() creation_args = mock_prisma_client.db.litellm_usertable.create.call_args[1]["data"] - assert ( - "user_email" in creation_args - ), "user_email should be included when upserting a new user" + assert "user_email" in creation_args, "user_email should be included when upserting a new user" assert creation_args["user_email"] == "test@example.com" assert creation_args["user_id"] == "new_test_user" @@ -962,12 +913,8 @@ async def test_get_user_object_backfills_null_email_from_cache_hit(): was returned unchanged and the DB was never updated. """ cache = UserApiKeyCache() - existing = LiteLLM_UserTable( - user_id="jwt-user-1", user_email=None, user_role="internal_user" - ) - await cache.async_set_cache( - key="jwt-user-1", value=existing, model_type=LiteLLM_UserTable - ) + existing = LiteLLM_UserTable(user_id="jwt-user-1", user_email=None, user_role="internal_user") + await cache.async_set_cache(key="jwt-user-1", value=existing, model_type=LiteLLM_UserTable) mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1) @@ -996,9 +943,7 @@ async def test_get_user_object_backfills_null_email_from_cache_hit(): assert update_kwargs["where"] == {"user_id": "jwt-user-1", "user_email": None} assert update_kwargs["data"]["user_email"] == "jwt-user-1@example.com" - refreshed = await cache.async_get_cache( - key="jwt-user-1", model_type=LiteLLM_UserTable - ) + refreshed = await cache.async_get_cache(key="jwt-user-1", model_type=LiteLLM_UserTable) assert refreshed is not None assert refreshed.user_email == "jwt-user-1@example.com" @@ -1010,9 +955,7 @@ async def test_get_user_object_backfills_null_email_from_db_read(): backfilled from the JWT-provided email before it is cached and returned. """ cache = UserApiKeyCache() - db_row = LiteLLM_UserTable( - user_id="jwt-user-3", user_email=None, user_role="internal_user" - ) + db_row = LiteLLM_UserTable(user_id="jwt-user-3", user_email=None, user_role="internal_user") backfilled_row = LiteLLM_UserTable( user_id="jwt-user-3", user_email="jwt-user-3@example.com", @@ -1020,15 +963,11 @@ async def test_get_user_object_backfills_null_email_from_db_read(): ) mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=[db_row, backfilled_row] - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=[db_row, backfilled_row]) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1) - with patch( - "litellm.proxy.auth.auth_checks._should_check_db", return_value=True - ): + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): result = await get_user_object( user_id="jwt-user-3", prisma_client=mock_prisma_client, @@ -1042,9 +981,7 @@ async def test_get_user_object_backfills_null_email_from_db_read(): assert result.user_email == "jwt-user-3@example.com" mock_prisma_client.db.litellm_usertable.update_many.assert_called_once() - refreshed = await cache.async_get_cache( - key="jwt-user-3", model_type=LiteLLM_UserTable - ) + refreshed = await cache.async_get_cache(key="jwt-user-3", model_type=LiteLLM_UserTable) assert refreshed is not None assert refreshed.user_email == "jwt-user-3@example.com" @@ -1062,9 +999,7 @@ async def test_get_user_object_does_not_overwrite_existing_email(): user_email="operator-set@example.com", user_role="internal_user", ) - await cache.async_set_cache( - key="jwt-user-2", value=existing, model_type=LiteLLM_UserTable - ) + await cache.async_set_cache(key="jwt-user-2", value=existing, model_type=LiteLLM_UserTable) mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=0) @@ -1091,12 +1026,8 @@ async def test_get_user_object_backfill_race_prefers_db_email(): with the value the DB accepted, not this request's proposed email. """ cache = UserApiKeyCache() - existing = LiteLLM_UserTable( - user_id="jwt-user-4", user_email=None, user_role="internal_user" - ) - await cache.async_set_cache( - key="jwt-user-4", value=existing, model_type=LiteLLM_UserTable - ) + existing = LiteLLM_UserTable(user_id="jwt-user-4", user_email=None, user_role="internal_user") + await cache.async_set_cache(key="jwt-user-4", value=existing, model_type=LiteLLM_UserTable) winner_row = LiteLLM_UserTable( user_id="jwt-user-4", @@ -1105,9 +1036,7 @@ async def test_get_user_object_backfill_race_prefers_db_email(): ) mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=winner_row - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=winner_row) result = await get_user_object( user_id="jwt-user-4", @@ -1121,9 +1050,7 @@ async def test_get_user_object_backfill_race_prefers_db_email(): assert result is not None assert result.user_email == "winner@example.com" - refreshed = await cache.async_get_cache( - key="jwt-user-4", model_type=LiteLLM_UserTable - ) + refreshed = await cache.async_get_cache(key="jwt-user-4", model_type=LiteLLM_UserTable) assert refreshed is not None assert refreshed.user_email == "winner@example.com" @@ -1138,12 +1065,8 @@ async def test_get_user_object_backfill_caches_persisted_email_not_proposed(): optimistically caching the proposed email would serve a stale value. """ cache = UserApiKeyCache() - existing = LiteLLM_UserTable( - user_id="jwt-user-5", user_email=None, user_role="internal_user" - ) - await cache.async_set_cache( - key="jwt-user-5", value=existing, model_type=LiteLLM_UserTable - ) + existing = LiteLLM_UserTable(user_id="jwt-user-5", user_email=None, user_role="internal_user") + await cache.async_set_cache(key="jwt-user-5", value=existing, model_type=LiteLLM_UserTable) persisted_row = LiteLLM_UserTable( user_id="jwt-user-5", @@ -1152,9 +1075,7 @@ async def test_get_user_object_backfill_caches_persisted_email_not_proposed(): ) mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1) - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=persisted_row - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=persisted_row) result = await get_user_object( user_id="jwt-user-5", @@ -1168,9 +1089,7 @@ async def test_get_user_object_backfill_caches_persisted_email_not_proposed(): assert result is not None assert result.user_email == "admin-edited@example.com" - refreshed = await cache.async_get_cache( - key="jwt-user-5", model_type=LiteLLM_UserTable - ) + refreshed = await cache.async_get_cache(key="jwt-user-5", model_type=LiteLLM_UserTable) assert refreshed is not None assert refreshed.user_email == "admin-edited@example.com" @@ -1224,10 +1143,7 @@ async def test_get_user_object_upsert_routes_default_team_to_membership(monkeypa mock_add_to_team.assert_awaited_once() passed_teams = mock_add_to_team.await_args[1]["teams"] assert [team.team_id for team in passed_teams] == ["default-team"] - assert ( - mock_add_to_team.await_args[1]["user_api_key_dict"].user_role - == LitellmUserRoles.PROXY_ADMIN - ) + assert mock_add_to_team.await_args[1]["user_api_key_dict"].user_role == LitellmUserRoles.PROXY_ADMIN def test_log_budget_lookup_failure_dry_run(): @@ -1252,9 +1168,7 @@ def test_log_budget_lookup_failure_skips_user_not_found(): @pytest.mark.asyncio -@patch( - "litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock -) +@patch("litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock) async def test_get_team_db_check_calls_new_team_on_upsert(mock_new_team, monkeypatch): """ Test that _get_team_db_check correctly calls the `new_team` function @@ -1288,12 +1202,8 @@ async def test_get_team_db_check_calls_new_team_on_upsert(mock_new_team, monkeyp @pytest.mark.asyncio -@patch( - "litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock -) -async def test_get_team_db_check_does_not_call_new_team_if_exists( - mock_new_team, monkeypatch -): +@patch("litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock) +async def test_get_team_db_check_does_not_call_new_team_if_exists(mock_new_team, monkeypatch): """ Test that _get_team_db_check does NOT call the `new_team` function if the team already exists in the database. @@ -1327,9 +1237,7 @@ async def test_get_team_db_check_does_not_call_new_team_if_exists( (MagicMock(), MagicMock(), True), # No vector stores to run ], ) -async def test_vector_store_access_check_early_returns( - prisma_client, vector_store_registry, expected_result -): +async def test_vector_store_access_check_early_returns(prisma_client, vector_store_registry, expected_result): """Test vector_store_access_check returns True for early exit conditions""" request_body = {"messages": [{"role": "user", "content": "test"}]} @@ -1411,9 +1319,7 @@ async def test_vector_store_access_check_skips_db_lookup_when_no_vector_stores_r ), # Partial access ], ) -def test_can_object_call_vector_stores_scenarios( - object_permissions, vector_store_ids, should_raise, error_type -): +def test_can_object_call_vector_stores_scenarios(object_permissions, vector_store_ids, should_raise, error_type): """Test _can_object_call_vector_stores with various permission scenarios""" # Convert dict to object if not None if object_permissions is not None: @@ -1421,11 +1327,7 @@ def test_can_object_call_vector_stores_scenarios( mock_permissions.vector_stores = object_permissions["vector_stores"] object_permissions = mock_permissions - object_type = ( - "key" - if error_type == ProxyErrorTypes.key_vector_store_access_denied - else "team" - ) + object_type = "key" if error_type == ProxyErrorTypes.key_vector_store_access_denied else "team" if should_raise: with pytest.raises(ProxyException) as exc_info: @@ -1460,9 +1362,7 @@ async def test_vector_store_access_check_with_permissions(): mock_prisma_client = MagicMock() mock_permissions = MagicMock() mock_permissions.vector_stores = ["store-1", "store-2"] - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=mock_permissions - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=mock_permissions) mock_vector_store_registry = MagicMock() mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-1"] @@ -1508,14 +1408,10 @@ async def test_vector_store_access_check_with_team_permissions(): mock_prisma_client = MagicMock() team_permissions = MagicMock() team_permissions.vector_stores = ["team-store-allowed"] - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=team_permissions - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=team_permissions) mock_vector_store_registry = MagicMock() - mock_vector_store_registry.get_vector_store_ids_to_run.return_value = [ - "team-store-allowed" - ] + mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["team-store-allowed"] with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), @@ -1529,9 +1425,7 @@ async def test_vector_store_access_check_with_team_permissions(): assert result is True - mock_vector_store_registry.get_vector_store_ids_to_run.return_value = [ - "team-store-denied" - ] + mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["team-store-denied"] with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), @@ -2092,9 +1986,7 @@ async def test_get_tag_objects_batch(): mock_cache.async_set_cache = AsyncMock() # Mock DB to return all uncached tags in ONE query - mock_prisma.db.litellm_tagtable.find_many = AsyncMock( - return_value=[uncached_tag_1, uncached_tag_2, uncached_tag_3] - ) + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(return_value=[uncached_tag_1, uncached_tag_2, uncached_tag_3]) # Call batch fetch tag_objects = await get_tag_objects_batch( @@ -2196,9 +2088,7 @@ async def test_get_tag_objects_batch_never_queries_db_for_unregistered_tags(): from litellm.proxy.auth.auth_checks import get_tag_objects_batch mock_prisma = MagicMock() - mock_prisma.db.litellm_tagtable.find_many = AsyncMock( - return_value=[_tag_registry_row("some-other-tag")] - ) + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(return_value=[_tag_registry_row("some-other-tag")]) cache = UserApiKeyCache() first = await get_tag_objects_batch( @@ -2209,9 +2099,7 @@ async def test_get_tag_objects_batch_never_queries_db_for_unregistered_tags(): assert first == {} # The only query is the names-only registry fetch; the tag itself is never looked up. - mock_prisma.db.litellm_tagtable.find_many.assert_called_once_with( - take=TAG_REGISTRY_MAX_SIZE + 1 - ) + mock_prisma.db.litellm_tagtable.find_many.assert_called_once_with(take=TAG_REGISTRY_MAX_SIZE + 1) second = await get_tag_objects_batch( tag_names=["unregistered-tag"], @@ -2380,9 +2268,7 @@ async def test_get_tag_objects_batch_oversized_registry_falls_back_and_stops_ref """Past the cap the registry is unusable: keep the old per-tag path, but stop rebuilding it.""" from litellm.proxy.auth.auth_checks import get_tag_objects_batch - oversized = [ - _tag_registry_row(f"tag-{index}") for index in range(TAG_REGISTRY_MAX_SIZE + 1) - ] + oversized = [_tag_registry_row(f"tag-{index}") for index in range(TAG_REGISTRY_MAX_SIZE + 1)] async def fake_find_many(**kwargs): if "where" not in kwargs: @@ -2399,10 +2285,7 @@ async def test_get_tag_objects_batch_oversized_registry_falls_back_and_stops_ref user_api_key_cache=cache, ) assert list(first) == ["tag-a"] - assert ( - await cache.async_get_cache(key=tag_registry_cache_key()) - == TAG_REGISTRY_OVERFLOW_SENTINEL - ) + assert await cache.async_get_cache(key=tag_registry_cache_key()) == TAG_REGISTRY_OVERFLOW_SENTINEL second = await get_tag_objects_batch( tag_names=["tag-b"], @@ -2427,17 +2310,12 @@ async def test_tag_max_budget_check_still_enforces_registered_tag_over_budget(): async def fake_find_many(**kwargs): if "where" not in kwargs: return [_tag_registry_row("paid-tag")] - return [ - _tag_db_row(name, max_budget=1.0) - for name in kwargs["where"]["tag_name"]["in"] - ] + return [_tag_db_row(name, max_budget=1.0) for name in kwargs["where"]["tag_name"]["in"]] mock_prisma = MagicMock() mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:paid-tag": return 1.5 return fallback_spend @@ -2846,8 +2724,7 @@ def _pass_through_request() -> Request: LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) - def pass_through_endpoint(): - ... + def pass_through_endpoint(): ... setattr(pass_through_endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True) return Request(scope={"type": "http", "headers": [], "endpoint": pass_through_endpoint}) @@ -2857,8 +2734,7 @@ def _builtin_request() -> Request: """A Request dispatched to a built-in (non-pass-through) handler, e.g. what a custom path colliding with a core route actually resolves to.""" - def chat_completions(): - ... + def chat_completions(): ... return Request(scope={"type": "http", "headers": [], "endpoint": chat_completions}) @@ -3023,9 +2899,7 @@ async def test_virtual_key_soft_budget_check_without_user_obj(): ], ) @pytest.mark.asyncio -async def test_virtual_key_soft_budget_check_scenarios( - spend, soft_budget, expect_alert -): +async def test_virtual_key_soft_budget_check_scenarios(spend, soft_budget, expect_alert): """Test _virtual_key_soft_budget_check with various spend and soft_budget scenarios""" alert_triggered = False @@ -3054,9 +2928,9 @@ async def test_virtual_key_soft_budget_check_scenarios( await asyncio.sleep(0.1) - assert ( - alert_triggered == expect_alert - ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + assert alert_triggered == expect_alert, ( + f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + ) @pytest.mark.asyncio @@ -3167,9 +3041,7 @@ async def test_virtual_key_max_budget_alert_check_without_user_obj(): ], ) @pytest.mark.asyncio -async def test_virtual_key_max_budget_alert_check_scenarios( - spend, max_budget, expect_alert -): +async def test_virtual_key_max_budget_alert_check_scenarios(spend, max_budget, expect_alert): """Test _virtual_key_max_budget_alert_check with various spend and max_budget scenarios""" alert_triggered = False @@ -3198,9 +3070,9 @@ async def test_virtual_key_max_budget_alert_check_scenarios( await asyncio.sleep(0.1) - assert ( - alert_triggered == expect_alert - ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}" + assert alert_triggered == expect_alert, ( + f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}" + ) @pytest.mark.asyncio @@ -3459,9 +3331,7 @@ async def test_custom_auth_common_checks_opt_in(): "prisma_client": None, "user_api_key_cache": MagicMock(), "proxy_logging_obj": MagicMock(), - "general_settings": ( - {"custom_auth_run_common_checks": True} if flag else {} - ), + "general_settings": ({"custom_auth_run_common_checks": True} if flag else {}), "llm_router": None, "user_custom_auth": user_custom_auth, "litellm_proxy_admin_name": "admin", @@ -3533,9 +3403,7 @@ async def test_virtual_key_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:key:test-hashed-token": return 1.5 return fallback_spend @@ -3569,9 +3437,7 @@ async def test_virtual_key_budget_check_fallback_no_counter(): proxy_logging_obj.budget_alerts = AsyncMock() # get_current_spend returns fallback_spend when no counter exists - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): return fallback_spend with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): @@ -3601,9 +3467,7 @@ def _over_budget_token(**overrides) -> UserAPIKeyAuth: def _patched_spend(value: float): - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): return value return patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend) @@ -3664,9 +3528,7 @@ async def test_budget_throttle_decision_cleared_before_caching(): otherwise it would re-apply (and compound) on every subsequent request.""" from litellm.proxy.auth.auth_checks import _copy_user_api_key_auth_for_cache - valid_token = _over_budget_token( - tpm_limit=1000, rpm_limit=100, metadata={"throttle_on_budget_exceeded": True} - ) + valid_token = _over_budget_token(tpm_limit=1000, rpm_limit=100, metadata={"throttle_on_budget_exceeded": True}) valid_token.budget_throttle_pct = 0.1 cached = _copy_user_api_key_auth_for_cache(user_api_key_obj=valid_token) @@ -3762,9 +3624,7 @@ async def test_team_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team:test-team": return 1.5 return fallback_spend @@ -3791,9 +3651,7 @@ async def test_end_user_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:end_user:customer-1": return 1.5 return fallback_spend @@ -3821,9 +3679,7 @@ async def test_tag_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:paid-tag": return 1.5 return fallback_spend @@ -3873,9 +3729,7 @@ async def test_team_member_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 1.5 return fallback_spend @@ -3915,9 +3769,7 @@ class TestGuardrailModificationCheck: team_object = MagicMock() team_object.metadata = {} # no permission - return _guardrail_modification_check( - request_body=request_body, team_object=team_object - ) + return _guardrail_modification_check(request_body=request_body, team_object=team_object) def test_noop_when_no_guardrail_keys_present(self): # no-op — should return silently @@ -3965,9 +3817,7 @@ class TestGuardrailModificationCheck: return_value=False, ): with pytest.raises(HTTPException) as exc: - self._call( - {"metadata": {"opted_out_global_guardrails": ["some_guardrail"]}} - ) + self._call({"metadata": {"opted_out_global_guardrails": ["some_guardrail"]}}) assert exc.value.status_code == 403 @pytest.mark.parametrize( @@ -4101,18 +3951,12 @@ async def test_team_member_budget_check_falls_back_to_team_default_budget_id(): fake_budget_row = MagicMock() fake_budget_row.max_budget = 50.0 - fake_budget_row.dict = MagicMock( - return_value={"budget_id": "budget-default", "max_budget": 50.0} - ) + fake_budget_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": 50.0}) prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_budget_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_budget_row) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 70.0 return fallback_spend @@ -4203,15 +4047,11 @@ async def test_team_member_budget_check_per_member_override_wins_over_team_defau fake_budget_row.max_budget = 50.0 prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_budget_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_budget_row) mocked_spend = 70.0 - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return mocked_spend return fallback_spend @@ -4292,18 +4132,12 @@ async def test_team_member_budget_check_null_clone_falls_back_to_team_default(): fake_default_row = MagicMock() fake_default_row.max_budget = 65.0 - fake_default_row.dict = MagicMock( - return_value={"budget_id": "budget-default", "max_budget": 65.0} - ) + fake_default_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": 65.0}) prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_default_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_default_row) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 500.0 return fallback_spend @@ -4361,18 +4195,12 @@ async def test_team_member_budget_check_null_clone_with_null_default_skips_enfor fake_default_row = MagicMock() fake_default_row.max_budget = None - fake_default_row.dict = MagicMock( - return_value={"budget_id": "budget-default", "max_budget": None} - ) + fake_default_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": None}) prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_default_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_default_row) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 1000.0 return fallback_spend @@ -4430,18 +4258,12 @@ async def test_team_member_budget_check_zero_team_default_treated_as_no_cap(): # Team default budget row with max_budget=0.0 (the regression trigger). fake_default_row = MagicMock() fake_default_row.max_budget = 0.0 - fake_default_row.dict = MagicMock( - return_value={"budget_id": "budget-default", "max_budget": 0.0} - ) + fake_default_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": 0.0}) prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_default_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_default_row) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend @@ -4499,9 +4321,7 @@ async def test_team_member_budget_check_zero_per_member_row_still_blocks(): prisma_client = MagicMock() prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend @@ -4549,19 +4369,13 @@ def _patch_validation_helpers(monkeypatch, *, end_user=None, user=None, fuzzy=No """Stub out the DB helpers resolve_and_validate_end_user_id delegates to.""" from litellm.proxy.auth import auth_checks - monkeypatch.setattr( - auth_checks, "get_end_user_object", AsyncMock(return_value=end_user) - ) + monkeypatch.setattr(auth_checks, "get_end_user_object", AsyncMock(return_value=end_user)) monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(return_value=user)) - monkeypatch.setattr( - auth_checks, "_get_fuzzy_user_object", AsyncMock(return_value=fuzzy) - ) + monkeypatch.setattr(auth_checks, "_get_fuzzy_user_object", AsyncMock(return_value=fuzzy)) @pytest.mark.asyncio -async def test_resolve_end_user_returns_none_for_none_input( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_returns_none_for_none_input(_validate_flag_on, monkeypatch): from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id _patch_validation_helpers(monkeypatch) @@ -4596,9 +4410,7 @@ async def test_resolve_end_user_passes_through_when_flag_disabled(monkeypatch): @pytest.mark.asyncio -async def test_resolve_end_user_passes_through_when_no_prisma_client( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_passes_through_when_no_prisma_client(_validate_flag_on, monkeypatch): from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id _patch_validation_helpers(monkeypatch) @@ -4632,9 +4444,7 @@ async def test_resolve_end_user_matches_end_user_table(_validate_flag_on, monkey @pytest.mark.asyncio -async def test_resolve_end_user_matches_user_table_by_user_id( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_matches_user_table_by_user_id(_validate_flag_on, monkeypatch): from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4652,9 +4462,7 @@ async def test_resolve_end_user_matches_user_table_by_user_id( @pytest.mark.asyncio -async def test_resolve_end_user_matches_user_table_by_email( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_matches_user_table_by_email(_validate_flag_on, monkeypatch): """Email-shaped ids route through get_user_object with user_email set. The fuzzy lookup must happen inside get_user_object so it shares the @@ -4682,9 +4490,7 @@ async def test_resolve_end_user_matches_user_table_by_email( @pytest.mark.asyncio -async def test_resolve_end_user_non_email_id_does_not_pass_user_email( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_non_email_id_does_not_pass_user_email(_validate_flag_on, monkeypatch): """Non-email ids skip the email fuzzy path to avoid a pointless DB hit.""" from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4703,9 +4509,7 @@ async def test_resolve_end_user_non_email_id_does_not_pass_user_email( @pytest.mark.asyncio -async def test_resolve_end_user_drops_codex_opaque_identifier( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_drops_codex_opaque_identifier(_validate_flag_on, monkeypatch): from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id _patch_validation_helpers(monkeypatch) # all helpers return None @@ -4727,9 +4531,7 @@ async def test_resolve_end_user_drops_codex_opaque_identifier( @pytest.mark.asyncio -async def test_resolve_end_user_preserves_id_when_default_budget_configured( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_preserves_id_when_default_budget_configured(_validate_flag_on, monkeypatch): """Don't drop unregistered ids when litellm.max_end_user_budget_id is set. The default end-user budget is applied downstream when the id is present @@ -4766,9 +4568,7 @@ async def test_resolve_end_user_drops_unknown_email(_validate_flag_on, monkeypat @pytest.mark.asyncio -async def test_resolve_end_user_uses_cached_valid_result( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_uses_cached_valid_result(_validate_flag_on, monkeypatch): from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4788,9 +4588,7 @@ async def test_resolve_end_user_uses_cached_valid_result( @pytest.mark.asyncio -async def test_resolve_end_user_uses_cached_invalid_result( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_uses_cached_invalid_result(_validate_flag_on, monkeypatch): from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4809,9 +4607,7 @@ async def test_resolve_end_user_uses_cached_invalid_result( @pytest.mark.asyncio -async def test_resolve_end_user_swallows_db_errors_and_returns_none( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_swallows_db_errors_and_returns_none(_validate_flag_on, monkeypatch): from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4932,19 +4728,13 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): ) # (1) team_id-keyed write fires with the refreshed object - written_keys = [ - (c.kwargs.get("key") or c.args[0]) - for c in cache.async_set_cache.await_args_list - ] + written_keys = [(c.kwargs.get("key") or c.args[0]) for c in cache.async_set_cache.await_args_list] assert written_keys == ["team_id:team-1234"], ( "Only the team_id-keyed write should fire; the alias key must be " "deleted, NOT written. " f"Got writes: {written_keys}" ) - written_value = ( - cache.async_set_cache.await_args.kwargs.get("value") - or cache.async_set_cache.await_args.args[1] - ) + written_value = cache.async_set_cache.await_args.kwargs.get("value") or cache.async_set_cache.await_args.args[1] assert written_value is team_table # (2) team_alias-keyed entry is deleted in BOTH the in-memory cache @@ -4978,10 +4768,7 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): logging_obj2.internal_usage_cache.dual_cache.async_delete_cache.assert_awaited_once_with( key="team_id:team-no-alias" ) - written_keys_aliasless = [ - (c.kwargs.get("key") or c.args[0]) - for c in cache2.async_set_cache.await_args_list - ] + written_keys_aliasless = [(c.kwargs.get("key") or c.args[0]) for c in cache2.async_set_cache.await_args_list] assert written_keys_aliasless == ["team_id:team-no-alias"] @@ -5061,9 +4848,7 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): await _cache_team_object( team_id=team_id, - team_table=LiteLLM_TeamTableCachedObj( - team_id=team_id, models=["model-a", "model-b"] - ), + team_table=LiteLLM_TeamTableCachedObj(team_id=team_id, models=["model-a", "model-b"]), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -5173,9 +4958,7 @@ async def test_cache_team_object_tolerates_cache_invalidation_failures(): cache.async_set_cache = AsyncMock() cache.delete_cache = MagicMock(side_effect=Exception("redis down")) logging_obj = MagicMock() - logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock( - side_effect=Exception("redis down") - ) + logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock(side_effect=Exception("redis down")) await _cache_team_object( team_id="team-cache-outage", @@ -5188,10 +4971,7 @@ async def test_cache_team_object_tolerates_cache_invalidation_failures(): proxy_logging_obj=logging_obj, ) - written_keys = [ - (c.kwargs.get("key") or c.args[0]) - for c in cache.async_set_cache.await_args_list - ] + written_keys = [(c.kwargs.get("key") or c.args[0]) for c in cache.async_set_cache.await_args_list] assert written_keys == ["team_id:team-cache-outage"] @@ -5467,8 +5247,9 @@ async def test_common_checks_budget_reads_run_concurrently(): probe = _BudgetSpendConcurrencyProbe(expected=4) - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", probe + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", probe), ): task = asyncio.create_task( common_checks( @@ -5536,8 +5317,9 @@ async def test_common_checks_budget_gather_raises_highest_priority_scope(): request=MagicMock(spec=Request), ) - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), ): # Both team and end-user over budget: team wins on priority. _spend_by_counter.team = 999.0 @@ -5572,8 +5354,9 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), ): with pytest.raises(litellm.BudgetExceededError) as over: await common_checks( @@ -5615,9 +5398,11 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): async def _no_membership(*args, **kwargs): return None - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter - ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), + patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership), + ): result = await common_checks( request_body={"messages": [{"role": "user", "content": "hi"}]}, team_object=team, @@ -5656,9 +5441,11 @@ async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag async def _no_membership(*args, **kwargs): return None - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter - ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), + patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership), + ): with pytest.raises(litellm.BudgetExceededError) as exc_info: await common_checks( request_body={"messages": [{"role": "user", "content": "hi"}]}, @@ -5689,8 +5476,9 @@ async def test_common_checks_personal_user_budget_still_enforced_on_personal_key async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), ): with pytest.raises(litellm.BudgetExceededError): await common_checks( @@ -5773,10 +5561,11 @@ async def test_budget_checks_only_run_on_llm_api_routes(scope, route, expect_blo request=MagicMock(spec=Request), ) - with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter - ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership), patch( - "litellm.proxy.auth.auth_checks.get_org_object", _get_org + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), + patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership), + patch("litellm.proxy.auth.auth_checks.get_org_object", _get_org), ): if expect_blocked: with pytest.raises(litellm.BudgetExceededError): @@ -6275,9 +6064,7 @@ async def test_get_end_user_object_token_budget_gate_keeps_fetching_unrestricted mock_prisma = MagicMock() mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( - return_value=_end_user_db_row("eu-anon-1", spend=100.0) - ) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1", spend=100.0)) cache = UserApiKeyCache() result = await get_end_user_object( @@ -6959,9 +6746,7 @@ async def test_common_checks_ignores_non_llm_route_when_enabled(monkeypatch): monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) router = _router_with_priced_and_unpriced_models() - result = await _run_common_checks( - model="unpriced-group", llm_router=router, route="/model/new" - ) + result = await _run_common_checks(model="unpriced-group", llm_router=router, route="/model/new") assert result is True @@ -7077,11 +6862,15 @@ def test_team_allowed_routes_exact_route_does_not_become_a_prefix_grant(): roles = LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/model-a"]) assert ( - allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-a", litellm_proxy_roles=roles) + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-a", litellm_proxy_roles=roles + ) is True ) assert ( - allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-b", litellm_proxy_roles=roles) + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-b", litellm_proxy_roles=roles + ) is False ) @@ -7159,8 +6948,7 @@ async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_cle assert await real_cache.async_get_cache(key="team_membership:user-1:team-1") is None assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 0.0 assert ( - real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1") - == 0.0 + real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1") == 0.0 ), "the DB-floor marker kept the pre-reset value; a stale-floor read can raise the counter right back up" @@ -7356,9 +7144,9 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t ) assert remote_spend_counter_in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0 - assert ( - remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0 - ), "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor" + assert remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0, ( + "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor" + ) @pytest.mark.asyncio @@ -7541,3 +7329,136 @@ async def test_key_budget_error_keeps_the_masked_key_name(key_name): names are just as valid as the alphanumeric ones.""" message = await _run_key_budget_check(key_name) assert f"Key=prod-key ({key_name}) Current cost" in message + + +class _UntouchedPrisma: + def __getattr__(self, name: str) -> object: + raise AssertionError(f"database reached through {name}") + + +class _MissingUserPrisma: + class db: + class litellm_usertable: + @staticmethod + async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: + return None + + +@pytest.mark.asyncio +async def test_enforced_model_allowlists_treats_a_missing_user_row_as_unrestricted(): + from litellm.proxy.auth.auth_checks import enforced_model_allowlists + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + cache = UserApiKeyCache() + scopes = await enforced_model_allowlists( + valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="default_user_id"), + prisma_client=_MissingUserPrisma(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert [list(scope) for scope in scopes] == [[], [], [], [], []] + + +class _UnreachableUserPrisma: + class db: + class litellm_usertable: + @staticmethod + async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: + raise RuntimeError("database gone") + + +@pytest.mark.asyncio +async def test_enforced_model_allowlists_surfaces_a_failed_user_lookup(): + from litellm.proxy.auth.auth_checks import enforced_model_allowlists + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + cache = UserApiKeyCache() + with pytest.raises(ValueError, match="database gone"): + await enforced_model_allowlists( + valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), + prisma_client=_UnreachableUserPrisma(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + +@pytest.mark.asyncio +async def test_enforced_model_allowlists_reads_every_level_from_cache(): + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamMembership, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + ) + from litellm.proxy.auth.auth_checks import enforced_model_allowlists + from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + team_membership_reservation_cache_key, + ) + from litellm.proxy.utils import ProxyLogging + + cache = UserApiKeyCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=cache) + await cache.async_set_cache( + key="team_id:team-fake", value=LiteLLM_TeamTableCachedObj(team_id="team-fake", models=["gpt-4o"]) + ) + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="user-fake", team_id="team-fake"), + value=LiteLLM_TeamMembership( + user_id="user-fake", + team_id="team-fake", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["gpt-4o-mini"]), + ), + ) + await cache.async_set_cache( + key="project_id:project-fake", + value=LiteLLM_ProjectTableCachedObj(project_id="project-fake", models=["gpt-4.1"]), + ) + await cache.async_set_cache(key="user-fake", value=LiteLLM_UserTable(user_id="user-fake", models=["o3"])) + prisma_client = _UntouchedPrisma() + + team_scoped = await enforced_model_allowlists( + valid_token=UserAPIKeyAuth( + token="hashed-fake", + models=["all-team-models"], + team_models=["gpt-4o", "gpt-4o-mini"], + user_id="user-fake", + team_id="team-fake", + project_id="project-fake", + ), + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + ) + personal = await enforced_model_allowlists( + valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + ) + without_database = await enforced_model_allowlists( + valid_token=UserAPIKeyAuth( + token="hashed-fake", + models=["gpt-4o"], + team_models=["gpt-4o-mini"], + user_id="user-fake", + team_id="team-fake", + ), + prisma_client=None, + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert [list(scope) for scope in team_scoped] == [ + ["gpt-4o", "gpt-4o-mini"], + ["gpt-4o"], + ["gpt-4o-mini"], + [], + ["gpt-4.1"], + ] + assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []] + assert [list(scope) for scope in without_database] == [["gpt-4o"], ["gpt-4o-mini"]] diff --git a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py index b22e202d9e0..7d79192b884 100644 --- a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py @@ -1,16 +1,39 @@ -"""OpenAI passthrough must register WebSocket catch-all routes (#36088).""" +"""OpenAI passthrough WebSocket route: registration, opt-in gating, and refusals.""" -from unittest.mock import AsyncMock, MagicMock, patch +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType, SimpleNamespace +from typing import Final +from unittest.mock import patch import pytest from starlette.routing import WebSocketRoute from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _OPENAI_WS_DISABLED_REFUSAL, + _OPENAI_WS_MODEL_RESTRICTED_REFUSAL, + _has_model_restrictions, + _openai_websocket_refusal, + _proxy_model_allowlists, openai_websocket_proxy_route, router, ) +Scopes = tuple[Sequence[str], ...] + +ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True}) +DISABLED_SETTINGS: Final = ( + MappingProxyType({}), + MappingProxyType({"enable_openai_websocket_passthrough": False}), + MappingProxyType({"enable_openai_websocket_passthrough": "false"}), + MappingProxyType({"enable_openai_websocket_passthrough": None}), +) +GET_CREDENTIALS: Final = ( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials" +) + def test_openai_websocket_passthrough_routes_registered(): ws_paths = {route.path for route in router.routes if isinstance(route, WebSocketRoute)} @@ -18,164 +41,258 @@ def test_openai_websocket_passthrough_routes_registered(): assert "/openai_passthrough/{endpoint:path}" in ws_paths -def _mock_websocket(path: str, query: str, headers: dict[str, str] | None = None) -> MagicMock: - websocket = MagicMock() - websocket.url.path = path - websocket.url.query = query - websocket.headers = headers or {} - websocket.accept = AsyncMock() - websocket.close = AsyncMock() - return websocket +class _FakeWebSocket: + def __init__(self, path: str, query: str, subprotocols: str | None = None) -> None: + self.url = SimpleNamespace(path=path, query=query) + self.headers = {"sec-websocket-protocol": subprotocols} if subprotocols else {} + self.accepts: list[str | None] = [] + self.sent: list[str] = [] + self.closed: tuple[int, str] | None = None + + async def accept(self, subprotocol: str | None = None) -> None: + self.accepts.append(subprotocol) + + async def send_text(self, data: str) -> None: + self.sent.append(data) + + async def close(self, code: int = 1000, reason: str = "") -> None: + self.closed = (code, reason) + + def error_message(self) -> str: + assert len(self.sent) == 1 + frame = json.loads(self.sent[0]) + assert frame["type"] == "error" + return frame["error"]["message"] + + +@dataclass(frozen=True, slots=True) +class _RelayCall: + target: str + custom_headers: Mapping[str, str] + forward_headers: bool + endpoint: str + accept_websocket: bool + + +class _FakeRelay: + def __init__(self) -> None: + self.calls: list[_RelayCall] = [] + + async def __call__( + self, + *, + websocket: _FakeWebSocket, + target: str, + custom_headers: dict[str, str], + user_api_key_dict: UserAPIKeyAuth, + forward_headers: bool, + endpoint: str, + accept_websocket: bool, + ) -> None: + self.calls.append( + _RelayCall( + target=target, + custom_headers=MappingProxyType(dict(custom_headers)), + forward_headers=forward_headers, + endpoint=endpoint, + accept_websocket=accept_websocket, + ) + ) + + +class _FakeModelAllowlists: + def __init__(self, scopes: Scopes) -> None: + self.scopes = scopes + self.calls: list[UserAPIKeyAuth] = [] + + async def __call__(self, valid_token: UserAPIKeyAuth, /) -> Scopes: + self.calls.append(valid_token) + return self.scopes + + +@dataclass(frozen=True, slots=True) +class _Served: + relay: _FakeRelay + allowlists: _FakeModelAllowlists + + +async def _serve( + websocket: _FakeWebSocket, + endpoint: str, + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], + scopes: Scopes = (), +) -> _Served: + served = _Served(relay=_FakeRelay(), allowlists=_FakeModelAllowlists(scopes)) + await openai_websocket_proxy_route( + websocket=websocket, + endpoint=endpoint, + user_api_key_dict=user_api_key_dict, + general_settings=general_settings, + relay=served.relay, + model_allowlists=served.allowlists, + ) + return served @pytest.mark.asyncio @pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) -async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix): - websocket = _mock_websocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") +async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, monkeypatch): + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value="sk-provider", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._join_url_paths", - return_value="https://api.openai.com/v1/realtime", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws, - ): - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/realtime", - user_api_key_dict=UserAPIKeyAuth(), + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) + + assert served.relay.calls == [ + _RelayCall( + target="wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview", + custom_headers=MappingProxyType({"Authorization": "Bearer sk-provider"}), + forward_headers=False, + endpoint=f"/{prefix}/v1/realtime", + accept_websocket=False, ) - - kwargs = mock_ws.await_args.kwargs - assert kwargs["target"] == "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview" - assert kwargs["custom_headers"] == {"Authorization": "Bearer sk-provider"} - assert kwargs["forward_headers"] is False - assert kwargs["endpoint"] == f"/{prefix}/v1/realtime" - assert kwargs["accept_websocket"] is False - websocket.accept.assert_awaited_once_with(subprotocol=None) - websocket.close.assert_not_awaited() + ] + assert websocket.accepts == [None] + assert websocket.sent == [] + assert websocket.closed is None @pytest.mark.asyncio async def test_openai_websocket_accepts_first_client_subprotocol(): - websocket = _mock_websocket( + websocket = _FakeWebSocket( "/openai/v1/realtime", "model=gpt-4o-realtime-preview", - headers={ - "sec-websocket-protocol": "realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1" - }, + subprotocols="realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1", ) - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value="sk-provider", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws, - ): - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/realtime", - user_api_key_dict=UserAPIKeyAuth(), - ) + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) - websocket.accept.assert_awaited_once_with(subprotocol="realtime") - assert mock_ws.await_args.kwargs["accept_websocket"] is False - websocket.close.assert_not_awaited() + assert websocket.accepts == ["realtime"] + assert [call.accept_websocket for call in served.relay.calls] == [False] + assert websocket.closed is None @pytest.mark.asyncio async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing(): - websocket = _mock_websocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") + websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value=None, - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws, - ): - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/realtime", - user_api_key_dict=UserAPIKeyAuth(), - ) + with patch(GET_CREDENTIALS, return_value=None): + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) - websocket.close.assert_awaited_once() - assert websocket.close.await_args.kwargs["code"] == 1011 - websocket.accept.assert_not_awaited() - mock_ws.assert_not_awaited() + assert websocket.closed is not None + assert websocket.closed[0] == 1011 + assert "OPENAI_API_KEY" in websocket.closed[1] + assert websocket.accepts == [] + assert served.relay.calls == [] @pytest.mark.asyncio -@pytest.mark.parametrize( - "user_api_key_dict", - [ - UserAPIKeyAuth(models=["gpt-4o"]), - UserAPIKeyAuth(team_models=["gpt-4o-realtime-preview"]), - UserAPIKeyAuth(models=["all-team-models"], team_models=["gpt-4o"]), - ], -) -async def test_openai_websocket_rejects_model_restricted_keys(user_api_key_dict): - websocket = _mock_websocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") +@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) +@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) +async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings): + websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws: - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/realtime", - user_api_key_dict=user_api_key_dict, - ) + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings) - websocket.close.assert_awaited_once() - assert websocket.close.await_args.kwargs["code"] == 1008 - websocket.accept.assert_not_awaited() - mock_ws.assert_not_awaited() + assert "enable_openai_websocket_passthrough" in websocket.error_message() + assert websocket.accepts == [None] + assert websocket.closed == (1008, _OPENAI_WS_DISABLED_REFUSAL.close_reason) + assert served.relay.calls == [] @pytest.mark.asyncio -@pytest.mark.parametrize( - "user_api_key_dict", - [ - UserAPIKeyAuth(), - UserAPIKeyAuth(models=["all-proxy-models"]), - UserAPIKeyAuth(models=["*"]), - UserAPIKeyAuth(models=["all-team-models"], team_models=["all-proxy-models"]), - ], +@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) +async def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): + refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), general_settings, _FakeModelAllowlists(())) + assert refusal is _OPENAI_WS_DISABLED_REFUSAL + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [True, "true", "True"]) +async def test_openai_websocket_refusal_is_none_for_truthy_settings(value): + settings = MappingProxyType({"enable_openai_websocket_passthrough": value}) + assert await _openai_websocket_refusal(UserAPIKeyAuth(), settings, _FakeModelAllowlists(())) is None + + +@pytest.mark.asyncio +async def test_openai_websocket_refusal_echoes_requested_subprotocol(): + websocket = _FakeWebSocket( + "/openai_passthrough/v1/realtime", + "model=gpt-4o-realtime-preview", + subprotocols="realtime, openai-beta.realtime-v1", + ) + + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), MappingProxyType({})) + + assert websocket.accepts == ["realtime"] + assert websocket.closed == (1008, _OPENAI_WS_DISABLED_REFUSAL.close_reason) + assert served.relay.calls == [] + + +RESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( + (("gpt-4o",),), + ((), ("gpt-4o-realtime-preview",)), + (("all-team-models",), ("gpt-4o",)), + ((), ("all-proxy-models",), ("gpt-4o",)), + ((), (), (), ("gpt-4o",)), + (("*",), (), (), (), ("gpt-4o",)), +) +UNRESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( + (), + ((),), + (("all-proxy-models",),), + (("*",),), + (("all-team-models",), ("all-proxy-models",)), + ((), (), (), (), ()), + (("*",), ("all-proxy-models",), ("all-team-models",), (), ()), ) -async def test_openai_websocket_allows_unrestricted_keys(user_api_key_dict): - websocket = _mock_websocket("/openai/v1/responses", "") - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value="sk-provider", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws, - ): - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/responses", - user_api_key_dict=user_api_key_dict, - ) - mock_ws.assert_awaited_once() - websocket.close.assert_not_awaited() +@pytest.mark.asyncio +@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) +async def test_openai_websocket_rejects_model_restricted_identities(scopes): + websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") + user_api_key_dict = UserAPIKeyAuth(token="hashed-fake", user_id="user-fake", team_id="team-fake") + + served = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED, scopes) + + assert "model restrictions" in websocket.error_message() + assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason) + assert served.relay.calls == [] + assert served.allowlists.calls == [user_api_key_dict] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) +async def test_openai_websocket_disabled_refusal_skips_allowlist_lookups(scopes): + allowlists = _FakeModelAllowlists(scopes) + + refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), MappingProxyType({}), allowlists) + + assert refusal is _OPENAI_WS_DISABLED_REFUSAL + assert allowlists.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scopes", UNRESTRICTED_SCOPES) +async def test_openai_websocket_allows_unrestricted_identities(scopes): + websocket = _FakeWebSocket("/openai/v1/responses", "") + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED, scopes) + + assert len(served.relay.calls) == 1 + assert websocket.sent == [] + assert websocket.closed is None + + +@pytest.mark.asyncio +async def test_proxy_model_allowlists_reads_the_token_scopes_without_a_database(): + token: Final = UserAPIKeyAuth(models=[], team_id="team-fake", team_models=["gpt-4o"]) + with patch("litellm.proxy.proxy_server.prisma_client", None): + scopes = await _proxy_model_allowlists()(token) + + assert tuple(tuple(scope) for scope in scopes) == ((), ("gpt-4o",)) + assert _has_model_restrictions(scopes) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d91928a203e..34216f7e1b9 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11568,14 +11568,10 @@ async def test_key_window_spend_row_is_enqueued_with_the_actual_cost(): reset_at = datetime.now(timezone.utc) + timedelta(days=10) key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] + key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}] with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) enqueued = await _drain(queue) assert len(enqueued) == 1 @@ -11594,14 +11590,10 @@ async def test_team_window_spend_row_is_enqueued(): reset_at = datetime.now(timezone.utc) + timedelta(days=3) team_obj = MagicMock() - team_obj.budget_limits = [ - {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} - ] + team_obj.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}] with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: - await increment_spend_counters( - token=None, team_id="team-1", user_id=None, response_cost=1.5 - ) + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=1.5) enqueued = await _drain(queue) assert len(enqueued) == 1 @@ -11620,9 +11612,7 @@ async def test_window_spend_row_is_enqueued_even_when_the_counter_was_reserved() reset_at = datetime.now(timezone.utc) + timedelta(days=10) key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] + key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}] reservation = { "entries": [ {"counter_key": "spend:key:hashed-token", "reserved": 1.0}, @@ -11660,9 +11650,7 @@ async def test_sliding_window_without_reset_at_is_not_enqueued(): key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0}] with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) enqueued = await _drain(queue) assert enqueued == [] @@ -11680,9 +11668,7 @@ async def test_each_configured_window_gets_its_own_row_enqueue(): ] with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) enqueued = await _drain(queue) assert sorted(item["window_duration"] for item in enqueued) == ["1d", "30d"] @@ -11697,9 +11683,7 @@ async def test_no_window_spend_row_enqueued_without_budget_limits(): key_obj.budget_limits = None with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) enqueued = await _drain(queue) assert enqueued == [] @@ -11713,9 +11697,7 @@ async def test_window_spend_row_carries_the_request_start_time(): reset_at = datetime.now(timezone.utc) + timedelta(days=10) key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] + key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}] with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: await increment_spend_counters( @@ -11736,9 +11718,7 @@ async def test_team_window_spend_row_carries_the_request_start_time(): reset_at = datetime.now(timezone.utc) + timedelta(days=3) team_obj = MagicMock() - team_obj.budget_limits = [ - {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} - ] + team_obj.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}] with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: await increment_spend_counters( @@ -12050,7 +12030,6 @@ async def test_init_guardrails_in_db_snapshots_and_reconciles_under_guardrail_re assert not GUARDRAIL_RECONCILE_LOCK.locked() - @pytest.mark.asyncio async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeypatch): from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY @@ -12094,7 +12073,9 @@ async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeyp await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) assert served_content() == "Begin every reply with AHOY" - prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row("Begin every reply with HOWDY")]) + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[db_row("Begin every reply with HOWDY")] + ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) assert served_content() == "Begin every reply with HOWDY" @@ -12547,3 +12528,40 @@ def test_disabling_docs_does_not_disable_other_routes(monkeypatch): assert client.get("/redoc").status_code == 404 assert client.get("/health/liveliness").status_code == 200 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_general_settings, expected", + [ + ({"enable_openai_websocket_passthrough": True}, True), + ({"enable_openai_websocket_passthrough": False}, False), + ({}, None), + ], +) +async def test_update_general_settings_propagates_openai_websocket_passthrough(db_general_settings, expected): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch("litellm.proxy.proxy_server.general_settings", {"enable_openai_websocket_passthrough": True}): + await proxy_config._update_general_settings(db_general_settings=db_general_settings) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["enable_openai_websocket_passthrough"] is expected + + +@pytest.mark.asyncio +async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough(): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_general_settings_keys = {"enable_openai_websocket_passthrough"} + + with patch("litellm.proxy.proxy_server.general_settings", {"enable_openai_websocket_passthrough": False}): + await proxy_config._update_general_settings(db_general_settings={"enable_openai_websocket_passthrough": True}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["enable_openai_websocket_passthrough"] is False diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0f9d937c3b2..94dc03123c0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25776,6 +25776,11 @@ export interface components { * @description If True and SSO is configured (MICROSOFT_CLIENT_ID, GOOGLE_CLIENT_ID, GENERIC_CLIENT_ID, or SAML_IDP_METADATA_URL/XML), disables username/password login on /login, /v2/login, and /v3/login so SSO is the only way to reach the Admin UI. An admin locked out of the UI can still administer the proxy over the API with the master key; unset this setting and restart the proxy to restore UI username/password login. Default is False. */ disable_password_login_when_sso_enabled?: boolean | null; + /** + * Enable Openai Websocket Passthrough + * @description Serve the OpenAI pass-through WebSocket route, which relays frames to OpenAI under the proxy's own provider credential without reading them. Off by default. + */ + enable_openai_websocket_passthrough?: boolean | null; /** * Enable Public Model Hub * @description Public model hub for users to see what models they have access to, supported openai params, etc.