From 1f102d177865dd933afa59e652aa8ac3a441cb1d Mon Sep 17 00:00:00 2001 From: VimarshDwivedi <167737761+VimarshDwivedi@users.noreply.github.com> Date: Tue, 13 Oct 2026 21:36:43 +0530 Subject: [PATCH] fix(guardrails): restore streamed PII placeholders on Anthropic /v1/messages With output_parse_pii enabled, Presidio's restore callback implements apply_guardrail, so on /v1/messages the proxy downgraded its native streaming hook to the unified scan-only path, which drops text rewrites. Clients got the raw placeholder back instead of the original value. Setting use_native_lifecycle_hooks keeps the restore callback on its own hook, matching every other guardrail that rewrites streamed content. That alone still corrupted a placeholder split across two SSE deltas, since each delta was unmasked in isolation. _unmask_sse_bytes_chunk now carries a pending buffer across chunks for the same stream, holding back only text that could still grow into a token, and flushes it with the truncation-fallback once the block actually ends (not on a keep-alive ping). Fixes #46087 --- .../guardrails/guardrail_hooks/presidio.py | 146 ++++++++++++++---- .../guardrail_hooks/test_presidio.py | 119 +++++++++++++- .../test_proxy_logging_hook_detection.py | 67 ++++++++ 3 files changed, 291 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index a60059fcc4d..232fb1d6da1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -16,7 +16,7 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaita from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast import aiohttp from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config @@ -178,6 +178,8 @@ async def _coalesce_first_sse_frame(stream: AsyncIterator[object]) -> AsyncGener class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): + use_native_lifecycle_hooks: ClassVar[bool] = True + user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None @@ -1526,12 +1528,62 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return anthropic_sse_chunks_from_response(assembled) @staticmethod - def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str]) -> bytes: + def _replace_complete_tokens(text: str, pii_tokens: dict[str, str]) -> str: + """ + Replace every complete occurrence of a PII token in *text* with its + original value, leaving a token that only partially appears (e.g. cut + off at the end by a streamed delta boundary) untouched. + """ + for token, original_text in pii_tokens.items(): + if token in text: + text = text.replace(token, original_text) + return text + + @staticmethod + def _longest_token_prefix_suffix(text: str, pii_tokens: dict[str, str]) -> str: + """ + The longest suffix of *text* that is a strict, non-empty prefix of a + PII token, i.e. text that may still grow into a full token with more + incoming stream data. Empty when no suffix qualifies. + """ + longest = "" + for token in pii_tokens: + max_len = min(len(token) - 1, len(text)) + for suffix_len in range(max_len, len(longest), -1): + if token.startswith(text[-suffix_len:]): + longest = text[-suffix_len:] + break + return longest + + @staticmethod + def _flush_pending_as_sse_text_delta(pending: str, pii_tokens: dict[str, str], index: int) -> str: + """ + A synthetic ``content_block_delta`` SSE line carrying *pending* once no + more stream data can arrive for it, so the truncated-token fallback in + ``_unmask_pii_text`` is safe to apply. + """ + flushed: Final = OPTIONAL_PresidioPIIMasking._unmask_pii_text(pending, pii_tokens) + event: Final = { + "type": "content_block_delta", + "index": index, + "delta": {"type": "text_delta", "text": flushed}, + } + return "data: " + json.dumps(event, ensure_ascii=False) + + @staticmethod + def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str], carried_in: str) -> tuple[bytes, str]: + """ + Unmask PII tokens in a raw Anthropic SSE bytes chunk, carrying over any + trailing text that might be the start of a token split across multiple + streamed deltas. Returns the rewritten chunk and the carry-over text to + pass into the next call for the same stream. + """ try: text: Final = chunk.decode("utf-8") except UnicodeDecodeError: - return chunk + return chunk, carried_in + pending = carried_in result_lines: Final[list[str]] = [] for line in text.split("\n"): line = line.rstrip("\r") @@ -1546,15 +1598,26 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 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 + combined = OPTIONAL_PresidioPIIMasking._replace_complete_tokens( + pending + delta["text"], pii_tokens + ) + carry = OPTIONAL_PresidioPIIMasking._longest_token_prefix_suffix(combined, pii_tokens) + emit_text = combined[: len(combined) - len(carry)] if carry else combined + pending = carry + if emit_text != delta["text"]: + event["delta"]["text"] = emit_text line = "data: " + json.dumps(event, ensure_ascii=False) + elif pending and isinstance(event, dict) and event.get("type") != "ping": + index = event.get("index", 0) if isinstance(event, dict) else 0 + result_lines.append( + OPTIONAL_PresidioPIIMasking._flush_pending_as_sse_text_delta(pending, pii_tokens, index) + ) + pending = "" except (json.JSONDecodeError, KeyError, TypeError): pass result_lines.append(line) - return "\n".join(result_lines).encode("utf-8") + return "\n".join(result_lines).encode("utf-8"), pending def _unmask_responses_api_completed_chunk(self, chunk: object, pii_tokens: dict[str, str]) -> None: """ @@ -1579,23 +1642,55 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): elif hasattr(content_block, "text") and isinstance(content_block.text, str): content_block.text = self._unmask_pii_text(content_block.text, pii_tokens) - async def _stream_pii_unmasking( + async def _finalize_non_sse_stream( self, - response: AsyncIterable[object], + remaining_chunks: list[ModelResponseStream], request_data: dict, ) -> AsyncGenerator[object, None]: - """Apply PII unmasking to streaming output (output_parse_pii=True path).""" + """ + Reassemble buffered ``ModelResponseStream`` chunks (the ``/chat/completions`` + path) into one response, unmask it, and replay it as a single chunk. + """ from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, ) from litellm.main import stream_chunk_builder from litellm.types.utils import ModelResponse + if not remaining_chunks: + return + + assembled_model_response: Final = stream_chunk_builder( + chunks=remaining_chunks, messages=request_data.get("messages") + ) + + if not isinstance(assembled_model_response, ModelResponse): + for chunk in remaining_chunks: + yield chunk + return + + self._preserve_usage_from_last_chunk(assembled_model_response, remaining_chunks) + + await self._process_response_for_pii( + response=assembled_model_response, + request_data=request_data, + mode="unmask", + ) + + yield convert_model_response_to_streaming(assembled_model_response) + + async def _stream_pii_unmasking( + self, + response: AsyncIterable[object], + request_data: dict, + ) -> AsyncGenerator[object, None]: + """Apply PII unmasking to streaming output (output_parse_pii=True path).""" metadata: Final = (request_data.get("metadata") or {}) if request_data else {} pii_tokens: Final[dict[str, str]] = metadata.get("pii_tokens", {}) remaining_chunks: list[ModelResponseStream] = [] saw_non_chat_chunk = False + pending_sse_text = "" try: async for chunk in response: if isinstance(chunk, ModelResponseStream): @@ -1605,7 +1700,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): remaining_chunks.append(chunk) elif isinstance(chunk, bytes): if pii_tokens: - yield self._unmask_sse_bytes_chunk(chunk, pii_tokens) + rewritten, pending_sse_text = self._unmask_sse_bytes_chunk(chunk, pii_tokens, pending_sse_text) + yield rewritten else: yield chunk continue @@ -1623,31 +1719,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): saw_non_chat_chunk = True yield chunk - if saw_non_chat_chunk: - return + if pending_sse_text: + line = self._flush_pending_as_sse_text_delta(pending_sse_text, pii_tokens, 0) + yield (line + "\n\n").encode("utf-8") - if not remaining_chunks: - return - - assembled_model_response: Final = stream_chunk_builder( - chunks=remaining_chunks, messages=request_data.get("messages") - ) - - if not isinstance(assembled_model_response, ModelResponse): - for chunk in remaining_chunks: + if not saw_non_chat_chunk: + async for chunk in self._finalize_non_sse_stream(remaining_chunks, request_data): yield chunk - return - - self._preserve_usage_from_last_chunk(assembled_model_response, remaining_chunks) - - await self._process_response_for_pii( - response=assembled_model_response, - request_data=request_data, - mode="unmask", - ) - - mock_response_stream: Final = convert_model_response_to_streaming(assembled_model_response) - yield mock_response_stream except Exception as e: verbose_proxy_logger.error("Error in PII streaming processing: %s", e) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 7d0c81af4c2..92f54c1e532 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -3377,11 +3377,12 @@ def test_unmask_sse_bytes_chunk_replaces_text_delta(): } chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") - result = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens) + result, pending = 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?" + assert pending == "" def test_unmask_sse_bytes_chunk_ignores_non_text_delta(): @@ -3392,8 +3393,9 @@ def test_unmask_sse_bytes_chunk_ignores_non_text_delta(): # 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) + result, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens, "") assert result == chunk + assert pending == "" # input_json_delta — should not be touched event2 = { @@ -3402,20 +3404,23 @@ def test_unmask_sse_bytes_chunk_ignores_non_text_delta(): "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) + result2, pending2 = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk2, pii_tokens, "") assert result2 == chunk2 + assert pending2 == "" 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"}) + result, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"": "Bobby"}, "") assert result == chunk + assert pending == "" 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"}) + result, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"": "Bobby"}, "") assert result == chunk + assert pending == "" def test_unmask_sse_bytes_chunk_non_ascii_pii_not_escaped(): @@ -3429,12 +3434,13 @@ def test_unmask_sse_bytes_chunk_non_ascii_pii_not_escaped(): } chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") - result = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens) + result, pending = 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é!" + assert pending == "" def test_unmask_sse_bytes_chunk_handles_crlf_line_endings(): @@ -3448,12 +3454,111 @@ def test_unmask_sse_bytes_chunk_handles_crlf_line_endings(): } 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) + result, pending = 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 + assert pending == "" + + +def test_unmask_sse_bytes_chunk_carries_token_split_across_deltas(): + """A placeholder split across two text_delta events is reassembled, not corrupted (issue #46087).""" + import json + + pii_tokens = {"": "Jane Smith"} + event1 = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Signature: , thanks"}, + } + chunk1 = ("data: " + json.dumps(event1) + "\n\n").encode("utf-8") + chunk2 = ("data: " + json.dumps(event2) + "\n\n").encode("utf-8") + + result1, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk1, pii_tokens, "") + parsed1 = json.loads(result1.decode("utf-8").split("data: ", 1)[1].strip()) + assert parsed1["delta"]["text"] == "Signature: " + assert pending == "": "Jane Smith"} + delta_event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hi ": "Jane Smith"} + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "just ordinary text"}, + } + chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + + result, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens, "") + + parsed = json.loads(result.decode("utf-8").split("data: ", 1)[1].strip()) + assert parsed["delta"]["text"] == "just ordinary text" + assert pending == "" + + +def test_unmask_sse_bytes_chunk_keeps_pending_across_a_keepalive_ping(): + """A keep-alive ping mid-block must not flush a still-growing token early.""" + import json + + pii_tokens = {"": "Jane Smith"} + delta_event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hi token + instead of the original value, including when the token was split across two deltas.""" + import json + + from litellm.caching.caching import DualCache + from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler + from litellm.types.guardrails import SupportedGuardrailIntegrations + + handler = InMemoryGuardrailHandler() + result = handler.initialize_guardrail( + guardrail={ + "guardrail_name": "presidio-pii-restore", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value, + "mode": ["pre_call", "post_call"], + "default_on": True, + "output_parse_pii": True, + "presidio_filter_scope": "input", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + "pii_entities_config": {"PERSON": "MASK"}, + "mock_redacted_text": {"text": "", "items": []}, + }, + } + ) + guardrail_id = result["guardrail_id"] + callbacks = [ + handler.guardrail_id_to_custom_guardrail[guardrail_id], + *handler.guardrail_id_to_sibling_callbacks[guardrail_id], + ] + monkeypatch.setattr(litellm, "callbacks", callbacks) + + # The pre_call scan step would have stored this mapping; set it directly + # so the test exercises only the restore step the issue is about. + pii_tokens = {"": "Jane Smith"} + chunks = _anthropic_stream_chunks(["Signature: , thanks"]) + + async def fake_stream(): + for chunk in chunks: + yield chunk + + delivered = [] + async for chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/messages"), + request_data={ + "model": "claude-sonnet-5", + "litellm_logging_obj": _streaming_logging_obj(), + "metadata": {"pii_tokens": pii_tokens}, + }, + ): + delivered.append(chunk) + + wire = b"".join(delivered).decode() + text_deltas = [ + json.loads(line[6:])["delta"]["text"] + for line in wire.split("\n") + if line.startswith("data: ") and json.loads(line[6:]).get("delta", {}).get("type") == "text_delta" + ] + assert "" not in wire, wire + assert "".join(text_deltas) == "Signature: Jane Smith, thanks", wire + assert wire.count("event: message_stop") == 1, wire + + class _AppliesGuardrail(CustomGuardrail): """Implements the unified interface only, so the proxy routes it to unified_guardrail."""