Merge pull request #26878 from BerriAI/litellm_presidio-responses-stream-passthrough

fix(guardrails): preserve responses event streams in presidio output masking

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -1160,14 +1160,38 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         from litellm.types.utils import ModelResponse
 
         all_chunks: List[ModelResponseStream] = []
+        passthrough_due_to_unknown_stream_shape = False
         try:
             async for chunk in response:
                 if isinstance(chunk, ModelResponseStream):
-                    all_chunks.append(chunk)
+                    if passthrough_due_to_unknown_stream_shape:
+                        yield chunk
+                    else:
+                        all_chunks.append(chunk)
                 elif isinstance(chunk, bytes):
                     yield chunk  # type: ignore[misc]
                     continue
-
+                else:
+                    if all_chunks:
+                        # Flush buffered chunks and switch to transparent passthrough for this stream shape.
+                        # NOTE: these buffered chunks are emitted unmasked because this
+                        # stream mixed chunk types and cannot be safely reconstructed.
+                        verbose_proxy_logger.warning(
+                            "Presidio apply_to_output: mixed stream detected (ModelResponseStream + unknown event). "
+                            "Flushing %d buffered chunks without PII masking and switching to transparent passthrough.",
+                            len(all_chunks),
+                        )
+                        for buffered_chunk in all_chunks:
+                            yield buffered_chunk
+                        all_chunks = []
+                    passthrough_due_to_unknown_stream_shape = True
+                    yield chunk
+            if passthrough_due_to_unknown_stream_shape:
+                verbose_proxy_logger.warning(
+                    "Presidio apply_to_output: streaming response contained unknown event objects "
+                    "(e.g. /v1/responses events). Output PII masking was skipped for this response."
+                )
+                return
             if not all_chunks:
                 verbose_proxy_logger.warning(
                     "Presidio apply_to_output: streaming response contained only "

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py

@@ -2119,6 +2119,103 @@ async def test_streaming_unmask_path_bytes_passthrough():
     assert chunks[0] == byte_chunk
 
 
+@pytest.mark.asyncio
+async def test_apply_to_output_streaming_unknown_events_passthrough():
+    """
+    Regression test: /v1/responses-style event objects (neither bytes nor
+    ModelResponseStream) must be preserved in order and not dropped.
+    """
+    guardrail = _OPTIONAL_PresidioPIIMasking(
+        mock_testing=True,
+        apply_to_output=True,
+    )
+
+    class FakeResponsesEvent:
+        def __init__(self, event_type: str):
+            self.type = event_type
+
+    events = [
+        FakeResponsesEvent("response.created"),
+        FakeResponsesEvent("response.output_text.delta"),
+        FakeResponsesEvent("response.completed"),
+    ]
+
+    async def mock_stream():
+        for event in events:
+            yield event
+
+    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
+    received = []
+    async for chunk in guardrail.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=mock_user_api_key,
+        response=mock_stream(),
+        request_data={},
+    ):
+        received.append(chunk)
+
+    # Preserve exact objects and ordering so clients receive full event lifecycle.
+    assert received == events
+    assert [e.type for e in received] == [
+        "response.created",
+        "response.output_text.delta",
+        "response.completed",
+    ]
+
+
+@pytest.mark.asyncio
+async def test_apply_to_output_streaming_mixed_chunks_flushes_and_warns():
+    """
+    Regression test for mixed stream shape:
+    a buffered ModelResponseStream chunk followed by unknown responses-style
+    events should be preserved, and masking skip should be visible via warnings.
+    """
+    guardrail = _OPTIONAL_PresidioPIIMasking(
+        mock_testing=True,
+        apply_to_output=True,
+    )
+
+    class FakeResponsesEvent:
+        def __init__(self, event_type: str):
+            self.type = event_type
+
+    model_chunk = ModelResponseStream(
+        id="chatcmpl-mixed-1",
+        choices=[],
+        created=1,
+        model="gpt-4",
+        object="chat.completion.chunk",
+        system_fingerprint=None,
+    )
+    response_completed = FakeResponsesEvent("response.completed")
+
+    async def mock_stream():
+        yield model_chunk
+        yield response_completed
+
+    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
+    received = []
+    with patch(
+        "litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger"
+    ) as mock_logger:
+        async for chunk in guardrail.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=mock_user_api_key,
+            response=mock_stream(),
+            request_data={},
+        ):
+            received.append(chunk)
+
+        # Preserve original ordering across mixed stream types.
+        assert received == [model_chunk, response_completed]
+
+        # Two warnings are expected:
+        # 1) mixed stream detected + unmasked flush
+        # 2) passthrough mode skipped output masking
+        assert mock_logger.warning.call_count == 2
+        warning_messages = [call.args[0] for call in mock_logger.warning.call_args_list]
+        assert any("mixed stream detected" in msg for msg in warning_messages)
+        assert any("unknown event objects" in msg for msg in warning_messages)
+
+
 # ---------------------------------------------------------------------------
 # Fix 4: apply_guardrail unmask path for input_type="response"
 # ---------------------------------------------------------------------------