fix(model_armor): scan each responses field once, not twice

Separating delta fields stopped the terminal body from matching the delta text, so a turn
with two visible fields sent Model Armor both copies. Only the delta fields the body does not
already carry are appended now.
This commit is contained in:
Yucheng Zhu 2026-09-02 17:04:55 -07:00
parent a0b0fec5f7
commit 669e70055e
2 changed files with 67 additions and 10 deletions

View file

@ -984,8 +984,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
return tuple(str(getattr(chunk, attr, None)) for attr in _RESPONSES_DELTA_FIELD_ATTRS)
@staticmethod
def _responses_delta_text(all_chunks: Sequence[object]) -> str:
"""Text a ``/v1/responses`` stream has already spelled out in its delta events.
def _responses_delta_field_texts(all_chunks: Sequence[object]) -> tuple[str, ...]:
"""Text each field of a ``/v1/responses`` turn has already spelled out in its delta events.
One field's deltas are joined as they streamed, since a finding can be split across them,
and separate fields stay apart, so a reasoning summary running into the visible answer
@ -997,7 +997,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
if getattr(chunk, "type", None) in _RESPONSES_DELTA_EVENT_TYPES
and isinstance(delta := getattr(chunk, "delta", None), str)
)
return "\n".join(
return tuple(
"".join(delta for field, delta in deltas if field == streamed_field)
for streamed_field in dict.fromkeys(field for field, _ in deltas)
)
@ -1011,17 +1011,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
"""Text to scan for a buffered stream, which is everything the client is about to receive.
A ``/v1/responses`` stream also spells out reasoning summaries and tool-call arguments in
delta events that its terminal body never repeats, so body and deltas are scanned together.
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_streaming_content(assembled_response)
if surface is not _StreamSurface.RESPONSES:
return content
delta_text: Final = self._responses_delta_text(all_chunks)
if delta_text in content:
return content
if content in delta_text:
return delta_text
return f"{content}\n{delta_text}"
unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content)
return "\n".join(part for part in (content, *unscanned) if part)
@staticmethod
def _apply_sanitized_content(assembled_response: ModelResponse, sanitized_content: str) -> None:

View file

@ -4855,6 +4855,66 @@ async def test_streaming_responses_one_fields_deltas_still_join_into_a_single_fi
assert "Streaming response blocked by Model Armor" in rendered
@pytest.mark.asyncio
async def test_streaming_responses_fields_the_body_repeats_are_not_scanned_a_second_time():
"""A turn whose visible fields all reach the terminal body is scanned once, not twice.
Two output_text fields stream as deltas and come back in the completed body, so scanning the
deltas on top of the body would send Model Armor two copies of everything the client sees.
"""
from litellm.types.llms.openai import (
OutputTextDeltaEvent,
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
paragraphs = ("the first thing to know", "a second and separate point")
text_deltas = tuple(
OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id=f"msg_{index}",
output_index=index,
content_index=0,
delta=paragraph,
)
for index, paragraph in enumerate(paragraphs)
)
completed = ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=ResponsesAPIResponse(
id="resp_1",
created_at=0,
model="gpt-5-mini",
object="response",
output=[
{
"type": "message",
"id": f"msg_{index}",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": paragraph, "annotations": []}],
}
for index, paragraph in enumerate(paragraphs)
],
parallel_tool_calls=False,
tool_choice="auto",
tools=[],
),
)
guardrail = _surface_guardrail()
post = _armor_post_mock(_MODEL_ARMOR_CLEAN)
with patch.object(guardrail.async_handler, "post", post):
delivered = await _drain_surface_hook(guardrail, (*text_deltas, completed))
post.assert_called_once()
scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"]
assert [scanned.count(paragraph) for paragraph in paragraphs] == [1, 1]
rendered = "".join(str(item) for item in delivered)
assert all(paragraph in rendered for paragraph in paragraphs)
def test_every_responses_delta_event_is_in_the_scanned_set():
"""Every ``.delta`` the Responses event enum defines is model output on its way to the client."""
from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import (