[Fix] End-of-stream-only guardrail bypass in action mode

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

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