From b6e3242865934f5394d0432039b98c6bc260b470 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:12:50 -0700 Subject: [PATCH 01/13] fix(proxy): keep a retried /v1/messages stream's spend log by logging the superseded attempt as a failure A /v1/messages stream the router retries before any content shares one logging object across both attempts. The superseded attempt's end-of-stream logging ran as a success (a usage-only or placeholder result) and claimed the request's one-shot success log, so the recovered stream's spend row was either missing or written under the superseded attempt's message id. The passthrough logging path now logs an Anthropic stream whose frames carry an `event: error` as a failure with whatever usage the provider reported, the way a dropped stream already does, and the router closes the superseded stream before it opens the retry so that logging is queued first. --- .../streaming_handler.py | 31 +++++- litellm/router.py | 8 +- ...thropic_messages_pre_content_retry_wire.py | 12 ++- .../proxy/hooks/test_prompt_cache_observer.py | 4 + .../test_streaming_handler_interrupt.py | 96 +++++++++++++++++++ tests/unit/test_router/test_router.py | 24 +++++ 6 files changed, 167 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 8199a96ceab..8ffbd4a5f50 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -1,5 +1,5 @@ import traceback -from collections.abc import Coroutine, Sequence +from collections.abc import Coroutine, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from typing import Final, Protocol @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception 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 @@ -297,6 +298,7 @@ class PassThroughStreamingHandler: from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( is_message_stop_chunk, # pyright: ignore[reportPrivateUsage] # both native stream paths share terminal-event detection is_provider_error_chunk, # pyright: ignore[reportPrivateUsage] # provider errors must not become cache evidence + parse_anthropic_error_event, ) # Transport reads can split event names and JSON payloads. Recognize terminal @@ -312,7 +314,18 @@ class PassThroughStreamingHandler: and is_message_stop_chunk(complete_frames) and not is_provider_error_chunk(complete_frames) ) + provider_error_event: Final = ( + parse_anthropic_error_event(complete_frames) if endpoint_type == EndpointType.ANTHROPIC else None + ) try: + if provider_error_event is not None: + await PassThroughStreamingHandler._log_anthropic_error_frame_as_failure( + litellm_logging_obj=litellm_logging_obj, + request_body=request_body, + raw_bytes=raw_bytes, + error_event=provider_error_event, + ) + return # TinyFish billing is owned by the detached poller; the $0 fallback below is only for streams with no run_id if endpoint_type == EndpointType.TINYFISH: if sse_poller_spawned(litellm_logging_obj): @@ -373,6 +386,22 @@ class PassThroughStreamingHandler: except Exception as e: verbose_proxy_logger.error("Error in _route_streaming_logging_to_handler: %s", e) + @staticmethod + async def _log_anthropic_error_frame_as_failure( + litellm_logging_obj: LiteLLMLoggingObj, + request_body: Mapping[str, object], + raw_bytes: Sequence[bytes], + error_event: tuple[str, str, int], + ) -> None: + error_type, message, status_code = error_event + await asyncify(AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure)( + litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=raw_bytes + ) + frame_error: Final = anthropic_error_frame_exception( + error_type, message, status_code, str(litellm_logging_obj.model) + ) + await litellm_logging_obj.dispatch_failure_handlers(frame_error, "", prefer_async_handlers=True) + @staticmethod def _build_passthrough_logging_result( litellm_logging_obj: LiteLLMLoggingObj, diff --git a/litellm/router.py b/litellm/router.py index ab9d8884a1d..90f3950cc61 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5652,6 +5652,10 @@ class Router: source_iterator: Final = response + async def close_source() -> None: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(source_iterator) + async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]: from litellm.exceptions import MidStreamFallbackError @@ -5731,6 +5735,7 @@ class Router: for buffered_chunk in buffered_lifecycle_chunks: yield buffered_chunk except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate + await close_source() async for item in self._aanthropic_messages_recover_stream_error( stream_error, has_generated_content, @@ -5741,8 +5746,7 @@ class Router: ): yield item finally: - with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): - await aclose_if_supported(source_iterator) + await close_source() # Referenced by stream_with_fallbacks via closure - assigned here, before # the generator body ever runs, so the reference resolves fine despite diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index cb7f4c2b9ce..51d8680f014 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -124,12 +124,12 @@ def _stream(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[int, tuple return response.status_code, parse_sse(response.text) +def _rows(request_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)) + + def _success_rows(request_id: str) -> list[dict[str, JsonValue]]: - return eventually( - lambda: read_rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)), - lambda rows: len(rows) >= 1, - seconds=70, - ) + return eventually(lambda: _rows(request_id), lambda rows: len(rows) >= 1, seconds=70) def _assert_completed_by_retry( @@ -141,6 +141,8 @@ def _assert_completed_by_retry( assert delta_text(events) == _TEXT, events assert [request.target for request in wire.drain()] == ["/v1/messages"] * attempts assert _success_rows(_served_id(prompt, attempts)) == [{"status": "success", "model_group": model}] + superseded_rows: Final = {attempt: _rows(_served_id(prompt, attempt)) for attempt in range(1, attempts)} + assert not any(superseded_rows.values()), superseded_rows def _assert_rejected_before_content( diff --git a/tests/unit/proxy/hooks/test_prompt_cache_observer.py b/tests/unit/proxy/hooks/test_prompt_cache_observer.py index 82af3e9a6ef..5626431ad2a 100644 --- a/tests/unit/proxy/hooks/test_prompt_cache_observer.py +++ b/tests/unit/proxy/hooks/test_prompt_cache_observer.py @@ -204,6 +204,10 @@ class RecordingObserver(PromptCacheObserver): await super().async_log_success_event(kwargs, response_obj, start_time, end_time) self.finished.set() + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + await super().async_log_failure_event(kwargs, response_obj, start_time, end_time) + self.finished.set() + def native_response(): return { diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index e6b19f4eec1..0c53ae7e3ae 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -749,6 +749,102 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception( assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) +def _anthropic_message_start_frame() -> bytes: + return _anthropic_sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 52, "output_tokens": 1}, + }, + }, + ) + + +def _anthropic_overloaded_error_frame() -> bytes: + return _anthropic_sse("error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}) + + +def _anthropic_stream_of(frames: tuple[bytes, ...]): + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + + async def _aiter_bytes(): + for frame in frames: + yield frame + + mock.aiter_bytes = _aiter_bytes + return mock + + +@pytest.mark.parametrize( + "frames, billed", + [ + pytest.param( + (_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()), True, id="error-after-message-start" + ), + pytest.param((_anthropic_overloaded_error_frame(),), False, id="error-frame-only"), + ], +) +@pytest.mark.asyncio +async def test_chunk_processor_logs_a_provider_error_frame_as_a_failure_when_closed_at_the_frame(frames, billed): + """The router retries a /v1/messages stream whose provider sent `event: error` before any content by closing + it at the frame. That attempt never completed: it logs once as a failure carrying whatever usage the provider + had reported, and the request's one-shot success log stays free for the retry that shares the logging object.""" + 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-error-frame", + function_id="test-error-frame", + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + stream = PassThroughStreamingHandler.chunk_processor( + response=_anthropic_stream_of(frames), + 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", + ) + received = [] + async for chunk in stream: + received.append(chunk) + if b"event: error" in chunk: + break + await stream.aclose() + + for _ in range(300): + if recorder.failure_kwargs: + break + await asyncio.sleep(0.01) + + assert received == list(frames) + assert recorder.success_kwargs == [] + assert logging_obj.model_call_details.get("has_dispatched_final_stream_success") is not True + assert len(recorder.failure_kwargs) == 1 + failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] + assert failure_payload["status"] == "failure" + assert isinstance(recorder.failure_kwargs[0]["exception"], litellm.InternalServerError) + if billed: + assert failure_payload["prompt_tokens"] == 52 + assert failure_payload["response_cost"] > 0 + else: + assert failure_payload["response_cost"] == 0 + + def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes: payload = { "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 34b41a0bbff..85b9b49301a 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -15407,6 +15407,30 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] +@pytest.mark.asyncio +async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_opens(): + """Both attempts log through the request's one logging object, so the superseded attempt's end-of-stream + logging must be queued before the retry starts producing its own: the router closes the dropped stream + before it asks the group for another.""" + router = _anthropic_messages_retry_router(num_retries=1) + superseded = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] + ) + superseded_closed_when_retry_opened: list[bool] = [] + + def retry_stream(): + superseded_closed_when_retry_opened.append(superseded.closed) + return _anthropic_messages_retried_stream() + + provider = _AnthropicMessagesScriptedProvider(lambda: superseded, retry_stream) + + stream = await _anthropic_messages_stream_through_router(router, provider) + body = [chunk async for chunk in stream] + + assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] + assert superseded_closed_when_retry_opened == [True] + + @pytest.mark.asyncio async def test_anthropic_messages_stream_dropped_after_content_keeps_the_error_and_is_not_retried(): """A drop once content reached the client cannot be retried without a second overlapping message From 67ea7b649b77e2bae301dc4f5a6a19f70befdffc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 7 Oct 2026 13:10:54 -0700 Subject: [PATCH 02/13] fix(proxy): log a provider error frame through an attempt-scoped logging copy The error-frame failure stashed the superseded attempt's usage on the logging object the retry shares, so the retry's success log carried the failed attempt's token counts. The failure now logs through a copy of the logging object with its own model_call_details, and the helper moves into the Anthropic passthrough logging handler so the proxy streaming handler only routes by endpoint type. --- litellm/litellm_core_utils/litellm_logging.py | 7 ++ .../anthropic_passthrough_logging_handler.py | 23 ++++++ .../streaming_handler.py | 23 +----- ...t_anthropic_passthrough_logging_handler.py | 54 ++++++++++++- .../test_streaming_handler_interrupt.py | 79 ++++++++++++++++++- 5 files changed, 162 insertions(+), 24 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 2ce05b9aa39..7ba7399bbd4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2239,6 +2239,13 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def attempt_scoped_copy(self) -> "Logging": + """A copy with its own model_call_details, so logging one attempt of a request through it leaves + the state the request's other attempts share untouched.""" + attempt: Final = copy.copy(self) + attempt.model_call_details = dict(self.model_call_details) + return attempt + def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None: """Bill a fully streamed response on the failure log when a post-call hook rejects it.""" usage: Final = getattr(assembled, "usage", None) 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 dfb731972ac..1728e52bcd1 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 @@ -10,6 +10,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ANTHROPIC_BATCHES_ROUTE +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model @@ -21,6 +22,7 @@ from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -253,6 +255,27 @@ class AnthropicPassthroughLoggingHandler: ), ) + @staticmethod + async def log_error_frame_as_failure( + litellm_logging_obj: LiteLLMLoggingObj, + request_body: Mapping[str, object], + all_chunks: Sequence[str | bytes], + error_event: tuple[str, str, int], + ) -> None: + """A stream that ended in the provider's `event: error` frame is one failed attempt of the request: + it logs through a copy of the logging object so the usage it consumed never reaches the success log + of the attempt the router opens in its place.""" + error_type, message, status_code = error_event + attempt_logging_obj: Final = litellm_logging_obj.attempt_scoped_copy() + await asyncify(AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure)( + litellm_logging_obj=attempt_logging_obj, request_body=request_body, all_chunks=all_chunks + ) + await attempt_logging_obj.dispatch_failure_handlers( + anthropic_error_frame_exception(error_type, message, status_code, str(attempt_logging_obj.model)), + "", + prefer_async_handlers=True, + ) + @staticmethod def _cost_partial_stream_or_zero( partial_response: ModelResponse | TextCompletionResponse, model: str, logging_obj: LiteLLMLoggingObj diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 8ffbd4a5f50..3d2c414d3a9 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -1,5 +1,5 @@ import traceback -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import Coroutine, Sequence from dataclasses import dataclass from datetime import datetime, timezone from typing import Final, Protocol @@ -12,7 +12,6 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception 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 @@ -319,10 +318,10 @@ class PassThroughStreamingHandler: ) try: if provider_error_event is not None: - await PassThroughStreamingHandler._log_anthropic_error_frame_as_failure( + await AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( litellm_logging_obj=litellm_logging_obj, request_body=request_body, - raw_bytes=raw_bytes, + all_chunks=raw_bytes, error_event=provider_error_event, ) return @@ -386,22 +385,6 @@ class PassThroughStreamingHandler: except Exception as e: verbose_proxy_logger.error("Error in _route_streaming_logging_to_handler: %s", e) - @staticmethod - async def _log_anthropic_error_frame_as_failure( - litellm_logging_obj: LiteLLMLoggingObj, - request_body: Mapping[str, object], - raw_bytes: Sequence[bytes], - error_event: tuple[str, str, int], - ) -> None: - error_type, message, status_code = error_event - await asyncify(AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure)( - litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=raw_bytes - ) - frame_error: Final = anthropic_error_frame_exception( - error_type, message, status_code, str(litellm_logging_obj.model) - ) - await litellm_logging_obj.dispatch_failure_handlers(frame_error, "", prefer_async_handlers=True) - @staticmethod def _build_passthrough_logging_result( litellm_logging_obj: LiteLLMLoggingObj, diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 1c9bd470e52..9c6eb6a6558 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - +import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( @@ -2535,6 +2535,58 @@ class TestRecordPartialUsageForFailure: assert "response_cost" not in logging_obj.model_call_details +class _FailureRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.failure_kwargs = [] + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self.failure_kwargs.append(kwargs) + + +class TestLogErrorFrameAsFailure: + """A stream the provider ended with `event: error` is one failed attempt: its failure log bills the usage + streamed before the frame, and the request's shared logging state stays as it was, so the attempt the + router opens next logs with nothing left behind.""" + + @pytest.mark.asyncio + async def test_logs_a_failure_with_the_usage_streamed_before_the_frame_and_leaves_the_request_state_alone(self): + recorder = _FailureRecorder() + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="test-error-frame-failure", + function_id="test-error-frame-failure", + dynamic_async_failure_callbacks=[recorder], + ) + chunks = [ + *TestRecordPartialUsageForFailure()._interrupted_chunks()[:1], + TestRecordPartialUsageForFailure._sse( + "error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}} + ), + ] + + await AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( + litellm_logging_obj=logging_obj, + request_body={"model": "claude-sonnet-5", "stream": True}, + all_chunks=chunks, + error_event=("overloaded_error", "Overloaded", 529), + ) + + 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"], litellm.InternalServerError) + assert "combined_usage_object" not in logging_obj.model_call_details + assert "response_cost" not in logging_obj.model_call_details + assert logging_obj.should_run_logging("async_failure") + + @pytest.mark.parametrize( ("all_chunks", "interrupted"), [ diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 0c53ae7e3ae..7fdaf0f6812 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -749,25 +749,51 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception( assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) -def _anthropic_message_start_frame() -> bytes: +def _anthropic_message_start_frame(message_id: str = "msg_1", input_tokens: int = 52) -> bytes: return _anthropic_sse( "message_start", { "type": "message_start", "message": { - "id": "msg_1", + "id": message_id, "type": "message", "role": "assistant", "model": "claude-sonnet-5", "content": [], "stop_reason": None, "stop_sequence": None, - "usage": {"input_tokens": 52, "output_tokens": 1}, + "usage": {"input_tokens": input_tokens, "output_tokens": 1}, }, }, ) +def _anthropic_completed_stream_tail() -> tuple[bytes, ...]: + return ( + _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}}, + ), + _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 3}}, + ), + _anthropic_sse("message_stop", {"type": "message_stop"}), + ) + + +async def _settle(recorded: list) -> None: + for _ in range(300): + if recorded: + return + await asyncio.sleep(0.01) + + def _anthropic_overloaded_error_frame() -> bytes: return _anthropic_sse("error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}) @@ -845,6 +871,53 @@ async def test_chunk_processor_logs_a_provider_error_frame_as_a_failure_when_clo assert failure_payload["response_cost"] == 0 +@pytest.mark.asyncio +async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_frame(): + """Both attempts of a retried /v1/messages stream log through the request's one logging object. The + superseded attempt's failure log must not leave the usage it consumed behind, or the retry's success + log bills the failed attempt's tokens instead of its own.""" + 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-retry-own-usage", + function_id="test-retry-own-usage", + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + + def attempt(frames: tuple[bytes, ...]): + return PassThroughStreamingHandler.chunk_processor( + response=_anthropic_stream_of(frames), + 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", + ) + + superseded = attempt((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame())) + async for chunk in superseded: + if b"event: error" in chunk: + break + await superseded.aclose() + await _settle(recorder.failure_kwargs) + retry = attempt((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) + async for _ in retry: + pass + await _settle(recorder.success_kwargs) + + failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] + assert failure_payload["prompt_tokens"] == 52 + success_payload = recorder.success_kwargs[0]["standard_logging_object"] + assert success_payload["status"] == "success" + assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) + + def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes: payload = { "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], From 6dac86f731bb4966973018dfa8c78c7dee84f70c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:40:02 -0700 Subject: [PATCH 03/13] fix(proxy): scope the error-frame failure log to streams closed at the frame A stream that ran to its end at the provider's error frame kept logging through the success route before this PR, and the attempt-scoped failure log is now taken only for a stream the router closed at that frame, with the copy made at the close so the retry's deployment never reaches it --- .../anthropic_passthrough_logging_handler.py | 9 +- .../streaming_handler.py | 47 +++++++--- ...thropic_messages_pre_content_retry_wire.py | 20 +++++ ...t_anthropic_passthrough_logging_handler.py | 8 +- .../test_streaming_handler_interrupt.py | 85 +++++++++++++++++++ 5 files changed, 147 insertions(+), 22 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 1728e52bcd1..d9ef333a345 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 @@ -257,16 +257,15 @@ class AnthropicPassthroughLoggingHandler: @staticmethod async def log_error_frame_as_failure( - litellm_logging_obj: LiteLLMLoggingObj, + attempt_logging_obj: LiteLLMLoggingObj, request_body: Mapping[str, object], all_chunks: Sequence[str | bytes], error_event: tuple[str, str, int], ) -> None: - """A stream that ended in the provider's `event: error` frame is one failed attempt of the request: - it logs through a copy of the logging object so the usage it consumed never reaches the success log - of the attempt the router opens in its place.""" + """A stream closed at the provider's `event: error` frame is one failed attempt of the request: it logs + through the attempt-scoped copy its caller took at the close, so the usage it consumed never reaches + the success log of the attempt the router opens in its place.""" error_type, message, status_code = error_event - attempt_logging_obj: Final = litellm_logging_obj.attempt_scoped_copy() await asyncify(AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure)( litellm_logging_obj=attempt_logging_obj, request_body=request_body, all_chunks=all_chunks ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 3d2c414d3a9..3e0f9545174 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -165,6 +165,7 @@ class PassThroughStreamingHandler: ) logging_scheduled = False + stream_ended = False # rebind-ok: flipped once the upstream ran to its end, so the finally can tell a close-at-frame from a completed stream model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection( request_body=request_body, url_route=url_route, @@ -224,6 +225,7 @@ class PassThroughStreamingHandler: ) if pending: yield pending + stream_ended = True # Stream completed cleanly. When the proxy armed deferred # dispatch (post-call guardrails active), park the logging # coroutine on logging_obj instead of enqueueing now, so @@ -267,13 +269,44 @@ class PassThroughStreamingHandler: # a success would double-log the same request. if not logging_scheduled and raw_bytes and response.status_code < 400: logging_scheduled = True + superseded_error_event: Final = ( + None + if stream_ended + else PassThroughStreamingHandler._error_frame_closed_at(endpoint_type, raw_bytes) + ) try: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine()) + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + async_coroutine=( + _build_logging_coroutine() + if superseded_error_event is None + else AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( + attempt_logging_obj=litellm_logging_obj.attempt_scoped_copy(), + request_body=resolved_request_body, + all_chunks=tuple(raw_bytes), + error_event=superseded_error_event, + ) + ) + ) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) else: bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params) + @staticmethod + def _error_frame_closed_at(endpoint_type: EndpointType, raw_bytes: Sequence[bytes]) -> tuple[str, str, int] | None: + """The provider `event: error` frame a stream was closed at before it ended: the router closes a + /v1/messages stream there when a retry or fallback takes the request over, and that attempt is logged + as its own failure, on a copy taken now so the retry's rebinding of the shared logging object never + reaches it. A stream that ran to its end, error frame or not, keeps logging through the success route.""" + if endpoint_type != EndpointType.ANTHROPIC: + return None + from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( + parse_anthropic_error_event, + ) + + complete_frames, _ = split_complete_sse_frames(b"".join(raw_bytes)) + return parse_anthropic_error_event(complete_frames) + @staticmethod async def _route_streaming_logging_to_handler( litellm_logging_obj: LiteLLMLoggingObj, @@ -297,7 +330,6 @@ class PassThroughStreamingHandler: from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( is_message_stop_chunk, # pyright: ignore[reportPrivateUsage] # both native stream paths share terminal-event detection is_provider_error_chunk, # pyright: ignore[reportPrivateUsage] # provider errors must not become cache evidence - parse_anthropic_error_event, ) # Transport reads can split event names and JSON payloads. Recognize terminal @@ -313,18 +345,7 @@ class PassThroughStreamingHandler: and is_message_stop_chunk(complete_frames) and not is_provider_error_chunk(complete_frames) ) - provider_error_event: Final = ( - parse_anthropic_error_event(complete_frames) if endpoint_type == EndpointType.ANTHROPIC else None - ) try: - if provider_error_event is not None: - await AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( - litellm_logging_obj=litellm_logging_obj, - request_body=request_body, - all_chunks=raw_bytes, - error_event=provider_error_event, - ) - return # TinyFish billing is owned by the detached poller; the $0 fallback below is only for streams with no run_id if endpoint_type == EndpointType.TINYFISH: if sse_poller_spawned(litellm_logging_obj): diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index 51d8680f014..e0b2e4a336f 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -40,6 +40,7 @@ FirstAttempt = Literal[ "drop_after_message_start", "error_frame", "error_frame_after_message_start", + "error_frame_after_content", "http_status", "drop_after_content", ] @@ -66,6 +67,8 @@ class _Failure: return status_reply(self.status) case "drop_after_content": return dropping_reply(chunks, abort_after=3) + case "error_frame_after_content": + return stream_reply((*chunks[:3], error_frame(self.status, f"scripted {self.status}"))) _MID_STREAM_FAILURES: Final = ( @@ -304,6 +307,23 @@ def test_stream_dropping_after_content_is_not_retried(gateway: Gateway) -> None: assert [request.target for request in wire.drain()] == ["/v1/messages"] +def test_stream_error_frame_after_content_is_not_retried_and_keeps_its_spend_row(gateway: Gateway) -> None: + prompt: Final = _prompt() + attempts: Final = Attempts() + with ( + wire_server(_upstream(prompt, _Failure("error_frame_after_content", 529), attempts)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + status, events = _stream(gateway, _body(model, prompt, "deployment")) + assert status == 200, events + assert event_types(events) == (*LIFECYCLE[:3], "error"), events + assert message_id(events) == _served_id(prompt, 1), events + assert delta_text(events) == _TEXT, events + assert [request.target for request in wire.drain()] == ["/v1/messages"] + assert [row["model_group"] for row in _success_rows(_served_id(prompt, 1))] == [model] + + def test_stream_rejected_with_401_before_it_opens_is_not_retried(gateway: Gateway) -> None: prompt: Final = _prompt() attempts: Final = Attempts() diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 9c6eb6a6558..197c0d43cd7 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2545,9 +2545,9 @@ class _FailureRecorder(CustomLogger): class TestLogErrorFrameAsFailure: - """A stream the provider ended with `event: error` is one failed attempt: its failure log bills the usage - streamed before the frame, and the request's shared logging state stays as it was, so the attempt the - router opens next logs with nothing left behind.""" + """A stream closed at the provider's `event: error` frame is one failed attempt: its failure log bills the + usage streamed before the frame on the attempt-scoped copy its caller hands it, and the request's shared + logging state stays as it was, so the attempt the router opens next logs with nothing left behind.""" @pytest.mark.asyncio async def test_logs_a_failure_with_the_usage_streamed_before_the_frame_and_leaves_the_request_state_alone(self): @@ -2570,7 +2570,7 @@ class TestLogErrorFrameAsFailure: ] await AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( - litellm_logging_obj=logging_obj, + attempt_logging_obj=logging_obj.attempt_scoped_copy(), request_body={"model": "claude-sonnet-5", "stream": True}, all_chunks=chunks, error_event=("overloaded_error", "Overloaded", 529), diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 7fdaf0f6812..4e20f23ae63 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -860,6 +860,9 @@ async def test_chunk_processor_logs_a_provider_error_frame_as_a_failure_when_clo assert received == list(frames) assert recorder.success_kwargs == [] assert logging_obj.model_call_details.get("has_dispatched_final_stream_success") is not True + assert "combined_usage_object" not in logging_obj.model_call_details + assert "response_cost" not in logging_obj.model_call_details + assert logging_obj.should_run_logging("async_failure") assert len(recorder.failure_kwargs) == 1 failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] assert failure_payload["status"] == "failure" @@ -871,6 +874,88 @@ async def test_chunk_processor_logs_a_provider_error_frame_as_a_failure_when_clo assert failure_payload["response_cost"] == 0 +@pytest.mark.asyncio +async def test_chunk_processor_logs_a_stream_that_ended_at_the_provider_error_frame_through_the_success_route(): + """A provider error frame that ends a stream nobody closed early is the request's final answer, and that + stream keeps logging the way it always has: once, through the success route, with the usage it carried.""" + 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-final-error-frame", + function_id="test-final-error-frame", + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + frames = (_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()) + received = [ + chunk + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=_anthropic_stream_of(frames), + 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", + ) + ] + await _settle(recorder.success_kwargs) + + assert received == list(frames) + assert recorder.failure_kwargs == [] + success_payload = recorder.success_kwargs[0]["standard_logging_object"] + assert success_payload["status"] == "success" + assert success_payload["prompt_tokens"] == 52 + assert logging_obj.model_call_details["prompt_cache_response_complete"] is False + + +@pytest.mark.asyncio +async def test_superseded_attempt_failure_log_keeps_its_own_deployment_when_the_retry_rebinds_the_logging_object(): + """The router rebinds the shared logging object to the retry's deployment right after closing the superseded + stream, before the logging worker runs. The failure log of the superseded attempt must still name the + deployment that served it, or the router's cooldown lands on the deployment serving the retry.""" + 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-superseded-deployment", + function_id="test-superseded-deployment", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="claude-sonnet-5", litellm_params={"model_info": {"id": "deployment-1"}}, optional_params={} + ) + superseded = PassThroughStreamingHandler.chunk_processor( + response=_anthropic_stream_of((_anthropic_message_start_frame(), _anthropic_overloaded_error_frame())), + 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", + ) + async for chunk in superseded: + if b"event: error" in chunk: + break + await superseded.aclose() + logging_obj.update_environment_variables( + model="claude-sonnet-5-retry", litellm_params={"model_info": {"id": "deployment-2"}}, optional_params={} + ) + await _settle(recorder.failure_kwargs) + + failure_kwargs = recorder.failure_kwargs[0] + assert failure_kwargs["litellm_params"]["model_info"]["id"] == "deployment-1" + assert failure_kwargs["model"] == "claude-sonnet-5" + assert logging_obj.model_call_details["litellm_params"]["model_info"]["id"] == "deployment-2" + + @pytest.mark.asyncio async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_frame(): """Both attempts of a retried /v1/messages stream log through the request's one logging object. The From e2607da785c9ecc7d62d4c778d98039cc1231ede Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 7 Oct 2026 15:20:24 -0700 Subject: [PATCH 04/13] fix(router): log each dropped /v1/messages attempt through its own logging copy and open the retry without its usage A /v1/messages stream dropped after message_start and retried in its model group shares one Logging object across attempts, so the retry's success log inherited the dropped attempt's stashed usage and exception, and a second drop never reached the failure handlers because the first had already marked the object as logged. The dropped attempt now dispatches its failure handlers through an attempt-scoped copy, and the router discards the stashed usage right before it opens the next attempt, on same-group retries and on every fallback hop, so a request that runs out of attempts still bills the last attempt's partial usage. --- litellm/litellm_core_utils/litellm_logging.py | 6 + .../streaming_handler.py | 3 +- litellm/router.py | 2 + .../router_utils/fallback_event_handlers.py | 22 +++- ...thropic_messages_pre_content_retry_wire.py | 74 +++++++++++- .../test_streaming_handler_interrupt.py | 106 +++++++++++++++++- .../test_fallback_event_handlers.py | 78 +++++++++++++ tests/unit/test_router/test_router.py | 34 +++++- 8 files changed, 317 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7ba7399bbd4..d38b613f948 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2239,6 +2239,12 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def discard_partial_usage_for_failure(self) -> None: + """The router is about to open the next attempt of this request: the usage a superseded attempt + stashed here belongs to that attempt's own failure log, and the next attempt's log carries its own.""" + self.model_call_details.pop("combined_usage_object", None) + self.model_call_details.pop("response_cost", None) + def attempt_scoped_copy(self) -> "Logging": """A copy with its own model_call_details, so logging one attempt of a request through it leaves the state the request's other attempts share untouched.""" diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 3e0f9545174..9dc9367da1f 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -85,9 +85,10 @@ class PassThroughStreamingHandler: raw_bytes=raw_bytes, stream_context=stream_context, ) + attempt_logging_obj: Final = litellm_logging_obj.attempt_scoped_copy() try: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=litellm_logging_obj.dispatch_failure_handlers( + async_coroutine=attempt_logging_obj.dispatch_failure_handlers( exception, traceback.format_exc(), prefer_async_handlers=True ) ) diff --git a/litellm/router.py b/litellm/router.py index 90f3950cc61..712fe4f14c6 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -221,6 +221,7 @@ from litellm.router_utils.fallback_event_handlers import ( check_non_standard_fallback_format, clear_pre_routing_selection, committed_retry_budget_for_request, + discard_superseded_attempt_usage, fallback_lookup_groups, fallbacks_disabled_for_request, get_fallback_model_group_for_lookup_groups, @@ -5965,6 +5966,7 @@ class Router: verbose_router_logger.debug( "Retrying anthropic_messages stream dropped before content, attempt %s of %s", attempt + 1, budget ) + discard_superseded_attempt_usage(retry_kwargs) try: response = await self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt(**retry_kwargs) except Exception as retry_error: # noqa: BLE001 # every failure of a retry before its stream opens is the fallback chain's to judge diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index e33897e3b63..43dd7383888 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -12,6 +12,7 @@ from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, @@ -378,6 +379,15 @@ def mid_stream_retry_kwargs( return {key: value for key, value in hop_kwargs.items() if key not in _MID_STREAM_RETRY_STRIPPED_KEYS} +def discard_superseded_attempt_usage(kwargs: Mapping[str, object]) -> None: + """The next attempt of this request is about to open. A stream the previous attempt dropped stashed the + usage it consumed on the request's logging object for its own failure log and for the failure row of a + request that runs out of attempts; the next attempt's log must carry its own usage instead.""" + logging_obj: Final = kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.discard_partial_usage_for_failure() + + def _request_metadata_bucket(kwargs: Mapping[str, object]) -> Mapping[str, object] | None: bucket: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) return bucket if isinstance(bucket, Mapping) else None @@ -743,6 +753,12 @@ async def run_async_fallback( ) continue attempted.record(attempt_key) + logging_obj: Final = kwargs.get("litellm_logging_obj") + failure_callbacks_blocked_logging_obj: Final = ( + logging_obj + if logging_obj is not None and logging_obj.model_call_details.get("has_logged_async_failure", False) is True + else None + ) try: # LOGGING kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception) @@ -766,6 +782,7 @@ async def run_async_fallback( kwargs["attempted_targets"] = attempted if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors + discard_superseded_attempt_usage(kwargs) response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) verbose_router_logger.info("Successful fallback b/w models.") response = add_fallback_headers_to_response( @@ -788,13 +805,12 @@ async def run_async_fallback( kwargs=kwargs, original_exception=original_exception, ) - logging_obj = kwargs.get("litellm_logging_obj") - if logging_obj is not None and logging_obj.model_call_details.get("has_logged_async_failure", False): + if failure_callbacks_blocked_logging_obj is not None: _trigger_cooldown_for_failed_deployment( litellm_router=litellm_router, kwargs=kwargs, exception=e, - model_call_details=logging_obj.model_call_details, + model_call_details=failure_callbacks_blocked_logging_obj.model_call_details, ) raise error_from_fallbacks diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index e0b2e4a336f..6609b1753a8 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -1,13 +1,15 @@ import json import os import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass +from pathlib import Path from typing import Final, Literal import anthropic import httpx import pytest +import yaml from integration._support.anthropic_sse import ( LIFECYCLE, Attempts, @@ -27,6 +29,7 @@ from integration._support.anthropic_sse import ( ) from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows +from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue from redis import Redis @@ -463,6 +466,75 @@ def test_retried_stream_spend_row_records_the_attempt_count(gateway: Gateway) -> assert rows == [{"attempted": "1", "budget": "1"}], rows +def _proxy_with_callback_sink(gateway: Gateway, tmp_path: Path, sink: Wire) -> Iterator[Gateway]: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) + path: Final = tmp_path / "callbacks.yaml" + path.write_text(yaml.safe_dump(config)) + return owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": sink.url, "GENERIC_LOGGER_HEADERS": "Authorization=Bearer synthetic-sink-secret"}, + config=path, + ) + + +def _delivered_usage(sink: Wire, batches: list[Request], model: str) -> tuple[tuple[str, int, int], ...]: + batches.extend(sink.drain()) + events: Final = tuple(event for batch in batches for event in json.loads(batch.body)) + return tuple( + (str(event["status"]), int(event["prompt_tokens"]), int(event["completion_tokens"])) + for event in events + if event.get("model_group") == model + ) + + +def _request_rows(model: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model=%s', (model,)) + + +def test_dropped_attempt_and_its_retry_each_log_their_own_usage(gateway: Gateway, tmp_path: Path) -> None: + prompt: Final = _prompt() + attempts: Final = Attempts() + with ( + wire_server(_upstream(prompt, _Failure("drop_after_message_start"), attempts)) as wire, + wire_server(lambda request: Reply()) as sink, + _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + status, events = _stream(candidate, _body(model, prompt, "deployment")) + _assert_completed_by_retry(status, events, prompt, model, wire) + batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + delivered: Final = eventually( + lambda: _delivered_usage(sink, batches, model), lambda values: len(values) >= 2, seconds=20 + ) + assert sorted(delivered) == [("failure", 5, 1), ("success", 5, 3)], delivered + + +def test_every_dropped_attempt_logs_its_own_failure_and_the_request_row_still_bills_the_dropped_usage( + gateway: Gateway, tmp_path: Path +) -> None: + prompt: Final = _prompt() + attempts: Final = Attempts() + with ( + wire_server(_upstream(prompt, _Failure("drop_after_message_start"), attempts, failing_attempts=99)) as wire, + wire_server(lambda request: Reply()) as sink, + _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + status, events = _stream(candidate, _body(model, prompt, "deployment")) + _assert_stream_failed_after_message_start(status, events, wire, attempts=2) + batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + delivered: Final = eventually( + lambda: _delivered_usage(sink, batches, model), lambda values: len(values) >= 2, seconds=20 + ) + assert delivered == (("failure", 5, 1), ("failure", 5, 1)), delivered + rows: Final = eventually(lambda: _request_rows(model), lambda values: len(values) >= 1, seconds=70) + assert rows == [{"status": "failure", "prompt_tokens": 5, "completion_tokens": 1}], rows + + def _string_values(cache: Redis) -> tuple[bytes, ...]: keys: Final = tuple(key for key in cache.scan_iter(count=1000) if cache.type(key) == b"string") return tuple(value for value in cache.mget(keys) if value is not None) if keys else () diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 4e20f23ae63..8a27e4429bc 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -787,9 +787,9 @@ def _anthropic_completed_stream_tail() -> tuple[bytes, ...]: ) -async def _settle(recorded: list) -> None: +async def _settle(recorded: list, count: int = 1) -> None: for _ in range(300): - if recorded: + if len(recorded) >= count: return await asyncio.sleep(0.01) @@ -1003,6 +1003,108 @@ async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_fr assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) +def _anthropic_stream_dropping_after(frames: tuple[bytes, ...]): + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + + async def _aiter_bytes(): + for frame in frames: + yield frame + raise httpx.ReadError("upstream closed the connection") + + mock.aiter_bytes = _aiter_bytes + return mock + + +def _anthropic_messages_attempt(logging_obj: LiteLLMLoggingObj, response): + return PassThroughStreamingHandler.chunk_processor( + response=response, + 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", + ) + + +def _anthropic_messages_logging_obj(recorder: _EventRecorder, call_id: str) -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + + +async def _drop_after_message_start(logging_obj: LiteLLMLoggingObj, message_id: str, input_tokens: int) -> None: + dropped = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_dropping_after((_anthropic_message_start_frame(message_id, input_tokens),)) + ) + with pytest.raises(httpx.ReadError): + async for _ in dropped: + pass + + +@pytest.mark.asyncio +async def test_dropped_attempt_logs_its_failure_without_claiming_the_shared_logging_object(): + """The router retries a /v1/messages stream the provider drops after message_start on the same logging + object. The dropped attempt's failure log bills its own usage and leaves the shared object free: the + retry's success log carries the retry's usage and no stale exception, while the stash stays for the + failure row of a request that runs out of retries until the router opens the next attempt.""" + recorder = _EventRecorder() + logging_obj = _anthropic_messages_logging_obj(recorder, "test-dropped-attempt") + + await _drop_after_message_start(logging_obj, "msg_a1", 52) + await _settle(recorder.failure_kwargs) + + assert recorder.failure_kwargs[0]["standard_logging_object"]["prompt_tokens"] == 52 + assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadError) + assert logging_obj.model_call_details.get("has_logged_async_failure") is not True + assert "exception" not in logging_obj.model_call_details + stashed = logging_obj.model_call_details.get("combined_usage_object") + assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 52 + + logging_obj.discard_partial_usage_for_failure() + retry = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) + ) + async for _ in retry: + pass + await _settle(recorder.success_kwargs) + + success_payload = recorder.success_kwargs[0]["standard_logging_object"] + assert success_payload["status"] == "success" + assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) + assert recorder.success_kwargs[0].get("exception") is None + assert len(recorder.failure_kwargs) == 1 + + +@pytest.mark.asyncio +async def test_every_dropped_attempt_logs_its_own_failure(): + """Two attempts of one /v1/messages request dropped in a row each log a failure billing their own usage; + the first attempt's log must not swallow the second's, and the shared object ends up holding the last + attempt's usage for the request's failure row.""" + recorder = _EventRecorder() + logging_obj = _anthropic_messages_logging_obj(recorder, "test-every-dropped-attempt") + + await _drop_after_message_start(logging_obj, "msg_a1", 52) + await _settle(recorder.failure_kwargs) + logging_obj.discard_partial_usage_for_failure() + await _drop_after_message_start(logging_obj, "msg_a2", 9) + await _settle(recorder.failure_kwargs, count=2) + + assert [kwargs["standard_logging_object"]["prompt_tokens"] for kwargs in recorder.failure_kwargs] == [52, 9] + assert all(isinstance(kwargs["exception"], httpx.ReadError) for kwargs in recorder.failure_kwargs) + stashed = logging_obj.model_call_details.get("combined_usage_object") + assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 9 + + def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes: payload = { "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 732248b0550..430a796ab9d 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -9,6 +9,8 @@ import pytest import litellm from litellm.litellm_core_utils import get_llm_provider_logic +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.types.utils import Usage from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure from litellm.router_utils.fallback_event_handlers import( MID_STREAM_FALLBACK_CONTROLS_KEY, @@ -129,6 +131,82 @@ class RecordingRouter: return StreamingWrapper() +def _request_logging_obj() -> Logging: + return Logging( + model="primary-model", + messages=[{"role": "user", "content": "ping"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="fallback-hop", + function_id="fallback-hop", + ) + + +class StashRecordingRouter(FakeRouter): + def __init__(self) -> None: + self.stash_when_hop_opened: list[object] = [] + + async def async_function_with_fallbacks(self, *args, **kwargs): + self.stash_when_hop_opened.append(kwargs["litellm_logging_obj"].model_call_details.get("combined_usage_object")) + return StreamingWrapper() + + +@pytest.mark.asyncio +async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_partial_usage(): + """A stream the failed attempt dropped stashed the usage it consumed on the request's logging object. The + hop that replaces it logs its own usage, so the stash is gone by the time the hop opens.""" + router = StashRecordingRouter() + logging_obj = _request_logging_obj() + logging_obj.record_partial_usage_for_failure(Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53), 0.01) + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + ) + + assert router.stash_when_hop_opened == [None] + assert "response_cost" not in logging_obj.model_call_details + + +class HopLoggingItsOwnFailureRouter(AlwaysFailRouter): + async def async_function_with_fallbacks(self, *args, **kwargs): + kwargs["litellm_logging_obj"].model_call_details["has_logged_async_failure"] = True + raise RuntimeError("fallback model also failed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_already_logged", [False, True]) +async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_logged: bool): + """The manual cooldown stands in for the failure callbacks the logging object blocks once a failure has + been logged. A hop whose own failure was the first one logged already ran them, so only a hop that failed + with the flag already set gets the manual cooldown.""" + logging_obj = _request_logging_obj() + if failure_already_logged: + logging_obj.model_call_details["has_logged_async_failure"] = True + + with ( + patch("litellm.router_utils.fallback_event_handlers._trigger_cooldown_for_failed_deployment") as trigger, + pytest.raises(RuntimeError), + ): + await run_async_fallback( + litellm_router=HopLoggingItsOwnFailureRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + ) + + assert trigger.call_count == (1 if failure_already_logged else 0) + + @pytest.mark.asyncio async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_call(): """A nested fallback (multi-hop) must keep collecting errors, so the opt-in diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 85b9b49301a..4affe642ba2 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -76,7 +76,7 @@ from litellm.llms.base_llm.vector_store.transformation import( LiteLLMVectorStoreEmbeddingExecutor, RouterVectorStoreEmbeddingExecutor, ) -from litellm.types.utils import CallTypes, CredentialItem +from litellm.types.utils import CallTypes, CredentialItem, Usage from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -15431,6 +15431,38 @@ async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_o assert superseded_closed_when_retry_opened == [True] +@pytest.mark.asyncio +async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_partial_usage(): + """The dropped attempt stashes the usage it consumed on the request's logging object, for its own failure + log and for the failure row of a request that runs out of retries. The router discards it before the retry + opens, so a retry that completes bills its own usage instead of the dropped attempt's.""" + router = _anthropic_messages_retry_router(num_retries=1) + logging_obj = LiteLLMLogging( + model="glm", + messages=[{"role": "user", "content": "ping"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="retry-own-usage", + function_id="retry-own-usage", + ) + logging_obj.record_partial_usage_for_failure(Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53), 0.01) + stash_when_retry_opened: list[object] = [] + + def retry_stream(): + stash_when_retry_opened.append(logging_obj.model_call_details.get("combined_usage_object")) + return _anthropic_messages_retried_stream() + + provider = _AnthropicMessagesScriptedProvider(_anthropic_messages_dropped_before_content, retry_stream) + + stream = await _anthropic_messages_stream_through_router(router, provider, litellm_logging_obj=logging_obj) + body = [chunk async for chunk in stream] + + assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] + assert stash_when_retry_opened == [None] + assert "response_cost" not in logging_obj.model_call_details + + @pytest.mark.asyncio async def test_anthropic_messages_stream_dropped_after_content_keeps_the_error_and_is_not_retried(): """A drop once content reached the client cannot be retried without a second overlapping message From 8ad8db07c94ab7136457a26dff62de45a3eab703 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:59:40 -0700 Subject: [PATCH 05/13] fix(router): dispatch chat stream failures through the attempt-scoped logging copy The router now clears a dropped attempt's stashed usage before every fallback hop, so the chat stream wrapper's queued failure log must read that usage from its own copy of the logging object, like the /v1/messages path already does. The stash test asserts through a real success log, the cooldown test reads a real Router's cooldown list, and the integration helpers carry full types. --- .../litellm_core_utils/streaming_handler.py | 30 +++---- ...thropic_messages_pre_content_retry_wire.py | 13 ++- .../test_max_streaming_duration.py | 4 +- .../test_streaming_handler.py | 60 +++++++++++++ .../test_fallback_event_handlers.py | 89 ++++++++++++++----- 5 files changed, 154 insertions(+), 42 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 43611eb3192..d30272ea70f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2119,12 +2119,7 @@ class CustomStreamWrapper: traceback_exception = traceback.format_exc() ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT traceback_exception += f"\nLiteLLM Default Request Timeout - {litellm.request_timeout}" - if self.logging_obj is not None: - self._record_partial_usage_for_failure() - ## LOGGING - asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) - ) + self._dispatch_stream_failure(e, traceback_exception) self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: if self.received_finish_reason is None: @@ -2248,23 +2243,26 @@ class CustomStreamWrapper: return processed_chunk def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn: - traceback_exception: Final = traceback.format_exc() - if self.logging_obj is not None: - self._record_partial_usage_for_failure() - ## LOGGING - asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) - ) + self._dispatch_stream_failure(e, traceback.format_exc()) self._handle_stream_fallback_error(e) + def _dispatch_stream_failure(self, e: Exception, traceback_exception: str) -> None: + if self.logging_obj is None: + return + self._record_partial_usage_for_failure() + attempt_logging_obj: Final = self.logging_obj.attempt_scoped_copy() + asyncio.create_task( + attempt_logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) + ) + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks already delivered. Recover that partial usage from the chunks seen so far and stash it, with its cost, on the logging object so the failure - handler records the real partial spend instead of zero. A request that - later recovers via a router fallback overwrites this with the combined - success log on the same request id, so this never double counts. + handler records the real partial spend instead of zero. The router + discards the stash right before it opens the next attempt, so a request + that later recovers via a fallback logs that hop's own usage. """ if self.logging_obj is None or not self.chunks: return diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index 6609b1753a8..fadbf780980 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -1,7 +1,8 @@ import json import os import uuid -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager from dataclasses import dataclass from pathlib import Path from typing import Final, Literal @@ -466,7 +467,7 @@ def test_retried_stream_spend_row_records_the_attempt_count(gateway: Gateway) -> assert rows == [{"attempted": "1", "budget": "1"}], rows -def _proxy_with_callback_sink(gateway: Gateway, tmp_path: Path, sink: Wire) -> Iterator[Gateway]: +def _proxy_with_callback_sink(gateway: Gateway, tmp_path: Path, sink: Wire) -> AbstractContextManager[Gateway]: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) path: Final = tmp_path / "callbacks.yaml" @@ -479,12 +480,16 @@ def _proxy_with_callback_sink(gateway: Gateway, tmp_path: Path, sink: Wire) -> I ) +def _sink_events(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]: + for batch in batches: + yield from json.loads(batch.body) + + def _delivered_usage(sink: Wire, batches: list[Request], model: str) -> tuple[tuple[str, int, int], ...]: batches.extend(sink.drain()) - events: Final = tuple(event for batch in batches for event in json.loads(batch.body)) return tuple( (str(event["status"]), int(event["prompt_tokens"]), int(event["completion_tokens"])) - for event in events + for event in _sink_events(batches) if event.get("model_group") == model ) diff --git a/tests/unit/litellm_core_utils/test_max_streaming_duration.py b/tests/unit/litellm_core_utils/test_max_streaming_duration.py index fc8daab3899..c387c9bf103 100644 --- a/tests/unit/litellm_core_utils/test_max_streaming_duration.py +++ b/tests/unit/litellm_core_utils/test_max_streaming_duration.py @@ -70,8 +70,10 @@ class TestCustomStreamWrapperMaxDuration: """__anext__ should check the limit before iterating, dispatching the same failure-callback/logging path every other stream failure goes through (dispatch_failure_handlers is async on the real Logging class, - so the mock needs to be awaitable too).""" + so the mock needs to be awaitable too, and the wrapper dispatches it + through the attempt-scoped copy of the logging object).""" wrapper = _make_custom_stream_wrapper() + wrapper.logging_obj.attempt_scoped_copy.return_value = wrapper.logging_obj wrapper.logging_obj.dispatch_failure_handlers = AsyncMock() wrapper._stream_created_time = time.time() - 20 with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index d4c6a781498..c7278b2f720 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -5,6 +5,8 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest import asyncio + +import httpx import traceback from typing import Final, Optional @@ -18,6 +20,8 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkEarlyReturn, _ProviderChunkParsed, ) +from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import ( CompletionTokensDetailsWrapper, Delta, @@ -3478,6 +3482,62 @@ def test_record_partial_usage_for_failure_stashes_usage_and_cost(): assert isinstance(logging_obj.model_call_details["response_cost"], float) +class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: list[tuple[int, int]] = [] + self.logged = asyncio.Event() + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + payload = kwargs["standard_logging_object"] + self.failures.append((payload["prompt_tokens"], payload["completion_tokens"])) + self.logged.set() + + +@pytest.mark.asyncio +async def test_log_stream_failure_bills_the_dropped_attempt_even_after_the_router_discards_its_usage(): + """The dropped attempt's failure log is queued, and the router discards the stashed usage as it opens the + fallback hop before that log runs. The log must still bill what the dropped attempt consumed, and the shared + logging object must not be marked as already failed, since the hop's own failure still has to reach the + integrations.""" + recorder = _FailureRecorder() + logging_obj = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="partial-usage-copy", + function_id="1245", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + wrapper = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-partial-copy", + created=1742056047, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="The", role="assistant"))], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + + with pytest.raises(MidStreamFallbackError): + wrapper._log_stream_failure_and_raise(httpx.ReadError("connection dropped")) + logging_obj.discard_partial_usage_for_failure() + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + + assert recorder.failures == [(30, 1)] + assert logging_obj.model_call_details.get("has_logged_async_failure") is not True + + def test_record_partial_usage_for_failure_noop_without_chunks(): """With no chunks delivered there is nothing billed to recover, so the failure stash must stay absent and not force a zero-usage row. diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 430a796ab9d..e39c49d313d 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -9,9 +9,10 @@ import pytest import litellm from litellm.litellm_core_utils import get_llm_provider_logic +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.types.utils import Usage -from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure +from litellm.types.utils import Choices, Message, ModelResponse, Usage +from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments, mark_advisor_orchestration_failure from litellm.router_utils.fallback_event_handlers import( MID_STREAM_FALLBACK_CONTROLS_KEY, AttemptedFallbackTargets, @@ -143,25 +144,51 @@ def _request_logging_obj() -> Logging: ) -class StashRecordingRouter(FakeRouter): +class _SuccessRecorder(CustomLogger): def __init__(self) -> None: - self.stash_when_hop_opened: list[object] = [] + super().__init__() + self.logged: list[tuple[int, float]] = [] + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + payload = kwargs["standard_logging_object"] + self.logged.append((payload["completion_tokens"], payload["response_cost"])) + + +class HopLoggingItsOwnSuccessRouter(FakeRouter): async def async_function_with_fallbacks(self, *args, **kwargs): - self.stash_when_hop_opened.append(kwargs["litellm_logging_obj"].model_call_details.get("combined_usage_object")) + hop_response = ModelResponse( + id="chatcmpl-hop", + model="gpt-4o-mini", + choices=[Choices(index=0, message=Message(role="assistant", content="served by the hop"))], + usage=Usage(prompt_tokens=52, completion_tokens=7, total_tokens=59), + ) + now = datetime.now() + await kwargs["litellm_logging_obj"].async_success_handler(result=hop_response, start_time=now, end_time=now) return StreamingWrapper() @pytest.mark.asyncio async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_partial_usage(): """A stream the failed attempt dropped stashed the usage it consumed on the request's logging object. The - hop that replaces it logs its own usage, so the stash is gone by the time the hop opens.""" - router = StashRecordingRouter() - logging_obj = _request_logging_obj() + hop that replaces it logs its own usage, so the integrations see the hop's tokens and cost, not the stash.""" + recorder = _SuccessRecorder() + logging_obj = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "ping"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="fallback-hop-success", + function_id="fallback-hop-success", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + logging_obj.model_call_details["litellm_params"] = {"acompletion": True} + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] logging_obj.record_partial_usage_for_failure(Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53), 0.01) await run_async_fallback( - litellm_router=router, + litellm_router=HopLoggingItsOwnSuccessRouter(), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("stream dropped before content"), @@ -170,14 +197,29 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ litellm_logging_obj=logging_obj, ) - assert router.stash_when_hop_opened == [None] - assert "response_cost" not in logging_obj.model_call_details + assert [tokens for tokens, _ in recorder.logged] == [7] + assert recorder.logged[0][1] != 0.01 -class HopLoggingItsOwnFailureRouter(AlwaysFailRouter): - async def async_function_with_fallbacks(self, *args, **kwargs): - kwargs["litellm_logging_obj"].model_call_details["has_logged_async_failure"] = True - raise RuntimeError("fallback model also failed") +class _RateLimitedHop(Exception): + def __init__(self) -> None: + super().__init__("fallback deployment rate limited") + self.status_code = 429 + self.failed_deployment_id = "fallback-dep-1" + + +def _two_deployment_router() -> Router: + return Router( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"}, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + num_retries=0, + ) @pytest.mark.asyncio @@ -186,25 +228,30 @@ async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_l """The manual cooldown stands in for the failure callbacks the logging object blocks once a failure has been logged. A hop whose own failure was the first one logged already ran them, so only a hop that failed with the flag already set gets the manual cooldown.""" + router = _two_deployment_router() logging_obj = _request_logging_obj() if failure_already_logged: logging_obj.model_call_details["has_logged_async_failure"] = True - with ( - patch("litellm.router_utils.fallback_event_handlers._trigger_cooldown_for_failed_deployment") as trigger, - pytest.raises(RuntimeError), - ): + async def rate_limited_hop(*args, **kwargs): + kwargs["litellm_logging_obj"].model_call_details["has_logged_async_failure"] = True + raise _RateLimitedHop() + + with pytest.raises(_RateLimitedHop): await run_async_fallback( - litellm_router=HopLoggingItsOwnFailureRouter(), + litellm_router=router, fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("stream dropped before content"), max_fallbacks=3, fallback_depth=0, litellm_logging_obj=logging_obj, + original_function=rate_limited_hop, + metadata={}, ) - assert trigger.call_count == (1 if failure_already_logged else 0) + cooled_down = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + assert cooled_down == (["fallback-dep-1"] if failure_already_logged else []) @pytest.mark.asyncio From 52ba8a9f7c09ef6879f9052cff7213e88acaf1e4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:11:48 -0700 Subject: [PATCH 06/13] test(vertex): give the stream wrapper's logging stub an attempt_scoped_copy --- .../gemini/test_vertex_and_google_ai_studio_gemini.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 02fe48c0e38..563418c436f 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2214,6 +2214,9 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped(): async def dispatch_failure_handlers(self, *args, **kwargs): return None + def attempt_scoped_copy(self): + return self + async def failing_make_call(client=None, **kwargs): raise VertexAIError(status_code=400, message="bad input", headers={}) From fe3fe19ece6ec380a6b2d440a128f485e73fe1ef Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:51:49 -0700 Subject: [PATCH 07/13] test: route the pass-through logging mock's attempt copy to itself and read the callback sink through a reader The pass-through failure test stubbed dispatch_failure_handlers on a spec'd MagicMock, which the attempt-scoped copy no longer reached, so the logging worker never finished. The integration tests kept earlier sink batches in a caller-owned list the helper mutated; a small reader keeps them instead. --- ...thropic_messages_pre_content_retry_wire.py | 32 ++++++++++++------- .../test_streaming_handler.py | 1 + 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index fadbf780980..94df5984cdc 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -485,17 +485,27 @@ def _sink_events(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]: yield from json.loads(batch.body) -def _delivered_usage(sink: Wire, batches: list[Request], model: str) -> tuple[tuple[str, int, int], ...]: - batches.extend(sink.drain()) - return tuple( - (str(event["status"]), int(event["prompt_tokens"]), int(event["completion_tokens"])) - for event in _sink_events(batches) - if event.get("model_group") == model - ) +class _SinkReader: + """Reads the callback sink across polls: drain() consumes the sink's queue, so earlier batches are kept here.""" + + def __init__(self, sink: Wire, model: str) -> None: + self._sink: Final = sink + self._model: Final = model + self._batches: tuple[Request, ...] = () + + def delivered_usage(self) -> tuple[tuple[str, int, int], ...]: + self._batches = (*self._batches, *self._sink.drain()) + return tuple( + (str(event["status"]), int(event["prompt_tokens"]), int(event["completion_tokens"])) + for event in _sink_events(self._batches) + if event.get("model_group") == self._model + ) def _request_rows(model: str) -> list[dict[str, JsonValue]]: - return read_rows('SELECT status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model=%s', (model,)) + return read_rows( + 'SELECT status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model=%s', (model,) + ) def test_dropped_attempt_and_its_retry_each_log_their_own_usage(gateway: Gateway, tmp_path: Path) -> None: @@ -510,9 +520,8 @@ def test_dropped_attempt_and_its_retry_each_log_their_own_usage(gateway: Gateway model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) status, events = _stream(candidate, _body(model, prompt, "deployment")) _assert_completed_by_retry(status, events, prompt, model, wire) - batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches delivered: Final = eventually( - lambda: _delivered_usage(sink, batches, model), lambda values: len(values) >= 2, seconds=20 + _SinkReader(sink, model).delivered_usage, lambda values: len(values) >= 2, seconds=20 ) assert sorted(delivered) == [("failure", 5, 1), ("success", 5, 3)], delivered @@ -531,9 +540,8 @@ def test_every_dropped_attempt_logs_its_own_failure_and_the_request_row_still_bi model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) status, events = _stream(candidate, _body(model, prompt, "deployment")) _assert_stream_failed_after_message_start(status, events, wire, attempts=2) - batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches delivered: Final = eventually( - lambda: _delivered_usage(sink, batches, model), lambda values: len(values) >= 2, seconds=20 + _SinkReader(sink, model).delivered_usage, lambda values: len(values) >= 2, seconds=20 ) assert delivered == (("failure", 5, 1), ("failure", 5, 1)), delivered rows: Final = eventually(lambda: _request_rows(model), lambda values: len(values) >= 1, seconds=70) diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py index fdd21afbc67..b5e6c2a14f3 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py @@ -218,6 +218,7 @@ async def test_failed_anthropic_stream_records_partial_usage_off_the_event_loop( logging_obj.model_call_details = {"model": model, "stream": True} logging_obj.litellm_params = {} logging_obj.get_router_model_id.return_value = None + logging_obj.attempt_scoped_copy.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() _, took, lags = await timed_with_loop_lags( From aa5918c0cb01fa5099f46bfb8398ff248f27cd02 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:40:54 +0000 Subject: [PATCH 08/13] test(router): use real Logging in fallback cooldown test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/router_utils/test_fallback_event_handlers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 70b15e85c13..72477279c5e 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -1224,9 +1224,9 @@ class TestRunAsyncFallbackTriggersCooldown: async def async_function_with_fallbacks(self, *args, **kwargs): raise RuntimeError("fallback model also failed") - def _logging_obj(self, has_logged_async_failure: bool) -> MagicMock: - logging_obj = MagicMock() - logging_obj.model_call_details = {"has_logged_async_failure": has_logged_async_failure} + def _logging_obj(self, has_logged_async_failure: bool) -> Logging: + logging_obj: Final = _request_logging_obj() + logging_obj.model_call_details["has_logged_async_failure"] = has_logged_async_failure return logging_obj @pytest.mark.asyncio From 4df47b58a765894e5e8254e56d300101b330270c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 19:11:54 +0000 Subject: [PATCH 09/13] fix(router): preserve superseded retry usage for billing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/streaming_handler.py | 4 ++-- litellm/router.py | 5 ++-- .../router_utils/fallback_event_handlers.py | 10 ++++---- ...thropic_messages_pre_content_retry_wire.py | 13 ++++++++-- .../test_fallback_event_handlers.py | 24 ++++++++++++++++++- tests/unit/test_router/test_router.py | 13 ++++++---- 6 files changed, 50 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index bbd22968a9e..6656c3d7a3d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2305,8 +2305,8 @@ class CustomStreamWrapper: already delivered. Recover that partial usage from the chunks seen so far and stash it, with its cost, on the logging object so the failure handler records the real partial spend instead of zero. The router - discards the stash right before it opens the next attempt, so a request - that later recovers via a fallback logs that hop's own usage. + moves it into terminal failure usage before opening the next attempt, + so a recovered attempt logs its own usage. """ if self.logging_obj is None or not self.chunks: return diff --git a/litellm/router.py b/litellm/router.py index 529b3d55b70..6fbca759565 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -233,11 +233,11 @@ from litellm.router_utils.fallback_event_handlers import ( attempted_retries_for_request, carry_over_pre_routing_selection, carry_over_routed_deployment, + carry_superseded_attempt_usage, check_non_standard_fallback_format, clear_pre_routing_selection, committed_retry_budget_for_request, continuation_fallback_targets, - discard_superseded_attempt_usage, fallback_lookup_groups, fallback_target_name, fallback_target_overrides_request_params, @@ -6163,7 +6163,6 @@ class Router: for buffered_chunk in buffered_lifecycle_chunks: yield buffered_chunk except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate - await close_source() async for item in self._aanthropic_messages_recover_stream_error( stream_error, has_generated_content, @@ -6407,7 +6406,7 @@ class Router: verbose_router_logger.debug( "Retrying stream that failed before content, attempt %s of %s", attempt + 1, budget ) - discard_superseded_attempt_usage(retry_kwargs) + carry_superseded_attempt_usage(retry_kwargs) try: response = await attempt_function(**retry_kwargs) except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index e8248e9e429..31a664a38bc 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -441,13 +441,11 @@ def mid_stream_retry_kwargs( return {key: value for key, value in hop_kwargs.items() if key not in _MID_STREAM_RETRY_STRIPPED_KEYS} -def discard_superseded_attempt_usage(kwargs: Mapping[str, object]) -> None: - """The next attempt of this request is about to open. A stream the previous attempt dropped stashed the - usage it consumed on the request's logging object for its own failure log and for the failure row of a - request that runs out of attempts; the next attempt's log must carry its own usage instead.""" +def carry_superseded_attempt_usage(kwargs: Mapping[str, object]) -> None: + """Carry the superseded stream's partial usage into billing if the request ultimately fails.""" logging_obj: Final = kwargs.get("litellm_logging_obj") if isinstance(logging_obj, Logging): - logging_obj.discard_partial_usage_for_failure() + logging_obj.carry_failed_attempt_usage() def _failure_callbacks_blocked_logging_object(kwargs: Mapping[str, object]) -> Logging | None: @@ -881,7 +879,7 @@ async def run_async_fallback( kwargs["attempted_targets"] = attempted if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors - discard_superseded_attempt_usage(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs)) + carry_superseded_attempt_usage(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs)) response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) verbose_router_logger.info("Successful fallback b/w models.") response = add_fallback_headers_to_response( diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index 94df5984cdc..e12b68374be 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -397,6 +397,8 @@ def test_retry_rejected_before_it_opens_counts_against_the_same_budget(gateway: model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) response: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "deployment")) _assert_rejected_before_content(response, wire, status=500, attempts=2, error="error") + rows: Final = eventually(lambda: _request_rows(model), lambda values: len(values) >= 1, seconds=70) + assert rows == [{"status": "failure", "prompt_tokens": 5, "completion_tokens": 1}], rows def test_non_stream_messages_rejected_with_500_is_retried_per_the_deployment_budget(gateway: Gateway) -> None: @@ -508,11 +510,18 @@ def _request_rows(model: str) -> list[dict[str, JsonValue]]: ) -def test_dropped_attempt_and_its_retry_each_log_their_own_usage(gateway: Gateway, tmp_path: Path) -> None: +@pytest.mark.parametrize( + "failure", + (_Failure("drop_after_message_start"), _Failure("error_frame_after_message_start", 500)), + ids=("transport-drop", "recoverable-error-frame"), +) +def test_superseded_attempt_and_its_retry_each_log_their_own_usage( + gateway: Gateway, tmp_path: Path, failure: _Failure +) -> None: prompt: Final = _prompt() attempts: Final = Attempts() with ( - wire_server(_upstream(prompt, _Failure("drop_after_message_start"), attempts)) as wire, + wire_server(_upstream(prompt, failure, attempts)) as wire, wire_server(lambda request: Reply()) as sink, _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, candidate.scenario() as scenario, diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 72477279c5e..29aecc4b26d 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -199,7 +199,8 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ logging_obj.model_call_details["custom_llm_provider"] = "openai" logging_obj.model_call_details["litellm_params"] = {"acompletion": True} logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] - logging_obj.record_partial_usage_for_failure(Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53), 0.01) + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) await run_async_fallback( litellm_router=HopLoggingItsOwnSuccessRouter(), @@ -213,6 +214,27 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ assert [tokens for tokens, _ in recorder.logged] == [7] assert recorder.logged[0][1] != 0.01 + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) + + +@pytest.mark.asyncio +async def test_run_async_fallback_preserves_superseded_usage_when_hop_fails_before_opening() -> None: + logging_obj: Final = _request_logging_obj() + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=AlwaysFailRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + ) + + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) class _RateLimitedHop(Exception): diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 069b5572a50..93894efbcb5 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -16096,6 +16096,7 @@ class _AnthropicMessagesFakeByteStream: self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}} self.closed = False self.closed_as_superseded = False + self.close_contexts: list[bool] = [] def __aiter__(self): return self @@ -16106,6 +16107,7 @@ class _AnthropicMessagesFakeByteStream: return self._chunks.pop(0) async def aclose(self) -> None: + self.close_contexts.append(closing_superseded_attempt()) self.closed_as_superseded = self.closed_as_superseded or closing_superseded_attempt() self.closed = True @@ -17350,10 +17352,8 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi @pytest.mark.asyncio -async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_partial_usage(): - """The dropped attempt stashes the usage it consumed on the request's logging object, for its own failure - log and for the failure row of a request that runs out of retries. The router discards it before the retry - opens, so a retry that completes bills its own usage instead of the dropped attempt's.""" +async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_partial_usage() -> None: + """The dropped attempt's usage stays available for terminal billing without reaching the retry's success log.""" router = _anthropic_messages_retry_router(num_retries=1) logging_obj = LiteLLMLogging( model="glm", @@ -17364,7 +17364,8 @@ async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_pa litellm_call_id="retry-own-usage", function_id="retry-own-usage", ) - logging_obj.record_partial_usage_for_failure(Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53), 0.01) + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) stash_when_retry_opened: list[object] = [] def retry_stream(): @@ -17379,6 +17380,7 @@ async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_pa assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] assert stash_when_retry_opened == [None] assert "response_cost" not in logging_obj.model_call_details + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) @pytest.mark.asyncio @@ -17403,6 +17405,7 @@ async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_o assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] assert superseded_closed_when_retry_opened == [True] + assert superseded.close_contexts[0] is True @pytest.mark.asyncio From 29389ce988b48a223a5c20629fb8ee9430b2e9e9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 19:49:40 +0000 Subject: [PATCH 10/13] fix(proxy): bill retry input after prior stream usage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 18 +++++++++--- ..._utils_model_creation_and_error_logging.py | 28 +++++++++++++++++++ 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 29749afb429..b8046e7a636 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1009,7 +1009,7 @@ def _billed_across_attempts( ) -def _failure_usage_to_lift( +def _current_failure_usage_to_lift( model_call_details: Mapping[str, object], request_body: Mapping[str, object], dispatched: bool, @@ -1023,9 +1023,6 @@ def _failure_usage_to_lift( chat-completions provider) without filling optional_params. Returns the (combined_usage_object, response_cost) pair to lift, or None.""" recovered_usage: Final = model_call_details.get("combined_usage_object") - failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) - if failed_attempts: - return _billed_across_attempts(failed_attempts, recovered_usage, model_call_details.get("response_cost")) if recovered_usage is not None: return recovered_usage, model_call_details.get("response_cost") if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL): @@ -1049,6 +1046,19 @@ def _failure_usage_to_lift( return estimated_usage, 0.0 +def _failure_usage_to_lift( + model_call_details: Mapping[str, object], + request_body: Mapping[str, object], + dispatched: bool, +) -> tuple[object, object] | None: + current_attempt: Final = _current_failure_usage_to_lift(model_call_details, request_body, dispatched) + failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) + if not failed_attempts: + return current_attempt + current_usage, current_cost = current_attempt if current_attempt is not None else (None, None) + return _billed_across_attempts(failed_attempts, current_usage, current_cost) + + _EMPTY_LIFT: Final = MappingProxyType({}) diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index ec7104fa459..75f064a541c 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1921,6 +1921,34 @@ def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost(): assert response_cost == 0.0 +def test_a_retry_failure_bills_its_input_after_prior_stream_usage() -> None: + from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_usage_to_lift + from litellm.types.utils import Usage + + prior_usage: Final = Usage(prompt_tokens=4, completion_tokens=2, total_tokens=6) + estimated_retry_usage: Final = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) + assert estimated_retry_usage is not None + + lifted: Final = _failure_usage_to_lift( + model_call_details={ + "failed_attempts_usage": ((prior_usage, 0.25),), + "call_type": "acompletion", + "model": FAILURE_USAGE_MODEL, + "messages": ONE_USER_MESSAGE, + }, + request_body={}, + dispatched=True, + ) + + assert lifted is not None + usage, response_cost = lifted + assert isinstance(usage, Usage) + assert usage.prompt_tokens == prior_usage.prompt_tokens + estimated_retry_usage.prompt_tokens + assert usage.completion_tokens == prior_usage.completion_tokens + assert usage.total_tokens == prior_usage.total_tokens + estimated_retry_usage.total_tokens + assert response_cost == 0.25 + + @pytest.mark.parametrize( "model_call_details, dispatched", [ From a78fd88d879114e3b2fcd4ae49222c4cf456f60a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:12:19 +0000 Subject: [PATCH 11/13] fix(proxy): track retry dispatch for terminal billing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 6 ++- litellm/proxy/utils.py | 2 +- ..._utils_model_creation_and_error_logging.py | 45 ++++++++++++------- .../test_fallback_event_handlers.py | 2 + 4 files changed, 37 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ca596655d82..99de79e7bcc 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -805,6 +805,7 @@ class Logging(LiteLLMLoggingBaseClass): "litellm_call_id": litellm_call_id, "start_time": start_time, "first_api_call_start_time": None, + "current_attempt_api_call_start_time": None, "litellm_params": litellm_params, } new_call.streaming_chunks = [] @@ -1496,7 +1497,9 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - self.model_call_details["api_call_start_time"] = datetime.datetime.now() + api_call_start_time: Final = datetime.datetime.now() + self.model_call_details["api_call_start_time"] = api_call_start_time + self.model_call_details["current_attempt_api_call_start_time"] = api_call_start_time # Set-once first provider-handoff instant. api_call_start_time # is overwritten on every retry, so it can't measure one-time # preprocessing; pinning the first attempt excludes retry loops @@ -2318,6 +2321,7 @@ class Logging(LiteLLMLoggingBaseClass): def carry_failed_attempt_usage(self) -> None: """Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its own usage while the request's failure log still bills every attempt.""" + self.model_call_details["current_attempt_api_call_start_time"] = None usage: Final = self.model_call_details.pop("combined_usage_object", None) response_cost: Final = self.model_call_details.pop("response_cost", None) if isinstance(usage, Usage): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b8046e7a636..91928a72243 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1186,7 +1186,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, _usage_to_lift: Final = _failure_usage_to_lift( model_call_details=_model_call_details, request_body=request_data, - dispatched=_first_handoff is not None, + dispatched=_model_call_details.get("current_attempt_api_call_start_time") is not None, ) _entries: Final = ( ("start_time", _model_call_details.get("start_time")), diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 75f064a541c..ebe661cf43f 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1921,32 +1921,45 @@ def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost(): assert response_cost == 0.0 -def test_a_retry_failure_bills_its_input_after_prior_stream_usage() -> None: - from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_usage_to_lift +@pytest.mark.parametrize( + "current_attempt_api_call_start_time", + [None, 1700000001.0], + ids=["retry-not-dispatched", "retry-dispatched"], +) +def test_a_retry_failure_bills_input_only_for_dispatched_attempts( + current_attempt_api_call_start_time: float | None, +) -> None: + from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_fields_to_lift from litellm.types.utils import Usage prior_usage: Final = Usage(prompt_tokens=4, completion_tokens=2, total_tokens=6) estimated_retry_usage: Final = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert estimated_retry_usage is not None - lifted: Final = _failure_usage_to_lift( - model_call_details={ - "failed_attempts_usage": ((prior_usage, 0.25),), - "call_type": "acompletion", - "model": FAILURE_USAGE_MODEL, - "messages": ONE_USER_MESSAGE, - }, - request_body={}, - dispatched=True, + lifted: Final = _failure_fields_to_lift( + { + "litellm_logging_obj": _LoggingObj( + { + "first_api_call_start_time": 1700000000.0, + "current_attempt_api_call_start_time": current_attempt_api_call_start_time, + "failed_attempts_usage": ((prior_usage, 0.25),), + "call_type": "acompletion", + "model": FAILURE_USAGE_MODEL, + "messages": ONE_USER_MESSAGE, + } + ) + } ) - assert lifted is not None - usage, response_cost = lifted + usage: Final = lifted["combined_usage_object"] assert isinstance(usage, Usage) - assert usage.prompt_tokens == prior_usage.prompt_tokens + estimated_retry_usage.prompt_tokens + expected_retry_tokens: Final = ( + estimated_retry_usage.prompt_tokens if current_attempt_api_call_start_time is not None else 0 + ) + assert usage.prompt_tokens == prior_usage.prompt_tokens + expected_retry_tokens assert usage.completion_tokens == prior_usage.completion_tokens - assert usage.total_tokens == prior_usage.total_tokens + estimated_retry_usage.total_tokens - assert response_cost == 0.25 + assert usage.total_tokens == prior_usage.total_tokens + expected_retry_tokens + assert lifted["response_cost"] == 0.25 @pytest.mark.parametrize( diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 29aecc4b26d..aaf32cb6e73 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -198,6 +198,7 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ ) logging_obj.model_call_details["custom_llm_provider"] = "openai" logging_obj.model_call_details["litellm_params"] = {"acompletion": True} + logging_obj.model_call_details["current_attempt_api_call_start_time"] = datetime.now() logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) @@ -215,6 +216,7 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ assert [tokens for tokens, _ in recorder.logged] == [7] assert recorder.logged[0][1] != 0.01 assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) + assert logging_obj.model_call_details["current_attempt_api_call_start_time"] is None @pytest.mark.asyncio From 3da3e194d60bb4f245b326f38726cc7c42360f4b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:18:53 +0000 Subject: [PATCH 12/13] test(router): use deterministic retry timestamp Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/router_utils/test_fallback_event_handlers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index aaf32cb6e73..258961e0030 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -198,7 +198,7 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ ) logging_obj.model_call_details["custom_llm_provider"] = "openai" logging_obj.model_call_details["litellm_params"] = {"acompletion": True} - logging_obj.model_call_details["current_attempt_api_call_start_time"] = datetime.now() + logging_obj.model_call_details["current_attempt_api_call_start_time"] = datetime(2026, 1, 1) logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) From 5d7a1b8b28c97e10bd732d3eabe0857d7f7ead76 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:43:22 +0000 Subject: [PATCH 13/13] test(router): avoid real clock in fallback tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/router_utils/test_fallback_event_handlers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 258961e0030..b6117886dcb 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -152,7 +152,7 @@ def _request_logging_obj() -> Logging: messages=[{"role": "user", "content": "ping"}], stream=True, call_type="anthropic_messages", - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), litellm_call_id="fallback-hop", function_id="fallback-hop", ) @@ -176,7 +176,7 @@ class HopLoggingItsOwnSuccessRouter(FakeRouter): choices=[Choices(index=0, message=Message(role="assistant", content="served by the hop"))], usage=Usage(prompt_tokens=52, completion_tokens=7, total_tokens=59), ) - now = datetime.now() + now = datetime(2026, 1, 1) await kwargs["litellm_logging_obj"].async_success_handler(result=hop_response, start_time=now, end_time=now) return StreamingWrapper() @@ -191,7 +191,7 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ messages=[{"role": "user", "content": "ping"}], stream=False, call_type="completion", - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), litellm_call_id="fallback-hop-success", function_id="fallback-hop-success", dynamic_async_success_callbacks=[recorder],