From 4ee39dac33f8ea7ab5c07c191314b8e6ae28d382 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:51:30 -0700 Subject: [PATCH] fix(proxy): keep reading running answers after a refused websocket frame so they are billed --- .../pass_through_endpoints.py | 17 ++++-------- .../test_pass_through_endpoints.py | 26 ++++++++++++------- 2 files changed, 21 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 2ccc4fe56d2..7b304331f28 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2793,13 +2793,6 @@ async def refuse_websocket(websocket: WebSocket, refusal: WebsocketFrameModelRef 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 refuse_websocket(websocket, refusal) - await upstream_ws.close() - - def _stamp_billed_turn(kwargs: Mapping[str, object], turn: _TurnAuthorization) -> None: """Charge a turn to the access groups that authorized its model, and its model budget to the model it asked for, since the provider may answer under a dated snapshot name the budget was never set on.""" @@ -2894,11 +2887,11 @@ async def _relay_client_frame( gate: WebsocketFrameModelGate | None, ledger: _WebsocketTurnLedger, ) -> _ForwardClientFrame | _RefuseClientFrame: - """Forward one client frame upstream, or refuse it and end the session.""" + """Forward one client frame upstream, or refuse it and close the client.""" verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate) match verdict: case _RefuseClientFrame(refusal=refusal): - await _refuse_client_frame(websocket, upstream_ws, refusal) + await refuse_websocket(websocket, refusal) case _ForwardClientFrame(starts_turn=True, turn=turn, lane=lane): ledger.queue_turn(lane, turn) await upstream_ws.send(forwarded_frame) @@ -2943,7 +2936,7 @@ async def websocket_passthrough_request( upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones that open the session, so each turn logs its own spend and keeps the session's provider-named model walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client - left, so the provider's usage for it still reaches spend + left or was refused, so the provider's usage for it still reaches spend builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider model: The model the caller named before connecting, logged on a failure before the setup frame arrives """ @@ -3174,8 +3167,6 @@ async def websocket_passthrough_request( message = await websocket.receive() message_type = message.get("type") if message_type == "websocket.disconnect": - if not _turn_in_flight(upstream_frame_kind, websocket_messages): - await upstream_ws.close() break text_data: str | None = message.get("text") @@ -3258,6 +3249,8 @@ async def websocket_passthrough_request( _RefuseClientFrame, ): break + if not _turn_in_flight(upstream_frame_kind, websocket_messages): + await upstream_ws.close() except asyncio.CancelledError: raise except Exception: 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 95629b7f6c5..217e8781c98 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 @@ -8052,14 +8052,13 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out class ScriptedUpstreamWebSocket: - """Answers each client frame with its scripted reply, and flushes what was still in flight when it is closed""" + """Answers each client frame with its scripted reply""" - def __init__(self, *opening: str, replies: Sequence[Sequence[str]] = (), in_flight: Sequence[str] = ()): + def __init__(self, *opening: str, replies: Sequence[Sequence[str]] = ()): self._frames: asyncio.Queue[str | None] = asyncio.Queue() for frame in opening: self._frames.put_nowait(frame) self._replies = iter(replies) - self._in_flight = in_flight self._closed = False self.send = AsyncMock(side_effect=self._reply) self.close = AsyncMock(side_effect=self._flush_and_close) @@ -8072,8 +8071,6 @@ class ScriptedUpstreamWebSocket: if self._closed: return self._closed = True - for frame in self._in_flight: - self._frames.put_nowait(frame) self._frames.put_nowait(None) while not self._frames.empty(): await asyncio.sleep(0) @@ -8586,8 +8583,8 @@ async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_lat """ A refused frame closes the client while the provider may still be sending the answer to an earlier, allowed frame. That answer cannot go to a closed client, - but it was paid for, so it has to reach spend instead of turning the whole - session into a failure that bills nothing. + but the provider bills it once it ends, so the gateway keeps reading until it + ends and its usage reaches spend. """ gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) client: Final = _turn_taking_client( @@ -8597,10 +8594,19 @@ async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_lat ], answered_by='"response.created"', ) - created, completed = _responses_turn("gpt-5.4-mini") - upstream_ws: Final = ScriptedUpstreamWebSocket(replies=[[created]], in_flight=[completed]) + refused: Final = asyncio.Event() + close_client: Final = client.close.side_effect - relayed: Final = await _relay_turns(client, upstream_ws, gate) + def close_client_and_note_the_refusal(*args: object, **kwargs: object) -> None: + close_client(*args, **kwargs) + refused.set() + + client.close = AsyncMock(side_effect=close_client_and_note_the_refusal) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + refused, _responses_turn_started("gpt-5.4-mini"), (_responses_turn_ended("gpt-5.4-mini"),) + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate, walkout_turn_drain_seconds=5) assert [json.loads(text.args[0])["type"] for text in client.send_text.await_args_list] == [ "response.created",