diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index bd24be1b808..d09ea8818fd 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2519,7 +2519,6 @@ class WebsocketClientFrameGate(Protocol): async def _relay_client_frame( - websocket: WebSocket, upstream_ws: ClientConnection, checked_frame: str | bytes, forwarded_frame: str | bytes, @@ -2529,14 +2528,12 @@ async def _relay_client_frame( """Forward one client frame upstream, or refuse it and hand back the denial that ended the session. The denial travels out so the caller reports a failed session instead of leaving the success hooks to - record a refusal as a success, and tearing the upstream down is the caller's job: awaiting that close - here lets the upstream reader finish first, which cancels this task mid-close and loses the denial + record a refusal as a success, and closing either socket is the caller's job: awaiting a close here lets + the upstream reader finish first, which cancels this task mid-close and loses the denial """ refusal: Final = None if gate is None else await gate(checked_frame, valid_token) if refusal is None: await upstream_ws.send(forwarded_frame) - return None - await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) return refusal @@ -2715,7 +2712,6 @@ async def websocket_passthrough_request( if text_data is not None: if ( text_refusal := await _relay_client_frame( - websocket, upstream_ws, text_data, _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), @@ -2783,7 +2779,6 @@ async def websocket_passthrough_request( elif bytes_data is not None: if ( bytes_refusal := await _relay_client_frame( - websocket, upstream_ws, bytes_data, bytes_data, @@ -2873,6 +2868,9 @@ async def websocket_passthrough_request( (result for result in task_results if isinstance(result, ProxyException)), None ) if client_frame_refusal is not None: + await websocket.close( + code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(client_frame_refusal.message) + ) await upstream_ws.close() if client_frame_refusal is not None and not websocket_messages: await log_failure(client_frame_refusal) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 8bd2106e8cd..e388c658200 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7555,7 +7555,7 @@ async def _run_vertex_live_gated_passthrough( """Drive /vertex_ai/live with the real frame gate, returning the client and upstream sockets""" upstream_ws = RecordingUpstreamWebSocket() websocket, _ = await _run_vertex_live_gated_session( - AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]), + _client_websocket(AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}])), upstream_ws, valid_token, llm_router or _vertex_live_router(), @@ -7564,7 +7564,7 @@ async def _run_vertex_live_gated_passthrough( async def _run_vertex_live_gated_session( - receive: Callable[[], Awaitable[Mapping[str, str | bytes]]], + websocket: MagicMock, upstream_ws: "RecordingUpstreamWebSocket | ScriptedUpstreamWebSocket", valid_token: UserAPIKeyAuth, llm_router: litellm.Router, @@ -7574,8 +7574,6 @@ async def _run_vertex_live_gated_session( _build_vertex_live_setup_model_rewriter, ) - websocket = _client_websocket(receive) - with _patched_websocket_passthrough_environment(upstream_ws) as environment: await websocket_passthrough_request( websocket=websocket, @@ -8330,7 +8328,7 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( yield _setup_frame("gemini-3-pro-live") websocket, environment = await _run_vertex_live_gated_session( - client_frames().__anext__, + _client_websocket(client_frames().__anext__), upstream_ws, UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), _vertex_live_router(), @@ -8346,6 +8344,66 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( environment.proxy_logging.post_call_failure_hook.assert_not_awaited() +class StreamingUpstreamWebSocket(ScriptedUpstreamWebSocket): + """An upstream that sends the given frames in order, then keeps streaming model output until it is closed""" + + async def recv(self, decode: bool = True) -> str | None: + frame: Final = next(self._frames, None) + if frame is not None: + return frame + self.sent_every_frame.set() + await asyncio.sleep(0) + return json.dumps({"serverContent": {"modelTurn": {"parts": [{"text": "still talking"}]}}}) + + +def _client_websocket_with_a_close_handshake(receive) -> MagicMock: + """A client socket that, like Starlette's, refuses sends once its close starts and yields while closing""" + from starlette.websockets import WebSocketState + + websocket: Final = _client_websocket(receive) + + async def close(*args: object, **kwargs: object) -> None: + websocket.application_state = WebSocketState.DISCONNECTED + for _ in range(ClosableUpstreamWebSocket.CLOSE_HANDSHAKE_YIELDS): + await asyncio.sleep(0) + + async def send_text(data: str) -> None: + if websocket.application_state == WebSocketState.DISCONNECTED: + raise RuntimeError('Cannot call "send" once a close message has been sent.') + + websocket.close = AsyncMock(side_effect=close) + websocket.send_text = AsyncMock(side_effect=send_text) + return websocket + + +@pytest.mark.asyncio +async def test_vertex_live_turns_served_before_a_frame_refused_mid_stream_are_still_billed(): + upstream_ws: Final = StreamingUpstreamWebSocket( + ( + json.dumps({"setupComplete": {}}), + json.dumps({"serverContent": {"turnComplete": True}, "usageMetadata": {"totalTokenCount": 60}}), + ) + ) + + async def client_frames() -> AsyncIterator[dict[str, str]]: + yield _setup_frame("gemini-live-2.5-flash") + await upstream_ws.sent_every_frame.wait() + yield _setup_frame("gemini-3-pro-live") + + websocket, environment = await _run_vertex_live_gated_session( + _client_websocket_with_a_close_handshake(client_frames().__anext__), + upstream_ws, + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + _vertex_live_router(), + ) + + reason: Final = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + upstream_ws.close.assert_awaited() + environment.proxy_logging.post_call_success_hook.assert_awaited_once() + environment.proxy_logging.post_call_failure_hook.assert_not_awaited() + + async def _run_passthrough_with_frame_gate( gate: WebsocketClientFrameGate, ) -> tuple[MagicMock, ClosableUpstreamWebSocket, SimpleNamespace]: