fix(guardrails): hold unscannable Responses windows and key terminal envelopes by output items

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -1175,13 +1175,11 @@ class OpenAIResponsesHandler(BaseTranslation):
         last_event_type: Final = stream_item_field(last_event, "type")
         if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
             return None
-        if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value:
+        if last_event_type in _TERMINAL_ENVELOPE_EVENT_TYPES:
             return self._completed_response_scan_key(stream_item_field(last_event, "response"))
-        stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
         return StreamingScanKey(
             texts=(self.get_streaming_string_so_far(responses_so_far),),
-            stream_ended=stream_ended,
-            tool_calls_in_flight=not stream_ended and self._has_streamed_tool_call_events(responses_so_far),
+            tool_calls_in_flight=self._has_streamed_tool_call_events(responses_so_far),
         )
 
     @staticmethod

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -1082,7 +1082,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                 scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
                 if scan_key is not None:
                     tool_calls_in_flight = scan_key.tool_calls_in_flight
-                hold_window = buffer_until_moderated and tool_calls_in_flight
+                hold_window = buffer_until_moderated and (scan_key is None or tool_calls_in_flight)
                 if _is_redundant_scan(scan_key, last_scan_key):
                     verbose_proxy_logger.debug(
                         "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
@@ -1159,7 +1159,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                     last_scan_key = scan_key
                 if hold_window:
                     verbose_proxy_logger.debug(
-                        "Holding %s buffered chunks for guardrail %s: streamed tool calls await the end-of-stream scan",
+                        "Holding %s buffered chunks for guardrail %s: this round could not scan the whole window",
                         len(withheld_items),
                         guardrail_to_apply.guardrail_name,
                     )

tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py

@@ -3212,6 +3212,24 @@ class TestOpenAIResponsesHandlerStreamingScanKey:
         done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}}
         assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None
 
+    @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"])
+    def test_non_completed_terminal_envelopes_key_their_output_items(self, terminal_type):
+        handler = OpenAIResponsesHandler()
+        arguments_delta = {
+            "type": "response.function_call_arguments.delta",
+            "sequence_number": 1,
+            "item_id": "fc_1",
+            "delta": '{"city":',
+        }
+        function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city":'}
+        terminal = {"type": terminal_type, "sequence_number": 2, "response": {"id": "resp_1", "output": [function_call]}}
+        mid_stream_key = handler.get_streaming_scan_key([arguments_delta])
+        ended_key = handler.get_streaming_scan_key([arguments_delta, terminal])
+        assert ended_key.stream_ended is True
+        assert ended_key.tool_calls_in_flight is False
+        assert len(ended_key.tool_calls) == 1
+        assert ended_key != mid_stream_key
+
     def test_streamed_tool_call_events_flag_tool_calls_in_flight_until_the_stream_ends(self):
         handler = OpenAIResponsesHandler()
         added = {

tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py

@@ -121,6 +121,32 @@ class _SecondScanBlockingGuardrail(_CountingPassingGuardrail):
         return inputs
 
 
+class _MarkerBlockingGuardrail(_CountingPassingGuardrail):
+    """Blocks as soon as the inspected input field (texts or tool_calls) carries the marker."""
+
+    def __init__(self, *args, marker: str, field: Literal["texts", "tool_calls"] = "texts", **kwargs):
+        super().__init__(*args, **kwargs)
+        self.marker = marker
+        self.field = field
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.scan_count += 1
+        if self.marker in json.dumps(inputs.get(self.field, [])):
+            raise ModifyResponseException(
+                message=BLOCK_MESSAGE,
+                model="gpt-4o",
+                request_data=request_data,
+                guardrail_name=self.guardrail_name,
+            )
+        return inputs
+
+
 def _sse_event(event_type: str, data: dict) -> bytes:
     return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
 
@@ -271,6 +297,76 @@ async def _run_windowed(
     return collected, yielded_count
 
 
+def _responses_message_stream_events(text_chunks: List[str]) -> List[dict]:
+    message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"}
+    content = [{"type": "output_text", "text": "".join(text_chunks), "annotations": []}]
+    return [
+        {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}},
+        *(
+            {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": text}
+            for text in text_chunks
+        ),
+        {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": content}},
+        {
+            "type": "response.completed",
+            "response": {"id": "resp_1", "model": "gpt-4o", "status": "completed", "output": [{**message, "content": content}]},
+        },
+    ]
+
+
+def _responses_truncated_function_call_events(text: str, argument_chunks: List[str]) -> List[dict]:
+    message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"}
+    content = [{"type": "output_text", "text": text, "annotations": []}]
+    function_call = {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "run_shell"}
+    return [
+        {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}},
+        {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": text},
+        {"type": "response.output_item.added", "output_index": 1, "item": {**function_call, "arguments": ""}},
+        *(
+            {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": arguments}
+            for arguments in argument_chunks
+        ),
+        {
+            "type": "response.incomplete",
+            "response": {
+                "id": "resp_1",
+                "model": "gpt-4o",
+                "status": "incomplete",
+                "output": [
+                    {**message, "content": content},
+                    {**function_call, "arguments": "".join(argument_chunks), "status": "incomplete"},
+                ],
+            },
+        },
+    ]
+
+
+async def _replay(events: List[dict]) -> AsyncGenerator[dict, None]:
+    for event in events:
+        yield event
+
+
+async def _run_windowed_responses(guardrail: CustomGuardrail, events: List[dict]) -> str:
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = True
+    guardrail.streaming_sampling_rate = 2
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses")
+    request_data = {
+        "input": "hi",
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": [guardrail.guardrail_name]},
+    }
+    collected: List[Any] = []
+    async for chunk in unified.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=user_api_key_dict,
+        response=_replay(events),
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    return json.dumps([chunk if isinstance(chunk, dict) else str(chunk) for chunk in collected])
+
+
 def _chat_text(chunks: List[Any]) -> str:
     return "".join(
         choice.delta.content or ""
@@ -365,6 +461,35 @@ async def test_windowed_buffer_holds_tool_call_windows_until_end_of_stream_scan(
     assert guardrail.tool_call_scan_indexes == [guardrail.scan_count]
 
 
+@pytest.mark.asyncio
+async def test_windowed_responses_output_item_done_round_keeps_text_window_withheld():
+    guardrail = _MarkerBlockingGuardrail(guardrail_name="windowed-responses", event_hook="post_call", marker=ORIGINAL_MARKER)
+    events = _responses_message_stream_events(["one ", f"{ORIGINAL_MARKER} "])
+
+    raw = await _run_windowed_responses(guardrail, events)
+
+    assert ORIGINAL_MARKER not in raw, f"unscanned window leaked: {raw!r}"
+    assert BLOCK_MESSAGE in raw
+    assert guardrail.scan_count >= 1
+
+
+@pytest.mark.asyncio
+async def test_windowed_responses_incomplete_stream_scans_tool_call_before_release():
+    guardrail = _MarkerBlockingGuardrail(
+        guardrail_name="windowed-responses-tools",
+        event_hook="post_call",
+        marker=TOOL_ARGUMENTS_MARKER,
+        field="tool_calls",
+    )
+    events = _responses_truncated_function_call_events("hi ", ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'])
+
+    raw = await _run_windowed_responses(guardrail, events)
+
+    assert '"hi "' in raw
+    assert TOOL_ARGUMENTS_MARKER not in raw, f"unscanned tool call leaked: {raw!r}"
+    assert BLOCK_MESSAGE in raw
+
+
 @pytest.mark.asyncio
 async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buffered():
     guardrail = _CountingPassingGuardrail(guardrail_name="windowed-eos", event_hook="post_call")