fix(guardrails): keep Model Armor's non-streamed Responses scan on message text

A masking Model Armor answered 400 on a Responses body with an email in forced tool call arguments,
where main answers 200. Scan only message text on that path, since one sanitized string cannot be
written back into tool call arguments, and spread that string over every part so a body with several
parts is masked instead of blocked
This commit is contained in:
mateo-berri 2026-10-30 16:39:22 -07:00
parent 1108ac2ab6
commit f671407132
4 changed files with 163 additions and 95 deletions

View file

@ -1439,21 +1439,40 @@ class OpenAIResponsesHandler(BaseTranslation):
rewrites_by_position=MappingProxyType(dict(zip(placeable_positions, chain((rewritten_text,), repeat(""))))),
)
async def deliver_text_rewrite(
def message_part_texts(self, response: "ResponsesAPIResponse") -> Mapping[tuple[int, int], str]:
"""Text of every message content part, keyed by ``(output_index, content_index)``."""
texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists
positions: Final[list[tuple[int, int]]] = [] # mutable-ok: the same extractor's position sink
for output_idx, output_item in enumerate(response.output or ()):
self._extract_output_text_and_images(
output_item=output_item,
output_idx=output_idx,
texts_to_check=texts,
images_to_check=[],
task_mappings=positions,
)
return MappingProxyType(
{position: text for position, text in zip(positions, texts) if position[1] != _COMPACTION_SUMMARY_POSITION}
)
async def spread_text_rewrite(
self,
response: "ResponsesAPIResponse",
position: tuple[int, int],
rewritten_text: str,
stream_events: Sequence[object] = (),
) -> None:
"""Write one rewritten text into the response at ``position`` and into any buffered events
that spell out the same text, so a client reading deltas sees what the body says."""
"""Replace all message text in ``response`` with one rewrite, for a guardrail that answers
one string for everything it scanned: the whole rewrite goes on the first part and every
later part is blanked, in the body and in any buffered events spelling out the same text."""
rewrites_by_position: Final = MappingProxyType(
dict(zip(self.message_part_texts(response), chain((rewritten_text,), repeat(""))))
)
await self._apply_guardrail_responses_to_output(
response=response, responses=[rewritten_text], task_mappings=[position]
)
self._sync_stream_events_with_rewrites(
stream_events=stream_events, rewrites_by_position=MappingProxyType({position: rewritten_text})
response=response,
responses=list(rewrites_by_position.values()),
task_mappings=list(rewrites_by_position),
)
self._sync_stream_events_with_rewrites(stream_events=stream_events, rewrites_by_position=rewrites_by_position)
@staticmethod
def _write_event_field(event: object, field: str, value: object) -> None:

View file

@ -1,7 +1,7 @@
import time
from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Protocol, runtime_checkable
import httpx
from fastapi import HTTPException
@ -10,6 +10,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
import json
@ -174,12 +175,6 @@ def _redact_scanned_content(payload: object, depth: int = 0) -> object:
return payload
class _ResponsesBodyScan(NamedTuple):
texts: tuple[str, ...]
positions: tuple[tuple[int, int], ...]
tool_calls: tuple[ChatCompletionToolCallChunk, ...]
class ModelArmorGuardrail(CustomGuardrail, VertexBase):
"""
Google Cloud Model Armor Guardrail integration for LiteLLM.
@ -548,7 +543,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
return False, None
if isinstance(assembled_response, ModelResponse) or (
isinstance(assembled_response, ResponsesAPIResponse)
and self._sole_text_position(assembled_response, content) is not None
and self._responses_message_text(assembled_response) == content
):
return False, sanitized_content
verbose_proxy_logger.warning(
@ -919,7 +914,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
return
# Extract content from response
content: Final = self._extract_scannable_content(response)
content: Final = (
self._responses_message_text(response)
if isinstance(response, ResponsesAPIResponse)
else self._extract_content_from_response(response)
)
if not content:
verbose_proxy_logger.debug("Model Armor: No text content to process in response, skipping guardrail")
return
@ -973,13 +972,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
if self.mask_response_content:
sanitized_content: Final = self._get_sanitized_content(armor_response)
if sanitized_content and sanitized_content != content:
if isinstance(response, ResponsesAPIResponse) and not await self._apply_sanitized_responses_content(
response, content, sanitized_content
):
raise HTTPException(
status_code=400,
detail=self._build_block_error_detail("Response blocked by Model Armor", armor_response),
)
if isinstance(response, ResponsesAPIResponse):
await self._responses_handler().spread_text_rewrite(response, sanitized_content)
if isinstance(response, litellm.ModelResponse):
self._apply_sanitized_content(response, sanitized_content)
@ -1039,26 +1033,17 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
)
@staticmethod
def _scan_responses_api_response(response: ResponsesAPIResponse) -> _ResponsesBodyScan:
"""Texts of a Responses API response, where each one sits, and its tool calls."""
def _responses_handler() -> "OpenAIResponsesHandler":
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists
positions: Final[list[tuple[int, int]]] = [] # mutable-ok: the same extractor's position sink
tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] # mutable-ok: the same extractor's tool-call sink
handler: Final = OpenAIResponsesHandler()
for output_idx, output_item in enumerate(response.output or ()):
handler._extract_output_text_and_images( # pyright: ignore[reportPrivateUsage] # the shared Responses output extractor; forking it would duplicate per-item parsing
output_item=output_item,
output_idx=output_idx,
texts_to_check=texts,
images_to_check=[],
task_mappings=positions,
tool_calls_to_check=tool_calls,
)
return _ResponsesBodyScan(texts=tuple(texts), positions=tuple(positions), tool_calls=tuple(tool_calls))
return OpenAIResponsesHandler()
@staticmethod
def _responses_message_text(response: ResponsesAPIResponse) -> str:
"""Message text of a Responses API response, every part joined as Model Armor scans it."""
return "".join(ModelArmorGuardrail._responses_handler().message_part_texts(response).values())
@staticmethod
def _responses_api_response_text(response: ResponsesAPIResponse) -> str:
@ -1067,47 +1052,25 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
Tool calls are folded in because ``get_content_from_model_response`` folds them into what
the chat surface scans, and a Responses turn can carry its whole payload in them.
"""
scan: Final = ModelArmorGuardrail._scan_responses_api_response(response)
return "".join((*scan.texts, *(json.dumps(tool_call) for tool_call in scan.tool_calls)))
texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists
tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] # mutable-ok: the same extractor's tool-call sink
handler: Final = ModelArmorGuardrail._responses_handler()
for output_idx, output_item in enumerate(response.output or ()):
handler._extract_output_text_and_images( # pyright: ignore[reportPrivateUsage] # the shared Responses output extractor; forking it would duplicate per-item parsing
output_item=output_item,
output_idx=output_idx,
texts_to_check=texts,
images_to_check=[],
task_mappings=[],
tool_calls_to_check=tool_calls,
)
return "".join((*texts, *(json.dumps(tool_call) for tool_call in tool_calls)))
@staticmethod
def _sole_text_position(response: ResponsesAPIResponse, content: str) -> tuple[int, int] | None:
"""Where one sanitized string can land: the part ``content`` came from, when that is one part.
A scan that also joined a second part, tool-call arguments or reasoning deltas comes back as
one string with no boundaries, so it has no single place to go.
"""
scan: Final = ModelArmorGuardrail._scan_responses_api_response(response)
return scan.positions[0] if scan.texts == (content,) else None
async def _apply_sanitized_responses_content(
self,
response: ResponsesAPIResponse,
content: str,
sanitized_content: str,
stream_events: Sequence[object] = (),
) -> bool:
"""Write the sanitized text into the response and its buffered events, if it has one place to go."""
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
position: Final = self._sole_text_position(response, content)
if position is None:
return False
await OpenAIResponsesHandler().deliver_text_rewrite(
response=response,
position=position,
rewritten_text=sanitized_content,
stream_events=stream_events,
)
return True
def _extract_scannable_content(self, response: object) -> str:
"""Text to scan from a response, for every endpoint shape this hook serves."""
if isinstance(response, ResponsesAPIResponse):
return self._responses_api_response_text(response)
return self._extract_content_from_response(response)
def _extract_streaming_content(self, assembled_response: object) -> str:
"""Text to scan from an assembled stream, for every endpoint shape this hook serves."""
if isinstance(assembled_response, ResponsesAPIResponse):
return self._responses_api_response_text(assembled_response)
return self._extract_content_from_response(assembled_response)
@staticmethod
def _responses_delta_field(chunk: object) -> tuple[str, ...]:
@ -1145,7 +1108,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
delta events that its terminal body never repeats, so every delta field the body does not
already carry is scanned after it.
"""
content: Final = self._extract_scannable_content(assembled_response)
content: Final = self._extract_streaming_content(assembled_response)
if surface is not _StreamSurface.RESPONSES:
return content
unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content)
@ -1316,8 +1279,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
)
if sanitized_content is not None and isinstance(assembled_response, ResponsesAPIResponse):
await self._apply_sanitized_responses_content(
assembled_response, content, sanitized_content, stream_events=all_chunks
await self._responses_handler().spread_text_rewrite(
assembled_response, sanitized_content, stream_events=all_chunks
)
if sanitized_content is not None and isinstance(assembled_response, ModelResponse):

View file

@ -4966,6 +4966,34 @@ class TestCompactionOutputSummaryReachesOutputGuardrails:
assert self._summary(result.output[0]) == "The user lives at 1 [MASKED]."
@pytest.mark.asyncio
async def test_spreading_one_rewrite_over_message_text_leaves_the_summary_alone(self) -> None:
response: Final = ResponsesAPIResponse(
id="resp_1",
created_at=1,
model="claude-sonnet-5-5",
object="response",
status="completed",
output=[
self._compaction_item(),
{
"type": "message",
"id": "msg_1",
"status": "completed",
"role": "assistant",
"content": [
{"type": "output_text", "text": "first", "annotations": []},
{"type": "refusal", "refusal": "second"},
],
},
],
)
await OpenAIResponsesHandler().spread_text_rewrite(response, "rewritten")
assert self._summary(response.output[0]) == self.SUMMARY
assert [response.output[1].content[0].text, response.output[1].content[1].refusal] == ["rewritten", ""]
@pytest.mark.asyncio
async def test_ended_stream_summary_is_masked_in_the_done_item_and_the_envelope(self) -> None:
events: Final = [

View file

@ -3892,9 +3892,17 @@ async def _anthropic_sse_stream():
_SECOND_TEXT_PART: Final = {"type": "output_text", "text": " and more"}
_TOOL_CALL_ITEM: Final = {
"type": "function_call",
"id": "fc_1",
"call_id": "call_1",
"name": "save_card",
"arguments": '{"card": "4111-1111-1111-1111"}',
"status": "completed",
}
def _responses_api_events(*extra_parts: dict):
def _responses_api_events(extra_parts: tuple[dict, ...] = (), extra_items: tuple[dict, ...] = ()):
from litellm.types.llms.openai import (
OutputTextDeltaEvent,
ResponseCompletedEvent,
@ -3914,7 +3922,8 @@ def _responses_api_events(*extra_parts: dict):
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "my card is 4111-1111-1111-1111"}, *extra_parts],
}
},
*extra_items,
],
parallel_tool_calls=False,
tool_choice="auto",
@ -4050,9 +4059,9 @@ async def test_streaming_masking_re_emits_anthropic_sse_with_sanitized_text():
@pytest.mark.asyncio
async def test_streaming_masking_blocks_a_responses_stream_with_no_single_part_to_rewrite():
"""One sanitized string cannot be split back over several parts, so releasing the stream would
ship the content the guardrail just rewrote. It is blocked instead."""
async def test_streaming_masking_blocks_a_responses_stream_whose_scan_includes_a_tool_call():
"""Tool-call arguments are scanned on a stream but one sanitized string has no way back into
them, so releasing the stream would ship the content the guardrail just rewrote."""
from litellm.types.llms.openai import ErrorEvent
guardrail = _surface_guardrail(mask_response_content=True)
@ -4060,7 +4069,7 @@ async def test_streaming_masking_blocks_a_responses_stream_with_no_single_part_t
with patch.object(
guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_SANITIZED)
):
delivered = await _drain_surface_hook(guardrail, _responses_api_events(_SECOND_TEXT_PART))
delivered = await _drain_surface_hook(guardrail, _responses_api_events(extra_items=(_TOOL_CALL_ITEM,)))
assert len(delivered) == 1
assert isinstance(delivered[0], ErrorEvent)
@ -4538,9 +4547,9 @@ async def test_streaming_deidentify_match_without_a_rewrite_fails_closed(chunks,
@pytest.mark.asyncio
async def test_streaming_status_records_a_surface_that_cannot_carry_the_rewrite_as_blocked():
"""A multi-part Responses body has no single part to rewrite, so a de-identify match ends as a
refusal. The status stamped on metadata feeds the spend log and has to say so rather than
reporting the success the block check alone would have implied."""
"""A Responses stream that scanned a tool call has nowhere to put the rewrite, so a de-identify
match ends as a refusal. The status stamped on metadata feeds the spend log and has to say so
rather than reporting the success the block check alone would have implied."""
guardrail = _surface_guardrail(mask_response_content=True)
request_data = {
"model": "gpt-4o-mini",
@ -4550,7 +4559,7 @@ async def test_streaming_status_records_a_surface_that_cannot_carry_the_rewrite_
with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)):
delivered = await _drain_surface_hook(
guardrail, _responses_api_events(_SECOND_TEXT_PART), request_data=request_data
guardrail, _responses_api_events(extra_items=(_TOOL_CALL_ITEM,)), request_data=request_data
)
body = _delivered_bytes(delivered)
@ -5477,7 +5486,7 @@ async def test_post_call_masks_plain_text_in_a_responses_body():
@pytest.mark.asyncio
async def test_post_call_blocks_a_responses_body_whose_parts_one_redaction_cannot_split_across():
async def test_post_call_masks_every_part_of_a_responses_body_through_the_first():
guardrail = _model_armor_post_call_guardrail(mask_response_content=True)
post = AsyncMock(return_value=_armor_http_response(_armor_redaction("card [REDACTED] and more")))
response = _responses_body(
@ -5485,11 +5494,38 @@ async def test_post_call_blocks_a_responses_body_whose_parts_one_redaction_canno
{"type": "output_text", "text": " and more", "annotations": []},
)
with patch.object(guardrail.async_handler, "post", post), pytest.raises(HTTPException):
with patch.object(guardrail.async_handler, "post", post):
await guardrail.async_post_call_success_hook(
data=dict(_POST_CALL_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=response
)
assert [part.text for part in response.output[0].content] == ["card [REDACTED] and more", ""]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("mask_response_content", "flagged"),
[(True, _armor_redaction("saving [REDACTED]")), (False, _armor_block())],
)
async def test_post_call_answers_a_tool_call_beside_responses_text_unchanged(mask_response_content, flagged):
from litellm.types.llms.openai import ResponsesAPIResponse
async def flag_the_card(**kwargs):
verdict = flagged if "4111" in str(kwargs["json"]) else {"filterMatchState": "NO_MATCH_FOUND"}
return _armor_http_response(verdict)
guardrail = _model_armor_post_call_guardrail(mask_response_content=mask_response_content)
body = _responses_body({"type": "output_text", "text": "saving it", "annotations": []}).model_dump()
response = ResponsesAPIResponse.model_validate({**body, "output": [*body["output"], _TOOL_CALL_ITEM]})
with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=flag_the_card)):
await guardrail.async_post_call_success_hook(
data=dict(_POST_CALL_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=response
)
assert response.output[0].content[0].text == "saving it"
assert response.output[1].arguments == _TOOL_CALL_ITEM["arguments"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
@ -5540,3 +5576,25 @@ async def test_streaming_masks_every_event_that_spells_out_a_responses_part(part
assert "4111" not in str(released)
assert "".join(event.delta for event in released[:2]) == "I cannot help card [REDACTED]"
assert getattr(released[2], field) == "I cannot help card [REDACTED]"
@pytest.mark.asyncio
async def test_streaming_masks_a_responses_body_with_several_parts():
from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents
guardrail = _surface_guardrail(mask_response_content=True)
first_delta, terminal = _responses_api_events(extra_parts=(_SECOND_TEXT_PART,))
second_delta = OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id="msg_1",
output_index=0,
content_index=1,
delta=_SECOND_TEXT_PART["text"],
)
with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)):
delivered = await _drain_surface_hook(guardrail, (first_delta, second_delta, terminal))
assert b"4111-1111-1111-1111" not in _delivered_bytes(delivered)
assert [delivered[0].delta, delivered[1].delta] == ["my card is [REDACTED]", ""]
assert [part["text"] for part in delivered[-1].response.output[0]["content"]] == ["my card is [REDACTED]", ""]