Merge pull request #40271 from BerriAI/litellm_lit_7174_stream_tool_call_rewrites

feat(guardrails): deliver tool-call rewrites into buffered chat, Responses, and Messages streams

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

@@ -13,10 +13,11 @@ Pattern Overview:
 """
 
 import json
-from collections.abc import Iterator, Mapping, Sequence
+from collections.abc import Mapping, MutableSequence, Sequence
 from copy import deepcopy
 from dataclasses import dataclass
 from itertools import chain, repeat
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
 
 from typing_extensions import ReadOnly, TypedDict, assert_never
@@ -41,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     merge_guardrailed_scoped_messages,
     merge_returned_tools_into_request_tools,
     scoped_structured_message_indices,
+    stream_item_field,
     stream_item_fingerprint,
 )
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
@@ -153,6 +155,46 @@ class ExtractedInput:
 EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
 
 
+@dataclass(frozen=True, slots=True)
+class _ToolCallShape:
+    name: str | None
+    arguments: str
+
+
+@dataclass(frozen=True, slots=True)
+class _SSEFieldRewrite:
+    """One field of one nested section of a buffered SSE event, rewritten."""
+
+    section: str
+    field: str
+    value: object
+
+
+class _SSEEventRewriter(Protocol):
+    def __call__(self, event: Mapping[str, object]) -> _SSEFieldRewrite | None: ...
+
+
+def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewriter) -> Mapping[str, object]:
+    rewrite: Final = rewrite_event(event)
+    section: Final = None if rewrite is None else event.get(rewrite.section)
+    if rewrite is None or not isinstance(section, Mapping):
+        return event
+    return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}}  # mutable-ok: json.dumps needs a dict
+
+
+def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]:
+    """The guardrail-visible shape of each tool call, whether the guardrail handed
+    back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts."""
+    functions: Final = tuple(stream_item_field(tool_call, "function") for tool_call in tool_calls)
+    return tuple(
+        _ToolCallShape(
+            name=name if isinstance(name := stream_item_field(function, "name"), str) else None,
+            arguments=arguments if isinstance(arguments := stream_item_field(function, "arguments"), str) else "",
+        )
+        for function in functions
+    )
+
+
 class _AnthropicSSEDelta(TypedDict, total=False):
     type: ReadOnly[str]
     text: ReadOnly[str]
@@ -170,7 +212,7 @@ class AnthropicMessagesHandler(BaseTranslation):
     them through guardrail rewrites; downstream provider handling is out of scope.
     """
 
-    delivers_ended_stream_text_rewrites = True
+    delivers_ended_stream_rewrites = True
 
     def __init__(self):
         super().__init__()
@@ -1050,6 +1092,7 @@ class AnthropicMessagesHandler(BaseTranslation):
                     first_choice.message.tool_calls,
                 )
                 string_so_far = first_choice.message.content
+                pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ())
                 guardrail_inputs: Final = GenericGuardrailAPIInputs()
                 if string_so_far:
                     guardrail_inputs["texts"] = [string_so_far]
@@ -1084,6 +1127,19 @@ class AnthropicMessagesHandler(BaseTranslation):
                     and guardrailed_texts[0] != string_so_far
                 ):
                     self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
+                if deliver_ended_stream_rewrites:
+                    returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
+                    self._write_ended_stream_tool_call_rewrites(
+                        responses_so_far,
+                        pre_guardrail_tool_calls=pre_guardrail_tool_calls,
+                        post_guardrail_tool_calls=_tool_call_shapes(
+                            returned_tool_calls
+                            if isinstance(returned_tool_calls, list)
+                            and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
+                            else tool_calls_list or ()
+                        ),
+                        guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
+                    )
             else:
                 verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
             return responses_so_far
@@ -1206,44 +1262,124 @@ class AnthropicMessagesHandler(BaseTranslation):
 
     @staticmethod
     def _write_ended_stream_text_rewrite(
-        responses_so_far: list[Any],  # mutable-ok: rewrites the caller's buffered chunks in place
+        responses_so_far: MutableSequence[object],  # mutable-ok: rewrites the caller's buffered chunks in place
         rewritten_text: str,
     ) -> None:
         """Deliver an ended-stream guardrail text rewrite by rewriting the
         buffered chunks in place: the first ``text_delta`` carries the full
         rewritten text and every later one is blanked, leaving the surrounding
-        message and content-block framing untouched. Handles both chunk formats
-        this stream carries (parsed event dicts and raw SSE bytes)."""
+        message and content-block framing untouched."""
         replacements: Final = chain((rewritten_text,), repeat(""))
-        for idx, item in enumerate(responses_so_far):
-            if isinstance(item, dict):
-                delta = item.get("delta")
-                if item.get("type") == "content_block_delta" and isinstance(delta, dict):
-                    if delta.get("type") == "text_delta":
-                        delta["text"] = next(replacements)
-            elif isinstance(item, (bytes, bytearray)):
-                responses_so_far[idx] = (  # rebind-ok: delivers the rewrite into the caller's buffer
-                    AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements)
-                )
+
+        def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
+            delta: Final = event.get("delta")
+            if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping):
+                return None
+            if delta.get("type") != "text_delta":
+                return None
+            return _SSEFieldRewrite("delta", "text", next(replacements))
+
+        AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
+
+    @classmethod
+    def _write_ended_stream_tool_call_rewrites(
+        cls,
+        responses_so_far: MutableSequence[object],  # mutable-ok: rewrites the caller's buffered chunks in place
+        *,
+        pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+        post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+        guardrail_name: str,
+    ) -> None:
+        """Deliver ended-stream guardrail tool-call rewrites by rewriting the
+        buffered chunks in place: the rebuilt response lists tool calls in the
+        order of the stream's ``tool_use`` blocks, so the nth rewritten call lands
+        on the nth block, its first ``input_json_delta`` carrying the full rewritten
+        arguments, every later one blanked, and ``content_block_start`` carrying the
+        rewritten name. Blocks that do not line up with the rebuilt tool calls make
+        the rewrite undeliverable, so the pipeline executor discards it and releases
+        the original chunks."""
+        if post_guardrail_tool_calls == pre_guardrail_tool_calls:
+            return
+        block_indices: Final = tuple(
+            index
+            for item in responses_so_far
+            for event in cls._iter_sse_events(item)
+            if event.get("type") == "content_block_start"
+            and isinstance(block := event.get("content_block"), Mapping)
+            and block.get("type") == "tool_use"
+            and isinstance(index := event.get("index"), int)
+        )
+        if len(block_indices) != len(post_guardrail_tool_calls):
+            from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+            raise UndeliverableStreamRewrite(guardrail_name)
+        rewrites_by_block: Final = MappingProxyType(
+            {
+                index: after
+                for index, before, after in zip(block_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls)
+                if after != before
+            }
+        )
+        argument_replacements: Final = MappingProxyType(
+            {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()}
+        )
+
+        def rewrite_tool_use(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
+            index: Final = event.get("index")
+            if not isinstance(index, int) or index not in rewrites_by_block:
+                return None
+            match event.get("type"):
+                case "content_block_start":
+                    name: Final = rewrites_by_block[index].name
+                    if name is None:
+                        return None
+                    return _SSEFieldRewrite("content_block", "name", name)
+                case "content_block_delta":
+                    delta: Final = event.get("delta")
+                    if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta":
+                        return None
+                    return _SSEFieldRewrite("delta", "partial_json", next(argument_replacements[index]))
+                case _:
+                    return None
+
+        cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use)
 
     @staticmethod
-    def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes:
-        """Rewrite every ``text_delta`` data line in one SSE chunk with the next
-        replacement text, leaving all other events and framing byte-identical."""
+    def _rewrite_ended_stream_events(
+        responses_so_far: MutableSequence[object],  # mutable-ok: rewrites the caller's buffered chunks in place
+        rewrite_event: _SSEEventRewriter,
+    ) -> None:
+        """Replace every buffered event ``rewrite_event`` returns a rewrite for, in
+        both chunk formats this stream carries (parsed event dicts and raw SSE
+        bytes), leaving every other event and the framing untouched."""
+        rewritten_items: Final = tuple(
+            AnthropicMessagesHandler._rewrite_buffered_item(item, rewrite_event) for item in responses_so_far
+        )
+        responses_so_far[:] = rewritten_items  # rebind-ok: delivers the rewrites into the caller's buffer
+
+    @staticmethod
+    def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object:
+        if isinstance(item, dict):
+            return _rewritten_event(_as_str_mapping(item), rewrite_event)
+        if isinstance(item, (bytes, bytearray)):
+            return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event)
+        return item
+
+    @staticmethod
+    def _rewrite_sse_events(sse_bytes: bytes, rewrite_event: _SSEEventRewriter) -> bytes:
+        """Rewrite the data lines of one SSE chunk that ``rewrite_event`` rewrites,
+        leaving all other events and framing byte-identical."""
         try:
             decoded: Final = sse_bytes.decode("utf-8")
         except UnicodeDecodeError:
             return sse_bytes
         return "\n\n".join(
-            AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n")
+            "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, rewrite_event) for line in block.split("\n"))
+            for block in decoded.split("\n\n")
         ).encode("utf-8")
 
     @staticmethod
-    def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str:
-        return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n"))
-
-    @staticmethod
-    def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str:
+    def _rewrite_sse_line(line: str, rewrite_event: _SSEEventRewriter) -> str:
         if not line.startswith("data:"):
             return line
         try:
@@ -1252,14 +1388,10 @@ class AnthropicMessagesHandler(BaseTranslation):
             )
         except json.JSONDecodeError:
             return line
-        if not isinstance(data, dict) or data.get("type") != "content_block_delta":
+        if not isinstance(data, dict):
             return line
-        delta: Final = data.get("delta")
-        if not isinstance(delta, dict) or delta.get("type") != "text_delta":
-            return line
-        return "data: " + json.dumps(
-            {**data, "delta": {**delta, "text": next(replacements)}}  # mutable-ok: json.dumps needs plain dicts
-        )
+        rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event)
+        return line if rewritten is data else "data: " + json.dumps(rewritten)
 
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
         stream_ended: Final = self._check_streaming_has_ended(responses_so_far)

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -52,13 +52,14 @@ class StreamingScanKey:
 
 
 class BaseTranslation(ABC):
-    delivers_ended_stream_text_rewrites: ClassVar[bool] = False
+    delivers_ended_stream_rewrites: ClassVar[bool] = False
     """Whether ``process_output_streaming_response`` accepts
     ``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered)
-    stream, writes guardrail text rewrites back across ``responses_so_far`` so
-    a buffered pipeline can release rewritten chunks. Tool-call rewrites, and
-    text rewrites on every other translation, are undeliverable: the pipeline
-    executor discards them and releases the original chunks."""
+    stream, writes guardrail text and tool-call rewrites back across
+    ``responses_so_far`` so a buffered pipeline can release rewritten chunks,
+    raising ``UndeliverableStreamRewrite`` for a shape it cannot place. Rewrites
+    on every other translation are undeliverable: the pipeline executor
+    discards them and releases the original chunks."""
 
     @staticmethod
     def transform_user_api_key_dict_to_metadata(
@@ -175,9 +176,9 @@ class BaseTranslation(ABC):
         transformations (see ``StreamTransformSink``); base handlers ignore it.
         ``deliver_ended_stream_rewrites`` is passed True only when the caller
         holds the whole buffered stream and the subclass declares
-        ``delivers_ended_stream_text_rewrites``: the handler then writes
-        guardrail text rewrites back across ``responses_so_far`` instead of
-        discarding them.
+        ``delivers_ended_stream_rewrites``: the handler then writes
+        guardrail text and tool-call rewrites back across ``responses_so_far``
+        instead of discarding them.
         """
         return responses_so_far
 

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

@@ -49,6 +49,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import
     coerce_stream_holdback_value,
 )
 from litellm.types.utils import (
+    ChatCompletionDeltaToolCall,
+    ChatCompletionMessageToolCall,
     Choices,
     GenericGuardrailAPIInputs,
     ModelResponse,
@@ -78,7 +80,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     Methods can be overridden to customize behavior for different message formats.
     """
 
-    delivers_ended_stream_text_rewrites = True
+    delivers_ended_stream_rewrites = True
 
     def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
         """
@@ -610,13 +612,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         deliver_ended_stream_rewrites: bool,
     ) -> None:
         """Ended-stream path: rebuild the full response, run the non-streaming
-        output guardrail against it, and (when opted in) write any text rewrite
-        back across the buffered chunks."""
+        output guardrail against it, and (when opted in) write any text or
+        tool-call rewrite back across the buffered chunks."""
         model_response: Final = cast(
             ModelResponse,
             stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
         )
         pre_guardrail_texts: Final = self._string_choice_contents(model_response)
+        pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response)
         await self.process_output_response(
             response=model_response,
             guardrail_to_apply=guardrail_to_apply,
@@ -624,13 +627,21 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             user_api_key_dict=user_api_key_dict,
             request_data=request_data,
         )
-        if deliver_ended_stream_rewrites:
-            await self._write_ended_stream_text_rewrites(
-                responses_so_far=responses_so_far,
-                guardrailed_response=model_response,
-                pre_guardrail_texts=pre_guardrail_texts,
-                guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
-            )
+        if not deliver_ended_stream_rewrites:
+            return
+        guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown"
+        await self._write_ended_stream_text_rewrites(
+            responses_so_far=responses_so_far,
+            guardrailed_response=model_response,
+            pre_guardrail_texts=pre_guardrail_texts,
+            guardrail_name=guardrail_name,
+        )
+        self._write_ended_stream_tool_call_rewrites(
+            responses_so_far=responses_so_far,
+            guardrailed_response=model_response,
+            pre_guardrail_tool_calls=pre_guardrail_tool_calls,
+            guardrail_name=guardrail_name,
+        )
 
     def build_stream_error_items(
         self,
@@ -1043,6 +1054,71 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             task_mappings=[(target_choice_index, None) for _ in changed],  # mutable-ok: callee takes lists
         )
 
+    @staticmethod
+    def _function_tool_call_shapes(response: "ModelResponse") -> tuple[tuple[str | None, str], ...]:
+        return tuple(
+            (tool_call.function.name, tool_call.function.arguments)
+            for choice in response.choices
+            for tool_call in choice.message.tool_calls or ()
+            if isinstance(tool_call, ChatCompletionMessageToolCall)
+        )
+
+    @staticmethod
+    def _function_tool_call_fragments(
+        responses_so_far: Sequence["ModelResponseStream"],
+    ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]:
+        """Group the stream's function tool-call fragments by their tool-call index, in
+        the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping
+        only the indices the builder keeps (an id and a name somewhere in the stream)."""
+        fragments: Final = tuple(
+            tool_call
+            for response in responses_so_far
+            for choice in response.choices
+            for tool_call in choice.delta.tool_calls or ()
+            if isinstance(tool_call, ChatCompletionDeltaToolCall)
+        )
+        identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id)
+        named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name)
+        return tuple(
+            tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named)
+        )
+
+    def _write_ended_stream_tool_call_rewrites(
+        self,
+        responses_so_far: list["ModelResponseStream"],  # mutable-ok: rewrites the caller's buffered chunks in place
+        guardrailed_response: "ModelResponse",
+        pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...],
+        guardrail_name: str,
+    ) -> None:
+        """Write ended-stream guardrail tool-call rewrites back across the buffered
+        chunks: the rewritten name and full arguments land in the tool call's first
+        fragment and the arguments of its later fragments are blanked, mirroring the
+        text write-back. A rewrite on a stream carrying more than one distinct choice
+        index, or whose fragments do not line up with the rebuilt tool calls, is
+        reported as undeliverable, so the pipeline executor discards it and releases
+        the original chunks."""
+        post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response)
+        if post_guardrail_tool_calls == pre_guardrail_tool_calls:
+            return
+        stream_choice_indices: Final = frozenset(
+            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):
+            from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+            raise UndeliverableStreamRewrite(guardrail_name)
+        for before, (name, arguments), fragments in zip(
+            pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call
+        ):
+            if (name, arguments) == before:
+                continue
+            head, *tail = fragments
+            head.function.name = name
+            head.function.arguments = arguments
+            for fragment in tail:
+                fragment.function.arguments = ""
+
     async def _apply_guardrail_responses_to_output_streaming(
         self,
         responses: list["ModelResponseStream"],

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

@@ -101,6 +101,18 @@ if TYPE_CHECKING:
     from litellm.types.llms.openai import ResponseInputParam
 
 
+class _ToolCallShape(NamedTuple):
+    name: str | None
+    arguments: str
+
+
+def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]:
+    return tuple(
+        _ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", ""))
+        for tool_call in tool_calls
+    )
+
+
 class ResponseOutputEnvelope(TypedDict, total=False):
     """Dict form of a Responses API response, as far as guardrail write-back reads it."""
 
@@ -128,6 +140,10 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
 )
 
 
+_FUNCTION_CALL_ARGUMENT_EVENT_TYPES: Final = frozenset(
+    {"response.function_call_arguments.delta", "response.function_call_arguments.done"}
+)
+_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
 _PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
     {"function_call_output": "output", "message": "content"}
 )
@@ -340,7 +356,7 @@ class OpenAIResponsesHandler(BaseTranslation):
     Methods can be overridden to customize behavior for different message formats.
     """
 
-    delivers_ended_stream_text_rewrites = True
+    delivers_ended_stream_rewrites = True
 
     def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
         """
@@ -754,6 +770,7 @@ class OpenAIResponsesHandler(BaseTranslation):
                 if response_model:
                     inputs["model"] = response_model
 
+                pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
                 guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
                     inputs=inputs,
                     request_data=request_data,
@@ -762,6 +779,12 @@ class OpenAIResponsesHandler(BaseTranslation):
                 )
 
                 guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
+                returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
+                post_guardrail_tool_calls: Final = _tool_call_shapes(
+                    returned_tool_calls
+                    if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check)
+                    else tool_calls_to_check
+                )
 
                 # Write guardrailed texts back into the output items in-place.
                 # final_chunk is a reference into responses_so_far so this
@@ -784,6 +807,13 @@ class OpenAIResponsesHandler(BaseTranslation):
                             stream_events=responses_so_far[:-1],
                             rewrites_by_position=rewrites_by_position,
                         )
+                    self._deliver_ended_stream_tool_call_rewrites(
+                        responses_so_far=responses_so_far,
+                        outputs=outputs,
+                        pre_guardrail_tool_calls=pre_guardrail_tool_calls,
+                        post_guardrail_tool_calls=post_guardrail_tool_calls,
+                        guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
+                    )
                 return responses_so_far
 
         # ------------------------------------------------------------------ #
@@ -894,6 +924,129 @@ class OpenAIResponsesHandler(BaseTranslation):
                 continue
             OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten)
 
+    def _deliver_ended_stream_tool_call_rewrites(
+        self,
+        responses_so_far: Sequence[object],
+        outputs: Sequence[object],
+        pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+        post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
+        guardrail_name: str,
+    ) -> None:
+        """Write ended-stream guardrail tool-call rewrites into the completed
+        envelope's ``function_call`` items and sync the earlier stream events,
+        keyed by ``call_id``. The guardrail sees the envelope's function calls
+        in output order, which is how a rewritten call finds its ``call_id``;
+        the stream events find their call through the ``call_id`` on
+        ``output_item`` events and the ``item_id`` on argument events, since an
+        event's ``output_index`` need not match the envelope's (the chat bridge
+        numbers tool calls from 1 while the envelope lists them after the
+        message). A rewrite whose calls do not line up with the envelope, or
+        whose events cannot be found, is reported as undeliverable, so the
+        pipeline executor discards it and releases the original events."""
+        if post_guardrail_tool_calls == pre_guardrail_tool_calls:
+            return
+        function_call_items: Final = tuple(
+            output_item for output_item in outputs if stream_item_field(output_item, "type") == "function_call"
+        )
+        call_ids: Final = tuple(
+            call_id
+            for output_item in function_call_items
+            if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id
+        )
+        stream_events: Final = responses_so_far[:-1]
+        call_id_by_item_id: Final = self._function_call_ids_by_item_id(stream_events)
+        event_call_ids: Final = tuple(
+            self._function_call_event_call_id(event, call_id_by_item_id) for event in stream_events
+        )
+        rewrites_by_call_id: Final = MappingProxyType(
+            {
+                call_id: after
+                for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls)
+                if after != before
+            }
+        )
+        unresolved_argument_event: Final = any(
+            call_id is None and stream_item_field(event, "type") in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES
+            for event, call_id in zip(stream_events, event_call_ids)
+        )
+        if (
+            len(call_ids) != len(function_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)
+        ):
+            from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+            raise UndeliverableStreamRewrite(guardrail_name)
+        for output_item, rewrite in (
+            (output_item, rewrites_by_call_id[call_id])
+            for output_item, call_id in zip(function_call_items, call_ids)
+            if call_id in rewrites_by_call_id
+        ):
+            self._write_function_call_item(output_item, rewrite.name, rewrite.arguments)
+        delta_replacements: Final = MappingProxyType(
+            {call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()}
+        )
+        for event, call_id in zip(stream_events, event_call_ids):
+            if call_id not in rewrites_by_call_id:
+                continue
+            match stream_item_field(event, "type"):
+                case "response.function_call_arguments.delta":
+                    self._write_event_field(event, "delta", next(delta_replacements[call_id]))
+                case "response.function_call_arguments.done":
+                    self._write_event_field(event, "arguments", rewrites_by_call_id[call_id].arguments)
+                case "response.output_item.added":
+                    self._write_function_call_item(
+                        stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None
+                    )
+                case "response.output_item.done":
+                    self._write_function_call_item(
+                        stream_item_field(event, "item"),
+                        rewrites_by_call_id[call_id].name,
+                        rewrites_by_call_id[call_id].arguments,
+                    )
+                case _:
+                    pass
+
+    @staticmethod
+    def _function_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
+        items: Final = tuple(
+            stream_item_field(event, "item")
+            for event in stream_events
+            if stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES
+        )
+        return MappingProxyType(
+            {
+                item_id: call_id
+                for item in items
+                if stream_item_field(item, "type") == "function_call"
+                and isinstance(item_id := stream_item_field(item, "id"), str)
+                and isinstance(call_id := stream_item_field(item, "call_id"), str)
+            }
+        )
+
+    @staticmethod
+    def _function_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
+        event_type: Final = stream_item_field(event, "type")
+        if event_type in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES:
+            item_id: Final = stream_item_field(event, "item_id")
+            return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None
+        if event_type not in _OUTPUT_ITEM_EVENT_TYPES:
+            return None
+        item: Final = stream_item_field(event, "item")
+        call_id: Final = stream_item_field(item, "call_id")
+        return call_id if stream_item_field(item, "type") == "function_call" and isinstance(call_id, str) else None
+
+    @staticmethod
+    def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None:
+        if item is None:
+            return
+        if name is not None:
+            OpenAIResponsesHandler._write_event_field(item, "name", name)
+        if arguments is not None:
+            OpenAIResponsesHandler._write_event_field(item, "arguments", arguments)
+
     def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
         """
         Check if the streaming has ended.

litellm/proxy/policy_engine/pipeline_executor.py

@@ -78,6 +78,10 @@ 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)
+
+
 _GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object])
 
 
@@ -89,10 +93,11 @@ def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT:
 class _StreamRewriteObserver(CustomGuardrail):
     """Stand-in handed to the endpoint translation in place of a streaming pipeline step's
     guardrail. It records whether the guardrail returned different output than it was given,
-    which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text
-    rewrites are deliverable on translations that write them back across the buffered chunks
-    (``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any
-    other translation are discarded by the executor, which releases the original chunks.
+    which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text and
+    tool-call rewrites are deliverable on translations that write them back across the
+    buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation,
+    and a rewrite that drops or adds a tool call on any translation, are discarded by the
+    executor, which releases the original chunks.
     The inner guardrail's ``apply_guardrail`` already records the guardrail information
     and span, so the observer's stays out of ``log_guardrail_information``."""
 
@@ -101,6 +106,7 @@ class _StreamRewriteObserver(CustomGuardrail):
         self.inner: Final = inner
         self.rewrote_texts = False
         self.rewrote_tool_calls = False
+        self.changed_tool_call_count = False
 
     def structured_messages_cover_full_request(self) -> bool:
         return self.inner.structured_messages_cover_full_request()
@@ -118,9 +124,11 @@ class _StreamRewriteObserver(CustomGuardrail):
         outputs: Final = await self.inner.apply_guardrail(
             inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj
         )
+        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, _tool_call_shapes(outputs.get("tool_calls"))
+        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(
+            sent_tool_shapes, returned_tool_shapes
         )
         return outputs
 
@@ -296,13 +304,13 @@ class PipelineExecutor:
         litellm_logging_obj: "LiteLLMLoggingObj | None",
     ) -> None:
         """Run one streaming post_call step through the endpoint translation, delivering
-        text rewrites on translations that support ended-stream write-back. A rewrite that
-        cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation
-        without write-back, or one the translation refused with
-        ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the
-        originals and the step passes, so the client gets the stream the merge base sent."""
+        text and tool-call rewrites on translations that support ended-stream write-back. A
+        rewrite that cannot reach the client yet (one on a translation without write-back, or
+        one the translation refused with ``UndeliverableStreamRewrite``) is discarded: the
+        buffered chunks go back to the originals and the step passes, so the client gets the
+        stream the merge base sent."""
         observer: Final = _StreamRewriteObserver(callback)
-        deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites
+        deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites
         originals: Final = copy.deepcopy(streaming_chunks)
         try:
             if deliver_rewrites:
@@ -325,7 +333,9 @@ class PipelineExecutor:
         except UndeliverableStreamRewrite:
             _release_original_chunks(step.guardrail, streaming_chunks, originals)
         else:
-            if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
+            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)
         if not callback.records_own_guardrail_information:
             add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail)

litellm/proxy/utils.py

@@ -197,6 +197,7 @@ if TYPE_CHECKING:
     from prisma.types import HttpConfig
 
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
     from litellm.models.team import LiteLLM_TeamTableCachedObj
     from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
     from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
@@ -674,7 +675,7 @@ def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") ->
 
 
 def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool:
-    return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None
+    return resolve_endpoint_translation(user_api_key_dict, None) is not None
 
 
 def _stream_gated_guardrail_names(
@@ -3564,12 +3565,16 @@ class ProxyLogging:
                     ),
                 )
 
-        if post_call_pipelines:
+        pipeline_translation: Final = (
+            resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None
+        )
+        if pipeline_translation is not None:
             current_response = self._pipeline_gated_stream(
                 response=current_response,
                 user_api_key_dict=user_api_key_dict,
                 request_data=request_data,
                 pipelines=post_call_pipelines,
+                translation=pipeline_translation,
             )
 
         try:
@@ -3593,6 +3598,7 @@ class ProxyLogging:
         user_api_key_dict: UserAPIKeyAuth,
         request_data: dict,  # mutable-ok: same request-payload shape the hooks mutate
         pipelines: "tuple[tuple[str, GuardrailPipeline], ...]",
+        translation: "tuple[str, BaseTranslation]",
     ) -> "AsyncGenerator[Any, None]":
         """
         Execute post_call policy pipelines against a streamed response.
@@ -3602,14 +3608,13 @@ class ProxyLogging:
         assembled output through the endpoint guardrail translation, the same
         machinery flat post_call guardrails use at end of stream. An allow
         releases the buffered chunks: verbatim when no guardrail rewrote the
-        output, rewritten in place when one rewrote text and the translation
-        delivers ended-stream rewrites (later steps then re-scan the rewritten
-        chunks, so rewrites chain). A rewrite the translation cannot deliver
-        yet (a tool-call rewrite, or a text rewrite on a route without
-        write-back) is discarded by the executor and the original chunks are
-        released, as is a buffered shape no translation resolves; a block or
-        modify_response terminates with the translation's block chunks or the
-        raised error.
+        output, rewritten in place when one rewrote text or a tool call and the
+        translation delivers ended-stream rewrites (later steps then re-scan the
+        rewritten chunks, so rewrites chain). A rewrite the translation cannot
+        deliver yet (one on a route without write-back, or a shape the route
+        refuses) is discarded by the executor and the original chunks are
+        released; a block or modify_response terminates with the translation's
+        block chunks or the raised error.
         """
         buffered: Final[list[object]] = []  # mutable-ok: accumulates the stream before the pipeline verdict
         async for item in response:
@@ -3617,17 +3622,7 @@ class ProxyLogging:
         if not buffered:
             return
 
-        resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0])
-        if resolved is None:
-            verbose_proxy_logger.warning(
-                "Policies with post_call guardrail pipelines cannot scan this streaming response shape yet; "
-                "the stream is released ungoverned by them: %s",
-                ", ".join(policy_name for policy_name, _pipeline in pipelines),
-            )
-            for buffered_item in buffered:
-                yield buffered_item
-            return
-        call_type, endpoint_translation = resolved
+        call_type, endpoint_translation = translation
 
         for policy_name, pipeline in pipelines:
             result: PipelineExecutionResult = await PipelineExecutor.execute_steps(

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

@@ -315,6 +315,120 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
         assert "event: message_start" in raw and "event: message_stop" in raw
         assert '"stop_reason": "end_turn"' in raw
 
+    @staticmethod
+    def _ended_tool_use_sse_chunks() -> list:
+        events = [
+            ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
+            ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}),
+            ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}),
+            ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}),
+            ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}),
+            ("content_block_stop", {"type": "content_block_stop", "index": 0}),
+            ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
+            ("message_stop", {"type": "message_stop"}),
+        ]
+        return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
+
+    @staticmethod
+    def _argument_masking_guardrail() -> CustomGuardrail:
+        class MaskArguments(CustomGuardrail):
+            async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+                for tool_call in inputs.get("tool_calls", []):
+                    tool_call.function.arguments = '{"fruit": "[MASKED]"}'
+                return inputs
+
+        return MaskArguments(guardrail_name="test")
+
+    @staticmethod
+    def _partial_jsons(chunks: list) -> list:
+        return [
+            json.loads(line[len("data:") :].strip())["delta"]["partial_json"]
+            for chunk in chunks
+            for line in chunk.decode().split("\n")
+            if line.startswith("data:") and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta"
+        ]
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_writes_tool_use_input_back_into_sse_chunks(self):
+        handler = AnthropicMessagesHandler()
+        chunks = self._ended_tool_use_sse_chunks()
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=self._argument_masking_guardrail(),
+            litellm_logging_obj=MagicMock(),
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert result is chunks
+        assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', "", ""]
+        raw = b"".join(chunks).decode()
+        assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw
+        assert '"stop_reason": "tool_use"' in raw
+        assert "persim" not in raw
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_writes_tool_use_name_back_into_sse_chunks(self):
+        class RenameTool(CustomGuardrail):
+            async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+                for tool_call in inputs.get("tool_calls", []):
+                    tool_call.function.name = "lookup_fruit_reviewed"
+                return inputs
+
+        handler = AnthropicMessagesHandler()
+        chunks = self._ended_tool_use_sse_chunks()
+
+        await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=RenameTool(guardrail_name="test"),
+            litellm_logging_obj=MagicMock(),
+            deliver_ended_stream_rewrites=True,
+        )
+
+        raw = b"".join(chunks).decode()
+        assert '"name": "lookup_fruit_reviewed"' in raw and '"id": "toolu_1"' in raw
+        assert '"name": "lookup_fruit"' not in raw
+        assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"}
+
+    @pytest.mark.asyncio
+    async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self):
+        handler = AnthropicMessagesHandler()
+        chunks = self._ended_tool_use_sse_chunks()
+        original = [bytes(chunk) for chunk in chunks]
+
+        await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=self._argument_masking_guardrail(),
+            litellm_logging_obj=MagicMock(),
+        )
+
+        assert chunks == original
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self):
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+        handler = AnthropicMessagesHandler()
+        server_tool_use = [
+            ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}),
+            ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}),
+            ("content_block_stop", {"type": "content_block_stop", "index": 0}),
+        ]
+        tool_use = self._ended_tool_use_sse_chunks()
+        chunks = (
+            tool_use[:1]
+            + [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in server_tool_use]
+            + [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]]
+        )
+
+        with pytest.raises(UndeliverableStreamRewrite):
+            await handler.process_output_streaming_response(
+                responses_so_far=chunks,
+                guardrail_to_apply=self._argument_masking_guardrail(),
+                litellm_logging_obj=MagicMock(),
+                deliver_ended_stream_rewrites=True,
+            )
+
     @pytest.mark.asyncio
     async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
         handler = AnthropicMessagesHandler()

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

@@ -1113,6 +1113,102 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
         assert chunks[1].choices[0].delta.content in (None, "")
         assert chunks[1].choices[0].finish_reason == "stop"
 
+    @staticmethod
+    def _ended_tool_call_stream_chunks() -> list:
+        from litellm.types.utils import (
+            ChatCompletionDeltaToolCall,
+            Delta,
+            Function,
+            ModelResponseStream,
+            StreamingChoices,
+        )
+
+        def chunk(tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None):
+            return ModelResponseStream(
+                id="chatcmpl-123",
+                created=1234567890,
+                model="gpt-4",
+                object="chat.completion.chunk",
+                choices=[
+                    StreamingChoices(
+                        index=0,
+                        delta=Delta(tool_calls=[tool_call] if tool_call else None),
+                        finish_reason=finish_reason,
+                    )
+                ],
+            )
+
+        def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None):
+            return ChatCompletionDeltaToolCall(
+                id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments)
+            )
+
+        return [
+            chunk(fragment("", name="lookup_fruit", call_id="call_1")),
+            chunk(fragment('{"fruit":')),
+            chunk(fragment(' "persimmon"}')),
+            chunk(None, finish_reason="tool_calls"),
+        ]
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_writes_tool_call_arguments_back_into_chunks(self):
+        handler = OpenAIChatCompletionsHandler()
+        guardrail = MockGuardrail(guardrail_name="test")
+        chunks = self._ended_tool_call_stream_chunks()
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=guardrail,
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert result is chunks
+        fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]]
+        assert [fragment[0].function.arguments for fragment in fragments] == ['{"fruit": "PERSIMMON"}', "", ""]
+        assert fragments[0][0].function.name == "lookup_fruit"
+        assert fragments[0][0].id == "call_1"
+        assert chunks[3].choices[0].delta.tool_calls is None
+        assert chunks[3].choices[0].finish_reason == "tool_calls"
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_writes_tool_call_name_back_into_chunks(self):
+        class RenameTool(CustomGuardrail):
+            async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+                for tool_call in inputs.get("tool_calls", []):
+                    tool_call["function"]["name"] = "lookup_fruit_reviewed"
+                return inputs
+
+        handler = OpenAIChatCompletionsHandler()
+        chunks = self._ended_tool_call_stream_chunks()
+
+        await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=RenameTool(guardrail_name="test"),
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        fragments = [chunk.choices[0].delta.tool_calls[0] for chunk in chunks[:3]]
+        assert [fragment.function.name for fragment in fragments] == ["lookup_fruit_reviewed", None, None]
+        assert json.loads("".join(fragment.function.arguments for fragment in fragments)) == {"fruit": "persimmon"}
+        assert fragments[0].id == "call_1"
+
+    @pytest.mark.asyncio
+    async def test_ended_stream_tool_call_rewrite_leaves_chunks_untouched_by_default(self):
+        handler = OpenAIChatCompletionsHandler()
+        guardrail = MockGuardrail(guardrail_name="test")
+        chunks = self._ended_tool_call_stream_chunks()
+
+        await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=guardrail,
+            litellm_logging_obj=None,
+        )
+
+        fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]]
+        assert [fragment[0].function.arguments for fragment in fragments] == ["", '{"fruit":', ' "persimmon"}']
+
     @pytest.mark.asyncio
     async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
         handler = OpenAIChatCompletionsHandler()
@@ -1179,6 +1275,62 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
                 deliver_ended_stream_rewrites=True,
             )
 
+    @staticmethod
+    def _two_choice_tool_call_stream_chunks() -> list:
+        from litellm.types.utils import (
+            ChatCompletionDeltaToolCall,
+            Delta,
+            Function,
+            ModelResponseStream,
+            StreamingChoices,
+        )
+
+        def chunk(
+            choice_index: int, tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None
+        ) -> ModelResponseStream:
+            return ModelResponseStream(
+                id="chatcmpl-123",
+                created=1234567890,
+                model="gpt-4",
+                object="chat.completion.chunk",
+                choices=[
+                    StreamingChoices(
+                        index=choice_index,
+                        delta=Delta(tool_calls=[tool_call] if tool_call else None),
+                        finish_reason=finish_reason,
+                    )
+                ],
+            )
+
+        def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None):
+            return ChatCompletionDeltaToolCall(
+                id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments)
+            )
+
+        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, None, finish_reason="tool_calls"),
+            chunk(1, None, finish_reason="tool_calls"),
+        ]
+
+    @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
+
+        handler = OpenAIChatCompletionsHandler()
+        chunks = self._two_choice_tool_call_stream_chunks()
+
+        with pytest.raises(UndeliverableStreamRewrite):
+            await handler.process_output_streaming_response(
+                responses_so_far=chunks,
+                guardrail_to_apply=MockGuardrail(guardrail_name="test"),
+                litellm_logging_obj=None,
+                deliver_ended_stream_rewrites=True,
+            )
+
     @pytest.mark.asyncio
     async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):
         handler = OpenAIChatCompletionsHandler()

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

@@ -17,6 +17,7 @@ from fastapi import HTTPException
 from openai.types.responses import ResponseFunctionToolCall
 
 from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms import get_guardrail_translation_mapping
 from litellm.llms.openai.responses.guardrail_translation.handler import (
     OpenAIResponsesHandler,
@@ -1195,6 +1196,230 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
         assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]"
         assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
 
+    @staticmethod
+    def _ended_function_call_stream_events() -> List[dict]:
+        def item(arguments: str, status: str) -> dict:
+            return {
+                "type": "function_call",
+                "id": "fc_123",
+                "call_id": "call_123",
+                "name": "lookup_fruit",
+                "arguments": arguments,
+                "status": status,
+            }
+
+        return [
+            {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")},
+            {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": '{"fruit":'},
+            {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": ' "persimmon"}'},
+            {
+                "type": "response.function_call_arguments.done",
+                "item_id": "fc_123",
+                "output_index": 0,
+                "arguments": '{"fruit": "persimmon"}',
+            },
+            {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")},
+            {
+                "type": "response.completed",
+                "response": {
+                    "id": "resp_123",
+                    "created_at": 1,
+                    "model": "gpt-4o",
+                    "output": [item('{"fruit": "persimmon"}', "completed")],
+                    "status": "completed",
+                },
+            },
+        ]
+
+    @staticmethod
+    def _argument_masking_guardrail() -> CustomGuardrail:
+        class MaskArguments(CustomGuardrail):
+            async def apply_guardrail(
+                self,
+                inputs: GenericGuardrailAPIInputs,
+                request_data: dict,
+                input_type: Literal["request", "response"],
+                logging_obj: LiteLLMLoggingObj | None = None,
+            ) -> GenericGuardrailAPIInputs:
+                tool_calls = [
+                    {**tool_call, "function": {**tool_call["function"], "arguments": '{"fruit": "[MASKED]"}'}}
+                    for tool_call in inputs.get("tool_calls", [])
+                ]
+                return {**inputs, "tool_calls": tool_calls}
+
+        return MaskArguments(guardrail_name="test-mask-arguments")
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_syncs_function_call_events(self):
+        handler = OpenAIResponsesHandler()
+        events = self._ended_function_call_stream_events()
+
+        result = 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 result is events
+        assert events[0]["item"]["arguments"] == ""
+        assert events[1]["delta"] == '{"fruit": "[MASKED]"}'
+        assert events[2]["delta"] == ""
+        assert events[3]["arguments"] == '{"fruit": "[MASKED]"}'
+        assert events[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}'
+        assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}'
+        assert events[5]["response"]["output"][0]["name"] == "lookup_fruit"
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_syncs_typed_function_call_events(self):
+        from litellm.types.llms.openai import (
+            FunctionCallArgumentsDeltaEvent,
+            FunctionCallArgumentsDoneEvent,
+            OutputItemAddedEvent,
+            OutputItemDoneEvent,
+            ResponseCompletedEvent,
+            ResponsesAPIResponse,
+        )
+
+        handler = OpenAIResponsesHandler()
+        typed_events: List[Any] = [
+            model.model_validate(event)
+            for model, event in zip(
+                (
+                    OutputItemAddedEvent,
+                    FunctionCallArgumentsDeltaEvent,
+                    FunctionCallArgumentsDeltaEvent,
+                    FunctionCallArgumentsDoneEvent,
+                    OutputItemDoneEvent,
+                    ResponseCompletedEvent,
+                ),
+                self._ended_function_call_stream_events(),
+            )
+        ]
+        completed_event = typed_events[5]
+        assert isinstance(completed_event, ResponseCompletedEvent)
+        assert isinstance(completed_event.response, ResponsesAPIResponse)
+        assert isinstance(completed_event.response.output[0], ResponseFunctionToolCall)
+
+        await handler.process_output_streaming_response(
+            responses_so_far=typed_events,
+            guardrail_to_apply=self._argument_masking_guardrail(),
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert typed_events[1].delta == '{"fruit": "[MASKED]"}'
+        assert typed_events[2].delta == ""
+        assert typed_events[3].arguments == '{"fruit": "[MASKED]"}'
+        assert typed_events[4].item.arguments == '{"fruit": "[MASKED]"}'
+        assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}'
+        assert completed_event.response.output[0].name == "lookup_fruit"
+
+    @staticmethod
+    def _bridged_function_call_stream_events() -> List[dict]:
+        reasoning = {"type": "reasoning", "id": "rs_1", "summary": []}
+        text = {"type": "output_text", "text": "Looking that up", "annotations": []}
+        message = {"type": "message", "id": "msg_1", "role": "assistant", "status": "completed", "content": [text]}
+
+        def function_call(arguments: str, status: str) -> dict:
+            return {
+                "type": "function_call",
+                "id": "fc_1",
+                "call_id": "call_1",
+                "name": "lookup_fruit",
+                "arguments": arguments,
+                "status": status,
+            }
+
+        return [
+            {"type": "response.output_item.added", "output_index": 0, "item": dict(reasoning)},
+            {"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning)},
+            {"type": "response.output_item.added", "output_index": 0, "item": {**message, "status": "in_progress", "content": []}},
+            {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": "Looking that up"},
+            {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": [dict(text)]}},
+            {"type": "response.output_item.added", "output_index": 1, "item": function_call("", "in_progress")},
+            {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": '{"fruit":'},
+            {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": ' "persimmon"}'},
+            {
+                "type": "response.function_call_arguments.done",
+                "item_id": "fc_1",
+                "output_index": 1,
+                "arguments": '{"fruit": "persimmon"}',
+            },
+            {"type": "response.output_item.done", "output_index": 1, "item": function_call('{"fruit": "persimmon"}', "completed")},
+            {
+                "type": "response.completed",
+                "response": {
+                    "id": "resp_1",
+                    "model": "claude-haiku-4-5",
+                    "output": [
+                        dict(reasoning),
+                        {**message, "content": [dict(text)]},
+                        function_call('{"fruit": "persimmon"}', "completed"),
+                    ],
+                },
+            },
+        ]
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_rewrites_keys_bridged_function_call_events_by_call_id(self):
+        handler = OpenAIResponsesHandler()
+        events = self._bridged_function_call_stream_events()
+
+        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 events[6]["delta"] == '{"fruit": "[MASKED]"}'
+        assert events[7]["delta"] == ""
+        assert events[8]["arguments"] == '{"fruit": "[MASKED]"}'
+        assert events[5]["item"]["name"] == "lookup_fruit"
+        assert events[9]["item"]["arguments"] == '{"fruit": "[MASKED]"}'
+        assert events[10]["response"]["output"][2]["arguments"] == '{"fruit": "[MASKED]"}'
+        assert events[3]["delta"] == "Looking that up"
+        assert events[4]["item"]["content"][0]["text"] == "Looking that up"
+        assert events[10]["response"]["output"][1]["content"][0]["text"] == "Looking that up"
+        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):
+        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+
+        handler = OpenAIResponsesHandler()
+        events = self._ended_function_call_stream_events()
+        envelope_item = events[5]["response"]["output"][0]
+        if mismatch == "orphan_call_id":
+            events[5]["response"]["output"] = [{**envelope_item, "call_id": "call_999"}]
+        else:
+            events[5]["response"]["output"] = [dict(envelope_item), dict(envelope_item)]
+
+        with pytest.raises(UndeliverableStreamRewrite):
+            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,
+            )
+
+    @pytest.mark.asyncio
+    async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self):
+        handler = OpenAIResponsesHandler()
+        events = self._ended_function_call_stream_events()
+
+        await handler.process_output_streaming_response(
+            responses_so_far=events,
+            guardrail_to_apply=self._argument_masking_guardrail(),
+            litellm_logging_obj=None,
+        )
+
+        assert events[1]["delta"] == '{"fruit":'
+        assert events[3]["arguments"] == '{"fruit": "persimmon"}'
+        assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "persimmon"}'
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"])
     async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type):

tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py

@@ -1065,7 +1065,7 @@ class _TextReturningGuardrail(CustomGuardrail):
 
 
 class _TextTranslation:
-    delivers_ended_stream_text_rewrites = False
+    delivers_ended_stream_rewrites = False
 
     def __init__(self):
         self.seen_guardrail_names = []
@@ -1087,7 +1087,7 @@ class _WritingTranslation:
     """Writes the guardrail's text (and tool-call) outputs back into the buffered chunks the way the
     chat/Responses/Messages handlers do on an ended stream."""
 
-    delivers_ended_stream_text_rewrites = True
+    delivers_ended_stream_rewrites = True
 
     async def process_output_streaming_response(
         self,
@@ -1106,12 +1106,13 @@ class _WritingTranslation:
             logging_obj=litellm_logging_obj,
         )
         responses_so_far[0]["text"] = outputs["texts"][0]
-        responses_so_far[0]["tool_call"] = outputs["tool_calls"][0]
+        if len(outputs["tool_calls"]) == 1:
+            responses_so_far[0]["tool_call"] = outputs["tool_calls"][0]
         return responses_so_far
 
 
 class _RefusingTranslation:
-    delivers_ended_stream_text_rewrites = True
+    delivers_ended_stream_rewrites = True
 
     async def process_output_streaming_response(
         self,
@@ -1229,13 +1230,47 @@ async def test_streaming_step_delivers_text_rewrite_through_writing_translation(
 
 
 @pytest.mark.asyncio
-async def test_streaming_step_discards_tool_call_rewrite_and_restores_written_text(monkeypatch, caplog):
+async def test_streaming_step_delivers_tool_call_rewrite_through_writing_translation(monkeypatch, caplog):
     monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)])
     chunks = [_chunk()]
 
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
         result = await _run_streaming_step(_WritingTranslation(), chunks)
 
+    assert result.terminal_action == "allow"
+    assert chunks[0]["text"] == "hello [MASKED]"
+    assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "[MASKED]"}'
+    assert not any("discarded" in record.getMessage() for record in caplog.records)
+
+
+class _ToolCallDroppingGuardrail(CustomGuardrail):
+    def __init__(self):
+        super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
+
+    async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+        return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": []}
+
+
+@pytest.mark.asyncio
+async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool_call(monkeypatch, caplog):
+    monkeypatch.setattr(litellm, "callbacks", [_ToolCallDroppingGuardrail()])
+    chunks = [_chunk()]
+
+    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+        result = await _run_streaming_step(_WritingTranslation(), chunks)
+
+    _assert_passed_with_discard_warning(result, caplog)
+    assert chunks == [_chunk()]
+
+
+@pytest.mark.asyncio
+async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_write_back(monkeypatch, caplog):
+    monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)])
+    chunks = [_chunk()]
+
+    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+        result = await _run_streaming_step(_TextTranslation(), chunks)
+
     _assert_passed_with_discard_warning(result, caplog)
     assert chunks == [_chunk()]
 

tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py

@@ -26,6 +26,7 @@ from litellm.integrations.custom_guardrail import (
     ModifyResponseException,
 )
 from litellm.integrations.prometheus import PrometheusLogger
+from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
 from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines
@@ -2065,7 +2066,7 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]:
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")])
-async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_tool_call_rewrite(
+async def test_streaming_iterator_hook_pipeline_delivers_runtime_tool_call_rewrite(
     proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error, caplog
 ):
     transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')}  # noqa: E731
@@ -2084,9 +2085,12 @@ async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_to
             delivered.append(item)
 
     assert len(delivered) == 2
-    assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}'
+    delivered_tool_call = delivered[0].choices[0].delta.tool_calls[0]
+    assert delivered_tool_call.function.arguments == '{"ssn": "[MASKED]"}'
+    assert delivered_tool_call.function.name == "lookup"
+    assert delivered_tool_call.id == "call_1"
     assert delivered[1].choices[0].finish_reason == "tool_calls"
-    assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))
+    assert not any("discarded" in message for message in _warnings(caplog))
 
 
 @pytest.mark.asyncio
@@ -2194,28 +2198,61 @@ async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_anothe
 
 
 @pytest.mark.asyncio
-async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvable_response_shape(
+async def test_streaming_iterator_hook_skips_pipeline_and_warns_without_request_route(
     proxy_logging, make_user_api_key_auth, monkeypatch, caplog
 ):
     seen: Dict[str, Any] = {}
     monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
     monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
     data = _post_call_pipeline_data(stream=True)
-    chunks = [object(), object()]
-    delivered: List[Any] = []
+    chunks = _stream_chunks()
 
     with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
-        async for item in proxy_logging.async_post_call_streaming_iterator_hook(
-            user_api_key_dict=make_user_api_key_auth(),
-            response=_async_chunk_iter(chunks),
-            request_data=data,
-        ):
-            delivered.append(item)
+        delivered = [
+            item
+            async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+                user_api_key_dict=make_user_api_key_auth(),
+                response=_async_chunk_iter(chunks),
+                request_data=data,
+            )
+        ]
 
     assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True]
     assert len(delivered) == 2
     assert seen.get("count") is None
-    assert any("response-governance" in message and "shape" in message for message in _warnings(caplog))
+    assert any("response-governance" in message and "route None" in message for message in _warnings(caplog))
+
+
+@pytest.mark.asyncio
+async def test_per_chunk_streaming_hook_runs_pipeline_managed_guardrail_without_request_route(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    seen: Dict[str, Any] = {}
+
+    class UnifiedRecordingGuardrail(CustomGuardrail):
+        async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+            seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1
+            return None
+
+        async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+            return inputs
+
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [UnifiedRecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)],
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data(stream=True)
+
+    result = await proxy_logging.async_post_call_streaming_hook(
+        data=data,
+        response=_stream_chunks()[0],
+        user_api_key_dict=make_user_api_key_auth(),
+    )
+
+    assert result is not None
+    assert seen["gr-post"] == 1
 
 
 def _anthropic_sse_chunks() -> List[bytes]:
@@ -2442,3 +2479,177 @@ async def test_per_chunk_streaming_hook_runs_guardrail_whose_pipeline_cannot_str
     assert result is not None
     assert seen["count"] == 1
     assert seen["response"] == "hello "
+
+
+def _mask_tool_call_arguments(inputs: Dict[str, Any]) -> Dict[str, Any]:
+    return {
+        "tool_calls": [
+            {
+                "id": stream_item_field(tool_call, "id"),
+                "type": "function",
+                "function": {
+                    "name": stream_item_field(stream_item_field(tool_call, "function"), "name"),
+                    "arguments": '{"fruit": "[MASKED]"}',
+                },
+            }
+            for tool_call in inputs.get("tool_calls", [])
+        ]
+    }
+
+
+def _anthropic_tool_use_sse_chunks() -> List[bytes]:
+    events = [
+        ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
+        ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}),
+        ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}),
+        ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}),
+        ("content_block_stop", {"type": "content_block_stop", "index": 0}),
+        ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
+        ("message_stop", {"type": "message_stop"}),
+    ]
+    return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
+
+
+@pytest.mark.asyncio
+async def test_streaming_iterator_hook_pipeline_delivers_tool_use_rewrite_on_anthropic_sse(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data(stream=True)
+
+    delivered = [
+        item
+        async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"),
+            response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()),
+            request_data=data,
+        )
+    ]
+
+    raw = b"".join(delivered).decode()
+    assert '{\\"fruit\\": \\"[MASKED]\\"}' in raw
+    assert "persim" not in raw
+    assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw
+    assert '"stop_reason": "tool_use"' in raw
+    assert raw.count("event: content_block_delta") == 2
+
+
+def _responses_function_call_events() -> List[Dict[str, Any]]:
+    def item(arguments: str, status: str) -> Dict[str, Any]:
+        return {
+            "type": "function_call",
+            "id": "fc_1",
+            "call_id": "call_1",
+            "name": "lookup_fruit",
+            "arguments": arguments,
+            "status": status,
+        }
+
+    return [
+        {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")},
+        {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": '{"fruit":'},
+        {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": ' "persimmon"}'},
+        {"type": "response.function_call_arguments.done", "item_id": "fc_1", "output_index": 0, "arguments": '{"fruit": "persimmon"}'},
+        {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")},
+        {
+            "type": "response.completed",
+            "response": {"id": "resp_1", "created_at": 1, "model": "m", "output": [item('{"fruit": "persimmon"}', "completed")], "status": "completed"},
+        },
+    ]
+
+
+@pytest.mark.asyncio
+async def test_streaming_iterator_hook_pipeline_delivers_function_call_rewrite_on_responses_events(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data(stream=True)
+
+    delivered = [
+        item
+        async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"),
+            response=_async_chunk_iter(_responses_function_call_events()),
+            request_data=data,
+        )
+    ]
+
+    assert [event["type"] for event in delivered] == [event["type"] for event in _responses_function_call_events()]
+    assert [event["delta"] for event in delivered if event["type"] == "response.function_call_arguments.delta"] == ['{"fruit": "[MASKED]"}', ""]
+    assert delivered[3]["arguments"] == '{"fruit": "[MASKED]"}'
+    assert delivered[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}'
+    assert delivered[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}'
+    assert "persimmon" not in json.dumps(delivered)
+
+
+def _drop_tool_calls(inputs: Dict[str, Any]) -> Dict[str, Any]:
+    return {"tool_calls": []}
+
+
+@pytest.mark.asyncio
+async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_chat_chunks(
+    proxy_logging, make_user_api_key_auth, monkeypatch, caplog
+):
+    monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data(stream=True)
+
+    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+        delivered = [
+            item
+            async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+                user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
+                response=_async_chunk_iter(_tool_call_stream_chunks()),
+                request_data=data,
+            )
+        ]
+
+    assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}'
+    assert delivered[1].choices[0].finish_reason == "tool_calls"
+    assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))
+
+
+@pytest.mark.asyncio
+async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_anthropic_sse(
+    proxy_logging, make_user_api_key_auth, monkeypatch, caplog
+):
+    monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data(stream=True)
+
+    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+        delivered = [
+            item
+            async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+                user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"),
+                response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()),
+                request_data=data,
+            )
+        ]
+
+    assert delivered == _anthropic_tool_use_sse_chunks()
+    assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))
+
+
+@pytest.mark.asyncio
+async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_responses_events(
+    proxy_logging, make_user_api_key_auth, monkeypatch, caplog
+):
+    monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data(stream=True)
+
+    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+        delivered = [
+            item
+            async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+                user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"),
+                response=_async_chunk_iter(_responses_function_call_events()),
+                request_data=data,
+            )
+        ]
+
+    assert delivered == _responses_function_call_events()
+    assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))