fix(policy_engine): withhold streams when a pipeline guardrail rewrites output at runtime

litellm/proxy/policy_engine/pipeline_executor.py

@@ -6,8 +6,11 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding.
 """
 
 import time
+from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, Literal
 
+from pydantic import BaseModel
+
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.integrations.custom_guardrail import (
@@ -24,8 +27,10 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
     PipelineStep,
     PipelineStepResult,
 )
+from litellm.types.utils import GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
     from litellm.llms.base_llm.guardrail_translation.base_translation import (
         BaseTranslation,
     )
@@ -36,6 +41,90 @@ except ImportError:
     HTTPException = None
 
 
+class UndeliverableStreamRewrite(Exception):
+    def __init__(self, guardrail_name: str) -> None:
+        super().__init__(
+            f"Guardrail '{guardrail_name}' rewrote the streamed response, which streaming pipelines cannot deliver"
+        )
+        self.guardrail_name: Final = guardrail_name
+
+
+def _tool_call_shape(tool_call: object) -> tuple[object, object]:
+    plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
+    function: Final = plain.get("function") if isinstance(plain, Mapping) else None
+    if not isinstance(function, Mapping):
+        return (None, None)
+    return (function.get("name"), function.get("arguments"))
+
+
+def _rewrote_texts(sent: Sequence[str] | None, returned: Sequence[str] | None) -> bool:
+    return sent is not None and returned is not None and list(returned) != list(sent)
+
+
+def _rewrote_tool_calls(sent: Sequence[object] | None, returned: Sequence[object] | None) -> bool:
+    if sent is None or returned is None:
+        return False
+    return [_tool_call_shape(tool_call) for tool_call in returned] != [
+        _tool_call_shape(tool_call) for tool_call in sent
+    ]
+
+
+class _StreamRewriteObserver(CustomGuardrail):
+    """Stand-in handed to the endpoint translation in place of a streaming pipeline step's
+    guardrail. Translations cannot rewrite every buffered chunk consistently, so the gate
+    withholds the stream whenever the guardrail returned different output than it was given,
+    which for guardrails like Bedrock's ANONYMIZED action is only known at runtime."""
+
+    def __init__(self, inner: CustomGuardrail) -> None:
+        super().__init__(guardrail_name=inner.guardrail_name)
+        self.inner: Final = inner
+        self.rewrote = False
+
+    def structured_messages_cover_full_request(self) -> bool:
+        return self.inner.structured_messages_cover_full_request()
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,  # mutable-ok: matches CustomGuardrail.apply_guardrail
+        input_type: Literal["request", "response"],
+        logging_obj: "LiteLLMLoggingObj | None" = None,
+    ) -> GenericGuardrailAPIInputs:
+        outputs: Final = await self.inner.apply_guardrail(
+            inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj
+        )
+        self.rewrote = (
+            self.rewrote
+            or _rewrote_texts(inputs.get("texts"), outputs.get("texts"))
+            or _rewrote_tool_calls(inputs.get("tool_calls"), outputs.get("tool_calls"))
+        )
+        return outputs
+
+
+def _prepare_hook_input(
+    step: PipelineStep,
+    callback: CustomLogger,
+    data: dict,  # mutable-ok: same request-payload shape the hooks mutate
+    raw_request_snapshot: dict | None,  # mutable-ok: same request-payload shape as data
+) -> tuple[dict, bool]:  # mutable-ok: returns that same request-payload dict
+    """Inject the step's guardrail name into metadata so should_run_guardrail() allows it,
+    and pick the payload the step scans: a scan_raw_request step evaluates the pristine
+    pre-pipeline snapshot instead of `data` (which earlier pass_data steps in this same
+    pipeline may have already rewritten), same reason the normal sequential/parallel
+    guardrail loops do this."""
+    if "metadata" not in data:
+        data["metadata"] = {}
+    data["metadata"]["guardrails"] = [step.guardrail]
+
+    scans_raw_request: Final = getattr(callback, "scan_raw_request", False)
+    hook_input: Final[dict] = (  # mutable-ok: same request-payload shape as data
+        independent_snapshot(raw_request_snapshot) if scans_raw_request and raw_request_snapshot is not None else data
+    )
+    if hook_input is not data:
+        hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
+    return hook_input, scans_raw_request
+
+
 class PipelineExecutor:
     """Executes guardrail pipelines with ordered, conditional step logic."""
 
@@ -195,23 +284,7 @@ class PipelineExecutor:
             return ("error", None, f"Guardrail '{step.guardrail}' not found", None)
 
         try:
-            # Inject guardrail name into metadata so should_run_guardrail() allows it
-            if "metadata" not in data:
-                data["metadata"] = {}
-            data["metadata"]["guardrails"] = [step.guardrail]
-
-            # A scan_raw_request step evaluates the pristine pre-pipeline
-            # snapshot instead of `data` (which earlier pass_data steps in
-            # this same pipeline may have already rewritten), same reason
-            # the normal sequential/parallel guardrail loops do this.
-            scans_raw_request: Final = getattr(callback, "scan_raw_request", False)
-            hook_input: Final[dict] = (  # mutable-ok: same request-payload shape as data
-                independent_snapshot(raw_request_snapshot)
-                if scans_raw_request and raw_request_snapshot is not None
-                else data
-            )
-            if hook_input is not data:
-                hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
+            hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot)
 
             # Use unified_guardrail path if callback implements apply_guardrail
             target: CustomLogger = callback
@@ -239,13 +312,16 @@ class PipelineExecutor:
                         f"Guardrail '{step.guardrail}' does not support streaming pipeline execution",
                         None,
                     )
+                observer: Final = _StreamRewriteObserver(callback)
                 await endpoint_translation.process_output_streaming_response(
                     responses_so_far=streaming_chunks,
-                    guardrail_to_apply=callback,
+                    guardrail_to_apply=observer,
                     litellm_logging_obj=data.get("litellm_logging_obj"),
                     user_api_key_dict=user_api_key_dict,
                     request_data=hook_input,
                 )
+                if observer.rewrote:
+                    raise UndeliverableStreamRewrite(step.guardrail)
                 response = None
             elif mode == "post_call":
                 response = await target.async_post_call_success_hook(
@@ -269,6 +345,8 @@ class PipelineExecutor:
                 return ("pass", {"response": response}, None, None)
             return ("pass", response if isinstance(response, dict) else None, None, None)
 
+        except UndeliverableStreamRewrite:
+            raise
         except Exception as e:
             if CustomGuardrail._is_guardrail_intervention(e):
                 error_msg: Final = _extract_error_message(e)

litellm/proxy/utils.py

@@ -155,7 +155,7 @@ from litellm.proxy.hooks.sensitive_data_routing import (
 )
 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
 from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
-from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
+from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite
 from litellm.repositories.budget_repository import BudgetRepository
 from litellm.repositories.config_repository import ConfigRepository
 from litellm.repositories.table_repositories import (
@@ -511,6 +511,23 @@ class _PipelineErrorDetail(TypedDict):
     error: ReadOnly[_PipelineErrorBody]
 
 
+def _undeliverable_stream_rewrite_error(policy_name: str, guardrail_name: str) -> HTTPException:
+    detail: Final[_PipelineErrorDetail] = {
+        "error": {
+            "message": (
+                f"Streaming response withheld by policy pipeline '{policy_name}' because guardrail "
+                f"'{guardrail_name}' rewrote the streamed output, and streaming pipelines cannot deliver "
+                "rewrites. Retry with stream=false, or drop it from the pipeline steps so guardrails.add "
+                "applies it to streamed output."
+            ),
+            "type": "guardrail_pipeline_error",
+            "policies": (policy_name,),
+            "guardrails": (guardrail_name,),
+        }
+    }
+    return HTTPException(status_code=400, detail=detail)
+
+
 def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> None:
     """
     Reject up front the requests whose post_call pipelines could never run.
@@ -3468,11 +3485,12 @@ class ProxyLogging:
         pipeline allows it), then runs each pipeline's steps against the
         assembled output through the endpoint guardrail translation, the same
         machinery flat post_call guardrails use at end of stream. An allow
-        releases the buffered chunks as that machinery left them (the
-        Responses and A2A translations write guardrail output back into the
-        final chunk, exactly as they do for flat guardrails); a block or
-        modify_response terminates with the translation's block chunks or the
-        raised error.
+        releases the buffered chunks verbatim; a step whose guardrail rewrote
+        the output withholds the stream with a 400 instead, since no
+        translation rewrites every buffered chunk consistently and some
+        rewrites (Bedrock's ANONYMIZED action, for one) are only decided at
+        runtime; a block or modify_response terminates with the translation's
+        block chunks or the raised error.
         """
         buffered: Final[list[object]] = []  # mutable-ok: accumulates the stream before the pipeline verdict
         async for item in response:
@@ -3495,16 +3513,26 @@ class ProxyLogging:
         call_type, endpoint_translation = resolved
 
         for policy_name, pipeline in pipelines:
-            result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
-                steps=pipeline.steps,
-                mode="post_call",
-                data=request_data,
-                user_api_key_dict=user_api_key_dict,
-                call_type=call_type,
-                policy_name=policy_name,
-                streaming_chunks=buffered,
-                endpoint_translation=endpoint_translation,
-            )
+            try:
+                result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
+                    steps=pipeline.steps,
+                    mode="post_call",
+                    data=request_data,
+                    user_api_key_dict=user_api_key_dict,
+                    call_type=call_type,
+                    policy_name=policy_name,
+                    streaming_chunks=buffered,
+                    endpoint_translation=endpoint_translation,
+                )
+            except UndeliverableStreamRewrite as rewrite:
+                async for error_chunk in unified_guardrail.emit_streaming_http_error(
+                    _undeliverable_stream_rewrite_error(policy_name, rewrite.guardrail_name),
+                    call_type,
+                    buffered,
+                    request_data,
+                ):
+                    yield error_chunk
+                return
             try:
                 ProxyLogging._handle_pipeline_result(
                     result, data=request_data, policy_name=policy_name, original_response=buffered

tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py

@@ -13,7 +13,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
     CustomCodeGuardrail,
 )
-from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
+from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite
 from litellm.types.proxy.policy_engine.pipeline_types import (
     GuardrailPipeline,
     PipelineStep,
@@ -811,3 +811,64 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch):
     assert outcome == "pass"
     assert guardrail.native_pre_call_ran is True
     assert "guardrail_to_apply" not in data
+
+
+class _TextReturningGuardrail(CustomGuardrail):
+    def __init__(self, returned_texts):
+        super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
+        self.returned_texts = returned_texts
+
+    async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+        return {**inputs, "texts": self.returned_texts}
+
+
+class _TextTranslation:
+    def __init__(self):
+        self.seen_guardrail_names = []
+
+    async def process_output_streaming_response(
+        self, responses_so_far, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
+    ):
+        self.seen_guardrail_names.append(guardrail_to_apply.guardrail_name)
+        await guardrail_to_apply.apply_guardrail(
+            inputs={"texts": ["hello world"]},
+            request_data=request_data or {},
+            input_type="response",
+            logging_obj=litellm_logging_obj,
+        )
+        return responses_so_far
+
+
+async def _run_streaming_step(returned_texts, translation):
+    return await PipelineExecutor.execute_steps(
+        steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail="next", on_error="next")],
+        mode="post_call",
+        data={"model": "m"},
+        user_api_key_dict=MagicMock(),
+        call_type="completion",
+        policy_name="p",
+        streaming_chunks=[object()],
+        endpoint_translation=translation,
+    )
+
+
+@pytest.mark.asyncio
+async def test_streaming_step_rewrite_escapes_execute_steps_regardless_of_step_actions(monkeypatch):
+    monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])
+    translation = _TextTranslation()
+
+    with pytest.raises(UndeliverableStreamRewrite) as info:
+        await _run_streaming_step(["hello [MASKED]"], translation)
+
+    assert info.value.guardrail_name == "masker"
+    assert translation.seen_guardrail_names == ["masker"]
+
+
+@pytest.mark.asyncio
+async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeypatch):
+    monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(("hello world",))])
+
+    result = await _run_streaming_step(("hello world",), _TextTranslation())
+
+    assert result.terminal_action == "allow"
+    assert [step.outcome for step in result.step_results] == ["pass"]

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

@@ -11,7 +11,7 @@ from __future__ import annotations
 
 import asyncio
 import json
-from typing import Any, Dict, List
+from typing import Any, Callable, Dict, List
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -878,10 +878,12 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p
 # ---------------------------------------------------------------------------
 
 
-def _post_call_pipeline_data(guardrail: str = "gr-post", **extra: Any) -> Dict[str, Any]:
+def _post_call_pipeline_data(
+    guardrail: str = "gr-post", step: PipelineStep | None = None, **extra: Any
+) -> Dict[str, Any]:
     pipeline = GuardrailPipeline(
         mode="post_call",
-        steps=[PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")],
+        steps=[step or PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")],
     )
     return {
         "model": "m",
@@ -1587,6 +1589,101 @@ async def test_streaming_iterator_hook_pipeline_block_withholds_all_chunks(
     assert "output blocked" in str(info.value.detail)
 
 
+def _rewriting_stream_guardrail(transform: Callable[[Dict[str, Any]], Dict[str, Any]]) -> CustomGuardrail:
+    class RewritingStreamGuardrail(CustomGuardrail):
+        async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+            return {**inputs, **transform(inputs)}
+
+    return RewritingStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)
+
+
+def _tool_call_stream_chunks() -> List[Any]:
+    tool_call = {
+        "index": 0,
+        "id": "call_1",
+        "type": "function",
+        "function": {"name": "lookup", "arguments": '{"ssn": "123"}'},
+    }
+    return [
+        litellm.ModelResponseStream(
+            choices=[{"index": 0, "delta": {"tool_calls": [tool_call]}, "finish_reason": None}]
+        ),
+        litellm.ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]),
+    ]
+
+
+def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]:
+    return [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}}]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")])
+@pytest.mark.parametrize(
+    "make_chunks, transform",
+    [
+        (_stream_chunks, lambda inputs: {"texts": ["hello [MASKED]"]}),
+        (_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')}),
+    ],
+    ids=["texts", "tool_calls"],
+)
+async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite(
+    proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform, on_fail, on_error
+):
+    monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    step = PipelineStep(guardrail="gr-post", on_pass="allow", on_fail=on_fail, on_error=on_error)
+    data = _post_call_pipeline_data(step=step, stream=True)
+    delivered: List[Any] = []
+
+    async def _drain() -> None:
+        async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
+            response=_async_chunk_iter(make_chunks()),
+            request_data=data,
+        ):
+            delivered.append(item)
+
+    with pytest.raises(HTTPException) as info:
+        await _drain()
+
+    error = info.value.detail["error"]
+    assert delivered == []
+    assert info.value.status_code == 400
+    assert error["type"] == "guardrail_pipeline_error"
+    assert error["policies"] == ("response-governance",)
+    assert error["guardrails"] == ("gr-post",)
+    assert "stream=false" in error["message"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "make_chunks, transform",
+    [
+        (_stream_chunks, lambda inputs: {"texts": tuple(inputs["texts"])}),
+        (_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "123"}')}),
+    ],
+    ids=["texts_as_tuple", "tool_calls_as_dicts"],
+)
+async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_another_shape(
+    proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform
+):
+    monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data(stream=True)
+    chunks = make_chunks()
+
+    delivered = [
+        item
+        async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
+            response=_async_chunk_iter(chunks),
+            request_data=data,
+        )
+    ]
+
+    assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks]
+
+
 @pytest.mark.asyncio
 async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_shape(
     proxy_logging, make_user_api_key_auth, monkeypatch