scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

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

Browse source 
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.
This commit is contained in:
Deepanshu 2026-08-27 20:34:00 -04:00
parent 2b0c1699d4
commit 1c6e2f7959
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

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