From a3ca8300e6097dcf1068a0dccbde6d8bbcb66d65 Mon Sep 17 00:00:00 2001 From: Jesus Velarde Date: Fri, 6 Nov 2026 17:16:15 -0800 Subject: [PATCH] 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 --- .../chat/guardrail_translation/handler.py | 9 ++- .../base_llm/guardrail_translation/utils.py | 32 +++++++++++ .../chat/guardrail_translation/handler.py | 18 +++--- .../guardrail_translation/handler.py | 49 ++++------------ .../unified_guardrail/unified_guardrail.py | 42 +++++++++----- .../test_streaming_buffer_until_moderated.py | 57 +++++++++++++++++++ 6 files changed, 140 insertions(+), 67 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index a6261026a40..b6363a6d5c9 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/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 diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 900a9c0712e..6550b6ae0b1 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/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) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 519e077919e..2b79aae7919 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/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) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 093cc46bc10..b01e6281aa6 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/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", diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 71f7e1841f3..f128b876af4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/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 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index cb33a0c1c35..1527a626698 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/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)