From 9e683e86f312f61eb0dc4ed3a565743948f7dd31 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 26 Sep 2026 15:54:45 -0700 Subject: [PATCH] test(guardrails): cover masking of every Gemini candidate merged by index --- .../guardrail_hooks/test_presidio.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) 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 9da3048a728..41fcf6806a0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2833,6 +2833,39 @@ async def test_apply_to_output_streaming_gemini_error_frame_mid_stream_is_carrie assert frame["error"] == {"code": 503, "message": "overloaded", "status": "UNAVAILABLE"} +@pytest.mark.asyncio +async def test_apply_to_output_streaming_gemini_every_candidate_is_merged_by_index_and_masked(): + """ + Gemini interleaves every candidate's fragments across frames, so each candidate + is joined on its own index and scanned, not only the first one. + """ + + def frame(*texts: tuple[int, str]) -> bytes: + candidates = [{"content": {"parts": [{"text": text}], "role": "model"}, "index": index} for index, text in texts] + return b"data: " + json.dumps({"candidates": candidates}).encode() + b"\n\n" + + frames = [ + frame((0, "The architect was John"), (1, "The engineer was Ada")), + frame((0, " Smith."), (1, " Lovelace.")), + ] + collected: list[object] = [] + + async def mock_stream(): + for chunk in frames: + yield chunk + + 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,) = _gemini_frames(collected) + assert [candidate["content"]["parts"] for candidate in masked["candidates"]] == [ + [{"text": "The architect was ."}], + [{"text": "The engineer was ."}], + ] + + @pytest.mark.asyncio async def test_apply_to_output_streaming_gemini_upstream_abort_mid_stream_forwards_no_frame(): guardrail = _OPTIONAL_PresidioPIIMasking(