From 46bd3d40d7abb7f44db19fb8d81b0d9871a314f1 Mon Sep 17 00:00:00 2001
From: yassin <yassin@berri.ai>
Date: Tue, 15 Sep 2026 20:23:50 +0000
Subject: [PATCH] refactor(logging): bill an assembled stream on the failure
 log via a public Logging method

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/litellm_core_utils/litellm_logging.py    |  8 ++++++--
 litellm/proxy/utils.py                           | 16 ++++------------
 .../proxy_logging/test_post_call_failure_hook.py |  1 -
 3 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index a7ad774b02d..abac624d5ec 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -639,8 +639,6 @@ class Logging(LiteLLMLoggingBaseClass):
         self._defer_async_logging: bool = False
         self._enqueue_deferred_logging: Callable[[], None] | None = None
         self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
-        self._on_deferred_stream_complete: Callable[..., Awaitable[None]] | None = None
-        self._deferred_stream_complete_args: tuple[object, ...] | None = None
 
     def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None:
         """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""
@@ -1993,6 +1991,12 @@ class Logging(LiteLLMLoggingBaseClass):
         self.model_call_details["combined_usage_object"] = usage
         self.model_call_details["response_cost"] = response_cost
 
+    def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None:
+        """Bill a fully streamed response on the failure log when a post-call hook rejects it."""
+        usage: Final = getattr(assembled, "usage", None)
+        if isinstance(usage, Usage):
+            self.record_partial_usage_for_failure(usage, self._response_cost_calculator(result=assembled) or 0.0)
+
     async def dispatch_failure_handlers(
         self,
         exception: Exception,
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index a12bb56f8f8..80cf6ba3c8a 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -3745,13 +3745,9 @@ class ProxyLogging:
 
     @staticmethod
     def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool:
-        """Drop the parked success dispatch when the stream ends in an error the proxy logs
-        as a failure (``_PROXY_ONLY_LLM_API_ERRORS``, e.g. a post_call guardrail block) and
-        the CSW parked an assembled ``ModelResponse``, carrying its usage onto the logging
-        object so the failure row bills what the stream consumed. Returns False, leaving the
-        parked dispatch for the caller to flush, for any other error and for the native
-        /v1/messages and responses shapes that park a logging coroutine with no usage.
-        """
+        """Drop the parked success dispatch for an assembled chat stream that ends in an error
+        ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead.
+        Returns False when the parked dispatch should still be flushed by the caller."""
         logging_obj: Final = request_data.get("litellm_logging_obj")
         if not isinstance(logging_obj, Logging):
             return False
@@ -3761,11 +3757,7 @@ class ProxyLogging:
             return False
         logging_obj._on_deferred_stream_complete = None
         logging_obj._deferred_stream_complete_args = None
-        usage: Final[Usage | None] = getattr(assembled, "usage", None)
-        if isinstance(usage, Usage):
-            logging_obj.record_partial_usage_for_failure(
-                usage, logging_obj._response_cost_calculator(result=assembled) or 0.0
-            )
+        logging_obj.record_assembled_response_for_failure(assembled)
         return True
 
     async def _arelease_max_parallel_requests_on_disconnect(
diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py
index 13fcccbad97..d7a6124dd97 100644
--- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py
+++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py
@@ -5,7 +5,6 @@ from __future__ import annotations
 
 import asyncio
 from datetime import datetime
-from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest