From 1bcf4a26912d947f4da297cef377f28bf7b031ab Mon Sep 17 00:00:00 2001 From: Avani Prajapati <143805019+Avani-prajapati@users.noreply.github.com> Date: Wed, 10 Jun 2026 15:30:02 +0530 Subject: [PATCH] fix(presidio): unmask PII tokens in Anthropic native SSE streaming bytes (#30028) * fix(presidio): unmask PII tokens in Anthropic native SSE streaming bytes When output_parse_pii=true on the Anthropic native path (anthropic/claude-*), response chunks arrive as raw bytes in SSE format. _stream_pii_unmasking was yielding those bytes unchanged, so tokens were never replaced with the original values before reaching the caller. Add _unmask_sse_bytes_chunk to parse each data: line, find content_block_delta / text_delta events, and apply _unmask_pii_text before re-encoding. Wire it into _stream_pii_unmasking so bytes chunks are unmasked when pii_tokens exist. * fix(presidio): handle CRLF line endings and non-ASCII PII in SSE unmask Strip trailing \r before the [DONE] guard so CRLF-terminated SSE chunks don't bypass it and silently swallow a JSONDecodeError. Add ensure_ascii=False to json.dumps so non-ASCII replacement values like accented names are preserved as UTF-8 on the wire rather than being \uXXXX-escaped. Add regression tests for both cases. --- .../guardrails/guardrail_hooks/presidio.py | 41 ++++- .../guardrail_hooks/test_presidio.py | 162 +++++++++++++++++- 2 files changed, 197 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index fc414ab7b54..033e1d0b8e7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1225,6 +1225,39 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): for chunk in all_chunks: yield chunk + @staticmethod + def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: Dict[str, str]) -> bytes: + try: + text = chunk.decode("utf-8") + except UnicodeDecodeError: + return chunk + + result_lines: List[str] = [] + for line in text.split("\n"): + line = line.rstrip("\r") + if line.startswith("data: ") and line != "data: [DONE]": + raw_json = line[6:] + try: + event = json.loads(raw_json) + delta = event.get("delta") if isinstance(event, dict) else None + if ( + isinstance(delta, dict) + and event.get("type") == "content_block_delta" + and delta.get("type") == "text_delta" + and isinstance(delta.get("text"), str) + ): + unmasked = _OPTIONAL_PresidioPIIMasking._unmask_pii_text( + delta["text"], pii_tokens + ) + if unmasked != delta["text"]: + event["delta"]["text"] = unmasked + line = "data: " + json.dumps(event, ensure_ascii=False) + except (json.JSONDecodeError, KeyError, TypeError): + pass + result_lines.append(line) + + return "\n".join(result_lines).encode("utf-8") + async def _stream_pii_unmasking( self, response: Any, @@ -1237,13 +1270,19 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): from litellm.main import stream_chunk_builder from litellm.types.utils import ModelResponse + metadata = (request_data.get("metadata") or {}) if request_data else {} + pii_tokens: Dict[str, str] = metadata.get("pii_tokens", {}) + remaining_chunks: List[ModelResponseStream] = [] try: async for chunk in response: if isinstance(chunk, ModelResponseStream): remaining_chunks.append(chunk) elif isinstance(chunk, bytes): - yield chunk # type: ignore[misc] + if pii_tokens: + yield self._unmask_sse_bytes_chunk(chunk, pii_tokens) # type: ignore[misc] + else: + yield chunk # type: ignore[misc] continue if not remaining_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 565bf83c6a2..8a5eeeff367 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2398,11 +2398,9 @@ async def test_anonymize_text_uses_correct_positions_no_parse_pii(): ) expected = "My name is , my email is , phone " - assert result == expected, ( - f"anonymize_text produced garbled output with PII remnants.\n" - f"Expected: {expected!r}\n" - f"Got: {result!r}" - ) + assert ( + result == expected + ), f"anonymize_text produced garbled output with PII remnants.\nExpected: {expected!r}\nGot: {result!r}" assert masked_entity_count == { "PERSON": 1, "EMAIL_ADDRESS": 1, @@ -2495,3 +2493,157 @@ async def test_anonymize_text_uses_correct_positions_with_parse_pii(): assert pii_tokens.get("") == "John Smith" assert pii_tokens.get("") == "john@example.com" assert pii_tokens.get("") == "555-867-5309" + + +def test_unmask_sse_bytes_chunk_replaces_text_delta(): + import json + + pii_tokens = {"": "Bobby"} + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello , how are you?"}, + } + chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens) + + decoded = result.decode("utf-8") + parsed = json.loads(decoded.split("data: ", 1)[1].strip()) + assert parsed["delta"]["text"] == "Hello Bobby, how are you?" + + +def test_unmask_sse_bytes_chunk_ignores_non_text_delta(): + import json + + pii_tokens = {"": "Bobby"} + + # message_start event — no delta + event = {"type": "message_start", "message": {"id": "msg_01", "role": "assistant"}} + chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens) + assert result == chunk + + # input_json_delta — should not be touched + event2 = { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": '{"name": ""}'}, + } + chunk2 = ("data: " + json.dumps(event2) + "\n\n").encode("utf-8") + result2 = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk2, pii_tokens) + assert result2 == chunk2 + + +def test_unmask_sse_bytes_chunk_handles_malformed_json(): + chunk = b"data: {not valid json}\n\n" + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk( + chunk, {"": "Bobby"} + ) + assert result == chunk + + +def test_unmask_sse_bytes_chunk_handles_unicode_decode_error(): + chunk = b"\xff\xfe invalid utf-8" + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk( + chunk, {"": "Bobby"} + ) + assert result == chunk + + +def test_unmask_sse_bytes_chunk_non_ascii_pii_not_escaped(): + import json + + pii_tokens = {"": "José"} + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello !"}, + } + chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens) + + decoded = result.decode("utf-8") + assert "Jos\\u" not in decoded + parsed = json.loads(decoded.split("data: ", 1)[1].strip()) + assert parsed["delta"]["text"] == "Hello José!" + + +def test_unmask_sse_bytes_chunk_handles_crlf_line_endings(): + import json + + pii_tokens = {"": "Bobby"} + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hi !"}, + } + crlf_chunk = ("data: " + json.dumps(event) + "\r\ndata: [DONE]\r\n").encode("utf-8") + + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk( + crlf_chunk, pii_tokens + ) + + decoded = result.decode("utf-8") + parsed = json.loads(decoded.split("data: ", 1)[1].split("\n")[0].strip()) + assert parsed["delta"]["text"] == "Hi Bobby!" + assert "data: [DONE]" in decoded + + +@pytest.mark.asyncio +async def test_stream_pii_unmasking_unmaskes_bytes_chunks(mock_user_api_key): + import json + + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + pii_tokens = {"": "Bobby"} + request_data = {"metadata": {"pii_tokens": pii_tokens}} + + def _make_sse_chunk(text: str) -> bytes: + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": text}, + } + return ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + + async def mock_stream(): + yield _make_sse_chunk("Hello !") + yield _make_sse_chunk(" How can I help?") + + chunks = [] + async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data): + chunks.append(chunk) + + assert len(chunks) == 2 + first = chunks[0].decode("utf-8") + first_event = json.loads(first.split("data: ", 1)[1].strip()) + assert first_event["delta"]["text"] == "Hello Bobby!" + + second = chunks[1].decode("utf-8") + second_event = json.loads(second.split("data: ", 1)[1].strip()) + assert second_event["delta"]["text"] == " How can I help?" + + +@pytest.mark.asyncio +async def test_stream_pii_unmasking_passthrough_when_no_tokens(mock_user_api_key): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + raw_chunk = b"data: {}\n\n" + request_data: dict = {"metadata": {}} + + async def mock_stream(): + yield raw_chunk + + chunks = [] + async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data): + chunks.append(chunk) + + assert chunks == [raw_chunk]