From c94a4f8696fe63dcbe42b6e25ad3cf1e9cddf8f4 Mon Sep 17 00:00:00 2001
From: kerry <kerry@berri.ai>
Date: Tue, 15 Sep 2026 22:49:57 +0000
Subject: [PATCH] test: restore xai reported-cost passthrough tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../test_xai_responses_transformation.py      | 38 ++++++++++++++
 .../llms/xai/test_xai_chat_transformation.py  | 49 +++++++++++++++++++
 2 files changed, 87 insertions(+)

diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py
index 994dea1b877..e3feb7d5342 100644
--- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py
+++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py
@@ -11,6 +11,7 @@ from unittest.mock import MagicMock, Mock
 
 import httpx
 
+from litellm.llms.xai.cost_calculator import cost_per_token
 from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig
 from litellm.responses.utils import ResponseAPILoggingUtils
 from litellm.types.llms.openai import (
@@ -484,6 +485,43 @@ class TestXAIResponsesReportedCost:
         )
         return response.usage
 
+    def test_reported_cost_reaches_the_cost_calculator(self):
+        usage = self._transformed_usage(
+            {
+                "input_tokens": 100,
+                "output_tokens": 200,
+                "total_tokens": 300,
+                "cost_in_usd_ticks": 37756000,
+            }
+        )
+
+        assert usage.cost == 0.0037756
+
+        chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage)
+        assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756)
+
+    def test_streamed_reported_cost_reaches_the_cost_calculator(self):
+        event = XAIResponsesAPIConfig().transform_streaming_response(
+            model="grok-4-latest",
+            parsed_chunk={
+                "type": "response.completed",
+                "sequence_number": 7,
+                "response": self._response_body(
+                    {
+                        "input_tokens": 100,
+                        "output_tokens": 200,
+                        "total_tokens": 300,
+                        "cost_in_usd_ticks": 37756000,
+                    }
+                ),
+            },
+            logging_obj=Mock(),
+        )
+
+        assert isinstance(event, ResponseCompletedEvent)
+        chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage)
+        assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756)
+
     def test_usage_without_a_reported_cost_is_left_alone(self):
         usage = self._transformed_usage({"input_tokens": 100, "output_tokens": 200, "total_tokens": 300})
 
diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py
index 7047408df5b..290cd3dcb3a 100644
--- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py
+++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py
@@ -2,9 +2,12 @@ from unittest.mock import Mock
 
 import httpx
 
+import litellm
 from litellm.llms.xai.chat.transformation import (
+    XAIChatCompletionStreamingHandler,
     XAIChatConfig,
 )
+from litellm.llms.xai.cost_calculator import cost_per_token
 from litellm.types.utils import (
     CompletionTokensDetailsWrapper,
     ModelResponse,
@@ -233,6 +236,19 @@ class TestXAIReportedCost:
         )
         return response.usage
 
+    def test_reported_cost_reaches_the_cost_calculator(self):
+        usage = self._transformed_usage(
+            {
+                "prompt_tokens": 100,
+                "completion_tokens": 200,
+                "total_tokens": 300,
+                "cost_in_usd_ticks": 37756000,
+            }
+        )
+
+        assert usage.cost == 0.0037756
+        assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0037756)
+
     def test_usage_without_a_reported_cost_is_left_alone(self):
         usage = self._transformed_usage({"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300})
 
@@ -250,3 +266,36 @@ class TestXAIReportedCost:
         )
 
         assert getattr(usage, "cost", None) is None
+
+    def test_streamed_reported_cost_survives_chunk_aggregation(self):
+        """Streamed spend only matches if the conversion happens on the chunk.
+
+        Chunk aggregation rebuilds usage from the fields it models plus ``cost``, so a
+        chunk still carrying only ``cost_in_usd_ticks`` loses the reported amount.
+        """
+        handler = XAIChatCompletionStreamingHandler(streaming_response=iter([]), sync_stream=True)
+
+        parsed = handler.chunk_parser(
+            {
+                "id": "chatcmpl-xai",
+                "object": "chat.completion.chunk",
+                "created": 0,
+                "model": "grok-4-latest",
+                "choices": [],
+                "usage": {
+                    "prompt_tokens": 100,
+                    "completion_tokens": 200,
+                    "total_tokens": 300,
+                    "cost_in_usd_ticks": 37756000,
+                },
+            }
+        )
+
+        assert parsed.usage.cost == 0.0037756
+
+        assembled = litellm.stream_chunk_builder(chunks=[parsed])
+        assert assembled.usage.cost == 0.0037756
+        assert cost_per_token(model="grok-4-latest", usage=assembled.usage) == (
+            0.0,
+            0.0037756,
+        )