fix(policy_engine): fail open when an unended Messages stream has no text delta to carry the rewrite

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -1311,7 +1311,11 @@ class AnthropicMessagesHandler(BaseTranslation):
                     and guardrailed_texts
                     and guardrailed_texts[0] != string_so_far
                 ):
-                    self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
+                    self._write_ended_stream_text_rewrite(
+                        responses_so_far,
+                        guardrailed_texts[0],
+                        guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
+                    )
                 if deliver_ended_stream_rewrites:
                     returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
                     self._write_ended_stream_tool_call_rewrites(
@@ -1353,7 +1357,11 @@ class AnthropicMessagesHandler(BaseTranslation):
             raise
         unended_texts: Final = _guardrailed_inputs.get("texts")
         if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,):
-            self._write_ended_stream_text_rewrite(responses_so_far, unended_texts[0])
+            self._write_ended_stream_text_rewrite(
+                responses_so_far,
+                unended_texts[0],
+                guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
+            )
         return responses_so_far
 
     def _prepare_request_data(
@@ -1447,26 +1455,40 @@ class AnthropicMessagesHandler(BaseTranslation):
             inputs["model"] = response_model
         return inputs
 
-    @staticmethod
+    @classmethod
     def _write_ended_stream_text_rewrite(
+        cls,
         responses_so_far: MutableSequence[object],  # mutable-ok: rewrites the caller's buffered chunks in place
         rewritten_text: str,
+        guardrail_name: str,
     ) -> None:
         """Deliver an ended-stream guardrail text rewrite by rewriting the
         buffered chunks in place: the first ``text_delta`` carries the full
         rewritten text and every later one is blanked, leaving the surrounding
-        message and content-block framing untouched."""
+        message and content-block framing untouched. A buffer with no
+        ``text_delta`` has nowhere to carry the rewrite, so the pipeline
+        executor discards it and releases the original chunks."""
+
+        def is_text_delta(event: Mapping[str, object]) -> bool:
+            delta: Final = event.get("delta")
+            return (
+                event.get("type") == "content_block_delta"
+                and isinstance(delta, Mapping)
+                and delta.get("type") == "text_delta"
+            )
+
+        if not any(is_text_delta(event) for item in responses_so_far for event in cls._iter_sse_events(item)):
+            from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+            raise UndeliverableStreamRewrite(guardrail_name)
         replacements: Final = chain((rewritten_text,), repeat(""))
 
         def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
-            delta: Final = event.get("delta")
-            if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping):
-                return None
-            if delta.get("type") != "text_delta":
+            if not is_text_delta(event):
                 return None
             return _SSEFieldRewrite("delta", "text", next(replacements))
 
-        AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
+        cls._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
 
     @classmethod
     def _write_ended_stream_tool_call_rewrites(

tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py

@@ -462,6 +462,28 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
         assert "event: message_start" in raw and "event: content_block_stop" in raw
         assert "event: message_stop" not in raw
 
+    @pytest.mark.asyncio
+    async def test_unended_stream_rewrite_with_no_text_delta_to_carry_it_fails_open(self):
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+        class FillEmpty(CustomGuardrail):
+            async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+                return {**inputs, "texts": ["[INJECTED]" for _ in inputs.get("texts", [])]}
+
+        handler = AnthropicMessagesHandler()
+        chunks = self._ended_sse_chunks()[:2]
+        original = [bytes(chunk) for chunk in chunks]
+
+        with pytest.raises(UndeliverableStreamRewrite):
+            await handler.process_output_streaming_response(
+                responses_so_far=chunks,
+                guardrail_to_apply=FillEmpty(guardrail_name="test"),
+                litellm_logging_obj=MagicMock(),
+                deliver_ended_stream_rewrites=True,
+            )
+
+        assert chunks == original
+
     @pytest.mark.asyncio
     async def test_unended_stream_without_rewrite_is_released_with_delivery_expected(self):
         handler = AnthropicMessagesHandler()

tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py

@@ -1304,32 +1304,39 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
         ]
 
     @pytest.mark.asyncio
-    async def test_deliver_rewrite_on_unfinished_stream_lands_in_the_buffered_deltas(self):
-        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+    async def test_deliver_ended_stream_rewrites_every_choice_when_a_usage_only_chunk_closes_the_stream(self):
+        from litellm.types.utils import ModelResponseStream, Usage
 
         handler = OpenAIChatCompletionsHandler()
-
-        def chunk(content: str) -> ModelResponseStream:
-            return ModelResponseStream(
-                id="chatcmpl-123",
-                created=1234567890,
-                model="gpt-4",
-                object="chat.completion.chunk",
-                choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=None)],
-            )
-
-        chunks = [chunk("hello "), chunk("world")]
+        guardrail = MockGuardrail(guardrail_name="test")
+        usage_chunk = ModelResponseStream(
+            id="chatcmpl-123",
+            created=1234567890,
+            model="gpt-4",
+            object="chat.completion.chunk",
+            choices=[],
+            usage=Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12),
+        )
+        chunks = [*self._two_choice_stream_chunks(), usage_chunk]
 
         result = await handler.process_output_streaming_response(
             responses_so_far=chunks,
-            guardrail_to_apply=self._world_masking_guardrail(),
+            guardrail_to_apply=guardrail,
             litellm_logging_obj=None,
             deliver_ended_stream_rewrites=True,
         )
 
         assert result is chunks
-        assert [c.choices[0].delta.content for c in chunks] == ["hello [MASKED]", ""]
-        assert [c.choices[0].finish_reason for c in chunks] == [None, None]
+        assert guardrail.last_inputs["texts"] == ["safe text", "hello world"]
+        assert [(c.choices[0].index, c.choices[0].delta.content) for c in chunks[:4]] == [
+            (0, "SAFE TEXT"),
+            (1, "HELLO WORLD"),
+            (0, ""),
+            (1, ""),
+        ]
+        assert [c.choices[0].finish_reason for c in chunks[:4]] == [None, None, "stop", "stop"]
+        assert chunks[4].choices == []
+        assert chunks[4].usage.completion_tokens == 7
 
     @staticmethod
     def _two_choice_tool_call_stream_chunks() -> list: