scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(proxy): restore keepalive-ping exclusion from streaming disconnect refund

Browse source 
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.
This commit is contained in:
Deepanshu 2026-08-26 11:15:21 -04:00
parent e2a0002850
commit 7a8ff014ee
1 changed files with 9 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/proxy/common_request_processing.py
View file

@ -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,
)