diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 22a44de6bd9..a19ac2e66ef 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2628,9 +2628,18 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): Upstream connection established successfully", endpoint ) + async def send_upstream(frame: str | bytes) -> None: + nonlocal client_text_count + billed: Final = 0 if count_client_text is None else count_client_text(frame) + try: + await upstream_ws.send(frame) + except asyncio.CancelledError: + client_text_count += billed # rebind-ok: the frame can already be on the wire when the cancel lands + raise + client_text_count += billed # rebind-ok: the session total grows with each delivered frame + async def forward_client_to_upstream() -> None: """Forward messages from client to upstream WebSocket""" - nonlocal client_text_count try: while True: message = await websocket.receive() @@ -2699,13 +2708,9 @@ async def websocket_passthrough_request( ) # Not a JSON message or doesn't contain setup data - await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) - if count_client_text is not None: - client_text_count += count_client_text(text_data) + await send_upstream(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) elif bytes_data is not None: - await upstream_ws.send(bytes_data) - if count_client_text is not None: - client_text_count += count_client_text(bytes_data) + await send_upstream(bytes_data) 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 d4f0b91ca92..5c879f9b97b 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 @@ -7943,6 +7943,35 @@ async def test_websocket_passthrough_counts_client_text_when_the_upstream_ends_w ] +class _DrainStallingUpstreamWebSocket(SpeakingUpstreamWebSocket): + """Takes the client's Close, then stalls in send the way a full write buffer does, so the session ends while the + relay is still inside that send""" + + async def send(self, message: str | bytes) -> None: + await super().send(message) + if message == SPEAK_CLOSE_FRAME: + await asyncio.Event().wait() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_model_cost_map") +async def test_websocket_passthrough_bills_a_client_frame_whose_send_is_cancelled_after_reaching_the_upstream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session: Final = await _relay_and_capture_spend( + monkeypatch, + _DrainStallingUpstreamWebSocket(), + _speak_client(), + DEEPGRAM_SPEAK_TARGET, + "/deepgram/v1/speak", + count_client_text=len, + ) + + assert [log["response_cost"] for log in session.spend_logs] == [ + pytest.approx(SPEAK_SESSION_CLIENT_LENGTH * _aura_2_usd_per_character()) + ] + + @pytest.mark.asyncio @pytest.mark.usefixtures("local_model_cost_map") async def test_deepgram_speak_session_logs_character_spend_once_and_never_the_spoken_text(