fix: report real usage on streaming blocks, disable buffered mode for content-rewriting guardrails

- _standalone_block_chunks and _block_continuation_chunks now read real
  token usage from ModifyResponseException.original_response instead of
  hardcoding zero, matching the non-streaming _blocked_response_usage path.
  Shared helper moved to guardrail_translation/utils.py.
- streaming_buffer_until_moderated is now forced off when the guardrail has
  mask_response_content=True, since buffered replay releases the withheld
  original chunks verbatim -- unsafe for a guardrail that rewrites content
  (e.g. PII masking).
- Fix inverted streaming-flag precedence comment.

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -104,6 +104,9 @@ class AnthropicMessagesHandler(BaseTranslation):
         from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
             FakeAnthropicMessagesStreamIterator,
         )
+        from litellm.llms.base_llm.guardrail_translation.utils import (
+            blocked_response_usage,
+        )
         from litellm.types.utils import AnthropicMessagesResponse
 
         block_response = AnthropicMessagesResponse(
@@ -113,7 +116,7 @@ class AnthropicMessagesHandler(BaseTranslation):
             content=[{"type": "text", "text": exc.message}],
             model=exc.model,
             stop_reason="end_turn",
-            usage={"input_tokens": 0, "output_tokens": 0},
+            usage=blocked_response_usage(getattr(exc, "original_response", None)),
         )
         return list(FakeAnthropicMessagesStreamIterator(response=block_response))
 
@@ -122,9 +125,16 @@ class AnthropicMessagesHandler(BaseTranslation):
         append the block message as a new text block, then end the message --
         without a second message_start."""
 
+        from litellm.llms.base_llm.guardrail_translation.utils import (
+            blocked_response_usage,
+        )
+
         def _sse(event_type: str, payload: dict) -> bytes:
             return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
 
+        output_tokens = blocked_response_usage(getattr(exc, "original_response", None))[
+            "output_tokens"
+        ]
         open_index, max_index = self._content_block_state(responses_so_far)
         new_index = (max_index + 1) if max_index is not None else 0
         chunks: list[bytes] = []
@@ -153,7 +163,7 @@ class AnthropicMessagesHandler(BaseTranslation):
                 {
                     "type": "message_delta",
                     "delta": {"stop_reason": "end_turn", "stop_sequence": None},
-                    "usage": {"output_tokens": 0},
+                    "usage": {"output_tokens": output_tokens},
                 },
             ),
             _sse("message_stop", {"type": "message_stop"}),

litellm/llms/base_llm/guardrail_translation/utils.py

@@ -1,10 +1,38 @@
 from __future__ import annotations
 
-from typing import Any, List
+from typing import Any, List, Optional
 
+from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
 from litellm.types.llms.openai import AllMessageValues
 
 
+def blocked_response_usage(original_response: Optional[Any]) -> AnthropicUsage:
+    """
+    Token usage for a synthetic guardrail-blocked response.
+
+    A post-call block replaces the LLM's response with the violation message,
+    but the upstream call already consumed tokens -- report that real usage
+    (carried on ``ModifyResponseException.original_response``) rather than
+    discarding it. Pre-call blocks never invoked the LLM (no original_response),
+    so usage is zero.
+    """
+    usage_obj: Any = None
+    if isinstance(original_response, dict):
+        usage_obj = original_response.get("usage")
+    elif original_response is not None:
+        usage_obj = getattr(original_response, "usage", None)
+
+    def _tokens(key: str) -> int:
+        if isinstance(usage_obj, dict):
+            return int(usage_obj.get(key, 0) or 0)
+        return int(getattr(usage_obj, key, 0) or 0)
+
+    return AnthropicUsage(
+        input_tokens=_tokens("input_tokens"),
+        output_tokens=_tokens("output_tokens"),
+    )
+
+
 def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool:
     per = getattr(guardrail_to_apply, "skip_system_message_in_guardrail", None)
     if per is not None:

litellm/proxy/anthropic_endpoints/endpoints.py

@@ -12,6 +12,9 @@ from litellm.integrations.custom_guardrail import ModifyResponseException
 from litellm.llms.anthropic.experimental_pass_through.context_management import (
     AnthropicContextManagementError,
 )
+from litellm.llms.base_llm.guardrail_translation.utils import (
+    blocked_response_usage as _blocked_response_usage,
+)
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_request_processing import (
@@ -19,7 +22,6 @@ from litellm.proxy.common_request_processing import (
     create_response,
 )
 from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
-from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
 from litellm.types.utils import TokenCountResponse
 
 router = APIRouter()
@@ -59,33 +61,6 @@ def _strip_total_tokens_from_anthropic_response(response: Any) -> None:
         usage.pop("total_tokens", None)
 
 
-def _blocked_response_usage(original_response: Optional[Any]) -> AnthropicUsage:
-    """
-    Token usage for a synthetic guardrail-blocked response.
-
-    A post-call block replaces the LLM's response with the violation message,
-    but the upstream call already consumed tokens -- report that real usage
-    (carried on ``ModifyResponseException.original_response``) rather than
-    discarding it. Pre-call blocks never invoked the LLM (no original_response),
-    so usage is zero.
-    """
-    usage_obj: Any = None
-    if isinstance(original_response, dict):
-        usage_obj = original_response.get("usage")
-    elif original_response is not None:
-        usage_obj = getattr(original_response, "usage", None)
-
-    def _tokens(key: str) -> int:
-        if isinstance(usage_obj, dict):
-            return int(usage_obj.get(key, 0) or 0)
-        return int(getattr(usage_obj, key, 0) or 0)
-
-    return AnthropicUsage(
-        input_tokens=_tokens("input_tokens"),
-        output_tokens=_tokens("output_tokens"),
-    )
-
-
 @router.post(
     "/v1/messages",
     tags=["[beta] Anthropic `/v1/messages`"],

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -319,8 +319,9 @@ class UnifiedLLMGuardrails(CustomLogger):
 
         guardrail_to_apply: CustomGuardrail = request_data.pop("guardrail_to_apply", None)
 
-        # Get streaming configuration, with precedence: guardrail attribute ->
-        # guardrail_config dict -> this callback's optional_params.
+        # Get streaming configuration. Resolution order (later wins): default
+        # < guardrail attribute < guardrail_config dict < this callback's
+        # optional_params.
         def _streaming_flag(name: str, default: Any) -> Any:
             value = default
             if guardrail_to_apply is not None:
@@ -336,11 +337,25 @@ class UnifiedLLMGuardrails(CustomLogger):
         # Withhold every chunk until end-of-stream moderation passes, then
         # release the original chunks (clean) or only the block message
         # (blocked) -- moderating the whole response *before* any content
-        # reaches the client. Intended for allow/block guardrails: on release
-        # the original chunks are replayed as-is, so content-rewriting
-        # guardrails (e.g. PII masking) are not applied to a buffered stream.
+        # reaches the client. Only safe for allow/block guardrails: on
+        # release the original chunks are replayed as-is, so a
+        # content-rewriting guardrail (e.g. PII masking) would leak
+        # unredacted content. Guarded below via mask_response_content.
         buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", False)
 
+        if (
+            buffer_until_moderated
+            and guardrail_to_apply is not None
+            and getattr(guardrail_to_apply, "mask_response_content", False)
+        ):
+            verbose_proxy_logger.warning(
+                "UnifiedLLMGuardrails: streaming_buffer_until_moderated is disabled for %s "
+                "because mask_response_content=True -- buffered replay would release "
+                "unredacted original chunks instead of the moderated output.",
+                guardrail_to_apply.guardrail_name,
+            )
+            buffer_until_moderated = False
+
         # Buffering can only moderate the assembled response, so it always
         # defers to end-of-stream.
         if buffer_until_moderated:

tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py

@@ -157,3 +157,19 @@ async def test_buffered_clean_releases_all_content():
         raw.rstrip().endswith('event: message_stop\ndata: {"type": "message_stop"}'.rstrip()) or "message_stop" in raw
     )
     assert BLOCK_MESSAGE not in raw
+
+
+@pytest.mark.asyncio
+async def test_buffered_mode_disabled_for_content_rewriting_guardrail():
+    """Buffered replay yields the withheld *original* chunks verbatim, which
+    is unsafe for a guardrail that rewrites response text (e.g. PII masking):
+    the client would get the unredacted original instead of the moderated
+    output. mask_response_content=True must force buffering off so the
+    request falls back to the (correctly moderated) non-buffered path."""
+    guardrail = _PassingGuardrail(
+        guardrail_name="masker", event_hook="post_call", mask_response_content=True
+    )
+    raw = await _run(guardrail)
+    assert guardrail.streaming_buffer_until_moderated is True  # request asked for buffering
+    assert ORIGINAL_MARKER in raw
+    assert BLOCK_MESSAGE not in raw