fix(streaming): rebuild each choice of an n>1 chat stream on its own (#49002)

* fix(streaming): rebuild each choice of an n>1 chat stream on its own

stream_chunk_builder folded every choice of an n>1 chat stream into choice 0, so post-call guardrails such as tool_permission re-emitted one choice holding every choice's tool calls, and spend logs and logging callbacks kept only the first choice's text. Each choice is now assembled from its own chunks, while usage, cost, and hidden params still read the whole stream. Cisco AI Defense compares each choice's streamed text on its own so safe n>1 streams still pass through unchanged.

* fix(streaming): keep bridged reasoning summaries on choice 0 and infer an unreported choice finish reason

* fix(streaming): match MCP follow-ups and bridged summary text to the rebuilt choices

MCP auto-execute now answers only the first choice's tool calls, since its
follow-up turn continues that choice alone. The Responses bridge streams a
space before each reasoning summary part after the first, matching the
non-streamed reply. build_base_response now takes the per-choice streams
instead of having its choices replaced afterwards

* fix(streaming): space every later reasoning summary part, deny per judged call, and rebuild in one text pass

* fix(guardrails): key tool_permission denials for custom tool calls too, and type the stream rebuild's messages

* fix(mcp): continue the first choice that called a tool when an n>1 reply auto-executes MCP tools

* fix(noma): scan each choice of a multi-choice reply on its own

* fix(noma): keep masked choices when a later choice's scan fails open on a stream

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -1979,6 +1979,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
         self._tool_call_index_map: dict[int, int] = {}  # mutable-ok: per-stream accumulator state
         self._structured_output: Final = structured_output
         self._first_message_item_done: bool = False
+        self._reasoning_summary_streamed: bool = False
 
     def _handle_string_chunk(
         self, str_line: Union[str, "BaseModel"]
@@ -2021,6 +2022,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
         tool_call_index_map: dict[int, int] | None = None,  # mutable-ok: per-stream state, remapped in place
         replays_openai_history: bool = False,
         structured_output: bool = False,
+        reasoning_summary_streamed: bool = False,
     ) -> "ModelResponseStream":
         """
         Translate a Responses API streaming chunk to OpenAI chat completion streaming format.
@@ -2030,6 +2032,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
             tool_call_index_map: Per-stream output_index -> sequential tool_call index map
             replays_openai_history: Record the completed output's item sequence for the next turn
             structured_output: Only the first message item became content
+            reasoning_summary_streamed: Whether earlier chunks already streamed reasoning summary text
 
         Returns:
             ModelResponseStream: OpenAI-formatted streaming chunk
@@ -2210,13 +2213,15 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                 )
             else:
                 raise ValueError(f"Chat provider: Invalid text delta {parsed_chunk}")
+        elif event_type == "response.reasoning_summary_part.added" and reasoning_summary_streamed:
+            return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(reasoning_content=" "))])
         elif event_type == "response.reasoning_summary_text.delta":
             content_part = parsed_chunk.get("delta", None)
             if content_part:
                 return ModelResponseStream(
                     choices=[
                         StreamingChoices(
-                            index=cast(int, parsed_chunk.get("summary_index")),
+                            index=0,
                             delta=Delta(reasoning_content=content_part),
                         )
                     ]
@@ -2317,10 +2322,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                 tool_call_index_map=self._tool_call_index_map,
                 replays_openai_history=self._replays_openai_history,
                 structured_output=self._structured_output,
+                reasoning_summary_streamed=self._reasoning_summary_streamed,
             )
         )
         if self._structured_output and _is_message_item_done(chunk):
             self._first_message_item_done = True
+        if _responses_event_type(chunk) == "response.reasoning_summary_text.delta" and chunk.get("delta"):
+            self._reasoning_summary_streamed = True
         return self._with_served_service_tier(self._with_stream_scoped_id(translated))
 
     def _is_text_delta_after_structured_answer(self, chunk: Mapping[str, object]) -> bool:

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -3,7 +3,7 @@ import builtins
 import time
 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
 from dataclasses import dataclass, replace
-from itertools import chain, groupby
+from itertools import chain, groupby, takewhile
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Optional, Protocol, TypeAlias, TypedDict, TypeVar, Union, cast
 
@@ -257,7 +257,7 @@ def apply_grounding_request_counts(
 
 
 class ChunkProcessor:
-    def __init__(self, chunks: list, messages: Sequence | None = None):
+    def __init__(self, chunks: list, messages: Sequence[object] | None = None):
         self.chunks = self._sort_chunks(chunks)
         self.messages = messages
         self.first_chunk = chunks[0]
@@ -390,6 +390,30 @@ class ChunkProcessor:
             case _:
                 return "assistant"
 
+    @staticmethod
+    def _finish_reason_of_chunks(chunks: Sequence["_BaseChunk"], unreported: str = "stop") -> object:
+        finish_reason = unreported
+        for chunk in chunks:
+            if "choices" in chunk and len(chunk["choices"]) > 0:
+                chunk_finish_reason = None
+                if hasattr(chunk["choices"][0], "finish_reason"):
+                    chunk_finish_reason = chunk["choices"][0].finish_reason
+                elif "finish_reason" in chunk["choices"][0]:
+                    chunk_finish_reason = _CHOICE_FIELD_VALUE.validate_python(chunk["choices"][0]["finish_reason"])
+                if chunk_finish_reason is not None:
+                    finish_reason = chunk_finish_reason
+        return finish_reason
+
+    @staticmethod
+    def build_base_choice(index: int, chunks: Sequence["_BaseChunk"]) -> Choices:
+        return Choices(
+            index=index,
+            finish_reason=ChunkProcessor._finish_reason_of_chunks(
+                chunks, unreported="tool_calls" if _carries_tool_calls(chunks) else "stop"
+            ),
+            message={"role": ChunkProcessor._get_role_from_chunks(chunks), "content": ""},
+        )
+
     @staticmethod
     def _get_model_from_chunks(chunks: Sequence["_BaseChunk"], first_chunk_model: str) -> str:
         """
@@ -414,7 +438,11 @@ class ChunkProcessor:
             None,
         )
 
-    def build_base_response(self, chunks: Sequence["_BaseChunk"]) -> ModelResponse:
+    def build_base_response(
+        self,
+        chunks: Sequence["_BaseChunk"],
+        choice_streams: Sequence[tuple[int, Sequence["_BaseChunk"]]],
+    ) -> ModelResponse:
         chunk = self.first_chunk
         id: Final = ChunkProcessor._get_chunk_id(chunks)
         object: Final = chunk["object"]
@@ -426,16 +454,7 @@ class ChunkProcessor:
         service_tier: Final = ChunkProcessor._get_service_tier_from_chunks(chunks)
 
         role: Final = ChunkProcessor._get_role_from_chunks(chunks)
-        finish_reason = "stop"
-        for chunk in chunks:
-            if "choices" in chunk and len(chunk["choices"]) > 0:
-                chunk_finish_reason = None
-                if hasattr(chunk["choices"][0], "finish_reason"):
-                    chunk_finish_reason = chunk["choices"][0].finish_reason
-                elif "finish_reason" in chunk["choices"][0]:
-                    chunk_finish_reason = _CHOICE_FIELD_VALUE.validate_python(chunk["choices"][0]["finish_reason"])
-                if chunk_finish_reason is not None:
-                    finish_reason = chunk_finish_reason
+        finish_reason: Final = ChunkProcessor._finish_reason_of_chunks(chunks)
 
         # Initialize the response dictionary
         response = ModelResponse(
@@ -455,13 +474,17 @@ class ChunkProcessor:
                     if service_tier is not None
                     else MappingProxyType({})
                 ),
-                "choices": [
-                    {
-                        "index": 0,
-                        "message": {"role": role, "content": ""},
-                        "finish_reason": finish_reason,
-                    }
-                ],
+                "choices": (
+                    [ChunkProcessor.build_base_choice(index, choice_chunks) for index, choice_chunks in choice_streams]
+                    if len(choice_streams) > 1
+                    else [
+                        {
+                            "index": 0,
+                            "message": {"role": role, "content": ""},
+                            "finish_reason": finish_reason,
+                        }
+                    ]
+                ),
                 "usage": {
                     "prompt_tokens": 0,  # Modify as needed
                     "completion_tokens": 0,  # Modify as needed
@@ -1520,11 +1543,188 @@ def _bedrock_trace_fields(chunk: object) -> Mapping[str, object]:
     return {field: fields[field] for field in BEDROCK_TRACE_FIELDS if field in fields}
 
 
+def _stream_choice_index(choice: object) -> int:
+    index: Final = _field_of(choice, "index")
+    return index if isinstance(index, int) else 0
+
+
+def _carries_tool_calls(chunks: Sequence[_BaseChunk]) -> bool:
+    return any(
+        "tool_calls" in choice["delta"] and choice["delta"]["tool_calls"]
+        for choice in chain.from_iterable(chunk.get("choices") or () for chunk in chunks)
+    )
+
+
+def _choice_streams(chunks: Sequence[_BaseChunk]) -> tuple[tuple[int, Sequence[_BaseChunk]], ...]:
+    indices: Final = sorted(
+        frozenset(map(_stream_choice_index, chain.from_iterable(chunk.get("choices") or () for chunk in chunks)))
+    )
+    if len(indices) <= 1:
+        return ((0, chunks),)
+    return tuple((index, _narrowed_to_choice(chunks, index)) for index in indices)
+
+
+def _narrowed_to_choice(chunks: Sequence[_BaseChunk], index: int) -> tuple[_BaseChunk, ...]:
+    return tuple(
+        _BaseChunk(choices=[choice for choice in chunk.get("choices") or () if _stream_choice_index(choice) == index])
+        for chunk in chunks
+    )
+
+
+def _simple_text(chunks: Sequence[_BaseChunk]) -> str | None:
+    choices_per_chunk: Final = tuple(chunk["choices"] for chunk in chunks if chunk.get("choices"))
+    parts: Final = tuple(takewhile(lambda part: part is not None, map(_simple_text_part, choices_per_chunk)))
+    return "".join(filter(None, parts)) if len(parts) == len(choices_per_chunk) else None
+
+
+def _assembled_message_fields(processor: ChunkProcessor, chunks: Sequence) -> Iterator[tuple[str, object]]:
+    tool_call_chunks: Final = [
+        chunk
+        for chunk in chunks
+        if any(
+            "tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] is not None
+            for choice in chunk.get("choices") or ()
+        )
+    ]
+
+    if tool_call_chunks:
+        tool_calls_list: Final = processor.get_combined_tool_content(tool_call_chunks)
+        yield "content", None
+        yield "tool_calls", tool_calls_list
+
+    function_call_chunks: Final = [
+        chunk
+        for chunk in chunks
+        if chunk.get("choices")
+        and "function_call" in chunk["choices"][0]["delta"]
+        and chunk["choices"][0]["delta"]["function_call"] is not None
+    ]
+
+    if function_call_chunks:
+        yield "content", None
+        yield "function_call", processor.get_combined_function_call_content(function_call_chunks)
+
+    content_chunks: Final = [
+        chunk
+        for chunk in chunks
+        if chunk.get("choices")
+        and "content" in chunk["choices"][0]["delta"]
+        and chunk["choices"][0]["delta"]["content"] is not None
+    ]
+
+    if content_chunks:
+        yield "content", processor.get_combined_content(content_chunks)
+
+    thinking_blocks: Final = [
+        chunk
+        for chunk in chunks
+        if chunk.get("choices")
+        and "thinking_blocks" in chunk["choices"][0]["delta"]
+        and chunk["choices"][0]["delta"]["thinking_blocks"] is not None
+    ]
+
+    if thinking_blocks:
+        yield "thinking_blocks", processor.get_combined_thinking_content(thinking_blocks)
+
+    reasoning_chunks: Final = [
+        chunk
+        for chunk in chunks
+        if chunk.get("choices")
+        and "reasoning_content" in chunk["choices"][0]["delta"]
+        and chunk["choices"][0]["delta"]["reasoning_content"] is not None
+    ]
+
+    if reasoning_chunks:
+        yield "reasoning_content", processor.get_combined_reasoning_content(reasoning_chunks)
+
+    typed_chunks: Final = cast(Sequence[object], chunks)  # cast-ok: chunks are untyped upstream
+    reasoning_items: Final = tuple(chain.from_iterable(_chunk_reasoning_items(chunk) for chunk in typed_chunks))
+    if reasoning_items:
+        yield "reasoning_items", list(reasoning_items)
+
+    annotation_chunks: Final = [
+        chunk
+        for chunk in chunks
+        if chunk.get("choices")
+        and "annotations" in chunk["choices"][0]["delta"]
+        and chunk["choices"][0]["delta"]["annotations"] is not None
+    ]
+
+    message_annotations, provider_annotations = (
+        partition_annotations(
+            list(chain.from_iterable(chunk["choices"][0]["delta"]["annotations"] for chunk in annotation_chunks))
+        )
+        if annotation_chunks
+        else (None, ())
+    )
+    if message_annotations is not None:
+        yield "annotations", message_annotations
+
+    audio_chunks: Final = [
+        chunk
+        for chunk in chunks
+        if chunk.get("choices")
+        and "audio" in chunk["choices"][0]["delta"]
+        and chunk["choices"][0]["delta"]["audio"] is not None
+    ]
+
+    if audio_chunks:
+        yield "audio", processor.get_combined_audio_content(audio_chunks)
+
+    # Handle image chunks from models like gemini-2.5-flash-image
+    # See: https://github.com/BerriAI/litellm/issues/19478
+    image_chunks: Final = [
+        chunk
+        for chunk in chunks
+        if chunk.get("choices")
+        and "images" in chunk["choices"][0]["delta"]
+        and chunk["choices"][0]["delta"]["images"] is not None
+    ]
+
+    if image_chunks:
+        # Images come complete in a single chunk, collect all images from all chunks
+        all_images: Final = []
+        for chunk in image_chunks:
+            all_images.extend(chunk["choices"][0]["delta"]["images"])
+        yield "images", all_images
+
+    # Combine provider_specific_fields from streaming chunks (e.g., web_search_results, citations)
+    # See: https://github.com/BerriAI/litellm/issues/17737
+    provider_specific_chunks: Final = [
+        chunk
+        for chunk in chunks
+        if chunk.get("choices")
+        and "provider_specific_fields" in chunk["choices"][0]["delta"]
+        and chunk["choices"][0]["delta"]["provider_specific_fields"] is not None
+    ]
+
+    provider_field_dicts: Final = tuple(
+        fields
+        for chunk in provider_specific_chunks
+        for fields in (chunk["choices"][0]["delta"]["provider_specific_fields"],)
+        if isinstance(fields, dict)
+    )
+    streamed_citations: Final = tuple(
+        fields["citation"] for fields in provider_field_dicts if fields.get("citation") is not None
+    )
+    citation_fields: Final = {"citations": _joined_streamed_citations(streamed_citations)} if streamed_citations else {}
+    provider_annotation_fields: Final = {"annotations": list(provider_annotations)} if provider_annotations else {}
+    combined_provider_fields: Final = {
+        key: value
+        for fields in (provider_annotation_fields, citation_fields, *provider_field_dicts)
+        for key, value in fields.items()
+        if key != "citation"
+    }
+
+    if combined_provider_fields:
+        yield "provider_specific_fields", combined_provider_fields
+
+
 def stream_chunk_builder(
     chunks: list,
     messages: Sequence | None = None,
-    start_time=None,
-    end_time=None,
+    start_time: object = None,
+    end_time: object = None,
     logging_obj: Optional["Logging"] = None,
     count_prompt_tokens: Callable[[], int] | None = None,
 ) -> ModelResponse | TextCompletionResponse | None:
@@ -1561,26 +1761,15 @@ def stream_chunk_builder(
         usage_model: Final = usage_tokenizer_model(
             model, get_base_model_from_metadata(logging_obj.model_call_details) if logging_obj is not None else None
         )
-        # Initialize the response dictionary
-        response: Final = processor.build_base_response(chunks)
+        choice_streams: Final = _choice_streams(chunks)
+        response: Final = processor.build_base_response(chunks, choice_streams)
+        rebuilt_choices: Final = tuple(zip(response.choices, (choice_chunks for _, choice_chunks in choice_streams)))
 
-        # Fast path for the common text-only streaming case:
-        # avoid repeated multi-pass list scans over chunks.
-        simple_content_parts: Final[list[str]] = []
-        is_simple_text_stream = True
-        for chunk in chunks:
-            if not chunk.get("choices"):
-                continue
-
-            if (part := _simple_text_part(chunk["choices"])) is None:
-                is_simple_text_stream = False
-                break
-            if part:
-                simple_content_parts.append(part)
-
-        if is_simple_text_stream:
-            if simple_content_parts:
-                response["choices"][0]["message"]["content"] = "".join(simple_content_parts)
+        simple_texts: Final = tuple(_simple_text(choice_chunks) for _, choice_chunks in choice_streams)
+        if None not in simple_texts:
+            for (choice, _), simple_text in zip(rebuilt_choices, simple_texts):
+                if simple_text:
+                    choice.message.content = simple_text
             completion_output = get_content_from_model_response(response)
             usage = processor.calculate_usage(
                 chunks=chunks,
@@ -1619,165 +1808,9 @@ def stream_chunk_builder(
             processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
             return response
 
-        tool_call_chunks: Final = [
-            chunk
-            for chunk in chunks
-            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:
-            tool_calls_list: Final = processor.get_combined_tool_content(tool_call_chunks)
-            _choice = response.choices[0]
-            _choice.message.content = None
-            _choice.message.tool_calls = tool_calls_list
-
-        function_call_chunks: Final = [
-            chunk
-            for chunk in chunks
-            if chunk.get("choices")
-            and "function_call" in chunk["choices"][0]["delta"]
-            and chunk["choices"][0]["delta"]["function_call"] is not None
-        ]
-
-        if len(function_call_chunks) > 0:
-            _choice = response.choices[0]
-            _choice.message.content = None
-            _choice.message.function_call = processor.get_combined_function_call_content(function_call_chunks)
-
-        content_chunks: Final = [
-            chunk
-            for chunk in chunks
-            if chunk.get("choices")
-            and "content" in chunk["choices"][0]["delta"]
-            and chunk["choices"][0]["delta"]["content"] is not None
-        ]
-
-        if len(content_chunks) > 0:
-            response["choices"][0]["message"]["content"] = processor.get_combined_content(content_chunks)
-
-        thinking_blocks: Final = [
-            chunk
-            for chunk in chunks
-            if chunk.get("choices")
-            and "thinking_blocks" in chunk["choices"][0]["delta"]
-            and chunk["choices"][0]["delta"]["thinking_blocks"] is not None
-        ]
-
-        if len(thinking_blocks) > 0:
-            response["choices"][0]["message"]["thinking_blocks"] = processor.get_combined_thinking_content(
-                thinking_blocks
-            )
-
-        reasoning_chunks: Final = [
-            chunk
-            for chunk in chunks
-            if chunk.get("choices")
-            and "reasoning_content" in chunk["choices"][0]["delta"]
-            and chunk["choices"][0]["delta"]["reasoning_content"] is not None
-        ]
-
-        if len(reasoning_chunks) > 0:
-            response["choices"][0]["message"]["reasoning_content"] = processor.get_combined_reasoning_content(
-                reasoning_chunks
-            )
-
-        typed_chunks: Final = cast(Sequence[object], chunks)  # cast-ok: chunks are untyped upstream
-        reasoning_items: Final = tuple(chain.from_iterable(_chunk_reasoning_items(chunk) for chunk in typed_chunks))
-        if reasoning_items:
-            response["choices"][0]["message"]["reasoning_items"] = list(reasoning_items)
-
-        annotation_chunks: Final = [
-            chunk
-            for chunk in chunks
-            if chunk.get("choices")
-            and "annotations" in chunk["choices"][0]["delta"]
-            and chunk["choices"][0]["delta"]["annotations"] is not None
-        ]
-
-        if len(annotation_chunks) > 0:
-            # Merge annotations from ALL chunks — providers may spread
-            # them across multiple streaming chunks or send them only in
-            # the final chunk.
-            all_annotations: Final[list[object]] = []
-            for ac in annotation_chunks:
-                all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
-            message_annotations, provider_annotations = partition_annotations(all_annotations)
-            if message_annotations is not None:
-                response["choices"][0]["message"]["annotations"] = message_annotations
-            if provider_annotations:
-                response.choices[0].message.provider_specific_fields = {
-                    **(response.choices[0].message.provider_specific_fields or {}),
-                    "annotations": list(provider_annotations),
-                }
-
-        audio_chunks: Final = [
-            chunk
-            for chunk in chunks
-            if chunk.get("choices")
-            and "audio" in chunk["choices"][0]["delta"]
-            and chunk["choices"][0]["delta"]["audio"] is not None
-        ]
-
-        if len(audio_chunks) > 0:
-            _choice = response.choices[0]
-            _choice.message.audio = processor.get_combined_audio_content(audio_chunks)
-
-        # Handle image chunks from models like gemini-2.5-flash-image
-        # See: https://github.com/BerriAI/litellm/issues/19478
-        image_chunks: Final = [
-            chunk
-            for chunk in chunks
-            if chunk.get("choices")
-            and "images" in chunk["choices"][0]["delta"]
-            and chunk["choices"][0]["delta"]["images"] is not None
-        ]
-
-        if len(image_chunks) > 0:
-            # Images come complete in a single chunk, collect all images from all chunks
-            all_images: Final = []
-            for chunk in image_chunks:
-                all_images.extend(chunk["choices"][0]["delta"]["images"])
-            response["choices"][0]["message"]["images"] = all_images
-
-        # Combine provider_specific_fields from streaming chunks (e.g., web_search_results, citations)
-        # See: https://github.com/BerriAI/litellm/issues/17737
-        provider_specific_chunks: Final = [
-            chunk
-            for chunk in chunks
-            if chunk.get("choices")
-            and "provider_specific_fields" in chunk["choices"][0]["delta"]
-            and chunk["choices"][0]["delta"]["provider_specific_fields"] is not None
-        ]
-
-        if len(provider_specific_chunks) > 0:
-            provider_field_dicts: Final = tuple(
-                fields
-                for chunk in provider_specific_chunks
-                for fields in (chunk["choices"][0]["delta"]["provider_specific_fields"],)
-                if isinstance(fields, dict)
-            )
-            streamed_citations: Final = tuple(
-                fields["citation"] for fields in provider_field_dicts if fields.get("citation") is not None
-            )
-            citation_fields: Final = (
-                {"citations": _joined_streamed_citations(streamed_citations)} if streamed_citations else {}
-            )
-            combined_provider_fields: Final = {
-                key: value
-                for fields in (citation_fields, *provider_field_dicts)
-                for key, value in fields.items()
-                if key != "citation"
-            }
-
-            if combined_provider_fields:
-                _choice = response.choices[0]
-                _choice.message.provider_specific_fields = {
-                    **(_choice.message.provider_specific_fields or {}),
-                    **combined_provider_fields,
-                }
+        for choice, choice_chunks in rebuilt_choices:
+            for field, value in _assembled_message_fields(processor, choice_chunks):
+                choice.message[field] = value
 
         completion_output = get_content_from_model_response(response)
 

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

@@ -687,11 +687,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         responses_so_far: Sequence["ModelResponseStream"],
         litellm_logging_obj: "LiteLLMLoggingObj | None",
     ) -> "ModelResponse":
-        """``stream_chunk_builder`` folds every choice of a stream into one index-0
-        choice, so the stream is rebuilt one choice index at a time (every chunk
-        kept, its choices narrowed to that index, so usage-only chunks still
-        count) and the rebuilt choices are stitched into one response, each
-        carrying the index the stream gave it."""
         choice_indices: Final = tuple(
             sorted(frozenset(choice.index for response in responses_so_far for choice in response.choices))
         )

litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py

@@ -22,6 +22,7 @@ import os
 from collections.abc import AsyncIterator, Mapping, Sequence
 from dataclasses import dataclass, replace
 from datetime import datetime
+from itertools import chain
 from typing import TYPE_CHECKING, Any, Final, Literal
 
 import httpx
@@ -120,6 +121,11 @@ class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
     """Base-class constructor options this guardrail forwards untouched to CustomGuardrail."""
 
 
+def _streamed_choice_index(choice: object) -> int:
+    index: Final[object] = getattr(choice, "index", None)
+    return index if isinstance(index, int) else 0
+
+
 class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail):
     """
     Cisco AI Defense guardrail integration.
@@ -585,29 +591,44 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail):
 
     @classmethod
     def _extract_streaming_chunk_scan_text(cls, chunks: Sequence[object]) -> str:
+        streamed_choices: Final = tuple(
+            choice
+            for choice in chain.from_iterable(getattr(chunk, "choices", None) or () for chunk in chunks)
+            if getattr(choice, "delta", None) is not None
+        )
+        choice_indices: Final = sorted(frozenset(map(_streamed_choice_index, streamed_choices)))
+        choice_texts: Final = (
+            cls._extract_streamed_deltas_scan_text(
+                tuple(
+                    getattr(choice, "delta", None)
+                    for choice in streamed_choices
+                    if _streamed_choice_index(choice) == index
+                )
+            )
+            for index in choice_indices
+        )
+        return " ".join(text for text in choice_texts if text)
+
+    @classmethod
+    def _extract_streamed_deltas_scan_text(cls, deltas: Sequence[object]) -> str:
         original_text = ""
         argument_text = ""
-        for chunk in chunks:
-            choices = getattr(chunk, "choices", None) or []
-            for c in choices:
-                delta: object | None = getattr(c, "delta", None)
-                if delta is None:
-                    continue
-                text = getattr(delta, "content", None)
-                if isinstance(text, str):
-                    original_text += text
-                reasoning_text = " ".join(cls._extract_message_reasoning_parts(delta))
-                if reasoning_text:
-                    original_text += reasoning_text
-                for tc in getattr(delta, "tool_calls", None) or []:
-                    args = cls._extract_tool_call_arguments(tc)
-                    if args:
-                        argument_text += args
-                fc: object | None = getattr(delta, "function_call", None)
-                if fc is not None:
-                    args = cls._extract_function_call_arguments(fc)
-                    if args:
-                        argument_text += args
+        for delta in deltas:
+            text = getattr(delta, "content", None)
+            if isinstance(text, str):
+                original_text += text
+            reasoning_text = " ".join(cls._extract_message_reasoning_parts(delta))
+            if reasoning_text:
+                original_text += reasoning_text
+            for tc in getattr(delta, "tool_calls", None) or []:
+                args = cls._extract_tool_call_arguments(tc)
+                if args:
+                    argument_text += args
+            fc: object | None = getattr(delta, "function_call", None)
+            if fc is not None:
+                args = cls._extract_function_call_arguments(fc)
+                if args:
+                    argument_text += args
         return " ".join(part for part in (original_text, argument_text) if part)
 
     # ------------------------------------------------------------------

litellm/proxy/guardrails/guardrail_hooks/noma/noma.py

@@ -246,20 +246,38 @@ class NomaGuardrail(CustomGuardrail):
     ) -> str | None:
         """Shared logic for processing LLM response checks"""
 
-        start_time: Final = datetime.now()
-        extra_data: Final = self.get_guardrail_dynamic_request_body_params(request_data)
-
         if not isinstance(response, litellm.ModelResponse):
             return None
 
-        content = None
-        for choice in response.choices:
-            if isinstance(choice, litellm.Choices) and choice.message.content:
-                content = choice.message.content
-                break
+        checked_contents: Final = tuple(
+            [
+                await self._process_choice_check(request_data, response, position, content, user_auth, event_type)
+                for position, content in self._choice_texts(response)
+            ]
+        )
+        return next(iter(checked_contents), None)
 
-        if not content or not isinstance(content, str):
-            return None
+    @staticmethod
+    def _choice_texts(response: litellm.ModelResponse) -> tuple[tuple[int, str], ...]:
+        return tuple(
+            (position, choice.message.content)
+            for position, choice in enumerate(response.choices)
+            if isinstance(choice, litellm.Choices)
+            and isinstance(choice.message.content, str)
+            and choice.message.content
+        )
+
+    async def _process_choice_check(
+        self,
+        request_data: dict[str, object],
+        response: litellm.ModelResponse,
+        position: int,
+        content: str,
+        user_auth: UserAPIKeyAuth,
+        event_type: GuardrailEventHooks | None,
+    ) -> str:
+        start_time: Final = datetime.now()
+        extra_data: Final = self.get_guardrail_dynamic_request_body_params(request_data)
 
         payload: Final = {
             "input": [
@@ -306,7 +324,7 @@ class NomaGuardrail(CustomGuardrail):
             anonymized_content: Final = self._extract_anonymized_content(response_json, ASSISTANT_ROLE)
             if anonymized_content:
                 # Replace the LLM response content with anonymized version
-                self._replace_llm_response_content(response, anonymized_content)
+                self._replace_llm_response_content(response, anonymized_content, position)
                 verbose_proxy_logger.debug("Noma guardrail anonymized LLM response: %s", anonymized_content)
                 return anonymized_content
 
@@ -486,21 +504,14 @@ class NomaGuardrail(CustomGuardrail):
                 messages[i]["content"] = anonymized_content
                 break
 
-    def _replace_llm_response_content(self, response: LLMResponse, anonymized_content: str):
-        """
-        Replace the LLM response content with anonymized version.
-
-        Args:
-            response: The original LLM response
-            anonymized_content: The anonymized content to replace with
-        """
+    def _replace_llm_response_content(self, response: LLMResponse, anonymized_content: str, position: int):
+        """Replace the content of the choice at ``position`` with its anonymized version."""
         if not isinstance(response, litellm.ModelResponse):
             return
 
-        # Replace content in all choices
-        for choice in response.choices:
-            if isinstance(choice, litellm.Choices) and choice.message.content:
-                choice.message.content = anonymized_content
+        choice: Final = response.choices[position]
+        if isinstance(choice, litellm.Choices):
+            choice.message.content = anonymized_content
 
     async def _check_user_message_background(
         self,
@@ -815,6 +826,7 @@ class NomaGuardrail(CustomGuardrail):
         )
 
         if isinstance(assembled_model_response, ModelResponse):
+            texts_before_check: Final = self._choice_texts(assembled_model_response)
             try:
                 processed_response: Final = await self._check_llm_response(
                     request_data,
@@ -828,7 +840,11 @@ class NomaGuardrail(CustomGuardrail):
                 if self.block_failures:
                     raise
                 verbose_proxy_logger.error("Noma streaming post-call hook failed: %s", e)
-                for chunk in all_chunks:
+                if self._choice_texts(assembled_model_response) == texts_before_check:
+                    for chunk in all_chunks:
+                        yield chunk
+                    return
+                async for chunk in MockResponseIterator(model_response=assembled_model_response):
                     yield chunk
                 return
 

litellm/proxy/guardrails/guardrail_hooks/tool_permission.py

@@ -36,6 +36,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
 )
 from litellm.types.utils import (
     CallTypesLiteral,
+    ChatCompletionMessageCustomToolCall,
     ChatCompletionMessageToolCall,
     Choices,
     Function,
@@ -744,6 +745,14 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
         return ToolResult(tool_use_id=tool_call.id, content=error_message, is_error=True)
 
+    @staticmethod
+    def _judged_tool_call(
+        tool_call: ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall,
+    ) -> tuple[str | None, ...]:
+        if isinstance(tool_call, ChatCompletionMessageCustomToolCall):
+            return (tool_call.id, tool_call.type, tool_call.custom.name, tool_call.custom.input)
+        return (tool_call.id, tool_call.type, tool_call.function.name, tool_call.function.arguments)
+
     def _modify_response_with_permission_errors(
         self,
         response: ModelResponse,
@@ -761,11 +770,10 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
         verbose_proxy_logger.info("Blocking %s unauthorized tool uses", len(denied_tools))
 
-        # Create a mapping of tool_use_id to error result
-        error_results: Final = {}
-        for tool_use, error in denied_tools:
-            error_result = self._create_permission_error_result(tool_use, error)
-            error_results[tool_use.id] = error_result
+        error_results: Final = {
+            self._judged_tool_call(tool_use): self._create_permission_error_result(tool_use, error)
+            for tool_use, error in denied_tools
+        }
 
         # Modify the response content
         for choice_index, choice in enumerate(response.choices):
@@ -775,9 +783,8 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
                 # Rewrite tool_calls
                 for tool_call in choice.message.tool_calls or []:
-                    tool_call_id = tool_call.id
-                    if tool_call_id in error_results:
-                        error_result = error_results[tool_call_id]
+                    error_result = error_results.get(self._judged_tool_call(tool_call))
+                    if error_result is not None:
                         error_messages.append(error_result.content)
                     else:
                         filtered_tool_calls.append(tool_call)
@@ -788,7 +795,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
                     getattr(choice.message, "function_call", None), choice_index
                 )
                 if legacy_tool_call is not None:
-                    legacy_error_result = error_results.get(legacy_tool_call.id)
+                    legacy_error_result = error_results.get(self._judged_tool_call(legacy_tool_call))
                     if legacy_error_result is not None:
                         choice.message.function_call = None
                         error_messages.append(legacy_error_result.content)

litellm/responses/mcp/litellm_proxy_mcp_handler.py

@@ -35,10 +35,7 @@ from litellm.types.llms.openai import ToolParam as ResponsesToolParam
 from litellm.types.mcp_server.mcp_server_manager import MCPServer
 from litellm.types.utils import (
     CallTypes,
-    ChatCompletionMessageCustomToolCall,
-    ChatCompletionMessageToolCall,
     Choices,
-    Message,
     ModelResponse,
     StandardLoggingMCPToolCall,
 )
@@ -587,29 +584,18 @@ class LiteLLM_Proxy_MCP_Handler:
 
     _extract_tool_calls_from_response = extract_tool_calls_from_response
 
+    @staticmethod
+    def _continued_choice(response: ModelResponse) -> Choices | None:
+        choices: Final = tuple(choice for choice in response.choices if isinstance(choice, Choices))
+        return next((choice for choice in choices if choice.message.tool_calls), next(iter(choices), None))
+
     @staticmethod
     def extract_tool_calls_from_chat_response(response: ModelResponse) -> list[object]:
-        """Extract tool calls from a chat completion response."""
-        tool_calls: Final[list[object]] = []
-
-        try:
-            for choice in response.choices:
-                message: Message | None = getattr(choice, "message", None)
-                if message is None:
-                    continue
-                tool_call_entries: (
-                    Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None
-                ) = getattr(message, "tool_calls", None)
-                if tool_call_entries:
-                    for tool_call in tool_call_entries:
-                        if hasattr(tool_call, "model_dump"):
-                            tool_calls.append(tool_call.model_dump())
-                        else:
-                            tool_calls.append(tool_call)
-        except Exception:
-            verbose_logger.exception("Failed to extract tool calls from chat completion response")
-
-        return tool_calls
+        """Extract the tool calls of the first choice that made any, since the follow-up turn continues it alone."""
+        continued_choice: Final = LiteLLM_Proxy_MCP_Handler._continued_choice(response)
+        if continued_choice is None:
+            return []
+        return [tool_call.model_dump() for tool_call in continued_choice.message.tool_calls or ()]
 
     _extract_tool_calls_from_chat_response = extract_tool_calls_from_chat_response
 
@@ -1018,9 +1004,9 @@ class LiteLLM_Proxy_MCP_Handler:
 
         message_to_append: dict | None = None
         try:
-            first_choice: Final = response.choices[0]
-            if isinstance(first_choice, Choices) and getattr(first_choice, "message", None):
-                message_to_append = first_choice.message.model_dump(exclude_none=True)
+            continued_choice: Final = LiteLLM_Proxy_MCP_Handler._continued_choice(response)
+            if continued_choice is not None:
+                message_to_append = continued_choice.message.model_dump(exclude_none=True)
                 # Ensure tool_calls have arguments field (required by OpenAI API)
                 if message_to_append.get("tool_calls"):
                     for tool_call in message_to_append["tool_calls"]:

tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py

@@ -418,6 +418,45 @@ def test_openai_responses_chunk_parser_reasoning_summary():
     assert delta.function_call is None
 
 
+def test_bridged_reasoning_summary_parts_rebuild_into_the_one_choice_they_streamed_on():
+    from litellm.completion_extras.litellm_responses_transformation.transformation import (
+        OpenAiResponsesToChatCompletionStreamIterator,
+    )
+
+    iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True)
+    events = (
+        {"type": "response.reasoning_summary_part.added", "output_index": 0, "summary_index": 0},
+        {
+            "type": "response.reasoning_summary_text.delta",
+            "delta": "First part.",
+            "output_index": 0,
+            "summary_index": 0,
+        },
+        {"type": "response.reasoning_summary_part.done", "output_index": 0, "summary_index": 0},
+        {"type": "response.reasoning_summary_part.added", "output_index": 0, "summary_index": 1},
+        {
+            "type": "response.reasoning_summary_text.delta",
+            "delta": "Second part.",
+            "output_index": 0,
+            "summary_index": 1,
+        },
+        {"type": "response.reasoning_summary_part.added", "output_index": 1, "summary_index": 0},
+        {
+            "type": "response.reasoning_summary_text.delta",
+            "delta": "Next item.",
+            "output_index": 1,
+            "summary_index": 0,
+        },
+        {"type": "response.output_text.delta", "delta": "Answer", "output_index": 2, "content_index": 0},
+    )
+
+    rebuilt = litellm.stream_chunk_builder(chunks=[iterator.chunk_parser(event) for event in events])
+
+    assert [(choice.index, choice.message.reasoning_content, choice.message.content) for choice in rebuilt.choices] == [
+        (0, "First part. Second part. Next item.", "Answer")
+    ]
+
+
 def test_chunk_parser_string_output_text_delta_produces_text():
     from litellm.completion_extras.litellm_responses_transformation.transformation import (
         OpenAiResponsesToChatCompletionStreamIterator,

tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py

@@ -1304,12 +1304,118 @@ def test_stream_chunk_builder_keeps_each_choices_tool_call_arguments_apart():
         ]
     )
 
-    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"}'),
+    assert [
+        (choice.index, [(tool_call.id, tool_call.function.arguments) for tool_call in choice.message.tool_calls])
+        for choice in response.choices
+    ] == [
+        (0, [("call_a", '{"fruit": "persimmon"}')]),
+        (1, [("call_b", '{"fruit": "durian"}')]),
     ]
 
 
+def test_stream_chunk_builder_rebuilds_each_choice_of_a_multi_choice_text_stream_on_its_own():
+    def chunk(choice_index: int, content: str, finish_reason: str | None = None) -> ModelResponseStream:
+        return ModelResponseStream(
+            id="chatcmpl-123",
+            object="chat.completion.chunk",
+            created=1234567890,
+            model="gpt-4.1-mini",
+            choices=[StreamingChoices(index=choice_index, delta=Delta(content=content), finish_reason=finish_reason)],
+        )
+
+    response = stream_chunk_builder(
+        chunks=[
+            chunk(0, "The persim"),
+            chunk(1, "The"),
+            chunk(1, " persimmon is"),
+            chunk(0, "mon is ripe.", finish_reason="stop"),
+            chunk(1, " ripe and", finish_reason="length"),
+        ]
+    )
+
+    assert [(choice.index, choice.message.content, choice.finish_reason) for choice in response.choices] == [
+        (0, "The persimmon is ripe.", "stop"),
+        (1, "The persimmon is ripe and", "length"),
+    ]
+
+
+def test_stream_chunk_builder_gives_a_choice_without_a_finish_chunk_the_finish_reason_its_content_implies():
+    def chunk(choice_index: int, delta: Delta, finish_reason: str | None = None) -> ModelResponseStream:
+        return ModelResponseStream(
+            id="chatcmpl-123",
+            object="chat.completion.chunk",
+            created=1234567890,
+            model="gpt-4.1-mini",
+            choices=[StreamingChoices(index=choice_index, delta=delta, finish_reason=finish_reason)],
+        )
+
+    def tool_call(call_id: str) -> list[ChatCompletionDeltaToolCall]:
+        return [
+            ChatCompletionDeltaToolCall(
+                id=call_id, index=0, type="function", function=Function(name="lookup_fruit", arguments="{}")
+            )
+        ]
+
+    response = stream_chunk_builder(
+        chunks=[
+            chunk(0, Delta(tool_calls=tool_call("call_a"))),
+            chunk(1, Delta(tool_calls=tool_call("call_b"))),
+            chunk(2, Delta(content="Kiwis are fuzzy.")),
+            chunk(0, Delta(), finish_reason="tool_calls"),
+        ]
+    )
+
+    assert [(choice.index, choice.finish_reason) for choice in response.choices] == [
+        (0, "tool_calls"),
+        (1, "tool_calls"),
+        (2, "stop"),
+    ]
+
+
+def test_stream_chunk_builder_keeps_each_choices_function_call_annotations_and_images_on_that_choice():
+    def chunk(choice_index: int, delta: Delta) -> ModelResponseStream:
+        return ModelResponseStream(
+            id="chatcmpl-123",
+            object="chat.completion.chunk",
+            created=1234567890,
+            model="gpt-4.1-mini",
+            choices=[StreamingChoices(index=choice_index, delta=delta, finish_reason=None)],
+        )
+
+    citation: Final = {
+        "type": "url_citation",
+        "url_citation": {"start_index": 0, "end_index": 5, "url": "https://example.com/kiwi", "title": "Kiwi"},
+    }
+    image: Final = {"type": "image_url", "index": 0, "image_url": {"url": "data:image/png;base64,a2l3aQ=="}}
+
+    response = stream_chunk_builder(
+        chunks=[
+            chunk(0, Delta(function_call={"name": "lookup_fruit", "arguments": '{"fruit": '})),
+            chunk(1, Delta(content="Kiwis are")),
+            chunk(0, Delta(function_call={"arguments": '"kiwi"}'})),
+            chunk(1, Delta(content=" fuzzy.", annotations=[citation])),
+            chunk(1, Delta(images=[image])),
+        ]
+    )
+
+    function_choice, cited_choice = response.choices
+    assert (
+        function_choice.index,
+        function_choice.message.content,
+        function_choice.message.function_call.name,
+        function_choice.message.function_call.arguments,
+        getattr(function_choice.message, "annotations", None),
+        getattr(function_choice.message, "images", None),
+    ) == (0, None, "lookup_fruit", '{"fruit": "kiwi"}', None, None)
+    assert (
+        cited_choice.index,
+        cited_choice.message.content,
+        cited_choice.message.function_call,
+        cited_choice.message.annotations,
+        cited_choice.message.images,
+    ) == (1, "Kiwis are fuzzy.", None, [citation], [image])
+
+
 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/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py

@@ -2108,6 +2108,62 @@ class TestCiscoAIDefenseStreamingBypass:
             f"Original: {chunks!r}, received: {received!r}"
         )
 
+    @pytest.mark.asyncio
+    async def test_streaming_safe_n2_response_inspects_each_choice_and_yields_original_chunks(self):
+        g = _make_guardrail(event_hook=["pre_call", "post_call"])
+        chunks = [
+            ModelResponseStream(
+                id="resp_1",
+                choices=[StreamingChoices(delta=Delta(content=part), finish_reason=finish_reason, index=index)],
+                created=1234567890,
+                model="gpt-4",
+                object="chat.completion.chunk",
+            )
+            for index, part, finish_reason in (
+                (0, "The persim", None),
+                (1, "The", None),
+                (0, "mon is ripe.", "stop"),
+                (1, " persimmon is ripe.", "stop"),
+            )
+        ]
+
+        received, post_mock = await _streaming_setup(g, chunks, cisco_response=_safe_response())
+
+        inspected = post_mock.call_args.kwargs["json"]["messages"]
+        assert [m["content"] for m in inspected if m["role"] == "assistant"] == [
+            "The persimmon is ripe.",
+            "The persimmon is ripe.",
+        ]
+        assert received == chunks
+
+    @pytest.mark.asyncio
+    async def test_streaming_safe_n2_response_with_a_function_call_choice_yields_original_chunks(self):
+        g = _make_guardrail(event_hook=["pre_call", "post_call"])
+        chunks = [
+            ModelResponseStream(
+                id="resp_1",
+                choices=[StreamingChoices(delta=delta, finish_reason=None, index=index)],
+                created=1234567890,
+                model="gpt-4",
+                object="chat.completion.chunk",
+            )
+            for index, delta in (
+                (0, Delta(function_call={"name": "lookup_fruit", "arguments": '{"fruit": '})),
+                (1, Delta(content="Kiwis are")),
+                (0, Delta(function_call={"arguments": '"kiwi"}'})),
+                (1, Delta(content=" fuzzy.")),
+            )
+        ]
+
+        received, post_mock = await _streaming_setup(g, chunks, cisco_response=_safe_response())
+
+        inspected = post_mock.call_args.kwargs["json"]["messages"]
+        assert [m["content"] for m in inspected if m["role"] == "assistant"] == [
+            '{"fruit": "kiwi"}',
+            "Kiwis are fuzzy.",
+        ]
+        assert received == chunks
+
     @pytest.mark.asyncio
     async def test_streaming_redact_does_not_replay_tool_call_arguments(self):
         g = _make_guardrail(

tests/unit/proxy/guardrails/guardrail_hooks/test_noma.py

@@ -1,4 +1,5 @@
 import copy
+from itertools import chain
 import os
 from typing import cast
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -18,7 +19,7 @@ import litellm.proxy.guardrails.guardrail_hooks.noma.noma as noma_legacy_module
 from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage
 from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
 from litellm.types.llms.openai import AllMessageValues
-from litellm.types.utils import Choices, Message
+from litellm.types.utils import Choices, Delta, Message, ModelResponseStream, StreamingChoices
 
 
 @pytest.fixture
@@ -1848,12 +1849,207 @@ class TestNomaAnonymizationLogic:
         )
 
         anonymize_guardrail._replace_llm_response_content(
-            response, "Your email is *******"
+            response, "Your email is *******", 0
         )
 
         assert response.choices[0].message.content == "Your email is *******"
 
 
+def _noma_reply_per_text(verdict_for_text):
+    def reply(*args, **kwargs):
+        response = MagicMock()
+        response.json.return_value = verdict_for_text(kwargs["json"]["input"][0]["content"][0]["text"])
+        response.raise_for_status = MagicMock()
+        return response
+
+    return reply
+
+
+class TestNomaMultiChoiceResponses:
+    @pytest.mark.asyncio
+    async def test_streamed_two_choice_reply_is_blocked_when_only_the_second_choice_is_flagged(self):
+        guardrail = NomaGuardrail(
+            api_key="test-api-key",
+            api_base="https://api.test.noma.security/",
+            monitor_mode=False,
+            block_failures=True,
+            guardrail_name="noma-post-call",
+            event_hook="post_call",
+            default_on=True,
+        )
+        scanned_texts = []
+
+        def verdict_for_text(text):
+            scanned_texts.append(text)
+            return {"aggregatedScanResult": "launch code" in text, "scanResult": []}
+
+        async def upstream():
+            for index, part in ((0, "Kiwis are fuzzy."), (1, "The launch code is 0000.")):
+                yield ModelResponseStream(
+                    id="resp_1",
+                    choices=[StreamingChoices(delta=Delta(content=part), finish_reason="stop", index=index)],
+                    created=1234567890,
+                    model="gpt-4.1-mini",
+                    object="chat.completion.chunk",
+                )
+
+        with patch.object(guardrail.async_handler, "post", side_effect=_noma_reply_per_text(verdict_for_text)):
+            with pytest.raises(NomaBlockedMessage):
+                async for _ in guardrail.async_post_call_streaming_iterator_hook(
+                    user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
+                    response=upstream(),
+                    request_data={"litellm_call_id": "call-1", "metadata": {}},
+                ):
+                    pass
+
+        assert scanned_texts == ["Kiwis are fuzzy.", "The launch code is 0000."]
+
+    @pytest.mark.asyncio
+    async def test_each_choice_gets_its_own_anonymized_text(self):
+        guardrail = NomaGuardrail(
+            api_key="test-api-key",
+            api_base="https://api.test.noma.security/",
+            anonymize_input=True,
+            monitor_mode=False,
+            block_failures=True,
+            guardrail_name="noma-post-call",
+            event_hook="post_call",
+            default_on=True,
+        )
+
+        def verdict_for_text(text):
+            return {
+                "aggregatedScanResult": False,
+                "scanResult": [
+                    {
+                        "role": "assistant",
+                        "type": "message",
+                        "results": {
+                            "anonymizedContent": {"anonymized": text.replace("test@example.com", "*******")},
+                            "sensitiveData": {"PII": {"result": False, "status": "SUCCESS"}},
+                        },
+                    }
+                ],
+            }
+
+        response = ModelResponse(
+            id="resp_1",
+            choices=[
+                Choices(index=0, finish_reason="stop", message=Message(role="assistant", content="Mail test@example.com")),
+                Choices(
+                    index=1, finish_reason="stop", message=Message(role="assistant", content="Write test@example.com")
+                ),
+            ],
+            model="gpt-4.1-mini",
+        )
+
+        with patch.object(guardrail.async_handler, "post", side_effect=_noma_reply_per_text(verdict_for_text)):
+            await guardrail.async_post_call_success_hook(
+                data={"litellm_call_id": "call-1", "metadata": {}},
+                user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
+                response=response,
+            )
+
+        assert [choice.message.content for choice in response.choices] == ["Mail *******", "Write *******"]
+
+    @pytest.mark.asyncio
+    async def test_fail_open_stream_keeps_redactions_made_before_a_later_choice_fails(self):
+        guardrail = NomaGuardrail(
+            api_key="test-api-key",
+            api_base="https://api.test.noma.security/",
+            anonymize_input=True,
+            monitor_mode=False,
+            block_failures=False,
+            guardrail_name="noma-post-call",
+            event_hook="post_call",
+            default_on=True,
+        )
+
+        def verdict_for_text(text):
+            if text.startswith("Write"):
+                raise httpx.ConnectError("Noma is unreachable")
+            return {
+                "aggregatedScanResult": False,
+                "scanResult": [
+                    {
+                        "role": "assistant",
+                        "type": "message",
+                        "results": {
+                            "anonymizedContent": {"anonymized": text.replace("test@example.com", "*******")},
+                            "sensitiveData": {"PII": {"result": False, "status": "SUCCESS"}},
+                        },
+                    }
+                ],
+            }
+
+        async def upstream():
+            for index, part in ((0, "Mail test@example.com"), (1, "Write back soon")):
+                yield ModelResponseStream(
+                    id="resp_1",
+                    choices=[StreamingChoices(delta=Delta(content=part), finish_reason="stop", index=index)],
+                    created=1234567890,
+                    model="gpt-4.1-mini",
+                    object="chat.completion.chunk",
+                )
+
+        with patch.object(guardrail.async_handler, "post", side_effect=_noma_reply_per_text(verdict_for_text)):
+            streamed = [
+                chunk
+                async for chunk in guardrail.async_post_call_streaming_iterator_hook(
+                    user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
+                    response=upstream(),
+                    request_data={"litellm_call_id": "call-1", "metadata": {}},
+                )
+            ]
+
+        streamed_choices = tuple(chain.from_iterable(part.choices for part in streamed))
+        text_by_choice = {
+            index: "".join(choice.delta.content or "" for choice in streamed_choices if choice.index == index)
+            for index in (0, 1)
+        }
+        assert text_by_choice == {0: "Mail *******", 1: "Write back soon"}
+
+    @pytest.mark.asyncio
+    async def test_fail_open_stream_replays_the_original_chunks_when_nothing_was_redacted(self):
+        guardrail = NomaGuardrail(
+            api_key="test-api-key",
+            api_base="https://api.test.noma.security/",
+            anonymize_input=True,
+            monitor_mode=False,
+            block_failures=False,
+            guardrail_name="noma-post-call",
+            event_hook="post_call",
+            default_on=True,
+        )
+        originals = [
+            ModelResponseStream(
+                id="resp_1",
+                choices=[StreamingChoices(delta=Delta(content=part), index=0)],
+                created=1234567890,
+                model="gpt-4.1-mini",
+                object="chat.completion.chunk",
+            )
+            for part in ("Kiwis", " are", " fuzzy.")
+        ]
+
+        async def upstream():
+            for chunk in originals:
+                yield chunk
+
+        with patch.object(guardrail.async_handler, "post", side_effect=httpx.ConnectError("Noma is unreachable")):
+            streamed = [
+                chunk
+                async for chunk in guardrail.async_post_call_streaming_iterator_hook(
+                    user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
+                    response=upstream(),
+                    request_data={"litellm_call_id": "call-1", "metadata": {}},
+                )
+            ]
+
+        assert len(streamed) == len(originals)
+        assert all(sent is original for sent, original in zip(streamed, originals))
+
+
 class TestNomaAnonymizationFlow:
     """Test full anonymization flow with real Noma response objects"""
 

tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py

@@ -5,6 +5,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics)
 import json
 import logging
 import re
+from itertools import chain
 from typing import Final
 from unittest.mock import patch
 
@@ -25,10 +26,16 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
     PermissionError,
 )
 from litellm.types.utils import (
+    ChatCompletionCustomToolCallPayload,
+    ChatCompletionDeltaToolCall,
+    ChatCompletionMessageCustomToolCall,
     ChatCompletionMessageToolCall,
     Choices,
+    Delta,
+    Function,
     ModelResponse,
     ModelResponseStream,
+    StreamingChoices,
 )
 
 
@@ -425,6 +432,41 @@ class TestToolPermissionGuardrail:
         assert isinstance(choice.message.content, str)
         assert "berri" in choice.message.content
 
+    @pytest.mark.asyncio
+    async def test_rewrite_denies_only_the_choice_whose_call_is_denied_when_choices_share_a_call_id(self):
+        guardrail = ToolPermissionGuardrail(
+            guardrail_name="choice-guardrail",
+            rules=[{"id": "deny_bash", "tool_name": "Bash", "decision": "deny"}],
+            default_action="allow",
+            on_disallowed_action="rewrite",
+            event_hook=GuardrailEventHooks.post_call,
+        )
+        response = ModelResponse(
+            choices=[
+                Choices(
+                    index=index,
+                    finish_reason="tool_calls",
+                    message={
+                        "tool_calls": [
+                            {"id": "call_shared", "type": "function", "function": {"name": name, "arguments": "{}"}}
+                        ]
+                    },
+                )
+                for index, name in enumerate(("Bash", "Read"))
+            ]
+        )
+
+        await guardrail.async_post_call_success_hook(
+            data={"guardrails": ["choice-guardrail"]}, user_api_key_dict=UserAPIKeyAuth(), response=response
+        )
+
+        kept: Final = [
+            [tool_call.function.name for tool_call in choice.message.tool_calls or ()]
+            for choice in response.choices
+            if isinstance(choice, Choices)
+        ]
+        assert kept == [[], ["Read"]]
+
     @pytest.mark.asyncio
     async def test_async_post_call_success_hook_missing_arguments_blocks_param_rule(
         self,
@@ -866,6 +908,46 @@ class TestToolPermissionGuardrail:
             f"got: {chunks[0].choices[0].delta.content!r}"
         )
 
+    @pytest.mark.asyncio
+    async def test_async_post_call_streaming_iterator_hook_re_emits_each_choice_of_an_n2_tool_call_stream(self):
+        def chunk(choice_index: int, tool_call: ChatCompletionDeltaToolCall) -> ModelResponseStream:
+            return ModelResponseStream(
+                id="chatcmpl-n2-tools",
+                created=1700000000,
+                model="gpt-4.1-mini",
+                object="chat.completion.chunk",
+                choices=[StreamingChoices(index=choice_index, delta=Delta(tool_calls=[tool_call]))],
+            )
+
+        def fragment(arguments: str, name: str | None = None) -> ChatCompletionDeltaToolCall:
+            return ChatCompletionDeltaToolCall(
+                id="call_shared" if name else None,
+                index=0,
+                type="function",
+                function=Function(name=name, arguments=arguments),
+            )
+
+        async def _stream():
+            for streamed in (
+                chunk(0, fragment("", name="Bash")),
+                chunk(1, fragment("", name="Bash")),
+                chunk(0, fragment('{"command": "ls"}')),
+                chunk(1, fragment('{"command": "pwd"}')),
+            ):
+                yield streamed
+
+        re_emitted: Final = [
+            chunk
+            async for chunk in self.guardrail.async_post_call_streaming_iterator_hook(
+                user_api_key_dict=UserAPIKeyAuth(), response=_stream(), request_data={}
+            )
+        ]
+
+        assert [
+            (choice.index, [(call.function.name, call.function.arguments) for call in choice.delta.tool_calls])
+            for choice in chain.from_iterable(streamed.choices for streamed in re_emitted)
+        ] == [(0, [("Bash", '{"command": "ls"}')]), (1, [("Bash", '{"command": "pwd"}')])]
+
     def test_modify_response_with_permission_errors(self):
         # Setup a response with one tool_call
         tool_call = ChatCompletionMessageToolCall(
@@ -952,6 +1034,35 @@ class TestToolPermissionGuardrail:
         assert "Permission denied" in choice.message.content
 
 
+    def test_modify_response_with_permission_errors_keeps_custom_tool_call_in_another_choice(self):
+        denied_call = ChatCompletionMessageToolCall(
+            id="call_bash", type="function", function=Function(name="Bash", arguments="{}")
+        )
+        custom_call = ChatCompletionMessageCustomToolCall(
+            id="call_patch", custom=ChatCompletionCustomToolCallPayload(name="apply_patch", input="*** Begin Patch")
+        )
+        response = ModelResponse(
+            choices=[
+                Choices(index=0, message={"role": "assistant", "content": None, "tool_calls": [denied_call]}),
+                Choices(index=1, message={"role": "assistant", "content": None, "tool_calls": [custom_call]}),
+            ]
+        )
+        denied_tools = [
+            (
+                denied_call,
+                PermissionError(tool_name="Bash", rule_id="deny_bash", message="Tool 'Bash' denied by rule 'deny_bash'"),
+            )
+        ]
+
+        self.guardrail._modify_response_with_permission_errors(response, denied_tools)
+
+        kept = [
+            [call.id for call in choice.message.tool_calls or []]
+            for choice in response.choices
+            if isinstance(choice, Choices)
+        ]
+        assert kept == [[], ["call_patch"]]
+
 class TestToolPermissionGuardrailIntegration:
     """Integration tests for Tool Permission Guardrail"""
 

tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py

@@ -131,6 +131,86 @@ def test_extract_tool_calls_from_chat_response_handles_tool_calls():
     assert tool_calls[0]["function"]["name"] == "foo"
 
 
+def test_chat_follow_up_answers_only_tool_calls_in_its_assistant_turn():
+    response = ModelResponse(
+        id="resp-n2",
+        choices=[
+            {
+                "index": index,
+                "message": {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {"id": call_id, "type": "function", "function": {"name": "foo", "arguments": "{}"}}
+                    ],
+                },
+                "finish_reason": "tool_calls",
+            }
+            for index, call_id in enumerate(("call-a", "call-b"))
+        ],
+        model="gpt",
+        created=0,
+        object="chat.completion",
+    )
+    tool_results = [
+        {"tool_call_id": tool_call["id"], "name": tool_call["function"]["name"], "result": "done"}
+        for tool_call in LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response)
+    ]
+
+    follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_messages_for_chat(
+        [{"role": "user", "content": "hi"}], response, tool_results
+    )
+
+    assistant_call_ids = [tool_call["id"] for tool_call in follow_up[1]["tool_calls"]]
+    answered_call_ids = [message["tool_call_id"] for message in follow_up if message["role"] == "tool"]
+    assert answered_call_ids == assistant_call_ids == ["call-a"]
+
+
+def test_chat_follow_up_continues_the_choice_that_called_a_tool_when_choice_0_answered():
+    response = ModelResponse(
+        id="resp-n2-second",
+        choices=[
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "no tool needed"},
+                "finish_reason": "stop",
+            },
+            {
+                "index": 1,
+                "message": {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {"id": "call-b", "type": "function", "function": {"name": "foo", "arguments": "{}"}}
+                    ],
+                },
+                "finish_reason": "tool_calls",
+            },
+        ],
+        model="gpt",
+        created=0,
+        object="chat.completion",
+    )
+    tool_results = [
+        {"tool_call_id": tool_call["id"], "name": tool_call["function"]["name"], "result": "done"}
+        for tool_call in LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response)
+    ]
+
+    follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_messages_for_chat(
+        [{"role": "user", "content": "hi"}], response, tool_results
+    )
+
+    assistant_call_ids = [tool_call["id"] for tool_call in follow_up[1]["tool_calls"]]
+    answered_call_ids = [message["tool_call_id"] for message in follow_up if message["role"] == "tool"]
+    assert answered_call_ids == assistant_call_ids == ["call-b"]
+
+
+def test_extract_tool_calls_from_chat_response_without_choices_finds_none():
+    response = ModelResponse(id="resp-empty", choices=[], model="gpt", created=0, object="chat.completion")
+
+    assert LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response) == []
+
+
 def test_create_follow_up_messages_for_chat_appends_tool_results():
     original_messages = [{"role": "user", "content": "hi"}]
     response = ModelResponse(

tests/unit/test_main.py

@@ -2831,9 +2831,10 @@ def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of
 
     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] == [
+    text_choice, tool_choice = response.choices
+    assert (text_choice.index, text_choice.message.content, text_choice.message.tool_calls) == (0, "hello world", None)
+    assert tool_choice.index == 1
+    assert [(call.id, call.function.name, call.function.arguments) for call in tool_choice.message.tool_calls] == [
         ("call_1", "lookup_fruit", '{"fruit":"kiwi"}')
     ]