fix(guardrails): fail closed on malformed tool-call rewrites and keep masking guardrails off undeliverable buffers

A buffered stream whose guardrail returns a tool call without a function now fails closed on the Responses,
chat, and Anthropic Messages translations through one shared check, instead of releasing masked text next to
the original tool arguments. A mask_response_content guardrail on a route whose translation cannot write
rewrites back into a buffered stream gets the old guard again: buffering is turned off with a warning rather
than withholding the stream and then replaying the unredacted chunks

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

@@ -61,6 +61,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     stream_item_field,
     stream_item_fingerprint,
     unappliable_request_rewrite,
+    undeliverable_tool_calls_reason,
 )
 from litellm.types.llms.anthropic import (
     ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX,
@@ -528,13 +529,11 @@ def _ended_stream_tool_calls(
     """A guardrail that adds or drops a tool call on an ended stream cannot be written back, block for block."""
     if returned_tool_calls is None:
         return streamed_tool_calls
-    if len(returned_tool_calls) != expected_count:
+    undeliverable_reason: Final = undeliverable_tool_calls_reason(returned_tool_calls, expected_count)
+    if undeliverable_reason is not None:
         from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
-        raise UndeliverableStreamRewrite(
-            guardrail_name,
-            f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried {expected_count}",
-        )
+        raise UndeliverableStreamRewrite(guardrail_name, undeliverable_reason)
     return returned_tool_calls
 
 

litellm/llms/base_llm/guardrail_translation/utils.py

@@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 from litellm.llms.custom_httpx.sse_lines import split_sse_lines
 from litellm.types.llms.anthropic import claude_hosted_web_tool_call
 from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
+from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage
 from litellm.types.responses.main import with_hosted_web_tool_arguments
 
@@ -132,6 +133,37 @@ def blocked_responses_api_usage(original_response: object) -> ResponseAPIUsage:
     )
 
 
+class ReturnedToolCallFunction(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    name: str | None = None
+    arguments: str = ""
+
+
+class _ReturnedToolCall(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    function: ReturnedToolCallFunction
+
+
+def returned_tool_call_function(tool_call: object) -> ReturnedToolCallFunction | None:
+    payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
+    try:
+        return _ReturnedToolCall.model_validate(payload).function
+    except ValidationError:
+        return None
+
+
+def undeliverable_tool_calls_reason(returned_tool_calls: Sequence[object] | None, scanned_count: int) -> str | None:
+    if returned_tool_calls is None:
+        return None
+    if len(returned_tool_calls) != scanned_count:
+        return f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried {scanned_count}"
+    if any(returned_tool_call_function(tool_call) is None for tool_call in returned_tool_calls):
+        return "the guardrail returned tool calls without a function name and arguments"
+    return None
+
+
 def stream_item_field(item: object, field: str) -> object | None:
     if isinstance(item, dict):
         return item.get(field)

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

@@ -57,6 +57,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     stream_item_fingerprint,
     stream_item_items,
     unappliable_request_rewrite,
+    undeliverable_tool_calls_reason,
 )
 from litellm.types.llms.openai import (
     AllMessageValues,
@@ -523,18 +524,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
 
             guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
             returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
-            if (
-                deliver_ended_stream_rewrites
-                and isinstance(returned_tool_calls, list)
-                and len(returned_tool_calls) != pre_guardrail_tool_call_count
-            ):
+            undeliverable_reason: Final = (
+                undeliverable_tool_calls_reason(returned_tool_calls, pre_guardrail_tool_call_count)
+                if deliver_ended_stream_rewrites and isinstance(returned_tool_calls, list)
+                else None
+            )
+            if undeliverable_reason is not None:
                 from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
-                raise UndeliverableStreamRewrite(
-                    guardrail_to_apply.guardrail_name or "unknown",
-                    f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried "
-                    f"{pre_guardrail_tool_call_count}",
-                )
+                raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown", undeliverable_reason)
             guardrailed_tool_calls: Final[list[dict[str, object]]] = (
                 cast(list[dict[str, object]], returned_tool_calls)
                 if isinstance(returned_tool_calls, list)

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

@@ -39,7 +39,7 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeVar, Union, cast
 
 from openai.types.responses.response_output_refusal import ResponseOutputRefusal
-from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
+from pydantic import BaseModel, TypeAdapter, ValidationError
 from typing_extensions import (
     ReadOnly,
     TypedDict,
@@ -65,12 +65,14 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     effective_skip_tool_message_for_guardrail,
     list_field,
     merge_guardrailed_scoped_messages,
+    returned_tool_call_function,
     role_out_of_guardrail_scope,
     scoped_structured_message_indices,
     stream_item_field,
     stream_item_fingerprint,
     stream_item_items,
     unappliable_request_rewrite,
+    undeliverable_tool_calls_reason,
     web_search_call,
     with_web_search_call_arguments,
 )
@@ -91,7 +93,6 @@ from litellm.responses.litellm_completion_transformation.transformation import (
     claude_web_search_call_tool_call,
     tool_output_to_chat_content,
 )
-from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.llms.openai import (
     AllMessageValues,
     BaseLiteLLMOpenAIResponseObject,
@@ -148,19 +149,6 @@ class _ToolCallShape(NamedTuple):
     arguments: str
 
 
-class _ToolCallFunctionFields(LiteLLMBaseModel):
-    model_config = ConfigDict(frozen=True)
-
-    name: str | None = None
-    arguments: str = ""
-
-
-class _ToolCallFields(LiteLLMBaseModel):
-    model_config = ConfigDict(frozen=True)
-
-    function: _ToolCallFunctionFields
-
-
 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", ""))
@@ -169,12 +157,8 @@ def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tupl
 
 
 def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None:
-    payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
-    try:
-        fields: Final = _ToolCallFields.model_validate(payload)
-    except ValidationError:
-        return None
-    return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments)
+    function: Final = returned_tool_call_function(tool_call)
+    return None if function is None else _ToolCallShape(name=function.name, arguments=function.arguments)
 
 
 def _post_guardrail_tool_call_shapes(
@@ -186,18 +170,15 @@ def _post_guardrail_tool_call_shapes(
 ) -> tuple[_ToolCallShape, ...]:
     if not pre_guardrail_tool_calls:
         return pre_guardrail_tool_calls
-    if (
-        deliver_ended_stream_rewrites
-        and returned_tool_calls is not None
-        and len(returned_tool_calls) != len(pre_guardrail_tool_calls)
-    ):
+    undeliverable_reason: Final = (
+        undeliverable_tool_calls_reason(returned_tool_calls, len(pre_guardrail_tool_calls))
+        if deliver_ended_stream_rewrites
+        else None
+    )
+    if undeliverable_reason is not None:
         from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
-        raise UndeliverableStreamRewrite(
-            guardrail_name or "unknown",
-            f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried "
-            f"{len(pre_guardrail_tool_calls)}",
-        )
+        raise UndeliverableStreamRewrite(guardrail_name or "unknown", undeliverable_reason)
     if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls):
         verbose_proxy_logger.warning(
             "OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, "
@@ -210,12 +191,6 @@ def _post_guardrail_tool_call_shapes(
     returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls)
     validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None)
     if len(validated_shapes) != len(returned_shapes):
-        if deliver_ended_stream_rewrites:
-            from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
-
-            raise UndeliverableStreamRewrite(
-                guardrail_name or "unknown", "the guardrail returned a tool call without a function name and arguments"
-            )
         verbose_proxy_logger.warning(
             "OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, "
             "leaving the tool call output items unchanged",

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -245,6 +245,14 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N
             data["litellm_metadata"] = user_metadata
 
 
+def _route_delivers_ended_stream_rewrites(
+    request_route: str | None, mappings: Mapping[CallTypes, type["BaseTranslation"]]
+) -> bool:
+    call_types: Final = get_call_types_for_route(request_route) if request_route is not None else None
+    translation: Final = mappings.get(call_types[0]) if call_types else None
+    return translation is not None and translation.delivers_ended_stream_rewrites
+
+
 async def _scan_ended_stream(
     endpoint_translation: "BaseTranslation",
     *,
@@ -1216,11 +1224,29 @@ class UnifiedLLMGuardrails(CustomLogger):
         # path; "incremental_diff" emits them as synthetic deltas (see
         # _run_incremental_transform_stream).
         streaming_transform_mode: Final = _streaming_setting("streaming_transform_mode", "block_only")
+        mappings: Final = load_guardrail_translation_mappings()
         # Withhold every chunk until end-of-stream moderation passes, then release the chunks
         # (with the guardrail's rewrite written into them) or only the block message.
-        buffer_until_moderated: Final = _streaming_setting(
+        buffering_requested: Final = _streaming_setting(
             "streaming_buffer_until_moderated", buffer_until_moderated_default
         )
+        undeliverable_masking_guardrail: Final = (
+            guardrail_to_apply
+            if buffering_requested
+            and isinstance(guardrail_to_apply, CustomGuardrail)
+            and guardrail_to_apply.mask_response_content
+            and not _route_delivers_ended_stream_rewrites(user_api_key_dict.request_route, mappings)
+            else None
+        )
+        if undeliverable_masking_guardrail is not None:
+            verbose_proxy_logger.warning(
+                "UnifiedLLMGuardrails: streaming_buffer_until_moderated is disabled for %s because "
+                "mask_response_content=True and route %s cannot write the masked output back into a buffered "
+                "stream, so buffered replay would release the unredacted original chunks",
+                undeliverable_masking_guardrail.guardrail_name,
+                user_api_key_dict.request_route,
+            )
+        buffer_until_moderated: Final = buffering_requested and undeliverable_masking_guardrail is None
         release_on_scan: Final = _streaming_setting("streaming_buffer_release_on_scan", False)
 
         if buffer_until_moderated and not release_on_scan:
@@ -1248,8 +1274,6 @@ class UnifiedLLMGuardrails(CustomLogger):
                 yield item
             return
 
-        mappings: Final = load_guardrail_translation_mappings()
-
         # Streaming text transformation (incremental_diff) diverges enough from the
         # block_only path that it runs as its own iterator. It requires a route we
         # can resolve up front to an OpenAI-chat handler (the only supported v1
@@ -1328,18 +1352,6 @@ class UnifiedLLMGuardrails(CustomLogger):
                     async for remaining_item in response:
                         yield remaining_item
                     return
-                if (
-                    chunk_counter == 1
-                    and buffer_until_moderated
-                    and guardrail_to_apply.mask_response_content
-                    and not mappings[CallTypes(call_type)]().delivers_ended_stream_rewrites
-                ):
-                    verbose_proxy_logger.warning(
-                        "UnifiedLLMGuardrails: guardrail %s masks response content, but the %s translation cannot write "
-                        "a rewrite back into a buffered stream, so the scan can only block or release the original chunks",
-                        guardrail_to_apply.guardrail_name,
-                        call_type,
-                    )
 
                 # If end_of_stream_only mode, yield chunks without processing.
                 # When buffering, withhold them instead -- they are released (or

tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py

@@ -966,3 +966,60 @@ def test_pipeline_live_streaming_resolves_the_flag_like_the_flat_streaming_path(
     monkeypatch.setattr(litellm, "callbacks", [guardrail])
 
     assert _pipelines_stream_live(_detect_only_pipeline(guardrail)) is streams_live
+
+
+async def _a2a_ndjson_stream(yielded_count: list[int], collected: list[Any]) -> AsyncGenerator[str, None]:
+    for text in ["Here is ", ORIGINAL_MARKER, " for you."]:
+        yielded_count.append(len(collected))
+        yield (
+            json.dumps(
+                {
+                    "jsonrpc": "2.0",
+                    "id": "1",
+                    "result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": text}]}},
+                }
+            )
+            + "\n"
+        )
+
+
+@pytest.mark.asyncio
+async def test_buffered_masking_guardrail_streams_live_on_a_translation_that_cannot_deliver_rewrites():
+    guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call", mask_response_content=True)
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = False
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/a2a/agent-1/message/stream")
+    request_data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": [guardrail.guardrail_name]},
+    }
+    collected: list[Any] = []
+    yielded_count: list[int] = []
+    async for chunk in unified.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=user_api_key_dict,
+        response=_a2a_ndjson_stream(yielded_count, collected),
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    assert len(collected) == 3
+    assert yielded_count == [0, 1, 2], "the stream was withheld although the rewrite cannot be delivered"
+
+
+@pytest.mark.asyncio
+async def test_buffered_chat_stream_withholds_everything_when_the_guardrail_returns_a_malformed_tool_call():
+    guardrail = _ToolManglingGuardrail(guardrail_name="mangler", event_hook="post_call")
+    with pytest.raises(HTTPException) as raised:
+        await _run_buffered_chat(guardrail, [f"call {ORIGINAL_MARKER}"], ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'])
+    assert raised.value.status_code == 500
+    assert "mangler" in str(raised.value.detail)
+
+
+@pytest.mark.asyncio
+async def test_buffered_messages_stream_withholds_everything_when_the_guardrail_returns_a_malformed_tool_call():
+    guardrail = _ToolManglingGuardrail(guardrail_name="mangler", event_hook="post_call")
+    with pytest.raises(HTTPException) as raised:
+        await _run_messages_stream(guardrail, _anthropic_tool_use_stream())
+    assert raised.value.status_code == 500
+    assert "mangler" in str(raised.value.detail)