diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 6e3277aef8d..b1366fdf120 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 9e3deb7fd3f..ca43aa5d1ac 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/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( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index c7c31938f5b..54a46cc5f9f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/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."""