From ef26590f72877a65947d2cb16df0bbbe3cd61892 Mon Sep 17 00:00:00 2001
From: milan <milan@berri.ai>
Date: Tue, 28 Jul 2026 23:12:06 +0000
Subject: [PATCH] test(streaming): assert logged cost from success callback for
 usage-only chunk

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../test_streaming_handler.py                 | 52 +++++++++++++------
 1 file changed, 35 insertions(+), 17 deletions(-)

diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
index 22daaf64dfd..f96a19a26c0 100644
--- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
@@ -1508,15 +1508,15 @@ async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Loggin
 
 
 @pytest.mark.asyncio
-async def test_openrouter_streaming_usage_only_chunk_without_stream_options(
-    logging_obj: Logging,
-):
+async def test_openrouter_streaming_usage_only_chunk_without_stream_options():
     """
     Regression: OpenRouter's post-finish chunk has `choices: []`. When the caller did not
     pass stream_options.include_usage it was dropped before cost tracking, so the
     provider-reported cost never reached the assembled response.
     """
-    from litellm.cost_calculator import get_response_cost_from_hidden_params
+    import time
+
+    from litellm.integrations.custom_logger import CustomLogger
     from litellm.utils import ModelResponseListIterator
 
     chunk1 = ModelResponseStream(
@@ -1549,30 +1549,48 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options(
         ),
     )
 
+    class MockCallback(CustomLogger):
+        pass
+
+    mock_callback = MockCallback()
+    litellm.success_callback = [mock_callback]
+    litellm._async_success_callback = [mock_callback]
+
+    stream_logging_obj = Logging(
+        model="openrouter/claude",
+        messages=[{"role": "user", "content": "Hey"}],
+        stream=True,
+        call_type="completion",
+        start_time=time.time(),
+        litellm_call_id="12345",
+        function_id="1245",
+    )
+    stream_logging_obj.update_environment_variables(
+        model="openrouter/claude",
+        optional_params={},
+        litellm_params={},
+        custom_llm_provider="openrouter",
+    )
+
     response = CustomStreamWrapper(
         completion_stream=ModelResponseListIterator(
             model_responses=[chunk1, chunk2, usage_only_chunk]
         ),
         model="openrouter/claude",
         custom_llm_provider="openrouter",
-        logging_obj=logging_obj,
+        logging_obj=stream_logging_obj,
     )
 
-    collected_chunks = [chunk async for chunk in response]
+    with patch.object(mock_callback, "async_log_success_event") as mock_success_event:
+        collected_chunks = [chunk async for chunk in response]
+        await asyncio.sleep(1)
 
     assert all(getattr(chunk, "usage", None) is None for chunk in collected_chunks)
 
-    complete_response = litellm.stream_chunk_builder(
-        chunks=response.chunks,
-        messages=[{"role": "user", "content": "Hey"}],
-    )
-    assert complete_response is not None
-    assert complete_response.usage.cost == 0.00025
-
-    CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response)
-    assert (
-        get_response_cost_from_hidden_params(complete_response._hidden_params) == 0.00025
-    )
+    mock_success_event.assert_called_once()
+    logged_kwargs = mock_success_event.call_args.kwargs["kwargs"]
+    assert logged_kwargs["response_cost"] == 0.00025
+    assert logged_kwargs["standard_logging_object"]["response_cost"] == 0.00025
 
 
 def test_openrouter_streaming_cost_propagates_to_hidden_params():