scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 178

Merge pull request #39589 from BerriAI/litellm_fix_v1_messages_midstream_timeout_failure_logging

Browse source 
fix(proxy): log mid-stream /v1/messages failures as failures with partial usage
This commit is contained in:
Mateo Wang 2026-09-04 13:20:30 -07:00 • committed by GitHub
commit 44b1cc7b0f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 1145 additions and 176 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/litellm_core_utils/litellm_logging.py
View file

@ -10,7 +10,7 @@ import subprocess
import sys
import time
import traceback
from collections.abc import Callable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from datetime import datetime as dt_object
from functools import lru_cache
from types import MappingProxyType, TracebackType
@ -576,6 +576,7 @@ class Logging(LiteLLMLoggingBaseClass):
# enqueue closure here instead of firing it immediately.
self._defer_async_logging: bool = False
self._enqueue_deferred_logging: Callable[[], None] | None = None
self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None:
"""Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""
@ -1894,6 +1895,11 @@ class Logging(LiteLLMLoggingBaseClass):
**kwargs,
)
def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None:
"""Stash what an interrupted stream already consumed so the failure log bills it instead of zero."""
self.model_call_details["combined_usage_object"] = usage
self.model_call_details["response_cost"] = response_cost
async def dispatch_failure_handlers(
self,
exception: Exception,

54
litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py
View file

@ -1,6 +1,6 @@
import asyncio
import json
from collections.abc import AsyncIterator, Mapping
from collections.abc import AsyncIterator, Mapping, Sequence
from datetime import datetime
from typing import Any, Final, Protocol, runtime_checkable
@ -177,12 +177,6 @@ def _try_claim_detached_drain_slot() -> bool:
def _exception_left_unconsumed(queue: "asyncio.Queue[bytes | None | BaseException]", exc: BaseException) -> bool:
"""After client detach the relay never reads the queue again, so drain it here.
The forwarded exception still sitting in the queue means the relay tore
down before re-raising it, so the proxy's failure handling never ran and
the caller must salvage spend itself.
"""
remaining: Final = tuple(queue.get_nowait() for _ in range(queue.qsize()))
return any(item is exc for item in remaining)
@ -671,27 +665,41 @@ class BaseAnthropicMessagesStreamingIterator:
self,
queue: "asyncio.Queue[bytes | None | BaseException]",
client_detached: "asyncio.Event",
collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks
exc: BaseException,
collected_chunks: Sequence[bytes],
exc: Exception,
) -> None:
"""Forward a provider error to a still-connected client, else salvage partial spend.
"""Log the request as failed with its partial usage, then make sure the proxy's failure hook runs once.
Handing the original exception to the client-facing generator lets it
re-raise so the proxy's failure handling keeps the provider status and
owns logging (no success-bill). If the client already went away, or
disconnects before ever consuming the queued exception, no failure hook
runs, so bill the partial instead of dropping the request.
A still-connected client gets the original exception through the queue,
the relay re-raises it, and the proxy's own failure handling records the
failed spend. When the client already left, or leaves before consuming
the queued exception, that handling never runs, so the detached-failure
hook the proxy armed on the logging object fires here instead.
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
PassThroughStreamingHandler.schedule_stream_failure_logging(
litellm_logging_obj=self.litellm_logging_obj,
endpoint_type=EndpointType.ANTHROPIC,
request_body=self.request_body,
raw_bytes=collected_chunks,
exception=exc,
)
if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc):
await client_detached.wait()
if not _exception_left_unconsumed(queue, exc):
return
verbose_proxy_logger.warning(
"async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)",
len(collected_chunks),
type(exc).__name__,
exc,
)
await self._bill_collected_chunks(collected_chunks, stream_teardown=True)
await self._fire_detached_failure_hook(exc)
async def _fire_detached_failure_hook(self, exc: Exception) -> None:
from litellm._logging import verbose_proxy_logger
on_detached_failure: Final = getattr(self.litellm_logging_obj, "_on_detached_stream_failure", None)
if on_detached_failure is None:
return
try:
await on_detached_failure(exc)
except Exception as hook_failure: # noqa: BLE001 # a failing proxy hook must not crash the detached pump
verbose_proxy_logger.warning(
"async_sse_wrapper detached failure hook raised: %s(%s)", type(hook_failure).__name__, hook_failure
)

33
litellm/proxy/common_request_processing.py
View file

@ -2520,6 +2520,11 @@ class ProxyBaseLLMRequestProcessing:
# This handles cases like websearch_interception agentic loop
# which returns a non-streaming dict even for streaming requests
if self._is_streaming_response(response):
self._arm_detached_stream_failure_hook(
logging_obj=logging_obj,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
selected_data_generator = ProxyBaseLLMRequestProcessing.async_sse_data_generator(
response=response,
user_api_key_dict=user_api_key_dict,
@ -2875,6 +2880,34 @@ class ProxyBaseLLMRequestProcessing:
),
)
def _arm_detached_stream_failure_hook(
self,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: "UserAPIKeyAuth",
proxy_logging_obj: ProxyLogging,
) -> None:
"""Let a stream that fails after the client left still reach ``post_call_failure_hook``.
The client-facing generator reports a mid-stream failure itself, but once
the client disconnects that generator is gone and the detached upstream
drain is the only code that sees the provider error. It fires this closure
so the failed spend is still written and the budget reservation released;
a replacement error the hook raises has no client left to reach.
"""
request_data: Final = self.data
async def _on_detached_stream_failure(exc: Exception) -> None:
try:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=exc,
request_data=request_data,
)
except HTTPException:
return
logging_obj._on_detached_stream_failure = _on_detached_stream_failure
def _is_streaming_response(self, response: Any) -> bool:
"""
Check if the response object is actually a streaming response by inspecting its type.

58
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -4518,12 +4518,25 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
statuses=statuses,
)
def _recovered_partial_usage_tokens(self, source: Mapping[str, object]) -> tuple[int, int, int]:
usage: Final = source.get("combined_usage_object")
if not isinstance(usage, Usage) or (usage.completion_tokens or 0) <= 0:
return 0, 0, 0
billable_input, completion_tokens, _ = self._resolve_io_token_reconcile_usage(usage)
return (
self._get_total_tokens_from_usage(usage=usage, rate_limit_type=self.get_rate_limit_type()),
billable_input,
completion_tokens,
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
"""
On failure: decrement max_parallel_requests and refund the upfront
TPM reservation only against the scopes the reservation actually
charged. Unreserved scopes were never incremented at pre-call, so
refunding them would drive their counter negative.
refunding them would drive their counter negative. A failed stream
whose partial usage was recovered settles the reservation at that
usage instead of refunding it.
"""
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
@ -4552,31 +4565,31 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if stash is None or stash.reservation_released
else (stash.reserved_tokens, stash.itpm_reserved_tokens, stash.otpm_reserved_tokens)
)
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(kwargs)
if stash is not None and reserved_tokens > 0:
verbose_proxy_logger.debug("Releasing reserved TPM tokens on failure: %s", reserved_tokens)
# Refund only against the scopes the reservation actually
# charged. _build_reservation_aware_tpm_ops with
# actual_tokens=0 emits -reserved on reserved scopes and 0
# on unreserved (skipped), so unreserved scopes can't drift
# negative.
verbose_proxy_logger.debug(
"Settling reserved TPM tokens on failure: reserved=%s actual=%s", reserved_tokens, tpm_actual
)
# Settle only against the scopes the reservation actually
# charged: unreserved scopes were never incremented, so a
# refund there would drive their counter negative.
pipeline_operations.extend(
self._build_reservation_aware_tpm_ops(
targets=list(stash.reserved_scopes),
reserved_scopes=stash.reserved_scopes,
actual_tokens=0,
actual_tokens=tpm_actual,
reserved_tokens=reserved_tokens,
)
)
# Refund project ITPM/OTPM reservations the same way -- full
# refund, since a failed call has no billable usage to reconcile
# against.
# Settle project ITPM/OTPM reservations the same way: at the
# recovered partial usage, or a full refund when there is none.
itpm_operations: Final = (
self._build_project_reservation_ops(
targets=tuple(stash.itpm_reserved_scopes),
reserved_scopes=stash.itpm_reserved_scopes,
actual_tokens=0,
actual_tokens=itpm_actual,
reserved_tokens=itpm_reserved,
reservation_window_identities=stash.itpm_reserved_window_identities,
)
@ -4584,7 +4597,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else self._build_reservation_aware_tpm_ops(
targets=tuple(stash.itpm_reserved_scopes),
reserved_scopes=stash.itpm_reserved_scopes,
actual_tokens=0,
actual_tokens=itpm_actual,
reserved_tokens=itpm_reserved,
)
if stash is not None and itpm_reserved > 0
@ -4595,7 +4608,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._build_project_reservation_ops(
targets=tuple(stash.otpm_reserved_scopes),
reserved_scopes=stash.otpm_reserved_scopes,
actual_tokens=0,
actual_tokens=otpm_actual,
reserved_tokens=otpm_reserved,
reservation_window_identities=stash.otpm_reserved_window_identities,
)
@ -4603,7 +4616,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else self._build_reservation_aware_tpm_ops(
targets=tuple(stash.otpm_reserved_scopes),
reserved_scopes=stash.otpm_reserved_scopes,
actual_tokens=0,
actual_tokens=otpm_actual,
reserved_tokens=otpm_reserved,
)
if stash is not None and otpm_reserved > 0
@ -4742,7 +4755,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
removal is a no-op ZREM on a second run), and the TPM/ITPM/OTPM
refund is guarded by the stash's ``reservation_released`` flag — if
both this hook and async_log_failure_event end up running in the same
flow, only the first release/refund applies.
flow, only the first release/refund applies. A mid-stream failure
relayed here with recovered partial usage settles the reservation at
that usage instead of refunding it.
"""
try:
stash: Final = get_request_stash()
@ -4769,12 +4784,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
otpm_reserved: Final = stash.otpm_reserved_tokens
if reserved_tokens <= 0 and itpm_reserved <= 0 and otpm_reserved <= 0:
return
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(request_data)
combined_ops: Final = (
self._build_reservation_aware_tpm_ops(
targets=tuple(stash.reserved_scopes),
reserved_scopes=stash.reserved_scopes,
actual_tokens=0,
actual_tokens=tpm_actual,
reserved_tokens=reserved_tokens,
)
if reserved_tokens > 0
@ -4784,7 +4800,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._build_project_reservation_ops(
targets=tuple(stash.itpm_reserved_scopes),
reserved_scopes=stash.itpm_reserved_scopes,
actual_tokens=0,
actual_tokens=itpm_actual,
reserved_tokens=itpm_reserved,
reservation_window_identities=stash.itpm_reserved_window_identities,
)
@ -4792,7 +4808,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else self._build_reservation_aware_tpm_ops(
targets=tuple(stash.itpm_reserved_scopes),
reserved_scopes=stash.itpm_reserved_scopes,
actual_tokens=0,
actual_tokens=itpm_actual,
reserved_tokens=itpm_reserved,
)
if itpm_reserved > 0
@ -4802,7 +4818,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._build_project_reservation_ops(
targets=tuple(stash.otpm_reserved_scopes),
reserved_scopes=stash.otpm_reserved_scopes,
actual_tokens=0,
actual_tokens=otpm_actual,
reserved_tokens=otpm_reserved,
reservation_window_identities=stash.otpm_reserved_window_identities,
)
@ -4810,7 +4826,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else self._build_reservation_aware_tpm_ops(
targets=tuple(stash.otpm_reserved_scopes),
reserved_scopes=stash.otpm_reserved_scopes,
actual_tokens=0,
actual_tokens=otpm_actual,
reserved_tokens=otpm_reserved,
)
if otpm_reserved > 0

227
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -37,6 +37,7 @@ from litellm.types.utils import (
Message,
ModelResponse,
TextCompletionResponse,
Usage,
)
if TYPE_CHECKING:
@ -148,6 +149,144 @@ class AnthropicPassthroughLoggingHandler:
return model_group.removeprefix("passthrough/")
return model
@staticmethod
def _resolve_logged_model(
litellm_logging_obj: LiteLLMLoggingObj,
request_body: Mapping[str, object],
all_chunks: Sequence[str | bytes],
) -> str:
request_model: Final = request_body.get("model")
logged_model: Final = (
request_model
if isinstance(request_model, str) and request_model
else str(litellm_logging_obj.model_call_details.get("model") or "")
)
if logged_model and logged_model != "unknown":
return logged_model
return AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(all_chunks) or logged_model
@staticmethod
def _usage_only_response_or_none(
all_chunks: Sequence[str | bytes], model: str, speed: str | None
) -> ModelResponse | None:
try:
return AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
all_chunks=all_chunks, model=model, speed=speed
)
except Exception as e: # noqa: BLE001 # the usage-only fallback must never raise out of failure logging
verbose_proxy_logger.warning("Anthropic passthrough: usage-only fallback failed (model=%s): %s", model, e)
return None
@staticmethod
def _assemble_streaming_response(
all_chunks: Sequence[str | bytes],
litellm_logging_obj: LiteLLMLoggingObj,
model: str,
speed: str | None,
) -> ModelResponse | TextCompletionResponse | None:
try:
assembled: Final = AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=speed,
)
except Exception as e: # noqa: BLE001 # any assembly error falls back to usage-only cost
verbose_proxy_logger.warning(
"Anthropic passthrough: stream assembly raised (model=%s): %s; falling "
"back to usage-only cost from raw SSE events.",
model,
e,
)
return AnthropicPassthroughLoggingHandler._usage_only_response_or_none(all_chunks, model, speed)
if assembled is not None:
return assembled
return AnthropicPassthroughLoggingHandler._usage_only_response_or_none(all_chunks, model, speed)
@staticmethod
def _build_streaming_response_for_logging(
litellm_logging_obj: LiteLLMLoggingObj,
request_body: Mapping[str, object],
all_chunks: Sequence[str | bytes],
model: str,
) -> ModelResponse | TextCompletionResponse | None:
response: Final = AnthropicPassthroughLoggingHandler._assemble_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=AnthropicPassthroughLoggingHandler._cost_relevant_speed(request_body),
)
if response is None:
return None
AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens(
response=response, all_chunks=all_chunks, model=model
)
return response
@staticmethod
def record_partial_usage_for_failure(
litellm_logging_obj: LiteLLMLoggingObj,
request_body: Mapping[str, object],
all_chunks: Sequence[str | bytes],
) -> None:
if not all_chunks:
return
model: Final = AnthropicPassthroughLoggingHandler._resolve_logged_model(
litellm_logging_obj, request_body, all_chunks
)
partial_response: Final = AnthropicPassthroughLoggingHandler._build_streaming_response_for_logging(
litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=all_chunks, model=model
)
usage: Final = cast(Usage | None, getattr(partial_response, "usage", None))
if partial_response is None or usage is None:
return
litellm_logging_obj.record_partial_usage_for_failure(
usage=usage,
response_cost=AnthropicPassthroughLoggingHandler._cost_partial_stream_or_zero(
partial_response=partial_response, model=model, logging_obj=litellm_logging_obj
),
)
@staticmethod
def _cost_partial_stream_or_zero(
partial_response: ModelResponse | TextCompletionResponse, model: str, logging_obj: LiteLLMLoggingObj
) -> float:
try:
return AnthropicPassthroughLoggingHandler._compute_response_cost(
litellm_model_response=partial_response,
model=AnthropicPassthroughLoggingHandler._resolve_costing_model(model, logging_obj),
logging_obj=logging_obj,
)
except Exception as e: # noqa: BLE001 # an uncostable partial stream still bills its tokens, at zero cost
verbose_proxy_logger.warning(
"Anthropic passthrough: could not cost the partial usage of a failed stream (model=%s): %s", model, e
)
return 0.0
@staticmethod
def _compute_response_cost(
litellm_model_response: ModelResponse | TextCompletionResponse,
model: str,
logging_obj: LiteLLMLoggingObj,
) -> float:
if logging_obj.model_call_details.get("cache_hit") is True:
return 0.0
custom_llm_provider: Final = logging_obj.model_call_details.get("custom_llm_provider")
model_for_cost: Final = (
f"{custom_llm_provider}/{model}"
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/")
else model
)
return litellm.completion_cost(
completion_response=litellm_model_response,
model=model_for_cost,
custom_llm_provider=custom_llm_provider,
custom_pricing=use_custom_pricing_for_model(
litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None)
),
router_model_id=logging_obj.get_router_model_id(),
)
@staticmethod
def _extract_message_start_field(
all_chunks: Sequence[str | bytes],
@ -278,31 +417,9 @@ class AnthropicPassthroughLoggingHandler:
if logging_obj.model_call_details.get("stream") is True:
logging_obj.model_call_details["complete_streaming_response"] = litellm_model_response
try:
# Get custom_llm_provider from logging object if available (e.g., azure_ai for Azure Anthropic)
custom_llm_provider: Final = logging_obj.model_call_details.get("custom_llm_provider")
model = AnthropicPassthroughLoggingHandler._resolve_costing_model(model, logging_obj)
# Prepend custom_llm_provider to model if not already present
model_for_cost = model
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
model_for_cost = f"{custom_llm_provider}/{model}"
router_model_id: Final = logging_obj.get_router_model_id()
custom_pricing: Final = use_custom_pricing_for_model(
litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None)
)
response_cost: Final = (
0.0
if logging_obj.model_call_details.get("cache_hit") is True
else litellm.completion_cost(
completion_response=litellm_model_response,
model=model_for_cost,
custom_llm_provider=custom_llm_provider,
custom_pricing=custom_pricing,
router_model_id=router_model_id,
)
response_cost: Final = AnthropicPassthroughLoggingHandler._compute_response_cost(
litellm_model_response=litellm_model_response, model=model, logging_obj=logging_obj
)
kwargs["response_cost"] = response_cost
@ -356,57 +473,12 @@ class AnthropicPassthroughLoggingHandler:
- Logs in litellm callbacks
"""
speed: Final = AnthropicPassthroughLoggingHandler._cost_relevant_speed(request_body)
model = request_body.get("model", "")
# Check if it's available in the logging object
if (
not model
and hasattr(litellm_logging_obj, "model_call_details")
and litellm_logging_obj.model_call_details.get("model")
):
model = cast(str, litellm_logging_obj.model_call_details.get("model"))
if not model or model == "unknown":
chunk_model: Final = AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(all_chunks)
if chunk_model:
model = chunk_model
try:
complete_streaming_response = AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=speed,
)
except Exception as e:
# stream_chunk_builder re-raises assembly failures (as litellm.APIError)
# on large agentic tool-use / thinking streams; treat that the same as a
# None result so the usage-only fallback below still recovers cost
verbose_proxy_logger.warning(
"Anthropic passthrough: stream assembly raised (model=%s): %s; falling "
"back to usage-only cost from raw SSE events.",
model,
e,
)
complete_streaming_response = None
if complete_streaming_response is None:
# stream_chunk_builder cannot always reassemble large agentic streams, but
# Anthropic still emits token usage in the message_start / message_delta SSE
# events regardless of content shape; recover usage-only so cost is tracked.
# Guard it too: a raise here would defeat the point and drop the request
try:
complete_streaming_response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
all_chunks=all_chunks,
model=model,
speed=speed,
)
except Exception as e:
verbose_proxy_logger.warning(
"Anthropic passthrough: usage-only fallback failed (model=%s): %s",
model,
e,
)
complete_streaming_response = None
model: Final = AnthropicPassthroughLoggingHandler._resolve_logged_model(
litellm_logging_obj, request_body, all_chunks
)
complete_streaming_response: Final = AnthropicPassthroughLoggingHandler._build_streaming_response_for_logging(
litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=all_chunks, model=model
)
if complete_streaming_response is None:
verbose_proxy_logger.error(
"Unable to build complete streaming response for Anthropic passthrough endpoint, not logging..."
@ -415,11 +487,6 @@ class AnthropicPassthroughLoggingHandler:
"result": None,
"kwargs": {},
}
AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens(
response=complete_streaming_response,
all_chunks=all_chunks,
model=model,
)
kwargs: Final = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
litellm_model_response=complete_streaming_response,
model=model,

81
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -870,6 +870,35 @@ async def _log_passthrough_upstream_failure(
)
async def _relay_reporting_failures(
stream: AsyncGenerator[bytes, None],
upstream_status: int,
user_api_key_dict: UserAPIKeyAuth,
request_payload: dict, # mutable-ok: post_call_failure_hook lifts fields onto request_data in place
) -> AsyncGenerator[bytes, None]:
from litellm.proxy.proxy_server import proxy_logging_obj
try:
async for chunk in stream:
yield chunk
except Exception as e:
if upstream_status >= 400:
raise
try:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=request_payload,
traceback_str=traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG),
)
except Exception: # noqa: BLE001 - a failing logging callback must never mask the upstream error
verbose_proxy_logger.warning(
"pass_through_endpoint: post_call_failure_hook raised for a mid-stream upstream error",
exc_info=True,
)
raise
from litellm.passthrough.timeout_utils import (
DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, # noqa: F401 - re-exported for backward compat
resolve_llm_passthrough_timeout, # noqa: F401 - re-exported for backward compat
@ -1293,14 +1322,24 @@ async def pass_through_request(
return StreamingResponse(
wrap_passthrough_sse_bytes_with_keepalive_pings(
stream=_own_streamed_managed_ids(
stream=PassThroughStreamingHandler.chunk_processor(
response=response,
request_body=_parsed_body,
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=start_time,
passthrough_success_handler_obj=pass_through_endpoint_logging,
url_route=str(url),
stream=_relay_reporting_failures(
stream=PassThroughStreamingHandler.chunk_processor(
response=response,
request_body=_parsed_body,
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=start_time,
passthrough_success_handler_obj=pass_through_endpoint_logging,
url_route=str(url),
),
upstream_status=response.status_code,
user_api_key_dict=user_api_key_dict,
request_payload=_build_passthrough_failure_request_payload(
parsed_body=_parsed_body,
kwargs=kwargs,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
),
),
managed_id_provider=_managed_id_provider,
request=request,
@ -1374,14 +1413,24 @@ async def pass_through_request(
return StreamingResponse(
wrap_passthrough_sse_bytes_with_keepalive_pings(
stream=_own_streamed_managed_ids(
stream=PassThroughStreamingHandler.chunk_processor(
response=response,
request_body=_parsed_body,
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=start_time,
passthrough_success_handler_obj=pass_through_endpoint_logging,
url_route=str(url),
stream=_relay_reporting_failures(
stream=PassThroughStreamingHandler.chunk_processor(
response=response,
request_body=_parsed_body,
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=start_time,
passthrough_success_handler_obj=pass_through_endpoint_logging,
url_route=str(url),
),
upstream_status=response.status_code,
user_api_key_dict=user_api_key_dict,
request_payload=_build_passthrough_failure_request_payload(
parsed_body=_parsed_body,
kwargs=kwargs,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
),
),
managed_id_provider=_managed_id_provider,
request=request,

110
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -1,5 +1,7 @@
from collections.abc import Coroutine
from datetime import datetime
import traceback
from collections.abc import Coroutine, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final, Protocol
import httpx
@ -12,7 +14,7 @@ from litellm.proxy._types import PassThroughEndpointLoggingResultValues
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
from litellm.types.utils import StandardPassThroughResponseObject
from litellm.types.utils import StandardPassThroughResponseObject, Usage
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -44,12 +46,85 @@ class RouteStreamingLogging(Protocol):
) -> Coroutine[None, None, None]: ...
@dataclass(frozen=True, slots=True)
class PassThroughStreamContext:
passthrough_success_handler_obj: PassThroughEndpointLogging
url_route: str
start_time: datetime
class PassThroughStreamingHandler:
@staticmethod
def _stamp_first_chunk_if_needed(litellm_logging_obj: LiteLLMLoggingObj) -> None:
if litellm_logging_obj.completion_start_time is None:
litellm_logging_obj._update_completion_start_time(completion_start_time=datetime.now())
@staticmethod
def schedule_stream_failure_logging(
litellm_logging_obj: LiteLLMLoggingObj,
endpoint_type: EndpointType,
request_body: dict[str, object],
raw_bytes: Sequence[bytes],
exception: Exception,
stream_context: PassThroughStreamContext | None = None,
) -> None:
PassThroughStreamingHandler._record_partial_usage_for_failure(
litellm_logging_obj=litellm_logging_obj,
endpoint_type=endpoint_type,
request_body=request_body,
raw_bytes=raw_bytes,
stream_context=stream_context,
)
try:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=litellm_logging_obj.dispatch_failure_handlers(
exception, traceback.format_exc(), prefer_async_handlers=True
)
)
except Exception as e:
verbose_proxy_logger.error("Error scheduling stream failure logging: %s", e)
@staticmethod
def _record_partial_usage_for_failure(
litellm_logging_obj: LiteLLMLoggingObj,
endpoint_type: EndpointType,
request_body: dict[str, object],
raw_bytes: Sequence[bytes],
stream_context: PassThroughStreamContext | None,
) -> None:
if endpoint_type == EndpointType.ANTHROPIC:
AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=raw_bytes
)
return
if stream_context is None or not raw_bytes:
return
try:
partial_response, kwargs = PassThroughStreamingHandler._build_passthrough_logging_result(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=stream_context.passthrough_success_handler_obj,
url_route=stream_context.url_route,
request_body=request_body,
endpoint_type=endpoint_type,
start_time=stream_context.start_time,
raw_bytes=raw_bytes,
end_time=datetime.now(timezone.utc),
model=None,
)
except Exception as e:
verbose_proxy_logger.warning(
"Could not recover the partial usage of a failed %s pass-through stream: %s", endpoint_type.value, e
)
return
usage: Final = getattr(partial_response, "usage", None)
if not isinstance(usage, Usage):
return
response_cost: Final = kwargs.get("response_cost")
litellm_logging_obj.record_partial_usage_for_failure(
usage=usage,
response_cost=float(response_cost) if isinstance(response_cost, (int, float)) else 0.0,
)
@staticmethod
async def chunk_processor(
response: httpx.Response,
@ -65,13 +140,14 @@ class PassThroughStreamingHandler:
route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler
)
raw_bytes: Final[list[bytes]] = []
resolved_request_body: Final[dict[str, object]] = request_body or {}
def _build_logging_coroutine() -> Coroutine[None, None, None]:
return resolved_route_streaming_logging(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body or {},
request_body=resolved_request_body,
endpoint_type=endpoint_type,
start_time=start_time,
raw_bytes=raw_bytes,
@ -132,9 +208,9 @@ class PassThroughStreamingHandler:
# coroutine on logging_obj instead of enqueueing now, so
# ProxyLogging._fire_deferred_stream_logging fires it after
# guardrail end-of-stream blocks populate guardrail_information.
# Disconnect/exception paths skip this and fall through to the
# immediate enqueue in ``finally`` to keep partial billing
# (LIT-2642).
# Disconnect paths skip this and fall through to the immediate
# enqueue in ``finally`` to keep partial billing (LIT-2642);
# upstream exceptions log a failure instead (LIT-3798).
if (
getattr(litellm_logging_obj, "_on_deferred_stream_complete", None) is not None
and raw_bytes
@ -144,6 +220,20 @@ class PassThroughStreamingHandler:
litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),)
except Exception as e:
verbose_proxy_logger.error("Error in chunk_processor: %s", e)
if response.status_code < 400:
logging_scheduled = True
PassThroughStreamingHandler.schedule_stream_failure_logging(
litellm_logging_obj=litellm_logging_obj,
endpoint_type=endpoint_type,
request_body=resolved_request_body,
raw_bytes=raw_bytes,
exception=e,
stream_context=PassThroughStreamContext(
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
start_time=start_time,
),
)
raise
finally:
# GeneratorExit (raised on client disconnect) is not caught by
@ -168,7 +258,7 @@ class PassThroughStreamingHandler:
request_body: dict,
endpoint_type: EndpointType,
start_time: datetime,
raw_bytes: list[bytes],
raw_bytes: Sequence[bytes],
end_time: datetime,
model: str | None = None,
):
@ -218,7 +308,7 @@ class PassThroughStreamingHandler:
request_body: dict,
endpoint_type: EndpointType,
start_time: datetime,
raw_bytes: list[bytes],
raw_bytes: Sequence[bytes],
end_time: datetime,
model: str | None,
) -> tuple[PassThroughEndpointLoggingResultValues, dict]:
@ -336,7 +426,7 @@ class PassThroughStreamingHandler:
return None
@staticmethod
def _convert_raw_bytes_to_str_lines(raw_bytes: list[bytes]) -> list[str]:
def _convert_raw_bytes_to_str_lines(raw_bytes: Sequence[bytes]) -> list[str]:
"""
Converts a list of raw bytes into a list of string lines, similar to aiter_lines()

108
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py
View file

@ -5,6 +5,7 @@ from datetime import datetime
import pytest
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.experimental_pass_through.messages import streaming_iterator as streaming_iterator_module
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
@ -32,7 +33,16 @@ class _RecordingLoggingIterator(BaseAnthropicMessagesStreamingIterator):
self.logging_call_count += 1
def _make_logging_obj(test_name: str) -> LiteLLMLoggingObj:
class _FailureRecorder(CustomLogger):
def __init__(self):
super().__init__()
self.failure_kwargs: list = []
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
self.failure_kwargs.append(kwargs)
def _make_logging_obj(test_name: str, failure_recorder: _FailureRecorder | None = None) -> LiteLLMLoggingObj:
return LiteLLMLoggingObj(
model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0",
messages=[{"role": "user", "content": "hi"}],
@ -41,9 +51,19 @@ def _make_logging_obj(test_name: str) -> LiteLLMLoggingObj:
start_time=datetime.now(),
litellm_call_id=test_name,
function_id=test_name,
dynamic_async_failure_callbacks=[failure_recorder] if failure_recorder is not None else None,
)
async def _wait_for_failure_event(recorder: _FailureRecorder) -> dict:
for _ in range(300):
if recorder.failure_kwargs:
break
await asyncio.sleep(0.01)
assert len(recorder.failure_kwargs) == 1, "expected exactly one failure event"
return recorder.failure_kwargs[0]
def _make_iterator(test_name: str) -> BaseAnthropicMessagesStreamingIterator:
return BaseAnthropicMessagesStreamingIterator(
litellm_logging_obj=_make_logging_obj(test_name),
@ -524,6 +544,25 @@ async def test_async_sse_wrapper_dispatches_deferred_logging_when_client_disconn
await asyncio.wait_for(deferred_fired.wait(), timeout=5)
class _DetachedFailureRecorder:
"""Stands in for the closure the proxy arms so a detached-stream failure still reaches its failure hook."""
def __init__(self):
self.exceptions = []
async def __call__(self, exc: Exception) -> None:
self.exceptions.append(exc)
async def _wait_for_detached_failure(recorder: _DetachedFailureRecorder) -> Exception:
for _ in range(200):
if recorder.exceptions:
await asyncio.sleep(0.02)
return recorder.exceptions[0]
await asyncio.sleep(0.01)
raise AssertionError("the detached failure hook never fired")
class _ProviderStreamError(Exception):
"""Stand-in for a provider-specific streaming failure carrying a status code."""
@ -539,7 +578,8 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client():
before message_stop must propagate the ORIGINAL provider exception to a
still-connected client, so the proxy's failure handling keeps the
provider-specific status. The pump must not swallow it into a generic
api_error event + normal termination.
api_error event + normal termination, and the request is logged as a
failure carrying the partial usage, never as a success.
"""
async def _failing_stream():
@ -547,10 +587,13 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client():
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}}
raise _ProviderStreamError("bedrock stream blew up", status_code=529)
recorder = _FailureRecorder()
iterator = _RecordingLoggingIterator(
litellm_logging_obj=_make_logging_obj("test_reraises_upstream_error"),
litellm_logging_obj=_make_logging_obj("test_reraises_upstream_error", recorder),
request_body={},
)
detached_hook = _DetachedFailureRecorder()
iterator.litellm_logging_obj._on_detached_stream_failure = detached_hook
received = []
@ -561,18 +604,25 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client():
with pytest.raises(_ProviderStreamError) as excinfo:
await _drain()
failure_kwargs = await _wait_for_failure_event(recorder)
assert excinfo.value.status_code == 529
assert received
assert not any(c.startswith(b"event: error\n") for c in received)
assert iterator.logged_chunks == []
assert failure_kwargs["standard_logging_object"]["status"] == "failure"
assert failure_kwargs["standard_logging_object"]["prompt_tokens"] == 52
await asyncio.sleep(0.05)
assert detached_hook.exceptions == [], "the relay re-raised the error, so the proxy failure hook already ran"
@pytest.mark.asyncio
async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_disconnect():
async def test_async_sse_wrapper_logs_failure_on_upstream_error_after_disconnect():
"""
When the upstream errors AFTER the client has already disconnected there is
no live client to re-raise to and no failure hook will run, so the pump
salvages partial spend from what it collected instead of dropping the row.
no live client to re-raise to and no proxy failure hook will run, so the
pump logs the failure itself with the partial usage it collected; it must
never bill the broken stream as a success.
"""
tail_gated = asyncio.Event()
@ -582,34 +632,38 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_
await tail_gated.wait()
raise _ProviderStreamError("late failure", status_code=500)
recorder = _FailureRecorder()
iterator = _RecordingLoggingIterator(
litellm_logging_obj=_make_logging_obj("test_salvage_partial_on_late_error"),
litellm_logging_obj=_make_logging_obj("test_failure_logged_on_late_error", recorder),
request_body={},
)
detached_hook = _DetachedFailureRecorder()
iterator.litellm_logging_obj._on_detached_stream_failure = detached_hook
gen = iterator.async_sse_wrapper(_gated_failing_stream())
received = [await gen.__anext__(), await gen.__anext__()]
await gen.aclose() # client disconnects before the upstream error
tail_gated.set() # let the upstream raise now, after disconnect
for _ in range(100):
if iterator.logged_chunks:
break
await asyncio.sleep(0.01)
failure_kwargs = await _wait_for_failure_event(recorder)
assert len(received) == 2
assert iterator.logged_chunks == received
assert iterator.logging_call_count == 0
assert failure_kwargs["standard_logging_object"]["status"] == "failure"
assert failure_kwargs["standard_logging_object"]["prompt_tokens"] == 52
assert isinstance(failure_kwargs["exception"], _ProviderStreamError)
assert await _wait_for_detached_failure(detached_hook) is failure_kwargs["exception"]
assert len(detached_hook.exceptions) == 1
@pytest.mark.asyncio
async def test_async_sse_wrapper_salvages_spend_when_queued_error_is_never_consumed():
async def test_async_sse_wrapper_logs_failure_when_queued_error_is_never_consumed():
"""
When the upstream errors while the client is still connected, the pump
forwards the exception through the queue expecting the relay to re-raise it
into the proxy's failure handling. If the client disconnects before
consuming that queued exception, the handoff never happens and no failure
hook runs, so the pump must notice the unconsumed exception at teardown and
salvage partial spend instead of dropping the row entirely.
forwards the exception through the queue for the relay to re-raise. If the
client disconnects before consuming that queued exception, no proxy failure
hook runs, so the failure logged by the pump itself is the only record of
the request; it must be a failure row, not a salvaged success.
"""
upstream_errored = asyncio.Event()
@ -619,23 +673,27 @@ async def test_async_sse_wrapper_salvages_spend_when_queued_error_is_never_consu
upstream_errored.set()
raise _ProviderStreamError("mid-stream failure", status_code=500)
recorder = _FailureRecorder()
iterator = _RecordingLoggingIterator(
litellm_logging_obj=_make_logging_obj("test_salvage_on_unconsumed_queued_error"),
litellm_logging_obj=_make_logging_obj("test_failure_logged_on_unconsumed_queued_error", recorder),
request_body={},
)
detached_hook = _DetachedFailureRecorder()
iterator.litellm_logging_obj._on_detached_stream_failure = detached_hook
gen = iterator.async_sse_wrapper(_failing_stream())
received = [await gen.__anext__(), await gen.__anext__()]
await upstream_errored.wait() # exception is now queued behind the consumed chunks
await gen.aclose() # client disconnects without ever consuming the queued exception
for _ in range(100):
if iterator.logged_chunks:
break
await asyncio.sleep(0.01)
failure_kwargs = await _wait_for_failure_event(recorder)
assert iterator.logging_call_count == 1
assert iterator.logged_chunks == received
assert len(received) == 2
assert iterator.logging_call_count == 0
assert failure_kwargs["standard_logging_object"]["status"] == "failure"
assert failure_kwargs["standard_logging_object"]["prompt_tokens"] == 52
assert await _wait_for_detached_failure(detached_hook) is failure_kwargs["exception"]
assert len(detached_hook.exceptions) == 1
@pytest.mark.asyncio

167
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -3647,6 +3647,173 @@ async def test_stash_applies_when_owner_or_callback_call_id_missing():
assert claimed.reservation_released is True
async def _reserve_tpm_for_owner_call(handler, local_cache, api_key: str, call_id: str) -> int:
await handler.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key=api_key, tpm_limit=10_000),
cache=local_cache,
data={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 50,
"litellm_call_id": call_id,
},
call_type="completion",
)
stash = get_request_stash()
assert stash is not None and stash.reserved_tokens > 0
return stash.reserved_tokens
@pytest.mark.asyncio
async def test_failure_event_settles_tpm_reservation_at_recovered_partial_usage_v3():
"""
A stream that fails mid-way after the model already produced tokens is
logged as a failure carrying the recovered partial usage. Those tokens
were consumed, so the TPM window must settle at them instead of refunding
the whole reservation (which would let repeated timeouts burn output
tokens for free).
"""
_api_key = hash_token("sk-partial-stream-failure")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
await _reserve_tpm_for_owner_call(handler, local_cache, _api_key, "partial-call")
await handler.async_log_failure_event(
kwargs={
"litellm_call_id": "partial-call",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 27
stash = get_request_stash()
assert stash is not None and stash.reservation_released is True
@pytest.mark.asyncio
async def test_failure_event_refunds_reservation_for_input_only_estimate_v3():
"""
A failure with no recovered output carries only the input-token estimate
the proxy lifts onto every failure; that is not consumed usage, so the
reservation is still refunded in full.
"""
_api_key = hash_token("sk-estimated-failure")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
await _reserve_tpm_for_owner_call(handler, local_cache, _api_key, "estimate-call")
await handler.async_log_failure_event(
kwargs={
"litellm_call_id": "estimate-call",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0
@pytest.mark.asyncio
async def test_post_call_failure_hook_settles_reservation_at_recovered_partial_usage_v3():
"""
Pass-through streams report a mid-stream failure through the proxy-level
failure hook first, with the recovered usage lifted onto request_data.
That hook must settle at the partial usage too, and the later failure
callback must not double-apply it.
"""
_api_key = hash_token("sk-partial-post-call")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000)
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
await _reserve_tpm_for_owner_call(handler, local_cache, _api_key, "post-call")
await handler.async_post_call_failure_hook(
request_data={
"model": "gpt-4o-mini",
"litellm_call_id": "post-call",
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27),
},
original_exception=Exception("upstream dropped the stream"),
user_api_key_dict=user_api_key_dict,
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 27
await handler.async_log_failure_event(
kwargs={
"litellm_call_id": "post-call",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 27
@pytest.mark.asyncio
async def test_failure_event_settles_project_itpm_otpm_at_recovered_partial_usage_v3():
"""
Project ITPM/OTPM reservations settle the same way: input at the billable
prompt tokens and output at the completion tokens the failed stream
actually produced.
"""
_api_key = hash_token("sk-partial-project-io")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(
api_key=_api_key,
project_id="proj-partial",
project_metadata={
"model_itpm_limit": {"gpt-4o-mini": 10_000},
"model_otpm_limit": {"gpt-4o-mini": 10_000},
},
)
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 50,
"litellm_call_id": "project-call",
},
call_type="completion",
)
stash = get_request_stash()
assert stash is not None and stash.itpm_reserved_tokens > 0 and stash.otpm_reserved_tokens > 0
itpm_key = handler.create_rate_limit_keys(
key="model_per_project_itpm", value="proj-partial:gpt-4o-mini", rate_limit_type="tokens"
)
otpm_key = handler.create_rate_limit_keys(
key="model_per_project_otpm", value="proj-partial:gpt-4o-mini", rate_limit_type="tokens"
)
await handler.async_log_failure_event(
kwargs={
"litellm_call_id": "project-call",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert int(await local_cache.async_get_cache(key=itpm_key) or 0) == 20
assert int(await local_cache.async_get_cache(key=otpm_key) or 0) == 7
# ----------------------- Per-MCP-server rate limiting (v3) -----------------------

88
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
View file

@ -2441,3 +2441,91 @@ class TestAnthropicPassthroughFastMode:
assert served_standard.usage.speed == "standard"
assert self._cost(served_standard) == pytest.approx(self._cost(standard))
class TestRecordPartialUsageForFailure:
"""A stream that dies mid-way still carries the usage the provider billed in
message_start; the failure row must keep it and its cost instead of logging
a zero-cost failure (or, worse, a success)."""
@staticmethod
def _sse(event, data):
return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
@staticmethod
def _make_logging_obj() -> LiteLLMLoggingObj:
return LiteLLMLoggingObj(
model="claude-sonnet-5",
messages=[{"role": "user", "content": "hello"}],
stream=True,
call_type="anthropic_messages",
start_time=datetime.now(),
litellm_call_id="test-partial-usage-failure",
function_id="test-partial-usage-failure",
)
def _interrupted_chunks(self):
return [
self._sse(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_abc",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 52, "output_tokens": 1},
},
},
),
self._sse(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
),
self._sse(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}},
),
]
def test_stashes_partial_usage_and_cost_from_interrupted_stream(self):
logging_obj = self._make_logging_obj()
AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
litellm_logging_obj=logging_obj,
request_body={"model": "claude-sonnet-5", "stream": True},
all_chunks=self._interrupted_chunks(),
)
usage = logging_obj.model_call_details["combined_usage_object"]
assert usage.prompt_tokens == 52
assert logging_obj.model_call_details["response_cost"] > 0
def test_stashes_partial_usage_at_zero_cost_when_model_is_unpriced(self):
logging_obj = self._make_logging_obj()
AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
litellm_logging_obj=logging_obj,
request_body={"model": "claude-unpriced-test-model", "stream": True},
all_chunks=self._interrupted_chunks(),
)
usage = logging_obj.model_call_details["combined_usage_object"]
assert usage.prompt_tokens == 52
assert logging_obj.model_call_details["response_cost"] == 0.0
def test_leaves_logging_obj_untouched_when_nothing_streamed(self):
logging_obj = self._make_logging_obj()
AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
litellm_logging_obj=logging_obj,
request_body={"model": "claude-sonnet-5", "stream": True},
all_chunks=[],
)
assert "combined_usage_object" not in logging_obj.model_call_details
assert "response_cost" not in logging_obj.model_call_details

72
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -3988,6 +3988,78 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged(
assert failure_call_kwargs["original_exception"].status_code == 403
class _UpstreamDroppingMidStream(httpx.AsyncByteStream):
async def __aiter__(self):
yield b'data: {"id": "chatcmpl-1", "choices": [{"delta": {"content": "hi"}}]}\n\n'
raise httpx.ReadError("upstream dropped the connection mid-stream")
async def _relay_everything(body_iterator) -> list:
return [chunk async for chunk in body_iterator]
@pytest.mark.asyncio
async def test_pass_through_request_mid_stream_upstream_drop_fires_failure_hook():
"""
Regression: a 200 stream whose upstream dies mid-body used to end with no
proxy-level failure hook at all, so the request left no spend row, no
failure metric, and no alert; the pre-stream 4xx/5xx path already fires it.
"""
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
def transport_handler(upstream_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, stream=_UpstreamDroppingMidStream(), headers={"content-type": "text/event-stream"})
real_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.PassThroughEndpoint,
params={"timeout": resolve_pass_through_request_timeout(None)},
)
cache_dict = litellm.in_memory_llm_clients_cache.cache_dict
cache_key = next(key for key, cached in cache_dict.items() if cached is real_handler)
cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler)))
mock_proxy_logging = MagicMock()
mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
mock_proxy_logging.get_proxy_hook = MagicMock(return_value=None)
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.scope = {"path": "/relay-chat"}
mock_request.url = MagicMock()
mock_request.url.path = "/relay-chat"
mock_request.body = AsyncMock(return_value=b'{"model": "gpt-5.6", "stream": true}')
mock_request.headers = Headers({"content-type": "application/json"})
mock_request.query_params = QueryParams({})
try:
with patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request; there is no injection seam
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging
):
response = await pass_through_request(
request=mock_request,
target="http://target-api.com/v1/chat/completions",
custom_headers={},
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
stream=True,
)
with pytest.raises(httpx.ReadError):
await _relay_everything(response.body_iterator)
await asyncio.sleep(0)
finally:
cache_dict[cache_key] = real_handler
mock_proxy_logging.post_call_failure_hook.assert_awaited_once()
failure_call_kwargs = mock_proxy_logging.post_call_failure_hook.call_args.kwargs
assert isinstance(failure_call_kwargs["original_exception"], httpx.ReadError)
request_data = failure_call_kwargs["request_data"]
assert request_data["litellm_call_id"]
assert request_data["model"] == "gpt-5.6"
assert isinstance(request_data["litellm_logging_obj"], LiteLLMLoggingObj)
@pytest.mark.asyncio
async def test_pass_through_request_non_streaming_success_unchanged():
"""Success (2xx) passthrough behavior must remain unchanged by the error fix."""

206
tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
View file

@ -9,6 +9,8 @@ import httpx
import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.proxy.pass_through_endpoints.streaming_handler import (
PassThroughStreamingHandler,
@ -632,3 +634,207 @@ async def test_chunk_processor_enqueues_immediately_on_disconnect_even_when_arme
mock_enqueue.assert_called_once()
assert logging_obj._deferred_stream_complete_args is None
class _EventRecorder(CustomLogger):
def __init__(self):
super().__init__()
self.failure_kwargs = []
self.success_kwargs = []
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
self.failure_kwargs.append(kwargs)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.success_kwargs.append(kwargs)
def _anthropic_sse(event: str, payload: dict) -> bytes:
return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode()
def _anthropic_stream_that_times_out_mid_stream():
mock = MagicMock(spec=httpx.Response)
mock.status_code = 200
async def _aiter_bytes():
yield _anthropic_sse(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 52, "output_tokens": 1},
},
},
)
yield _anthropic_sse(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
)
yield _anthropic_sse(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}},
)
raise httpx.ReadTimeout("Timeout on reading data from socket")
mock.aiter_bytes = _aiter_bytes
return mock
@pytest.mark.asyncio
async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception():
"""A stream that dies after the first chunks is a failed request: the failure
callbacks must fire once with the partial usage and cost, and the success
routing must never run for it."""
recorder = _EventRecorder()
logging_obj = LiteLLMLoggingObj(
model="claude-sonnet-5",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="anthropic_messages",
start_time=datetime.now(),
litellm_call_id="test-mid-stream-timeout",
function_id="test-mid-stream-timeout",
dynamic_async_success_callbacks=[recorder],
dynamic_async_failure_callbacks=[recorder],
)
success_routes = []
async def _record_success_route(**kwargs):
success_routes.append(kwargs)
received = []
async def _consume_stream():
async for chunk in PassThroughStreamingHandler.chunk_processor(
response=_anthropic_stream_that_times_out_mid_stream(),
request_body={"model": "claude-sonnet-5", "stream": True},
litellm_logging_obj=logging_obj,
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
route_streaming_logging=_record_success_route,
):
received.append(chunk)
with pytest.raises(httpx.ReadTimeout):
await _consume_stream()
for _ in range(300):
if recorder.failure_kwargs:
break
await asyncio.sleep(0.01)
assert len(received) == 3
assert success_routes == []
assert recorder.success_kwargs == []
assert len(recorder.failure_kwargs) == 1
failure_payload = recorder.failure_kwargs[0]["standard_logging_object"]
assert failure_payload["status"] == "failure"
assert failure_payload["prompt_tokens"] == 52
assert failure_payload["response_cost"] > 0
assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout)
def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes:
payload = {
"candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}],
"usageMetadata": {
"promptTokenCount": prompt_tokens,
"candidatesTokenCount": completion_tokens,
"totalTokenCount": prompt_tokens + completion_tokens,
},
"modelVersion": "gemini-3.8-flash",
}
return f"data: {json.dumps(payload)}\r\n\r\n".encode()
def _google_stream_that_times_out_mid_stream():
mock = MagicMock(spec=httpx.Response)
mock.status_code = 200
async def _aiter_bytes():
yield _google_sse(9, 4, "The sea")
yield _google_sse(9, 12, " is wide and restless")
raise httpx.ReadTimeout("Timeout on reading data from socket")
mock.aiter_bytes = _aiter_bytes
return mock
@pytest.mark.parametrize(
"endpoint_type, url_route",
[
(EndpointType.GEMINI, "/gemini/v1beta/models/gemini-3.8-flash:streamGenerateContent?alt=sse"),
(
EndpointType.VERTEX_AI,
"/vertex_ai/v1/projects/p/locations/us-central1/publishers/google/models/gemini-3.8-flash:streamGenerateContent?alt=sse",
),
],
)
@pytest.mark.asyncio
async def test_chunk_processor_bills_partial_google_usage_on_mid_stream_exception(endpoint_type, url_route):
"""Google streams carry cumulative usage on every chunk, so a stream that
dies mid-way must log a failure billed at what was already delivered rather
than a failure at zero usage."""
recorder = _EventRecorder()
logging_obj = LiteLLMLoggingObj(
model="gemini-3.8-flash",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="pass_through_endpoint",
start_time=datetime.now(),
litellm_call_id=f"test-google-mid-stream-timeout-{endpoint_type.value}",
function_id="test-google-mid-stream-timeout",
dynamic_async_success_callbacks=[recorder],
dynamic_async_failure_callbacks=[recorder],
)
logging_obj.update_environment_variables(
model="gemini-3.8-flash",
user="unknown",
optional_params={},
litellm_params={"metadata": {}},
call_type="pass_through_endpoint",
)
success_routes = []
async def _record_success_route(**kwargs):
success_routes.append(kwargs)
async def _consume_stream():
async for _ in PassThroughStreamingHandler.chunk_processor(
response=_google_stream_that_times_out_mid_stream(),
request_body={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
litellm_logging_obj=logging_obj,
endpoint_type=endpoint_type,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route=url_route,
route_streaming_logging=_record_success_route,
):
pass
with pytest.raises(httpx.ReadTimeout):
await _consume_stream()
for _ in range(300):
if recorder.failure_kwargs:
break
await asyncio.sleep(0.01)
assert success_routes == []
assert recorder.success_kwargs == []
assert len(recorder.failure_kwargs) == 1
failure_payload = recorder.failure_kwargs[0]["standard_logging_object"]
assert failure_payload["status"] == "failure"
assert failure_payload["prompt_tokens"] == 9
assert failure_payload["completion_tokens"] == 12
assert failure_payload["response_cost"] > 12 * 3.75e-06
assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout)

109
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -7836,3 +7836,112 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_
records = [r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage()]
assert len(records) == 1
assert (records[0].exc_info is not None) is expect_traceback
class _FailureHookRecorder:
"""Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it."""
def __init__(self, raises: Optional[Exception] = None):
self.calls = []
self._raises = raises
async def post_call_failure_hook(self, **kwargs):
self.calls.append(kwargs)
if self._raises is not None:
raise self._raises
class TestDetachedStreamFailureHook:
"""
Regression for LIT-3798. A streaming /v1/messages request whose client disconnected
before the provider failed mid-stream never reached the proxy's failure hook: the
client-facing generator was gone, and the detached upstream drain only fired the
logging object's callbacks, so no failure spend row was written and the budget
reservation stayed held. base_process_llm_request now arms a closure on the logging
object that the detached drain awaits, and that closure runs post_call_failure_hook
with the request's key and data.
"""
@staticmethod
def _logging_obj():
logging_obj = MagicMock()
logging_obj.litellm_call_id = "call-lit3798"
logging_obj.model_call_details = {}
logging_obj._enqueue_deferred_logging = None
logging_obj._on_deferred_stream_complete = None
logging_obj._on_detached_stream_failure = None
return logging_obj
@staticmethod
def _proxy_logging_obj(recorder: _FailureHookRecorder):
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
proxy_logging_obj.update_request_status = AsyncMock(return_value=None)
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
proxy_logging_obj.post_call_failure_hook = recorder.post_call_failure_hook
return proxy_logging_obj
@pytest.mark.asyncio
async def test_streaming_messages_arms_the_detached_failure_hook(self, monkeypatch):
import litellm.proxy.common_request_processing as crp
from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth
async def _stream():
yield b"event: message_start\n\n"
async def fake_route_request(**kwargs):
async def _llm_call():
return _stream()
return _llm_call()
monkeypatch.setattr(crp, "route_request", fake_route_request)
monkeypatch.setattr(litellm, "callbacks", [])
recorder = _FailureHookRecorder()
logging_obj = self._logging_obj()
user_api_key_dict = RealUserAPIKeyAuth(api_key="sk-test")
processing_obj = ProxyBaseLLMRequestProcessing(
data={"litellm_logging_obj": logging_obj, "model": "claude-sonnet-4-5"}
)
await processing_obj.base_process_llm_request(
request=MagicMock(spec=Request, headers={}),
fastapi_response=Response(),
user_api_key_dict=user_api_key_dict,
route_type="anthropic_messages",
proxy_logging_obj=self._proxy_logging_obj(recorder),
general_settings={},
proxy_config=MagicMock(spec=ProxyConfig),
select_data_generator=None,
llm_router=None,
skip_pre_call_logic=True,
)
failure = RuntimeError("upstream died after the client left")
await logging_obj._on_detached_stream_failure(failure)
assert recorder.calls == [
{
"user_api_key_dict": user_api_key_dict,
"original_exception": failure,
"request_data": processing_obj.data,
}
]
@pytest.mark.asyncio
async def test_detached_failure_hook_drops_the_replacement_error_it_cannot_deliver(self):
from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth
recorder = _FailureHookRecorder(raises=HTTPException(status_code=429, detail="budget exceeded"))
logging_obj = self._logging_obj()
processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj})
processing_obj._arm_detached_stream_failure_hook(
logging_obj=logging_obj,
user_api_key_dict=RealUserAPIKeyAuth(api_key="sk-test"),
proxy_logging_obj=self._proxy_logging_obj(recorder),
)
failure = RuntimeError("upstream died after the client left")
await logging_obj._on_detached_stream_failure(failure)
assert [call["original_exception"] for call in recorder.calls] == [failure]