From 58660e1c55462acde9ad8fb2db4a07a79b5c5483 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 19 Jun 2026 12:03:02 -0700 Subject: [PATCH 1/5] fix(passthrough): recover output tokens for interrupted anthropic streams (#30787) (cherry picked from commit bd74c62ff188d65e46e9e0a1a6c930aaf74bf9a2) --- .../anthropic_passthrough_logging_handler.py | 86 +++++++++++ ...t_anthropic_passthrough_logging_handler.py | 141 ++++++++++++++++++ 2 files changed, 227 insertions(+) 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 a912a88a993..22ed7b76d35 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 @@ -8,6 +8,9 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + get_content_from_model_response, +) from litellm.llms.anthropic import get_anthropic_config from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, @@ -136,6 +139,84 @@ class AnthropicPassthroughLoggingHandler: return model return None + @staticmethod + def _stream_was_interrupted( + all_chunks: Sequence[Union[str, bytes]], + ) -> bool: + """ + Anthropic ends a stream with ``content_block_stop`` -> ``message_delta`` + -> ``message_stop``; a client disconnect leaves the last event mid + ``content_block_delta``. Scan from the tail and decide on the first + terminal-region event, so the common completed case is O(1) rather than + re-deserializing every line of the stream. + """ + for raw in reversed(all_chunks): + text = raw.decode("utf-8") if isinstance(raw, bytes) else raw + for line in reversed(text.splitlines()): + if not line.startswith("data:"): + continue + try: + data = json.loads(line[len("data:") :].strip()) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(data, dict): + continue + etype = data.get("type") + if etype == "message_delta": + return False + if etype in ( + "content_block_delta", + "content_block_stop", + "message_start", + ): + return True + return True + + @staticmethod + def _recover_interrupted_stream_output_tokens( + response: Union[ModelResponse, TextCompletionResponse], + all_chunks: Sequence[Union[str, bytes]], + model: str, + ) -> None: + """ + An Anthropic stream interrupted before its terminal ``message_delta`` + (client disconnect) carries only the ``message_start`` ``output_tokens`` + placeholder (typically 1-3), so completion tokens and spend are + undercounted ~20x. Re-tokenize the buffered output text to recover a + realistic ``output_tokens`` for usage/cost. Completed streams are + untouched because their terminal ``message_delta`` short-circuits here. + """ + if not isinstance(response, ModelResponse): + return + if not AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks): + return + usage = getattr(response, "usage", None) + if usage is None: + return + output_text = get_content_from_model_response(response) + if not output_text: + return + try: + recovered_output_tokens = litellm.token_counter( + model=model, text=output_text, count_response_tokens=True + ) + except Exception: + verbose_proxy_logger.warning( + "Could not re-tokenize interrupted stream output; " + "keeping placeholder completion token count." + ) + return + if recovered_output_tokens <= (usage.completion_tokens or 0): + return + usage.completion_tokens = recovered_output_tokens + usage.total_tokens = (usage.prompt_tokens or 0) + recovered_output_tokens + # Anthropic costing reads completion_tokens_details.text_tokens, so the + # stale message_start placeholder there must be corrected too or spend + # stays undercounted even after completion_tokens is fixed. + details = getattr(usage, "completion_tokens_details", None) + if details is not None and getattr(details, "text_tokens", None) is not None: + details.text_tokens = recovered_output_tokens + @staticmethod def _create_anthropic_response_logging_payload( litellm_model_response: Union[ModelResponse, TextCompletionResponse], @@ -277,6 +358,11 @@ class AnthropicPassthroughLoggingHandler: "result": None, "kwargs": {}, } + AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( + response=complete_streaming_response, + all_chunks=all_chunks, + model=model, + ) kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( litellm_model_response=complete_streaming_response, model=model, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 1083c3e9252..224659f6079 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -1376,6 +1376,147 @@ class TestPureTextFastPathParity: ) +class TestInterruptedStreamOutputTokenRecovery: + """ + When an Anthropic pass-through stream is interrupted (client disconnect) + before the terminal ``message_delta``, the only usage signal is the + ``message_start`` ``output_tokens`` placeholder (typically 1-3), so + completion tokens and spend are undercounted ~20x. The handler must + re-tokenize the buffered ``content_block_delta`` text to recover a + realistic ``output_tokens``; completed streams must stay untouched. + """ + + @staticmethod + def _sse(event, data): + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + _MODEL = "claude-3-5-haiku-20241022" + _OUTPUT_TEXT = ( + "The history of computing spans centuries, beginning with mechanical " + "calculators and the abacus, advancing through Charles Babbage's " + "analytical engine, Ada Lovelace's first algorithm, Alan Turing's " + "theoretical machine, and the electronic computers of the twentieth " + "century that gave rise to the modern information age." + ) + + def _interrupted_chunks(self, *, placeholder_output_tokens: int = 2): + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + + words = self._OUTPUT_TEXT.split(" ") + frames = [ + self._sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_interrupted", + "type": "message", + "role": "assistant", + "model": self._MODEL, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": { + "input_tokens": 29, + "output_tokens": placeholder_output_tokens, + }, + }, + }, + ), + self._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + ] + for i, word in enumerate(words): + text = word if i == 0 else " " + word + frames.append( + self._sse( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": text}, + }, + ) + ) + # Client disconnects here: no content_block_stop / message_delta / + # message_stop are ever received. + return list(PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames)) + + def _completed_chunks(self, *, final_output_tokens: int = 80): + chunks = self._interrupted_chunks() + chunks.append( + "data: " + + json.dumps( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": final_output_tokens}, + } + ) + ) + chunks.append('data: {"type": "message_stop"}') + return chunks + + def _run(self, all_chunks): + logging_obj = MagicMock() + logging_obj.model_call_details = {"model": self._MODEL, "stream": True} + logging_obj.litellm_call_id = "test-call-id" + logging_obj.litellm_params = {} + logging_obj.get_router_model_id.return_value = None + + return AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": self._MODEL, "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=all_chunks, + end_time=datetime.now(), + ) + + def test_interrupted_stream_retokenizes_buffered_output(self): + import litellm + + placeholder = 2 + result = self._run( + self._interrupted_chunks(placeholder_output_tokens=placeholder) + ) + usage = result["result"].usage + + expected = litellm.token_counter( + model=self._MODEL, + text=self._OUTPUT_TEXT, + count_response_tokens=True, + ) + + assert expected > placeholder * 5 + assert usage.completion_tokens == expected + assert usage.completion_tokens > placeholder + assert usage.total_tokens == usage.prompt_tokens + expected + # Anthropic spend is priced off completion_tokens_details.text_tokens; if the + # placeholder leaks through here, cost stays undercounted even though + # completion_tokens looks right. + assert usage.completion_tokens_details.text_tokens == expected + + def test_completed_stream_keeps_message_delta_tokens(self): + final = 80 + result = self._run(self._completed_chunks(final_output_tokens=final)) + usage = result["result"].usage + + # Terminal message_delta present: recovery must not fire; the authoritative + # provider count is preserved verbatim. + assert usage.completion_tokens == final + + class TestStreamFalseDeduplication: """ Regression tests for the duplicate-callback bug where a streaming pass-through From 433d016f0c25f18b638060ecb7246fda8f1aac49 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 19 Jun 2026 12:03:15 -0700 Subject: [PATCH 2/5] fix(proxy): record partial spend on the failure row for interrupted streams (#30788) A streaming request that breaks mid-flight, for example on a mid-stream read timeout, still bills the provider for the chunks already delivered, yet the proxy recorded that interrupted request as a zero-spend failure. An earlier revision logged the recovered partial usage through the success path, which mislabeled a failed request as a success and produced a misleading spend row This recovers the partial usage where the failure is actually logged. The streaming handler assembles the usage from the chunks seen so far and stashes it, with its cost, on the logging object before firing the failure handlers. The proxy failure hook lifts that usage and cost onto request_data before the non-serialisable logging object is popped, and the spend-log writer records the real partial spend on the failure row instead of a hardcoded zero; get_logging_payload honors the recovered usage for the token columns and _failure_handler_helper_fn preserves the recovered cost so the non-DB failure loggers stay consistent A request that recovers via a successful fallback is unaffected: the failure hook only fires when the whole request fails, so the fallback's combined-usage success row stays the single source of truth and there is no double counting Resolves LIT-3825 Co-authored-by: veria-ai[bot] <224490171+veria-ai[bot]@users.noreply.github.com> (cherry picked from commit 4847fa5dd5991496a071d235781e07d39857b0f7) --- litellm/litellm_core_utils/litellm_logging.py | 7 +- .../litellm_core_utils/streaming_handler.py | 29 ++++ .../proxy/hooks/proxy_track_cost_callback.py | 15 +- .../spend_tracking/spend_tracking_utils.py | 7 + litellm/proxy/utils.py | 15 +- .../test_litellm_logging.py | 43 ++++++ .../test_streaming_handler.py | 76 +++++++++ .../hooks/test_proxy_track_cost_callback.py | 37 +++++ .../test_spend_tracking_utils.py | 47 ++++++ tests/test_litellm/proxy/test_proxy_utils.py | 49 ++++++ tests/test_litellm/test_router.py | 145 ++++++++++++++++++ 11 files changed, 464 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 18321c70f6e..482bf88cbc4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2951,7 +2951,12 @@ class Logging(LiteLLMLoggingBaseClass): ) self.model_call_details["end_time"] = end_time self.model_call_details.setdefault("original_response", None) - self.model_call_details["response_cost"] = 0 + # A stream interrupted mid-flight still billed the provider for the + # chunks already delivered; the router stashes that recovered usage as + # ``combined_usage_object`` and pre-computes its cost, so preserve it + # here instead of zeroing the spend on an otherwise-failed request. + if self.model_call_details.get("combined_usage_object") is None: + self.model_call_details["response_cost"] = 0 if hasattr(exception, "headers") and isinstance(exception.headers, dict): self.model_call_details.setdefault("litellm_params", {}) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 29c0d0629e8..3d04b183176 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2233,6 +2233,7 @@ class CustomStreamWrapper: litellm.request_timeout ) if self.logging_obj is not None: + self._record_partial_usage_for_failure() ## LOGGING threading.Thread( target=self.logging_obj.failure_handler, @@ -2246,6 +2247,7 @@ class CustomStreamWrapper: except Exception as e: traceback_exception = traceback.format_exc() if self.logging_obj is not None: + self._record_partial_usage_for_failure() ## LOGGING threading.Thread( target=self.logging_obj.failure_handler, @@ -2257,6 +2259,33 @@ class CustomStreamWrapper: ) self._handle_stream_fallback_error(e) + 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. + """ + if self.logging_obj is None or not self.chunks: + return + try: + partial_response = litellm.stream_chunk_builder(chunks=self.chunks) + usage = cast(Optional[Usage], getattr(partial_response, "usage", None)) + if usage is None: + return + self.logging_obj.model_call_details["combined_usage_object"] = usage + self.logging_obj.model_call_details["response_cost"] = ( + self.logging_obj._response_cost_calculator(result=partial_response) + or 0.0 + ) + except Exception as recover_error: + verbose_logger.debug( + "could not recover partial usage for interrupted stream: %s", + recover_error, + ) + def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn": """ Common error handling for both __next__ and __anext__. diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3688f25ac44..0a85b181fba 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -40,7 +40,7 @@ class _ProxyDBLogger(CustomLogger): kwargs, response_obj, start_time, end_time ) - async def async_post_call_failure_hook( + async def async_post_call_failure_hook( # noqa: PLR0915 self, request_data: dict, original_exception: Exception, @@ -162,9 +162,20 @@ class _ProxyDBLogger(CustomLogger): if obj_start is not None: actual_start_time = obj_start + # A stream that broke mid-flight still billed the provider for the + # chunks already delivered. ``post_call_failure_hook`` lifts that + # recovered cost onto request_data (the usage rides along in + # ``combined_usage_object`` for the token columns), so attribute the + # real partial spend to this failure row instead of zero. + recovered_response_cost = 0.0 + if isinstance(request_data.get("combined_usage_object"), litellm.Usage): + recovered_response_cost = max( + float(request_data.get("response_cost") or 0.0), 0.0 + ) + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, - response_cost=0.0, + response_cost=recovered_response_cost, user_id=user_api_key_dict.user_id, end_user_id=user_api_key_dict.end_user_id, team_id=user_api_key_dict.team_id, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index e2881faca0d..26d8d748d2e 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -265,6 +265,13 @@ def get_logging_payload( # noqa: PLR0915 elif isinstance(_usage, dict): usage = _usage + # A request that failed mid-stream has no usable response_obj usage, but the + # streaming handler may have recovered the usage from the chunks already + # delivered. Honor that override so the partial usage lands in spend tracking. + _combined_usage = kwargs.get("combined_usage_object") + if not usage and isinstance(_combined_usage, litellm.Usage): + usage = _combined_usage.model_dump() + id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs) standard_logging_payload = cast( Optional[StandardLoggingPayload], kwargs.get("standard_logging_object", None) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index effd62da717..f0da87e1271 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2006,12 +2006,21 @@ class ProxyLogging: # compute preprocessing latency after the logging object is popped. _logging_obj = request_data.get("litellm_logging_obj") if _logging_obj is not None: - _first_handoff = getattr(_logging_obj, "model_call_details", {}).get( - "first_api_call_start_time" - ) + _model_call_details = getattr(_logging_obj, "model_call_details", {}) + _first_handoff = _model_call_details.get("first_api_call_start_time") if _first_handoff is not None: request_data["first_api_call_start_time"] = _first_handoff + # A stream that broke mid-flight still billed the provider for the + # chunks already delivered; the streaming handler stashes that + # recovered usage and cost here. Lift them onto request_data so the + # failure-path spend callbacks (which run after the logging object + # is popped) record the real partial spend instead of zero. + _recovered_usage = _model_call_details.get("combined_usage_object") + if _recovered_usage is not None: + request_data["combined_usage_object"] = _recovered_usage + request_data["response_cost"] = _model_call_details.get("response_cost") + # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index b64cb7c6905..f5677a3fcf2 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3078,3 +3078,46 @@ class TestFirstApiCallStartTimeSetOnce: assert obj.model_call_details["api_call_start_time"] > first assert obj.model_call_details["first_api_call_start_time"] == first assert user_meta == {} + + +def test_failure_handler_records_recovered_partial_spend(logging_obj): + """A stream interrupted mid-flight still billed the provider for the chunks + already delivered. When the router stashes that recovered usage as + ``combined_usage_object`` and pre-computes ``response_cost``, the failure + handler must preserve them so the failure row carries the real partial + spend instead of zero. + """ + from litellm.types.utils import Usage + + logging_obj.model_call_details["combined_usage_object"] = Usage( + prompt_tokens=17, completion_tokens=9, total_tokens=26 + ) + logging_obj.model_call_details["response_cost"] = 0.00012 + + logging_obj._failure_handler_helper_fn( + exception=Exception("Connection lost"), + traceback_exception="Traceback ...", + ) + + payload = logging_obj.model_call_details["standard_logging_object"] + assert payload["status"] == "failure" + assert payload["response_cost"] == 0.00012 + assert payload["prompt_tokens"] == 17 + assert payload["completion_tokens"] == 9 + assert payload["total_tokens"] == 26 + + +def test_failure_handler_zeroes_spend_without_recovered_usage(logging_obj): + """A failure with no recovered partial usage keeps the existing behavior of + recording zero spend, so the partial-spend preservation does not leak into + ordinary failures. + """ + logging_obj._failure_handler_helper_fn( + exception=Exception("boom"), + traceback_exception="Traceback ...", + ) + + payload = logging_obj.model_call_details["standard_logging_object"] + assert payload["status"] == "failure" + assert payload["response_cost"] == 0 + assert payload["total_tokens"] == 0 diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 63e2cb7f35c..09bb8532ec2 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2118,3 +2118,79 @@ def test_gemini_legacy_vertex_tool_calls_finish_reason_with_stop_enum(): f"Expected 'tool_calls' but got {final.choices[0].finish_reason!r}. " "STOP enum was not normalised through map_finish_reason()." ) + + +def test_record_partial_usage_for_failure_stashes_usage_and_cost(): + """A stream that breaks mid-flight must surface the usage assembled from the + chunks already delivered, plus its cost, on the logging object so the + failure handler records the real partial spend instead of zero. + """ + 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-1", + function_id="1245", + ) + 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-1", + created=1742056047, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="The Roman Empire began when", role="assistant" + ), + ) + ], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.prompt_tokens == 30 + assert stashed.completion_tokens == 1 + assert stashed.total_tokens == 31 + assert isinstance(logging_obj.model_call_details["response_cost"], float) + + +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. + """ + 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-2", + function_id="1245", + ) + wrapper = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [] + + wrapper._record_partial_usage_for_failure() + + assert "combined_usage_object" not in logging_obj.model_call_details diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 771e10a54a0..0cbf308076c 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1067,3 +1067,40 @@ async def test_failure_hook_drops_error_information_traceback_when_env_set( assert "traceback" not in error_information assert error_information["error_class"] == "RuntimeError" assert error_information["error_message"] == "boom-with-traceback" + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_records_recovered_partial_spend(): + """A stream that broke mid-flight still billed the provider. The failure + hook lifts the recovered cost onto request_data as ``response_cost``; this + hook must pass it through to update_database so the failure row records the + real partial spend instead of the hardcoded zero. + """ + from litellm.types.utils import Usage + + logger = _ProxyDBLogger() + user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key", user_id="u", team_id="t") + + request_data = { + "model": "anthropic/claude-haiku-4-5", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "proxy_server_request": {"request_id": "rid"}, + "response_cost": 3.5e-05, + "combined_usage_object": Usage( + prompt_tokens=30, completion_tokens=1, total_tokens=31 + ), + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("MidStreamFallbackError: read timeout"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05 diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 5ca058fc8d9..5272b105eb5 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2009,3 +2009,50 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form( assert sanitized is not None assert "leaked-via-pydantic-msg" not in sanitized["error_message"] assert REDACTED_BY_LITELM_STRING in sanitized["error_message"] + + +def test_get_logging_payload_uses_recovered_combined_usage_on_failure(): + """A request that fails mid-stream has no usable response_obj usage, but the + streaming handler recovers the usage from the chunks already delivered and + the failure hook surfaces it as ``combined_usage_object``. The spend-log + payload must record those token counts instead of zero. + """ + from litellm.types.utils import Usage + + kwargs = { + "model": "anthropic/claude-haiku-4-5", + "call_type": "acompletion", + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + "combined_usage_object": Usage( + prompt_tokens=30, completion_tokens=1, total_tokens=31 + ), + } + response_obj = Exception("MidStreamFallbackError: read timeout") + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert payload["prompt_tokens"] == 30 + assert payload["completion_tokens"] == 1 + assert payload["total_tokens"] == 31 + + +def test_get_logging_payload_failure_without_recovered_usage_is_zero(): + """A failure with no recovered usage keeps zero token counts, so the + combined-usage override never invents tokens for ordinary failures. + """ + kwargs = { + "model": "anthropic/claude-haiku-4-5", + "call_type": "acompletion", + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = Exception("BadRequestError") + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert payload["total_tokens"] == 0 diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 539a32db57d..b2a81a0cfaf 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -435,3 +435,52 @@ class TestPostCallFailureHookProxyExceptionLogging: ) is False ) +class TestPostCallFailureHookLiftsRecoveredPartialSpend: + """A stream that broke mid-flight still billed the provider for the chunks + already delivered. The streaming handler stashes that recovered usage and + cost on the logging object; post_call_failure_hook must lift them onto + request_data before the logging object is popped, so the failure-path spend + callbacks (which run after the pop) record the real partial spend. + """ + + async def _run(self, request_data): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import UserAPIKeyAuth + + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=Exception("boom"), + user_api_key_dict=UserAPIKeyAuth(), + ) + + @pytest.mark.asyncio + async def test_lifts_recovered_usage_and_cost(self): + from litellm.types.utils import Usage + + recovered_usage = Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31) + logging_obj = MagicMock() + logging_obj.model_call_details = { + "combined_usage_object": recovered_usage, + "response_cost": 3.5e-05, + } + request_data = {"litellm_logging_obj": logging_obj, "metadata": {}} + await self._run(request_data) + + assert request_data["combined_usage_object"] is recovered_usage + assert request_data["response_cost"] == 3.5e-05 + assert "litellm_logging_obj" not in request_data + + @pytest.mark.asyncio + async def test_no_recovered_usage_is_noop(self): + logging_obj = MagicMock() + logging_obj.model_call_details = {} + request_data = {"litellm_logging_obj": logging_obj, "metadata": {}} + await self._run(request_data) + assert "combined_usage_object" not in request_data + assert "response_cost" not in request_data + + diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5e636b86ed6..7a28eacca8d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3340,6 +3340,151 @@ def test_combine_fallback_usage(): assert chunk.usage.total_tokens == 15 +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure(): + """A mid-stream failure with no successful fallback raises and is logged as + a failure, so the router must never dispatch it as a success. Partial-spend + recovery for the failure row happens in the streaming handler, not here, so + this guards only against reintroducing a success log for a failed stream. + """ + from litellm.exceptions import MidStreamFallbackError + from litellm.types.utils import Delta, StreamingChoices, Usage + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, + }, + ], + set_verbose=True, + ) + + error = MidStreamFallbackError( + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="The Roman Empire began when", + ) + + def _make_interrupted_model_response(): + partial_chunk = litellm.ModelResponseStream( + id="chatcmpl-partial-1", + created=1742056047, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="The Roman Empire began when", role="assistant" + ), + ) + ], + usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), + ) + + class _RaisingStream: + def __init__(self): + self.index = 0 + self.chunks = [partial_chunk] + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index == 0: + self.index += 1 + return partial_chunk + raise error + + stream = _RaisingStream() + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.model_call_details = {} + setattr(stream, "model", "gpt-4") + setattr(stream, "custom_llm_provider", "openai") + setattr(stream, "logging_obj", logging_obj) + return stream, logging_obj + + messages = [{"role": "user", "content": "Hello"}] + initial_kwargs = {"model": "gpt-4", "stream": True} + + # Terminal path: no successful fallback -> the error propagates and the + # router never dispatches a success for the failed stream. + model_response, logging_obj = _make_interrupted_model_response() + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=error), + ): + result = await router._acompletion_streaming_iterator( + model_response=model_response, + messages=messages, + initial_kwargs=dict(initial_kwargs), + ) + collected = [] + with pytest.raises(MidStreamFallbackError): + async for chunk in result: + collected.append(chunk) + + assert len(collected) == 1 + logging_obj.dispatch_success_handlers.assert_not_called() + + # Fallback success: the fallback stream owns success accounting via + # _combine_fallback_usage, so this iterator must not dispatch its own. + model_response, logging_obj = _make_interrupted_model_response() + + class _FallbackStream: + def __init__(self, items): + self.items = items + self.index = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index >= len(self.items): + raise StopAsyncIteration + item = self.items[self.index] + self.index += 1 + return item + + fallback_stream = _FallbackStream( + [ + litellm.ModelResponseStream( + id="chatcmpl-fallback-1", + model="gpt-3.5-turbo", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=" continued", role="assistant"), + ) + ], + ) + ] + ) + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ): + result = await router._acompletion_streaming_iterator( + model_response=model_response, + messages=messages, + initial_kwargs=dict(initial_kwargs), + ) + collected = [] + async for chunk in result: + collected.append(chunk) + + assert len(collected) == 2 + logging_obj.dispatch_success_handlers.assert_not_called() + + @pytest.mark.asyncio async def test_team_scoped_model_fallback(): """ From dabf43b594cf8802f74b4b7216e702cc5d86a9a4 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 21:20:11 -0700 Subject: [PATCH 3/5] fix: completion_cost AttributeError on streaming Anthropic web_search responses (#26153) (#27346) * fix: coerce server_tool_use dict to ServerToolUse in Usage.__init__ (#26153) * fix: coerce server_tool_use to ServerToolUse in stream_chunk_builder (#26153) * fix: dict/pydantic-tolerant access in tool_call_cost_tracking (#26153) * fix: dict/pydantic-tolerant access in anthropic cost_calculation (#26153) * test: assert ServerToolUse type in existing stream_chunk_builder anthropic web search test * test: regression test for #26153 (stream_chunk_builder server_tool_use type) * test: dict/pydantic safety for tool_call_cost_tracking helper * test: dict/pydantic safety for anthropic web_search cost * refactor: consolidate _get_web_search_requests into shared cost-calc utils * test(realtime): use gpt-realtime; openai retired gpt-4o-realtime-preview OpenAI shut down the gpt-4o-realtime-preview family (incl. the undated alias) on 2026-05-07, causing the live realtime test to fail with a 4000 invalid_request_error.invalid_model close. gpt-realtime is the GA successor; switch the live-call tests to it, matching the base branch. * refactor(types): drop redundant server_tool_use coercion in Usage.__init__ --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> (cherry picked from commit 4a3860df1f148486d76093cf95b631e39f888510) [backport note -- stable/1.87.x] Aggregator-provenance pick, user-approved: #27346's content reached litellm_internal_staging only via the squashed OSS sync (#29932 / 32c88ca74f), so merge 4a3860df1f1 is not an ancestor of staging tip; its post-image is content-identical to staging. Also restored the 2-line server_tool_use dict->ServerToolUse coercion in litellm/types/utils.py: it is #27346's own first commit, dropped from the squash ("refactor(types): drop redundant server_tool_use coercion") only because staging already carried it via the OSS sync. 1.87.x predates that, and #31035's usage-only fallback needs it. Mirrors the stable/1.88.x precedent (24b9655cd42). --- .../llm_cost_calc/tool_call_cost_tracking.py | 7 +- .../litellm_core_utils/llm_cost_calc/utils.py | 22 ++- .../streaming_chunk_builder_utils.py | 13 +- litellm/llms/anthropic/cost_calculation.py | 14 +- litellm/types/utils.py | 5 +- ...est_tool_call_cost_tracking_dict_safety.py | 88 ++++++++++++ ...streaming_chunk_builder_server_tool_use.py | 130 ++++++++++++++++++ .../test_streaming_chunk_builder_utils.py | 5 +- .../test_cost_calculation_dict_safety.py | 94 +++++++++++++ 9 files changed, 364 insertions(+), 14 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py create mode 100644 tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py create mode 100644 tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 8da66d4600d..413ddb71bf8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -6,6 +6,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS +from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests from litellm.types.llms.openai import ( FileSearchTool, ResponsesAPIResponse, @@ -339,8 +340,7 @@ class StandardBuiltInToolCostTracking: # and _handle_web_search_cost() is never called. if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True return False @@ -352,8 +352,7 @@ class StandardBuiltInToolCostTracking: elif usage is not None: if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True elif ( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 59d0465e6d4..ba8b0764741 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1,7 +1,7 @@ # What is this? ## Helper utilities for cost_per_token() -from typing import Literal, Optional, Tuple, TypedDict, cast +from typing import Any, Literal, Optional, Tuple, TypedDict, cast import litellm from litellm._logging import verbose_logger @@ -30,6 +30,26 @@ _IMAGE_RESPONSE_CALL_TYPES = frozenset( ) +def _get_web_search_requests(server_tool_use: Any) -> Optional[int]: + """ + Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value + that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance, + or any other object supporting attribute access. + + Returns ``None`` when the value cannot be resolved — callers can + distinguish "absent" from "zero" using ``is None``. + + See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder`` + historically left this as a plain ``dict``, which broke direct attribute + access in cost calculation. + """ + if server_tool_use is None: + return None + if isinstance(server_tool_use, dict): + return server_tool_use.get("web_search_requests") + return getattr(server_tool_use, "web_search_requests", None) + + def _is_above_128k(tokens: float) -> bool: if tokens > 128000: return True diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fe7c62c3842..c2a17ae8dcc 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -588,7 +588,18 @@ class ChunkProcessor: hasattr(usage_chunk, "server_tool_use") and usage_chunk.server_tool_use is not None ): - server_tool_use = usage_chunk.server_tool_use + # Coerce dict to ServerToolUse so downstream cost-calc code + # (which accesses .web_search_requests as an attribute) + # doesn't raise AttributeError. Some providers / streaming + # paths leave server_tool_use as a plain dict on the chunk. + if isinstance(usage_chunk.server_tool_use, dict): + server_tool_use = ServerToolUse(**usage_chunk.server_tool_use) + elif isinstance(usage_chunk.server_tool_use, ServerToolUse): + server_tool_use = usage_chunk.server_tool_use + else: + server_tool_use = ServerToolUse.model_validate( + usage_chunk.server_tool_use + ) if ( usage_chunk_dict["prompt_tokens_details"] is not None and getattr( diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index 3882d8f978c..6a031498dae 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Optional, Tuple from litellm.litellm_core_utils.llm_cost_calc.utils import ( _get_token_base_cost, + _get_web_search_requests, _parse_prompt_tokens_details, calculate_cache_writing_cost, generic_cost_per_token, @@ -110,11 +111,12 @@ def get_cost_for_anthropic_web_search( if model_info is None: return 0.0 - if ( - usage is None - or usage.server_tool_use is None - or usage.server_tool_use.web_search_requests is None - ): + if usage is None: + return 0.0 + web_search_requests = _get_web_search_requests( + getattr(usage, "server_tool_use", None) + ) + if web_search_requests is None: return 0.0 ## Get the cost per web search request @@ -128,5 +130,5 @@ def get_cost_for_anthropic_web_search( return 0.0 ## Calculate the total cost - total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests + total_cost = cost_per_web_search_request * web_search_requests return total_cost diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c3bf729238f..3f0fe99d836 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1561,7 +1561,7 @@ class Usage(SafeAttributeModel, CompletionUsage): completion_tokens_details: Optional[ Union[CompletionTokensDetailsWrapper, dict] ] = None, - server_tool_use: Optional[ServerToolUse] = None, + server_tool_use: Optional[Union[ServerToolUse, dict]] = None, cost: Optional[float] = None, **params, ): @@ -1662,6 +1662,9 @@ class Usage(SafeAttributeModel, CompletionUsage): prompt_tokens_details=_prompt_tokens_details or None, ) + if isinstance(server_tool_use, dict): + server_tool_use = ServerToolUse(**server_tool_use) + if server_tool_use is not None: self.server_tool_use = server_tool_use else: # maintain openai compatibility in usage object if possible diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py new file mode 100644 index 00000000000..4eee6b59d34 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py @@ -0,0 +1,88 @@ +""" +Tests that the cost-tracking call sites tolerate ``server_tool_use`` being +either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( + StandardBuiltInToolCostTracking, + _get_web_search_requests, +) +from litellm.types.utils import ModelResponse, ServerToolUse, Usage + + +class _UsageWithDictServerToolUse: + """ + Tiny stand-in that mimics the broken streaming-rebuild shape: + ``server_tool_use`` is a plain dict. + """ + + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + self.prompt_tokens_details = None + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 5}) == 5 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + stu = ServerToolUse(web_search_requests=7) + assert _get_web_search_requests(stu) == 7 + + +def test_get_web_search_requests_handles_pydantic_with_none_value(): + stu = ServerToolUse() + assert _get_web_search_requests(stu) is None + + +def test_response_object_includes_web_search_call_with_dict_server_tool_use(): + """ + The exact bug: ``usage.server_tool_use`` is a dict and the check in + ``response_object_includes_web_search_call`` used to crash with + ``AttributeError``. + """ + response = ModelResponse() + usage = _UsageWithDictServerToolUse({"web_search_requests": 2}) + + # Must not raise — and must correctly detect the web search call. + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_pydantic_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(ServerToolUse(web_search_requests=2)) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_none_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(None) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is False diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py new file mode 100644 index 00000000000..4e28d5ba7d2 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py @@ -0,0 +1,130 @@ +""" +Regression tests for https://github.com/BerriAI/litellm/issues/26153 + +``stream_chunk_builder`` used to leave ``usage.server_tool_use`` as a plain +``dict`` when reconstructing a streaming response. Downstream cost-calculation +code (``StandardBuiltInToolCostTracking.response_object_includes_web_search_call`` +and ``get_cost_for_anthropic_web_search``) accesses +``usage.server_tool_use.web_search_requests`` as an attribute, which raised +``AttributeError: 'dict' object has no attribute 'web_search_requests'``. + +These tests reconstruct streaming chunks for an Anthropic-style web_search +response and assert: + +1. ``stream_chunk_builder`` returns ``ServerToolUse`` (not ``dict``) for + ``usage.server_tool_use``. +2. ``completion_cost`` runs end-to-end on the rebuilt response without + raising ``AttributeError``. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm import completion_cost, stream_chunk_builder +from litellm.types.utils import ( + Delta, + ModelResponseStream, + ServerToolUse, + StreamingChoices, + Usage, +) + + +def _make_text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(role="assistant", content=text), + ) + ], + ) + + +def _make_finish_chunk_with_usage_dict_server_tool_use() -> ModelResponseStream: + """Final chunk where server_tool_use is a *dict* — reproduces the bug shape.""" + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ], + usage=Usage( + prompt_tokens=42, + completion_tokens=11, + total_tokens=53, + # NOTE: passed as a dict on purpose — this is the shape that + # historically slipped through stream_chunk_builder unchanged. + server_tool_use={"web_search_requests": 3}, + ), + ) + + +def test_stream_chunk_builder_coerces_server_tool_use_to_pydantic(): + """ + Regression: stream_chunk_builder must produce ServerToolUse, not dict. + """ + chunks = [ + _make_text_chunk("Otters "), + _make_text_chunk("are great."), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + + assert rebuilt is not None + assert rebuilt.usage is not None # type: ignore[attr-defined] + server_tool_use = rebuilt.usage.server_tool_use # type: ignore[attr-defined] + + assert ( + server_tool_use is not None + ), "server_tool_use should be carried through from the final chunk" + assert isinstance(server_tool_use, ServerToolUse), ( + f"expected ServerToolUse, got {type(server_tool_use).__name__}: " + f"{server_tool_use!r}" + ) + # Attribute access must not raise (this is exactly what was broken). + assert server_tool_use.web_search_requests == 3 + + +def test_completion_cost_does_not_raise_on_streaming_web_search_response(): + """ + Regression: completion_cost(...) must not raise AttributeError when the + response was reconstructed by stream_chunk_builder from a streaming + Anthropic web_search call. + """ + chunks = [ + _make_text_chunk("hello"), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + assert rebuilt is not None + + # The exact dollar amount depends on the model-pricing table; what matters + # for this regression is that it does NOT raise AttributeError on + # `dict has no attribute 'web_search_requests'`. + try: + cost = completion_cost(completion_response=rebuilt) + except AttributeError as e: # pragma: no cover - regression guard + pytest.fail( + "completion_cost raised AttributeError after stream_chunk_builder " + f"(issue #26153 regression): {e}" + ) + + assert isinstance(cost, (int, float)) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd9..35aca525f6c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -520,7 +520,10 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 assert usage.total_tokens == 77 - assert usage.server_tool_use["web_search_requests"] == 2 + # server_tool_use must be a ServerToolUse pydantic so downstream cost-calc + # (which uses attribute access) works. See issue #26153. + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 2 def test_sort_chunks_handles_dict_hidden_params_created_at(): diff --git a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py new file mode 100644 index 00000000000..70fef0162e6 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py @@ -0,0 +1,94 @@ +""" +Tests that ``get_cost_for_anthropic_web_search`` tolerates ``server_tool_use`` +being either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.anthropic.cost_calculation import ( + _get_web_search_requests, + get_cost_for_anthropic_web_search, +) +from litellm.types.utils import ModelInfo, ServerToolUse + + +class _UsageWithServerToolUse: + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + + +def _make_model_info(cost_per_query: float = 0.01) -> ModelInfo: + info: ModelInfo = { # type: ignore[typeddict-item] + "search_context_cost_per_query": { + "search_context_size_low": cost_per_query, + "search_context_size_medium": cost_per_query, + "search_context_size_high": cost_per_query, + } + } + return info + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 4}) == 4 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + assert _get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2 + + +def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use(): + """ + Regression: ``server_tool_use`` was a dict from ``stream_chunk_builder`` and + direct attribute access on it raised ``AttributeError``. + """ + usage = _UsageWithServerToolUse({"web_search_requests": 3}) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use(): + usage = _UsageWithServerToolUse(ServerToolUse(web_search_requests=3)) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_none_server_tool_use(): + usage = _UsageWithServerToolUse(None) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == 0.0 + + +def test_get_cost_for_anthropic_web_search_with_no_usage(): + info = _make_model_info(cost_per_query=0.01) + cost = get_cost_for_anthropic_web_search(model_info=info, usage=None) + assert cost == 0.0 From b77cbbdcc95b83db4a76689b5f130d0ec4482846 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 22 Jun 2026 18:51:13 -0700 Subject: [PATCH 4/5] fix(passthrough,streaming): recover cost on interrupted and agentic Anthropic streams (#31035) Streaming and pass-through requests could be logged with $0 cost or dropped from SpendLogs entirely while the upstream provider still billed every token. This closes the leak paths not already covered by #30160, #30787 and #30788. - Catch a stream_chunk_builder raise in the core CustomStreamWrapper (sync and async). Large agentic tool-use / thinking streams can make assembly re-raise as APIError from inside the except-StopIteration handler, where the sibling except does not catch it, so it escaped __next__/__anext__ and dropped the request; recover best-effort usage from the raw chunks instead - Add a usage-only fallback for Anthropic streaming pass-through: when stream_chunk_builder returns None or raises, rebuild usage from the message_start / message_delta SSE events via AnthropicConfig.calculate_usage so cache, web-search and geo tokens are priced instead of left at $0 - Decode buffered pass-through bytes with errors="replace" so a stream cut mid-multibyte-sequence still logs the usage events already received - Record response_cost into model_call_details on the pass-through success path (it is read from there, not from kwargs), matching the gemini/cohere/openai handlers - Name the key (alias + masked key) in the virtual-key BudgetExceededError so operators don't have to reverse-map spend back to a key (cherry picked from commit b24b964e0482fe45d32bbffd379906b4464cd307) --- .../litellm_core_utils/streaming_handler.py | 54 ++- litellm/proxy/auth/auth_checks.py | 9 + .../anthropic_passthrough_logging_handler.py | 189 ++++++++++- .../base_passthrough_logging_handler.py | 3 + .../streaming_handler.py | 6 +- .../test_streaming_handler.py | 128 ++++++++ .../proxy/auth/test_auth_checks.py | 51 +++ ...t_anthropic_passthrough_logging_handler.py | 307 ++++++++++++++++++ .../test_streaming_handler_interrupt.py | 17 + 9 files changed, 745 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3d04b183176..5ed8cea9122 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1923,11 +1923,29 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages, - logging_obj=self.logging_obj, - ) + try: + complete_streaming_response = litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, + ) + except Exception as e: + # stream_chunk_builder can re-raise (as APIError) on large agentic + # streams. The raise originates inside this except-StopIteration block, + # so the sibling `except Exception` below does not catch it; it would + # escape __next__ and drop the request from SpendLogs. Recover + # best-effort usage from the raw chunks so cost is still tracked + verbose_logger.warning( + "stream_chunk_builder raised at end-of-stream (%s); logging " + "best-effort usage from chunks.", + str(e), + ) + try: + complete_streaming_response = self.model_response_creator( + chunk={"usage": calculate_total_usage(chunks=self.chunks)} + ) + except Exception: + complete_streaming_response = None response = self.model_response_creator() if complete_streaming_response is not None: @@ -2152,11 +2170,27 @@ class CustomStreamWrapper: except (StopAsyncIteration, StopIteration): if self.sent_last_chunk is True: # log the final chunk with accurate streaming values - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages, - logging_obj=self.logging_obj, - ) + try: + complete_streaming_response = litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, + ) + except Exception as e: + # see sync __next__: a raise from stream_chunk_builder inside this + # except handler escapes __anext__ and drops the request from SpendLogs. + # Recover best-effort usage from the raw chunks so cost is still tracked + verbose_logger.warning( + "stream_chunk_builder raised at end-of-stream (%s); logging " + "best-effort usage from chunks.", + str(e), + ) + try: + complete_streaming_response = self.model_response_creator( + chunk={"usage": calculate_total_usage(chunks=self.chunks)} + ) + except Exception: + complete_streaming_response = None response = self.model_response_creator() if complete_streaming_response is not None: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b1a109169c2..c07e510452d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3448,9 +3448,18 @@ async def _virtual_key_max_budget_check( # so a NaN max_budget would silently disable enforcement. Treat a # non-finite max_budget as "no configured limit" rather than as a bypass. if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: + # name the key in the error so operators don't have to reverse-map + # spend back to a key; key_name is the masked form (last 4 chars) + key_label = valid_token.key_alias or "key" + key_descriptor = ( + f"{key_label} ({valid_token.key_name})" + if valid_token.key_name + else key_label + ) raise litellm.BudgetExceededError( current_cost=spend, max_budget=valid_token.max_budget, + message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}", ) 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 22ed7b76d35..49aa8c0072c 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 @@ -6,6 +6,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger +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 from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -15,12 +16,19 @@ from litellm.llms.anthropic import get_anthropic_config from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) +from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) -from litellm.types.utils import LiteLLMBatch, ModelResponse, TextCompletionResponse +from litellm.types.utils import ( + Choices, + LiteLLMBatch, + Message, + ModelResponse, + TextCompletionResponse, +) if TYPE_CHECKING: from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType @@ -272,6 +280,9 @@ class AnthropicPassthroughLoggingHandler: kwargs["response_cost"] = response_cost kwargs["model"] = model + # the pass-through success path reads spend from + # model_call_details["response_cost"], not from kwargs + logging_obj.model_call_details["response_cost"] = response_cost passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = ( # type: ignore kwargs.get("passthrough_logging_payload") ) @@ -343,13 +354,42 @@ class AnthropicPassthroughLoggingHandler: if chunk_model: model = chunk_model - complete_streaming_response = ( - AnthropicPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=all_chunks, - litellm_logging_obj=litellm_logging_obj, - model=model, + try: + complete_streaming_response = ( + AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + ) ) - ) + except Exception as e: + # stream_chunk_builder re-raises assembly failures (as litellm.APIError) + # on large agentic tool-use / thinking streams; treat that the same as a + # None result so the usage-only fallback below still recovers cost + verbose_proxy_logger.warning( + "Anthropic passthrough: stream assembly raised (model=%s): %s; falling " + "back to usage-only cost from raw SSE events.", + model, + e, + ) + complete_streaming_response = None + if complete_streaming_response is None: + # stream_chunk_builder cannot always reassemble large agentic streams, but + # Anthropic still emits token usage in the message_start / message_delta SSE + # events regardless of content shape; recover usage-only so cost is tracked. + # Guard it too: a raise here would defeat the point and drop the request + try: + complete_streaming_response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=all_chunks, + model=model, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Anthropic passthrough: usage-only fallback failed (model=%s): %s", + model, + e, + ) + complete_streaming_response = None if complete_streaming_response is None: verbose_proxy_logger.error( "Unable to build complete streaming response for Anthropic passthrough endpoint, not logging..." @@ -636,6 +676,141 @@ class AnthropicPassthroughLoggingHandler: ) return complete_streaming_response + @staticmethod + def _extract_sse_data(event_str: str) -> Optional[dict]: + """Parse the JSON object from the ``data:`` line of an Anthropic SSE event.""" + for line in event_str.splitlines(): + stripped = line.strip() + if stripped.startswith("data:"): + payload = stripped[len("data:") :].strip() + if not payload or payload == "[DONE]": + return None + try: + return cast(dict, json.loads(payload)) + except (ValueError, TypeError): + return None + return None + + @staticmethod + def _build_usage_only_response_from_chunks( # noqa: PLR0915 + all_chunks: Sequence[Union[str, bytes]], + model: str, + ) -> Optional[ModelResponse]: + """ + Build a usage-bearing ModelResponse from Anthropic SSE token-usage events, for + cost tracking when stream_chunk_builder cannot reassemble the stream. + + Anthropic emits usage in ``message_start`` (uncached input + cache tokens, and an + initial output_tokens) and the final ``message_delta`` (cumulative output_tokens) + regardless of the content/tool shape, so cost is recoverable even when full + content assembly fails. Returns ``None`` if no usage event is found. + """ + input_tokens = 0 + cache_read = 0 + cache_creation = 0 + cache_creation_5m: Optional[int] = None + cache_creation_1h: Optional[int] = None + output_tokens = 0 + web_search_requests: Optional[int] = None + tool_search_requests: Optional[int] = None + inference_geo: Optional[str] = None + stop_reason: Optional[str] = None + found_usage = False + resolved_model = model + for _chunk_str in all_chunks: + for ( + event_str + ) in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events( + _chunk_str + ): + data = AnthropicPassthroughLoggingHandler._extract_sse_data(event_str) + if not data: + continue + event_type = data.get("type") + if event_type == "message_start": + message = data.get("message") or {} + if not resolved_model or resolved_model == "unknown": + resolved_model = message.get("model") or resolved_model + usage = message.get("usage") or {} + input_tokens = usage.get("input_tokens") or input_tokens + cache_read = usage.get("cache_read_input_tokens") or cache_read + cache_creation = ( + usage.get("cache_creation_input_tokens") or cache_creation + ) + _cc = usage.get("cache_creation") + if isinstance(_cc, dict): + cache_creation_5m = _cc.get("ephemeral_5m_input_tokens") + cache_creation_1h = _cc.get("ephemeral_1h_input_tokens") + if usage.get("inference_geo") is not None: + inference_geo = usage.get("inference_geo") + if usage.get("output_tokens") is not None: + output_tokens = usage.get("output_tokens") + found_usage = True + elif event_type == "message_delta": + _delta_stop = (data.get("delta") or {}).get("stop_reason") + if _delta_stop: + stop_reason = _delta_stop + usage = data.get("usage") or {} + if usage.get("output_tokens") is not None: + output_tokens = usage.get("output_tokens") + _stu = usage.get("server_tool_use") + if isinstance(_stu, dict): + if _stu.get("web_search_requests") is not None: + web_search_requests = _stu.get("web_search_requests") + if _stu.get("tool_search_requests") is not None: + tool_search_requests = _stu.get("tool_search_requests") + if usage.get("cache_read_input_tokens") is not None: + cache_read = usage.get("cache_read_input_tokens") + if usage.get("inference_geo") is not None: + inference_geo = usage.get("inference_geo") + found_usage = True + if not found_usage: + return None + # If only the 5m/1h split was provided, derive the cache_creation total from it. + if not cache_creation and (cache_creation_5m or cache_creation_1h): + cache_creation = (cache_creation_5m or 0) + (cache_creation_1h or 0) + # build usage via the same AnthropicConfig.calculate_usage path the success + # cases use, so prompt_tokens are cache-inclusive and cache / server_tool_use / + # inference_geo tokens are priced instead of left at $0 + usage_object: dict = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } + if cache_read: + usage_object["cache_read_input_tokens"] = cache_read + if cache_creation: + usage_object["cache_creation_input_tokens"] = cache_creation + if cache_creation_5m is not None or cache_creation_1h is not None: + usage_object["cache_creation"] = { + "ephemeral_5m_input_tokens": cache_creation_5m or 0, + "ephemeral_1h_input_tokens": cache_creation_1h or 0, + } + if web_search_requests is not None or tool_search_requests is not None: + _server_tool_use: dict = {} + if web_search_requests is not None: + _server_tool_use["web_search_requests"] = web_search_requests + if tool_search_requests is not None: + _server_tool_use["tool_search_requests"] = tool_search_requests + usage_object["server_tool_use"] = _server_tool_use + if inference_geo is not None: + usage_object["inference_geo"] = inference_geo + usage_obj = AnthropicConfig().calculate_usage( + usage_object=usage_object, reasoning_content=None + ) + return ModelResponse( + model=resolved_model, + choices=[ + Choices( + finish_reason=( + map_finish_reason(stop_reason) if stop_reason else "stop" + ), + index=0, + message=Message(role="assistant", content=""), + ) + ], + usage=usage_obj, + ) + @staticmethod def batch_creation_handler( # noqa: PLR0915 httpx_response: httpx.Response, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py index b9df8ecede3..a7ec2f0d368 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py @@ -116,6 +116,9 @@ class BasePassthroughLoggingHandler(ABC): kwargs["response_cost"] = response_cost kwargs["model"] = model + # the pass-through success path reads spend from + # model_call_details["response_cost"], not from kwargs + logging_obj.model_call_details["response_cost"] = response_cost passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = ( # type: ignore kwargs.get("passthrough_logging_payload") ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 33a6b719280..7a725472dd7 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -285,8 +285,10 @@ class PassThroughStreamingHandler: Returns: List of string lines, with each line being a complete data: {} chunk """ - # Combine all bytes and decode to string - combined_str = b"".join(raw_bytes).decode("utf-8") + # errors="replace" so a stream cut mid-multibyte-sequence (client disconnect) + # still decodes and logs the usage events already received, instead of raising + # and dropping the whole request from SpendLogs + combined_str = b"".join(raw_bytes).decode("utf-8", errors="replace") # Split by newlines and filter out empty lines lines = [line.strip() for line in combined_str.split("\n") if line.strip()] diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 09bb8532ec2..6988a2f50ad 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2194,3 +2194,131 @@ def test_record_partial_usage_for_failure_noop_without_chunks(): wrapper._record_partial_usage_for_failure() assert "combined_usage_object" not in logging_obj.model_call_details + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_chunk_builder_raise_at_end_of_stream_still_recovers_usage( + sync_mode, +): + """stream_chunk_builder re-raises (as APIError) on large agentic tool-use + streams. That raise originates inside the except-StopIteration handler, so + before the fix it escaped __next__/__anext__ and the request was dropped from + SpendLogs while the provider billed the tokens. The wrapper must catch it and + recover usage from the raw chunks so cost is still tracked.""" + final_usage_block = Usage( + completion_tokens=392, prompt_tokens=1799, total_tokens=2191 + ) + final_chunk = ModelResponseStream( + id="chatcmpl-raise-test", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="", role="assistant"), + ) + ], + usage=final_usage_block, + ) + test_chunks = bedrock_chunks + [final_chunk] + + logging_obj = Logging( + model="bedrock/claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="raise-test", + function_id="1245", + ) + + response = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=test_chunks), + model="bedrock/claude-haiku-4-5-20251001-v1:0", + custom_llm_provider="bedrock", + logging_obj=logging_obj, + stream_options={"include_usage": True}, + ) + + seen_usage = [] + with patch.object( + litellm, + "stream_chunk_builder", + side_effect=Exception("simulated assembly failure"), + ): + # before the fix this raised and dropped the request; it must not raise now + if sync_mode: + for chunk in response: + if getattr(chunk, "usage", None) is not None: + seen_usage.append(chunk.usage) + else: + async for chunk in response: + if getattr(chunk, "usage", None) is not None: + seen_usage.append(chunk.usage) + + assert any( + u.total_tokens == final_usage_block.total_tokens for u in seen_usage + ), "usage recovered from raw chunks was not emitted after stream_chunk_builder raised" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_chunk_builder_raise_and_usage_recovery_failure_does_not_crash( + sync_mode, +): + """If end-of-stream assembly raises AND best-effort usage recovery from the raw + chunks also fails, the stream must still complete cleanly rather than propagate + the exception to the consumer.""" + from litellm.litellm_core_utils import streaming_handler as sh_module + + final_chunk = ModelResponseStream( + id="chatcmpl-raise-recover-fail", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="", role="assistant"), + ) + ], + usage=Usage(completion_tokens=1, prompt_tokens=1, total_tokens=2), + ) + + response = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=bedrock_chunks + [final_chunk] + ), + model="bedrock/claude-haiku-4-5-20251001-v1:0", + custom_llm_provider="bedrock", + logging_obj=Logging( + model="bedrock/claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="raise-recover-fail", + function_id="1245", + ), + stream_options={"include_usage": True}, + ) + + with ( + patch.object( + litellm, "stream_chunk_builder", side_effect=Exception("assembly failed") + ), + patch.object( + sh_module, "calculate_total_usage", side_effect=Exception("recovery failed") + ), + ): + # must not raise even though both assembly and recovery fail + if sync_mode: + chunks = [c for c in response] + else: + chunks = [c async for c in response] + + assert len(chunks) > 0 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 116ba83f42e..d151e267b47 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -3370,3 +3370,54 @@ async def test_resolve_end_user_reraises_budget_exceeded( prisma_client=MagicMock(), user_api_key_cache=cache, ) + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_error_names_the_key(): + """BudgetExceededError for a virtual key must name the key (alias + masked key) + so operators don't have to reverse-map a spend figure back to a key.""" + valid_token = UserAPIKeyAuth( + token="hashed-token", + key_alias="payments-prod", + key_name="sk-...um_g", + max_budget=10.0, + spend=0.0, + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=25.0), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + + message = str(exc_info.value) + assert "payments-prod" in message + assert "sk-...um_g" in message + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_not_exceeded_does_not_raise(): + """Spend below the configured budget must not raise.""" + valid_token = UserAPIKeyAuth( + token="hashed-token", + key_alias="payments-prod", + max_budget=10.0, + spend=0.0, + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=1.0), + ): + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 224659f6079..f20fb4c40ec 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -1847,3 +1847,310 @@ class TestNonStreamingResponseRedaction: leaked = logging_obj.model_call_details.get("complete_streaming_response") assert leaked is None assert redacted.choices[0].message.content == "redacted-by-litellm" + + +def _sse_bytes(data: dict) -> bytes: + return f"event: {data['type']}\ndata: {json.dumps(data)}\n\n".encode() + + +class TestAnthropicUsageOnlyFallback: + """When stream_chunk_builder cannot reassemble a large/agentic stream (returns + None or raises), Anthropic still emits token usage in the message_start / + message_delta SSE events. The handler must recover usage-only so the request is + priced instead of being dropped from SpendLogs while Anthropic billed the tokens.""" + + _CHUNKS = [ + _sse_bytes( + { + "type": "message_start", + "message": { + "model": "claude-3-5-haiku-20241022", + "usage": { + "input_tokens": 100, + "cache_read_input_tokens": 40, + "cache_creation_input_tokens": 20, + "output_tokens": 1, + }, + }, + } + ), + _sse_bytes( + { + "type": "message_delta", + "usage": { + "output_tokens": 55, + "server_tool_use": {"web_search_requests": 2}, + }, + } + ), + ] + + def test_build_usage_only_recovers_cache_inclusive_usage(self): + response = ( + AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=self._CHUNKS, model="claude-3-5-haiku-20241022" + ) + ) + assert response is not None + usage = response.usage + # prompt_tokens must be cache-inclusive (input + cache_read + cache_creation) + assert usage.prompt_tokens == 160 + assert usage.completion_tokens == 55 + assert usage._cache_read_input_tokens == 40 + assert usage._cache_creation_input_tokens == 20 + assert usage.prompt_tokens_details.cached_tokens == 40 + assert usage.server_tool_use.web_search_requests == 2 + + def test_build_usage_only_returns_none_without_usage_events(self): + chunks = [_sse_bytes({"type": "content_block_delta", "delta": {"text": "hi"}})] + assert ( + AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=chunks, model="claude-3-5-haiku-20241022" + ) + is None + ) + + def test_build_usage_only_recovers_cache_split_server_tools_and_model(self): + # the model is "unknown" up-front and only the 5m/1h cache split is sent + # (no flat cache_creation_input_tokens); web/tool-search and geo arrive in + # message_delta. All must be recovered and priced, not left at $0. + chunks = [ + "event: ping\ndata: [DONE]\n\n", # ignored sentinel between real events + _sse_bytes( + { + "type": "message_start", + "message": { + "model": "claude-opus-4-6", + "usage": { + "input_tokens": 80, + "output_tokens": 1, + "cache_creation": { + "ephemeral_5m_input_tokens": 12, + "ephemeral_1h_input_tokens": 8, + }, + "inference_geo": "us", + }, + }, + } + ), + _sse_bytes( + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use"}, + "usage": { + "output_tokens": 40, + "cache_read_input_tokens": 5, + "inference_geo": "us", + "server_tool_use": { + "web_search_requests": 1, + "tool_search_requests": 3, + }, + }, + } + ), + ] + response = ( + AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=chunks, model="unknown" + ) + ) + assert response is not None + assert response.model == "claude-opus-4-6" + # the real stop_reason is surfaced, not a hardcoded "stop" + assert response.choices[0].finish_reason == "tool_calls" + usage = response.usage + # 80 input + 20 cache_creation (derived from 12+8) + 5 cache_read + assert usage.prompt_tokens == 105 + assert usage.completion_tokens == 40 + assert usage._cache_creation_input_tokens == 20 + assert usage._cache_read_input_tokens == 5 + assert usage.server_tool_use.web_search_requests == 1 + assert usage.server_tool_use.tool_search_requests == 3 + + @pytest.mark.parametrize( + "event_str,expected", + [ + ("data: [DONE]", None), + ("data: ", None), + ("data: {not-json", None), + ("event: ping", None), + ('data: {"a": 1}', {"a": 1}), + ], + ) + def test_extract_sse_data_handles_malformed_and_sentinel_lines( + self, event_str, expected + ): + assert ( + AnthropicPassthroughLoggingHandler._extract_sse_data(event_str) == expected + ) + + def _real_logging_obj(self): + from litellm.litellm_core_utils.litellm_logging import Logging as RealLoggingObj + + logging_obj = RealLoggingObj( + model="claude-3-5-haiku-20241022", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1", + ) + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + return logging_obj + + @patch("litellm.completion_cost") + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response" + ) + def test_handler_falls_back_when_assembly_returns_none( + self, mock_assemble, mock_cost + ): + mock_assemble.return_value = None + mock_cost.return_value = 0.0021 + logging_obj = self._real_logging_obj() + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-haiku-20241022", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=list(self._CHUNKS), + end_time=datetime.now(), + ) + + assert result["result"] is not None + assert result["result"].usage.completion_tokens == 55 + assert result["kwargs"]["response_cost"] == 0.0021 + + @patch("litellm.completion_cost") + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response" + ) + def test_handler_falls_back_when_assembly_raises(self, mock_assemble, mock_cost): + import litellm + + mock_assemble.side_effect = litellm.APIError( + status_code=500, + message="boom", + llm_provider="anthropic", + model="claude-3-5-haiku-20241022", + ) + mock_cost.return_value = 0.0021 + logging_obj = self._real_logging_obj() + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-haiku-20241022", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=list(self._CHUNKS), + end_time=datetime.now(), + ) + + # a raise from stream_chunk_builder must be treated like a None result, + # not propagate out and drop the request from SpendLogs + assert result["result"] is not None + assert result["result"].usage.completion_tokens == 55 + assert result["kwargs"]["response_cost"] == 0.0021 + + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response" + ) + def test_handler_returns_none_when_no_usage_recoverable(self, mock_assemble): + # assembly fails AND the chunks carry no usage event, so there is nothing + # to price; the handler must return None rather than fabricate a response + mock_assemble.return_value = None + logging_obj = self._real_logging_obj() + chunks = [_sse_bytes({"type": "content_block_delta", "delta": {"text": "hi"}})] + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-haiku-20241022", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=chunks, + end_time=datetime.now(), + ) + + assert result["result"] is None + assert result["kwargs"] == {} + + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_usage_only_response_from_chunks" + ) + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response" + ) + def test_handler_does_not_crash_when_usage_only_fallback_raises( + self, mock_assemble, mock_fallback + ): + # if the usage-only fallback itself raises, it must be treated as None and + # drop gracefully, not propagate out and crash the success handler + mock_assemble.return_value = None + mock_fallback.side_effect = Exception("fallback boom") + logging_obj = self._real_logging_obj() + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-haiku-20241022", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=list(self._CHUNKS), + end_time=datetime.now(), + ) + + assert result["result"] is None + assert result["kwargs"] == {} + + +class TestAnthropicResponseCostRecordedOnModelCallDetails: + """The pass-through success path reads spend from + model_call_details["response_cost"], not from kwargs, so the streaming payload + builder must record it there or streaming pass-through logs $0.""" + + def test_create_payload_records_response_cost_on_model_call_details(self): + from litellm.types.utils import Choices, Message, ModelResponse + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.get_router_model_id.return_value = None + logging_obj.litellm_params = {} + logging_obj.litellm_call_id = "test-call-id" + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="hello", role="assistant"), + ) + ], + created=1234567890, + model="claude-3-7-sonnet-20250219", + usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + ) + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-7-sonnet-20250219", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert ( + logging_obj.model_call_details["response_cost"] == kwargs["response_cost"] + ) + assert logging_obj.model_call_details["response_cost"] > 0 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index f73aee77cc1..38990644154 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -118,3 +118,20 @@ async def test_chunk_processor_does_not_schedule_logging_when_no_chunks(): assert received == [] mock_route.assert_not_called() + + +def test_convert_raw_bytes_survives_truncated_multibyte_sequence(): + """A stream cut mid-multibyte-sequence (client disconnect) must still decode + via errors="replace" so the usage events already received are logged, instead + of raising UnicodeDecodeError and dropping the whole request from SpendLogs.""" + # the 3-byte "☃" (E2 98 83) is cut after 2 bytes, leaving an invalid sequence + # that strict utf-8 decode would raise on, discarding the message_delta line too + truncated_codepoint = "☃".encode("utf-8")[:2] + raw_bytes = [ + b'data: {"text": "' + truncated_codepoint, + b'\ndata: {"type": "message_delta"}\n', + ] + + lines = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes) + + assert any('"type": "message_delta"' in line for line in lines) From 36031e6c39e4e02f7c14ab2b2b370da419f5160b Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 23 Jun 2026 17:51:25 -0700 Subject: [PATCH 5/5] fix(docker): bump wolfi-base digest to patch openssl CVE-2026-34182 (#31133) Re-pins LITELLM_BUILD_IMAGE and LITELLM_RUNTIME_IMAGE across all 6 Dockerfiles from the prior digests (openssl 3.6.2-r3) to the current chainguard wolfi-base digest c61ac691 (openssl 3.6.3-r2, >= the fixed 3.6.3-r0). The runtime stage is the shipped image, so the runtime digest is what actually resolves the customer-facing CVE; the build image is bumped too for hygiene. Two Dockerfiles tracked a second equally-stale digest; both are unified onto the patched one. (cherry picked from commit fda08dd727aabe50582191e31ab239a811cda3a0) --- Dockerfile | 4 ++-- backend/Dockerfile | 4 ++-- docker/Dockerfile.database | 4 ++-- docker/Dockerfile.non_root | 4 ++-- gateway/Dockerfile | 4 ++-- migrations/Dockerfile | 4 ++-- 6 files changed, 12 insertions(+), 12 deletions(-) diff --git a/Dockerfile b/Dockerfile index 9ad9ab31b65..68d7b14d19f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,8 +1,8 @@ # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/backend/Dockerfile b/backend/Dockerfile index c08014fc0ef..d969be69a20 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index c84003a065f..94e53bafdd9 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,8 +1,8 @@ # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 2729babb6d6..9304dd3784f 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,6 +1,6 @@ # Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG PROXY_EXTRAS_SOURCE=published ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a diff --git a/gateway/Dockerfile b/gateway/Dockerfile index a2ca3d3f83f..041fd11678b 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 2160514251a..6e79922a97a 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin