From 7a8ff014eeaf1743094005926bbd0670c75d264b Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:15:21 -0400 Subject: [PATCH] fix(proxy): restore keepalive-ping exclusion from streaming disconnect refund The disconnect-state-release hook added in the previous commit dropped the existing has_buffered_provider_output guard and the STREAM_SSE_KEEPALIVE_PING_BYTES exclusion while rewiring the streaming generator's cleanup path, so a client disconnecting after only keepalive pings (or while an agentic stream holds back real output) got refunded to input cost even when billable output had already been generated. Restores both checks; veria-ai caught this on review, and the existing test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_cost regression test now passes again. --- litellm/proxy/common_request_processing.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f743987669f..a03ffa9905d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -30,6 +30,7 @@ from litellm.constants import ( RETURN_RAW_MODEL_NAME_METADATA_KEY, ROUTER_MODEL_NAME_RESPONSE_FIELD, STREAM_SSE_DATA_PREFIX, + STREAM_SSE_KEEPALIVE_PING_BYTES, UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -203,6 +204,10 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma } +def _withheld_provider_output(response: object) -> bool: + return getattr(response, "has_buffered_provider_output", False) is True + + def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: return any( isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True @@ -3491,8 +3496,9 @@ class ProxyBaseLLMRequestProcessing: # so a GeneratorExit on client disconnect is raised there and any # statement after the yield never runs. The slow-path hook is # awaited above, so a cancellation during it still leaves this - # False and refunds. - delivered_chunk = True + # False and refunds. A keepalive ping carries no provider output, + # so it must not suppress that refund. + delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES yield serialize_chunk(chunk) stream_completed = True except (asyncio.CancelledError, GeneratorExit): @@ -3506,7 +3512,7 @@ class ProxyBaseLLMRequestProcessing: # only sees GeneratorExit on GC) cannot own the refund. if not stream_completed: client_disconnected = True - if not delivered_chunk: + if not delivered_chunk and not _withheld_provider_output(response): from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, )