fix(proxy): eliminate race condition in streaming guardrail_information logging

litellm/litellm_core_utils/streaming_handler.py

@@ -2174,12 +2174,16 @@ async def __anext__(self) -> "ModelResponseStream":  # noqa: PLR0915
                     None,
                 )
                 if _deferred_cb is not None:
-                    # Proxy has post-call guardrails — let the closure
-                    # run guardrails on the assembled response, then
-                    # fire logging with guardrail_information populated.
-                    self.logging_obj._on_deferred_stream_complete = None  # type: ignore[attr-defined]
-                    asyncio.create_task(
-                        _deferred_cb(complete_streaming_response, cache_hit)
+                    # Proxy has post-call guardrails. Store the assembled
+                    # response so the outer streaming consumer
+                    # (ProxyLogging.async_post_call_streaming_iterator_hook)
+                    # can fire the deferred callback AFTER all guardrail
+                    # end-of-stream blocks complete.  Scheduling here via
+                    # create_task would race with unified_guardrail's
+                    # end-of-stream block for short-stream providers.
+                    self.logging_obj._deferred_stream_complete_args = (  # type: ignore[attr-defined]
+                        complete_streaming_response,
+                        cache_hit,
                     )
                 else:
                     asyncio.create_task(

litellm/proxy/common_request_processing.py

@@ -1059,11 +1059,13 @@ async def base_process_llm_request(  # noqa: PLR0915
                         "_litellm_client_requested_model"
                     ] = requested_model_from_client
 
-                # Streaming: attach a closure that CSW.__anext__ will call
-                # at stream end instead of firing logging directly.  The
-                # closure runs ONLY guardrail hooks (not all callbacks) on
-                # the assembled response so guardrail_information is
-                # populated, then fires both logging handlers.
+                # Streaming: attach a closure that fires after all guardrail
+                # end-of-stream blocks complete.  CSW.__anext__ stores the
+                # assembled response on logging_obj; the outer consumer
+                # (ProxyLogging._fire_deferred_stream_logging) fires the
+                # closure after the full streaming pipeline finishes.
+                # The closure runs non-apply_guardrail hooks on the
+                # assembled response, then fires both logging handlers.
                 # Only for CustomStreamWrapper — raw async generators from
                 # passthrough routes bypass CSW and would orphan the closure.
                 from litellm.litellm_core_utils.streaming_handler import (
@@ -1383,16 +1385,19 @@ async def _run_deferred_stream_guardrails(
         cache_hit: Any,
     ) -> None:
         """
-        Run only post-call guardrail hooks on an assembled streaming response,
-        then fire both async and sync logging handlers.
+        Run non-streaming post-call guardrail hooks on an assembled streaming
+        response, then fire both async and sync logging handlers.
 
-        Called by CSW.__anext__ at stream end via a closure stored on
-        logging_obj._on_deferred_stream_complete.
+        Called by ProxyLogging._fire_deferred_stream_logging after the full
+        streaming pipeline (including unified_guardrail end-of-stream blocks)
+        has completed.
+
+        Guardrails with apply_guardrail are skipped — they already ran via
+        unified_guardrail's streaming iterator.  Only guardrails that override
+        async_post_call_success_hook directly (without apply_guardrail) run
+        here.
 
         This is audit-only — content has already been delivered to the client.
-        Blocking guardrails that raise HTTPException cannot prevent content
-        delivery for streaming.  Per-chunk filtering should use
-        async_post_call_streaming_hook instead.
 
         Extracted as a static method so tests can call the production
         implementation directly rather than reimplementing the closure.
@@ -1405,7 +1410,6 @@ async def _run_deferred_stream_guardrails(
             from litellm.proxy.utils import (
                 _check_and_merge_model_level_guardrails,
             )
-            from litellm.proxy.utils import unified_guardrail as _unified_guardrail
 
             guardrail_data = _check_and_merge_model_level_guardrails(
                 data=captured_data, llm_router=_global_llm_router
@@ -1421,14 +1425,12 @@ async def _run_deferred_stream_guardrails(
                 try:
                     guardrail_result = None
                     if "apply_guardrail" in type(cb).__dict__:
-                        guardrail_data["guardrail_to_apply"] = cb
-                        guardrail_result = (
-                            await _unified_guardrail.async_post_call_success_hook(
-                                user_api_key_dict=captured_user_api_key_dict,
-                                data=guardrail_data,
-                                response=_response,
-                            )
-                        )
+                        # Skip — apply_guardrail guardrails already ran via
+                        # unified_guardrail's end-of-stream block in the
+                        # streaming iterator pipeline.  Running them again
+                        # here would duplicate the guardrail API call
+                        # (e.g. double OpenAI Moderation charges).
+                        continue
                     else:
                         guardrail_result = await cb.async_post_call_success_hook(
                             user_api_key_dict=captured_user_api_key_dict,

litellm/proxy/utils.py

@@ -2235,6 +2235,32 @@ async def async_post_call_streaming_iterator_hook(
         async for chunk in current_response:
             yield chunk
 
+        # Fire deferred logging AFTER all guardrail end-of-stream blocks
+        # completed.  unified_guardrail writes guardrail_information during
+        # its end-of-stream block (inside current_response), so by the time
+        # we reach this point the metadata is fully populated.
+        ProxyLogging._fire_deferred_stream_logging(request_data)
+
+    @staticmethod
+    def _fire_deferred_stream_logging(request_data: dict) -> None:
+        """
+        Fire the deferred streaming logging callback after the full streaming
+        pipeline (including guardrail end-of-stream blocks) has completed.
+
+        CSW.__anext__ stores the callback and args on logging_obj instead of
+        scheduling via create_task (which would race with unified_guardrail's
+        end-of-stream block).  This method retrieves and fires them.
+        """
+        logging_obj = request_data.get("litellm_logging_obj")
+        if logging_obj is None:
+            return
+        _deferred_cb = getattr(logging_obj, "_on_deferred_stream_complete", None)
+        _args = getattr(logging_obj, "_deferred_stream_complete_args", None)
+        if _deferred_cb is not None and _args is not None:
+            logging_obj._on_deferred_stream_complete = None
+            logging_obj._deferred_stream_complete_args = None
+            asyncio.create_task(_deferred_cb(*_args))
+
     def _init_response_taking_too_long_task(self, data: Optional[dict] = None):
         """
         Initialize the response taking too long task if user is using slack alerting