From 1f102d177865dd933afa59e652aa8ac3a441cb1d Mon Sep 17 00:00:00 2001 From: VimarshDwivedi <167737761+VimarshDwivedi@users.noreply.github.com> Date: Tue, 13 Oct 2026 21:36:43 +0530 Subject: [PATCH 1/7] fix(guardrails): restore streamed PII placeholders on Anthropic /v1/messages With output_parse_pii enabled, Presidio's restore callback implements apply_guardrail, so on /v1/messages the proxy downgraded its native streaming hook to the unified scan-only path, which drops text rewrites. Clients got the raw placeholder back instead of the original value. Setting use_native_lifecycle_hooks keeps the restore callback on its own hook, matching every other guardrail that rewrites streamed content. That alone still corrupted a placeholder split across two SSE deltas, since each delta was unmasked in isolation. _unmask_sse_bytes_chunk now carries a pending buffer across chunks for the same stream, holding back only text that could still grow into a token, and flushes it with the truncation-fallback once the block actually ends (not on a keep-alive ping). Fixes #46087 --- .../guardrails/guardrail_hooks/presidio.py | 146 ++++++++++++++---- .../guardrail_hooks/test_presidio.py | 119 +++++++++++++- .../test_proxy_logging_hook_detection.py | 67 ++++++++ 3 files changed, 291 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index a60059fcc4d..232fb1d6da1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -16,7 +16,7 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaita from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast import aiohttp from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config @@ -178,6 +178,8 @@ async def _coalesce_first_sse_frame(stream: AsyncIterator[object]) -> AsyncGener class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): + use_native_lifecycle_hooks: ClassVar[bool] = True + user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None @@ -1526,12 +1528,62 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return anthropic_sse_chunks_from_response(assembled) @staticmethod - def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str]) -> bytes: + def _replace_complete_tokens(text: str, pii_tokens: dict[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. + """ + for token, original_text in pii_tokens.items(): + if token in text: + text = text.replace(token, original_text) + return text + + @staticmethod + def _longest_token_prefix_suffix(text: str, pii_tokens: dict[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 = "" + 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: dict[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, 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 _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str], carried_in: str) -> tuple[bytes, str]: + """ + Unmask PII tokens in a raw Anthropic SSE bytes chunk, carrying over any + trailing text that might be the start of a token split across multiple + streamed deltas. Returns the rewritten chunk and the carry-over text to + pass into the next call for the same stream. + """ try: text: Final = chunk.decode("utf-8") except UnicodeDecodeError: - return chunk + return chunk, carried_in + pending = carried_in result_lines: Final[list[str]] = [] for line in text.split("\n"): line = line.rstrip("\r") @@ -1546,15 +1598,26 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): and delta.get("type") == "text_delta" and isinstance(delta.get("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"]: + event["delta"]["text"] = emit_text line = "data: " + json.dumps(event, ensure_ascii=False) + elif pending and isinstance(event, dict) and event.get("type") != "ping": + index = event.get("index", 0) if isinstance(event, dict) else 0 + result_lines.append( + OPTIONAL_PresidioPIIMasking._flush_pending_as_sse_text_delta(pending, pii_tokens, index) + ) + pending = "" except (json.JSONDecodeError, KeyError, TypeError): pass result_lines.append(line) - return "\n".join(result_lines).encode("utf-8") + return "\n".join(result_lines).encode("utf-8"), pending def _unmask_responses_api_completed_chunk(self, chunk: object, pii_tokens: dict[str, str]) -> None: """ @@ -1579,23 +1642,55 @@ 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], + remaining_chunks: list[ModelResponseStream], request_data: 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 + + 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", + ) + + yield convert_model_response_to_streaming(assembled_model_response) + + 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 = "" try: async for chunk in response: if isinstance(chunk, ModelResponseStream): @@ -1605,7 +1700,8 @@ 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 = self._unmask_sse_bytes_chunk(chunk, pii_tokens, pending_sse_text) + yield rewritten else: yield chunk continue @@ -1623,31 +1719,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): saw_non_chat_chunk = True yield chunk - if saw_non_chat_chunk: - return + if pending_sse_text: + line = self._flush_pending_as_sse_text_delta(pending_sse_text, pii_tokens, 0) + yield (line + "\n\n").encode("utf-8") - 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: + if not saw_non_chat_chunk: + async for chunk in self._finalize_non_sse_stream(remaining_chunks, request_data): 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 except Exception as e: verbose_proxy_logger.error("Error in PII streaming processing: %s", e) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 7d0c81af4c2..92f54c1e532 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,111 @@ 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 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.""" From d0a135182351c347d5bf69d79bad965312b04a8f Mon Sep 17 00:00:00 2001 From: VimarshDwivedi <167737761+VimarshDwivedi@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:32:57 +0530 Subject: [PATCH 2/7] fix(guardrails): satisfy strict type-discipline rules and close patch coverage gaps The new helpers used dict in signatures (LIT001), rebound a bare local without Final (LIT010), and rebound a function parameter in place (LIT011). Switch to Mapping/Sequence where the callee only reads, suppress the two cases where in-place rebinding across a loop is the whole point, and bind a fresh local instead of rebinding a parameter where that was free to do. Also add direct unit tests for _finalize_non_sse_stream's chat- completions reassembly path and the end-of-stream safety-net flush, both reachable before this PR but newly surfaced as uncovered patch lines by the refactor that split them out. --- .../guardrails/guardrail_hooks/presidio.py | 31 ++++----- .../guardrail_hooks/test_presidio.py | 63 +++++++++++++++++++ 2 files changed, 79 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 232fb1d6da1..3274bc7ca8f 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 @@ -1528,25 +1528,26 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return anthropic_sse_chunks_from_response(assembled) @staticmethod - def _replace_complete_tokens(text: str, pii_tokens: dict[str, str]) -> str: + 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 text: - text = text.replace(token, original_text) - return text + if token in result: + result = result.replace(token, original_text) + return result @staticmethod - def _longest_token_prefix_suffix(text: str, pii_tokens: dict[str, str]) -> str: + 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 = "" + 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): @@ -1556,7 +1557,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return longest @staticmethod - def _flush_pending_as_sse_text_delta(pending: str, pii_tokens: dict[str, str], index: int) -> str: + 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 @@ -1571,7 +1572,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return "data: " + json.dumps(event, ensure_ascii=False) @staticmethod - def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str], carried_in: str) -> tuple[bytes, str]: + def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: Mapping[str, str], carried_in: str) -> tuple[bytes, str]: """ Unmask PII tokens in a raw Anthropic SSE bytes chunk, carrying over any trailing text that might be the start of a token split across multiple @@ -1583,7 +1584,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except UnicodeDecodeError: return chunk, carried_in - pending = carried_in + 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"): line = line.rstrip("\r") @@ -1644,8 +1645,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def _finalize_non_sse_stream( self, - remaining_chunks: list[ModelResponseStream], - request_data: dict, + remaining_chunks: Sequence[ModelResponseStream], + request_data: dict, # mutable-ok: forwarded verbatim to _process_response_for_pii, which requires dict ) -> AsyncGenerator[object, None]: """ Reassemble buffered ``ModelResponseStream`` chunks (the ``/chat/completions`` @@ -1661,7 +1662,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return assembled_model_response: Final = stream_chunk_builder( - chunks=remaining_chunks, messages=request_data.get("messages") + chunks=list(remaining_chunks), messages=request_data.get("messages") ) if not isinstance(assembled_model_response, ModelResponse): @@ -1690,7 +1691,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): remaining_chunks: list[ModelResponseStream] = [] saw_non_chat_chunk = False - pending_sse_text = "" + pending_sse_text = "" # rebind-ok: carries a split token across SSE chunks try: async for chunk in response: if isinstance(chunk, ModelResponseStream): @@ -1720,7 +1721,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk if pending_sse_text: - line = self._flush_pending_as_sse_text_delta(pending_sse_text, pii_tokens, 0) + line: Final = self._flush_pending_as_sse_text_delta(pending_sse_text, pii_tokens, 0) yield (line + "\n\n").encode("utf-8") if not saw_non_chat_chunk: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 92f54c1e532..12dd0620b97 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -3619,6 +3619,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 = {"": "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 Date: Tue, 13 Oct 2026 23:13:34 +0530 Subject: [PATCH 3/7] fix(guardrails): close basedpyright gate growth from the earlier type-discipline fixes Fixing LIT001 (Mapping/Sequence instead of dict/list) left a few new basedpyright diagnostics: an Unknown-typed delta/event payload fed a longer chain of new call sites, a Mapping couldn't satisfy a dict-typed callee, and a Sequence couldn't satisfy a list-typed one. Parse the SSE JSON payload through a TypeAdapter(dict[str, JsonValue]) instead of raw json.loads, which gives delta/event concrete types instead of Unknown and nets fewer unknown-type diagnostics than before this PR, not just no growth. Materialize remaining_chunks into one list up front so every call downstream of it gets the list type it expects. The two residual argument-type diagnostics are the same pre-existing untyped-request_data pattern already priced in elsewhere in this file, just reached from new call sites; suppressed with a reason, matching this file's own existing convention for the identical case. --- .../guardrails/guardrail_hooks/presidio.py | 50 ++++++++++++------- 1 file changed, 31 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 3274bc7ca8f..4804fa68b29 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -19,7 +19,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, ClassVar, 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 @@ -83,6 +83,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): @@ -1563,7 +1564,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 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, pii_tokens) + flushed: Final = OPTIONAL_PresidioPIIMasking._unmask_pii_text(pending, dict(pii_tokens)) event: Final = { "type": "content_block_delta", "index": index, @@ -1591,30 +1592,33 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 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) ): combined = OPTIONAL_PresidioPIIMasking._replace_complete_tokens( - pending + delta["text"], pii_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"]: - event["delta"]["text"] = emit_text + if emit_text != delta_text: + delta_dict["text"] = emit_text line = "data: " + json.dumps(event, ensure_ascii=False) - elif pending and isinstance(event, dict) and event.get("type") != "ping": - index = event.get("index", 0) if isinstance(event, dict) else 0 + 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, KeyError, TypeError): + except (json.JSONDecodeError, ValidationError, KeyError, TypeError): pass result_lines.append(line) @@ -1646,7 +1650,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def _finalize_non_sse_stream( self, remaining_chunks: Sequence[ModelResponseStream], - request_data: dict, # mutable-ok: forwarded verbatim to _process_response_for_pii, which requires dict + request_data: dict[str, Any], # mutable-ok: forwarded to _process_response_for_pii, which needs dict ) -> AsyncGenerator[object, None]: """ Reassemble buffered ``ModelResponseStream`` chunks (the ``/chat/completions`` @@ -1661,16 +1665,17 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if not remaining_chunks: return + chunks_list: Final = list(remaining_chunks) assembled_model_response: Final = stream_chunk_builder( - chunks=list(remaining_chunks), messages=request_data.get("messages") + chunks=chunks_list, messages=request_data.get("messages") ) if not isinstance(assembled_model_response, ModelResponse): - for chunk in remaining_chunks: + for chunk in chunks_list: yield chunk return - self._preserve_usage_from_last_chunk(assembled_model_response, remaining_chunks) + self._preserve_usage_from_last_chunk(assembled_model_response, chunks_list) await self._process_response_for_pii( response=assembled_model_response, @@ -1721,11 +1726,18 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk if pending_sse_text: - line: Final = self._flush_pending_as_sse_text_delta(pending_sse_text, pii_tokens, 0) + line: Final = self._flush_pending_as_sse_text_delta( + pending_sse_text, + pii_tokens, # pyright: ignore[reportUnknownArgumentType] # request_data is an untyped dict + 0, + ) yield (line + "\n\n").encode("utf-8") if not saw_non_chat_chunk: - async for chunk in self._finalize_non_sse_stream(remaining_chunks, request_data): + async for chunk in self._finalize_non_sse_stream( + remaining_chunks, + request_data, # pyright: ignore[reportUnknownArgumentType] # request_data is an untyped dict + ): yield chunk except Exception as e: From e4b9a512199ed8d9dd517d8f216f2343289951b0 Mon Sep 17 00:00:00 2001 From: VimarshDwivedi <167737761+VimarshDwivedi@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:48:30 +0530 Subject: [PATCH 4/7] fix(guardrails): scope the native streaming route to mask_response_content, not the whole class use_native_lifecycle_hooks=True on _OPTIONAL_PresidioPIIMasking fixed the streaming restore path but is a class-level flag read in several other places: during_call routing (presidio has no native async_moderation_hook, so this silently no-ops PII scanning), and the deployment pre_call hook (presidio's native async_pre_call_hook only reads data["messages"], so /v1/responses requests bypass pre-call masking entirely). Flagged by an automated reviewer on the open PR. Revert the class-level flag and instead set mask_response_content=True on just the restore callback, matching the already-correct mask_output_callback pattern. That flag is checked by the same streaming-route condition but has no effect on pre_call/during_call routing, so every other hook keeps going through the generic apply_guardrail path that already handles /v1/responses translation and during_call scanning correctly. --- .../guardrails/guardrail_hooks/presidio.py | 4 +-- .../guardrails/guardrail_initializers.py | 1 + .../guardrail_hooks/test_presidio.py | 26 +++++++++++++++++++ 3 files changed, 28 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 4804fa68b29..d0f103ae810 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -16,7 +16,7 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaita from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast import aiohttp from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config @@ -179,8 +179,6 @@ async def _coalesce_first_sse_frame(stream: AsyncIterator[object]) -> AsyncGener class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): - use_native_lifecycle_hooks: ClassVar[bool] = True - user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None 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 12dd0620b97..7875f29965b 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -4449,6 +4449,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: From 4b4019abce9aff046bfb845abf6d2e21c186b909 Mon Sep 17 00:00:00 2001 From: VimarshDwivedi <167737761+VimarshDwivedi@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:04:32 +0530 Subject: [PATCH 5/7] fix(guardrails): close three streaming-unmask gaps found by PR #46272 A contributor closed their overlapping PR #46272 in favor of this one and flagged three edge cases its own fix didn't cover: - A raw network chunk can cut a `data: {...}` SSE line mid-JSON. Each chunk was decoded and parsed independently, so the split line's JSON failed to parse and passed through to the client still masked. `_unmask_sse_bytes_chunk` now carries the incomplete line's raw bytes to the next call instead of treating a partial line as a complete one. - If the upstream iterator raised mid-stream, the token text and partial SSE line `_unmask_sse_bytes_chunk` was holding back were dropped instead of flushed. `_stream_pii_unmasking`'s exception handler now flushes both via a shared `_flush_pending_sse` helper. - `_unmask_pii_text`'s truncated-token fallback picked whichever token happened first in dict iteration order when a truncated suffix (e.g. `= min_overlap: - text = text[:i] + original_text - break + + start = text.rfind("<") + if start == -1: + return text + suffix = text[start:] + matches = [token for token in pii_tokens if token.startswith(suffix)] + if len(matches) == 1: + token = matches[0] + if len(suffix) >= min(20, len(token) // 2): + return text[:start] + pii_tokens[token] return text @staticmethod @@ -1571,21 +1574,50 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return "data: " + json.dumps(event, ensure_ascii=False) @staticmethod - def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: Mapping[str, str], carried_in: str) -> tuple[bytes, str]: + def _flush_pending_sse(pending_text: str, pending_bytes: bytes, pii_tokens: Mapping[str, str]) -> list[bytes]: """ - Unmask PII tokens in a raw Anthropic SSE bytes chunk, carrying over any - trailing text that might be the start of a token split across multiple - streamed deltas. Returns the rewritten chunk and the carry-over text to - pass into the next call for the same stream. + 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. """ + flushed: list[bytes] = [] + if pending_text: + line = OPTIONAL_PresidioPIIMasking._flush_pending_as_sse_text_delta(pending_text, pii_tokens, 0) + flushed.append((line + "\n\n").encode("utf-8")) + if pending_bytes: + flushed.append(pending_bytes) + return flushed + + @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 = chunk.decode("utf-8") + text: Final = data.decode("utf-8") except UnicodeDecodeError: - return chunk, carried_in + 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 = 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:] @@ -1620,7 +1652,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): pass result_lines.append(line) - return "\n".join(result_lines).encode("utf-8"), pending + 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: """ @@ -1695,6 +1728,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 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): @@ -1704,8 +1738,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): remaining_chunks.append(chunk) elif isinstance(chunk, bytes): if pii_tokens: - rewritten, pending_sse_text = self._unmask_sse_bytes_chunk(chunk, pii_tokens, pending_sse_text) - yield rewritten + 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 @@ -1723,13 +1760,12 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): saw_non_chat_chunk = True yield chunk - if pending_sse_text: - line: Final = self._flush_pending_as_sse_text_delta( - pending_sse_text, - pii_tokens, # pyright: ignore[reportUnknownArgumentType] # request_data is an untyped dict - 0, - ) - yield (line + "\n\n").encode("utf-8") + 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 not saw_non_chat_chunk: async for chunk in self._finalize_non_sse_stream( @@ -1740,6 +1776,12 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except Exception as e: verbose_proxy_logger.error("Error in PII streaming processing: %s", e) + 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 for chunk in remaining_chunks: yield chunk diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 7875f29965b..6a230ffe7ac 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -3377,7 +3377,7 @@ def test_unmask_sse_bytes_chunk_replaces_text_delta(): } chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") - result, pending = 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()) @@ -3393,7 +3393,7 @@ 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, pending = 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 == "" @@ -3404,21 +3404,21 @@ 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, pending2 = 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, pending = 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, pending = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"": "Bobby"}, "") + result, pending, _ = OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"": "Bobby"}, "") assert result == chunk assert pending == "" @@ -3434,7 +3434,7 @@ def test_unmask_sse_bytes_chunk_non_ascii_pii_not_escaped(): } chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") - result, pending = 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 @@ -3454,7 +3454,7 @@ 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, pending = 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()) @@ -3481,12 +3481,12 @@ def test_unmask_sse_bytes_chunk_carries_token_split_across_deltas(): 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, "") + 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 == "": "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 Date: Fri, 16 Oct 2026 12:10:46 +0530 Subject: [PATCH 6/7] fix(guardrails): keep _stream_pii_unmasking under the complexity gate Threading the byte-level SSE carry and the shared exception/success flush through the function pushed its cyclomatic complexity from 14 to 16, tripping the ratchet (ceiling 15). Extracting the /v1/responses event handling into _handle_non_chat_stream_chunk and merging the duplicated flush-on-exit into a single try/except/finally brings it to 12. --- .../guardrails/guardrail_hooks/presidio.py | 60 +++++++++++-------- 1 file changed, 35 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index c84a6dca770..e9f71296b2c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1716,6 +1716,28 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield convert_model_response_to_streaming(assembled_model_response) + async def _handle_non_chat_stream_chunk( + self, + chunk: object, + pii_tokens: dict[str, str], + remaining_chunks: list[ModelResponseStream], + ) -> 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 = 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], @@ -1747,43 +1769,31 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 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 - - 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 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 + had_error = False except Exception as e: verbose_proxy_logger.error("Error in PII streaming processing: %s", e) + had_error = True + 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, From 5c3de90b0b0abe4037923186cf04b54bd61bc2fd Mon Sep 17 00:00:00 2001 From: VimarshDwivedi <167737761+VimarshDwivedi@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:28:17 +0530 Subject: [PATCH 7/7] fix(guardrails): annotate new locals/params to satisfy the type-discipline gate The previous commit's new code left several local assignments without a Final declaration (LIT010) and a few mutable-collection annotations without a suppression reason (LIT001), pushing the repo-wide counts for both rules past their merge-base ceiling. - _unmask_pii_text's truncation-fallback locals are now Final. - _flush_pending_sse returns tuple[bytes, ...] built from a single Final expression instead of an appended-to list, dropping both the list-typed return annotation and the local list variable. - _unmask_sse_bytes_chunk's `lines` is Final (mutating it via .pop() after is unaffected, same as the file's existing Final-list locals). - _handle_non_chat_stream_chunk's two mutable-collection parameters are marked mutable-ok (one is threaded straight through to a dict-typed callee, the other is deliberately cleared in place); its `chunk_type` local is Final. - _stream_pii_unmasking's `had_error` flag is marked rebind-ok: it's set once in each of two mutually exclusive branches. --- .../guardrails/guardrail_hooks/presidio.py | 41 ++++++++++--------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index e9f71296b2c..ef6313ccb7a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1266,15 +1266,15 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if token in text: text = text.replace(token, original_text) - start = text.rfind("<") + start: Final = text.rfind("<") if start == -1: return text - suffix = text[start:] - matches = [token for token in pii_tokens if token.startswith(suffix)] + suffix: Final = text[start:] + matches: Final = [token for token in pii_tokens if token.startswith(suffix)] if len(matches) == 1: - token = matches[0] - if len(suffix) >= min(20, len(token) // 2): - return text[:start] + pii_tokens[token] + unique_token: Final = matches[0] + if len(suffix) >= min(20, len(unique_token) // 2): + return text[:start] + pii_tokens[unique_token] return text @staticmethod @@ -1574,20 +1574,21 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return "data: " + json.dumps(event, ensure_ascii=False) @staticmethod - def _flush_pending_sse(pending_text: str, pending_bytes: bytes, pii_tokens: Mapping[str, str]) -> list[bytes]: + 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. """ - flushed: list[bytes] = [] - if pending_text: - line = OPTIONAL_PresidioPIIMasking._flush_pending_as_sse_text_delta(pending_text, pii_tokens, 0) - flushed.append((line + "\n\n").encode("utf-8")) - if pending_bytes: - flushed.append(pending_bytes) - return flushed + 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( @@ -1612,7 +1613,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # No complete line yet in this chunk; wait for the rest of it. return b"", carried_in, data - lines = text.split("\n") + 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 @@ -1719,8 +1720,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def _handle_non_chat_stream_chunk( self, chunk: object, - pii_tokens: dict[str, str], - remaining_chunks: list[ModelResponseStream], + 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 @@ -1733,7 +1734,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): for buffered_chunk in remaining_chunks: yield buffered_chunk remaining_chunks.clear() - chunk_type = getattr(chunk, "type", None) + 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 @@ -1773,10 +1774,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield out saw_non_chat_chunk = True - had_error = False + 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 + had_error = True # rebind-ok: distinguishes the except path for the branch below finally: for flushed in self._flush_pending_sse( pending_sse_text,