From f10ed0f651cd914378fb991c1c2f23abfb66a4c7 Mon Sep 17 00:00:00 2001 From: loitragg Date: Mon, 5 Oct 2026 16:54:58 +0700 Subject: [PATCH] fix(presidio): restore masked PII in streamed /v1/messages responses With output_parse_pii, streamed Anthropic /v1/messages responses came back with the tokens still in them, for two reasons The restore callback was routed through unified_guardrail, whose block_only streaming mode drops text rewrites. It now sets mask_response_content so Presidio's own streaming hook runs The SSE unmask also worked one transport chunk at a time, so a token split across deltas () leaked. AnthropicSSEUnmasker now reassembles events across chunks, holds back a partial trailing token per content block, and flushes it before content_block_stop. It also restores input_json_delta tool arguments and leaves signed thinking_delta untouched --- .../guardrails/guardrail_hooks/presidio.py | 12 +- .../guardrail_hooks/presidio_anthropic_sse.py | 158 ++++++++++++ .../guardrails/guardrail_initializers.py | 2 + .../test_presidio_anthropic_sse.py | 224 ++++++++++++++++++ 4 files changed, 394 insertions(+), 2 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/presidio_anthropic_sse.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_presidio_anthropic_sse.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index d0006f1a091..65618d8a59d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -48,6 +48,7 @@ from litellm.proxy.guardrails.anthropic_sse import ( is_anthropic_sse_stream, model_response_text, ) +from litellm.proxy.guardrails.guardrail_hooks.presidio_anthropic_sse import AnthropicSSEUnmasker from litellm.types.guardrails import ( GuardrailEventHooks, LitellmParams, @@ -1588,6 +1589,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): remaining_chunks: list[ModelResponseStream] = [] saw_non_chat_chunk = False + sse_unmasker: Final = AnthropicSSEUnmasker(pii_tokens) if pii_tokens else None try: async for chunk in response: if isinstance(chunk, ModelResponseStream): @@ -1596,8 +1598,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): else: remaining_chunks.append(chunk) elif isinstance(chunk, bytes): - if pii_tokens: - yield self._unmask_sse_bytes_chunk(chunk, pii_tokens) + if sse_unmasker is not None: + if unmasked_bytes := sse_unmasker.feed(chunk): + yield unmasked_bytes else: yield chunk continue @@ -1615,6 +1618,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): saw_non_chat_chunk = True yield chunk + if sse_unmasker is not None and (tail := sse_unmasker.flush()): + yield tail + if saw_non_chat_chunk: return @@ -1645,6 +1651,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): verbose_proxy_logger.error("Error in PII streaming processing: %s", e) for chunk in remaining_chunks: yield chunk + if sse_unmasker is not None and (tail := sse_unmasker.flush()): + yield tail async def async_post_call_streaming_iterator_hook( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio_anthropic_sse.py b/litellm/proxy/guardrails/guardrail_hooks/presidio_anthropic_sse.py new file mode 100644 index 00000000000..0dbed4ef80b --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio_anthropic_sse.py @@ -0,0 +1,158 @@ +""" +Stateful Presidio PII unmasking for Anthropic ``/v1/messages`` SSE streams. + +Claude splits tokens like ```` across deltas, and byte chunks +need not end on an event boundary, so unmasking each chunk on its own misses +them. This reassembles events, holds back a trailing partial token per content +block until it completes, and flushes held text before ``content_block_stop``. + +``thinking_delta`` is left alone: thinking blocks are signed, and rewriting +them would invalidate the signature when the client sends them back. +""" + +import codecs +import json +import re +from collections.abc import Mapping +from functools import reduce +from typing import Final + +_EVENT_SPLIT: Final = re.compile(r"(\r?\n\r?\n)") + +# delta.type -> (field holding the text, whether values land inside a JSON string) +_UNMASKABLE_DELTAS: Final[Mapping[str, tuple[str, bool]]] = { + "text_delta": ("text", False), + "input_json_delta": ("partial_json", True), +} + +_FLUSH_ALL_EVENTS: Final = frozenset({"message_delta", "message_stop", "error"}) + +_MAX_BUFFERED_CHARS: Final = 1_000_000 + + +class AnthropicSSEUnmasker: + """Feed raw SSE byte chunks in order, then call ``flush`` once at end of stream.""" + + def __init__(self, pii_tokens: Mapping[str, str]) -> None: + self._tokens: Final = dict(pii_tokens) + self._max_token_len: Final = max((len(t) for t in self._tokens), default=0) + self._decoder: Final = codecs.getincrementaldecoder("utf-8")(errors="surrogateescape") + self._buffer: str = "" + self._held: dict[int, tuple[str, str]] = {} # mutable-ok: per-block stream state, index -> (delta type, text) + + def feed(self, chunk: bytes) -> bytes: + pieces: Final = _EVENT_SPLIT.split(self._buffer + self._decoder.decode(chunk)) + rest: Final = pieces[-1] + processed: Final = "".join( + tuple(self._process_event(event, terminator) for event, terminator in zip(pieces[:-1:2], pieces[1::2])) + ) + overflow: Final = len(rest) >= _MAX_BUFFERED_CHARS + self._buffer = "" if overflow else rest + return _encode(processed + (rest if overflow else "")) + + def flush(self) -> bytes: + held: Final = self._flush_all_held() + tail: Final = self._buffer + self._decoder.decode(b"", final=True) + self._buffer = "" + return _encode(held + tail) + + def _process_event(self, event_text: str, terminator: str) -> str: + original: Final = event_text + terminator + lines: Final = tuple(event_text.split("\n")) + data_indexes: Final = tuple(i for i, line in enumerate(lines) if line.startswith("data:")) + if len(data_indexes) != 1: + return original + data_index: Final = data_indexes[0] + line_ending: Final = "\r" if lines[data_index].endswith("\r") else "" + try: + event: Final[object] = json.loads(lines[data_index][5:].rstrip("\r").removeprefix(" ")) + except ValueError: + return original + if not isinstance(event, dict): + return original + + match event.get("type"): + case "content_block_stop" if isinstance(event.get("index"), int): + return self._flush_held(event["index"]) + original + case event_type if event_type in _FLUSH_ALL_EVENTS: + return self._flush_all_held() + original + case "content_block_delta": + rewritten: Final = self._unmask_delta(event) + if rewritten is None: + return original + data_line: Final = "data: " + json.dumps(rewritten, ensure_ascii=False) + line_ending + return "\n".join((*lines[:data_index], data_line, *lines[data_index + 1 :])) + terminator + case _: + return original + + def _unmask_delta(self, event: dict[str, object]) -> dict[str, object] | None: + """Return the rewritten event, or None to pass the original through untouched.""" + index: Final = event.get("index") + delta: Final = event.get("delta") + if not isinstance(index, int) or not isinstance(delta, dict): + return None + delta_type: Final = delta.get("type") + spec: Final = _UNMASKABLE_DELTAS.get(delta_type) if isinstance(delta_type, str) else None + if spec is None or not isinstance(delta.get(spec[0]), str): + return None + field, escape = spec + text: Final[str] = delta[field] + + _, previously_held = self._held.pop(index, ("", "")) + emit, hold = self._split_partial_token(previously_held + text) + if hold: + self._held[index] = (delta_type, hold) + unmasked: Final = self._replace_tokens(emit, escape) + if unmasked == text: + return None + return {**event, "delta": {**delta, field: unmasked}} + + def _split_partial_token(self, text: str) -> tuple[str, str]: + """Split off a trailing prefix of a known token, e.g. ``"" not in suffix + and len(suffix) < self._max_token_len + and any(token.startswith(suffix) for token in self._tokens) + ) + return (text[:start], suffix) if is_partial_token else (text, "") + + def _replace_tokens(self, text: str, escape: bool) -> str: + return reduce( + lambda acc, item: acc.replace(item[0], _render(item[1], escape)), + self._tokens.items(), + text, + ) + + def _flush_all_held(self) -> str: + return "".join(tuple(self._flush_held(index) for index in sorted(self._held))) + + def _flush_held(self, index: int) -> str: + held: Final = self._held.pop(index, None) + if held is None: + return "" + delta_type, text = held + field, escape = _UNMASKABLE_DELTAS[delta_type] + # Block ended mid-token (e.g. max_tokens): same truncation rule as + # _OPTIONAL_PresidioPIIMasking._unmask_pii_text. + truncated_value: Final = next( + ( + value + for token, value in self._tokens.items() + if token.startswith(text) and len(text) >= min(20, len(token) // 2) + ), + None, + ) + final_text: Final = text if truncated_value is None else _render(truncated_value, escape) + event: Final = {"type": "content_block_delta", "index": index, "delta": {"type": delta_type, field: final_text}} + return "event: content_block_delta\ndata: " + json.dumps(event, ensure_ascii=False) + "\n\n" + + +def _render(value: str, escape: bool) -> str: + return json.dumps(value, ensure_ascii=False)[1:-1] if escape else value + + +def _encode(text: str) -> bytes: + return text.encode("utf-8", errors="surrogateescape") diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index b7f3726d017..964994a7928 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -176,6 +176,8 @@ 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", + # Forces its own streaming hook: unified_guardrail's block_only mode drops the unmask on /v1/messages + 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_anthropic_sse.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio_anthropic_sse.py new file mode 100644 index 00000000000..a4db6f0f76c --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio_anthropic_sse.py @@ -0,0 +1,224 @@ +"""Tests for stateful Presidio unmasking of Anthropic /v1/messages SSE streams.""" + +import json + +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking +from litellm.proxy.guardrails.guardrail_hooks.presidio_anthropic_sse import AnthropicSSEUnmasker + +TOKENS = {"": "Maria Gonzalez", "": "maria@example.com"} + + +def _event(payload: dict) -> bytes: + return f"event: {payload['type']}\ndata: {json.dumps(payload)}\n\n".encode() + + +def _text_delta(text: str, index: int = 0) -> bytes: + return _event({"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}}) + + +def _json_delta(partial: str, index: int = 1) -> bytes: + return _event( + {"type": "content_block_delta", "index": index, "delta": {"type": "input_json_delta", "partial_json": partial}} + ) + + +def _stop(index: int = 0) -> bytes: + return _event({"type": "content_block_stop", "index": index}) + + +def _run(chunks: list[bytes], tokens: dict[str, str] = TOKENS) -> bytes: + unmasker = AnthropicSSEUnmasker(tokens) + return b"".join(unmasker.feed(c) for c in chunks) + unmasker.flush() + + +def _deltas(raw: bytes, field: str = "text", index: int | None = None) -> str: + """Concatenate the given delta field across every event in an SSE stream.""" + out = [] + for line in raw.decode().splitlines(): + if not line.startswith("data: "): + continue + event = json.loads(line[6:]) + if event.get("type") == "content_block_delta" and (index is None or event["index"] == index): + out.append(event["delta"].get(field, "")) + return "".join(out) + + +def test_whole_token_in_one_delta(): + raw = _run([_text_delta("Hello !"), _stop()]) + assert _deltas(raw) == "Hello Maria Gonzalez!" + + +def test_token_split_across_deltas(): + # Observed live from claude-sonnet-4-6: " and `` but" + raw = _run([_text_delta("Hi , mail now"), _stop()]) + assert _deltas(raw) == "Hi Maria Gonzalez, mail maria@example.com now" + assert "<" not in _deltas(raw) + + +def test_token_split_one_char_per_delta(): + token = "" + raw = _run([_text_delta("x ")] + [_text_delta(c) for c in token] + [_text_delta(" y"), _stop()]) + assert _deltas(raw) == "x maria@example.com y" + + +def test_held_text_is_never_emitted_early(): + unmasker = AnthropicSSEUnmasker(TOKENS) + first = unmasker.feed(_text_delta("Dear restore, matching + # _OPTIONAL_PresidioPIIMasking._unmask_pii_text's truncation fallback. + raw = _run([_text_delta("mail ", index=0), + _text_delta("ADDRESS_2>", index=2), + _stop(0), + _stop(2), + ] + ) + assert _deltas(raw, index=0) == "A Maria Gonzalez" + assert _deltas(raw, index=2) == "B maria@example.com" + + +def test_input_json_delta_split_token_is_json_escaped(): + tokens = {"": 'Bob "B" \\ Smith'} + raw = _run([_json_delta('{"to": ""}'), _stop(1)], tokens) + assert json.loads(_deltas(raw, field="partial_json")) == {"to": 'Bob "B" \\ Smith'} + + +def test_thinking_delta_untouched(): + chunk = _event( + {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": ""}} + ) + assert _run([chunk]) == chunk + + +def test_unchanged_events_pass_through_byte_for_byte(): + chunks = [ + _event({"type": "message_start", "message": {"id": "msg_1"}}), + b'event: ping\ndata: {"type": "ping"}\n\n', + _text_delta("no pii here"), + _stop(), + _event({"type": "message_stop"}), + ] + assert _run(chunks) == b"".join(chunks) + + +def test_event_split_across_byte_chunks(): + whole = _text_delta("Hello ") + _stop() + pieces = [whole[i : i + 7] for i in range(0, len(whole), 7)] + assert _deltas(_run(pieces)) == "Hello Maria Gonzalez" + + +def test_multibyte_character_split_across_byte_chunks(): + tokens = {"": "José"} + event = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "¡Hola ! ñ"}} + whole = f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + _stop() + cut = whole.index("ñ".encode()) + 1 # split inside the 2-byte ñ + raw = _run([whole[:cut], whole[cut:]], tokens) + assert _deltas(raw) == "¡Hola José! ñ" + assert "Jos\\u" not in raw.decode() + + +def test_crlf_framing(): + event = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi "}} + chunk = ("event: content_block_delta\r\ndata: " + json.dumps(event) + "\r\n\r\n").encode() + raw = _run([chunk]) + assert raw.endswith(b"\r\n\r\n") + assert json.loads(raw.decode().split("data: ", 1)[1].strip())["delta"]["text"] == "Hi Maria Gonzalez" + + +def test_malformed_and_non_sse_bytes_pass_through(): + chunks = [b"data: {not json}\n\n", b"\xff\xfe raw tail"] + assert _run(chunks) == b"".join(chunks) + + +@pytest.mark.asyncio +async def test_stream_pii_unmasking_handles_split_tokens(): + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True) + request_data = {"metadata": {"pii_tokens": dict(TOKENS)}} + + async def stream(): + yield _text_delta("Dear , see .") + yield _stop() + + out = b"".join([c async for c in guardrail._stream_pii_unmasking(stream(), request_data)]) + assert _deltas(out) == "Dear Maria Gonzalez, see maria@example.com." + + +@pytest.mark.asyncio +async def test_stream_pii_unmasking_flushes_when_stream_ends_mid_token(): + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True) + request_data = {"metadata": {"pii_tokens": dict(TOKENS)}} + + async def stream(): + yield _text_delta("x