fix(guardrails): fail closed on streamed non-chat routes the Airia rewrite cannot reach

Veria finding on #40784. incremental_diff, the mode that carries a rewrite into the stream,
exists only for OpenAI chat completions; UnifiedLLMGuardrails sends every other streamed route
(/v1/responses, /v1/messages) to block_only, which streams the original text live and drops
rewrites. So with post_call on, a streamed Responses or Anthropic call could receive text that
Airia then blocked or redacted.

Two changes, both confined to this guardrail:

- streaming_buffer_until_moderated is set, so on the block_only fallback the framework holds
  every chunk until the end-of-stream verdict instead of releasing it live. The flag is never
  read on the incremental_diff path, so streamed chat completions are unchanged.
- A REDACT verdict on a streamed response whose route cannot carry a rewrite is enforced as a
  block, since block_only would otherwise replay the original. The route check mirrors
  _resolve_transform_call_type so the two cannot disagree; an unresolvable route fails closed.
  Request-side rewrites and non-streamed calls are untouched.

The end-to-end test drives the real UnifiedLLMGuardrails iterator over an Anthropic
/v1/messages SSE stream and asserts not one original chunk is released for either verdict.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

litellm/proxy/guardrails/guardrail_hooks/airia/airia.py

@@ -1,5 +1,6 @@
 import os
 import uuid
+from collections.abc import Mapping
 from typing import (
     TYPE_CHECKING,
     Any,  # noqa: TID251  # the only type CustomGuardrail.__init__ accepts for its open-ended kwargs
@@ -15,6 +16,7 @@ from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
     log_guardrail_information,
 )
+from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
@@ -38,6 +40,56 @@ DEFAULT_BLOCKED_MESSAGE: Final = "Blocked by your organization's content policy.
 SUPPORTED_EVENT_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call)
 
 
+def _stream_cannot_carry_a_rewrite(request_data: Mapping[str, object] | None) -> bool:
+    """Whether this is a streamed response on a route the framework can only block, not rewrite.
+
+    ``incremental_diff``, the mode that writes a rewrite into the stream, exists only for OpenAI chat
+    completions: ``UnifiedLLMGuardrails._resolve_transform_call_type`` sends every other streamed route
+    to ``block_only``, which drops the rewrite and replays the buffered original. On those routes a
+    REDACT verdict can only be honoured as a block. The route check mirrors that resolver so the two
+    cannot disagree, and a route it cannot resolve fails closed the same way the resolver does.
+    """
+    if request_data is None or not _is_streamed(request_data):
+        return False
+    route: Final = _request_route(request_data)
+    call_types: Final = get_call_types_for_route(route) if route else None
+    if not call_types:
+        return True
+    # Deferred like the resolver's own import: the handler package imports the proxy.
+    from litellm.llms import load_guardrail_translation_mappings  # noqa: PLC0415
+    from litellm.llms.openai.chat.guardrail_translation.handler import (  # noqa: PLC0415
+        OpenAIChatCompletionsHandler,
+    )
+
+    handler: Final[object] = load_guardrail_translation_mappings().get(call_types[0])
+    return not (isinstance(handler, type) and issubclass(handler, OpenAIChatCompletionsHandler))
+
+
+def _is_streamed(request_data: Mapping[str, object]) -> bool:
+    if request_data.get("stream") is True:
+        return True
+    proxy_request: Final = _as_object(request_data.get("proxy_server_request"))
+    body: Final = _as_object(proxy_request.get("body")) if proxy_request is not None else None
+    return body is not None and body.get("stream") is True
+
+
+def _request_route(request_data: Mapping[str, object]) -> str | None:
+    """The route the proxy stamps as ``user_api_key_request_route``, in whichever metadata bucket the
+    route uses: ``metadata`` for chat completions, ``litellm_metadata`` for Responses and Anthropic."""
+    for bucket_name in ("metadata", "litellm_metadata"):
+        bucket = _as_object(request_data.get(bucket_name))
+        route = bucket.get("user_api_key_request_route") if bucket is not None else None
+        if isinstance(route, str) and route:
+            return route
+    return None
+
+
+def _as_object(value: object) -> dict[str, object] | None:
+    """`value` as a JSON object, or None. States the key type once so the checker does not see
+    `dict[Unknown, Unknown]` at every read."""
+    return value if isinstance(value, dict) else None
+
+
 class AiriaGuardrail(CustomGuardrail):
     """Evaluates prompts and responses against your Airia guardrail policy."""
 
@@ -61,6 +113,11 @@ class AiriaGuardrail(CustomGuardrail):
         self.api_key = api_key or os.getenv("AIRIA_API_KEY")
         self.streaming_transform_mode: Final[Literal["block_only", "incremental_diff"]] = "incremental_diff"
         self.streaming_end_of_stream_only: Final = True
+        # incremental_diff exists only for streamed chat completions. Every other streamed route falls
+        # back to block_only, where the framework releases the original text live unless told to hold
+        # it; holding it keeps the text behind the end-of-stream verdict, and apply_guardrail turns a
+        # REDACT that block_only cannot deliver into a block. Never read on the incremental_diff path.
+        self.streaming_buffer_until_moderated: Final = True
 
         if not self.api_base:
             raise ValueError("AiriaGuardrail requires api_base, or the AIRIA_GATEWAY_URL environment variable.")
@@ -130,6 +187,8 @@ class AiriaGuardrail(CustomGuardrail):
             )
 
         if action == ACTION_INTERVENED:
+            if input_type == "response" and _stream_cannot_carry_a_rewrite(request_data):
+                raise self._blocked()
             return self._rewritten(body, inputs)
 
         if action != ACTION_NONE:

tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py

@@ -1,3 +1,4 @@
+import json
 import uuid
 from collections.abc import AsyncGenerator
 from typing import Final, Literal
@@ -367,7 +368,8 @@ async def test_streaming_redaction_actually_reaches_the_client() -> None:
         user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"),
         response=source_stream(),
         request_data={  # mutable-ok: one-shot request_data the framework requires as a plain dict
-            "messages": [{"role": "user", "content": "share your contact"}]
+            "messages": [{"role": "user", "content": "share your contact"}],
+            "stream": True,
         },
         guardrail_to_apply=guardrail,
     ):
@@ -380,6 +382,166 @@ async def test_streaming_redaction_actually_reaches_the_client() -> None:
     guardrail.async_handler.post.assert_called_once()
 
 
+def _streamed_request(route: str | None) -> dict[str, object]:
+    """A proxy request_data as the streaming hook sees it: `stream` set, and the route the proxy
+    stamped into the metadata bucket (`litellm_metadata` for the non-chat routes under test)."""
+    request_data: dict[str, object] = {"stream": True}  # mutable-ok: built once per test
+    if route is not None:
+        request_data["litellm_metadata"] = {"user_api_key_request_route": route}
+    return request_data
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ["/v1/messages", "/v1/responses", "/responses", "/unknown/route", None])
+async def test_streamed_redaction_on_a_route_without_stream_rewrites_is_enforced_as_a_block(
+    route: str | None,
+) -> None:
+    """incremental_diff exists only for chat completions; every other streamed route runs block_only,
+    which drops rewrites and replays the original text. There a REDACT verdict must block instead."""
+    guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
+
+    with pytest.raises(GuardrailRaisedException) as excinfo:
+        await guardrail.apply_guardrail(
+            inputs=_inputs(texts=["secret"]), request_data=_streamed_request(route), input_type="response"
+        )
+
+    assert excinfo.value.blocked_content is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ["/chat/completions", "/v1/chat/completions"])
+async def test_streamed_redaction_on_chat_completions_is_still_delivered(route: str) -> None:
+    guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
+    request_data: dict[str, object] = {"stream": True, "metadata": {"user_api_key_request_route": route}}
+
+    result = await guardrail.apply_guardrail(
+        inputs=_inputs(texts=["secret"]), request_data=request_data, input_type="response"
+    )
+
+    assert result["texts"] == ["[Redacted]"]
+
+
+@pytest.mark.asyncio
+async def test_unstreamed_redaction_on_a_non_chat_route_is_delivered() -> None:
+    """Only the stream cannot carry a rewrite; a non-streamed Responses call is rewritten as usual."""
+    guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
+    request_data: dict[str, object] = {"litellm_metadata": {"user_api_key_request_route": "/v1/responses"}}
+
+    result = await guardrail.apply_guardrail(
+        inputs=_inputs(texts=["secret"]), request_data=request_data, input_type="response"
+    )
+
+    assert result["texts"] == ["[Redacted]"]
+
+
+@pytest.mark.asyncio
+async def test_streamed_request_side_redaction_is_unaffected() -> None:
+    """A pre_call rewrite lands in the prompt before the model is called; streaming is irrelevant there."""
+    guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
+
+    result = await guardrail.apply_guardrail(
+        inputs=_inputs(texts=["secret"]), request_data=_streamed_request("/v1/messages"), input_type="request"
+    )
+
+    assert result["texts"] == ["[Redacted]"]
+
+
+@pytest.mark.asyncio
+async def test_stream_flag_is_also_read_from_the_proxied_request_body() -> None:
+    guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]})
+    request_data: dict[str, object] = {
+        "proxy_server_request": {"body": {"stream": True}},
+        "litellm_metadata": {"user_api_key_request_route": "/v1/messages"},
+    }
+
+    with pytest.raises(GuardrailRaisedException) as excinfo:
+        await guardrail.apply_guardrail(
+            inputs=_inputs(texts=["secret"]), request_data=request_data, input_type="response"
+        )
+
+    assert excinfo.value.blocked_content is True
+
+
+def _anthropic_sse(event_type: str, data: dict[str, object]) -> bytes:
+    return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
+
+
+async def _anthropic_messages_stream() -> AsyncGenerator[bytes, None]:
+    """A complete Anthropic /v1/messages SSE stream, the shape the proxy relays on that route."""
+    yield _anthropic_sse(
+        "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 _anthropic_sse(
+        "content_block_start",
+        {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+    )
+    for text in ("Contact ", "me at ", "ada@example.com"):
+        yield _anthropic_sse(
+            "content_block_delta",
+            {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}},
+        )
+    yield _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0})
+    yield _anthropic_sse(
+        "message_delta",
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "usage": {"output_tokens": 5},
+        },
+    )
+    yield _anthropic_sse("message_stop", {"type": "message_stop"})
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "verdict",
+    [
+        {"action": "BLOCKED", "blocked_reason": "policy"},
+        {"action": "GUARDRAIL_INTERVENED", "texts": ["Contact me at [EmailAddress1]"]},
+    ],
+)
+async def test_non_chat_streamed_route_releases_nothing_before_the_verdict(verdict: dict[str, object]) -> None:
+    """Drives the real UnifiedLLMGuardrails iterator on /v1/messages, which incremental_diff does not
+    support, so the framework takes its block_only fallback. streaming_buffer_until_moderated holds
+    the text until the end-of-stream verdict; a block, and a redaction block_only could not deliver,
+    both refuse the response without one original chunk reaching the client."""
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
+        UnifiedLLMGuardrails,
+    )
+
+    guardrail: Final = _make_guardrail_with_response(verdict, event_hook="post_call", default_on=True)
+
+    released: list[object] = []  # mutable-ok: what the client received before the stream was refused
+    with pytest.raises(GuardrailRaisedException) as excinfo:
+        async for chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
+            user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/messages"),
+            response=_anthropic_messages_stream(),
+            request_data={  # mutable-ok: one-shot request_data the framework requires as a plain dict
+                "messages": [{"role": "user", "content": "share your contact"}],
+                "stream": True,
+            },
+            guardrail_to_apply=guardrail,
+        ):
+            released.append(chunk)  # noqa: PERF401  # a comprehension would lose what leaked before the raise
+
+    assert excinfo.value.blocked_content is True
+    assert "ada@example.com" not in "".join(c.decode() if isinstance(c, bytes) else str(c) for c in released)
+    guardrail.async_handler.post.assert_called_once()
+
+
 def test_config_model_is_exposed() -> None:
     assert AiriaGuardrail.get_config_model() is AiriaGuardrailConfigModel
     assert AiriaGuardrailConfigModel.ui_friendly_name() == "Airia Guardrail"