From 8d947db6fa31f627c3dc9cb7d87545a3669845ff Mon Sep 17 00:00:00 2001
From: VimarshDwivedi <167737761+VimarshDwivedi@users.noreply.github.com>
Date: Fri, 16 Oct 2026 12:10:46 +0530
Subject: [PATCH] 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,