Merge pull request #40986 from BerriAI/litellm_lit_7346_multi_choice_stream_guardrails

fix(guardrails): scan each choice's tool-call arguments apart on n>1 streams and log why a rewrite was discarded

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -138,9 +138,13 @@ class _ToolCallDelta(TypedDict, total=False):
 
 
 class _ToolCallChoice(TypedDict, total=False):
+    index: ReadOnly[int]
     delta: ReadOnly[_ToolCallDelta]
 
 
+_ToolCallKey: TypeAlias = tuple[int, int]
+
+
 class _ToolCallChunk(TypedDict):
     choices: ReadOnly[Sequence[_ToolCallChoice]]
 
@@ -417,40 +421,41 @@ class ChunkProcessor:
     @staticmethod
     def _iter_tool_call_fragments(
         tool_call_chunks: Sequence["_ToolCallChunk"],
-    ) -> Iterator[tuple[int, str, str]]:
+    ) -> Iterator[tuple[_ToolCallKey, str, str]]:
         for chunk in tool_call_chunks:
             for choice in chunk["choices"]:
                 delta = choice.get("delta")
                 if not delta:
                     continue
-                for tool_call in delta.get("tool_calls", ()):
+                choice_index = choice.get("index", 0)
+                for tool_call in delta.get("tool_calls") or ():
                     if not tool_call:
                         continue
                     if isinstance(tool_call, dict):
-                        index = tool_call.get("index", 0)
+                        key = (choice_index, tool_call.get("index", 0))
                         function = tool_call.get("function")
                         if isinstance(function, dict):
                             if fragment_arguments := function.get("arguments"):
-                                yield index, "arguments", fragment_arguments
+                                yield key, "arguments", fragment_arguments
                         elif function_arguments := getattr(function, "arguments", None):
-                            yield index, "arguments", function_arguments
+                            yield key, "arguments", function_arguments
                         custom = tool_call.get("custom")
                         if isinstance(custom, dict) and (custom_input := custom.get("input")):
-                            yield index, "custom_input", custom_input
+                            yield key, "custom_input", custom_input
                     else:
-                        index = getattr(tool_call, "index", 0)
+                        key = (choice_index, getattr(tool_call, "index", 0))
                         function = getattr(tool_call, "function", None)
                         if object_arguments := getattr(function, "arguments", None):
-                            yield index, "arguments", object_arguments
+                            yield key, "arguments", object_arguments
                         custom = getattr(tool_call, "custom", None)
                         if object_custom_input := getattr(custom, "input", None):
-                            yield index, "custom_input", object_custom_input
+                            yield key, "custom_input", object_custom_input
 
     @staticmethod
-    def _join_fragments_by_index_and_field(
-        fragment_records: Iterator[tuple[int, str, str]],
-    ) -> Mapping[tuple[int, str], str]:
-        def group_key(record: tuple[int, str, str]) -> tuple[int, str]:
+    def _join_fragments_by_key_and_field(
+        fragment_records: Iterator[tuple[_ToolCallKey, str, str]],
+    ) -> Mapping[tuple[_ToolCallKey, str], str]:
+        def group_key(record: tuple[_ToolCallKey, str, str]) -> tuple[_ToolCallKey, str]:
             return record[0], record[1]
 
         return MappingProxyType(
@@ -468,13 +473,14 @@ class ChunkProcessor:
         tool_calls_list: list[
             ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall
         ] = []  # mutable-ok: see return type
-        tool_call_map: Final[dict[int, dict[str, Any]]] = {}  # Map to store tool calls by index
+        tool_call_map: Final[dict[_ToolCallKey, dict[str, Any]]] = {}
 
         for chunk in tool_call_chunks:
             choices = chunk["choices"]
             for choice in choices:
                 delta = choice.get("delta", {})
-                tool_calls = delta.get("tool_calls", [])
+                tool_calls = delta.get("tool_calls") or ()
+                choice_index = choice.get("index", 0)
 
                 for tool_call in tool_calls:
                     # Handle both dict and object formats
@@ -496,9 +502,9 @@ class ChunkProcessor:
 
                     # Get index (handle both dict and object)
                     if isinstance(tool_call, dict):
-                        index = tool_call.get("index", 0)
+                        index = (choice_index, tool_call.get("index", 0))
                     else:
-                        index = getattr(tool_call, "index", 0)
+                        index = (choice_index, getattr(tool_call, "index", 0))
 
                     if index not in tool_call_map:
                         tool_call_map[index] = {
@@ -573,7 +579,7 @@ class ChunkProcessor:
                         if isinstance(provider_fields, dict):
                             merged_provider_fields.update(provider_fields)
 
-        joined_fragments: Final = self._join_fragments_by_index_and_field(
+        joined_fragments: Final = self._join_fragments_by_key_and_field(
             self._iter_tool_call_fragments(tool_call_chunks)
         )
 

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

@@ -1457,7 +1457,9 @@ class AnthropicMessagesHandler(BaseTranslation):
         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)
+            raise UndeliverableStreamRewrite(
+                guardrail_name, "the buffered stream carries no text_delta event to land the text rewrite on"
+            )
         replacements: Final = chain((rewritten_text,), repeat(""))
 
         def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
@@ -1498,7 +1500,11 @@ class AnthropicMessagesHandler(BaseTranslation):
         if len(block_indices) != len(post_guardrail_tool_calls):
             from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
-            raise UndeliverableStreamRewrite(guardrail_name)
+            raise UndeliverableStreamRewrite(
+                guardrail_name,
+                f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried "
+                f"{len(block_indices)} tool_use blocks",
+            )
         rewrites_by_block: Final = MappingProxyType(
             {
                 index: after

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

@@ -1169,10 +1169,22 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             choice.index for response in responses_so_far for choice in response.choices
         )
         fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far)
-        if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls):
+        if len(stream_choice_indices) != 1:
             from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
-            raise UndeliverableStreamRewrite(guardrail_name)
+            raise UndeliverableStreamRewrite(
+                guardrail_name,
+                f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written "
+                "back on single-choice streams",
+            )
+        if len(fragments_by_tool_call) != len(post_guardrail_tool_calls):
+            from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+            raise UndeliverableStreamRewrite(
+                guardrail_name,
+                f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried "
+                f"{len(fragments_by_tool_call)}",
+            )
         for before, (name, arguments), fragments in zip(
             pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call
         ):

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

@@ -167,6 +167,34 @@ def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCa
     return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments)
 
 
+def _undeliverable_tool_call_rewrite_reason(
+    call_ids: Sequence[str],
+    tool_call_item_count: int,
+    post_guardrail_tool_call_count: int,
+    unresolved_argument_event: bool,
+    rewritten_call_ids: frozenset[str],
+    event_call_ids: frozenset[str],
+) -> str | None:
+    if len(call_ids) != tool_call_item_count:
+        return (
+            f"{tool_call_item_count - len(call_ids)} of the stream's {tool_call_item_count} tool call items "
+            "carry no call_id"
+        )
+    if len(frozenset(call_ids)) != len(call_ids):
+        return "the stream's tool call items repeat a call_id"
+    if len(call_ids) != post_guardrail_tool_call_count:
+        return (
+            f"the guardrail returned {post_guardrail_tool_call_count} tool calls for the stream's "
+            f"{len(call_ids)} tool call items"
+        )
+    if unresolved_argument_event:
+        return "a tool call argument event names an item_id that no output_item event introduced"
+    missing_call_ids: Final = sorted(rewritten_call_ids - event_call_ids)
+    if missing_call_ids:
+        return f"no stream event carries the rewritten call_id {', '.join(missing_call_ids)}"
+    return None
+
+
 class ResponseOutputEnvelope(TypedDict, total=False):
     """Dict form of a Responses API response, as far as guardrail write-back reads it."""
 
@@ -999,7 +1027,11 @@ class OpenAIResponsesHandler(BaseTranslation):
         ):
             from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
-            raise UndeliverableStreamRewrite(guardrail_name)
+            raise UndeliverableStreamRewrite(
+                guardrail_name,
+                "the scanned text events are not all output_text deltas with an integer output_index and "
+                "content_index, so the text rewrite has nowhere to land",
+            )
         self._sync_stream_events_with_rewrites(
             stream_events=stream_events,
             rewrites_by_position=MappingProxyType(dict(zip(placeable_positions, chain((rewritten_text,), repeat(""))))),
@@ -1106,16 +1138,18 @@ class OpenAIResponsesHandler(BaseTranslation):
             call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
             for event, call_id in zip(stream_events, event_call_ids)
         )
-        if (
-            len(call_ids) != len(tool_call_items)
-            or len(frozenset(call_ids)) != len(call_ids)
-            or len(call_ids) != len(post_guardrail_tool_calls)
-            or unresolved_argument_event
-            or not rewrites_by_call_id.keys() <= frozenset(event_call_ids)
-        ):
+        undeliverable_reason: Final = _undeliverable_tool_call_rewrite_reason(
+            call_ids=call_ids,
+            tool_call_item_count=len(tool_call_items),
+            post_guardrail_tool_call_count=len(post_guardrail_tool_calls),
+            unresolved_argument_event=unresolved_argument_event,
+            rewritten_call_ids=frozenset(rewrites_by_call_id),
+            event_call_ids=frozenset(call_id for call_id in event_call_ids if call_id is not None),
+        )
+        if undeliverable_reason is not None:
             from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
-            raise UndeliverableStreamRewrite(guardrail_name)
+            raise UndeliverableStreamRewrite(guardrail_name, undeliverable_reason)
         for output_item, rewrite in (
             (output_item, rewrites_by_call_id[call_id])
             for output_item, call_id in zip(tool_call_items, call_ids)

litellm/main.py

@@ -8759,6 +8759,39 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o
         setattr(usage, "cost", computed_cost)
 
 
+_NON_TEXT_DELTA_FIELDS: Final = (
+    "tool_calls",
+    "function_call",
+    "reasoning_content",
+    "thinking_blocks",
+    "annotations",
+    "audio",
+    "images",
+    "provider_specific_fields",
+)
+
+
+def _stream_choice_delta(choice: object) -> Mapping[str, object]:
+    delta: Final = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {})
+    if isinstance(delta, Mapping):
+        return delta
+    if isinstance(delta, BaseModel):
+        return delta.model_dump()
+    return {}
+
+
+def _delta_carries_more_than_text(delta: Mapping[str, object]) -> bool:
+    return any(delta.get(field) is not None for field in _NON_TEXT_DELTA_FIELDS)
+
+
+def _simple_text_part(choices: Sequence[object]) -> str | None:
+    deltas: Final = tuple(_stream_choice_delta(choice) for choice in choices)
+    if any(_delta_carries_more_than_text(delta) for delta in deltas):
+        return None
+    content: Final = deltas[0].get("content")
+    return content if isinstance(content, str) else ""
+
+
 def stream_chunk_builder(
     chunks: list,
     messages: Sequence | None = None,
@@ -8803,31 +8836,11 @@ def stream_chunk_builder(
             if not chunk.get("choices"):
                 continue
 
-            choice = chunk["choices"][0]
-            delta_obj = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {})
-            if isinstance(delta_obj, dict):
-                delta = delta_obj
-            elif hasattr(delta_obj, "model_dump"):
-                delta = cast(dict[str, Any], delta_obj.model_dump())
-            else:
-                delta = {}
-
-            if (
-                delta.get("tool_calls") is not None
-                or delta.get("function_call") is not None
-                or delta.get("reasoning_content") is not None
-                or delta.get("thinking_blocks") is not None
-                or delta.get("annotations") is not None
-                or delta.get("audio") is not None
-                or delta.get("images") is not None
-                or delta.get("provider_specific_fields") is not None
-            ):
+            if (part := _simple_text_part(chunk["choices"])) is None:
                 is_simple_text_stream = False
                 break
-
-            content = delta.get("content")
-            if isinstance(content, str) and content:
-                simple_content_parts.append(content)
+            if part:
+                simple_content_parts.append(part)
 
         if is_simple_text_stream:
             if simple_content_parts:
@@ -8864,9 +8877,10 @@ def stream_chunk_builder(
         tool_call_chunks: Final = [
             chunk
             for chunk in chunks
-            if chunk.get("choices")
-            and "tool_calls" in chunk["choices"][0]["delta"]
-            and chunk["choices"][0]["delta"]["tool_calls"] is not None
+            if any(
+                "tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] is not None
+                for choice in chunk.get("choices") or ()
+            )
         ]
 
         if len(tool_call_chunks) > 0:

litellm/proxy/policy_engine/pipeline_executor.py

@@ -50,12 +50,16 @@ except ImportError:
 
 
 class UndeliverableStreamRewrite(Exception):
-    def __init__(self, guardrail_name: str) -> None:
-        super().__init__(
-            f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's "
-            "streaming pipeline cannot deliver"
-        )
+    def __init__(self, guardrail_name: str, reason: str) -> None:
+        super().__init__(guardrail_name, reason)
         self.guardrail_name: Final = guardrail_name
+        self.reason: Final = reason
+
+    def __str__(self) -> str:
+        return (
+            f"Guardrail '{self.guardrail_name}' rewrote the streamed response but the rewrite cannot be written "
+            f"back to the stream: {self.reason}"
+        )
 
 
 def _tool_call_shape(tool_call: object) -> tuple[object, object]:
@@ -82,8 +86,22 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non
     return sent is not None and returned is not None and returned != sent
 
 
-def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool:
-    return sent is not None and returned is not None and len(returned) != len(sent)
+def _count_change(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> tuple[int, int] | None:
+    if sent is None or returned is None or len(returned) == len(sent):
+        return None
+    return (len(sent), len(returned))
+
+
+def _tool_call_mismatch_reason(
+    sent: tuple[tuple[object, object], ...] | None, returned: tuple[tuple[object, object], ...] | None
+) -> str | None:
+    if sent == returned:
+        return None
+    sent_count: Final = len(sent or ())
+    returned_count: Final = len(returned or ())
+    if sent_count == returned_count:
+        return "the legacy hook changed a tool call's name or arguments, which this path cannot write back"
+    return f"the legacy hook returned {returned_count} tool calls for a stream that carried {sent_count}"
 
 
 _GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object])
@@ -110,7 +128,7 @@ class _StreamRewriteObserver(CustomGuardrail):
         self.inner: Final = inner
         self.rewrote_texts = False
         self.rewrote_tool_calls = False
-        self.changed_tool_call_count = False
+        self.tool_call_count_change: tuple[int, int] | None = None
 
     def structured_messages_cover_full_request(self) -> bool:
         return self.inner.structured_messages_cover_full_request()
@@ -131,11 +149,22 @@ class _StreamRewriteObserver(CustomGuardrail):
         returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls"))
         self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts")))
         self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes)
-        self.changed_tool_call_count = self.changed_tool_call_count or _changed_count(
+        self.tool_call_count_change = self.tool_call_count_change or _count_change(
             sent_tool_shapes, returned_tool_shapes
         )
         return outputs
 
+    def discard_reason(self, deliver_rewrites: bool) -> str | None:
+        if self.tool_call_count_change is not None:
+            sent, returned = self.tool_call_count_change
+            return (
+                f"the guardrail returned {returned} tool calls for a stream that carried {sent}, and a rewrite "
+                "that drops or adds a tool call cannot be written back"
+            )
+        if not deliver_rewrites and (self.rewrote_texts or self.rewrote_tool_calls):
+            return "this endpoint's streaming pipeline does not write ended-stream rewrites back yet"
+        return None
+
 
 class _ScannedTextRecorder(CustomGuardrail):
     def __init__(self, guardrail_name: str) -> None:
@@ -200,13 +229,24 @@ class _LegacyHookStreamAdapter(CustomGuardrail):
         if rewrite is None:
             return inputs
         rescanned: Final = await self._rescan(rewrite, logging_obj)
+        guardrail_name: Final = self.guardrail_name or "unknown"
         if rescanned is None:
-            raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
+            raise UndeliverableStreamRewrite(
+                guardrail_name, "the legacy hook's response could not be rescanned by this endpoint's translation"
+            )
         rewritten: Final = rescanned.get("texts")
-        if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))):
-            raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
-        if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")):
-            raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
+        returned_text_count: Final = len(_scanned_texts(rewritten))
+        sent_text_count: Final = len(_scanned_texts(inputs.get("texts")))
+        if returned_text_count != sent_text_count:
+            raise UndeliverableStreamRewrite(
+                guardrail_name,
+                f"the legacy hook returned {returned_text_count} texts for a stream that carried {sent_text_count}",
+            )
+        tool_call_mismatch: Final = _tool_call_mismatch_reason(
+            _tool_call_shapes(inputs.get("tool_calls")), _tool_call_shapes(rescanned.get("tool_calls"))
+        )
+        if tool_call_mismatch is not None:
+            raise UndeliverableStreamRewrite(guardrail_name, tool_call_mismatch)
         if not rewritten:
             return inputs
         rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten}
@@ -253,14 +293,16 @@ def _prepare_hook_input(
 
 def _release_original_chunks(
     guardrail_name: str,
+    reason: str,
     streaming_chunks: list[object],  # mutable-ok: shared buffered-stream chunks, restored in place
     originals: Sequence[object],
 ) -> None:
     streaming_chunks[:] = originals  # rebind-ok: the caller's buffer is the stream the client receives
     verbose_proxy_logger.warning(
-        "Pipeline: guardrail '%s' rewrote the streamed response in a way this endpoint's streaming "
-        "pipeline cannot deliver yet; the rewrite was discarded and the original stream released",
+        "Pipeline: guardrail '%s' rewrote the streamed response but the rewrite could not be written back to "
+        "the stream: %s. The whole rewrite, text rewrites included, was discarded and the original stream released",
         guardrail_name,
+        reason,
     )
 
 
@@ -433,13 +475,12 @@ class PipelineExecutor:
                     user_api_key_dict=user_api_key_dict,
                     request_data=hook_input,
                 )
-        except UndeliverableStreamRewrite:
-            _release_original_chunks(step.guardrail, streaming_chunks, originals)
+        except UndeliverableStreamRewrite as undeliverable:
+            _release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals)
             return
-        if observer.changed_tool_call_count or (
-            not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls)
-        ):
-            _release_original_chunks(step.guardrail, streaming_chunks, originals)
+        discard_reason: Final = observer.discard_reason(deliver_rewrites)
+        if discard_reason is not None:
+            _release_original_chunks(step.guardrail, discard_reason, streaming_chunks, originals)
             return
         if not callback.records_own_guardrail_information:
             add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail)

tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py

@@ -1278,6 +1278,61 @@ def _tool_call_delta_chunk(tool_call: dict[str, object] | ChatCompletionDeltaToo
     return {"choices": [{"delta": {"tool_calls": [tool_call]}}]}
 
 
+def _choice_tool_call_delta_chunk(choice_index: int, tool_call: dict[str, object]) -> dict[str, object]:
+    return {"choices": [{"index": choice_index, "delta": {"tool_calls": [tool_call]}}]}
+
+
+def test_get_combined_tool_content_keeps_each_choices_arguments_apart_when_choices_share_a_tool_index():
+    processor = ChunkProcessor.__new__(ChunkProcessor)
+    chunks = [
+        _choice_tool_call_delta_chunk(0, {"index": 0, "id": "call_a", "type": "function", "function": {"name": "f"}}),
+        _choice_tool_call_delta_chunk(1, {"index": 0, "id": "call_b", "type": "function", "function": {"name": "f"}}),
+        _choice_tool_call_delta_chunk(0, {"index": 0, "function": {"arguments": '{"fruit": "pers'}}),
+        _choice_tool_call_delta_chunk(1, {"index": 0, "function": {"arguments": '{"fruit": "dur'}}),
+        _choice_tool_call_delta_chunk(0, {"index": 0, "function": {"arguments": 'immon"}'}}),
+        _choice_tool_call_delta_chunk(1, {"index": 0, "function": {"arguments": 'ian"}'}}),
+    ]
+
+    combined = processor.get_combined_tool_content(chunks)
+
+    assert [(tool_call.id, tool_call.function.arguments) for tool_call in combined] == [
+        ("call_a", '{"fruit": "persimmon"}'),
+        ("call_b", '{"fruit": "durian"}'),
+    ]
+
+
+def test_stream_chunk_builder_keeps_each_choices_tool_call_arguments_apart():
+    def chunk(choice_index: int, tool_call: ChatCompletionDeltaToolCall) -> ModelResponseStream:
+        return ModelResponseStream(
+            id="chatcmpl-123",
+            object="chat.completion.chunk",
+            created=1234567890,
+            model="gpt-4.1-mini",
+            choices=[StreamingChoices(index=choice_index, delta=Delta(tool_calls=[tool_call]), finish_reason=None)],
+        )
+
+    def fragment(arguments: str, name: str | None = None, call_id: str | None = None) -> ChatCompletionDeltaToolCall:
+        return ChatCompletionDeltaToolCall(
+            id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments)
+        )
+
+    response = stream_chunk_builder(
+        chunks=[
+            chunk(0, fragment("", name="lookup_fruit", call_id="call_a")),
+            chunk(1, fragment("", name="lookup_fruit", call_id="call_b")),
+            chunk(0, fragment('{"fruit": "pers')),
+            chunk(1, fragment('{"fruit": "dur')),
+            chunk(0, fragment('immon"}')),
+            chunk(1, fragment('ian"}')),
+        ]
+    )
+
+    assert [(tool_call.id, tool_call.function.arguments) for tool_call in response.choices[0].message.tool_calls] == [
+        ("call_a", '{"fruit": "persimmon"}'),
+        ("call_b", '{"fruit": "durian"}'),
+    ]
+
+
 def test_get_combined_tool_content_joins_many_dict_shaped_argument_fragments_in_order():
     processor = ChunkProcessor.__new__(ChunkProcessor)
     first_fragments = [f"a{i};" for i in range(300)]

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

@@ -6,6 +6,7 @@ with guardrail transformations, including tool calls.
 """
 
 import json
+from collections.abc import Mapping
 from typing import Any, Literal, Optional
 
 import pytest
@@ -1372,12 +1373,51 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
         return [
             chunk(0, fragment("", name="lookup_fruit", call_id="call_1")),
             chunk(1, fragment("", name="lookup_fruit", call_id="call_2")),
-            chunk(0, fragment('{"fruit": "persimmon"}')),
-            chunk(1, fragment('{"fruit": "durian"}')),
+            chunk(0, fragment('{"fruit": "pers')),
+            chunk(1, fragment('{"fruit": "dur')),
+            chunk(0, fragment('immon"}')),
+            chunk(1, fragment('ian"}')),
             chunk(0, None, finish_reason="tool_calls"),
             chunk(1, None, finish_reason="tool_calls"),
         ]
 
+    @staticmethod
+    def _recording_guardrail() -> CustomGuardrail:
+        class Recorder(CustomGuardrail):
+            def __init__(self) -> None:
+                super().__init__(guardrail_name="recorder")
+                self.seen_inputs: list[GenericGuardrailAPIInputs] = []
+
+            async def apply_guardrail(
+                self,
+                inputs: GenericGuardrailAPIInputs,
+                request_data: Mapping[str, object],
+                input_type: Literal["request", "response"],
+                logging_obj: object = None,
+            ) -> GenericGuardrailAPIInputs:
+                self.seen_inputs.append(inputs)
+                return inputs
+
+        return Recorder()
+
+    @pytest.mark.asyncio
+    async def test_ended_multi_choice_stream_scans_each_choices_tool_call_arguments_apart(self):
+        handler = OpenAIChatCompletionsHandler()
+        chunks = self._two_choice_tool_call_stream_chunks()
+        guardrail = self._recording_guardrail()
+
+        await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=guardrail,
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert [
+            (tool_call["id"], tool_call["function"]["arguments"])
+            for tool_call in guardrail.seen_inputs[-1]["tool_calls"]
+        ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')]
+
     @pytest.mark.asyncio
     async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self):
         from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
@@ -1385,7 +1425,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
         handler = OpenAIChatCompletionsHandler()
         chunks = self._two_choice_tool_call_stream_chunks()
 
-        with pytest.raises(UndeliverableStreamRewrite):
+        with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised:
             await handler.process_output_streaming_response(
                 responses_so_far=chunks,
                 guardrail_to_apply=MockGuardrail(guardrail_name="test"),
@@ -1393,6 +1433,11 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
                 deliver_ended_stream_rewrites=True,
             )
 
+        assert raised.value.guardrail_name == "test"
+        assert raised.value.reason == (
+            "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams"
+        )
+
     @pytest.mark.asyncio
     async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):
         handler = OpenAIChatCompletionsHandler()

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

@@ -1610,13 +1610,14 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
         events = self._ended_custom_tool_call_stream_events()
         events[5]["response"]["output"] = [{**events[5]["response"]["output"][0], "call_id": "call_999"}]
 
-        with pytest.raises(UndeliverableStreamRewrite):
+        with pytest.raises(UndeliverableStreamRewrite) as undeliverable:
             await handler.process_output_streaming_response(
                 responses_so_far=events,
                 guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
                 litellm_logging_obj=None,
                 deliver_ended_stream_rewrites=True,
             )
+        assert undeliverable.value.reason == "no stream event carries the rewritten call_id call_999"
 
     @staticmethod
     def _bridged_function_call_stream_events() -> List[dict]:
@@ -1688,8 +1689,21 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
         assert events[1]["item"] == {"type": "reasoning", "id": "rs_1", "summary": []}
 
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("mismatch", ["orphan_call_id", "duplicate_call_id"])
-    async def test_deliver_ended_stream_function_call_rewrite_without_matching_events_fails_closed(self, mismatch):
+    @pytest.mark.parametrize(
+        ("mismatch", "expected_reason"),
+        [
+            ("orphan_call_id", "no stream event carries the rewritten call_id call_999"),
+            ("duplicate_call_id", "the stream's tool call items repeat a call_id"),
+            ("missing_call_id", "1 of the stream's 1 tool call items carry no call_id"),
+            (
+                "unknown_argument_item_id",
+                "a tool call argument event names an item_id that no output_item event introduced",
+            ),
+        ],
+    )
+    async def test_deliver_ended_stream_function_call_rewrite_without_matching_events_fails_closed(
+        self, mismatch, expected_reason
+    ):
         from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
         handler = OpenAIResponsesHandler()
@@ -1697,16 +1711,22 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
         envelope_item = events[5]["response"]["output"][0]
         if mismatch == "orphan_call_id":
             events[5]["response"]["output"] = [{**envelope_item, "call_id": "call_999"}]
-        else:
+        elif mismatch == "duplicate_call_id":
             events[5]["response"]["output"] = [dict(envelope_item), dict(envelope_item)]
+        elif mismatch == "missing_call_id":
+            events[5]["response"]["output"] = [{key: value for key, value in envelope_item.items() if key != "call_id"}]
+        else:
+            events[1]["item_id"] = "fc_unknown"
 
-        with pytest.raises(UndeliverableStreamRewrite):
+        with pytest.raises(UndeliverableStreamRewrite) as undeliverable:
             await handler.process_output_streaming_response(
                 responses_so_far=events,
                 guardrail_to_apply=self._argument_masking_guardrail(),
                 litellm_logging_obj=None,
                 deliver_ended_stream_rewrites=True,
             )
+        assert undeliverable.value.reason == expected_reason
+        assert str(undeliverable.value).endswith(f"cannot be written back to the stream: {expected_reason}")
 
     @pytest.mark.asyncio
     async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self):

tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py

@@ -6,6 +6,7 @@ Uses mock guardrails to validate pipeline execution without external services.
 
 import copy
 import logging
+import pickle
 from typing import Literal
 from unittest.mock import MagicMock
 
@@ -1122,7 +1123,7 @@ class _RefusingTranslation:
         deliver_ended_stream_rewrites=False,
     ):
         responses_so_far[0]["text"] = "half-written"
-        raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name)
+        raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name, "the translation refused it")
 
 
 def _chunk():
@@ -1143,10 +1144,22 @@ async def _run_streaming_step(translation, streaming_chunks=None):
     )
 
 
-def _assert_passed_with_discard_warning(result, caplog):
+NO_WRITE_BACK_REASON = "this endpoint's streaming pipeline does not write ended-stream rewrites back yet"
+
+
+def _assert_passed_with_discard_warning(result, caplog, reason):
     assert result.terminal_action == "allow"
     assert [step.outcome for step in result.step_results] == ["pass"]
-    assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records)
+    discard_warnings = [
+        record.getMessage()
+        for record in caplog.records
+        if record.levelno == logging.WARNING
+        and "'masker'" in record.getMessage()
+        and "discarded" in record.getMessage()
+    ]
+    assert len(discard_warnings) == 1
+    assert reason in discard_warnings[0]
+    assert "text rewrites included" in discard_warnings[0]
     assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", [])
 
 
@@ -1159,7 +1172,7 @@ async def test_streaming_step_discards_text_rewrite_when_translation_lacks_write
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         result = await _run_streaming_step(translation, chunks)
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON)
     assert chunks == [_chunk()]
     assert translation.seen_guardrail_names == ["masker"]
 
@@ -1196,7 +1209,7 @@ async def test_streaming_step_in_place_rewrite_is_discarded_without_write_back(m
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         result = await _run_streaming_step(_TextTranslation(), chunks)
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON)
     assert chunks == [_chunk()]
 
 
@@ -1258,7 +1271,9 @@ async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         result = await _run_streaming_step(_WritingTranslation(), chunks)
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(
+        result, caplog, "the guardrail returned 0 tool calls for a stream that carried 1"
+    )
     assert chunks == [_chunk()]
 
 
@@ -1270,7 +1285,7 @@ async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         result = await _run_streaming_step(_TextTranslation(), chunks)
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON)
     assert chunks == [_chunk()]
 
 
@@ -1362,7 +1377,7 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         result = await _run_streaming_step(_RefusingTranslation(), chunks)
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(result, caplog, "the translation refused it")
     assert chunks == [_chunk()]
 
 
@@ -1597,7 +1612,7 @@ async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(result, caplog, "the legacy hook returned 2 texts for a stream that carried 1")
     assert chunks == [_chunk()]
 
 
@@ -1612,7 +1627,7 @@ async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(m
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments")
     assert chunks == [_chunk()]
 
 
@@ -1624,7 +1639,9 @@ async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(
+        result, caplog, "the legacy hook returned 0 tool calls for a stream that carried 1"
+    )
     assert chunks == [_chunk()]
 
 
@@ -1656,7 +1673,7 @@ async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only
             monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation()
         )
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments")
     assert chunks == [_tool_only_chunk()]
 
 
@@ -1694,7 +1711,7 @@ async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_r
 
     result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation())
 
-    _assert_passed_with_discard_warning(result, caplog)
+    _assert_passed_with_discard_warning(result, caplog, "the legacy hook's response could not be rescanned")
     assert chunks == [_chunk()]
 
 
@@ -1711,3 +1728,15 @@ async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(mon
     assert chunks[0]["text"] == "[REWRITTEN] hello world"
     assert [call["response"] for call in masker.calls] == [_native("hello world")]
     assert [call["response"] for call in auditor.calls] == [_native("[REWRITTEN] hello world")]
+
+
+@pytest.mark.parametrize("clone", [copy.deepcopy, lambda exc: pickle.loads(pickle.dumps(exc))], ids=["deepcopy", "pickle"])
+def test_undeliverable_stream_rewrite_keeps_its_reason_through_a_copy(clone):
+    original = UndeliverableStreamRewrite("masker", "the translation refused it")
+
+    copied = clone(original)
+
+    assert copied.guardrail_name == "masker"
+    assert copied.reason == "the translation refused it"
+    assert str(copied) == str(original)
+    assert str(copied).endswith("cannot be written back to the stream: the translation refused it")

tests/test_litellm/test_main.py

@@ -1698,6 +1698,68 @@ async def test_async_mock_delay():
     assert delay >= 0.01
 
 
+def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of_a_multi_choice_chunk():
+    from litellm import stream_chunk_builder
+    from litellm.types.utils import (
+        ChatCompletionDeltaToolCall,
+        Delta,
+        Function,
+        ModelResponseStream,
+        StreamingChoices,
+    )
+
+    def chunk(choices: list[StreamingChoices]) -> ModelResponseStream:
+        return ModelResponseStream(
+            id="chatcmpl-multi-choice",
+            created=1751934860,
+            model="gpt-4.1-mini",
+            object="chat.completion.chunk",
+            choices=choices,
+        )
+
+    chunks = [
+        chunk(
+            [
+                StreamingChoices(index=0, delta=Delta(role="assistant", content="hello")),
+                StreamingChoices(
+                    index=1,
+                    delta=Delta(
+                        role="assistant",
+                        tool_calls=[
+                            ChatCompletionDeltaToolCall(
+                                id="call_1",
+                                index=0,
+                                type="function",
+                                function=Function(name="lookup_fruit", arguments='{"fruit":'),
+                            )
+                        ],
+                    ),
+                ),
+            ]
+        ),
+        chunk(
+            [
+                StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop"),
+                StreamingChoices(
+                    index=1,
+                    delta=Delta(
+                        tool_calls=[ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"kiwi"}'))]
+                    ),
+                    finish_reason="tool_calls",
+                ),
+            ]
+        ),
+    ]
+
+    response = stream_chunk_builder(chunks=chunks)
+
+    tool_calls = response.choices[0].message.tool_calls
+    assert tool_calls is not None
+    assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [
+        ("call_1", "lookup_fruit", '{"fruit":"kiwi"}')
+    ]
+
+
 def test_stream_chunk_builder_thinking_blocks():
     from litellm import stream_chunk_builder
     from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices