fix(model_armor): scan responses delta fields apart from each other

A Responses turn spells out its reasoning summary, its visible answer and its tool-call
arguments in separate delta events. Joining every delta into one string let a finding form
across the boundary between two fields that each carry nothing to find, so a safe stream
could be blocked. Group the deltas by the field they belong to, join a field's own deltas
as they streamed, and keep the fields apart.
This commit is contained in:
Yucheng Zhu 2026-09-02 16:33:06 -07:00
parent efa6c87a95
commit a0b0fec5f7
2 changed files with 133 additions and 3 deletions

View file

@ -69,6 +69,10 @@ _RESPONSES_DELTA_EVENT_TYPES: Final = frozenset(
event.value for event in ResponsesAPIStreamEvents if event.value.endswith(".delta")
)
# What makes two delta events part of the same field of the turn, rather than two fields that merely
# streamed next to each other
_RESPONSES_DELTA_FIELD_ATTRS: Final = ("type", "item_id", "output_index", "content_index", "summary_index")
class _StreamSurface(Enum):
"""Wire format of a buffered streaming response, which decides how it is read and how it is refused."""
@ -974,15 +978,29 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
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, ...]:
"""Which field of the turn a delta event belongs to."""
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."""
return "".join(
delta
"""Text a ``/v1/responses`` stream 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
cannot spell out a finding that neither of them carries.
"""
deltas: Final = tuple(
(ModelArmorGuardrail._responses_delta_field(chunk), delta)
for chunk in all_chunks
if getattr(chunk, "type", None) in _RESPONSES_DELTA_EVENT_TYPES
and isinstance(delta := getattr(chunk, "delta", None), str)
)
return "\n".join(
"".join(delta for field, delta in deltas if field == streamed_field)
for streamed_field in dict.fromkeys(field for field, _ in deltas)
)
def _streaming_content_to_scan(
self,

View file

@ -4743,6 +4743,118 @@ async def test_streaming_responses_reasoning_summary_deltas_are_scanned_alongsid
assert "Streaming response blocked by Model Armor" in rendered
@pytest.mark.asyncio
async def test_streaming_responses_deltas_of_separate_fields_do_not_form_a_finding_across_their_boundary():
"""Two fields of a turn are separate text, so what runs across their boundary is not model output.
A reasoning summary ending in half a card number and an answer opening with the other half
each carry nothing to find, and joining them without a break would invent one.
"""
from litellm.types.llms.openai import (
OutputTextDeltaEvent,
ReasoningSummaryTextDeltaEvent,
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
answer = "1111-1111 is not a full card"
summary_delta = ReasoningSummaryTextDeltaEvent(
type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA,
item_id="rs_1",
output_index=0,
delta="the prefix they gave me is 4111-1111-",
)
text_delta = OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id="msg_1",
output_index=1,
content_index=0,
delta=answer,
)
completed = ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=ResponsesAPIResponse(
id="resp_1",
created_at=0,
model="gpt-5-mini",
object="response",
output=[
{
"type": "message",
"id": "msg_1",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": answer, "annotations": []}],
}
],
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, (summary_delta, text_delta, completed))
post.assert_called_once()
scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"]
assert "4111-1111-" in scanned
assert answer in scanned
assert "4111-1111-1111-1111" not in scanned
rendered = "".join(str(item) for item in delivered)
assert "Streaming response blocked by Model Armor" not in rendered
@pytest.mark.asyncio
async def test_streaming_responses_one_fields_deltas_still_join_into_a_single_finding():
"""A card number split across two deltas of one field is still one card number to scan."""
from litellm.types.llms.openai import (
OutputTextDeltaEvent,
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
halves = ("my card is 4111-1111-", "1111-1111")
text_deltas = tuple(
OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id="msg_1",
output_index=0,
content_index=0,
delta=half,
)
for half in halves
)
completed = ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=ResponsesAPIResponse(
id="resp_1",
created_at=0,
model="gpt-5-mini",
object="response",
output=[],
parallel_tool_calls=False,
tool_choice="auto",
tools=[],
),
)
guardrail = _surface_guardrail()
post = _armor_post_mock(_MODEL_ARMOR_BLOCK)
with patch.object(guardrail.async_handler, "post", post):
delivered = await _drain_surface_hook(guardrail, (*text_deltas, completed))
post.assert_called_once()
assert "4111-1111-1111-1111" in post.call_args.kwargs["json"]["modelResponseData"]["text"]
rendered = "".join(str(item) for item in delivered)
assert "4111-1111-1111-1111" not in rendered
assert "Streaming response blocked by Model Armor" in rendered
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 (