diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4c7613e1edd..dd6633ea371 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4398,6 +4398,36 @@ async def can_key_call_resolved_model( ) +async def can_personal_user_call_model( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, +) -> None: + """Apply the personal user's model allowlist, the check ``common_checks`` runs for a key with no team.""" + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if valid_token.team_id is not None or valid_token.user_id is None or prisma_client is None: + return + try: + user_object: Final = 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, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except UserNotFoundError: + return + if user_object is None: + return + await can_user_call_model(model=model, llm_router=llm_router, user_object=user_object) + + def can_org_access_model( model: str, org_object: LiteLLM_OrganizationTable | None, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 3c0033659ef..4dc613185ea 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -37,7 +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 can_key_call_resolved_model +from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, can_personal_user_call_model 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 ( @@ -2406,6 +2406,7 @@ def _proxy_frame_model_gate() -> WebsocketFrameModelGate: await can_key_call_resolved_model( model=model, llm_model_list=llm_model_list, valid_token=valid_token, llm_router=llm_router ) + await can_personal_user_call_model(model=model, valid_token=valid_token, llm_router=llm_router) except ProxyException as denial: return denial.message return None diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py index eec6ae102d2..c7df3259ce2 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -41,6 +41,7 @@ class _ModelNamingEvent(BaseModel): class _CompletedResponse(BaseModel): model_config = ConfigDict(extra="ignore") + model: str | None = None usage: ResponseAPIUsage @@ -61,6 +62,12 @@ class _Billing: cost: float +@dataclass(frozen=True, slots=True) +class _CompletedTurn: + model: str + usage: Usage + + def _event_model(event: Mapping[str, object]) -> str | None: if event.get("type") not in _MODEL_NAMING_EVENTS: return None @@ -78,9 +85,7 @@ def _event_model(event: Mapping[str, object]) -> str | None: ) -def observed_websocket_model(client_model: object, messages: WebsocketMessages) -> str | None: - if isinstance(client_model, str) and client_model: - return client_model +def observed_websocket_model(messages: WebsocketMessages) -> str | None: return next((model for model in map(_event_model, messages) if model is not None), None) @@ -101,7 +106,7 @@ def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: Lite return _Billing(usage=usage, cost=cost) -def _responses_usage(event: Mapping[str, object]) -> Usage | None: +def _responses_turn(event: Mapping[str, object], session_model: str) -> _CompletedTurn | None: if event.get("type") != _RESPONSES_COMPLETED: return None try: @@ -111,19 +116,22 @@ def _responses_usage(event: Mapping[str, object]) -> Usage | None: "OpenAI websocket passthrough: a response.completed frame carried no readable usage, so it is not billed" ) return None - return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter - completed.response.usage + return _CompletedTurn( + model=completed.response.model or session_model, + usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter + completed.response.usage + ), ) def _responses_billing(messages: WebsocketMessages, model: str) -> _Billing | None: - turns: Final = tuple(usage for usage in map(_responses_usage, messages) if usage is not None) + turns: Final = tuple(turn for turn in (_responses_turn(event, model) for event in messages) if turn is not None) if not turns: return None usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects( - [*turns] # mutable-ok: combine_usage_objects takes a concrete list + [turn.usage for turn in turns] # mutable-ok: combine_usage_objects takes a concrete list ) - return _Billing(usage=usage, cost=_responses_cost(usage, model)) + return _Billing(usage=usage, cost=sum(_responses_cost(turn.usage, turn.model) for turn in turns)) def _responses_cost(usage: Usage, model: str) -> float: @@ -149,7 +157,7 @@ class OpenAIWebsocketPassthroughLoggingHandler: start_time: datetime, kwargs: Mapping[str, object], ) -> PassThroughEndpointLoggingTypedDict: - model: Final = observed_websocket_model(kwargs.get("model"), websocket_messages) + model: Final = observed_websocket_model(websocket_messages) if model is None: verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no model named, skipping cost", url_route) unpriced: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {**kwargs}} diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index d041f7e16ab..a1863b3002a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2137,7 +2137,7 @@ class WebsocketFrameModelGate(Protocol): @dataclass(frozen=True, slots=True) class _ForwardClientFrame: - model: str | None + pass @dataclass(frozen=True, slots=True) @@ -2145,7 +2145,8 @@ class _RefuseClientFrame: message: str -_MODEL_FRAME_MARKER: Final = '"model"' +_PLAIN_MODEL_KEY: Final = '"model"' +_JSON_STRING_ESCAPE: Final = "\\" _CLIENT_FRAME_MODEL_PLACEMENTS: Final = ( ("model",), ("response", "model"), @@ -2156,13 +2157,29 @@ _CLIENT_FRAME_MODEL_PLACEMENTS: Final = ( _JSON_OBJECT_FRAME: Final = TypeAdapter(dict[str, object]) -def _json_object_frame(text_data: str) -> Mapping[str, object] | None: +def _json_object_frame(frame_data: str | bytes) -> Mapping[str, object] | None: try: - return _JSON_OBJECT_FRAME.validate_json(text_data) + return _JSON_OBJECT_FRAME.validate_json(frame_data) except ValidationError: return None +def _may_name_a_model(text_data: str) -> bool: + return _PLAIN_MODEL_KEY in text_data or _JSON_STRING_ESCAPE in text_data + + +def _recorded_frames(upstream_message: str | bytes) -> tuple[Mapping[str, object], ...]: + frame: Final = _json_object_frame(upstream_message) + return () if frame is None else (frame,) + + +async def _relay_upstream_frame(websocket: WebSocket, upstream_message: str | bytes) -> None: + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + else: + await websocket.send_text(upstream_message) + + def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None: if frame is None: return None @@ -2181,10 +2198,10 @@ def _placed_model(frame: Mapping[str, object], placement: Sequence[str]) -> str return value if isinstance(value, str) else None -def _client_frame_models(text_data: str) -> tuple[str, ...]: - if _MODEL_FRAME_MARKER not in text_data: +def _client_frame_models(frame_data: str | bytes) -> tuple[str, ...]: + if isinstance(frame_data, str) and not _may_name_a_model(frame_data): return () - frame: Final = _json_object_frame(text_data) + frame: Final = _json_object_frame(frame_data) if frame is None: return () placed: Final = (_placed_model(frame, placement) for placement in _CLIENT_FRAME_MODEL_PLACEMENTS) @@ -2192,16 +2209,15 @@ def _client_frame_models(text_data: str) -> tuple[str, ...]: async def _client_frame_verdict( - text_data: str, gate: WebsocketFrameModelGate | None, valid_token: UserAPIKeyAuth + frame_data: str | bytes, gate: WebsocketFrameModelGate | None, valid_token: UserAPIKeyAuth ) -> _ForwardClientFrame | _RefuseClientFrame: if gate is None: - return _ForwardClientFrame(model=None) - models: Final = _client_frame_models(text_data) - for model in models: + return _ForwardClientFrame() + for model in _client_frame_models(frame_data): refusal = await gate(model, valid_token) if refusal is not None: return _RefuseClientFrame(message=refusal) - return _ForwardClientFrame(model=models[0] if models else None) + return _ForwardClientFrame() async def _refuse_client_frame(websocket: WebSocket, upstream_ws: ClientConnection, message: str) -> None: @@ -2210,6 +2226,24 @@ async def _refuse_client_frame(websocket: WebSocket, upstream_ws: ClientConnecti await upstream_ws.close() +async def _relay_client_frame( + websocket: WebSocket, + upstream_ws: ClientConnection, + checked_frame: str | bytes, + forwarded_frame: str | bytes, + gate: WebsocketFrameModelGate | None, + valid_token: UserAPIKeyAuth, +) -> bool: + """Forward one client frame upstream, or refuse it and report that the session is over.""" + match await _client_frame_verdict(checked_frame, gate, valid_token): + case _RefuseClientFrame(message=refusal_message): + await _refuse_client_frame(websocket, upstream_ws, refusal_message) + return False + case _ForwardClientFrame(): + await upstream_ws.send(forwarded_frame) + return True + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2243,7 +2277,7 @@ async def websocket_passthrough_request( # Initialize tracking variables start_time: Final = datetime.now() - websocket_messages: Final[list[dict[str, object]]] = [] + websocket_messages: Final[list[Mapping[str, object]]] = [] litellm_call_id: Final = str(uuid.uuid4()) verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target) @@ -2431,20 +2465,25 @@ async def websocket_passthrough_request( ) # Not a JSON message or doesn't contain setup data - match await _client_frame_verdict(text_data, client_frame_model_gate, user_api_key_dict): - case _RefuseClientFrame(message=refusal_message): - await _refuse_client_frame(websocket, upstream_ws, refusal_message) - break - case _ForwardClientFrame(model=frame_model): - if frame_model is not None: - kwargs["model"] = frame_model - logging_obj.model = frame_model - logging_obj.model_call_details["model"] = frame_model - await upstream_ws.send( - _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter) - ) + if not await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_model_gate, + user_api_key_dict, + ): + break elif bytes_data is not None: - await upstream_ws.send(bytes_data) + if not await _relay_client_frame( + websocket, + upstream_ws, + bytes_data, + bytes_data, + client_frame_model_gate, + user_api_key_dict, + ): + break except asyncio.CancelledError: raise except Exception: @@ -2457,11 +2496,8 @@ async def websocket_passthrough_request( """Forward messages from upstream to client WebSocket, returning the upstream's close frame""" try: # Wait for the first response from upstream - raw_response = await upstream_ws.recv(decode=False) - # Ensure raw_response is bytes before decoding - if isinstance(raw_response, str): - raw_response = raw_response.encode("utf-8") - setup_response: Final = _JSON_OBJECT_FRAME.validate_json(raw_response) + raw_response: Final = await upstream_ws.recv() + setup_response: Final = _json_object_frame(raw_response) verbose_proxy_logger.debug("Setup response: %s", setup_response) # Extract model and provider from setup response for Vertex AI Live @@ -2470,7 +2506,9 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): Processing server setup response for model extraction", endpoint, ) - extracted_model: Final = _extract_model_from_vertex_ai_setup(setup_response) + extracted_model: Final = ( + None if setup_response is None else _extract_model_from_vertex_ai_setup(setup_response) + ) if extracted_model: kwargs["model"] = extracted_model kwargs["custom_llm_provider"] = "vertex_ai_language_models" @@ -2496,27 +2534,13 @@ async def websocket_passthrough_request( ) # Send the setup response to the client - await websocket.send_text(json.dumps(setup_response)) - websocket_messages.append(setup_response) + websocket_messages.extend(_recorded_frames(raw_response)) + await _relay_upstream_frame(websocket, raw_response) # Now continuously forward messages from upstream to client async for upstream_message in upstream_ws: - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - # Parse and collect for cost tracking - try: - message_data: dict[str, object] = json.loads(upstream_message.decode()) - websocket_messages.append(message_data) - except (json.JSONDecodeError, UnicodeDecodeError): - pass - else: - await websocket.send_text(upstream_message) - # Parse and collect for cost tracking - try: - message_data = json.loads(upstream_message) - websocket_messages.append(message_data) - except json.JSONDecodeError: - pass + websocket_messages.extend(_recorded_frames(upstream_message)) + await _relay_upstream_frame(websocket, upstream_message) except (ConnectionClosedOK, ConnectionClosedError) as e: verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py index 81f8436335d..ca0268950f7 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -12,6 +12,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket REALTIME_MODEL: Final = "gpt-realtime-2.1" RESPONSES_MODEL: Final = "gpt-5.4-mini" +REASONING_MODEL: Final = "gpt-5.4" def _logging_obj() -> MagicMock: @@ -142,15 +143,66 @@ def test_responses_cost_covers_every_turn_of_the_connection(): assert handled["result"].usage.total_tokens == 45 -def test_the_model_a_client_frame_named_prices_a_session_the_provider_never_named(): +def test_a_model_the_caller_asked_for_never_prices_a_session_the_provider_never_named(): + """ + The caller picks what to ask for, and the provider picks what to run. Only + the provider's own frames may set the price of a session, so a model the + caller put in the request cannot stand in for one the provider never named. + """ + logging_obj = _logging_obj() + handled = _handle( [_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0)], kwargs=MappingProxyType({"model": REALTIME_MODEL}), + logging_obj=logging_obj, ) - assert handled["kwargs"]["model"] == REALTIME_MODEL - assert handled["result"].model == REALTIME_MODEL - assert handled["kwargs"]["response_cost"] == pytest.approx(_realtime_cost(REALTIME_MODEL, 10, 0, 5, 0)) + assert handled["result"] is None + assert handled["kwargs"] == {"model": REALTIME_MODEL} + assert logging_obj.model_call_details == {} + + +def test_each_responses_turn_is_priced_on_the_model_that_turn_ran(): + """ + One socket can run several models one after another, so a connection that + starts on a big model and finishes on a small one has to pay the big + model's price for the turn that used it. + """ + handled = _handle( + [ + _responses_completed(REASONING_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7), + ], + url_route="/openai_passthrough/v1/responses", + ) + + first = litellm.model_cost[REASONING_MODEL] + second = litellm.model_cost[RESPONSES_MODEL] + expected_cost = ( + 13 * first["input_cost_per_token"] + + 5 * first["output_cost_per_token"] + + 20 * second["input_cost_per_token"] + + 7 * second["output_cost_per_token"] + ) + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 45 + + +def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model(): + handled = _handle( + [ + _session_created(RESPONSES_MODEL), + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 13, "output_tokens": 5, "total_tokens": 18}}, + }, + ], + url_route="/openai_passthrough/v1/responses", + ) + + prices = litellm.model_cost[RESPONSES_MODEL] + expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) def test_a_session_nobody_named_a_model_for_is_left_unpriced(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index ae139d2a70b..c955bd86639 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4767,8 +4767,8 @@ class FakeUpstreamWebSocket: self._first_frame = first_frame self.close = AsyncMock() - async def recv(self, decode: bool = True): - return self._first_frame + async def recv(self, decode: bool | None = None): + return self._first_frame if decode is False else self._first_frame.decode("utf-8") def __aiter__(self): return self @@ -5130,13 +5130,16 @@ class RecordingFrameModelGate: return self.refusals.get(model) +def _received_frame(frame): + if isinstance(frame, bytes): + return {"type": "websocket.receive", "bytes": frame} + return {"type": "websocket.receive", "text": frame} + + async def _relay_client_frames(frames, gate, user_api_key_dict=None): upstream_ws = RecordingUpstreamWebSocket() websocket = _client_websocket( - AsyncMock( - side_effect=[{"type": "websocket.receive", "text": frame} for frame in frames] - + [{"type": "websocket.disconnect"}] - ) + AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}]) ) with ( @@ -5185,7 +5188,6 @@ async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(fr assert gate.checked == [expected_model] assert [call.args[0] for call in relayed.upstream.send.await_args_list] == [json.dumps(frame)] assert relayed.client.close.await_args_list == [call()] - assert relayed.success_handler.call_args.kwargs["model"] == expected_model @pytest.mark.asyncio @@ -5260,17 +5262,71 @@ async def test_websocket_passthrough_forwards_frames_that_name_no_model(frame): @pytest.mark.asyncio -async def test_websocket_passthrough_bills_the_latest_model_a_client_frame_named(): +async def test_websocket_passthrough_never_bills_the_model_a_client_frame_named(): + """ + A frame is what the caller asked for, not what the provider ran. A caller + who names a cheap model in a frame the provider ignores must not move the + session onto that model's price, so only the provider's own frames decide + what a session is billed as. + """ gate = RecordingFrameModelGate() frames = [ json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), - json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1-mini"}}), + json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}), ] relayed = await _relay_client_frames(frames, gate) assert gate.checked == ["gpt-realtime-2.1", "gpt-realtime-2.1-mini"] - assert relayed.success_handler.call_args.kwargs["model"] == "gpt-realtime-2.1-mini" + assert "model" not in relayed.success_handler.call_args.kwargs + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame, expected_model", + [ + ('{"type": "response.create", "\\u006dodel": "gpt-5.4"}', "gpt-5.4"), + ('{"type": "session.update", "session": {"\\u006d\\u006fdel": "gpt-realtime-2.1"}}', "gpt-realtime-2.1"), + ], +) +async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(frame, expected_model): + """ + JSON lets a caller spell the same key several ways, so a frame carrying + "\\u006dodel" asks for a model just as plainly as one carrying "model" and + has to face the same check. + """ + refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." + gate = RecordingFrameModelGate(MappingProxyType({expected_model: refusal})) + + relayed = await _relay_client_frames([frame], gate) + + assert gate.checked == [expected_model] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for(): + refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode() + + relayed = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-realtime-2.1"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_forwards_binary_audio_untouched(): + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) + audio = bytes(range(256)) * 4 + + relayed = await _relay_client_frames([audio], gate) + + assert gate.checked == [] + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [audio] @pytest.mark.asyncio @@ -5284,13 +5340,13 @@ async def test_websocket_passthrough_without_a_frame_gate_forwards_every_frame() class StreamingUpstreamWebSocket: - def __init__(self, setup_frame: str, streamed_frames: Sequence[str]): + def __init__(self, setup_frame: str | bytes, streamed_frames: Sequence[str]): self._setup_frame = setup_frame self._streamed_frames = iter(streamed_frames) self.close = AsyncMock() self.send = AsyncMock() - async def recv(self, decode: bool = True) -> str: + async def recv(self, decode: bool = True) -> str | bytes: return self._setup_frame def __aiter__(self): @@ -5320,7 +5376,8 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out A realtime client closes the moment it has the answer it asked for, so the next frame the provider sends lands on a socket that is already gone. That is how a session ends, not a gateway failure, and the turns the caller - already used still have to reach spend. + already used still have to reach spend, including the last frame the client + never collected. """ upstream_ws = StreamingUpstreamWebSocket( REALTIME_SESSION_CREATED, @@ -5343,6 +5400,71 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out accept_websocket=False, ) + relayed_frames: Final = success_handler.call_args.kwargs["response_body"] + assert [frame["type"] for frame in relayed_frames] == [ + "session.created", + "response.done", + "response.output_text.delta", + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_relays_a_first_frame_that_is_not_a_json_object(): + """ + A passthrough carries whatever the upstream sends. An upstream that opens + with a JSON array has nothing to bill from, and it still has to reach the + caller instead of taking the connection down. + """ + upstream_ws = StreamingUpstreamWebSocket("[1, 2]", [REALTIME_SESSION_CREATED, REALTIME_RESPONSE_DONE]) + websocket = _client_websocket(_pending_receive) + websocket.send_text = AsyncMock(side_effect=[None, None, WebSocketDisconnect(code=1006)]) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/openai_passthrough/v1/realtime", + accept_websocket=False, + ) + + assert websocket.send_text.await_args_list[0] == call("[1, 2]") + relayed_frames: Final = success_handler.call_args.kwargs["response_body"] + assert [frame["type"] for frame in relayed_frames] == ["session.created", "response.done"] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_relays_a_binary_first_frame_as_binary(): + """ + Some sessions open with audio rather than JSON, so the opening frame has to + reach the caller as the same bytes the provider sent instead of being read + as text and taking the connection down. + """ + audio: Final = bytes(range(256)) + upstream_ws = StreamingUpstreamWebSocket(audio, [REALTIME_SESSION_CREATED, REALTIME_RESPONSE_DONE]) + websocket = _client_websocket(_pending_receive) + websocket.send_text = AsyncMock(side_effect=[None, WebSocketDisconnect(code=1006)]) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/openai_passthrough/v1/realtime", + accept_websocket=False, + ) + + assert websocket.send_bytes.await_args_list == [call(audio)] relayed_frames: Final = success_handler.call_args.kwargs["response_body"] assert [frame["type"] for frame in relayed_frames] == ["session.created", "response.done"] 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 562599cb945..11fc2311795 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,17 @@ """OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks.""" import json -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from contextlib import contextmanager from dataclasses import dataclass from types import MappingProxyType, SimpleNamespace from typing import Final -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from starlette.routing import WebSocketRoute -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _OPENAI_WS_DISABLED_REFUSAL, _openai_websocket_refusal, @@ -267,6 +268,66 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_mo assert requested_model in refusal +GET_USER_OBJECT: Final = "litellm.proxy.auth.auth_checks.get_user_object" + + +def _personal_user(models: list[str]) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id="user-lit7014", models=models, max_budget=None, spend=0.0) + + +@contextmanager +def _personal_user_allowlist(models: list[str]) -> Iterator[AsyncMock]: + """Stand up the database-backed owner lookup the allowlist check runs, and report every call it made.""" + lookup: Final = AsyncMock(return_value=_personal_user(models)) + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: read at call time + patch(GET_USER_OBJECT, lookup), + ): + yield lookup + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_a_model_the_person_behind_the_key_lacks(): + """ + A key with no team of its own carries its owner's model allowlist, the same + one every HTTP route applies, so a frame naming a model that person cannot + use is refused even when the key itself names no models at all. + """ + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=[], user_id="user-lit7014") + + with _personal_user_allowlist(["gpt-5.4-mini"]): + refusal = await gate("gpt-5.4", token) + + assert refusal is not None + assert "gpt-5.4" in refusal + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_allows_a_model_the_person_behind_the_key_owns(): + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=[], user_id="user-lit7014") + + with _personal_user_allowlist(["gpt-5.4-mini", "gpt-realtime-2.1"]): + assert await gate("gpt-realtime-2.1", token) is None + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_leaves_a_team_key_to_the_teams_allowlist(): + """ + A key that belongs to a team is governed by the team's models, not by the + personal allowlist of whoever created it, so the owner's own restrictions + must not close a session the team is entitled to. + """ + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=[], user_id="user-lit7014", team_id="team-lit7014") + + with _personal_user_allowlist(["gpt-5.4-mini"]) as owner_lookup: + assert await gate("gpt-5.4", token) is None + + assert owner_lookup.await_args_list == [] + + @pytest.mark.asyncio @pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]]) async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models):