feat(guardrails): buffer + cleanly terminate streamed responses on block (#31389)

Streaming moderation improvements for the unified guardrail post-call
streaming iterator hook:

- streaming_buffer_until_moderated: withhold all chunks until end-of-stream
  moderation passes, then release the original response (clean) or only the
  block message (blocked) -- the original content is never delivered on a
  block. Snapshot chunks with a shallow list() copy (end-of-stream builds a
  separate assembled response; chunks aren't mutated in place).
- Clean Anthropic SSE on block: synthesize a well-formed termination sequence
  instead of a bare data: {"error": ...} blob that truncates the stream.
  Provider-specific synthesis lives in AnthropicMessagesHandler via
  build_block_sse_chunks (format-agnostic routing stays in the hook).
- Mid-stream blocks continue the in-progress message (close open content
  block, append block message, terminate) rather than emitting a second
  message_start, which clients reject. Standalone envelope only when no chunks
  were sent (buffered path).
- ModifyResponseException imported under TYPE_CHECKING + locally at runtime to
  avoid a module-level cyclic import.

Adds regression tests for buffering (content withheld on block) and mid-stream
continuation (single message_start).

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

@@ -48,7 +48,10 @@ from litellm.types.utils import (
 )
 
 if TYPE_CHECKING:
-    from litellm.integrations.custom_guardrail import CustomGuardrail
+    from litellm.integrations.custom_guardrail import (
+        CustomGuardrail,
+        ModifyResponseException,
+    )
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
     from litellm.types.llms.anthropic_messages.anthropic_response import (
         AnthropicMessagesResponse,
@@ -70,6 +73,146 @@ class AnthropicMessagesHandler(BaseTranslation):
         super().__init__()
         self.adapter = LiteLLMAnthropicMessagesAdapter()
 
+    def build_block_sse_chunks(
+        self,
+        exc: "ModifyResponseException",
+        stream_started: bool = False,
+        responses_so_far: Optional[list[Any]] = None,
+    ) -> list[bytes]:
+        """
+        Build an Anthropic SSE sequence delivering the guardrail block message
+        and terminating the stream cleanly.
+
+        - ``stream_started`` False (buffered / pre-stream): nothing has been
+          sent, so emit a complete standalone message (message_start ->
+          content_block_* -> message_delta -> message_stop) via
+          FakeAnthropicMessagesStreamIterator, the same converter the
+          /v1/messages pre-stream block handler uses.
+        - ``stream_started`` True (sampling / detect-only end-of-stream): real
+          chunks were already sent, so *continue* the in-progress message --
+          close the open content block, append the block message as a new text
+          block, then end the message. Emitting a second ``message_start`` here
+          would make Anthropic clients reject the stream.
+        """
+        if stream_started:
+            return self._block_continuation_chunks(exc, responses_so_far or [])
+        return self._standalone_block_chunks(exc)
+
+    def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]:
+        import uuid
+
+        from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
+            FakeAnthropicMessagesStreamIterator,
+        )
+        from litellm.types.utils import AnthropicMessagesResponse
+
+        block_response = AnthropicMessagesResponse(
+            id=f"msg_{uuid.uuid4()}",
+            type="message",
+            role="assistant",
+            content=[{"type": "text", "text": exc.message}],
+            model=exc.model,
+            stop_reason="end_turn",
+            usage={"input_tokens": 0, "output_tokens": 0},
+        )
+        return list(FakeAnthropicMessagesStreamIterator(response=block_response))
+
+    def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[Any]) -> list[bytes]:
+        """Continue an already-started message: close the open content block,
+        append the block message as a new text block, then end the message --
+        without a second message_start."""
+
+        def _sse(event_type: str, payload: dict) -> bytes:
+            return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
+
+        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] = []
+        if open_index is not None:
+            chunks.append(_sse("content_block_stop", {"type": "content_block_stop", "index": open_index}))
+        chunks += [
+            _sse(
+                "content_block_start",
+                {
+                    "type": "content_block_start",
+                    "index": new_index,
+                    "content_block": {"type": "text", "text": ""},
+                },
+            ),
+            _sse(
+                "content_block_delta",
+                {
+                    "type": "content_block_delta",
+                    "index": new_index,
+                    "delta": {"type": "text_delta", "text": exc.message},
+                },
+            ),
+            _sse("content_block_stop", {"type": "content_block_stop", "index": new_index}),
+            _sse(
+                "message_delta",
+                {
+                    "type": "message_delta",
+                    "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                    "usage": {"output_tokens": 0},
+                },
+            ),
+            _sse("message_stop", {"type": "message_stop"}),
+        ]
+        return chunks
+
+    @staticmethod
+    def _content_block_state(
+        responses_so_far: list[Any],
+    ) -> tuple[Optional[int], Optional[int]]:
+        """From the SSE chunks already sent to the client, return (open
+        content-block index or None, highest content-block index seen or None).
+
+        A single streamed item may bundle multiple SSE events (raw bytes) or be
+        an already-parsed event dict, so every event across every item is
+        considered -- matching how ``get_streaming_string_so_far`` reads the
+        same stream."""
+        open_indices: set[int] = set()
+        max_index: Optional[int] = None
+        for item in responses_so_far:
+            for data in AnthropicMessagesHandler._iter_sse_events(item):
+                event_type = data.get("type")
+                index = data.get("index")
+                if not isinstance(index, int):
+                    continue
+                if event_type == "content_block_start":
+                    open_indices.add(index)
+                    max_index = index if max_index is None else max(max_index, index)
+                elif event_type == "content_block_stop":
+                    open_indices.discard(index)
+        open_index = max(open_indices) if open_indices else None
+        return open_index, max_index
+
+    @staticmethod
+    def _iter_sse_events(item: Any) -> list[dict]:
+        """Yield the event-data dicts in one stream chunk.
+
+        Handles both formats this stream can carry (see
+        ``get_streaming_string_so_far``): raw SSE ``bytes`` -- which may bundle
+        several events separated by a blank line -- and an already-parsed event
+        ``dict``."""
+        if isinstance(item, dict):
+            return [item]
+        if not isinstance(item, (bytes, bytearray)):
+            return []
+        events: list[dict] = []
+        for block in item.decode("utf-8", errors="replace").split("\n\n"):
+            for line in block.split("\n"):
+                line = line.strip()
+                if not line.startswith("data:"):
+                    continue
+                try:
+                    parsed = json.loads(line[len("data:") :].strip())
+                except json.JSONDecodeError:
+                    continue
+                if isinstance(parsed, dict):
+                    events.append(parsed)
+        return events
+
     def _translate_to_openai(self, data: dict) -> ChatCompletionRequest:
         """Translate Anthropic request to OpenAI chat completion format."""
         (

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -2,7 +2,10 @@ from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 if TYPE_CHECKING:
-    from litellm.integrations.custom_guardrail import CustomGuardrail
+    from litellm.integrations.custom_guardrail import (
+        CustomGuardrail,
+        ModifyResponseException,
+    )
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
     from litellm.proxy._types import UserAPIKeyAuth
     from litellm.types.llms.openai import AllMessageValues
@@ -98,6 +101,30 @@ class BaseTranslation(ABC):
         """
         return responses_so_far
 
+    def build_block_sse_chunks(
+        self,
+        exc: "ModifyResponseException",
+        stream_started: bool = False,
+        responses_so_far: Optional[list[Any]] = None,
+    ) -> Optional[list[bytes]]:
+        """
+        Build the streaming chunks that deliver a guardrail block message and
+        cleanly terminate the stream in this provider's wire format.
+
+        ``stream_started`` is True when real chunks were already sent to the
+        client: the result must *continue* the in-progress message (e.g. close
+        the open content block and append the block message) rather than start
+        a new one, which clients reject. ``responses_so_far`` provides the prior
+        chunks needed to do so. When False, nothing has been sent and a
+        standalone block message is emitted.
+
+        Returns None when the format has no safe terminator; the caller then
+        re-raises ``exc`` so the proxy can surface a clean error instead.
+        Override in provider subclasses that support synthesizing a block
+        stream.
+        """
+        return None
+
     def get_structured_messages(self, data: dict) -> Optional[List["AllMessageValues"]]:
         """
         Convert request data to OpenAI-spec structured messages.

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

@@ -8,7 +8,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint
 
 import copy
 import json
-from typing import Any, AsyncGenerator, List, Optional, Union
+from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Optional, Union
 
 from fastapi import HTTPException
 
@@ -23,6 +23,11 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import CallTypes, CallTypesLiteral
 
+if TYPE_CHECKING:
+    # Imported lazily at runtime (inside the streaming hook) to avoid a
+    # module-level cyclic import with litellm.integrations.custom_guardrail.
+    from litellm.integrations.custom_guardrail import ModifyResponseException
+
 # Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error
 A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message)
 
@@ -263,6 +268,30 @@ class UnifiedLLMGuardrails(CustomLogger):
 
         return response
 
+    async def _handle_streaming_block(
+        self,
+        exc: "ModifyResponseException",
+        endpoint_translation: Any,
+        stream_started: bool,
+        responses_so_far: list[Any],
+    ) -> AsyncGenerator[Any, None]:
+        """
+        Terminate a streamed response cleanly when a guardrail blocks it.
+
+        Format-agnostic routing: delegates to the provider translation handler's
+        ``build_block_sse_chunks`` (see ``BaseTranslation.build_block_sse_chunks``
+        for the ``stream_started`` / ``responses_so_far`` contract). When the
+        format has no safe terminator the handler returns None and we re-raise
+        ``exc`` so the proxy can surface a clean error.
+        """
+        block_chunks = endpoint_translation.build_block_sse_chunks(
+            exc, stream_started=stream_started, responses_so_far=responses_so_far
+        )
+        if block_chunks is None:
+            raise exc
+        for chunk in block_chunks:
+            yield chunk
+
     async def async_post_call_streaming_iterator_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
@@ -284,26 +313,38 @@ class UnifiedLLMGuardrails(CustomLogger):
 
         global endpoint_guardrail_translation_mappings
 
+        # Local import avoids a module-level cyclic import with
+        # litellm.integrations.custom_guardrail.
+        from litellm.integrations.custom_guardrail import ModifyResponseException
+
         guardrail_to_apply: CustomGuardrail = request_data.pop("guardrail_to_apply", None)
 
-        # Get streaming configuration from guardrail or optional_params
-        sampling_rate = 5
-        end_of_stream_only = False  # If True, only apply guardrail at end of stream
+        # Get streaming configuration, with precedence: 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:
+                value = getattr(guardrail_to_apply, name, value)
+                config = getattr(guardrail_to_apply, "guardrail_config", {})
+                if isinstance(config, dict):
+                    value = config.get(name, value)
+            return self.optional_params.get(name, value)
 
-        if guardrail_to_apply is not None:
-            # Check direct attributes on guardrail first
-            sampling_rate = getattr(guardrail_to_apply, "streaming_sampling_rate", sampling_rate)
-            end_of_stream_only = getattr(guardrail_to_apply, "streaming_end_of_stream_only", end_of_stream_only)
+        sampling_rate = _streaming_flag("streaming_sampling_rate", 5)
+        # Only apply the guardrail at end of stream (not per chunk).
+        end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False)
+        # 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.
+        buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", False)
 
-            # Also check guardrail_config dict if present
-            guardrail_config = getattr(guardrail_to_apply, "guardrail_config", {})
-            if isinstance(guardrail_config, dict):
-                sampling_rate = guardrail_config.get("streaming_sampling_rate", sampling_rate)
-                end_of_stream_only = guardrail_config.get("streaming_end_of_stream_only", end_of_stream_only)
-
-        # Also check optional_params as fallback
-        sampling_rate = self.optional_params.get("streaming_sampling_rate", sampling_rate)
-        end_of_stream_only = self.optional_params.get("streaming_end_of_stream_only", end_of_stream_only)
+        # Buffering can only moderate the assembled response, so it always
+        # defers to end-of-stream.
+        if buffer_until_moderated:
+            end_of_stream_only = True
 
         if guardrail_to_apply is None:
             async for item in response:
@@ -328,6 +369,10 @@ class UnifiedLLMGuardrails(CustomLogger):
         call_type = None
         chunk_counter = 0
         responses_so_far: List[Any] = []
+        # Whether any real response chunk has been forwarded to the client.
+        # Drives how a block terminates the stream: continue the in-progress
+        # message (True) vs emit a standalone block message (False, buffered).
+        chunks_yielded = False
 
         async for item in response:
             chunk_counter += 1
@@ -349,9 +394,14 @@ class UnifiedLLMGuardrails(CustomLogger):
                     yield remaining_item
                 return
 
-            # If end_of_stream_only mode, yield chunks without processing
+            # If end_of_stream_only mode, yield chunks without processing.
+            # When buffering, withhold them instead -- they are released (or
+            # replaced by the block message) only after end-of-stream
+            # moderation runs below.
             if end_of_stream_only:
-                yield item
+                if not buffer_until_moderated:
+                    chunks_yielded = True
+                    yield item
                 continue
 
             # Process chunk based on sampling rate
@@ -381,6 +431,24 @@ class UnifiedLLMGuardrails(CustomLogger):
                         user_api_key_dict=user_api_key_dict,
                         request_data=request_data,
                     )
+                except ModifyResponseException as e:
+                    # Guardrail blocked the response mid-stream. Emit a clean
+                    # terminating SSE sequence delivering the block message
+                    # instead of letting the exception propagate into a bare
+                    # `data: {"error": ...}` blob (which truncates the stream).
+                    # Chunks have already been forwarded here, so the block
+                    # continues the in-progress message (stream_started=True).
+                    # The current chunk was appended to responses_so_far but not
+                    # yet yielded, so exclude it: the continuation must reflect
+                    # only what the client has actually received.
+                    async for block_chunk in self._handle_streaming_block(
+                        e,
+                        endpoint_translation,
+                        stream_started=chunks_yielded,
+                        responses_so_far=responses_so_far[:-1],
+                    ):
+                        yield block_chunk
+                    return
                 except HTTPException as e:
                     # Response already started (we already yielded chunks); cannot send 400.
                     # For A2A (NDJSON), yield an in-stream JSON-RPC error so the client sees it.
@@ -407,8 +475,10 @@ class UnifiedLLMGuardrails(CustomLogger):
                         yield error_chunk
                         return
                     raise
+                chunks_yielded = True
                 yield original_item
             else:
+                chunks_yielded = True
                 yield item
 
         # Stream has ended - do final processing with all collected chunks
@@ -421,6 +491,15 @@ class UnifiedLLMGuardrails(CustomLogger):
 
             endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
 
+            # When buffering, snapshot the original chunks before moderation.
+            # A shallow copy suffices: end-of-stream
+            # process_output_streaming_response builds a separate assembled
+            # response (it does not mutate the individual chunks in place), and
+            # the chunks themselves are replayed verbatim -- so we only need to
+            # preserve the list, not clone every chunk (deepcopy would double
+            # peak memory for large responses).
+            buffered_items = list(responses_so_far) if buffer_until_moderated else None
+
             try:
                 await endpoint_translation.process_output_streaming_response(
                     responses_so_far=responses_so_far,
@@ -429,6 +508,23 @@ class UnifiedLLMGuardrails(CustomLogger):
                     user_api_key_dict=user_api_key_dict,
                     request_data=request_data,
                 )
+                # Moderation passed: release the withheld original chunks.
+                if buffered_items is not None:
+                    for buffered_item in buffered_items:
+                        yield buffered_item
+            except ModifyResponseException as e:
+                # Block detected during end-of-stream processing. Emit a clean
+                # terminating SSE sequence with the block message rather than
+                # propagating into a bare error blob that truncates the stream.
+                # The withheld original chunks are never released.
+                async for block_chunk in self._handle_streaming_block(
+                    e,
+                    endpoint_translation,
+                    stream_started=chunks_yielded,
+                    responses_so_far=responses_so_far,
+                ):
+                    yield block_chunk
+                return
             except HTTPException as e:
                 if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
                     request_id = _get_a2a_request_id(responses_so_far, request_data)

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

@@ -0,0 +1,275 @@
+"""
+Regression tests for blocking an Anthropic streaming response from the
+unified guardrail post-call streaming iterator hook.
+
+When a guardrail's ``apply_guardrail`` raises ``ModifyResponseException`` while
+(or at the end of) an Anthropic ``/v1/messages`` stream is being relayed, the
+hook must emit a well-formed Anthropic SSE termination sequence carrying the
+block message - NOT a bare ``data: {"error": ...}`` blob that truncates the
+stream and causes the Anthropic SDK parser to discard the response.
+"""
+
+import json
+from typing import Any, List, Literal, Optional
+
+import pytest
+
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    ModifyResponseException,
+)
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
+    UnifiedLLMGuardrails,
+)
+from litellm.types.utils import GenericGuardrailAPIInputs
+
+BLOCK_MESSAGE = "Blocked by policy: this response was withheld."
+
+
+class _BlockingGuardrail(CustomGuardrail):
+    """Mock guardrail that always blocks by raising ModifyResponseException."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        raise ModifyResponseException(
+            message=BLOCK_MESSAGE,
+            model="claude-3-5-sonnet",
+            request_data=request_data,
+            guardrail_name=self.guardrail_name,
+        )
+
+
+def _sse_event(event_type: str, data: dict) -> bytes:
+    return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
+
+
+async def _anthropic_stream(end: bool):
+    """Yield Anthropic SSE byte chunks. If end=True, include a terminating
+    message_delta (stop_reason set) so the hook's end-of-stream path runs."""
+    yield _sse_event(
+        "message_start",
+        {
+            "type": "message_start",
+            "message": {
+                "id": "msg_orig",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude-3-5-sonnet",
+                "content": [],
+                "stop_reason": None,
+                "usage": {"input_tokens": 1, "output_tokens": 0},
+            },
+        },
+    )
+    yield _sse_event(
+        "content_block_start",
+        {
+            "type": "content_block_start",
+            "index": 0,
+            "content_block": {"type": "text", "text": ""},
+        },
+    )
+    for text in ["This ", "is ", "the ", "original ", "answer."]:
+        yield _sse_event(
+            "content_block_delta",
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "text_delta", "text": text},
+            },
+        )
+    if end:
+        yield _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+        yield _sse_event(
+            "message_delta",
+            {
+                "type": "message_delta",
+                "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                "usage": {"output_tokens": 5},
+            },
+        )
+        yield _sse_event("message_stop", {"type": "message_stop"})
+
+
+def _decode(chunks: List[Any]) -> str:
+    parts = []
+    for chunk in chunks:
+        parts.append(chunk.decode() if isinstance(chunk, bytes) else str(chunk))
+    return "".join(parts)
+
+
+def _parse_sse_event_types(raw: str) -> List[str]:
+    event_types = []
+    for block in raw.split("\n\n"):
+        for line in block.strip().split("\n"):
+            if line.startswith("data:"):
+                payload = line[len("data:") :].strip()
+                try:
+                    event_types.append(json.loads(payload).get("type"))
+                except json.JSONDecodeError:
+                    pass
+    return event_types
+
+
+async def _run_hook(end: bool, sampling_rate: int = 1) -> str:
+    guardrail = _BlockingGuardrail(guardrail_name="test-blocking-guardrail", event_hook="post_call")
+    # sampling_rate controls how many chunks are forwarded before the block
+    # fires: 1 blocks on the first chunk (nothing sent yet); >1 forwards earlier
+    # chunks first, exercising the mid-stream "continue the message" path.
+    guardrail.streaming_sampling_rate = sampling_rate
+
+    unified_guardrail = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/messages")
+    request_data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": ["test-blocking-guardrail"]},
+    }
+
+    collected: List[Any] = []
+    async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=user_api_key_dict,
+        response=_anthropic_stream(end=end),
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    return _decode(collected)
+
+
+def _assert_clean_block_termination(raw: str) -> None:
+    # No bare error blob that would truncate the stream.
+    assert '"error"' not in raw, f"unexpected error blob in stream: {raw!r}"
+    # The block message is delivered as assistant text.
+    assert BLOCK_MESSAGE in raw, f"block message missing from stream: {raw!r}"
+    # A complete, parseable Anthropic SSE termination sequence is present.
+    event_types = _parse_sse_event_types(raw)
+    assert "message_start" in event_types
+    assert "content_block_delta" in event_types
+    # Exactly one message_start: a block must never inject a second
+    # message envelope into an already-started stream (clients reject it).
+    assert event_types.count("message_start") == 1, f"expected a single message_start, got: {event_types}"
+    assert event_types[-1] == "message_stop", f"stream did not end cleanly: {event_types}"
+    # message_delta carries a stop_reason.
+    assert any('"stop_reason"' in block and "message_delta" in block for block in raw.split("\n\n"))
+
+
+@pytest.mark.asyncio
+async def test_mid_stream_block_emits_clean_anthropic_sse():
+    """Per-chunk block: a clean SSE termination with the block message, no error blob."""
+    raw = await _run_hook(end=False)
+    _assert_clean_block_termination(raw)
+
+
+@pytest.mark.asyncio
+async def test_end_of_stream_block_emits_clean_anthropic_sse():
+    """End-of-stream block: same clean SSE termination guarantees."""
+    raw = await _run_hook(end=True)
+    _assert_clean_block_termination(raw)
+
+
+@pytest.mark.asyncio
+async def test_mid_stream_block_after_prior_chunks_continues_message():
+    """Regression: when real chunks were already forwarded (sampling_rate>1),
+    the block must continue the in-progress message, not start a second one."""
+    raw = await _run_hook(end=False, sampling_rate=5)
+    # Some original content was forwarded before the block...
+    assert "message_start" in raw
+    # ...and the block continues that same message (single message_start) with
+    # the block message appended, ending cleanly.
+    _assert_clean_block_termination(raw)
+
+
+class TestContentBlockState:
+    """`_content_block_state` must reflect the true open/last block index across
+    the two chunk formats the stream can carry (multi-event bytes, parsed dict),
+    so a mid-stream block closes/opens the right indices."""
+
+    def _handler(self):
+        from litellm.llms.anthropic.chat.guardrail_translation.handler import (
+            AnthropicMessagesHandler,
+        )
+
+        return AnthropicMessagesHandler()
+
+    def test_multi_event_bytes_chunk_is_fully_parsed(self):
+        # One item bundles start(0) + delta + stop(0): the block is already
+        # closed, so open_index is None (not 0) and max_index is 0.
+        bundled = (
+            _sse_event(
+                "content_block_start",
+                {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+            )
+            + _sse_event(
+                "content_block_delta",
+                {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}},
+            )
+            + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+        )
+        open_index, max_index = self._handler()._content_block_state([bundled])
+        assert open_index is None
+        assert max_index == 0
+
+    def test_open_block_across_separate_chunks(self):
+        chunks = [
+            _sse_event(
+                "content_block_start",
+                {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+            ),
+            _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}),
+            _sse_event(
+                "content_block_start",
+                {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}},
+            ),
+        ]
+        open_index, max_index = self._handler()._content_block_state(chunks)
+        assert open_index == 1
+        assert max_index == 1
+
+    def test_dict_format_chunks_are_parsed(self):
+        # The backwards-compat parsed-dict format must be understood too.
+        chunks = [
+            {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+            {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}},
+        ]
+        open_index, max_index = self._handler()._content_block_state(chunks)
+        assert open_index == 0
+        assert max_index == 0
+
+    def test_continuation_closes_open_block_and_appends_after_it(self):
+        from litellm.integrations.custom_guardrail import ModifyResponseException
+
+        handler = self._handler()
+        # Client has seen an open text block at index 0.
+        seen = [
+            _sse_event(
+                "content_block_start",
+                {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+            ),
+            _sse_event(
+                "content_block_delta",
+                {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}},
+            ),
+        ]
+        exc = ModifyResponseException(
+            message=BLOCK_MESSAGE, model="claude-3-5-sonnet", request_data={}, guardrail_name="g"
+        )
+        raw = b"".join(handler.build_block_sse_chunks(exc, stream_started=True, responses_so_far=seen)).decode()
+        events = _parse_sse_event_types(raw)
+        # No new message envelope, closes block 0, appends block text at index 1.
+        assert "message_start" not in events
+        assert events == [
+            "content_block_stop",
+            "content_block_start",
+            "content_block_delta",
+            "content_block_stop",
+            "message_delta",
+            "message_stop",
+        ]
+        assert BLOCK_MESSAGE in raw
+        assert '"index": 1' in raw

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

@@ -0,0 +1,159 @@
+"""
+Tests for ``streaming_buffer_until_moderated`` on the unified guardrail
+post-call streaming iterator hook.
+
+With this flag set, the hook must withhold every upstream chunk until
+end-of-stream moderation has run. The decisive guarantee versus the
+detect-only ``streaming_end_of_stream_only`` behavior: when the guardrail
+blocks, the original (objectionable) content is NEVER yielded to the client --
+only the block message is. On a clean response, all original chunks are
+released unchanged after moderation passes.
+"""
+
+import json
+from typing import Any, List, Literal, Optional
+
+import pytest
+
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    ModifyResponseException,
+)
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
+    UnifiedLLMGuardrails,
+)
+from litellm.types.utils import GenericGuardrailAPIInputs
+
+BLOCK_MESSAGE = "Blocked by policy: this response was withheld."
+ORIGINAL_MARKER = "ORIGINAL-SECRET-ANSWER"
+
+
+class _BlockingGuardrail(CustomGuardrail):
+    """Always blocks at moderation time."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        raise ModifyResponseException(
+            message=BLOCK_MESSAGE,
+            model="claude-3-5-sonnet",
+            request_data=request_data,
+            guardrail_name=self.guardrail_name,
+        )
+
+
+class _PassingGuardrail(CustomGuardrail):
+    """Never blocks; returns inputs unchanged."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        return inputs
+
+
+def _sse_event(event_type: str, data: dict) -> bytes:
+    return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
+
+
+async def _anthropic_stream():
+    """A complete Anthropic /v1/messages SSE stream whose assistant text
+    contains ORIGINAL_MARKER so leakage is unambiguous to assert."""
+    yield _sse_event(
+        "message_start",
+        {
+            "type": "message_start",
+            "message": {
+                "id": "msg_orig",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude-3-5-sonnet",
+                "content": [],
+                "stop_reason": None,
+                "usage": {"input_tokens": 1, "output_tokens": 0},
+            },
+        },
+    )
+    yield _sse_event(
+        "content_block_start",
+        {
+            "type": "content_block_start",
+            "index": 0,
+            "content_block": {"type": "text", "text": ""},
+        },
+    )
+    for text in ["Here is ", "the ", ORIGINAL_MARKER, " for you."]:
+        yield _sse_event(
+            "content_block_delta",
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "text_delta", "text": text},
+            },
+        )
+    yield _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+    yield _sse_event(
+        "message_delta",
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "usage": {"output_tokens": 5},
+        },
+    )
+    yield _sse_event("message_stop", {"type": "message_stop"})
+
+
+def _decode(chunks: List[Any]) -> str:
+    return "".join(c.decode() if isinstance(c, bytes) else str(c) for c in chunks)
+
+
+async def _run(guardrail: CustomGuardrail) -> str:
+    # Rubrik's real config: end-of-stream-only moderation. Without buffering
+    # this releases every chunk before moderation runs (content leaks on
+    # block); the buffer flag must change that to moderate-then-release.
+    guardrail.streaming_end_of_stream_only = True
+    guardrail.streaming_buffer_until_moderated = True
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/messages")
+    request_data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": [guardrail.guardrail_name]},
+    }
+    collected: List[Any] = []
+    async for chunk in unified.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=user_api_key_dict,
+        response=_anthropic_stream(),
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    return _decode(collected)
+
+
+@pytest.mark.asyncio
+async def test_buffered_block_withholds_original_content():
+    raw = await _run(_BlockingGuardrail(guardrail_name="blk", event_hook="post_call"))
+    # The original content must never reach the client...
+    assert ORIGINAL_MARKER not in raw, f"original content leaked: {raw!r}"
+    # ...only the block message, in a clean terminating stream.
+    assert BLOCK_MESSAGE in raw
+    assert '"error"' not in raw
+
+
+@pytest.mark.asyncio
+async def test_buffered_clean_releases_all_content():
+    raw = await _run(_PassingGuardrail(guardrail_name="pass", event_hook="post_call"))
+    # A clean response is released in full after moderation passes.
+    assert ORIGINAL_MARKER in raw
+    assert (
+        raw.rstrip().endswith('event: message_stop\ndata: {"type": "message_stop"}'.rstrip()) or "message_stop" in raw
+    )
+    assert BLOCK_MESSAGE not in raw