Merge pull request #40191 from BerriAI/litellm_fix_streaming_guardrail_block_logging

fix(proxy): log blocked streaming guardrail responses as failures, not success

litellm/litellm_core_utils/litellm_logging.py

@@ -1991,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,

litellm/proxy/utils.py

@@ -85,7 +85,11 @@ from litellm._logging import _redact_string, verbose_proxy_logger
 from litellm._service_logger import ServiceLogging, ServiceTypes
 from litellm.caching.caching import DualCache, RedisCache
 from litellm.caching.dual_cache import LimitedSizeOrderedDict
-from litellm.exceptions import RejectedRequestError, SensitiveDataRouteException
+from litellm.exceptions import (
+    GuardrailRaisedException,
+    RejectedRequestError,
+    SensitiveDataRouteException,
+)
 from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
     ModifyResponseException,
@@ -901,6 +905,9 @@ def _call_type_for_route(route: str | None) -> str | None:
     return call_types[0].value if len(operations) == 1 else None
 
 
+_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException)
+
+
 def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]:
     """Failure-path callbacks run after ``litellm_logging_obj`` is popped from
     request_data (it is not serialisable), so the caller merges these fields
@@ -2991,6 +2998,7 @@ class ProxyLogging:
             - Authentication Errors from user_api_key_auth
             - HTTP HTTPException (rate limit errors)
             - ProxyException (guardrail blocks, budget / rate-limit errors)
+            - GuardrailRaisedException (guardrail blocks / guardrail failures)
         """
 
         #########################################################
@@ -3005,9 +3013,7 @@ class ProxyLogging:
         if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)):
             return False
 
-        return isinstance(original_exception, (HTTPException, ProxyException)) or (
-            error_type == ProxyErrorTypes.auth_error
-        )
+        return isinstance(original_exception, _PROXY_ONLY_LLM_API_ERRORS) or (error_type == ProxyErrorTypes.auth_error)
 
     async def _handle_logging_proxy_only_error(
         self,
@@ -3563,8 +3569,9 @@ class ProxyLogging:
                     yield chunk
             except (GeneratorExit, asyncio.CancelledError):
                 raise
-            except Exception:
-                ProxyLogging._fire_deferred_stream_logging(request_data)
+            except Exception as e:
+                if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e):
+                    ProxyLogging._fire_deferred_stream_logging(request_data)
                 raise
             ProxyLogging._fire_deferred_stream_logging(request_data)
             return
@@ -3638,8 +3645,9 @@ class ProxyLogging:
                 yield chunk
         except (GeneratorExit, asyncio.CancelledError):
             raise
-        except Exception:
-            ProxyLogging._fire_deferred_stream_logging(request_data)
+        except Exception as e:
+            if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e):
+                ProxyLogging._fire_deferred_stream_logging(request_data)
             raise
 
         # Fire deferred logging AFTER all guardrail end-of-stream blocks
@@ -3735,6 +3743,23 @@ class ProxyLogging:
             logging_obj._deferred_stream_complete_args = None
             asyncio.create_task(_deferred_cb(*_args))
 
+    @staticmethod
+    def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool:
+        """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
+        _args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None)
+        assembled: Final = _args[0] if _args else None
+        if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse):
+            return False
+        logging_obj._on_deferred_stream_complete = None
+        logging_obj._deferred_stream_complete_args = None
+        logging_obj.record_assembled_response_for_failure(assembled)
+        return True
+
     async def _arelease_max_parallel_requests_on_disconnect(
         self,
         user_api_key_dict: UserAPIKeyAuth,

tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py

@@ -4,13 +4,14 @@ and ``_handle_logging_proxy_only_error``."""
 from __future__ import annotations
 
 import asyncio
-from typing import Any
+from datetime import datetime
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 from fastapi import HTTPException
 
 import litellm
+from litellm.exceptions import GuardrailRaisedException
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.proxy._types import AlertType, ProxyErrorTypes
 from litellm.proxy.utils import ProxyLogging
@@ -47,12 +48,17 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging):
             error_type=ProxyErrorTypes.auth_error,
             route="/chat/completions",
         ),
+        "guardrail_raised_on_llm_route": proxy_logging._is_proxy_only_llm_api_error(
+            original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"),
+            route="/chat/completions",
+        ),
     }
     assert snapshot == {
         "no_route": False,
         "non_llm_route": False,
         "http_on_llm_route": True,
         "auth_short_circuit": True,
+        "guardrail_raised_on_llm_route": True,
     }
 
 
@@ -318,3 +324,50 @@ async def test_post_call_failure_hook_keeps_the_route_for_multi_operation_routes
         route=route,
     )
     assert request_data["call_type"] == route
+
+
+@pytest.mark.asyncio
+async def test_post_call_failure_hook_guardrail_block_fires_failure_callback(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    """A ``GuardrailRaisedException`` on an LLM route must reach the logging
+    object's ``async_failure_handler`` so custom loggers see a ``failure``
+    status - without this, guardrail blocks produce only
+    ``post_call_failure_hook`` and no failure logging event."""
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+
+    recorded: list[object] = []
+
+    class _StatusRecorder(CustomLogger):
+        async def async_log_failure_event(
+            self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
+        ) -> None:
+            standard_logging_object = kwargs.get("standard_logging_object")
+            recorded.append(standard_logging_object.get("status") if isinstance(standard_logging_object, dict) else None)
+
+    monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()])
+    logging_obj = LiteLLMLoggingObj(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=False,
+        call_type="acompletion",
+        start_time=datetime.now(),
+        litellm_call_id="test_guardrail_block_failure_cb",
+        function_id="test_guardrail_block_failure_cb",
+    )
+    request_data = {
+        "litellm_logging_obj": logging_obj,
+        "litellm_call_id": "test_guardrail_block_failure_cb",
+        "model": "gpt-4o-mini",
+        "messages": [{"role": "user", "content": "hi"}],
+        "metadata": {},
+    }
+    proxy_logging.alert_types = []
+    await proxy_logging.post_call_failure_hook(
+        request_data=request_data,
+        original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"),
+        user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"),
+    )
+    await asyncio.sleep(0)
+    await asyncio.sleep(0)
+    assert recorded == ["failure"]

tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py

@@ -10,6 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``,
 from __future__ import annotations
 
 import asyncio
+from collections.abc import AsyncGenerator, AsyncIterator
 from datetime import datetime
 from typing import Any, Dict, List
 from unittest.mock import AsyncMock, MagicMock
@@ -18,12 +19,15 @@ import pytest
 from fastapi import HTTPException
 
 import litellm
+from litellm.exceptions import GuardrailRaisedException
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
     BaseAnthropicMessagesStreamingIterator,
 )
+from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.utils import ProxyLogging
+from litellm.types.utils import Usage
 
 
 @pytest.fixture(autouse=True)
@@ -479,6 +483,135 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_
     assert logging_obj._deferred_stream_complete_args is None
 
 
+def _armed_chat_stream(
+    test_name: str, request_data: dict[str, object], events: list[str]
+) -> tuple[LiteLLMLoggingObj, AsyncIterator[dict[str, object]]]:
+    """A /chat/completions stream whose CSW shape parks ``(assembled ModelResponse, cache_hit)``
+    at upstream exhaustion, with the deferred dispatch recording into ``events``."""
+    logging_obj = LiteLLMLoggingObj(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=True,
+        call_type="acompletion",
+        start_time=datetime.now(),
+        litellm_call_id=test_name,
+        function_id=test_name,
+    )
+    logging_obj.optional_params = {}
+    logging_obj.litellm_params = {}
+    logging_obj.standard_built_in_tools_params = None
+
+    async def _dispatch_deferred_logging(*args: object) -> None:
+        events.append("success_dispatched")
+
+    logging_obj._on_deferred_stream_complete = _dispatch_deferred_logging
+    request_data["litellm_logging_obj"] = logging_obj
+
+    assembled = litellm.ModelResponse(
+        model="gpt-4o-mini",
+        choices=[{"index": 0, "message": {"role": "assistant", "content": "BANANA"}}],
+        usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8),
+    )
+
+    async def _upstream() -> AsyncIterator[dict[str, object]]:
+        yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "BAN"}}]}
+        yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "ANA"}}]}
+        logging_obj._deferred_stream_complete_args = (assembled, False)
+
+    return logging_obj, _upstream()
+
+
+def _raising_at_end_of_stream(error: Exception) -> CustomLogger:
+    class _EndOfStreamRaiser(CustomLogger):
+        async def async_post_call_streaming_iterator_hook(
+            self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object]
+        ) -> AsyncGenerator[object, None]:
+            async for chunk in response:
+                yield chunk
+            raise error
+
+    return _EndOfStreamRaiser()
+
+
+@pytest.mark.asyncio
+async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    """
+    A guardrail that raises ``GuardrailRaisedException`` at end of a
+    /chat/completions stream must NOT dispatch the parked success logging:
+    the request is logged via the failure path instead, with the consumed
+    usage carried over so the failure row bills correctly.
+    """
+    events: list[str] = []
+    request_data: dict[str, object] = {"metadata": {}}
+    logging_obj, upstream = _armed_chat_stream("test_chat_stream_guardrail_block", request_data, events)
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [_raising_at_end_of_stream(GuardrailRaisedException(guardrail_name="g", message="blocked"))],
+    )
+
+    with pytest.raises(GuardrailRaisedException):
+        async for _ in proxy_logging.async_post_call_streaming_iterator_hook(
+            response=upstream,
+            user_api_key_dict=make_user_api_key_auth(),
+            request_data=request_data,
+        ):
+            pass
+    await asyncio.sleep(0)
+    await asyncio.sleep(0)
+
+    snapshot = {
+        "events": events,
+        "callback_cleared": logging_obj._on_deferred_stream_complete is None,
+        "args_cleared": logging_obj._deferred_stream_complete_args is None,
+        "combined_usage_total_tokens": logging_obj.model_call_details["combined_usage_object"].total_tokens,
+        "response_cost_positive": logging_obj.model_call_details["response_cost"] > 0,
+    }
+    assert snapshot == {
+        "events": [],
+        "callback_cleared": True,
+        "args_cleared": True,
+        "combined_usage_total_tokens": 8,
+        "response_cost_positive": True,
+    }
+
+
+@pytest.mark.asyncio
+async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes_success_logging(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    """
+    ``post_call_failure_hook`` only routes proxy-level errors (HTTPException,
+    ProxyException, GuardrailRaisedException) through failure logging. A
+    callback that dies with any other exception after the stream completed
+    must keep flushing the parked success dispatch, or the request ends with
+    no terminal log at all.
+    """
+    events: list[str] = []
+    request_data: dict[str, object] = {"metadata": {}}
+    logging_obj, upstream = _armed_chat_stream("test_chat_stream_generic_callback_error", request_data, events)
+    monkeypatch.setattr(litellm, "callbacks", [_raising_at_end_of_stream(RuntimeError("callback crashed"))])
+
+    with pytest.raises(RuntimeError):
+        async for _ in proxy_logging.async_post_call_streaming_iterator_hook(
+            response=upstream,
+            user_api_key_dict=make_user_api_key_auth(),
+            request_data=request_data,
+        ):
+            pass
+    await asyncio.sleep(0)
+    await asyncio.sleep(0)
+
+    snapshot = {
+        "events": events,
+        "args_cleared": logging_obj._deferred_stream_complete_args is None,
+        "failure_usage_recorded": "combined_usage_object" in logging_obj.model_call_details,
+    }
+    assert snapshot == {"events": ["success_dispatched"], "args_cleared": True, "failure_usage_recorded": False}
+
+
 # ---------------------------------------------------------------------------
 # _fire_deferred_stream_logging
 # ---------------------------------------------------------------------------