scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 98

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)

Browse source 
* 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 <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-25 00:48:42 -07:00 • committed by GitHub
parent 2701e2008b
commit c19ce71bcd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 351 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/guardrails/anthropic_sse.py
View file

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

72
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -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:

61
tests/integration/observability/test_presidio_streaming_output.py
View file

@ -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(

241
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
View file

@ -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": "<PERSON>"},
)
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)) == "<PERSON>"
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": "<PERSON>"},
)
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)) == "<PERSON>"
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": "<PERSON>"},
)
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)) == "<PERSON>"
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": "<PERSON>"},
)
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)) == "<PERSON>"
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": "<PERSON>"},
)
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)) == "<PERSON>"
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": "<PERSON>"},
)
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(