From 74daaf4e6b4e70335d6cb5223dd758ba8f563873 Mon Sep 17 00:00:00 2001 From: Airia Open Source Date: Wed, 7 Oct 2026 00:50:49 +0000 Subject: [PATCH] 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 --- .../guardrails/guardrail_hooks/airia/airia.py | 59 +++++++ .../guardrails/guardrail_hooks/test_airia.py | 164 +++++++++++++++++- 2 files changed, 222 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py index 8f5423b65ca..6a5eee02e56 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py +++ b/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: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py index 37170b6e826..f5c547253e8 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py +++ b/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"