From 47d53cad6de63c69d2af3eebf088eeba7a1ccc91 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 09:54:16 -0700 Subject: [PATCH] fix(accounting): avoid adjusting stamped streaming costs twice --- litellm/proxy/spend_tracking/cost_compiler.py | 4 +- .../spend_tracking/test_cost_compiler.py | 52 ++++++++++++++++++- 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/spend_tracking/cost_compiler.py b/litellm/proxy/spend_tracking/cost_compiler.py index 3580d2854e2..5a801039b50 100644 --- a/litellm/proxy/spend_tracking/cost_compiler.py +++ b/litellm/proxy/spend_tracking/cost_compiler.py @@ -136,7 +136,7 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: } } ], - usage=Usage.model_validate(captured.usage or {}), + usage=Usage.model_validate({name: value for name, value in (captured.usage or {}).items() if name != "cost"}), service_tier=response_inputs.get("service_tier"), meta={"billed_units": response_inputs.get("billed_units")}, hidden_params={ @@ -182,7 +182,7 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: custom_llm_provider=provider, litellm_params={**deployment, **request, "metadata": {"model_info": _fields(pricing.get("model_info"))}}, additional_response_cost=request.get("additional_response_cost"), - usage=Usage.model_validate(captured.usage or {}), + usage=Usage.model_validate({name: value for name, value in (captured.usage or {}).items() if name != "cost"}), end_time=event.ended_at, mcp_tool_call_metadata={ "name": _fields(pricing.get("tool")).get("name"), diff --git a/tests/unit/proxy/spend_tracking/test_cost_compiler.py b/tests/unit/proxy/spend_tracking/test_cost_compiler.py index bd093d6e012..f5e35ec7efc 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_compiler.py +++ b/tests/unit/proxy/spend_tracking/test_cost_compiler.py @@ -149,8 +149,16 @@ def test_retained_token_pricing_preserves_cache_reasoning_and_service_tier() -> assert float(result.total_cost) == pytest.approx(60 * 0.03 + 40 * 0.002 + 15 * 0.05 + 5 * 0.07) -def test_upstream_reported_price_takes_precedence_without_custom_pricing() -> None: - result: Final = price_usage(json.dumps(_retained_envelope(logging={"provider_reported_cost": 0.037}))) +@pytest.mark.parametrize("stamped_cost", [None, 0.9]) +def test_upstream_reported_price_takes_precedence_without_custom_pricing(stamped_cost: float | None) -> None: + result: Final = price_usage( + json.dumps( + _retained_envelope( + logging={"provider_reported_cost": 0.037}, + usage={"prompt_tokens": 10, "completion_tokens": 4, "cost": stamped_cost}, + ) + ) + ) assert isinstance(result, PricingSnapshot) assert float(result.total_cost) == pytest.approx(0.037) @@ -847,3 +855,43 @@ def test_explicit_zero_retained_token_rates_are_known_free() -> None: assert isinstance(result, PricingSnapshot) assert result.total_cost == Decimal("0") assert result.pricing_source == "cost_map" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("completion_tokens", [1, 4]) +async def test_stream_replay_applies_adjustments_once_to_legacy_stamped_usage( + monkeypatch: pytest.MonkeyPatch, completion_tokens: int +) -> None: + import litellm + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from litellm.types.utils import ModelResponse, Usage + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.05}) + monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"percentage": 0.1, "fixed_amount": 0.001}}) + retained: Final = _capture_logger().request_pricing + assert retained is not None + monkeypatch.setitem(litellm.model_cost, "openai/capture-model", retained.models()["capture-model"]) + monkeypatch.setitem(litellm.model_cost, "capture-model", retained.models()["capture-model"]) + logger: Final = _logger(stream=True) + response: Final = ModelResponse( + id="terminal", model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=completion_tokens) + ) + expected: Final = _logger(stream=True).response_cost_calculator(response, ignore_stamped_cost=True) + assert expected is not None + stamped: Final = response.model_copy( + update={"usage": Usage(prompt_tokens=10, completion_tokens=completion_tokens, cost=expected)} + ) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, logger, stamped, _AT, _AT) + await meter.close() + priced: Final = price_usage(database.rows[-1].envelope) + assert isinstance(priced, PricingSnapshot) + assert float(priced.total_cost) == pytest.approx(expected) + assert float(priced.original_cost or 0) == pytest.approx(0.1 + completion_tokens * 0.02) + finally: + await meter.stop(1)