fix(passthrough): bill a client frame whose upstream send is cancelled

When the upstream ends the session while the relay is still inside a client frame's send, the task is cancelled after the frame may already be on the wire. Each send now goes through one helper that adds the frame's count on success and on cancellation, but not when the send itself fails

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:

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(