From 78f3cd9d3e6cc365d9635fb2a74509dd016f23e4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 12:11:29 -0700 Subject: [PATCH] fix(spend): keep the dropped-cost flag out of reach of provider usage keys A provider usage key named dropped_non_finite_cost shadowed the Usage method on the instance, so the stream builder raised and dropped the request's spend. The builder now reads the flag through the class, and Usage sets the flag after copying provider keys, so a _dropped_non_finite_cost key cannot set it either. --- .../streaming_chunk_builder_utils.py | 2 +- litellm/types/utils.py | 10 ++++++---- .../test_streaming_chunk_builder_utils.py | 16 ++++++++++++++++ 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0fa9ef55aa5..10fc77f7b04 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -991,7 +991,7 @@ class ChunkProcessor: prompt_tokens_details, cache_creation_token_details ) - if usage_chunk.dropped_non_finite_cost(): + if Usage.dropped_non_finite_cost(usage_chunk): cost = None elif usage_chunk_dict["cost"] is not None: cost = usage_chunk_dict["cost"] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2b960ff2f0e..21a43f9cad8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2076,11 +2076,11 @@ class Usage(SafeAttributeModel, CompletionUsage): else: # maintain openai compatibility in usage object if possible del self.server_tool_use + dropped_non_finite_cost: Final = cost is not None and _carries_non_finite_number(cost) if cost is None: del self.cost - elif _carries_non_finite_number(cost): + elif dropped_non_finite_cost: verbose_logger.warning("Dropping a non-finite provider usage cost") - self._dropped_non_finite_cost = True del self.cost else: self.cost = cost @@ -2098,9 +2098,11 @@ class Usage(SafeAttributeModel, CompletionUsage): for k, v in params.items(): setattr(self, k, v) + self._dropped_non_finite_cost = dropped_non_finite_cost - def dropped_non_finite_cost(self) -> bool: - return self._dropped_non_finite_cost + @staticmethod + def dropped_non_finite_cost(usage: "Usage") -> bool: + return usage._dropped_non_finite_cost def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 88e21bb262d..6125b30d984 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1102,6 +1102,22 @@ def test_a_non_finite_final_stream_cost_does_not_fall_back_to_an_earlier_chunk_c assert getattr(response.usage, "cost", None) is None assert response.usage.total_tokens == 27 + +@pytest.mark.parametrize("provider_key", ["dropped_non_finite_cost", "_dropped_non_finite_cost"]) +def test_a_provider_usage_key_named_like_the_dropped_cost_flag_keeps_the_stream_cost(provider_key: str): + chunk: Final = ModelResponseStream( + id="chatcmpl-flag-key", + created=1745513206, + model="sonar", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="ok"))], + usage={"prompt_tokens": 26, "completion_tokens": 1, "total_tokens": 27, "cost": 0.00371, provider_key: True}, + ) + + response: Final = stream_chunk_builder(chunks=[chunk]) + + assert response is not None + assert getattr(response.usage, "cost", None) == 0.00371 + def test_stream_chunk_builder_tolerates_trailing_chunk_without_choices(): """Regression for https://github.com/BerriAI/litellm/issues/32051