From d48e964b6e406512ffb2e3c3c9c369804bf78416 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:57:08 -0700 Subject: [PATCH] fix(proxy): close the Vertex AI Live sockets only after the refusal is kept Closing the client socket inside the frame relay let the upstream reader finish first, which cancelled the relay mid-close, lost the refusal, and logged a session that served turns as a failure. The relay now returns the refusal and the handler closes both sockets. --- .../pass_through_endpoints.py | 12 ++-- .../test_pass_through_endpoints.py | 68 +++++++++++++++++-- 2 files changed, 68 insertions(+), 12 deletions(-) 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]: