From aa617a9f09ade41ce23fac5862a873a720a1c656 Mon Sep 17 00:00:00 2001 From: Ben Langfeld Date: Wed, 16 Sep 2026 21:39:50 -0400 Subject: [PATCH] Make the usage.cost call site one testable entry point The call site was a two-line conditional inside _process_llm_request, so the injecting branch was only reachable through a full request and went uncovered. Collapsing it into _maybe_set_usage_cost leaves one unconditional line at the call site and puts the opt-in decision next to the write, where both branches can be exercised directly. Co-Authored-By: Claude Opus 5 --- litellm/proxy/common_request_processing.py | 14 ++++++++++++-- .../proxy/test_common_request_processing.py | 12 ++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 23dc64f0c1a..fe873f2c2c1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2823,8 +2823,7 @@ class ProxyBaseLLMRequestProcessing: # Same value the x-litellm-response-cost header carries below, so the body and # the header cannot report different costs for one request. - if self._should_include_cost_in_usage(request): - self._set_usage_cost(response, response_cost_for_headers) + self._maybe_set_usage_cost(request, response, response_cost_for_headers) # Always return the client-requested model name (not provider-prefixed internal identifiers) # for OpenAI-compatible responses. @@ -4097,6 +4096,17 @@ class ProxyBaseLLMRequestProcessing: return cost_from_logging_obj return ProxyBaseLLMRequestProcessing._completion_cost_or_none(model_response, model_name, service_tier) + @staticmethod + def _maybe_set_usage_cost(request: Request, response: object, cost: float | str | None) -> None: + """ + Record the gateway's cost on ``usage.cost``, if this request asked for it. + + Kept as one entry point so the decision and the write are exercised together + rather than only through a full request. + """ + if ProxyBaseLLMRequestProcessing._should_include_cost_in_usage(request): + ProxyBaseLLMRequestProcessing._set_usage_cost(response, cost) + @staticmethod def _should_include_cost_in_usage(request: Request) -> bool: """ diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 2c298413586..55948c92ff8 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -9384,6 +9384,18 @@ class TestIncludeCostInUsage: response: Final = _response_with_usage(prompt_tokens=11, completion_tokens=5, total_tokens=16) assert "cost" not in response.model_dump()["usage"] + def test_opted_in_request_gets_the_cost_recorded(self): + """The decision and the write, exercised together as the request path runs them.""" + response: Final = _response_with_usage(prompt_tokens=11, completion_tokens=5, total_tokens=16) + request: Final = _request_with_headers(x_litellm_include_cost_in_usage="true") + ProxyBaseLLMRequestProcessing._maybe_set_usage_cost(request, response, 5.85e-06) + assert response.model_dump()["usage"]["cost"] == 5.85e-06 + + def test_opted_out_request_is_left_untouched(self): + response: Final = _response_with_usage(prompt_tokens=11, completion_tokens=5, total_tokens=16) + ProxyBaseLLMRequestProcessing._maybe_set_usage_cost(_request_with_headers(), response, 5.85e-06) + assert "cost" not in response.model_dump()["usage"] + def test_a_response_without_usage_is_left_alone(self): sentinel: Final = SimpleNamespace(usage=None) ProxyBaseLLMRequestProcessing._set_usage_cost(sentinel, 5.85e-06)