scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

fix(proxy): keep reading running answers after a refused websocket frame so they are billed

Browse source
This commit is contained in:
mateo-berri 2026-10-30 18:51:30 -07:00
parent 6998e0f9ab
commit 4ee39dac33
2 changed files with 21 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -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:

26
tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -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",