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 <PERSON_1> 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

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)

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": "<PERSON_1>"}'},
     }
     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, {"<PERSON_1>": "Bobby"})
+    result, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"<PERSON_1>": "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, {"<PERSON_1>": "Bobby"})
+    result, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"<PERSON_1>": "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 = {"<PERSON_1>": "Jane Smith"}
+    event1 = {
+        "type": "content_block_delta",
+        "index": 0,
+        "delta": {"type": "text_delta", "text": "Signature: <PERS"},
+    }
+    event2 = {
+        "type": "content_block_delta",
+        "index": 0,
+        "delta": {"type": "text_delta", "text": "ON_1>, 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 == "<PERS"
+
+    result2, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk2, pii_tokens, pending)
+    parsed2 = json.loads(result2.decode("utf-8").split("data: ", 1)[1].strip())
+    assert parsed2["delta"]["text"] == "Jane Smith, thanks"
+    assert pending == ""
+
+
+def test_unmask_sse_bytes_chunk_flushes_pending_before_block_stop():
+    """A token truncated by max_tokens (never completes) still unmasks via the end-of-block flush."""
+    import json
+
+    pii_tokens = {"<PERSON_1>": "Jane Smith"}
+    delta_event = {
+        "type": "content_block_delta",
+        "index": 0,
+        "delta": {"type": "text_delta", "text": "Hi <PERSON"},
+    }
+    stop_event = {"type": "content_block_stop", "index": 0}
+    delta_chunk = ("data: " + json.dumps(delta_event) + "\n\n").encode("utf-8")
+    stop_chunk = ("data: " + json.dumps(stop_event) + "\n\n").encode("utf-8")
+
+    result1, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(delta_chunk, pii_tokens, "")
+    parsed1 = json.loads(result1.decode("utf-8").split("data: ", 1)[1].strip())
+    assert parsed1["delta"]["text"] == "Hi "
+    assert pending == "<PERSON"
+
+    result2, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(stop_chunk, pii_tokens, pending)
+    lines = [line for line in result2.decode("utf-8").split("\n") if line.startswith("data: ")]
+    flushed = json.loads(lines[0][len("data: ") :])
+    stopped = json.loads(lines[1][len("data: ") :])
+    assert flushed["delta"]["text"] == "Jane Smith"
+    assert stopped["type"] == "content_block_stop"
+    assert pending == ""
+
+
+def test_unmask_sse_bytes_chunk_does_not_delay_text_with_no_token_overlap():
+    """Ordinary text sharing no prefix with any token is emitted immediately, with no added latency."""
+    import json
+
+    pii_tokens = {"<PERSON_1>": "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 = {"<PERSON_1>": "Jane Smith"}
+    delta_event = {
+        "type": "content_block_delta",
+        "index": 0,
+        "delta": {"type": "text_delta", "text": "Hi <PERS"},
+    }
+    ping_event = {"type": "ping"}
+    delta_chunk = ("data: " + json.dumps(delta_event) + "\n\n").encode("utf-8")
+    ping_chunk = ("data: " + json.dumps(ping_event) + "\n\n").encode("utf-8")
+
+    result1, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(delta_chunk, pii_tokens, "")
+    assert pending == "<PERS"
+
+    result2, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(ping_chunk, pii_tokens, pending)
+    assert result2 == ping_chunk
+    assert pending == "<PERS"
 
 
 @pytest.mark.asyncio

tests/unit/proxy/test_proxy_logging_hook_detection.py

@@ -605,6 +605,73 @@ async def test_post_call_stream_presidio_output_masking_masks_anthropic_messages
     assert wire.count("event: message_stop") == 1, wire
 
 
+@pytest.mark.asyncio
+async def test_post_call_stream_presidio_output_parse_pii_restores_anthropic_messages_stream(monkeypatch):
+    """Regression (issue #46087): the presidio output_parse_pii restore callback was rerouted
+    onto the unified scan-only path on /v1/messages, so the client saw the raw <PERSON_1> 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": "<PERSON_1>", "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 = {"<PERSON_1>": "Jane Smith"}
+    chunks = _anthropic_stream_chunks(["Signature: <PERS", "ON_1>, 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 "<PERSON_1>" 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."""