fix(guardrails): key stream rewrites by choice index and scan delta-only responses buffers

Chat streaming write-backs now match chunks by the choice's index field
instead of its list position, delivering rewrites to the right choice on
n>1 streams; an ended-stream rewrite on a multi-choice buffer fails
closed since stream_chunk_builder collapses the choices. The Responses
fallback joins output_text.delta events when delivery is expected, so a
delta-only buffer is guardrail-checked instead of released raw.

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

@@ -620,6 +620,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 responses_so_far=responses_so_far,
                 guardrailed_response=model_response,
                 pre_guardrail_texts=pre_guardrail_texts,
+                guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
             )
 
     def build_stream_error_items(
@@ -747,8 +748,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         """
         combined_texts: Final[dict[tuple[int, int | None], str]] = {}
 
-        for response_idx, response in enumerate(responses_so_far):
-            for choice_idx, choice in enumerate(response.choices):
+        for response in responses_so_far:
+            for choice in response.choices:
                 if isinstance(choice, litellm.StreamingChoices):
                     content = choice.delta.content
                 elif isinstance(choice, litellm.Choices):
@@ -761,7 +762,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
 
                 if isinstance(content, str):
                     # String content - accumulate for this choice
-                    str_key: tuple[int, int | None] = (choice_idx, None)
+                    str_key: tuple[int, int | None] = (choice.index, None)
                     if str_key not in combined_texts:
                         combined_texts[str_key] = ""
                     combined_texts[str_key] += content
@@ -772,7 +773,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                         text_str = content_item.get("text")
                         if text_str:
                             list_key: tuple[int, int | None] = (
-                                choice_idx,
+                                choice.index,
                                 content_idx,
                             )
                             if list_key not in combined_texts:
@@ -973,24 +974,38 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         responses_so_far: list["ModelResponseStream"],  # mutable-ok: rewrites the caller's buffered chunks in place
         guardrailed_response: "ModelResponse",
         pre_guardrail_texts: tuple[str | None, ...],
+        guardrail_name: str,
     ) -> None:
         """Write ended-stream guardrail text rewrites back across the buffered
-        chunks: each rewritten choice's full text lands in its first
+        chunks: the full rewritten text lands in the choice's first
         content-carrying chunk and the rest are blanked, the same shape the
         in-flight write-back uses. Chunks carrying only finish_reason or usage
-        stay untouched."""
+        stay untouched. A rewrite on a stream carrying more than one distinct
+        choice index fails closed."""
         post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response)
         changed: Final = tuple(
-            (choice_idx, after)
-            for choice_idx, (before, after) in enumerate(zip(pre_guardrail_texts, post_guardrail_texts))
+            after
+            for before, after in zip(pre_guardrail_texts, post_guardrail_texts)
             if before is not None and after is not None and after != before
         )
         if not changed:
             return
+        stream_choice_indices: Final = frozenset(
+            choice.index for response in responses_so_far for choice in response.choices
+        )
+        if len(stream_choice_indices) != 1:
+            # stream_chunk_builder collapses every choice into one index-0
+            # choice, so a rewrite of the rebuilt response cannot be attributed
+            # back to a single choice on an n>1 stream: withhold the stream
+            # rather than deliver the rewrite on the wrong choice
+            from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+            raise UndeliverableStreamRewrite(guardrail_name)
+        target_choice_index: Final = next(iter(stream_choice_indices))
         await self._apply_guardrail_responses_to_output_streaming(
             responses=responses_so_far,
-            guardrailed_texts=[after for _choice_idx, after in changed],  # mutable-ok: callee takes lists
-            task_mappings=[(choice_idx, None) for choice_idx, _after in changed],  # mutable-ok: callee takes lists
+            guardrailed_texts=list(changed),  # mutable-ok: callee takes lists
+            task_mappings=[(target_choice_index, None) for _ in changed],  # mutable-ok: callee takes lists
         )
 
     async def _apply_guardrail_responses_to_output_streaming(
@@ -1008,7 +1023,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         Args:
             responses: List of ModelResponseStream objects to modify
             guardrailed_texts: List of guardrailed text responses (combined from all chunks)
-            task_mappings: List of tuples (choice_idx, content_idx)
+            task_mappings: List of tuples (choice_idx, content_idx), where choice_idx
+                is the choice's ``index`` field, not its position in a chunk's list
 
         Override this method to customize how responses are applied to streaming responses.
         """
@@ -1024,9 +1040,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         # Key: (choice_idx, content_idx), Value: boolean (True if already set)
         already_set: Final[dict[tuple[int, int | None], bool]] = {}
 
-        # Iterate through all responses and update content
-        for response_idx, response in enumerate(responses):
-            for choice_idx_in_response, choice in enumerate(response.choices):
+        # Iterate through all responses and update content, matching each chunk's
+        # choice by its index field: on n>1 streams a chunk usually carries one
+        # choice at list position 0 whose index names the logical choice.
+        for response in responses:
+            for choice in response.choices:
                 if isinstance(choice, litellm.StreamingChoices):
                     content = choice.delta.content
                 elif isinstance(choice, litellm.Choices):
@@ -1039,7 +1057,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
 
                 if isinstance(content, str):
                     # String content
-                    str_key: tuple[int, int | None] = (choice_idx_in_response, None)
+                    str_key: tuple[int, int | None] = (choice.index, None)
                     if str_key in guardrail_map:
                         if str_key not in already_set:
                             # First chunk - set the complete guardrailed text
@@ -1060,7 +1078,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                     for content_idx, content_item in enumerate(content):
                         if "text" in content_item:
                             list_key: tuple[int, int | None] = (
-                                choice_idx_in_response,
+                                choice.index,
                                 content_idx,
                             )
                             if list_key in guardrail_map:

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -634,14 +634,20 @@ class OpenAIResponsesHandler(BaseTranslation):
             return responses_so_far
 
         # ------------------------------------------------------------------ #
-        # Fallback: apply guardrail to the accumulated text string.           #
-        # No structured write-back is possible here; guardrails that only     #
-        # need to block/flag (not rewrite) still work correctly, and a        #
-        # rewrite a caller expects delivered fails closed instead.            #
+        # Fallback: apply guardrail to the accumulated text string. When a    #
+        # caller expects rewrites delivered and only output_text.delta events #
+        # carried the text (a stream cut off before any .done or terminal     #
+        # envelope), the delta text is scanned instead so nothing escapes     #
+        # unchecked. No structured write-back is possible here; guardrails    #
+        # that only need to block/flag (not rewrite) still work correctly,    #
+        # and a rewrite a caller expects delivered fails closed instead.      #
         # ------------------------------------------------------------------ #
         string_so_far: Final = self.get_streaming_string_so_far(responses_so_far)
-        if string_so_far:
-            fallback_inputs: Final = GenericGuardrailAPIInputs(texts=[string_so_far])
+        text_to_check: Final = string_so_far or (
+            self._delta_text_so_far(responses_so_far) if deliver_ended_stream_rewrites else ""
+        )
+        if text_to_check:
+            fallback_inputs: Final = GenericGuardrailAPIInputs(texts=[text_to_check])
             response_model = (
                 final_chunk.get("response", {}).get("model") if isinstance(final_chunk.get("response"), dict) else None
             )
@@ -654,7 +660,7 @@ class OpenAIResponsesHandler(BaseTranslation):
                 logging_obj=litellm_logging_obj,
             )
             fallback_texts: Final = fallback_outputs.get("texts")
-            if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,):
+            if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (text_to_check,):
                 from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
                 raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown")
@@ -754,6 +760,17 @@ class OpenAIResponsesHandler(BaseTranslation):
         """
         return "".join([response.get("text", "") for response in responses_so_far])
 
+    @staticmethod
+    def _delta_text_so_far(responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
+        """Accumulate the text carried by ``response.output_text.delta`` events,
+        for buffers where no ``.done`` event or terminal envelope repeats it."""
+        deltas: Final = (
+            response.get("delta")
+            for response in responses_so_far
+            if response.get("type") == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA.value
+        )
+        return "".join(delta for delta in deltas if isinstance(delta, str))
+
     def _has_text_content(self, response: "ResponsesAPIResponse") -> bool:
         """
         Check if response has any text content to process.

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

@@ -1128,6 +1128,99 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
         assert chunks[1].choices[0].delta.content == " world"
         assert chunks[1].choices[0].finish_reason == "stop"
 
+    @staticmethod
+    def _two_choice_stream_chunks() -> list:
+        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+        def chunk(index: int, content: str, finish_reason: Optional[str] = None) -> ModelResponseStream:
+            return ModelResponseStream(
+                id="chatcmpl-123",
+                created=1234567890,
+                model="gpt-4",
+                object="chat.completion.chunk",
+                choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)],
+            )
+
+        return [
+            chunk(0, "safe "),
+            chunk(1, "hello "),
+            chunk(0, "text", "stop"),
+            chunk(1, "world", "stop"),
+        ]
+
+    @staticmethod
+    def _world_masking_guardrail() -> CustomGuardrail:
+        class MaskWorld(CustomGuardrail):
+            async def apply_guardrail(
+                self,
+                inputs: GenericGuardrailAPIInputs,
+                request_data: dict,
+                input_type: Literal["request", "response"],
+                logging_obj: Optional[Any] = None,
+            ) -> GenericGuardrailAPIInputs:
+                texts = inputs.get("texts", [])
+                return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]}
+
+        return MaskWorld(guardrail_name="test-mask")
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrite_on_multi_choice_stream_fails_closed(self):
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+        handler = OpenAIChatCompletionsHandler()
+        chunks = self._two_choice_stream_chunks()
+
+        with pytest.raises(UndeliverableStreamRewrite):
+            await handler.process_output_streaming_response(
+                responses_so_far=chunks,
+                guardrail_to_apply=self._world_masking_guardrail(),
+                litellm_logging_obj=None,
+                deliver_ended_stream_rewrites=True,
+            )
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):
+        handler = OpenAIChatCompletionsHandler()
+        chunks = self._two_choice_stream_chunks()
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"),
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert result is chunks
+        assert [c.choices[0].delta.content for c in chunks] == ["safe ", "hello ", "text", "world"]
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrite_lands_on_nonzero_choice_index(self):
+        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+        handler = OpenAIChatCompletionsHandler()
+
+        def chunk(content: str, finish_reason: Optional[str]) -> ModelResponseStream:
+            return ModelResponseStream(
+                id="chatcmpl-123",
+                created=1234567890,
+                model="gpt-4",
+                object="chat.completion.chunk",
+                choices=[StreamingChoices(index=1, delta=Delta(content=content), finish_reason=finish_reason)],
+            )
+
+        chunks = [chunk("hello ", None), chunk("world", "stop")]
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=self._world_masking_guardrail(),
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert result is chunks
+        assert chunks[0].choices[0].delta.content == "hello [MASKED]"
+        assert chunks[1].choices[0].delta.content in (None, "")
+
 
 class TestGetStructuredMessages:
     """Test the get_structured_messages method."""

tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py

@@ -1212,6 +1212,77 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
                 deliver_ended_stream_rewrites=True,
             )
 
+    @staticmethod
+    def _recording_guardrail() -> "tuple[CustomGuardrail, List[List[str]]]":
+        seen: List[List[str]] = []
+
+        class Recorder(CustomGuardrail):
+            async def apply_guardrail(
+                self,
+                inputs: GenericGuardrailAPIInputs,
+                request_data: dict,
+                input_type: Literal["request", "response"],
+                logging_obj: Optional[Any] = None,
+            ) -> GenericGuardrailAPIInputs:
+                seen.append(list(inputs.get("texts", [])))
+                return inputs
+
+        return Recorder(guardrail_name="recorder"), seen
+
+    @pytest.mark.asyncio
+    async def test_fallback_delta_only_rewrite_with_delivery_expected_fails_closed(self):
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+        handler = OpenAIResponsesHandler()
+        events = [
+            {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
+            {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"},
+        ]
+
+        with pytest.raises(UndeliverableStreamRewrite):
+            await handler.process_output_streaming_response(
+                responses_so_far=events,
+                guardrail_to_apply=self._masking_guardrail(),
+                litellm_logging_obj=None,
+                deliver_ended_stream_rewrites=True,
+            )
+
+    @pytest.mark.asyncio
+    async def test_fallback_scans_delta_text_when_delivery_expected(self):
+        handler = OpenAIResponsesHandler()
+        guardrail, seen = self._recording_guardrail()
+        events = [
+            {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
+            {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "there"},
+        ]
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=events,
+            guardrail_to_apply=guardrail,
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert result is events
+        assert seen == [["hello there"]]
+
+    @pytest.mark.asyncio
+    async def test_fallback_ignores_delta_text_without_delivery_expected(self):
+        handler = OpenAIResponsesHandler()
+        guardrail, seen = self._recording_guardrail()
+        events = [
+            {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello world"},
+        ]
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=events,
+            guardrail_to_apply=guardrail,
+            litellm_logging_obj=None,
+        )
+
+        assert result is events
+        assert seen == []
+
     @pytest.mark.asyncio
     async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self):
         handler = OpenAIResponsesHandler()