diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index ee10dc91c1d..9cc5bfcf5b6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1446,10 +1446,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): elif isinstance(chunk, _SsePreface): yield chunk.raw elif isinstance(chunk, bytes): - if all_chunks or passthrough_due_to_unknown_stream_shape or is_sse_error_stream((chunk,)): - passthrough_due_to_unknown_stream_shape = ( - passthrough_due_to_unknown_stream_shape or not all_chunks - ) + if all_chunks or passthrough_due_to_unknown_stream_shape: yield chunk continue for masked_chunk in await self._mask_raw_sse_stream(chunk, stream, request_data): @@ -1492,14 +1489,18 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """The whole raw SSE stream masked as one response, or a raised refusal when it cannot be read. Raw frames are buffered to the end because PII can span frames, so no frame is forwarded - before the joined text was scanned. A stream carrying a frame the parser cannot read, whose - surface is unknown, or that its surface's assembler cannot rebuild, is withheld rather than - forwarded unmasked. + before the joined text was scanned, and the surface is decided on the whole stream rather + than on its first frame. A stream that is nothing but the refusal an earlier guardrail in + the chain emitted is forwarded as it arrived. A stream carrying a frame the parser cannot + read, whose surface is unknown, or that its surface's assembler cannot rebuild, is withheld + rather than forwarded unmasked. """ rest_chunks: Final = tuple([chunk async for chunk in rest]) chunks: Final = (first_chunk, *rest_chunks) if has_unreadable_sse_frames(chunks): raise self._withheld_stream_error("could not read every streamed frame") + if is_sse_error_stream(chunks): + return chunks if is_anthropic_sse_stream(chunks): return await self._mask_anthropic_sse_stream(chunks, request_data) if is_gemini_sse_stream(chunks): 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 2dc4ff53bce..a23eb0805bf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2834,6 +2834,59 @@ async def test_apply_to_output_streaming_gemini_error_frame_mid_stream_keeps_its assert error == {"error": {"code": 503, "message": "overloaded", "status": "UNAVAILABLE"}} +@pytest.mark.asyncio +async def test_apply_to_output_streaming_gemini_content_behind_a_leading_error_frame_is_still_masked(): + """ + The stream surface is decided on the whole buffered stream, not on its first frame, so an + error frame ahead of content frames does not switch masking off for the content behind it. + """ + frames = [ + b'data: {"error": {"code": 429, "message": "quota", "status": "RESOURCE_EXHAUSTED"}}\n\n', + _gemini_frame({"text": "The architect was John Smith."}), + ] + collected: list[object] = [] + + async def mock_stream(): + for frame in frames: + yield frame + + async with TestServer(_fake_presidio_app()) as server: + guardrail = _gemini_fake_masking_guardrail(server) + await _collect_masked_output(guardrail, mock_stream(), collected) + await guardrail._close_http_session() + + masked, error = _gemini_frames(collected) + assert masked["candidates"][0]["content"]["parts"] == [{"text": "The architect was ."}] + assert error == {"error": {"code": 429, "message": "quota", "status": "RESOURCE_EXHAUSTED"}} + assert not any(b"John Smith" in chunk for chunk in collected if isinstance(chunk, bytes)) + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_error_only_raw_stream_is_forwarded_as_it_arrived(): + """ + A post_call chain hands this hook the refusal frames an earlier guardrail emitted. They carry + nothing to scan, and withholding or rewriting them would hide the refusal the client is owed. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + frames = [ + b'data: {"error": {"message": "Violated guardrail policy", "type": "guardrail_error"}}\n\n', + b"data: [DONE]\n\n", + ] + collected: list[object] = [] + + async def mock_stream(): + for frame in frames: + yield frame + + await _collect_masked_output(guardrail, mock_stream(), collected) + + assert collected == frames + + @pytest.mark.asyncio async def test_apply_to_output_streaming_gemini_every_candidate_is_merged_by_index_and_masked(): """