From 1c6e2f795919ca405ef2cba37899fbfb0afe7c1d Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 20:34:00 -0400 Subject: [PATCH] fix(proxy): restore keepalive-ping exclusion from streaming disconnect refund An earlier round's rebase replayed a stale, pre-fix version of a commit this branch had already picked up the fix for, reverting delivered_chunk's keepalive-exclusion check and the _withheld_provider_output guard back to their original, buggy form. A client disconnecting after only keepalive pings (or while an agentic stream holds back real output) was refunded to input cost even though billable output had already been generated, exactly the veria-ai finding already fixed once upstream. Restores both checks; test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_cost 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 a687aaf2c5c..f98ff8e8e62 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -29,6 +29,7 @@ from litellm.constants import ( NON_INFERENCE_CALL_TYPES, RETURN_RAW_MODEL_NAME_METADATA_KEY, 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 @@ -3497,8 +3502,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): @@ -3512,7 +3518,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, )