From cf7abf81367ccf800c561e855ddcb5bad3730a86 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:25:19 -0700 Subject: [PATCH 1/7] fix(proxy): log mid-stream /v1/messages failures as failures with partial usage A provider read timeout after the 200 was already committed on a streamed /v1/messages request used to run the success logging path, so the failure callbacks never fired and the failure metrics stayed flat. The pass-through stream handler and the Bedrock relay iterator now dispatch the failure handlers instead, with the usage and cost of the chunks already delivered stashed on the logging object so the failure row still bills them. --- litellm/litellm_core_utils/litellm_logging.py | 5 + .../messages/streaming_iterator.py | 49 ++-- .../anthropic_passthrough_logging_handler.py | 217 +++++++++++------- .../streaming_handler.py | 39 +++- .../messages/test_streaming_iterator.py | 77 +++++-- ...t_anthropic_passthrough_logging_handler.py | 75 ++++++ .../test_streaming_handler_interrupt.py | 109 +++++++++ 7 files changed, 431 insertions(+), 140 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f54eeca5178..e3941d6fbe1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1888,6 +1888,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, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 66e36dab2ba..34286e2171f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -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 @@ -176,17 +176,6 @@ def _try_claim_detached_drain_slot() -> bool: return True -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) - - def _sse_event(event_type: str, payload: Mapping[str, object]) -> bytes: return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() @@ -671,27 +660,25 @@ 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. + """Forward a provider error to a still-connected client and log the request as failed. - 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. + The relay re-raises the forwarded exception so the proxy's failure hook + keeps the provider status; the logging object's failure handlers fire + here either way, carrying the partial usage the provider already + billed, so a client that left before consuming the exception still + gets a failure row rather than a success one. """ - from litellm._logging import verbose_proxy_logger + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler - 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, + if not client_detached.is_set(): + await self._enqueue_for_client(queue, client_detached, exc) + PassThroughStreamingHandler.schedule_stream_failure_logging( + litellm_logging_obj=self.litellm_logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + request_body=self.request_body or {}, + raw_bytes=collected_chunks, + exception=exc, ) - await self._bill_collected_chunks(collected_chunks, stream_teardown=True) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index a36a365f39a..4ce840ce6f9 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -37,6 +37,7 @@ from litellm.types.utils import ( Message, ModelResponse, TextCompletionResponse, + Usage, ) if TYPE_CHECKING: @@ -147,6 +148,134 @@ 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: + 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: + 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 + try: + response_cost: Final = AnthropicPassthroughLoggingHandler._compute_response_cost( + litellm_model_response=partial_response, + model=AnthropicPassthroughLoggingHandler._resolve_costing_model(model, litellm_logging_obj), + logging_obj=litellm_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Anthropic passthrough: could not cost the partial usage of a failed stream (model=%s): %s", model, e + ) + return + litellm_logging_obj.record_partial_usage_for_failure(usage=usage, response_cost=response_cost) + + @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_model_from_anthropic_chunks( all_chunks: Sequence[str | bytes], @@ -263,31 +392,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 @@ -342,57 +449,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..." @@ -401,11 +463,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, diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 022a1ecbac4..ba2717ef119 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -1,4 +1,5 @@ -from collections.abc import Coroutine +import traceback +from collections.abc import Coroutine, Mapping, Sequence from datetime import datetime from typing import Final, Protocol @@ -50,6 +51,27 @@ class PassThroughStreamingHandler: 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: Mapping[str, object], + raw_bytes: Sequence[bytes], + exception: Exception, + ) -> 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 + ) + 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 async def chunk_processor( response: httpx.Response, @@ -132,9 +154,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 +166,15 @@ 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=request_body or {}, + raw_bytes=raw_bytes, + exception=e, + ) raise finally: # GeneratorExit (raised on client disconnect) is not caught by diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 11a048edc1f..3d41d0942e5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -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), @@ -539,7 +559,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,8 +568,9 @@ 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={}, ) @@ -561,18 +583,23 @@ 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 @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,8 +609,9 @@ 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={}, ) @@ -592,24 +620,23 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_ 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) @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,8 +646,9 @@ 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={}, ) @@ -629,13 +657,12 @@ async def test_async_sse_wrapper_salvages_spend_when_queued_error_is_never_consu 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 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 19bca05fb84..480da1d040e 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2441,3 +2441,78 @@ 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_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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 56c89fed79a..c4ae0c81d6e 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -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,110 @@ 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) From 501f47ba2fa59ad950c4e122fe9cfddf890ec106 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:31:35 -0700 Subject: [PATCH 2/7] fix(proxy): run the failure hook when a pass-through stream dies mid-body --- .../pass_through_endpoints.py | 84 +++++++++++++++---- .../test_pass_through_endpoints.py | 72 ++++++++++++++++ 2 files changed, 140 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 79d5d0a016f..37c8cfb09d6 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -868,6 +868,38 @@ 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]: + """An upstream that dies mid-stream leaves the client a truncated body and the proxy no record, so run + ``post_call_failure_hook`` (spend row, alerting, failure metric) the way the unified endpoints' generators do. + Error statuses were already reported by ``_log_passthrough_upstream_failure`` and relay untouched.""" + 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 @@ -1291,14 +1323,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, @@ -1372,14 +1414,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, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d3f17c73499..442adbca08c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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.""" From f3021937c62a24e743adbef004a1e45942951c77 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:29:49 -0700 Subject: [PATCH 3/7] refactor(proxy): drop the docstring restating the pass-through failure relay --- litellm/proxy/pass_through_endpoints/pass_through_endpoints.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 37c8cfb09d6..51a579b27cd 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -874,9 +874,6 @@ async def _relay_reporting_failures( user_api_key_dict: UserAPIKeyAuth, request_payload: dict, # mutable-ok: post_call_failure_hook lifts fields onto request_data in place ) -> AsyncGenerator[bytes, None]: - """An upstream that dies mid-stream leaves the client a truncated body and the proxy no record, so run - ``post_call_failure_hook`` (spend row, alerting, failure metric) the way the unified endpoints' generators do. - Error statuses were already reported by ``_log_passthrough_upstream_failure`` and relay untouched.""" from litellm.proxy.proxy_server import proxy_logging_obj try: From 5acb81888d6f62108194b57beccbb83c3e302ee6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:52:58 -0700 Subject: [PATCH 4/7] fix(proxy): settle rate-limit reservations at a failed stream's partial usage --- .../hooks/parallel_request_limiter_v3.py | 58 +++--- .../hooks/test_parallel_request_limiter_v3.py | 167 ++++++++++++++++++ 2 files changed, 204 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 63129602082..31437af7770 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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 diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index fc0088b28d7..4003286d887 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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) ----------------------- From 0f759c56f002b511be497b7769041d03c791cee2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:00:08 -0700 Subject: [PATCH 5/7] fix(proxy): bill partial usage on failed Vertex and Gemini pass-through streams --- .../streaming_handler.py | 70 +++++++++++-- .../test_streaming_handler_interrupt.py | 97 +++++++++++++++++++ 2 files changed, 161 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ba2717ef119..88c14c9348c 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -1,6 +1,7 @@ import traceback from collections.abc import Coroutine, Mapping, Sequence -from datetime import datetime +from dataclasses import dataclass +from datetime import datetime, timezone from typing import Final, Protocol import httpx @@ -13,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, @@ -45,6 +46,13 @@ 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: @@ -58,11 +66,15 @@ class PassThroughStreamingHandler: request_body: Mapping[str, object], raw_bytes: Sequence[bytes], exception: Exception, + stream_context: PassThroughStreamContext | None = 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 - ) + 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( @@ -72,6 +84,47 @@ class PassThroughStreamingHandler: 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: Mapping[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=dict(request_body), + endpoint_type=endpoint_type, + start_time=stream_context.start_time, + raw_bytes=list(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, @@ -174,6 +227,11 @@ class PassThroughStreamingHandler: request_body=request_body or {}, 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: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index c4ae0c81d6e..ea6adc35b9a 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -741,3 +741,100 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception( 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) From fde676dc386188f4acf8f991548cbfd61745d298 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:42:08 -0700 Subject: [PATCH 6/7] fix(anthropic): run the proxy failure hook when a detached /v1/messages stream fails --- litellm/litellm_core_utils/litellm_logging.py | 3 +- .../messages/streaming_iterator.py | 37 ++++-- litellm/proxy/common_request_processing.py | 33 ++++++ .../messages/test_streaming_iterator.py | 31 +++++ .../proxy/test_common_request_processing.py | 109 ++++++++++++++++++ 5 files changed, 204 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 463cbf7cdbe..d6f2387ac71 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -10,7 +10,7 @@ import subprocess import sys import time import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, 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``.""" diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 275608fcccc..7d01aee5d98 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -176,6 +176,11 @@ def _try_claim_detached_drain_slot() -> bool: return True +def _exception_left_unconsumed(queue: "asyncio.Queue[bytes | None | BaseException]", exc: BaseException) -> bool: + remaining: Final = tuple(queue.get_nowait() for _ in range(queue.qsize())) + return any(item is exc for item in remaining) + + def _sse_event(event_type: str, payload: Mapping[str, object]) -> bytes: return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() @@ -663,18 +668,16 @@ class BaseAnthropicMessagesStreamingIterator: collected_chunks: Sequence[bytes], exc: Exception, ) -> None: - """Forward a provider error to a still-connected client and log the request as failed. + """Log the request as failed with its partial usage, then make sure the proxy's failure hook runs once. - The relay re-raises the forwarded exception so the proxy's failure hook - keeps the provider status; the logging object's failure handlers fire - here either way, carrying the partial usage the provider already - billed, so a client that left before consuming the exception still - gets a failure row rather than a success one. + 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.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler - if not client_detached.is_set(): - await self._enqueue_for_client(queue, client_detached, exc) PassThroughStreamingHandler.schedule_stream_failure_logging( litellm_logging_obj=self.litellm_logging_obj, endpoint_type=EndpointType.ANTHROPIC, @@ -682,3 +685,21 @@ class BaseAnthropicMessagesStreamingIterator: 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 + 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 + ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6542842f5e4..f25fa46197e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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. diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 3d41d0942e5..be33b2ee3b1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -544,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.""" @@ -573,6 +592,8 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client(): 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 = [] @@ -591,6 +612,8 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client(): 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 @@ -614,6 +637,8 @@ async def test_async_sse_wrapper_logs_failure_on_upstream_error_after_disconnect 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__()] @@ -627,6 +652,8 @@ async def test_async_sse_wrapper_logs_failure_on_upstream_error_after_disconnect 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 @@ -651,6 +678,8 @@ async def test_async_sse_wrapper_logs_failure_when_queued_error_is_never_consume 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__()] @@ -663,6 +692,8 @@ async def test_async_sse_wrapper_logs_failure_when_queued_error_is_never_consume 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 diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index ea665b60b19..f7fe6ad9d39 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -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] From 6c27754455b384a267bae21923a4f50a28d9d6f6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 09:26:51 -0700 Subject: [PATCH 7/7] fix(anthropic): bill an uncostable partial pass-through stream at zero cost instead of dropping its usage --- .../anthropic_passthrough_logging_handler.py | 22 ++++++++++++++----- ...t_anthropic_passthrough_logging_handler.py | 13 +++++++++++ 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index dae52bab956..30b75a7b482 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -240,18 +240,28 @@ class AnthropicPassthroughLoggingHandler: 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: - response_cost: Final = AnthropicPassthroughLoggingHandler._compute_response_cost( + return AnthropicPassthroughLoggingHandler._compute_response_cost( litellm_model_response=partial_response, - model=AnthropicPassthroughLoggingHandler._resolve_costing_model(model, litellm_logging_obj), - logging_obj=litellm_logging_obj, + model=AnthropicPassthroughLoggingHandler._resolve_costing_model(model, logging_obj), + logging_obj=logging_obj, ) - except Exception as e: # noqa: BLE001 # an uncostable partial stream must still log as a failure + 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 - litellm_logging_obj.record_partial_usage_for_failure(usage=usage, response_cost=response_cost) + return 0.0 @staticmethod def _compute_response_cost( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 480da1d040e..d721be62efe 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2505,6 +2505,19 @@ class TestRecordPartialUsageForFailure: 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()