From d70e10982a46f728c6d5a431fd8692a85b3ebf23 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:21:58 -0700 Subject: [PATCH] fix(guardrails): keep tool-results-only scans off function definitions and merge scoped write-backs Gate the OpenAI handler's tools forwarding behind scan_only_tool_results, matching the Anthropic handler, so a tool-results-only scan can no longer evaluate or rewrite trusted function definitions. When a guardrail returns a replacement structured_messages list, substitute the returned messages back into the positions their scoped originals came from instead of installing the scoped list as the whole conversation, so out-of-scope messages (system prompt, prior turns) survive redaction on both the OpenAI and Anthropic paths. --- .../chat/guardrail_translation/handler.py | 32 +++++--- .../base_llm/guardrail_translation/utils.py | 58 ++++++++++--- .../chat/guardrail_translation/handler.py | 26 +++--- .../test_anthropic_guardrail_handler.py | 54 +++++++++++++ .../test_openai_guardrail_handler.py | 81 +++++++++++++++++++ 5 files changed, 216 insertions(+), 35 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 535f4b7ae61..c25fa624f7f 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -29,7 +29,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, - filtered_structured_messages, + merge_guardrailed_scoped_messages, + scoped_structured_message_indices, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -330,17 +331,17 @@ class AnthropicMessagesHandler(BaseTranslation): chat_completion_compatible_request: Final = self._translate_to_openai(data) - structured_messages: Final = list( - filtered_structured_messages( - cast( - list[AllMessageValues], - chat_completion_compatible_request.get("messages", []), - ), - scan_only_tool_results=scan_only_tool_results, - skip_system=skip_system, - skip_tool=skip_tool, - ) + full_structured_messages: Final = cast( + list[AllMessageValues], + chat_completion_compatible_request.get("messages", []), ) + scoped_message_indices: Final = scoped_structured_message_indices( + full_structured_messages, + scan_only_tool_results=scan_only_tool_results, + skip_system=skip_system, + skip_tool=skip_tool, + ) + structured_messages: Final = [full_structured_messages[index] for index in scoped_message_indices] tools_to_check: Final[list[ChatCompletionToolParam]] = ( [] if scan_only_tool_results else chat_completion_compatible_request.get("tools", []) @@ -402,7 +403,14 @@ class AnthropicMessagesHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - self._write_back_structured_messages(data, guardrailed_structured_messages) + self._write_back_structured_messages( + data, + merge_guardrailed_scoped_messages( + full_messages=full_structured_messages, + scoped_indices=scoped_message_indices, + guardrailed_scoped=guardrailed_structured_messages, + ), + ) else: # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index e365913f2e1..fcd504fee2f 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from typing import Any, Final from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage @@ -130,12 +130,6 @@ def openai_messages_without_tool( return tuple(m for m in messages if _message_role(m) != "tool") -def openai_messages_only_tool( - messages: Sequence[AllMessageValues], -) -> tuple[AllMessageValues, ...]: - return tuple(m for m in messages if _message_role(m) == "tool") - - def effective_scan_only_tool_results_for_guardrail(guardrail_to_apply: Any) -> bool: return getattr(guardrail_to_apply, "scan_only_tool_results", None) is True @@ -154,13 +148,53 @@ def role_out_of_guardrail_scope( return scan_only_tool_results and role != "tool" -def filtered_structured_messages( +def scoped_structured_message_indices( messages: Sequence[AllMessageValues], *, scan_only_tool_results: bool, skip_system: bool, skip_tool: bool, -) -> tuple[AllMessageValues, ...]: - scoped: Final = openai_messages_only_tool(messages) if scan_only_tool_results else tuple(messages) - without_system: Final = openai_messages_without_system(scoped) if skip_system else scoped - return openai_messages_without_tool(without_system) if skip_tool else without_system +) -> tuple[int, ...]: + return tuple( + index + for index, message in enumerate(messages) + if not role_out_of_guardrail_scope( + _message_role(message), + skip_system_message=skip_system, + skip_tool_message=skip_tool, + scan_only_tool_results=scan_only_tool_results, + ) + ) + + +def merge_guardrailed_scoped_messages( + full_messages: Sequence[AllMessageValues], + scoped_indices: Sequence[int], + guardrailed_scoped: Sequence[AllMessageValues], +) -> list[AllMessageValues]: + """Substitute guardrail-returned messages back into the full conversation. + + Guardrails only ever see the scoped subset of messages, so a replacement + list they hand back describes that subset, not the whole request. Writing + it over ``data["messages"]`` wholesale would silently drop every + out-of-scope message (system prompt, prior turns). Instead, swap each + returned message into the position its scoped original came from; extra + returned messages land after the last scoped position, and scoped + originals without a counterpart are treated as removed by the guardrail. + When nothing was filtered out this degenerates to the returned list + itself, preserving wholesale-replacement behavior for unscoped guardrails. + """ + replacements: Final = dict(zip(scoped_indices, guardrailed_scoped)) + removed: Final = frozenset(scoped_indices[len(guardrailed_scoped) :]) + appended: Final = tuple(guardrailed_scoped[len(scoped_indices) :]) + last_scoped_index: Final = scoped_indices[-1] if scoped_indices else None + + def _merged() -> Iterator[AllMessageValues]: + for index, message in enumerate(full_messages): + if index in removed: + continue + yield replacements.get(index, message) + if index == last_scoped_index: + yield from appended + + return list(_merged()) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 67550890d2d..9d7fe6ce2a8 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -26,8 +26,9 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, - filtered_structured_messages, + merge_guardrailed_scoped_messages, role_out_of_guardrail_scope, + scoped_structured_message_indices, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -114,18 +115,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check structured_messages: Final = self.get_structured_messages(data) + scoped_message_indices: Final = scoped_structured_message_indices( + structured_messages or [], + scan_only_tool_results=scan_only_tool_results, + skip_system=skip_system, + skip_tool=skip_tool, + ) if structured_messages: - inputs["structured_messages"] = list( - filtered_structured_messages( - structured_messages, - scan_only_tool_results=scan_only_tool_results, - skip_system=skip_system, - skip_tool=skip_tool, - ) - ) + inputs["structured_messages"] = [structured_messages[index] for index in scoped_message_indices] # Pass tools (function definitions) to the guardrail tools: Final = data.get("tools") - if tools: + if tools and not scan_only_tool_results: inputs["tools"] = tools # Include model information if available model: Final = data.get("model") @@ -151,7 +151,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - data["messages"] = guardrailed_structured_messages + data["messages"] = merge_guardrailed_scoped_messages( + full_messages=structured_messages or [], + scoped_indices=scoped_message_indices, + guardrailed_scoped=guardrailed_structured_messages, + ) else: # Step 3: Map guardrail responses back to original message structure if guardrailed_texts and texts_to_check: diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index e90ae579d6d..a016e1a2deb 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -5,6 +5,7 @@ Tests the handler's ability to process streaming output for Anthropic Messages A with guardrail transformations, specifically testing edge cases with empty choices. """ +import json import os import sys from typing import Any, Literal, Optional @@ -778,12 +779,65 @@ class InputsRecordingGuardrail(MockMaskingGuardrail): return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) +class StructuredMessagesRewritingGuardrail(CustomGuardrail): + """Returns a new structured_messages list with a canary redacted, like redaction guardrails do.""" + + def __init__(self): + super().__init__(guardrail_name="structured-rewrite") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + structured = inputs.get("structured_messages") or [] + inputs["structured_messages"] = [ + json.loads(json.dumps(message).replace("POISON", "[BLOCKED]")) for message in structured + ] + return inputs + + class TestAnthropicMessagesScanOnlyToolResults: def _guardrail(self): guardrail = InputsRecordingGuardrail() guardrail.scan_only_tool_results = True return guardrail + @pytest.mark.asyncio + async def test_structured_write_back_merges_into_the_full_conversation(self): + handler = AnthropicMessagesHandler() + guardrail = StructuredMessagesRewritingGuardrail() + guardrail.scan_only_tool_results = True + data = { + "model": "claude-sonnet-4-5", + "system": "You are a careful agent harness.", + "messages": [ + {"role": "user", "content": "fetch the page"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu1", "name": "Bash", "input": {"cmd": "curl"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "tu1", "content": "fetched POISON page"}], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["system"] == "You are a careful agent harness." + assert [m["role"] for m in data["messages"]] == ["user", "assistant", "user"], ( + "a redacting guardrail must not strip out-of-scope turns from the request" + ) + serialized = json.dumps(data["messages"]) + assert "fetch the page" in serialized + assert "tool_use" in serialized + assert "fetched [BLOCKED] page" in serialized + assert "POISON" not in serialized + @pytest.mark.asyncio async def test_scan_narrows_to_tool_results_and_write_back_stays_aligned(self): handler = AnthropicMessagesHandler() diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index c8a1b98aa82..907da66e5bf 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1231,6 +1231,28 @@ class TestIncrementalScanRespectsSkipFlags: assert scanned == ["It is sunny in Paris.", "And tomorrow?"] +class StructuredRedactionGuardrail(CustomGuardrail): + """Captures inputs and returns a new structured_messages list with a canary redacted.""" + + def __init__(self): + super().__init__(guardrail_name="structured-redaction") + self.captured_inputs: Optional[GenericGuardrailAPIInputs] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.captured_inputs = inputs + structured = inputs.get("structured_messages") or [] + inputs["structured_messages"] = [ + {**m, "content": str(m.get("content", "")).replace("POISON", "[BLOCKED]")} for m in structured + ] + return inputs + + class TestScanOnlyToolResults: def _bedrock_guardrail(self): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail @@ -1297,3 +1319,62 @@ class TestScanOnlyToolResults: assert scanned == ["USER-PROMPT", "TOOL-RESULT"], ( "anything but an explicit True must leave the whole request in scope" ) + + @pytest.mark.parametrize("scan_only_tool_results", [True, False]) + @pytest.mark.asyncio + async def test_function_definitions_are_scoped_out_with_the_tool_results_flag(self, scan_only_tool_results): + handler = OpenAIChatCompletionsHandler() + guardrail = StructuredRedactionGuardrail() + guardrail.scan_only_tool_results = scan_only_tool_results + tools = [ + { + "type": "function", + "function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}}, + } + ] + data = { + "messages": [ + {"role": "user", "content": "read the report"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}, + ], + "tools": tools, + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + expected_tools = None if scan_only_tool_results else tools + assert guardrail.captured_inputs.get("tools") == expected_tools, ( + "function definitions must stay out of a tool-results-only scan" + ) + + @pytest.mark.asyncio + async def test_structured_write_back_keeps_out_of_scope_messages(self): + handler = OpenAIChatCompletionsHandler() + guardrail = StructuredRedactionGuardrail() + guardrail.scan_only_tool_results = True + data = { + "messages": [ + {"role": "system", "content": "SYSTEM-PROMPT"}, + {"role": "user", "content": "fetch the page"}, + { + "role": "assistant", + "content": "fetching", + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "fetch", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"}, + {"role": "user", "content": "and then?"}, + ] + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["system", "user", "assistant", "tool", "user"], ( + "a redacting guardrail must not strip out-of-scope messages from the request" + ) + assert data["messages"][0]["content"] == "SYSTEM-PROMPT" + assert data["messages"][3]["content"] == "page says [BLOCKED] here" + assert data["messages"][3]["tool_call_id"] == "call_1" + assert data["messages"][4]["content"] == "and then?"