From 537cd2e6d41ea24005117cdc1e2c4955dca8e4c0 Mon Sep 17 00:00:00 2001
From: shrey kharbanda <shrey@berri.ai>
Date: Mon, 26 Oct 2026 19:40:37 -0700
Subject: [PATCH] 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
---
 .../pass_through_endpoints.py                 | 19 +++++++-----
 .../test_pass_through_endpoints.py            | 29 +++++++++++++++++++
 2 files changed, 41 insertions(+), 7 deletions(-)

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(