From 5e8d58dff583bbed7475f0e47380b6cf5c238af0 Mon Sep 17 00:00:00 2001 From: mikamikasuki <211269698+mikamikasuki@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:59:20 -0700 Subject: [PATCH] fix(proxy): omit incomplete stream cost headers --- litellm/proxy/common_request_processing.py | 45 ++++++++++++------- .../proxy/test_common_request_processing.py | 42 +++++++++++++++++ 2 files changed, 70 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6fda5780f63..76064baa5f7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2678,25 +2678,36 @@ class ProxyBaseLLMRequestProcessing: version: str | None, callback_headers: Mapping[str, str], ) -> Mapping[str, str]: - """The streaming response headers describing `hidden_params`' deployment.""" + """Build stream headers without publishing a cost before the stream finishes.""" + raw_cache_key: Final[object] = hidden_params.get("cache_key") + raw_api_base: Final[object] = hidden_params.get("api_base") + raw_fastest_response_batch_completion: Final[object] = hidden_params.get("fastest_response_batch_completion") + raw_additional_headers: Final[object] = hidden_params.get("additional_headers") + headers: Final[Mapping[str, str]] = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=logging_obj.litellm_call_id, + model_id=self._get_model_id_from_response(hidden_params, self.data), + cache_key=raw_cache_key if isinstance(raw_cache_key, str) else "", + api_base=raw_api_base if isinstance(raw_api_base, str) else "", + version=version, + response_cost=None, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=( + raw_fastest_response_batch_completion + if isinstance(raw_fastest_response_batch_completion, bool) + else None + ), + request_data=self.data, + hidden_params=hidden_params, + litellm_logging_obj=logging_obj, + **(raw_additional_headers if isinstance(raw_additional_headers, Mapping) else MappingProxyType({})), + ) + merged_headers: Final[Mapping[str, str]] = {**headers, **callback_headers} return MappingProxyType( { - **ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - call_id=logging_obj.litellm_call_id, - model_id=self._get_model_id_from_response(hidden_params, self.data), - cache_key=hidden_params.get("cache_key") or "", - api_base=hidden_params.get("api_base") or "", - version=version, - response_cost=_response_cost_header_value(hidden_params), - model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"), - request_data=self.data, - hidden_params=hidden_params, - litellm_logging_obj=logging_obj, - **(hidden_params.get("additional_headers") or MappingProxyType({})), - ), - **callback_headers, + key: value + for key, value in merged_headers.items() + if not key.lower().startswith("x-litellm-response-cost") } ) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index a226ce8a4a9..2a509f81bdb 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -1353,6 +1353,48 @@ class TestProxyBaseLLMRequestProcessing: assert "x-litellm-response-cost-discount-amount" in headers assert float(headers["x-litellm-response-cost-discount-amount"]) == 0.000005 + def test_stream_response_headers_omit_incomplete_cost(self): + processing = ProxyBaseLLMRequestProcessing.__new__(ProxyBaseLLMRequestProcessing) + processing.data = {} + user = MagicMock(spec=UserAPIKeyAuth) + user.tpm_limit = None + user.rpm_limit = None + user.max_budget = None + user.spend = 0 + logging_obj = LiteLLMLoggingObj( + model="gpt-4", + messages=[{"role": "user", "content": "test"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + logging_obj.set_cost_breakdown( + input_cost=0.0, + output_cost=0.0, + total_cost=0.001, + cost_for_built_in_tools_cost_usd_dollar=0.0, + original_cost=0.0, + margin_percent=0.1, + margin_total_amount=0.001, + ) + + headers = processing._stream_response_headers( + hidden_params={"response_cost": 0.001}, + user_api_key_dict=user, + logging_obj=logging_obj, + version=None, + callback_headers={ + "x-litellm-response-cost": "0.067", + "x-custom-header": "custom-value", + }, + ) + + assert "x-litellm-call-id" in headers + assert headers["x-custom-header"] == "custom-value" + assert not any(key.startswith("x-litellm-response-cost") for key in headers) + def test_get_custom_headers_without_discount_info(self): """ Test that when no discount is applied, discount headers are not included.