fix(responses): keep the client's usage shape when the logging copy cannot re-validate the response

litellm/responses/streaming_iterator.py

@@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti
 
 import httpx
 from openai._streaming import SSEDecoder
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 from typing_extensions import TypeIs
 
 import litellm
@@ -439,18 +439,7 @@ class BaseResponsesAPIStreamingIterator:
         if self._persist_completed_response_before_logging:
             self._persist_completed_response_to_cache(is_async=is_async)
 
-        # Create a copy for logging to avoid modifying the response object that will be returned to the user
-        # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
-        # to chat completion format (prompt_tokens/completion_tokens) for internal logging
-        # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
-        # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
-        logging_response = self.completed_response
-        if self.completed_response is not None and hasattr(self.completed_response, "model_dump"):
-            try:
-                logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump())
-            except Exception:
-                # Fallback to original if serialization fails
-                pass
+        logging_response: Final[object] = _logging_copy(self.completed_response)
         self._restore_provider_response_headers(logging_response)
 
         end_time: Final = datetime.now()
@@ -489,10 +478,10 @@ class BaseResponsesAPIStreamingIterator:
     def _restore_provider_response_headers(self, logging_response: object) -> None:
         """Re-apply the provider's response headers to the copy handed to logging callbacks.
 
-        ``model_validate(model_dump())`` above drops pydantic private attributes, so the
+        ``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the
         ``_hidden_params`` the provider transform set on the nested response are lost. Returns early
-        when that copy fell back to the original event, so logging-only state never lands on the
-        object the caller is iterating.
+        when the event was not a pydantic model and logging got the original, so logging-only state
+        never lands on the object the caller is iterating.
         """
         if logging_response is self.completed_response:
             return
@@ -1294,6 +1283,26 @@ def _add_text_like_part_events(
         )
 
 
+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
+    deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow
+    copies of the event and its nested response still keep the caller's ``usage`` attribute separate."""
+    if not isinstance(event, BaseModel):
+        return event
+    try:
+        return type(event).model_validate(event.model_dump())
+    except Exception:
+        return _detached_shallow_copy(event)
+
+
+def _detached_shallow_copy(event: BaseModel) -> BaseModel:
+    nested: Final[object] = getattr(event, "response", None)
+    if isinstance(nested, BaseModel):
+        return event.model_copy(update={"response": nested.model_copy()})
+    return event.model_copy()
+
+
 def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
     if isinstance(usage, ResponseAPIUsage):
         return usage

tests/test_litellm/responses/test_streaming_iterator.py

@@ -18,6 +18,7 @@ from litellm.responses.streaming_iterator import (
     SyncResponsesAPIStreamingIterator,
 )
 from litellm.types.llms.openai import (
+    ResponseAPIUsage,
     ResponseCompletedEvent,
     ResponsesAPIResponse,
     ResponsesAPIStreamEvents,
@@ -329,8 +330,6 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead():
 
 
 def _responses_api_response_with_usage() -> ResponsesAPIResponse:
-    from litellm.types.llms.openai import ResponseAPIUsage
-
     return ResponsesAPIResponse(
         id="resp_lit6427",
         created_at=int(datetime(2025, 1, 1).timestamp()),
@@ -383,8 +382,6 @@ def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse:
 
 def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
     from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
-    from litellm.types.llms.openai import ResponseAPIUsage
-
     response = _unvalidated_response_with_dict_usage(
         {
             "input_tokens": 29,
@@ -406,8 +403,6 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage():
 
 def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost():
     from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
-    from litellm.types.llms.openai import ResponseAPIUsage
-
     response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149})
     logging_obj = Mock(spec=LiteLLMLoggingObj)
     logging_obj._response_cost_calculator.return_value = 0.000704
@@ -586,5 +581,50 @@ async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched():
     with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")):
         iterator._log_completed_response(is_async=True)
 
-    assert logged == [iterator.completed_response]
+    assert len(logged) == 1
+    assert logged[0] is not iterator.completed_response
+    assert logged[0].response is not iterator.completed_response.response
+    assert logged[0].response._hidden_params["headers"]["apim-request-id"] == "azure-correlation-1"
     assert iterator.completed_response.response._hidden_params == {}
+
+
+def _unvalidated_completed_config() -> Mock:
+    """Config whose completed event carries a Perplexity-style response that fails validation
+    (``truncation: ""``) and already holds the stamped ``ResponseAPIUsage``."""
+    mock_config = Mock(spec=BaseResponsesAPIConfig)
+
+    def _transform(model, parsed_chunk, logging_obj):
+        response = _unvalidated_response_with_dict_usage(
+            ResponseAPIUsage(input_tokens=29, output_tokens=373, total_tokens=402, cost={"total_cost": 0.0001})
+        )
+        return ResponseCompletedEvent(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response)
+
+    mock_config.transform_streaming_response.side_effect = _transform
+    return mock_config
+
+
+@pytest.mark.asyncio
+async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_validation():
+    """LIT-7391: the logging copy cannot round-trip a response that fails validation, and logging
+    rewrites the assembled response's usage to chat shape in place, so the event handed to logging
+    must never be the one the caller receives."""
+    logging_obj = _logging_obj_stub()
+    logging_obj.stream = True
+    logged: list[object] = []
+    logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
+    logging_obj._on_deferred_stream_complete = None
+
+    iterator = _make_header_iterator(headers={}, config=_unvalidated_completed_config(), logging_obj=logging_obj)
+    events = [event async for event in iterator]
+
+    assert len(logged) == 1
+    now = datetime.now()
+    LiteLLMLoggingObj._get_assembled_streaming_response(
+        logging_obj, logged[0], start_time=now, end_time=now, is_async=True, streaming_chunks=[]
+    )
+    assert logged[0].response.usage["prompt_tokens"] == 29
+
+    client_usage = events[-1].response.usage
+    assert isinstance(client_usage, ResponseAPIUsage)
+    assert client_usage.input_tokens == 29
+    assert client_usage.cost == pytest.approx(0.0001)