Skip to content

fix(guardrails): scan tool results and enforce tool policy on passthrough streams #36000

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mateo-berri wants to merge 1 commit into from
Closed

fix(guardrails): scan tool results and enforce tool policy on passthrough streams #36000

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 111 additions & 35 deletions litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,10 @@
"""

import json
from typing import TYPE_CHECKING, Any, Final, cast
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast

from typing_extensions import assert_never

from litellm._logging import verbose_proxy_logger
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
Expand All @@ -22,10 +25,11 @@
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
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,
openai_messages_without_system,
openai_messages_without_tool,
filtered_structured_messages,
role_out_of_guardrail_scope,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
Expand Down Expand Up @@ -58,6 +62,18 @@
)


InputTextKind = Literal["message_string", "text_block", "tool_result_string", "tool_result_block"]


class InputTextLocation(NamedTuple):
"""Where one scanned text lives inside an Anthropic request, so masked text can be written back."""

kind: InputTextKind
msg_idx: int
content_idx: int = 0
block_idx: int = 0


class AnthropicMessagesHandler(BaseTranslation):
"""
Handler for processing Anthropic messages with guardrails.
Expand Down Expand Up @@ -278,22 +294,23 @@ async def process_input_messages(

skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)

chat_completion_compatible_request: Final = self._translate_to_openai(data)

structured_messages = cast(
list[AllMessageValues],
chat_completion_compatible_request.get("messages", []),
structured_messages: Final = 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,
)
if skip_system:
structured_messages = openai_messages_without_system(structured_messages)
if skip_tool:
structured_messages = openai_messages_without_tool(structured_messages)

texts_to_check: Final[list[str]] = []
images_to_check: Final[list[str]] = []
tools_to_check: Final[list[ChatCompletionToolParam]] = chat_completion_compatible_request.get("tools", [])
task_mappings: Final[list[tuple[int, int | None]]] = []
tools_to_check: Final[list[ChatCompletionToolParam]] = (
[] if scan_only_tool_results else chat_completion_compatible_request.get("tools", [])
)
task_mappings: Final[list[InputTextLocation]] = []

# Step 1: Extract all text content and images
for msg_idx, message in enumerate(messages):
Expand All @@ -305,6 +322,7 @@ async def process_input_messages(
task_mappings=task_mappings,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
scan_only_tool_results=scan_only_tool_results,
)

# Step 2: Apply guardrail to all texts in batch
Expand All @@ -314,9 +332,9 @@ async def process_input_messages(
inputs["images"] = images_to_check
if tools_to_check:
inputs["tools"] = tools_to_check
original_structured_messages: Final = structured_messages
if structured_messages:
inputs["structured_messages"] = structured_messages
original_structured_messages: Final = list(structured_messages)
if original_structured_messages:
inputs["structured_messages"] = original_structured_messages
# Include model information if available
model: Final = data.get("model")
if model:
Expand Down Expand Up @@ -411,19 +429,21 @@ def _extract_input_text_and_images(
msg_idx: int,
texts_to_check: list[str],
images_to_check: list[str],
task_mappings: list[tuple[int, int | None]],
task_mappings: list[InputTextLocation],
skip_system_message: bool = False,
skip_tool_message: bool = False,
scan_only_tool_results: bool = False,
) -> None:
"""
Extract text content and images from a message.

Override this method to customize text/image extraction logic.
"""
role: Final = str(message.get("role") or "").lower()
if skip_system_message and role == "system":
return
if skip_tool_message and role == "tool":
if role_out_of_guardrail_scope(
str(message.get("role") or "").lower(),
skip_system_message=skip_system_message,
skip_tool_message=skip_tool_message,
):
return

content: Final = message.get("content", None)
Expand All @@ -434,17 +454,34 @@ def _extract_input_text_and_images(
## CHECK FOR TEXT + IMAGES
if content is not None and isinstance(content, str):
# Simple string content
if scan_only_tool_results:
return
texts_to_check.append(content)
task_mappings.append((msg_idx, None))
task_mappings.append(InputTextLocation(kind="message_string", msg_idx=msg_idx))

elif content is not None and isinstance(content, list):
# List content (e.g., multimodal with text and images)
for content_idx, content_item in enumerate(content):
if content_item.get("type") == "tool_result":
for tool_result_text, location in self._extract_tool_result_text(
tool_result=content_item,
msg_idx=msg_idx,
content_idx=int(content_idx),
):
texts_to_check.append(tool_result_text)
task_mappings.append(location)
continue

if scan_only_tool_results:
continue

# Extract text
text_str = content_item.get("text", None)
if text_str is not None:
texts_to_check.append(text_str)
task_mappings.append((msg_idx, int(content_idx)))
task_mappings.append(
InputTextLocation(kind="text_block", msg_idx=msg_idx, content_idx=int(content_idx))
)

# Extract images
if content_item.get("type") == "image":
Expand All @@ -455,6 +492,40 @@ def _extract_input_text_and_images(
if data:
images_to_check.append(data)

@staticmethod
def _extract_tool_result_text(
tool_result: Mapping[str, Any],
msg_idx: int,
content_idx: int,
) -> tuple[tuple[str, InputTextLocation], ...]:
"""
Extract the text a tool returned to the model, paired with where it came from.

Anthropic tool_result blocks carry their payload under ``content`` (a string or a list of
blocks), never under ``text``, so tool output reaches the model unscanned unless it is
pulled out here. It is the least trusted content in an agent request.
"""
payload: Final = tool_result.get("content")
if isinstance(payload, str):
return ((payload, InputTextLocation(kind="tool_result_string", msg_idx=msg_idx, content_idx=content_idx)),)

if not isinstance(payload, list):
return ()

return tuple(
(
block["text"],
InputTextLocation(
kind="tool_result_block",
msg_idx=msg_idx,
content_idx=content_idx,
block_idx=int(block_idx),
),
)
for block_idx, block in enumerate(payload)
if isinstance(block, dict) and block.get("text") is not None
)

def _extract_input_tools(
self,
tools: list[dict[str, Any]],
Expand All @@ -475,29 +546,34 @@ async def _apply_guardrail_responses_to_input(
self,
messages: list[dict[str, Any]],
responses: list[str],
task_mappings: list[tuple[int, int | None]],
task_mappings: list[InputTextLocation],
) -> None:
"""
Apply guardrail responses back to input messages.

Override this method to customize how responses are applied.
"""
for task_idx, guardrail_response in enumerate(responses):
mapping = task_mappings[task_idx]
msg_idx = cast(int, mapping[0])
content_idx_optional = cast(int | None, mapping[1])

content = messages[msg_idx].get("content", None)
location = task_mappings[task_idx]
content = messages[location.msg_idx].get("content", None)
if content is None:
continue

if isinstance(content, str) and content_idx_optional is None:
# Replace string content with guardrail response
messages[msg_idx]["content"] = guardrail_response

elif isinstance(content, list) and content_idx_optional is not None:
# Replace specific text item in list content
messages[msg_idx]["content"][content_idx_optional]["text"] = guardrail_response
match location.kind:
case "message_string":
if isinstance(content, str):
messages[location.msg_idx]["content"] = guardrail_response
case "text_block":
if isinstance(content, list):
content[location.content_idx]["text"] = guardrail_response
case "tool_result_string":
if isinstance(content, list):
content[location.content_idx]["content"] = guardrail_response
case "tool_result_block":
if isinstance(content, list):
content[location.content_idx]["content"][location.block_idx]["text"] = guardrail_response
case _:
assert_never(location.kind)

async def process_output_response(
self,
Expand Down
55 changes: 49 additions & 6 deletions litellm/llms/base_llm/guardrail_translation/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import json
from collections.abc import Sequence
from typing import Any, Final

from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
Expand Down Expand Up @@ -113,13 +114,55 @@ def effective_skip_tool_message_for_guardrail(guardrail_to_apply: Any) -> bool:
return bool(getattr(litellm, "skip_tool_message_in_guardrail", False))


def _message_role(message: AllMessageValues) -> str:
return str((message or {}).get("role") or "").lower()


def openai_messages_without_system(
messages: list[AllMessageValues],
) -> list[AllMessageValues]:
return [m for m in messages if str((m or {}).get("role") or "").lower() != "system"]
messages: Sequence[AllMessageValues],
) -> tuple[AllMessageValues, ...]:
return tuple(m for m in messages if _message_role(m) != "system")


def openai_messages_without_tool(
messages: list[AllMessageValues],
) -> list[AllMessageValues]:
return [m for m in messages if str((m or {}).get("role") or "").lower() != "tool"]
messages: Sequence[AllMessageValues],
) -> tuple[AllMessageValues, ...]:
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


def role_out_of_guardrail_scope(
role: str,
*,
skip_system_message: bool,
skip_tool_message: bool,
scan_only_tool_results: bool = False,
) -> bool:
"""Whether a message role falls outside what this guardrail is configured to scan."""
if skip_system_message and role == "system":
return True
if skip_tool_message and role == "tool":
return True
return scan_only_tool_results and role != "tool"


def filtered_structured_messages(
messages: Sequence[AllMessageValues],
*,
scan_only_tool_results: bool,
skip_system: bool,
skip_tool: bool,
) -> tuple[AllMessageValues, ...]:
"""Narrow the structured messages a guardrail sees, per its skip/scope settings."""
scoped: Final = openai_messages_only_tool(messages) if scan_only_tool_results else tuple(messages)

@veria-ai veria-ai Bot Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Low: Tool-result guardrail bypass

When scan_only_tool_results and Bedrock's experimental_use_latest_role_message_only are both enabled, this leaves only tool messages, while Bedrock's selector requires a user message and returns skip_scan=True if none exists. An authenticated user can therefore submit malicious tool-result content without invoking the configured scanner. Preserve tool-result eligibility when latest-role filtering is enabled, and add a test covering these options together for both OpenAI and Anthropic requests.

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
33 changes: 21 additions & 12 deletions litellm/llms/openai/chat/guardrail_translation/handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,11 @@
StreamTransformSink,
)
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,
openai_messages_without_system,
openai_messages_without_tool,
filtered_structured_messages,
role_out_of_guardrail_scope,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
Expand Down Expand Up @@ -82,6 +83,7 @@ async def process_input_messages(

skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)

texts_to_check: Final[list[str]] = []
images_to_check: Final[list[str]] = []
Expand All @@ -101,6 +103,7 @@ async def process_input_messages(
tool_call_task_mappings=tool_call_task_mappings,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
scan_only_tool_results=scan_only_tool_results,
)

# Step 2: Apply guardrail to all texts and tool calls in batch
Expand All @@ -110,13 +113,16 @@ async def process_input_messages(
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
structured_messages = self.get_structured_messages(data)
structured_messages: Final = self.get_structured_messages(data)
if structured_messages:
if skip_system:
structured_messages = openai_messages_without_system(structured_messages)
if skip_tool:
structured_messages = openai_messages_without_tool(structured_messages)
inputs["structured_messages"] = 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,
)
)
# Pass tools (function definitions) to the guardrail
tools: Final = data.get("tools")
if tools:
Expand Down Expand Up @@ -194,16 +200,19 @@ def _extract_inputs(
tool_call_task_mappings: list[tuple[int, int]],
skip_system_message: bool = False,
skip_tool_message: bool = False,
scan_only_tool_results: bool = False,
) -> None:
"""
Extract text content, images, and tool calls from a message.

Override this method to customize text/image/tool call extraction logic.
"""
role: Final = str(message.get("role") or "").lower()
if skip_system_message and role == "system":
return
if skip_tool_message and role == "tool":
if role_out_of_guardrail_scope(
str(message.get("role") or "").lower(),
skip_system_message=skip_system_message,
skip_tool_message=skip_tool_message,
scan_only_tool_results=scan_only_tool_results,
):
return

content: Final = message.get("content", None)
Expand Down
Loading
Loading