diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index cadc9ca689b..b2fd598f02c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/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. ``= 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, diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 0b38f2d9488..4e80cedc671 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/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 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index bd1e9ae4068..e30760ac197 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,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 = {"": "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 ": "Bobby"} + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello !"}, + } + 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 + (`": "Bobby", "": "Alice"} + assert OPTIONAL_PresidioPIIMasking._unmask_pii_text("Hi ": "Jane Smith"} + request_data = {"metadata": {"pii_tokens": pii_tokens}} + + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hi ": "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: , 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 = {"": "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 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: diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index 182ef833069..4183bc946e8 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/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 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."""