Address greptile findings

litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py

@@ -548,6 +548,7 @@ class GenericGuardrailAPI(CustomGuardrail):
             guardrail_response=guardrail_response,
         )
 
+    @log_guardrail_information
     async def apply_guardrail_action(
         self,
         inputs: GenericGuardrailAPIInputs,

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -152,6 +152,13 @@ def _chunk_tool_call_deltas(chunk: Any) -> List[Any]:
     return list(tcs)
 
 
+def _get_attr_or_key(obj: Any, key: str) -> Any:
+    """Read `key` off `obj` whether it's a dict (subscript) or a model (attribute)."""
+    if isinstance(obj, dict):
+        return obj.get(key)
+    return getattr(obj, key, None)
+
+
 def _accumulate_tool_calls(chunks: List[Any]) -> List[dict]:
     """
     Reduce all tool_call deltas across chunks into per-index accumulated state.
@@ -166,6 +173,7 @@ def _accumulate_tool_calls(chunks: List[Any]) -> List[dict]:
     a complete payload should return WAIT until they can parse it.
     """
     by_index: dict[int, dict] = {}
+    _g = _get_attr_or_key
     for chunk in chunks:
         for tc in _chunk_tool_call_deltas(chunk):
             idx = getattr(tc, "index", None)
@@ -178,11 +186,6 @@ def _accumulate_tool_calls(chunks: List[Any]) -> List[dict]:
                 {"id": None, "type": None, "function": {"name": None, "arguments": ""}},
             )
 
-            def _g(obj: Any, key: str) -> Any:
-                if isinstance(obj, dict):
-                    return obj.get(key)
-                return getattr(obj, key, None)
-
             if slot["id"] is None:
                 slot["id"] = _g(tc, "id")
             if slot["type"] is None:
@@ -944,23 +947,27 @@ class UnifiedLLMGuardrails(CustomLogger):
                     should_wrap_with_default_message=False,
                 )
 
-        # At EOS we can't retract bytes already emitted. Soft-fall back to
-        # the raw accumulated text rather than failing the whole stream.
+        # At EOS we can't retract bytes already emitted. If the guardrail
+        # returns text shorter than what's been emitted, emit nothing more
+        # — bytes through `cursor` were sent (with whatever substitutions
+        # earlier calls applied) and falling back to `accumulated_text`
+        # here would leak the raw, unmodified tail (e.g. unredacted PII)
+        # for guardrails that rewrote mid-stream but returned short at EOS.
+        # The terminal chunk below still carries the finish_reason.
         if len(new_text) < cursor:
             verbose_proxy_logger.error(
                 "UnifiedLLMGuardrails action mode (EOS): %s returned text "
-                "shorter than already-emitted (cursor=%d, new=%d) — falling "
-                "back to raw accumulated text",
+                "shorter than already-emitted (cursor=%d, new=%d) — "
+                "stopping emission to avoid leaking unmodified tail",
                 guardrail_name,
                 cursor,
                 len(new_text),
             )
-            new_text = accumulated_text
-
-        delta = new_text[cursor:]
-        cursor = len(new_text)
-        if delta and template_chunk is not None:
-            yield _build_delta_chunk(template_chunk, delta)
+        else:
+            delta = new_text[cursor:]
+            cursor = len(new_text)
+            if delta and template_chunk is not None:
+                yield _build_delta_chunk(template_chunk, delta)
         # Replay any tool_call deltas that arrived between the last sample-point
         # emit and EOS so the client sees the complete tool-call stream.
         for tc_chunk in _replay_tool_call_chunks(

tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py

@@ -635,6 +635,57 @@ class TestUnifiedLLMGuardrails:
                     )
                 )
 
+        @pytest.mark.asyncio
+        async def test_action_mode_eos_shrink_does_not_leak_raw_tail(self):
+            """At EOS, a guardrail returning text shorter than the cursor
+            must not cause the raw, unmodified tail to leak.
+
+            Regression scenario: a redaction guardrail substitutes cleanly
+            on every mid-stream sample, then at is_final=True returns short
+            text (e.g. due to a bug or backend timeout). Falling back to
+            `accumulated_text` would emit `accumulated_text[cursor:]` —
+            the raw, unredacted tail of the upstream stream. The correct
+            behavior is to emit nothing further; the bytes through cursor
+            were already sent (with substitutions) and we cannot extend
+            without potentially leaking source content past cursor.
+            """
+            handler = UnifiedLLMGuardrails()
+            # Mid-stream sample at cursor=2 (after chunk 2): substitute
+            # "ab" with "AB". Cursor advances to 2.
+            # Mid-stream sample at cursor=4: substitute with "ABCD". Cursor=4.
+            # EOS sample: guardrail returns short text "AB" (len 2 < cursor 4)
+            # — should not leak the raw tail "efgh".
+            guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail(
+                [
+                    ("GUARDRAIL_INTERVENED", "AB", None),
+                    ("GUARDRAIL_INTERVENED", "ABCD", None),
+                    ("GUARDRAIL_INTERVENED", "AB", None),  # EOS shrink
+                ],
+                sampling_rate=2,
+            )
+            user = UserAPIKeyAuth(
+                api_key="k", request_route="/v1/chat/completions"
+            )
+            out = await TestUnifiedLLMGuardrails.TestActionMode._collect(
+                handler.async_post_call_streaming_iterator_hook(
+                    user_api_key_dict=user,
+                    response=TestUnifiedLLMGuardrails.TestActionMode._content_chunks(
+                        ["ab", "cd", "ef", "gh"]
+                    ),
+                    request_data={"guardrail_to_apply": guardrail},
+                )
+            )
+            text = "".join(c.choices[0].delta.content or "" for c in out)
+            # Stream emitted "AB" + "CD" through mid-stream samples.
+            # EOS shrink: nothing more is emitted. The raw tail "efgh"
+            # MUST NOT appear.
+            assert "ef" not in text and "gh" not in text, (
+                f"raw upstream tail leaked at EOS shrink: {text!r}"
+            )
+            assert text == "ABCD", text
+            # Stream still terminates with finish_reason.
+            assert out[-1].choices[0].finish_reason == "stop"
+
         @pytest.mark.asyncio
         async def test_action_mode_chunk_straddling_surrogate(self):
             """Surrogate token spans chunks: WAIT until complete, then INTERVENED."""