diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0222e756ba1..b61cad35041 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2,6 +2,7 @@ import contextlib import os import sys import asyncio +from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -3158,6 +3159,47 @@ async def test_async_success_handler_preserves_response_cost_for_pass_through_en assert slo["response_cost"] > 0 +@pytest.mark.asyncio +async def test_sync_success_handler_preserves_streaming_cost_computed_by_async_handler(): + """Regression for #34875: the sync success_handler used to reset response_cost to None + for streams, so custom callbacks reading kwargs["response_cost"] saw None whenever it ran + after the async handler had already computed the cost on the shared model_call_details.""" + import litellm + + litellm.callbacks = [] + litellm.success_callback = [] + + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + mock_response="hello there", + stream_options={"include_usage": True}, + ) + logging_obj = response.logging_obj + async for _chunk in response: + pass + + async def wait_for_async_handler_cost() -> Optional[float]: + for _ in range(200): + cost = logging_obj.model_call_details.get("response_cost") + if cost: + return cost + await asyncio.sleep(0.05) + return None + + cost_from_async_handler = await wait_for_async_handler_cost() + assert cost_from_async_handler is not None and cost_from_async_handler > 0 + + assembled = logging_obj.model_call_details.get( + "async_complete_streaming_response" + ) or logging_obj.model_call_details.get("complete_streaming_response") + + for _ in range(2): + logging_obj.success_handler(assembled, None, None, False) + assert logging_obj.model_call_details["response_cost"] == cost_from_async_handler + + def test_process_hidden_params_recalculates_cost_after_failure_handler_zero(): """ Regression: PR #21844 preserved response_cost=0 set by failure_handler on failed