fix(proxy): gate the OpenAI websocket passthrough behind an explicit opt-in

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.",

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -14,8 +14,9 @@ import json
 import os
 import re
 from collections.abc import AsyncGenerator, Callable, Mapping
+from dataclasses import dataclass
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Annotated, Final, cast
+from typing import TYPE_CHECKING, Annotated, Final, Protocol, cast
 
 import httpx
 from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket
@@ -2345,19 +2346,99 @@ def _key_has_model_restrictions(user_api_key_dict: UserAPIKeyAuth) -> bool:
     return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for model in scoped_models)
 
 
+@dataclass(frozen=True, slots=True)
+class _OpenAIWebsocketRefusal:
+    close_reason: str
+    message: str
+
+
+_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
+
+
+def _openai_websocket_refusal(
+    user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object]
+) -> _OpenAIWebsocketRefusal | None:
+    if not _is_openai_websocket_passthrough_enabled(general_settings):
+        return _OPENAI_WS_DISABLED_REFUSAL
+    if _key_has_model_restrictions(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
+
+
 @router.websocket("/openai_passthrough/{endpoint:path}")
 @router.websocket("/openai/{endpoint:path}")
 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)],
 ) -> 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 = _openai_websocket_refusal(user_api_key_dict, general_settings)
+    if refusal is not None:
+        await websocket.accept(subprotocol=negotiated_subprotocol)
+        await websocket.send_text(
+            json.dumps(
+                {
+                    "type": "error",
+                    "error": {"type": "invalid_request_error", "message": refusal.message},
+                }
+            )
         )
+        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 +2474,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,

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
   }
 }

tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py

@@ -1,16 +1,35 @@
-"""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
+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,
+    _openai_websocket_refusal,
     openai_websocket_proxy_route,
     router,
 )
 
+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 +37,213 @@ 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,
+            )
+        )
+
+
+async def _serve(
+    websocket: _FakeWebSocket,
+    endpoint: str,
+    user_api_key_dict: UserAPIKeyAuth,
+    general_settings: Mapping[str, object],
+) -> _FakeRelay:
+    relay = _FakeRelay()
+    await openai_websocket_proxy_route(
+        websocket=websocket,
+        endpoint=endpoint,
+        user_api_key_dict=user_api_key_dict,
+        general_settings=general_settings,
+        relay=relay,
+    )
+    return relay
 
 
 @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"):
+        relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED)
+
+    assert 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"):
+        relay = 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 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):
+        relay = 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 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"]),
-    ],
+@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")
+
+    relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings)
+
+    assert "enable_openai_websocket_passthrough" in websocket.error_message()
+    assert websocket.accepts == [None]
+    assert websocket.closed == (1008, _OPENAI_WS_DISABLED_REFUSAL.close_reason)
+    assert relay.calls == []
+
+
+@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS)
+def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings):
+    assert _openai_websocket_refusal(UserAPIKeyAuth(), general_settings) is _OPENAI_WS_DISABLED_REFUSAL
+
+
+@pytest.mark.parametrize("value", [True, "true", "True"])
+def test_openai_websocket_refusal_is_none_for_truthy_settings(value):
+    settings = MappingProxyType({"enable_openai_websocket_passthrough": value})
+    assert _openai_websocket_refusal(UserAPIKeyAuth(), settings) 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",
+    )
+
+    relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), MappingProxyType({}))
+
+    assert websocket.accepts == ["realtime"]
+    assert websocket.closed == (1008, _OPENAI_WS_DISABLED_REFUSAL.close_reason)
+    assert relay.calls == []
+
+
+RESTRICTED_KEYS: Final = (
+    UserAPIKeyAuth(models=["gpt-4o"]),
+    UserAPIKeyAuth(team_models=["gpt-4o-realtime-preview"]),
+    UserAPIKeyAuth(models=["all-team-models"], team_models=["gpt-4o"]),
 )
+UNRESTRICTED_KEYS: Final = (
+    UserAPIKeyAuth(),
+    UserAPIKeyAuth(models=["all-proxy-models"]),
+    UserAPIKeyAuth(models=["*"]),
+    UserAPIKeyAuth(models=["all-team-models"], team_models=["all-proxy-models"]),
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("user_api_key_dict", RESTRICTED_KEYS)
 async def test_openai_websocket_rejects_model_restricted_keys(user_api_key_dict):
-    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.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,
-        )
+    relay = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED)
 
-    websocket.close.assert_awaited_once()
-    assert websocket.close.await_args.kwargs["code"] == 1008
-    websocket.accept.assert_not_awaited()
-    mock_ws.assert_not_awaited()
+    assert "model restrictions" in websocket.error_message()
+    assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason)
+    assert relay.calls == []
+
+
+@pytest.mark.parametrize("user_api_key_dict", RESTRICTED_KEYS)
+def test_openai_websocket_refusal_prefers_disabled_over_model_restriction(user_api_key_dict):
+    assert _openai_websocket_refusal(user_api_key_dict, MappingProxyType({})) is _OPENAI_WS_DISABLED_REFUSAL
 
 
 @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("user_api_key_dict", UNRESTRICTED_KEYS)
 async def test_openai_websocket_allows_unrestricted_keys(user_api_key_dict):
-    websocket = _mock_websocket("/openai/v1/responses", "")
+    websocket = _FakeWebSocket("/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,
-        )
+    with patch(GET_CREDENTIALS, return_value="sk-provider"):
+        relay = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED)
 
-    mock_ws.assert_awaited_once()
-    websocket.close.assert_not_awaited()
+    assert len(relay.calls) == 1
+    assert websocket.sent == []
+    assert websocket.closed is None

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -25731,6 +25731,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.