From d25807b7c1fec173012ba1f3677d18cb53960e63 Mon Sep 17 00:00:00 2001 From: Cyrill Bannwart Date: Tue, 5 May 2026 21:25:03 +0000 Subject: [PATCH] [Fix] End-of-stream-only guardrail bypass in action mode --- .../unified_guardrail/unified_guardrail.py | 17 ++++ .../test_unified_guardrail.py | 78 +++++++++++++++++++ 2 files changed, 95 insertions(+) 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 ca43aa5d1ac..e7e7eb9a607 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -587,6 +587,7 @@ class UnifiedLLMGuardrails(CustomLogger): request_data=request_data, user_api_key_dict=user_api_key_dict, sampling_rate=max(1, int(sampling_rate)), + end_of_stream_only=bool(end_of_stream_only), ): yield chunk return @@ -811,6 +812,7 @@ class UnifiedLLMGuardrails(CustomLogger): request_data: dict, user_api_key_dict: UserAPIKeyAuth, sampling_rate: int, + end_of_stream_only: bool = False, ) -> AsyncGenerator[Any, None]: """ Streaming action protocol state machine. @@ -819,6 +821,15 @@ class UnifiedLLMGuardrails(CustomLogger): on every chunk while in WAIT state), emits delta chunks past the cursor on NONE/GUARDRAIL_INTERVENED, terminates on BLOCKED, and always makes a final is_final=True call after upstream EOS. + + When `end_of_stream_only=True`, the per-sample mid-stream calls are + skipped entirely: chunks accumulate without invoking the guardrail + and without emitting anything to the client; a single is_final=True + call decides the whole response. This preserves the operator's + `streaming_end_of_stream_only` config flag for action-mode + guardrails — without it, an action-mode guardrail would still be + sampled mid-stream and could emit content past the cursor before + the final inspection had a chance to BLOCK. """ cursor = 0 in_wait_state = False @@ -838,6 +849,12 @@ class UnifiedLLMGuardrails(CustomLogger): if template_chunk is None: template_chunk = item + # end_of_stream_only forces buffer-all behavior: skip per-sample + # guardrail calls and per-sample emissions; a single is_final=True + # call below decides the whole response. + if end_of_stream_only: + continue + sample_due = (chunk_counter % sampling_rate == 0) or in_wait_state if not sample_due: continue 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 54a46cc5f9f..32d56993654 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,84 @@ class TestUnifiedLLMGuardrails: ) ) + @pytest.mark.asyncio + async def test_action_mode_end_of_stream_only_blocks_before_emit( + self, + ): + """ + With `streaming_end_of_stream_only=True`, action-mode guardrails + must not emit any content past cursor before the EOS decision. + A BLOCKED at is_final=true terminates the stream cleanly without + anything having reached the client mid-stream. + """ + from litellm.exceptions import GuardrailRaisedException + + handler = UnifiedLLMGuardrails() + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + # Only one decision — should only be called once at EOS. + [("BLOCKED", None, "policy violation at EOS")], + sampling_rate=1, # would normally fire every chunk + ) + guardrail.streaming_end_of_stream_only = True + + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + with pytest.raises( + GuardrailRaisedException, match="policy violation at EOS" + ): + 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}, + ) + ) + # Single guardrail call, at is_final=True. Mid-stream calls + # were suppressed by end_of_stream_only. + assert len(guardrail.calls) == 1, ( + f"expected single EOS call, got {len(guardrail.calls)}: " + f"{guardrail.calls}" + ) + assert guardrail.calls[0]["is_final"] is True + + @pytest.mark.asyncio + async def test_action_mode_end_of_stream_only_emits_modified_at_eos( + self, + ): + """ + With `streaming_end_of_stream_only=True`, mid-stream samples are + suppressed. A successful GUARDRAIL_INTERVENED at is_final=true + emits the modified text in one delta and terminates. + """ + handler = UnifiedLLMGuardrails() + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + [("GUARDRAIL_INTERVENED", "REWRITTEN", None)], + sampling_rate=1, + ) + guardrail.streaming_end_of_stream_only = True + + 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) + assert text == "REWRITTEN", text + assert out[-1].choices[0].finish_reason == "stop" + # Single guardrail call, at is_final=True. + assert len(guardrail.calls) == 1 + assert guardrail.calls[0]["is_final"] is True + @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