diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 03a754fa5df..3ee28ec42a7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1160,14 +1160,27 @@ 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. + 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: + return if not all_chunks: verbose_proxy_logger.warning( "Presidio apply_to_output: streaming response contained only " diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 38ea42285c1..a8c57e97b4a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2119,6 +2119,49 @@ 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", + ] + + # --------------------------------------------------------------------------- # Fix 4: apply_guardrail unmask path for input_type="response" # ---------------------------------------------------------------------------