From b173eb3565beaa9bc1d8b403e832961eb4cd5c33 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 1 Sep 2026 15:18:52 -0700 Subject: [PATCH] fix: read the Responses delta event types off the event enum The hand-listed set left out response.mcp_call_arguments.delta, so a turn that streamed only MCP tool arguments and then reported an empty body still took the no-content shortcut and forwarded those chunks unscanned. Deriving the set from ResponsesAPIStreamEvents keeps it complete as the enum grows, and the str guard in the reader already covers any event whose delta is not text. --- .../model_armor/model_armor.py | 12 ++-- .../guardrail_hooks/test_model_armor.py | 65 +++++++++++++++++++ 2 files changed, 69 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 3870a46b2ad..ed3a1b12747 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -45,6 +45,7 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolCallChunk, ResponsesAPIResponse, + ResponsesAPIStreamEvents, ) from litellm.types.utils import ( CallTypes, @@ -62,15 +63,10 @@ GUARDRAIL_NAME: Final = "model_armor" # Only these carry the finished output; response.created carries an empty body _RESPONSES_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) -# Every event whose ``delta`` is model output already on its way to the client +# Every event whose ``delta`` is model output already on its way to the client. Read off the event +# enum rather than listed, so an event added there cannot quietly fall out of the scan _RESPONSES_DELTA_EVENT_TYPES: Final = frozenset( - { - "response.output_text.delta", - "response.refusal.delta", - "response.function_call_arguments.delta", - "response.custom_tool_call_input.delta", - "response.reasoning_summary_text.delta", - } + event.value for event in ResponsesAPIStreamEvents if event.value.endswith(".delta") ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index efcb2617174..526d3e6ff8e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -4615,3 +4615,68 @@ async def test_streaming_responses_terminal_event_without_a_body_still_scans_the 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 + + +@pytest.mark.asyncio +async def test_streaming_responses_mcp_argument_deltas_are_scanned_when_the_body_is_empty(): + """A turn that only streamed MCP tool arguments still handed the client a payload. + + The delta fallback is read off the event enum rather than listed by hand, so an argument event + that carries no `output_text` cannot fall out of the scan. + """ + from litellm.types.llms.openai import ( + MCPCallArgumentsDeltaEvent, + ResponseIncompleteEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + chunks = ( + MCPCallArgumentsDeltaEvent( + type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, + output_index=0, + item_id="mcp_1", + delta='{"note": "my card is 4111-1111-1111-1111"}', + sequence_number=0, + ), + ResponseIncompleteEvent( + type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-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, chunks) + + 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 ( + _RESPONSES_DELTA_EVENT_TYPES, + ) + from litellm.types.llms.openai import ResponsesAPIStreamEvents + + missing = { + event.value + for event in ResponsesAPIStreamEvents + if event.value.endswith(".delta") and event.value not in _RESPONSES_DELTA_EVENT_TYPES + } + assert not missing + assert "response.mcp_call_arguments.delta" in _RESPONSES_DELTA_EVENT_TYPES