diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index fd50f2ae13f..0cd9f5caf53 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/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, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 0a591f9c2f8..2ccc4fe56d2 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/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() diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index 587faa02e35..839ab49408f 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/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."""