Merge 5c3de90b0b into c58fc0c996

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -12,7 +12,7 @@ import asyncio
 import json
 import re
 import threading
-from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Iterator, Sequence
+from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Iterator, Mapping, Sequence
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime
@@ -20,7 +20,7 @@ from itertools import accumulate, count, pairwise
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
 
 import aiohttp
-from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config
+from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 import litellm
@@ -85,6 +85,7 @@ class _PresidioAnonymizeResponse(TypedDict):
 _PRESIDIO_ANONYMIZE_ADAPTER: Final[TypeAdapter[_PresidioAnonymizeResponse | None]] = TypeAdapter(
     _PresidioAnonymizeResponse | None
 )
+_SSE_EVENT_ADAPTER: Final[TypeAdapter[dict[str, JsonValue]]] = TypeAdapter(dict[str, JsonValue])
 
 
 class _JsonResponse(Protocol):
@@ -1306,24 +1307,27 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         Replace PII tokens in *text* with their original values.
 
         Includes a fallback for tokens that were truncated by ``max_tokens``:
-        if the *end* of ``text`` matches the *beginning* of a token and the
-        overlap is long enough, the truncated suffix is replaced with the
-        original value.  The minimum overlap length is
-        ``min(20, len(token) // 2)`` to reduce the risk of false positives
-        when multiple tokens share a common prefix.
+        if the *end* of ``text`` matches the *beginning* of exactly one token
+        and the overlap is long enough, the truncated suffix is replaced with
+        that token's original value. A suffix shared by more than one token
+        (e.g. ``<PERSON_`` of two people) is left as-is rather than guessing
+        which one it was. A token only ever starts with ``<``, so the one
+        candidate suffix is the text after the last ``<`` — checking just
+        that keeps this linear in the number of tokens.
         """
         for token, original_text in pii_tokens.items():
             if token in text:
                 text = text.replace(token, original_text)
-            else:
-                # FALLBACK: Handle truncated tokens (token cut off by max_tokens)
-                # Only check at the very end of the text.
-                min_overlap = min(20, len(token) // 2)
-                for i in range(max(0, len(text) - len(token)), len(text)):
-                    sub = text[i:]
-                    if token.startswith(sub) and len(sub) >= min_overlap:
-                        text = text[:i] + original_text
-                        break
+
+        start: Final = text.rfind("<")
+        if start == -1:
+            return text
+        suffix: Final = text[start:]
+        matches: Final = [token for token in pii_tokens if token.startswith(suffix)]
+        if len(matches) == 1:
+            unique_token: Final = matches[0]
+            if len(suffix) >= min(20, len(unique_token) // 2):
+                return text[:start] + pii_tokens[unique_token]
         return text
 
     @staticmethod
@@ -1583,35 +1587,131 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         return anthropic_sse_chunks_from_response(assembled, chunks)
 
     @staticmethod
-    def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str]) -> bytes:
-        try:
-            text: Final = chunk.decode("utf-8")
-        except UnicodeDecodeError:
-            return chunk
+    def _replace_complete_tokens(text: str, pii_tokens: Mapping[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.
+        """
+        result = text  # rebind-ok: accumulates token replacements across the loop
+        for token, original_text in pii_tokens.items():
+            if token in result:
+                result = result.replace(token, original_text)
+        return result
 
+    @staticmethod
+    def _longest_token_prefix_suffix(text: str, pii_tokens: Mapping[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 = ""  # rebind-ok: iterative best-match search across every token
+        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: Mapping[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, dict(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 _flush_pending_sse(pending_text: str, pending_bytes: bytes, pii_tokens: Mapping[str, str]) -> tuple[bytes, ...]:
+        """
+        Whatever ``_unmask_sse_bytes_chunk`` was still holding when the stream
+        ended or the upstream iterator raised: a token held back in case it
+        kept growing, and/or a ``data: ...`` line whose trailing bytes never
+        arrived. Flushing both here means neither is silently dropped.
+        """
+        text_chunk: Final = (
+            (OPTIONAL_PresidioPIIMasking._flush_pending_as_sse_text_delta(pending_text, pii_tokens, 0) + "\n\n").encode(
+                "utf-8"
+            )
+            if pending_text
+            else b""
+        )
+        return tuple(chunk for chunk in (text_chunk, pending_bytes) if chunk)
+
+    @staticmethod
+    def _unmask_sse_bytes_chunk(
+        chunk: bytes, pii_tokens: Mapping[str, str], carried_in: str, carried_bytes: bytes = b""
+    ) -> tuple[bytes, str, bytes]:
+        """
+        Unmask PII tokens in a raw Anthropic SSE bytes chunk, carrying state
+        across calls for two kinds of boundary splits: trailing text that
+        might be the start of a token split across multiple streamed deltas
+        (*carried_in*), and a ``data: ...`` line cut mid-JSON by a raw network
+        chunk boundary (*carried_bytes*, prepended to this chunk before
+        anything is parsed). Returns the rewritten chunk plus both carry-overs
+        to pass into the next call for the same stream.
+        """
+        data: Final = carried_bytes + chunk
+        try:
+            text: Final = data.decode("utf-8")
+        except UnicodeDecodeError:
+            return data, carried_in, b""
+
+        if "\n" not in text:
+            # No complete line yet in this chunk; wait for the rest of it.
+            return b"", carried_in, data
+
+        lines: Final = text.split("\n")
+        incomplete_tail = lines.pop()  # rebind-ok: the segment after the last "\n", not a full line yet
+
+        pending = carried_in  # rebind-ok: grows/flushes per delta as the chunk's lines are scanned
         result_lines: Final[list[str]] = []
-        for line in text.split("\n"):
+        for line in lines:
             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
+                    event = _SSE_EVENT_ADAPTER.validate_python(json.loads(raw_json))
+                    delta = event.get("delta")
+                    delta_dict = delta if isinstance(delta, dict) else None
+                    delta_text = delta_dict.get("text") if delta_dict is not None else None
                     if (
-                        isinstance(delta, dict)
+                        delta_dict is not None
                         and event.get("type") == "content_block_delta"
-                        and delta.get("type") == "text_delta"
-                        and isinstance(delta.get("text"), str)
+                        and delta_dict.get("type") == "text_delta"
+                        and isinstance(delta_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:
+                            delta_dict["text"] = emit_text
                             line = "data: " + json.dumps(event, ensure_ascii=False)
-                except (json.JSONDecodeError, KeyError, TypeError):
+                    elif pending and event.get("type") != "ping":
+                        index_value = event.get("index", 0)
+                        index = index_value if isinstance(index_value, int) else 0
+                        result_lines.append(
+                            OPTIONAL_PresidioPIIMasking._flush_pending_as_sse_text_delta(pending, pii_tokens, index)
+                        )
+                        pending = ""
+                except (json.JSONDecodeError, ValidationError, KeyError, TypeError):
                     pass
             result_lines.append(line)
 
-        return "\n".join(result_lines).encode("utf-8")
+        rewritten: Final = ("\n".join(result_lines) + "\n").encode("utf-8")
+        return rewritten, pending, incomplete_tail.encode("utf-8")
 
     def _unmask_responses_api_completed_chunk(self, chunk: object, pii_tokens: dict[str, str]) -> None:
         """
@@ -1636,23 +1736,79 @@ 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],
-        request_data: dict,
+        remaining_chunks: Sequence[ModelResponseStream],
+        request_data: dict[str, Any],  # mutable-ok: forwarded to _process_response_for_pii, which needs 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
+
+        chunks_list: Final = list(remaining_chunks)
+        assembled_model_response: Final = stream_chunk_builder(
+            chunks=chunks_list, messages=request_data.get("messages")
+        )
+
+        if not isinstance(assembled_model_response, ModelResponse):
+            for chunk in chunks_list:
+                yield chunk
+            return
+
+        self._preserve_usage_from_last_chunk(assembled_model_response, chunks_list)
+
+        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 _handle_non_chat_stream_chunk(
+        self,
+        chunk: object,
+        pii_tokens: dict[str, str],  # mutable-ok: passed straight through to _unmask_responses_api_completed_chunk
+        remaining_chunks: list[ModelResponseStream],  # mutable-ok: cleared in place once flushed below
+    ) -> AsyncGenerator[object, None]:
+        """
+        A chunk that is neither a ``ModelResponseStream`` nor raw SSE bytes
+        (a ``/v1/responses`` event). A mixed stream can't be reassembled, so
+        any chat-completion chunks buffered so far are flushed in order
+        first rather than dropped (a no-op once already flushed, since
+        nothing more is ever buffered after this point), then this chunk is
+        unmasked in place if needed.
+        """
+        for buffered_chunk in remaining_chunks:
+            yield buffered_chunk
+        remaining_chunks.clear()
+        chunk_type: Final = getattr(chunk, "type", None)
+        if pii_tokens and chunk_type in ("response.completed", "response.incomplete"):
+            self._unmask_responses_api_completed_chunk(chunk, pii_tokens)
+        yield chunk
+
+    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 = ""  # rebind-ok: carries a split token across SSE chunks
+        pending_sse_bytes = b""  # rebind-ok: carries a data: line split across byte chunks
         try:
             async for chunk in response:
                 if isinstance(chunk, ModelResponseStream):
@@ -1662,54 +1818,40 @@ 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, pending_sse_bytes = self._unmask_sse_bytes_chunk(
+                            chunk, pii_tokens, pending_sse_text, pending_sse_bytes
+                        )
+                        if rewritten:
+                            yield rewritten
                     else:
                         yield chunk
                     continue
                 else:
-                    # /v1/responses events: unmask response.completed text in-place.
-                    # A mixed stream can't be reassembled, so flush buffered chat
-                    # chunks in order before passthrough instead of dropping them.
-                    if remaining_chunks and not saw_non_chat_chunk:
-                        for buffered_chunk in remaining_chunks:
-                            yield buffered_chunk
-                        remaining_chunks = []
-                    chunk_type = getattr(chunk, "type", None)
-                    if chunk_type in ("response.completed", "response.incomplete") and pii_tokens:
-                        self._unmask_responses_api_completed_chunk(chunk, pii_tokens)
+                    async for out in self._handle_non_chat_stream_chunk(chunk, pii_tokens, remaining_chunks):
+                        yield out
                     saw_non_chat_chunk = True
-                    yield chunk
-
-            if saw_non_chat_chunk:
-                return
-
-            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",
-            )
-
-            mock_response_stream: Final = convert_model_response_to_streaming(assembled_model_response)
-            yield mock_response_stream
 
+            had_error = False  # rebind-ok: set in the except branch below when the upstream iterator raises
         except Exception as e:
             verbose_proxy_logger.error("Error in PII streaming processing: %s", e)
+            had_error = True  # rebind-ok: distinguishes the except path for the branch below
+        finally:
+            for flushed in self._flush_pending_sse(
+                pending_sse_text,
+                pending_sse_bytes,
+                pii_tokens,  # pyright: ignore[reportUnknownArgumentType]  # request_data is an untyped dict
+            ):
+                yield flushed
+
+        if had_error:
             for chunk in remaining_chunks:
                 yield chunk
+        elif not saw_non_chat_chunk:
+            async for chunk in self._finalize_non_sse_stream(
+                remaining_chunks,
+                request_data,  # pyright: ignore[reportUnknownArgumentType]  # request_data is an untyped dict
+            ):
+                yield chunk
 
     async def async_post_call_streaming_iterator_hook(
         self,

litellm/proxy/guardrails/guardrail_initializers.py

@@ -177,6 +177,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail) ->
             output_parse_pii=True,
             event_hook=_presidio_output_mode(litellm_params.mode, include_mcp=True),
             _callback_role="restore",
+            mask_response_content=True,
         )
         if run_input and litellm_params.output_parse_pii
         else None

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,187 @@ 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"
+
+
+def test_unmask_sse_bytes_chunk_carries_a_data_line_split_across_byte_chunks():
+    """A raw network chunk can cut a `data: {...}` line in the middle of its
+    JSON. The first half must be held, not passed through still masked."""
+    import json
+
+    pii_tokens = {"<PERSON_1>": "Bobby"}
+    event = {
+        "type": "content_block_delta",
+        "index": 0,
+        "delta": {"type": "text_delta", "text": "Hello <PERSON_1>!"},
+    }
+    whole = ("data: " + json.dumps(event) + "\n\n").encode("utf-8")
+    split_at = len(whole) - 10
+    first_half, second_half = whole[:split_at], whole[split_at:]
+
+    result1, pending, carried_bytes = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(
+        first_half, pii_tokens, ""
+    )
+    assert result1 == b""
+    assert pending == ""
+    assert carried_bytes == first_half
+
+    result2, pending, carried_bytes = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(
+        second_half, pii_tokens, pending, carried_bytes
+    )
+    decoded = result2.decode("utf-8")
+    parsed = json.loads(decoded.split("data: ", 1)[1].strip())
+    assert parsed["delta"]["text"] == "Hello Bobby!"
+    assert carried_bytes == b""
+
+
+def test_unmask_pii_text_does_not_guess_an_ambiguous_truncated_token():
+    """When a truncated token's prefix matches more than one candidate
+    (`<PERSON_` of two people), leave it rather than restoring the wrong one."""
+    tokens = {"<PERSON_1>": "Bobby", "<PERSON_2>": "Alice"}
+    assert OPTIONAL_PresidioPIIMasking._unmask_pii_text("Hi <PERSON_", tokens) == "Hi <PERSON_"
+    # Unambiguous once enough of the suffix narrows it to a single token.
+    assert OPTIONAL_PresidioPIIMasking._unmask_pii_text("Hi <PERSON_2", tokens) == "Hi Alice"
+
+
+@pytest.mark.asyncio
+async def test_stream_pii_unmasking_flushes_pending_sse_state_when_upstream_raises(mock_user_api_key):
+    """A dropped upstream connection must not silently drop a token that was
+    being held back, or a data: line fragment still waiting for its rest."""
+    import json
+
+    guardrail = OPTIONAL_PresidioPIIMasking(
+        mock_testing=True,
+        output_parse_pii=True,
+    )
+
+    pii_tokens = {"<PERSON_1>": "Jane Smith"}
+    request_data = {"metadata": {"pii_tokens": pii_tokens}}
+
+    event = {
+        "type": "content_block_delta",
+        "index": 0,
+        "delta": {"type": "text_delta", "text": "Hi <PERSON"},
+    }
+    whole = ("data: " + json.dumps(event) + "\n\n").encode("utf-8")
+    first_half = whole[: len(whole) - 10]
+
+    async def mock_stream():
+        yield first_half
+        raise RuntimeError("upstream connection dropped")
+
+    chunks = []
+    async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data):
+        chunks.append(chunk)
+
+    # The exception is handled internally (logged, not re-raised); the held
+    # partial data: line bytes must still come through rather than being
+    # dropped silently.
+    assert b"".join(chunks) == first_half
 
 
 @pytest.mark.asyncio
@@ -3514,6 +3695,69 @@ async def test_stream_pii_unmasking_passthrough_when_no_tokens(mock_user_api_key
     assert chunks == [raw_chunk]
 
 
+@pytest.mark.asyncio
+async def test_stream_pii_unmasking_reassembles_chat_completion_chunks(mock_user_api_key):
+    """Structured chat completion chunks (no Anthropic SSE bytes) are buffered,
+    assembled into one response, and the placeholder is restored on it."""
+    guardrail = OPTIONAL_PresidioPIIMasking(
+        mock_testing=True,
+        output_parse_pii=True,
+    )
+
+    pii_tokens = {"<PERSON_1>": "Jane Smith"}
+    request_data = {
+        "messages": [{"role": "user", "content": "who is this"}],
+        "metadata": {"pii_tokens": pii_tokens},
+    }
+
+    async def mock_stream():
+        yield _chat_delta_chunk("Signature: <PERS")
+        yield _chat_delta_chunk("ON_1>, thanks")
+        yield _chat_delta_chunk("", finish_reason="stop")
+
+    chunks = []
+    async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data):
+        chunks.append(chunk)
+
+    assert len(chunks) == 1
+    assert isinstance(chunks[0], ModelResponseStream)
+    assert chunks[0].choices[0].delta.content == "Signature: Jane Smith, thanks"
+
+
+@pytest.mark.asyncio
+async def test_stream_pii_unmasking_flushes_pending_at_stream_end_with_no_block_stop(mock_user_api_key):
+    """A stream that ends with no content_block_stop (e.g. a dropped connection)
+    still flushes whatever pending text it was carrying, rather than losing it."""
+    import json
+
+    guardrail = OPTIONAL_PresidioPIIMasking(
+        mock_testing=True,
+        output_parse_pii=True,
+    )
+
+    pii_tokens = {"<PERSON_1>": "Jane Smith"}
+    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("Hi <PERSON")
+
+    chunks = []
+    async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data):
+        chunks.append(chunk)
+
+    assert len(chunks) == 2
+    flushed_event = json.loads(chunks[1].decode("utf-8").split("data: ", 1)[1].strip())
+    assert flushed_event["delta"]["text"] == "Jane Smith"
+
+
 def test_new_entities_pass_through_analyze_payload():
     """
     Newly added upstream entities (e.g. German DE_*) must reach the analyzer
@@ -4521,6 +4765,32 @@ async def test_restoration_never_contacts_presidio(has_tokens: bool) -> None:
     assert analyzed == []
 
 
+@pytest.mark.asyncio
+async def test_output_parse_pii_restore_callback_routes_natively_via_mask_response_content() -> None:
+    """Regression guard (issue #46087): the restore callback must get its native
+    streaming route through the per-instance ``mask_response_content`` flag, not
+    a class-level ``use_native_lifecycle_hooks`` override. The latter would also
+    silently skip during_call scans (no native ``async_moderation_hook``) and
+    /v1/responses pre_call masking (the native ``async_pre_call_hook`` only reads
+    ``data["messages"]``) for every Presidio guardrail instance, not just this one.
+    """
+    from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking
+    from litellm.proxy.guardrails.guardrail_initializers import initialize_presidio
+
+    params: Final = LitellmParams(
+        guardrail="presidio",
+        mode=["pre_call", "post_call"],
+        output_parse_pii=True,
+        presidio_filter_scope="input",
+        presidio_analyzer_api_base="http://test-analyzer/",
+        presidio_anonymizer_api_base="http://test-anonymizer/",
+    )
+    scan_callback, restore_callback = initialize_presidio(params, {"guardrail_name": "restore-mask"})
+    assert scan_callback.mask_response_content is False
+    assert restore_callback.mask_response_content is True
+    assert _OPTIONAL_PresidioPIIMasking.use_native_lifecycle_hooks is False
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("event_hook", ["pre_call", ["pre_call"], ["pre_call", "post_call"]])
 async def test_standalone_restoration_preserves_post_call_selection(event_hook: str | list[str]) -> None:

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."""