Merge pull request #27012 from BerriAI/litellm_fix-post-call-guardrail

fix(guardrails): post-call guardrail must only fire once

litellm/proxy/common_request_processing.py

@@ -1604,6 +1604,12 @@ class ProxyBaseLLMRequestProcessing:
                         # here would duplicate the guardrail API call
                         # (e.g. double OpenAI Moderation charges).
                         continue
+                    if "async_post_call_streaming_iterator_hook" in type(cb).__dict__:
+                        # Skip — the guardrail already scanned the assembled
+                        # response via its own streaming iterator hook in the
+                        # streaming pipeline. re running this function async_post_call_success_hook
+                        # here would duplicate the scan and can spuriously block the guardrail that already passed / failed.
+                        continue
                     else:
                         guardrail_result = await cb.async_post_call_success_hook(
                             user_api_key_dict=captured_user_api_key_dict,

tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py

@@ -758,6 +758,60 @@ class TestDeferredStreamingClosure:
             apply_guardrail_called is False
         ), "apply_guardrail guardrails must be SKIPPED in deferred path"
 
+    @pytest.mark.asyncio
+    async def test_streaming_iterator_hook_skipped_in_deferred_path(self):
+        """regression test: guardrails that define async_post_call_streaming_iterator_hook must be SKIPPED in _run_deferred_stream_guardrails.
+        The iterator hook already scanned the assembled response in the streaming
+        pipeline"""
+        success_hook_called = False
+
+        class IteratorHookGuardrail(CustomGuardrail):
+            def __init__(self):
+                super().__init__(
+                    guardrail_name="iterator-hook",
+                    default_on=True,
+                    event_hook=GuardrailEventHooks.post_call,
+                )
+
+            async def async_post_call_streaming_iterator_hook(
+                self, user_api_key_dict, response, request_data
+            ):
+                async for chunk in response:
+                    yield chunk
+
+            async def async_post_call_success_hook(
+                self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any
+            ) -> Any:
+                nonlocal success_hook_called
+                success_hook_called = True
+                return response
+
+        mock_logging_obj = MagicMock()
+        mock_logging_obj.model_call_details = {"metadata": {}}
+
+        async def track_async_success(*args, **kwargs):
+            pass
+
+        mock_logging_obj.async_success_handler = track_async_success
+
+        guardrail = IteratorHookGuardrail()
+
+        with patch("litellm.callbacks", [guardrail]):
+            await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
+                captured_data={"model": "gpt-4", "metadata": {}},
+                captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"),
+                captured_logging_obj=mock_logging_obj,
+                assembled_response=MagicMock(),
+                cache_hit=False,
+            )
+
+        await asyncio.sleep(0)
+
+        assert success_hook_called is False, (
+            "Guardrails that implement async_post_call_streaming_iterator_hook "
+            "must be SKIPPED in deferred path — the iterator hook already ran"
+        )
+
     @pytest.mark.asyncio
     async def test_hooks_receive_merged_guardrail_data(self):
         """Hooks must receive guardrail_data (the merged dict from