fix(cost): keep upstream_usage_cost on the Responses stream logging copy so bridged /v1/messages rows price their tokens (#49365)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/responses/streaming_iterator.py

@@ -41,6 +41,7 @@ from litellm.constants import (
 from litellm.exceptions import MidStreamFallbackError, RateLimitError, ServiceUnavailableError
 from litellm.litellm_core_utils.asyncify import run_async_function
 from litellm.litellm_core_utils.core_helpers import (
+    UPSTREAM_USAGE_COST,
     process_response_headers,
     provider_reported_usage_cost,
     set_response_cost_in_hidden_params,
@@ -184,6 +185,7 @@ def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams:
 _SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache"
 _UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text"
 _DECODED_JSON: Final = TypeAdapter(object)
+_HIDDEN_PARAMS: Final = TypeAdapter(dict[str, object])
 
 
 def _load_json_object(payload: str | bytes) -> dict[str, object]:
@@ -693,7 +695,9 @@ class BaseResponsesAPIStreamingIterator:
         source_hidden: Final[object] = getattr(
             getattr(self.completed_response, "response", None), "_hidden_params", None
         )
-        source: Final[Mapping[str, object]] = source_hidden if isinstance(source_hidden, Mapping) else EMPTY_MAPPING
+        source: Final[Mapping[str, object]] = (
+            _HIDDEN_PARAMS.validate_python(source_hidden) if isinstance(source_hidden, Mapping) else {}
+        )
         processed: Final[object] = source.get("additional_headers") or self._hidden_params.get("additional_headers")
         raw: Final[object] = source.get("headers") or self._raw_response_headers
         headers: Final[Mapping[str, object]] = processed if isinstance(processed, Mapping) else EMPTY_MAPPING
@@ -703,6 +707,7 @@ class BaseResponsesAPIStreamingIterator:
         target.hidden_params = {
             "additional_headers": {**headers},
             "headers": {**raw_headers},
+            **_stamped_bill_hidden_params(source),
             **existing,
         }
 
@@ -1749,6 +1754,11 @@ def _dump_json_safely(response: BaseModel) -> str | None:
         return None
 
 
+def _stamped_bill_hidden_params(source: Mapping[str, object]) -> Mapping[str, object]:
+    stamped: Final = UPSTREAM_USAGE_COST in source
+    return {key: source[key] for key in (UPSTREAM_USAGE_COST, "response_cost") if stamped and key in source}
+
+
 def _logging_copy(event: object) -> object:
     """Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never
     reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the

tests/integration/cost_calculation/test_messages_bridged_stream_margin_discount.py

@@ -0,0 +1,123 @@
+"""A streamed `/v1/messages` request served through the Responses bridge bills the same row as the plain request
+when a cost margin or a cost discount is configured. The streamed terminal event carries litellm's own bill in
+`usage.cost`, so the spend row must price the tokens the way the plain leg does instead of reading that bill back
+as the provider's charge, which books input 0 / output total and applies the margin or discount a second time.
+
+The stream_parity runner checks every leg against the suite's unadjusted expected row on the shared proxy, and a
+margin or discount set there would reprice every other case running beside it, so each pricing boots its own proxy."""
+
+from collections.abc import Iterator, Mapping
+from contextlib import ExitStack
+from dataclasses import dataclass
+from hashlib import sha256
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+import yaml
+from integration._support.client import Gateway, gateway_from_environment
+from integration._support.process import OwnedProxy, owned_proxy_process
+from integration._support.workers import cost_map_environment
+from integration.cost_calculation.cases.openai import (
+    GPT_6_LUNA_MESSAGES_BRIDGED_JSON,
+    GPT_6_LUNA_MESSAGES_BRIDGED_STREAM,
+)
+from integration.cost_calculation.conftest import CostRow, poll_rows, register_scenario_deployment
+from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase, ExactExpected
+from pydantic import JsonValue, TypeAdapter
+
+pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state]
+
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_STOCK_CONFIG: Final = Path("tests/integration/proxy_config.yaml")
+_LEGS: Final[tuple[CostTrackingTestCase, ...]] = (GPT_6_LUNA_MESSAGES_BRIDGED_JSON, GPT_6_LUNA_MESSAGES_BRIDGED_STREAM)
+
+
+@dataclass(frozen=True, slots=True)
+class _Pricing:
+    name: str
+    litellm_settings: Mapping[str, JsonValue]
+    spend_multiplier: float
+
+
+PRICINGS: Final = (
+    _Pricing("margin", MappingProxyType({"cost_margin_config": {"global": 0.5}}), 1.5),
+    _Pricing("discount", MappingProxyType({"cost_discount_config": {"openai": 0.1}}), 0.9),
+)
+
+
+def _config(directory: Path, pricing: _Pricing) -> Path:
+    stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(_STOCK_CONFIG.read_text()))
+    settings: Final = _JSON_OBJECT.validate_python(stock["litellm_settings"])
+    path: Final = directory / f"{pricing.name}-config.yaml"
+    path.write_text(yaml.safe_dump({**stock, "litellm_settings": {**settings, **pricing.litellm_settings}}))
+    return path
+
+
+@pytest.fixture(scope="module")
+def proxies(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Mapping[str, OwnedProxy]]:
+    with gateway_from_environment() as gateway, ExitStack() as stack:
+        overrides: Final = cost_map_environment(gateway.upstream_url)
+
+        def boot(pricing: _Pricing) -> OwnedProxy:
+            directory: Final = tmp_path_factory.mktemp(f"messages-bridged-stream-{pricing.name}-rig")
+            return stack.enter_context(
+                owned_proxy_process(
+                    gateway,
+                    directory,
+                    overrides,
+                    config=_config(directory, pricing),
+                    remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",),
+                )
+            )
+
+        yield MappingProxyType({pricing.name: boot(pricing) for pricing in PRICINGS})
+
+
+def _bill_legs(gateway: Gateway) -> Mapping[str, CostRow]:
+    with gateway.scenario() as scenario:
+        key: Final = scenario.key()
+        deployments: Final = {
+            leg.name: register_scenario_deployment(scenario, leg, sha256(leg.name.encode()).hexdigest()[:12], key)
+            for leg in _LEGS
+        }
+        responses: Final = {
+            leg.name: gateway.request(
+                "POST", leg.endpoint, {**leg.request, "model": deployments[leg.name].model_name}, key=key
+            )
+            for leg in _LEGS
+        }
+        for leg in _LEGS:
+            assert responses[leg.name].status_code == 200, (leg.name, responses[leg.name].text[:400])
+        rows: Final = poll_rows(key, len(_LEGS))
+        return MappingProxyType(
+            {leg.name: next(row for row in rows if row.model_id == deployments[leg.name].identity) for leg in _LEGS}
+        )
+
+
+def _priced(row: CostRow) -> tuple[float | None, float | None, float | None, int | None, int | None]:
+    breakdown: Final = row.breakdown
+    assert breakdown is not None, row
+    return (row.spend, breakdown.input_cost, breakdown.output_cost, row.prompt_tokens, row.completion_tokens)
+
+
+@pytest.mark.parametrize("pricing", PRICINGS, ids=[pricing.name for pricing in PRICINGS])
+def test_streamed_messages_bridged_row_matches_the_plain_row_under_a_margin_or_discount(
+    proxies: Mapping[str, OwnedProxy], pricing: _Pricing
+) -> None:
+    billed: Final = _bill_legs(proxies[pricing.name].gateway)
+    plain: Final = billed[GPT_6_LUNA_MESSAGES_BRIDGED_JSON.name]
+    streamed: Final = billed[GPT_6_LUNA_MESSAGES_BRIDGED_STREAM.name]
+    expected: Final = GPT_6_LUNA_MESSAGES_BRIDGED_JSON.expected
+    assert isinstance(expected, ExactExpected), expected
+    assert _priced(plain) == pytest.approx(
+        (
+            expected.spend * pricing.spend_multiplier,
+            expected.input_cost,
+            expected.output_cost,
+            expected.prompt_tokens,
+            expected.completion_tokens,
+        )
+    ), plain
+    assert _priced(streamed) == pytest.approx(_priced(plain)), (plain, streamed)

tests/unit/responses/test_streaming_iterator.py

@@ -22,6 +22,7 @@ from pydantic_core import PydanticSerializationError
 import litellm
 from litellm.exceptions import MidStreamFallbackError
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.core_helpers import UPSTREAM_USAGE_COST, billed_usage_cost
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
 from litellm.responses.streaming_iterator import (
@@ -955,7 +956,7 @@ def _capture_dispatch(logged: list):
     return _dispatch
 
 
-def _headers_config(*, transform_hidden_params: Optional[dict] = None) -> Mock:
+def _headers_config(*, transform_hidden_params: dict | None = None, usage: dict | None = None) -> Mock:
     """Config whose completed event carries a real ResponsesAPIResponse, so the logging copy
     performs a genuine model_dump/model_validate round trip."""
     mock_config = Mock(spec=BaseResponsesAPIConfig)
@@ -973,6 +974,7 @@ def _headers_config(*, transform_hidden_params: Optional[dict] = None) -> Mock:
             parallel_tool_calls=False,
             tool_choice="auto",
             tools=[],
+            usage=usage,
         )
         if transform_hidden_params is not None:
             response._hidden_params.update(transform_hidden_params)
@@ -1070,6 +1072,36 @@ async def test_streaming_logging_copy_preserves_transform_hidden_params():
     assert "response_cost" not in hidden_params
 
 
+@pytest.mark.asyncio
+async def test_streaming_logging_copy_keeps_the_stamped_bill_as_litellms_own():
+    """The terminal event's usage.cost is litellm's bill. The logging copy must carry upstream_usage_cost
+    beside it, or the spend row reads that bill back as an upstream charge (input 0, output total)."""
+    logging_obj = _logging_obj_stub()
+    logging_obj.response_cost_calculator.return_value = 0.00039
+    logged: list[object] = []
+    logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
+    logging_obj._on_deferred_stream_complete = None
+
+    iterator = _make_header_iterator(
+        headers={},
+        config=_headers_config(usage={"input_tokens": 1840, "output_tokens": 412, "total_tokens": 2252}),
+        logging_obj=logging_obj,
+    )
+    async for _ in iterator:
+        pass
+
+    assert len(logged) == 1
+    event: Final = logged[0]
+    assert isinstance(event, ResponseCompletedEvent), event
+    logging_response: Final = event.response
+    assert logging_response is not iterator.completed_response.response
+    assert logging_response.usage is not None
+    assert logging_response.usage.cost == 0.00039
+    assert logging_response._hidden_params[UPSTREAM_USAGE_COST] is None
+    assert logging_response._hidden_params["response_cost"] == 0.00039
+    assert billed_usage_cost(logging_response, "openai") is None
+
+
 @pytest.mark.asyncio
 async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched():
     """LIT-6055: when the logging copy falls back to the original event, the header restore must