From 78ed5126a5089ac37e4631aeeb732b09045db612 Mon Sep 17 00:00:00 2001
From: Cesar Garcia <128240629+Chesars@users.noreply.github.com>
Date: Tue, 4 Nov 2025 20:57:59 -0300
Subject: [PATCH] fix: Fix Responses API streaming tests usage field names and
 cost (#16236)

This commit fixes two bugs in Responses API streaming tests:

1. **Usage field naming bug**: Tests were using `input_tokens` and
   `output_tokens` but the Usage object uses `prompt_tokens` and
   `completion_tokens`.

2. **Missing cost in streaming usage**: When `include_cost_in_streaming_usage`
   was enabled, the cost was calculated and added to ResponseAPIUsage, but was
   lost during the transformation to the Usage object.

Changes:
- Updated test assertions to use correct field names (prompt_tokens, completion_tokens)
- Added cost preservation logic in FakeStreamerResponsesAPIIterator
- Modified _transform_response_api_usage_to_chat_usage() to preserve cost attribute

All streaming tests now pass successfully.
---
 litellm/responses/streaming_iterator.py          | 16 ++++++++++++++++
 litellm/responses/utils.py                       |  8 +++++++-
 .../base_responses_api.py                        | 12 ++++++------
 3 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py
index b78913402ff..3754959f6ab 100644
--- a/litellm/responses/streaming_iterator.py
+++ b/litellm/responses/streaming_iterator.py
@@ -325,6 +325,22 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
             for i in range(0, len(full_text), self.CHUNK_SIZE)
         ]
 
+        # Add cost to usage object if include_cost_in_streaming_usage is True
+        if litellm.include_cost_in_streaming_usage and logging_obj is not None:
+            usage_obj: Optional[ResponseAPIUsage] = getattr(
+                transformed, "usage", None
+            )
+            if usage_obj is not None:
+                try:
+                    cost: Optional[float] = logging_obj._response_cost_calculator(
+                        result=transformed
+                    )
+                    if cost is not None:
+                        setattr(usage_obj, "cost", cost)
+                except Exception:
+                    # If cost calculation fails, continue without cost
+                    pass
+
         # append the completed event
         self._events = deltas + [
             ResponseCompletedEvent(
diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py
index fdbc692a977..64db4d3e020 100644
--- a/litellm/responses/utils.py
+++ b/litellm/responses/utils.py
@@ -381,9 +381,15 @@ class ResponseAPILoggingUtils:
                 cached_tokens=response_api_usage.input_tokens_details.cached_tokens,
                 audio_tokens=response_api_usage.input_tokens_details.audio_tokens,
             )
-        return Usage(
+        usage = Usage(
             prompt_tokens=prompt_tokens,
             completion_tokens=completion_tokens,
             total_tokens=prompt_tokens + completion_tokens,
             prompt_tokens_details=prompt_tokens_details,
         )
+
+        # Preserve cost attribute if it exists on ResponseAPIUsage
+        if hasattr(response_api_usage, "cost") and response_api_usage.cost is not None:
+            setattr(usage, "cost", response_api_usage.cost)
+
+        return usage
diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py
index f68b373feae..0941e064a7f 100644
--- a/tests/llm_responses_api_testing/base_responses_api.py
+++ b/tests/llm_responses_api_testing/base_responses_api.py
@@ -189,12 +189,12 @@ class BaseResponsesAPITest(ABC):
             response_completed_event.response.usage,
         )
         assert (
-            response_completed_event.response.usage.input_tokens > 0
-            and response_completed_event.response.usage.input_tokens < 100
+            response_completed_event.response.usage.prompt_tokens > 0
+            and response_completed_event.response.usage.prompt_tokens < 100
         )
         assert (
-            response_completed_event.response.usage.output_tokens > 0
-            and response_completed_event.response.usage.output_tokens < 2000
+            response_completed_event.response.usage.completion_tokens > 0
+            and response_completed_event.response.usage.completion_tokens < 2000
         )
         assert (
             response_completed_event.response.usage.total_tokens > 0
@@ -204,8 +204,8 @@ class BaseResponsesAPITest(ABC):
         # total tokens should be the sum of input and output tokens
         assert (
             response_completed_event.response.usage.total_tokens
-            == response_completed_event.response.usage.input_tokens
-            + response_completed_event.response.usage.output_tokens
+            == response_completed_event.response.usage.prompt_tokens
+            + response_completed_event.response.usage.completion_tokens
         )
 
         # assert the response completed event includes cost when include_cost_in_streaming_usage is True