fix(proxy): check every model the realtime URL names before the OpenAI websocket relay connects

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -142,6 +142,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
     YieldsToConfiguredPassThroughRoute,
     create_pass_through_route,
     create_websocket_passthrough_route,
+    refuse_websocket,
     websocket_passthrough_request,
 )
 from litellm.proxy.utils import ProxyLogging as ProxyLoggingType
@@ -3560,6 +3561,14 @@ def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None:
     return requested_subprotocols[0] if requested_subprotocols else None
 
 
+async def _opening_model_refusal(
+    websocket: WebSocket, query_string: str, frame_model_gate: WebsocketFrameModelGate
+) -> WebsocketFrameModelRefused | None:
+    opening_models: Final = dict.fromkeys(httpx.QueryParams(query_string).get_list("model"))
+    verdicts: Final = tuple([await frame_model_gate(websocket, model) for model in opening_models])
+    return next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None)
+
+
 @router.websocket("/openai_passthrough/{endpoint:path}")
 @router.websocket("/openai/{endpoint:path}")
 async def openai_websocket_proxy_route(
@@ -3619,6 +3628,11 @@ async def openai_websocket_proxy_route(
 
     await websocket.accept(subprotocol=negotiated_subprotocol)
 
+    opening_refusal: Final = await _opening_model_refusal(websocket, query_string, frame_model_gate)
+    if opening_refusal is not None:
+        await refuse_websocket(websocket, opening_refusal)
+        return
+
     await relay(
         websocket=websocket,
         target=wss_target,

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -2788,11 +2788,15 @@ def _refusal_close_reason(message: str) -> str:
     return _truncated_close_reason(message if first_sentence is None else first_sentence.group(0))
 
 
+async def refuse_websocket(websocket: WebSocket, refusal: WebsocketFrameModelRefused) -> None:
+    await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type))
+    await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message))
+
+
 async def _refuse_client_frame(
     websocket: WebSocket, upstream_ws: ClientConnection, refusal: WebsocketFrameModelRefused
 ) -> None:
-    await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type))
-    await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message))
+    await refuse_websocket(websocket, refusal)
     await upstream_ws.close()
 
 

tests/unit/proxy/test_openai_ws_passthrough_routes.py

@@ -121,8 +121,13 @@ class _FakeRelay:
 
 
 class _FakeFrameModelGate:
+    def __init__(self, refusals: Mapping[str, WebsocketFrameModelRefused] = MappingProxyType({})) -> None:
+        self.refusals = refusals
+        self.checked: list[str] = []
+
     async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict:
-        return WebsocketFrameModelAllowed(access_groups=())
+        self.checked.append(model)
+        return self.refusals.get(model) or WebsocketFrameModelAllowed(access_groups=())
 
 
 @dataclass(frozen=True, slots=True)
@@ -136,8 +141,9 @@ async def _serve(
     endpoint: str,
     user_api_key_dict: UserAPIKeyAuth,
     general_settings: Mapping[str, object],
+    frame_model_gate: _FakeFrameModelGate | None = None,
 ) -> _Served:
-    served: Final = _Served(relay=_FakeRelay(), frame_model_gate=_FakeFrameModelGate())
+    served: Final = _Served(relay=_FakeRelay(), frame_model_gate=frame_model_gate or _FakeFrameModelGate())
     await openai_websocket_proxy_route(
         websocket=websocket,
         endpoint=endpoint,
@@ -332,6 +338,55 @@ async def test_openai_websocket_connects_model_restricted_keys(key_models):
     assert websocket.closed is None
 
 
+ALIAS_REFUSAL: Final = WebsocketFrameModelRefused(
+    message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "query",
+    [
+        "model=gpt-realtime-2.1",
+        "model=gpt-realtime-2.1&model=gpt-realtime-2.1-mini",
+        "model=gpt-realtime-2.1-mini&model=gpt-realtime-2.1",
+    ],
+    ids=["one_model", "refused_model_first", "refused_model_last"],
+)
+async def test_openai_websocket_refuses_an_opening_model_the_frame_gate_refuses(query):
+    """
+    The relay sends the connection URL to OpenAI unchanged, so every model it names has to pass the
+    same gate as a frame, or an alias the upgrade resolves to an allowed model opens the session on the
+    alias name itself
+    """
+    websocket: Final = _FakeWebSocket("/openai/v1/realtime", query)
+    gate: Final = _FakeFrameModelGate(MappingProxyType({"gpt-realtime-2.1": ALIAS_REFUSAL}))
+
+    with patch(GET_CREDENTIALS, return_value="sk-provider"):
+        served: Final = await _serve(
+            websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED, gate
+        )
+
+    assert websocket.error_message() == _OPENAI_WS_ALIASED_MODEL_REFUSAL
+    assert websocket.closed == (1008, "This key cannot use a model alias over OpenAI websocket passthrough.")
+    assert served.relay.calls == []
+    assert "gpt-realtime-2.1" in gate.checked
+
+
+@pytest.mark.asyncio
+async def test_openai_websocket_relays_a_session_once_every_opening_model_passes_the_frame_gate():
+    websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1-mini")
+
+    with patch(GET_CREDENTIALS, return_value="sk-provider"):
+        served: Final = await _serve(
+            websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED
+        )
+
+    assert served.frame_model_gate.checked == ["gpt-realtime-2.1-mini"]
+    assert len(served.relay.calls) == 1
+    assert websocket.closed is None
+
+
 class _FrameAuthorizer:
     """Stands in for the full key, team, user and budget auth the gate re-runs for every model a frame names."""