Merge 194c3de142 into 6fa0477d84

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

@@ -519,6 +519,26 @@ def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewrit
     return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}}
 
 
+def _ended_stream_tool_calls(
+    returned_tool_calls: Sequence[object] | None,
+    *,
+    streamed_tool_calls: Sequence[object],
+    expected_count: int,
+    guardrail_name: str,
+) -> Sequence[object]:
+    """A guardrail that adds or drops a tool call on an ended stream cannot be written back, block for block."""
+    if returned_tool_calls is None:
+        return streamed_tool_calls
+    if len(returned_tool_calls) != expected_count:
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+        raise UndeliverableStreamRewrite(
+            guardrail_name,
+            f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried {expected_count}",
+        )
+    return returned_tool_calls
+
+
 def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]:
     """The guardrail-visible shape of each tool call, whether the guardrail handed
     back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts."""
@@ -1634,10 +1654,12 @@ class AnthropicMessagesHandler(BaseTranslation):
                         pre_guardrail_tool_calls=pre_guardrail_tool_calls,
                         post_guardrail_tool_calls=_with_query_rewrites(
                             _tool_call_shapes(
-                                returned_tool_calls
-                                if isinstance(returned_tool_calls, list)
-                                and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
-                                else tool_calls_list or ()
+                                _ended_stream_tool_calls(
+                                    returned_tool_calls,
+                                    streamed_tool_calls=tool_calls_list or (),
+                                    expected_count=len(pre_guardrail_tool_calls),
+                                    guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
+                                )
                             ),
                             tool_call_blocks,
                             _query_rewrites(

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

@@ -393,6 +393,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
         user_api_key_dict: "UserAPIKeyAuth | None" = None,
         request_data: dict | None = None,
+        *,
+        deliver_ended_stream_rewrites: bool = False,
     ) -> ModelResponse:
         """
         Process output response by applying guardrails to text content.
@@ -458,6 +460,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             if hasattr(response, "model") and response.model:
                 inputs["model"] = response.model
 
+            pre_guardrail_tool_call_count: Final = len(tool_calls_to_check)
             guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
                 inputs=inputs,
                 request_data=request_data,
@@ -467,6 +470,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
 
             guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
             returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
+            if (
+                deliver_ended_stream_rewrites
+                and isinstance(returned_tool_calls, list)
+                and len(returned_tool_calls) != pre_guardrail_tool_call_count
+            ):
+                from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+                raise UndeliverableStreamRewrite(
+                    guardrail_to_apply.guardrail_name or "unknown",
+                    f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried "
+                    f"{pre_guardrail_tool_call_count}",
+                )
             guardrailed_tool_calls: Final[list[dict[str, object]]] = (
                 cast(list[dict[str, object]], returned_tool_calls)
                 if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check)
@@ -667,6 +682,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             litellm_logging_obj=litellm_logging_obj,
             user_api_key_dict=user_api_key_dict,
             request_data=request_data,
+            deliver_ended_stream_rewrites=deliver_ended_stream_rewrites,
         )
         if not deliver_ended_stream_rewrites:
             return

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -177,9 +177,23 @@ def _post_guardrail_tool_call_shapes(
     returned_tool_calls: Sequence[object] | None,
     pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
     guardrail_name: str | None,
+    *,
+    deliver_ended_stream_rewrites: bool = False,
 ) -> tuple[_ToolCallShape, ...]:
     if not pre_guardrail_tool_calls:
         return pre_guardrail_tool_calls
+    if (
+        deliver_ended_stream_rewrites
+        and returned_tool_calls is not None
+        and len(returned_tool_calls) != len(pre_guardrail_tool_calls)
+    ):
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+        raise UndeliverableStreamRewrite(
+            guardrail_name or "unknown",
+            f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried "
+            f"{len(pre_guardrail_tool_calls)}",
+        )
     if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls):
         verbose_proxy_logger.warning(
             "OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, "
@@ -1284,6 +1298,7 @@ class OpenAIResponsesHandler(BaseTranslation):
                     returned_tool_calls=guardrailed_inputs.get("tool_calls"),
                     pre_guardrail_tool_calls=pre_guardrail_tool_calls,
                     guardrail_name=guardrail_to_apply.guardrail_name,
+                    deliver_ended_stream_rewrites=deliver_ended_stream_rewrites,
                 )
 
                 # Write guardrailed texts back into the output items in-place.

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

@@ -9,6 +9,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint
 import asyncio
 import contextlib
 import copy
+import functools
 import json
 from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime_checkable
@@ -234,6 +235,57 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N
             data["litellm_metadata"] = user_metadata
 
 
+async def _scan_ended_stream(
+    endpoint_translation: "BaseTranslation",
+    *,
+    guardrail_to_apply: CustomGuardrail,
+    responses_so_far: list[object],  # mutable-ok: the translation writes the rewrites into the caller's buffer
+    user_api_key_dict: UserAPIKeyAuth,
+    request_data: dict[str, object],  # mutable-ok: handed through to the translation's dict parameter
+    buffered_items: tuple[object, ...] | None,
+    deliver_rewrites: bool,
+    shield: bool,
+) -> tuple[object, ...] | None:
+    """Run the end-of-stream scan and return the chunks to release to the client.
+
+    A fully withheld stream holds every chunk in ``responses_so_far``, so the guardrail's text and tool-call rewrites
+    are released from there, whether the translation rewrote the chunks in place or replaced them (SSE bytes).
+    A rewrite the translation cannot place (a changed tool-call count) fails closed: nothing is released and the
+    client gets a 500, since the originals carry the very content the guardrail rewrote.
+    """
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+    logging_obj: Final = request_data.get("litellm_logging_obj")
+    scan: Final = functools.partial(
+        endpoint_translation.process_output_streaming_response,
+        responses_so_far=responses_so_far,
+        guardrail_to_apply=guardrail_to_apply,
+        litellm_logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
+        user_api_key_dict=user_api_key_dict,
+        request_data=request_data,
+    )
+    try:
+        with anyio.CancelScope(shield=shield):
+            await (scan(deliver_ended_stream_rewrites=True) if deliver_rewrites else scan())
+    except UndeliverableStreamRewrite as undeliverable:
+        verbose_proxy_logger.warning(
+            "UnifiedLLMGuardrails: guardrail %s rewrote the buffered stream but the rewrite cannot be written back "
+            "(%s). The response was withheld",
+            undeliverable.guardrail_name,
+            undeliverable.reason,
+        )
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "error": f"Guardrail {undeliverable.guardrail_name} rewrote the response but the rewrite cannot be "
+                f"delivered on this stream ({undeliverable.reason}). The response was withheld",
+                "guardrail": undeliverable.guardrail_name,
+            },
+        ) from undeliverable
+    return tuple(responses_so_far) if deliver_rewrites else buffered_items
+
+
 class UnifiedLLMGuardrails(CustomLogger):
     def __init__(
         self,
@@ -1158,31 +1210,13 @@ class UnifiedLLMGuardrails(CustomLogger):
         # path; "incremental_diff" emits them as synthetic deltas (see
         # _run_incremental_transform_stream).
         streaming_transform_mode: Final = _streaming_setting("streaming_transform_mode", "block_only")
-        # 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. 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: object = _streaming_setting(
+        # Withhold every chunk until end-of-stream moderation passes, then release the chunks
+        # (with the guardrail's rewrite written into them) or only the block message.
+        buffer_until_moderated: Final = _streaming_setting(
             "streaming_buffer_until_moderated", buffer_until_moderated_default
         )
         release_on_scan: Final = _streaming_setting("streaming_buffer_release_on_scan", 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
-
         if buffer_until_moderated and not release_on_scan:
             end_of_stream_only = True
 
@@ -1476,19 +1510,25 @@ class UnifiedLLMGuardrails(CustomLogger):
                     return
 
                 try:
-                    with anyio.CancelScope(shield=chunks_yielded):
-                        if pending_scan is not None:
-                            pending_scan.result()  # a failed background scan raises into the handlers below
-                        await endpoint_translation.process_output_streaming_response(
-                            responses_so_far=responses_so_far,
-                            guardrail_to_apply=guardrail_to_apply,
-                            litellm_logging_obj=request_data.get("litellm_logging_obj"),
-                            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:
+                    if pending_scan is not None:
+                        pending_scan.result()  # a failed background scan raises into the handlers below
+                    released_items: Final = await _scan_ended_stream(
+                        endpoint_translation,
+                        guardrail_to_apply=guardrail_to_apply,
+                        responses_so_far=responses_so_far,
+                        user_api_key_dict=user_api_key_dict,
+                        request_data=request_data,
+                        buffered_items=buffered_items,
+                        deliver_rewrites=bool(
+                            buffer_until_moderated
+                            and end_of_stream_only
+                            and not responses_yielded
+                            and endpoint_translation.delivers_ended_stream_rewrites
+                        ),
+                        shield=chunks_yielded,
+                    )
+                    if released_items is not None:
+                        for buffered_item in released_items:
                             yield buffered_item
                     for pending_item in pending_end_of_stream_items:
                         responses_yielded.append(pending_item)

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

@@ -14,6 +14,7 @@ import json
 from typing import Any, AsyncGenerator, List, Literal, Optional
 
 import pytest
+from fastapi import HTTPException
 
 import litellm
 from litellm.integrations.custom_guardrail import (
@@ -30,6 +31,8 @@ from litellm.proxy.utils import _pipelines_stream_live
 from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
 from litellm.types.utils import (
     ChatCompletionDeltaToolCall,
+    ChatCompletionMessageToolCall,
+    ChatCompletionToolCallChunk,
     Delta,
     Function,
     FunctionCall,
@@ -562,20 +565,378 @@ async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buff
     assert guardrail.scan_count == 1
 
 
+MASKED_MARKER = "[MASKED]"
+
+
+class _MaskingGuardrail(_CountingPassingGuardrail):
+    """Rewrites the markers out of the response text and tool-call arguments, as a PII masker does."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Any | None = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.scan_count += 1
+        masked: GenericGuardrailAPIInputs = {
+            "texts": [text.replace(ORIGINAL_MARKER, MASKED_MARKER) for text in inputs.get("texts") or []]
+        }
+        tool_calls = inputs.get("tool_calls")
+        if tool_calls:
+            masked["tool_calls"] = [_mask_tool_call(tool_call) for tool_call in tool_calls]
+        return masked
+
+
+def _mask_tool_call(
+    tool_call: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall,
+) -> ChatCompletionToolCallChunk | ChatCompletionMessageToolCall:
+    """Handlers hand chat tool calls over as dicts and /v1/messages tool calls as pydantic objects."""
+    if isinstance(tool_call, dict):
+        function = tool_call["function"]
+        arguments = function["arguments"].replace(TOOL_ARGUMENTS_MARKER, MASKED_MARKER)
+        return {**tool_call, "function": {**function, "arguments": arguments}}
+    tool_call.function.arguments = tool_call.function.arguments.replace(TOOL_ARGUMENTS_MARKER, MASKED_MARKER)
+    return tool_call
+
+
+class _ToolDroppingGuardrail(_MaskingGuardrail):
+    """Masks like _MaskingGuardrail but returns no tool calls at all, as a guardrail that strips a disallowed tool does."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Any | None = None,
+    ) -> GenericGuardrailAPIInputs:
+        masked = await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
+        return {**masked, "tool_calls": []}
+
+
+class _ToolClearingGuardrail(_MaskingGuardrail):
+    """Empties the handler's own tool-call list in place and hands that same list back."""
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Any | None = None,
+    ) -> GenericGuardrailAPIInputs:
+        masked = await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
+        tool_calls = inputs.get("tool_calls")
+        if tool_calls is not None:
+            tool_calls.clear()
+        return {**masked, "tool_calls": tool_calls or []}
+
+
+async def _anthropic_tool_use_stream():
+    """A complete /v1/messages SSE stream with a text block and a tool_use block whose input carries the marker."""
+    yield _sse_event(
+        "message_start",
+        {
+            "type": "message_start",
+            "message": {
+                "id": "msg_tool",
+                "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": ""}},
+    )
+    yield _sse_event(
+        "content_block_delta",
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "text_delta", "text": f"Sending {ORIGINAL_MARKER}"},
+        },
+    )
+    yield _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+    yield _sse_event(
+        "content_block_start",
+        {
+            "type": "content_block_start",
+            "index": 1,
+            "content_block": {"type": "tool_use", "id": "toolu_1", "name": "send_email", "input": {}},
+        },
+    )
+    for partial in ['{"to": "', TOOL_ARGUMENTS_MARKER, '"}']:
+        yield _sse_event(
+            "content_block_delta",
+            {"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": partial}},
+        )
+    yield _sse_event("content_block_stop", {"type": "content_block_stop", "index": 1})
+    yield _sse_event(
+        "message_delta",
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "tool_use", "stop_sequence": None},
+            "usage": {"output_tokens": 5},
+        },
+    )
+    yield _sse_event("message_stop", {"type": "message_stop"})
+
+
+async def _run_messages_stream(guardrail: CustomGuardrail, stream: AsyncGenerator[bytes, None]) -> str:
+    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=stream,
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    return _decode(collected)
+
+
+async def _run_buffered_chat(
+    guardrail: CustomGuardrail, content_chunks: list[str], tool_argument_chunks: list[str] | None = None
+) -> list[Any]:
+    """The plain ``streaming_buffer_until_moderated: true`` config: hold the whole stream, scan once at the end."""
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = False
+    guardrail.streaming_end_of_stream_only = False
+    guardrail.streaming_sampling_rate = 2
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions")
+    request_data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": [guardrail.guardrail_name]},
+    }
+    collected: list[Any] = []
+    yielded_count: list[int] = []
+    async for chunk in unified.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=user_api_key_dict,
+        response=_windowed_chat_stream(yielded_count, collected, content_chunks, tool_argument_chunks, False),
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    assert yielded_count == [0] * len(yielded_count), "buffered stream released a chunk before the end scan"
+    return collected
+
+
 @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)
+async def test_buffered_messages_stream_delivers_masked_text_to_the_client():
+    """A masking guardrail on a fully buffered /v1/messages stream, the client gets the masked text, never the original."""
+    guardrail = _MaskingGuardrail(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 guardrail.scan_count == 1
+    assert MASKED_MARKER in raw, raw
+    assert ORIGINAL_MARKER not in raw, raw
     assert BLOCK_MESSAGE not in raw
 
 
+@pytest.mark.asyncio
+async def test_buffered_messages_stream_delivers_masked_tool_call_input():
+    guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call")
+    raw = await _run_messages_stream(guardrail, _anthropic_tool_use_stream())
+    assert guardrail.scan_count == 1
+    assert TOOL_ARGUMENTS_MARKER not in raw, raw
+    assert ORIGINAL_MARKER not in raw, raw
+    assert raw.count(MASKED_MARKER) == 2, raw
+    assert "message_stop" in raw
+
+
+@pytest.mark.asyncio
+async def test_buffered_chat_stream_delivers_masked_text_and_tool_arguments():
+    guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call")
+    collected = await _run_buffered_chat(
+        guardrail, [f"call {ORIGINAL_MARKER} now"], ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']
+    )
+    assert _chat_text(collected) == f"call {MASKED_MARKER} now"
+    assert _tool_argument_text(collected) == f'{{"cmd": "{MASKED_MARKER}"}}'
+    assert collected[-1].choices[0].finish_reason == "tool_calls"
+
+
+def _two_choice_tool_call_chunk(finish_reason: str | None = None) -> ModelResponseStream:
+    """An ``n=2`` chunk: tool-call rewrites cannot be written back onto a multi-choice stream."""
+    return ModelResponseStream(
+        id="chatcmpl-two-choices",
+        created=1724900000,
+        model="gpt-4",
+        choices=[
+            StreamingChoices(
+                index=index,
+                finish_reason=finish_reason,
+                delta=Delta(
+                    role="assistant",
+                    content=f"call {ORIGINAL_MARKER} now",
+                    tool_calls=[
+                        ChatCompletionDeltaToolCall(
+                            id=f"call_{index}",
+                            type="function",
+                            index=0,
+                            function=Function(name="run_shell", arguments=f'{{"cmd": "{TOOL_ARGUMENTS_MARKER}"}}'),
+                        )
+                    ],
+                ),
+            )
+            for index in (0, 1)
+        ],
+    )
+
+
+async def _two_choice_tool_call_stream() -> AsyncGenerator[ModelResponseStream, None]:
+    yield _two_choice_tool_call_chunk()
+    yield _two_choice_tool_call_chunk(finish_reason="tool_calls")
+
+
+@pytest.mark.asyncio
+async def test_buffered_messages_stream_withholds_everything_when_the_guardrail_drops_a_tool_call():
+    """A tool call the guardrail removed must not reach the client as the original, the request fails with a 500."""
+    guardrail = _ToolDroppingGuardrail(guardrail_name="dropper", event_hook="post_call")
+    with pytest.raises(HTTPException) as raised:
+        await _run_messages_stream(guardrail, _anthropic_tool_use_stream())
+    assert raised.value.status_code == 500
+    assert "dropper" in str(raised.value.detail)
+    assert guardrail.scan_count == 1
+
+
+@pytest.mark.asyncio
+async def test_buffered_chat_stream_withholds_everything_when_the_tool_call_rewrite_cannot_be_placed():
+    """Neither the originals nor a half-applied rewrite reach the client, the request fails with a 500 instead."""
+    guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call")
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = False
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions")
+    request_data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": [guardrail.guardrail_name]},
+    }
+    collected: list[Any] = []
+
+    async def drain() -> None:
+        async for chunk in unified.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=user_api_key_dict, response=_two_choice_tool_call_stream(), request_data=request_data
+        ):
+            collected.append(chunk)
+
+    with pytest.raises(HTTPException) as raised:
+        await drain()
+    assert raised.value.status_code == 500
+    assert "masker" in str(raised.value.detail)
+    assert guardrail.scan_count == 1
+    assert collected == []
+
+
+@pytest.mark.asyncio
+async def test_buffered_responses_stream_delivers_masked_text_to_the_client():
+    guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call")
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = False
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses")
+    request_data = {
+        "input": "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=_replay(_responses_message_stream_events(["Here is ", ORIGINAL_MARKER])),
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    raw = json.dumps([chunk if isinstance(chunk, dict) else str(chunk) for chunk in collected])
+    assert guardrail.scan_count == 1
+    assert MASKED_MARKER in raw, raw
+    assert ORIGINAL_MARKER not in raw, raw
+    assert "response.completed" in raw
+
+
+@pytest.mark.asyncio
+async def test_buffered_responses_stream_withholds_everything_when_the_guardrail_drops_a_tool_call():
+    """A function call the guardrail removed must not reach the client as the original, the request fails with a 500."""
+    guardrail = _ToolDroppingGuardrail(guardrail_name="dropper", event_hook="post_call")
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = False
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses")
+    request_data = {
+        "input": "hi",
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": [guardrail.guardrail_name]},
+    }
+    collected: list[Any] = []
+
+    async def drain() -> None:
+        async for chunk in unified.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=user_api_key_dict,
+            response=_replay(
+                _responses_truncated_function_call_events("Running", ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'])
+            ),
+            request_data=request_data,
+        ):
+            collected.append(chunk)
+
+    with pytest.raises(HTTPException) as raised:
+        await drain()
+
+    assert raised.value.status_code == 500
+    assert "dropper" in str(raised.value.detail)
+    assert guardrail.scan_count == 1
+    assert collected == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "guardrail_cls", [_ToolDroppingGuardrail, _ToolClearingGuardrail], ids=["returns fewer", "clears in place"]
+)
+async def test_buffered_chat_stream_withholds_everything_when_the_guardrail_drops_a_tool_call(guardrail_cls: type):
+    guardrail = guardrail_cls(guardrail_name="dropper", event_hook="post_call")
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = False
+    with pytest.raises(HTTPException) as raised:
+        await _run_buffered_chat(guardrail, ["Running"], ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'])
+    assert raised.value.status_code == 500
+    assert "dropper" in str(raised.value.detail)
+    assert guardrail.scan_count == 1
+
+
+@pytest.mark.asyncio
+async def test_live_messages_stream_keeps_block_only_for_a_masking_guardrail():
+    """Without buffering the chunks are already out, so a rewrite cannot reach the client and nothing is replaced."""
+    guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call")
+    guardrail.streaming_buffer_until_moderated = False
+    guardrail.streaming_end_of_stream_only = 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)
+    raw = _decode(collected)
+    assert ORIGINAL_MARKER in raw
+    assert MASKED_MARKER not in raw
+
+
 def _detect_only_pipeline(guardrail: CustomGuardrail) -> tuple[tuple[str, GuardrailPipeline], ...]:
     step = PipelineStep(guardrail=guardrail.guardrail_name, on_pass="allow", on_fail="next")
     return (("live-policy", GuardrailPipeline(mode="post_call", steps=[step])),)