From c19ce71bcdad234ca87a7af10cee45375ce811a8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 00:48:42 -0700 Subject: [PATCH] fix(presidio): mask streamed /v1/messages output when the first upstream read is a keepalive, a data-less ping, or a split utf8 character (#43023) * test(presidio): cover first-frame utf8 split, comment keepalive and data-less ping in streaming output masking Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(presidio): classify the streaming output shape on a frame with a data line and tolerate a split utf8 boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(presidio): relay leading data-less sse frames before classifying the stream shape Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/guardrails/anthropic_sse.py | 3 +- .../guardrails/guardrail_hooks/presidio.py | 72 ++++-- .../test_presidio_streaming_output.py | 61 ++++- .../guardrail_hooks/test_presidio.py | 241 ++++++++++++++++++ 4 files changed, 351 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index 28220f09f00..26dd2a95dc4 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -7,6 +7,7 @@ hook scan such a stream, and re-emit it when the guardrail rewrote the response. from __future__ import annotations +import codecs import json from collections.abc import Mapping, Sequence from typing import Final @@ -38,7 +39,7 @@ def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None: if isinstance(chunk, (str, bytes)) ) try: - return raw.decode("utf-8") + return codecs.getincrementaldecoder("utf-8")().decode(raw, final=False) except UnicodeDecodeError: return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 794bf08729e..f5e24c501f1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -10,9 +10,11 @@ import asyncio import json +import re import threading -from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Iterator, Sequence from contextlib import asynccontextmanager +from dataclasses import dataclass from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypedDict, cast @@ -39,6 +41,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, assemble_anthropic_sse_stream, @@ -97,16 +100,47 @@ def _json_escaped_len(text: str) -> int: _MAX_FIRST_SSE_FRAME_BYTES: Final = 64 * 1024 -def _holds_complete_sse_frame(raw: bytes) -> bool: - """Whether ``raw`` holds one blank-line terminated SSE event, or is too large to keep joining.""" - return b"\n\n" in raw or b"\r\n\r\n" in raw or len(raw) >= _MAX_FIRST_SSE_FRAME_BYTES +@dataclass(frozen=True, slots=True) +class _SsePreface: + """Complete leading SSE frames with no ``data:`` line, relayed verbatim before the stream shape is decided.""" + + raw: bytes + + +_SSE_FRAME_END: Final = re.compile(rb"\r\n\r\n|\n\n|\r\r") + + +def _split_sse_preface(complete_frames: bytes) -> tuple[bytes, bytes]: + """Split complete frames into ``(frames before the first data-bearing frame, that frame and everything after)``.""" + start = 0 + for end in _SSE_FRAME_END.finditer(complete_frames): + frame = complete_frames[start : end.end()] + if any(line.startswith(b"data:") for line in frame.splitlines()): + return complete_frames[:start], complete_frames[start:] + start = end.end() + return complete_frames, b"" + + +def _flush_unmaskable_buffer(all_chunks: list[ModelResponseStream]) -> Iterator[ModelResponseStream]: + """Buffered chunks flushed unmasked when a mixed stream shape makes reconstruction impossible.""" + if not all_chunks: + return + verbose_proxy_logger.warning( + "Presidio apply_to_output: mixed stream detected (ModelResponseStream + unknown event). " + "Flushing %d buffered chunks without PII masking and switching to transparent passthrough.", + len(all_chunks), + ) + yield from all_chunks async def _coalesce_first_sse_frame(stream: AsyncIterator[object]) -> AsyncGenerator[object, None]: """ - Join leading raw ``bytes`` chunks until they hold one complete SSE event, so - the stream shape is decided on a whole frame rather than a transport fragment. - Everything after that first frame is forwarded untouched. + Relay leading data-less SSE frames (comment keepalives, events without a + ``data:`` line) as they complete, and join raw ``bytes`` chunks until they + hold one complete SSE event with a data line, so the stream shape is + decided on a whole frame rather than a transport fragment. Everything + after that first frame is forwarded untouched. The byte cap can only be + reached by a single unterminated frame. """ pending = b"" try: @@ -115,7 +149,12 @@ async def _coalesce_first_sse_frame(stream: AsyncIterator[object]) -> AsyncGener yield chunk continue pending += chunk - if _holds_complete_sse_frame(pending): + complete_frames, tail = split_complete_sse_frames(pending) + preface, classifiable = _split_sse_preface(complete_frames) + if preface: + yield _SsePreface(preface) + pending = classifiable + tail + if classifiable or len(pending) >= _MAX_FIRST_SSE_FRAME_BYTES: break else: if pending: @@ -1400,6 +1439,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk else: all_chunks.append(chunk) + elif isinstance(chunk, _SsePreface): + yield chunk.raw elif isinstance(chunk, bytes): first_frame_is_anthropic = ( not passthrough_due_to_unknown_stream_shape @@ -1416,18 +1457,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield masked_chunk return else: - if all_chunks: - # Flush buffered chunks and switch to transparent passthrough for this stream shape. - # NOTE: these buffered chunks are emitted unmasked because this - # stream mixed chunk types and cannot be safely reconstructed. - verbose_proxy_logger.warning( - "Presidio apply_to_output: mixed stream detected (ModelResponseStream + unknown event). " - "Flushing %d buffered chunks without PII masking and switching to transparent passthrough.", - len(all_chunks), - ) - for buffered_chunk in all_chunks: - yield buffered_chunk - all_chunks = [] + for buffered_chunk in _flush_unmaskable_buffer(all_chunks): + yield buffered_chunk + all_chunks = [] passthrough_due_to_unknown_stream_shape = True yield chunk if passthrough_due_to_unknown_stream_shape: diff --git a/tests/integration/observability/test_presidio_streaming_output.py b/tests/integration/observability/test_presidio_streaming_output.py index 5bc0a46427b..680617f1652 100644 --- a/tests/integration/observability/test_presidio_streaming_output.py +++ b/tests/integration/observability/test_presidio_streaming_output.py @@ -181,7 +181,7 @@ class Rig: "model": self.anthropic, "max_tokens": 64, "stream": True, - "messages": [{"role": "user", "content": "who designed it"}], + "messages": [{"role": "user", "content": f"who designed it {uuid.uuid4().hex}"}], **({"guardrails": list(guardrails)} if guardrails is not None else {}), } @@ -193,6 +193,13 @@ def anthropic_text(received: Received) -> str: return "".join(event["delta"]["text"] for event in events if event.get("type") == "content_block_delta") +def anthropic_message_id(received: Received) -> str: + events: Final = tuple( + json.loads(line.removeprefix("data: ")) for line in received.text.split("\n") if line.startswith("data: ") + ) + return "".join(event["message"]["id"] for event in events if event.get("type") == "message_start") + + @contextmanager def presidio_rig( gateway: Gateway, @@ -339,10 +346,10 @@ def test_native_gemini_unauthenticated_request_is_rejected_before_upstream(gatew assert rig.upstream.drain() == () -def anthropic_provider(chunks: tuple[bytes, ...]) -> Callable[[Request], Reply]: +def anthropic_provider(chunks: tuple[bytes, ...], *, pause_between_chunks: float = 0) -> Callable[[Request], Reply]: def provider(request: Request) -> Reply: assert request.target == "/v1/messages", request.target - return Reply(content_type="text/event-stream", chunks=chunks) + return Reply(content_type="text/event-stream", chunks=chunks, pause_between_chunks=pause_between_chunks) return provider @@ -376,6 +383,50 @@ def test_anthropic_messages_first_frame_split_across_transport_chunks_is_still_m assert received.text.count("event: message_start") == 1 +def test_anthropic_messages_first_frame_split_inside_a_utf8_character_is_still_masked( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + whole: Final = anthropic_stream(identity, f"{PERSON} designed the caf\u00e9.") + delta: Final = whole[2].replace("\\u00e9".encode(), "\u00e9".encode()) + split_at: Final = delta.index("\u00e9".encode()) + 1 + assert delta[split_at - 1 : split_at] == b"\xc3", delta + chunks: Final = (whole[0] + whole[1] + delta[:split_at], delta[split_at:], *whole[3:]) + with presidio_rig(gateway, tmp_path, anthropic_provider(chunks, pause_between_chunks=0.5)) as rig: + received: Final = rig.stream("/v1/messages", rig.messages_body()) + assert received.status == 200, received.text + assert anthropic_text(received) == f"{MASK} designed the caf\u00e9." + assert PERSON not in received.text, received.text + assert anthropic_message_id(received) == identity, received.text + + +def test_anthropic_messages_stream_led_by_sse_comment_keepalive_is_still_masked( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + chunks: Final = (b": keepalive\n\n", *anthropic_stream(identity, f"{PERSON} designed it.")) + with presidio_rig(gateway, tmp_path, anthropic_provider(chunks, pause_between_chunks=0.5)) as rig: + received: Final = rig.stream("/v1/messages", rig.messages_body()) + assert received.status == 200, received.text + assert received.text.startswith(": keepalive"), received.text[:200] + assert anthropic_text(received) == f"{MASK} designed it." + assert PERSON not in received.text, received.text + assert identity in received.text + + +def test_anthropic_messages_stream_led_by_data_less_ping_event_is_still_masked( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + chunks: Final = (b"event: ping\n\n", *anthropic_stream(identity, f"{PERSON} designed it.")) + with presidio_rig(gateway, tmp_path, anthropic_provider(chunks, pause_between_chunks=0.5)) as rig: + received: Final = rig.stream("/v1/messages", rig.messages_body()) + assert received.status == 200, received.text + assert anthropic_text(received) == f"{MASK} designed it." + assert PERSON not in received.text, received.text + assert identity in received.text + + def test_anthropic_messages_stream_fails_closed_when_analyzer_is_down(gateway: Gateway, tmp_path: Path) -> None: identity: Final = "msg_" + uuid.uuid4().hex provider: Final = anthropic_provider(anthropic_stream(identity, f"{PERSON} designed it.")) @@ -468,7 +519,7 @@ def test_mixed_burst_survives_anonymizer_outage_and_recovers(gateway: Gateway, t return Reply( content_type="text/event-stream", chunks=(gemini_frame(f"{PERSON} "), gemini_frame("designed it.")), - pause_between_chunks=0.05, + pause_between_chunks=0.5, ) with presidio_rig(gateway, tmp_path, provider, anonymize=flaky_anonymizer) as rig: @@ -514,7 +565,7 @@ def test_mixed_burst_survives_anonymizer_outage_and_recovers(gateway: Gateway, t def test_native_gemini_keeps_streaming_after_one_worker_is_killed(gateway: Gateway, tmp_path: Path) -> None: frames: Final = (gemini_frame("alive "), gemini_frame("still.")) - provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=frames, pause_between_chunks=0.05)) + provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=frames, pause_between_chunks=0.5)) with presidio_rig(gateway, tmp_path, provider) as rig: workers: Final = eventually( lambda: tuple( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index fd48688df84..0a4ffbaef26 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2604,6 +2604,247 @@ async def test_apply_to_output_streaming_anthropic_first_frame_split_across_tran assert joined.count("event: message_start") == 1 +def _anthropic_stream_tail() -> list[bytes]: + return [ + _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {}}), + _anthropic_sse("message_stop", {"type": "message_stop"}), + ] + + +def _anthropic_stream_head() -> list[bytes]: + return [ + _anthropic_sse( + "message_start", + {"type": "message_start", "message": {"id": "msg_1", "model": "claude", "content": [], "usage": {}}}, + ), + _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ] + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_anthropic_first_frame_split_inside_a_utf8_character_is_still_masked(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + delta = ( + "event: content_block_delta\n" + + "data: " + + json.dumps( + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "John Smith designed the café."}, + }, + ensure_ascii=False, + ) + + "\n\n" + ).encode() + cut = delta.index("é".encode()) + 1 + assert delta[cut - 1 : cut] == b"\xc3", delta + byte_chunks = [*_anthropic_stream_head(), delta[:cut], delta[cut:], *_anthropic_stream_tail()] + + async def mock_stream(): + yield b"".join(byte_chunks[:2]) + byte_chunks[2] + for b in byte_chunks[3:]: + yield b + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + joined = b"".join(collected).decode() + assert "John Smith" not in joined, joined + assert "".join(text for _, text in _anthropic_text_deltas(collected)) == "" + assert joined.count("event: message_start") == 1 + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_anthropic_stream_led_by_sse_comment_keepalive_is_still_masked(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + byte_chunks = [ + b": keepalive\n\n", + *_anthropic_stream_head(), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "John Smith"}}, + ), + *_anthropic_stream_tail(), + ] + + async def mock_stream(): + for b in byte_chunks: + yield b + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + raw = b"".join(collected) + assert raw.startswith(b": keepalive\n\n"), raw[:200] + joined = raw.decode() + assert "John Smith" not in joined, joined + assert "".join(text for _, text in _anthropic_text_deltas(collected)) == "" + assert joined.count("event: message_start") == 1 + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_anthropic_stream_led_by_data_less_ping_event_is_still_masked(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + byte_chunks = [ + b"event: ping\n\n", + *_anthropic_stream_head(), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "John Smith"}}, + ), + *_anthropic_stream_tail(), + ] + + async def mock_stream(): + for b in byte_chunks: + yield b + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + raw = b"".join(collected) + assert raw.startswith(b"event: ping\n\n"), raw[:200] + joined = raw.decode() + assert "John Smith" not in joined, joined + assert "".join(text for _, text in _anthropic_text_deltas(collected)) == "" + assert joined.count("event: message_start") == 1 + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_leading_keepalive_is_forwarded_before_upstream_data_arrives(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + gate = asyncio.Event() + byte_chunks = [ + *_anthropic_stream_head(), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "John Smith"}}, + ), + *_anthropic_stream_tail(), + ] + + async def mock_stream(): + yield b": keepalive\n\n" + await gate.wait() + for b in byte_chunks: + yield b + + out = guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ) + assert await asyncio.wait_for(anext(out), 1) == b": keepalive\n\n" + assert not gate.is_set() + + gate.set() + collected = [chunk async for chunk in out] + + joined = b"".join(collected).decode() + assert "John Smith" not in joined, joined + assert "".join(text for _, text in _anthropic_text_deltas(collected)) == "" + assert joined.count("event: message_start") == 1 + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_leading_comments_over_the_frame_cap_are_still_masked(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + keepalives = [b": keepalive\n\n" * 512] * 12 # ~72 KiB of complete comment frames, over the 64 KiB cap + byte_chunks = [ + *keepalives[:-1], + keepalives[-1] + + b"".join( + [ + *_anthropic_stream_head(), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "John Smith"}}, + ), + ] + ), + *_anthropic_stream_tail(), + ] + + async def mock_stream(): + for b in byte_chunks: + yield b + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + joined = b"".join(collected).decode() + assert "John Smith" not in joined, joined + assert "".join(text for _, text in _anthropic_text_deltas(collected)) == "" + assert joined.count("event: message_start") == 1 + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_comment_only_stream_is_forwarded_unchanged(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + + async def mock_stream(): + yield b": keepalive\n\n" + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + assert collected == [b": keepalive\n\n"] + + @pytest.mark.asyncio async def test_apply_to_output_streaming_gemini_first_frame_split_across_transport_chunks_streams_incrementally(): guardrail = _OPTIONAL_PresidioPIIMasking(