diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 27d96e415fd..d98e6c6c911 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,16 +3,16 @@ "limit": 29204 }, "reportArgumentType": { - "limit": 2635 + "limit": 2634 }, "reportAssignmentType": { "limit": 329 }, "reportAttributeAccessIssue": { - "limit": 516 + "limit": 514 }, "reportCallIssue": { - "limit": 123 + "limit": 117 }, "reportConstantRedefinition": { "limit": 40 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9227 + "limit": 9225 }, "reportFunctionMemberAccess": { "limit": 7 @@ -99,34 +99,34 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45242 + "limit": 45145 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40340 + "limit": 39881 }, "reportUnknownParameterType": { - "limit": 20293 + "limit": 20258 }, "reportUnknownVariableType": { - "limit": 31796 + "limit": 31429 }, "reportUnnecessaryCast": { "limit": 122 }, "reportUnnecessaryComparison": { - "limit": 703 + "limit": 701 }, "reportUnnecessaryContains": { "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 865 + "limit": 864 }, "reportUntypedBaseClass": { - "limit": 72 + "limit": 0 }, "reportUntypedFunctionDecorator": { "limit": 33 diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 20f3aa430e9..2e91e082bd4 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -714,6 +714,29 @@ async def async_post_call_success_deployment_hook( return result + def supports_scan_only_tool_results(self) -> bool: + """Whether this guardrail can scan tool-result content. + + Guardrails whose own role filtering only ever scans human-authored + messages override this to return False, so configuring them with + ``scan_only_tool_results`` is rejected at initialization instead of + silently scanning nothing on every request. + """ + return True + + def structured_messages_cover_full_request(self) -> bool: + """Whether returned ``structured_messages`` span the whole request. + + Translation handlers hand guardrails only the in-scope subset of the + conversation and merge a returned ``structured_messages`` list back + into the full request. A guardrail that already rebuilds the complete + conversation itself (like CrowdStrike AIDR with its skip filters + active) overrides this to return True so the handler installs the + returned list as-is instead of merging it a second time, which would + duplicate the out-of-scope messages. + """ + return False + def should_run_guardrail( self, data, diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 3662389900b..88db9fae912 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -26,10 +26,13 @@ ) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( + anthropic_tool_name, + 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, + merge_guardrailed_scoped_messages, + merge_returned_tools_into_request_tools, + scoped_structured_message_indices, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -326,19 +329,25 @@ 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( + full_structured_messages: Final = cast( list[AllMessageValues], chat_completion_compatible_request.get("messages", []), ) - if skip_system: - structured_messages = openai_messages_without_system(structured_messages) - if skip_tool: - structured_messages = openai_messages_without_tool(structured_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]] = chat_completion_compatible_request.get("tools", []) + tools_to_check: Final[list[ChatCompletionToolParam]] = ( + [] if scan_only_tool_results else chat_completion_compatible_request.get("tools", []) + ) # Step 1: Extract all text content and images extracted: Final = tuple( @@ -347,6 +356,7 @@ async def process_input_messages( msg_idx=msg_idx, skip_system_message=skip_system, skip_tool_message=skip_tool, + scan_only_tool_results=scan_only_tool_results, ) for msg_idx, message in enumerate(messages) ) @@ -388,14 +398,31 @@ async def process_input_messages( if converted_tool is not None: anthropic_tools.append(converted_tool) # Note: MCP servers are handled separately in the main transformation - data["tools"] = anthropic_tools + data["tools"] = ( + merge_returned_tools_into_request_tools( + request_tools=data.get("tools"), + returned_tools=anthropic_tools, + tool_name=anthropic_tool_name, + ) + if scan_only_tool_results + else anthropic_tools + ) guardrailed_structured_messages: Final = guardrailed_inputs.get("structured_messages") if ( 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, + guardrailed_structured_messages + if guardrail_to_apply.structured_messages_cover_full_request() + else 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( @@ -461,6 +488,7 @@ def _extract_input_text_and_images( msg_idx: int, skip_system_message: bool = False, skip_tool_message: bool = False, + scan_only_tool_results: bool = False, ) -> ExtractedInput: """ Extract text content and images from a message. @@ -471,6 +499,8 @@ def _extract_input_text_and_images( content: Final = message.get("content", None) if isinstance(content, str): + if scan_only_tool_results: + return EMPTY_EXTRACTED_INPUT return ExtractedInput(scanned=(ScannedText(content, MessageContentTarget(msg_idx)),), images=()) if not isinstance(content, list): return EMPTY_EXTRACTED_INPUT @@ -481,6 +511,7 @@ def _extract_input_text_and_images( msg_idx=msg_idx, content_idx=content_idx, skip_tool_message=skip_tool_message, + scan_only_tool_results=scan_only_tool_results, ) for content_idx, content_item in enumerate(content) if isinstance(content_item, dict) @@ -497,12 +528,16 @@ def _extract_content_block( msg_idx: int, content_idx: int, skip_tool_message: bool, + scan_only_tool_results: bool = False, ) -> ExtractedInput: if content_item.get("type") == "tool_result": if skip_tool_message: return EMPTY_EXTRACTED_INPUT return cls._extract_tool_result(content_item=content_item, msg_idx=msg_idx, content_idx=content_idx) + if scan_only_tool_results: + return EMPTY_EXTRACTED_INPUT + text_str: Final = content_item.get("text", None) return ExtractedInput( scanned=( @@ -551,22 +586,6 @@ def _image_sources(block: Mapping[str, Any]) -> tuple[str, ...]: data: Final = source.get("data") return (data,) if data else () - def _extract_input_tools( - self, - tools: list[dict[str, Any]], - tools_to_check: list[ChatCompletionToolParam], - ) -> None: - """ - Extract tools from a message. - """ - ## CHECK FOR TOOLS - if tools is not None and isinstance(tools, list): - # TRANSFORM ANTHROPIC TOOLS TO OPENAI TOOLS - openai_tools: Final = self.adapter.translate_anthropic_tools_to_openai( - tools=cast(list[AllAnthropicToolsValues], tools) - ) - tools_to_check.extend(openai_tools) - async def _apply_guardrail_responses_to_input( self, messages: list[dict[str, Any]], diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 17cc0f118d6..f1ddf21cd3c 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,7 +1,8 @@ from __future__ import annotations import json -from typing import Any, Final +from collections.abc import Callable, Iterator, Sequence +from typing import Any, Final, TypeVar from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues @@ -113,13 +114,131 @@ 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], + 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: object) -> 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: + 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 not in ("tool", "function") + + +def scoped_structured_message_indices( + messages: Sequence[AllMessageValues], + *, + scan_only_tool_results: bool, + skip_system: bool, + skip_tool: bool, +) -> 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, + ) + ) + + +ToolT = TypeVar("ToolT") + + +def openai_tool_name(tool: object) -> str | None: + if not isinstance(tool, dict): + return None + function: Final = tool.get("function") + if isinstance(function, dict): + function_name: Final = function.get("name") + return function_name if isinstance(function_name, str) else None + flat_name: Final = tool.get("name") + return flat_name if isinstance(flat_name, str) else None + + +def anthropic_tool_name(tool: object) -> str | None: + name: Final = tool.get("name") if isinstance(tool, dict) else None + return name if isinstance(name, str) else None + + +def merge_returned_tools_into_request_tools( + request_tools: Sequence[ToolT] | None, + returned_tools: Sequence[ToolT], + tool_name: Callable[[ToolT], str | None], +) -> list[ToolT]: + """Union of the request's tools and guardrail-returned tools, keyed by name. + + Under ``scan_only_tool_results`` the guardrail never saw the request's + tools, so a returned list can neither replace them (it would drop every + user-defined function) nor be discarded (it may carry a tool the guardrail + synthesized and told the model to call, like Compresr's retrieve tool). + Keep every request tool and append only returned tools whose names aren't + already taken by a request tool or an earlier returned tool. + """ + originals: Final = tuple(request_tools or ()) + taken_names: Final = frozenset(name for tool in originals if (name := tool_name(tool)) is not None) + additions: Final = tuple( + tool + for index, tool in enumerate(returned_tools) + if (name := tool_name(tool)) not in taken_names + and (name is None or all(tool_name(earlier) != name for earlier in returned_tools[:index])) + ) + return [*originals, *additions] + + +def merge_guardrailed_scoped_messages( + full_messages: Sequence[AllMessageValues], + scoped_indices: Sequence[int], + guardrailed_scoped: Sequence[AllMessageValues], ) -> list[AllMessageValues]: - return [m for m in messages if str((m or {}).get("role") or "").lower() != "tool"] + """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 3988326f2c2..e411dc497fc 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -23,10 +23,14 @@ 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, + merge_guardrailed_scoped_messages, + merge_returned_tools_into_request_tools, + openai_tool_name, + role_out_of_guardrail_scope, + scoped_structured_message_indices, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -82,6 +86,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]] = [] @@ -101,6 +106,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 @@ -110,16 +116,18 @@ 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) + 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: - 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"] = [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") @@ -138,14 +146,30 @@ async def process_input_messages( guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", []) guardrailed_tools: Final = guardrailed_inputs.get("tools") if guardrailed_tools is not None: - data["tools"] = guardrailed_tools + data["tools"] = ( + merge_returned_tools_into_request_tools( + request_tools=tools, + returned_tools=guardrailed_tools, + tool_name=openai_tool_name, + ) + if scan_only_tool_results + else guardrailed_tools + ) guardrailed_structured_messages: Final = guardrailed_inputs.get("structured_messages") if ( guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - data["messages"] = guardrailed_structured_messages + data["messages"] = ( + guardrailed_structured_messages + if guardrail_to_apply.structured_messages_cover_full_request() + else 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: @@ -194,16 +218,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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index f7a5c7559b1..e9e729fb118 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -26,6 +26,9 @@ from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, +) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -402,6 +405,9 @@ def _collect_grounding_blocks(self, messages: list[AllMessageValues] | None) -> grounding.append(block) return grounding + def supports_scan_only_tool_results(self) -> bool: + return self.experimental_use_latest_role_message_only is not True + def _prepare_guardrail_messages_for_role( self, messages: list[AllMessageValues] | None, @@ -523,6 +529,11 @@ def _select_messages_for_apply_guardrail( latest_user_index: Final = self._find_latest_message_index(structured_messages, target_role="user") if latest_user_index is None: + if effective_scan_only_tool_results_for_guardrail(self): + verbose_proxy_logger.warning( + "Bedrock Guardrail: experimental_use_latest_role_message_only scans only the latest " + "user message, so scan_only_tool_results leaves nothing to scan for this request" + ) verbose_proxy_logger.debug("Bedrock Guardrail: no user-role message in request, skipping INPUT scan") return ApplyGuardrailMessageSelection(None, None, True, skip_scan=True) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 15ddd5e3458..b1bf9159607 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -362,6 +362,10 @@ def _extract_transformed_texts(self, guard_output: _GuardInput, num_assistant_me tail: Final = guard_output.messages[-num_assistant_messages:] if num_assistant_messages > 0 else [] return [_extract_text_from_message(msg) for msg in tail] + @override + def structured_messages_cover_full_request(self) -> bool: + return effective_skip_system_message_for_guardrail(self) or effective_skip_tool_message_for_guardrail(self) + def _writeback_messages( self, structured_messages: list[AllMessageValues], diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index ae1478a9210..13ced0ac06c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -22,6 +22,9 @@ CustomGuardrail, log_guardrail_information, ) +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -1561,6 +1564,9 @@ def _get_latest_user_text_indices( return scannable + def supports_scan_only_tool_results(self) -> bool: + return False + @staticmethod def _get_scannable_text_indices( texts: list[str], @@ -1716,6 +1722,15 @@ async def apply_guardrail( # - latest-user extraction returned None (no user / count mismatch) if scannable_indices is None: scannable_indices = self._get_scannable_text_indices(texts, structured_messages) + if ( + scannable_indices is not None + and not scannable_indices + and effective_scan_only_tool_results_for_guardrail(self) + ): + verbose_proxy_logger.warning( + "PANW Prisma AIRS scans only user, system, and developer messages, " + "so scan_only_tool_results leaves nothing to scan for this request" + ) for i, text in enumerate(texts): if not text or not text.strip(): diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 743ad888949..1a2c46f306c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -74,6 +74,9 @@ def __init__( super().__init__(**kwargs) + def supports_scan_only_tool_results(self) -> bool: + return self.check_tool_results + @log_guardrail_information async def apply_guardrail( self, diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index f77588cf087..9f70ed63dcb 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -14,6 +14,10 @@ from litellm._uuid import uuid from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, + effective_skip_tool_message_for_guardrail, +) from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, ) @@ -487,16 +491,27 @@ def initialize_guardrail( raise ValueError(f"Unsupported guardrail: {guardrail_type}") if custom_guardrail_callback is not None: - setattr( - custom_guardrail_callback, + for scoping_param in ( "skip_system_message_in_guardrail", - getattr(litellm_params, "skip_system_message_in_guardrail", None), - ) - setattr( - custom_guardrail_callback, "skip_tool_message_in_guardrail", - getattr(litellm_params, "skip_tool_message_in_guardrail", None), + "scan_only_tool_results", + ): + setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None)) + scan_only_tool_results_enabled: Final = effective_scan_only_tool_results_for_guardrail( + custom_guardrail_callback ) + if scan_only_tool_results_enabled and not custom_guardrail_callback.supports_scan_only_tool_results(): + raise ValueError( + f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results is enabled, but this " + "guardrail's role filtering never scans tool results, so no request content would ever " + "be scanned. Remove scan_only_tool_results or the guardrail's role-filtering option." + ) + if scan_only_tool_results_enabled and effective_skip_tool_message_for_guardrail(custom_guardrail_callback): + raise ValueError( + f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results and " + "skip_tool_message_in_guardrail are enabled together, which excludes every message from " + "scanning, so no request content would ever be scanned. Remove one of the two." + ) configured_run_in_parallel: Final = getattr(litellm_params, "run_in_parallel", None) if configured_run_in_parallel is not None: custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 60c3830fbef..6b354a39101 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -753,6 +753,16 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up ), ) + scan_only_tool_results: bool | None = Field( + default=None, + description=( + "When True, unified guardrails only evaluate tool results, the untrusted data an " + "agent feeds back into the model, and skip system, user, and assistant content. " + "Intended for agent harnesses whose own prompt scaffolding is trusted but often " + "trips prompt-attack detectors." + ), + ) + # Lakera specific params category_thresholds: LakeraCategoryThresholds | None = Field( default=None, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 421b424757b..65c98f6aab3 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -42,7 +42,7 @@ "limit": 81 }, "B010": { - "limit": 194 + "limit": 190 }, "B018": { "limit": 2 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 dff3390af12..c7a30f7f954 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 @@ with guardrail transformations, specifically testing edge cases with empty choices. """ +import json import os import sys from typing import Any, Literal, Optional @@ -760,3 +761,205 @@ async def test_tool_result_is_skipped_when_guardrail_skips_tool_messages(self): assert "skip me POISON" not in guardrail.seen_texts assert messages[1]["content"][0]["content"] == "skip me POISON" assert messages[0]["content"] == "keep me [BLOCKED]" + + +class InputsRecordingGuardrail(MockMaskingGuardrail): + def __init__(self): + super().__init__(guardrail_name="scan-only-capture") + 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 + 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() + guardrail = self._guardrail() + data = { + "model": "claude-sonnet-4-5", + "system": "You are a trusted agent harness with POISON heuristics.", + "tools": [ + { + "name": "Bash", + "description": "run a command", + "input_schema": {"type": "object", "properties": {}}, + } + ], + "messages": [ + {"role": "user", "content": "scaffolding POISON prompt"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu1", "name": "Bash", "input": {"cmd": "curl"}}], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "sibling POISON text"}, + {"type": "tool_result", "tool_use_id": "tu1", "content": "fetched POISON page"}, + ], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["fetched POISON page"], ( + "only the tool_result payload may reach the guardrail" + ) + assert guardrail.captured_inputs is not None + assert guardrail.captured_inputs.get("tools") is None + assert [m["role"] for m in guardrail.captured_inputs["structured_messages"]] == ["tool"] + assert data["messages"][2]["content"][1]["content"] == "fetched [BLOCKED] page" + assert data["messages"][0]["content"] == "scaffolding POISON prompt", ( + "out-of-scope content must come back untouched, not masked or dropped" + ) + assert data["messages"][2]["content"][0]["text"] == "sibling POISON text" + + @pytest.mark.asyncio + async def test_guardrail_synthesized_tools_are_appended_without_replacing_request_tools(self): + handler = AnthropicMessagesHandler() + guardrail = ToolAppendingGuardrail(guardrail_name="tool-appending") + guardrail.scan_only_tool_results = True + original_tools = [ + { + "name": "get_weather", + "description": "Get the weather at a specific location", + "input_schema": {"type": "object", "properties": {"location": {"type": "string"}}}, + } + ] + data = { + "model": "claude-sonnet-4-5", + "tools": original_tools, + "messages": [ + {"role": "user", "content": "what's the weather?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu1", "name": "get_weather", "input": {}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "tu1", "content": "sunny"}], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [t["name"] for t in data["tools"]] == ["get_weather", "injected_tool"], ( + "a tool the guardrail synthesized must reach the model, converted to Anthropic format, " + "without the request's own tools being replaced or dropped" + ) + assert data["tools"][0] == original_tools[0] + + @pytest.mark.asyncio + async def test_guardrail_is_not_called_when_the_request_has_no_tool_results(self): + handler = AnthropicMessagesHandler() + guardrail = self._guardrail() + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "What is 2 plus 2?"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is None + assert guardrail.seen_texts == [] + + @pytest.mark.asyncio + async def test_images_are_scoped_the_same_way_as_texts(self): + handler = AnthropicMessagesHandler() + guardrail = self._guardrail() + data = { + "model": "claude-sonnet-4-5", + "messages": [ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "base64", "data": "USER_IMG"}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu1", + "content": [ + {"type": "text", "text": "screenshot POISON"}, + {"type": "image", "source": {"type": "base64", "data": "TOOL_IMG"}}, + ], + } + ], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"] 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 7730b664c5e..2e75f29b1c5 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 @@ -1229,3 +1229,338 @@ async def test_second_turn_scans_only_new_eligible_content(self): assert mock_api.call_count == 1 scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] 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 ToolSynthesizingGuardrail(CustomGuardrail): + """Appends its own function tool to whatever tools it was given, like a + retrieval/recovery guardrail that injects a tool the model can later call.""" + + def __init__(self): + super().__init__(guardrail_name="tool-synthesizing") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + tools = list(inputs.get("tools") or []) + tools.append( + { + "type": "function", + "function": {"name": "injected_retrieve", "parameters": {"type": "object", "properties": {}}}, + } + ) + inputs["tools"] = tools + return inputs + + +class ToolNameCollidingGuardrail(CustomGuardrail): + """Returns a tool reusing a request tool's name plus a genuinely new tool.""" + + def __init__(self): + super().__init__(guardrail_name="tool-name-colliding") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + inputs["tools"] = [ + { + "type": "function", + "function": { + "name": "read_file", + "parameters": {"type": "object", "properties": {"hijacked": {"type": "string"}}}, + }, + }, + { + "type": "function", + "function": {"name": "injected_retrieve", "parameters": {"type": "object", "properties": {}}}, + }, + ] + return inputs + + +class DuplicateToolReturningGuardrail(CustomGuardrail): + """Returns the same synthesized tool name twice, second copy with a different schema.""" + + def __init__(self): + super().__init__(guardrail_name="duplicate-tool-returning") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + inputs["tools"] = [ + { + "type": "function", + "function": { + "name": "injected_retrieve", + "parameters": {"type": "object", "properties": {"first": {"type": "string"}}}, + }, + }, + { + "type": "function", + "function": { + "name": "injected_retrieve", + "parameters": {"type": "object", "properties": {"second": {"type": "string"}}}, + }, + }, + ] + return inputs + + +class TestScanOnlyToolResults: + def _bedrock_guardrail(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + guardrail = BedrockGuardrail( + guardrail_name="bedrock-scan-only-tool-results", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + ) + guardrail.scan_only_tool_results = True + return guardrail + + @pytest.mark.asyncio + async def test_only_tool_role_content_is_scanned(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + data = { + "messages": [ + {"role": "system", "content": "SYSTEM-PROMPT-not-scanned"}, + {"role": "user", "content": "USER-PROMPT-not-scanned"}, + { + "role": "assistant", + "content": "ASSISTANT-not-scanned", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "read_file", "arguments": '{"path": "report.html"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-scanned"}, + ] + } + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["TOOL-RESULT-scanned"] + + @pytest.mark.asyncio + async def test_legacy_function_role_results_are_scanned(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + data = { + "messages": [ + {"role": "user", "content": "USER-PROMPT-not-scanned"}, + {"role": "function", "name": "read_file", "content": "FUNCTION-RESULT-scanned"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-scanned"}, + ] + } + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["FUNCTION-RESULT-scanned", "TOOL-RESULT-scanned"], ( + "a tool result sent with the legacy function role must not bypass the scoped scan" + ) + + @pytest.mark.parametrize("flag_value", [None, "false", 0, object()]) + @pytest.mark.asyncio + async def test_scope_narrows_only_when_the_flag_is_actually_true(self, flag_value): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + guardrail.scan_only_tool_results = flag_value + data = { + "messages": [ + {"role": "user", "content": "USER-PROMPT"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}, + ] + } + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + 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.parametrize("scan_only_tool_results", [True, False]) + @pytest.mark.asyncio + async def test_guardrail_synthesized_tools_are_appended_without_replacing_request_tools( + self, scan_only_tool_results + ): + handler = OpenAIChatCompletionsHandler() + guardrail = ToolSynthesizingGuardrail() + guardrail.scan_only_tool_results = scan_only_tool_results + original_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": original_tools, + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"], ( + "a tool the guardrail synthesized (like a recovery/retrieve tool) must reach the model " + "without the request's own tools being replaced or dropped" + ) + assert data["tools"][0] == original_tools[0] + + @pytest.mark.asyncio + async def test_returned_tool_name_collisions_keep_the_request_schema(self): + handler = OpenAIChatCompletionsHandler() + guardrail = ToolNameCollidingGuardrail() + guardrail.scan_only_tool_results = True + original_read_file = { + "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": [original_read_file], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"] + assert data["tools"][0] == original_read_file, ( + "a returned tool reusing a request tool's name must not replace the request's schema" + ) + + @pytest.mark.asyncio + async def test_duplicate_returned_tool_names_keep_only_the_first(self): + handler = OpenAIChatCompletionsHandler() + guardrail = DuplicateToolReturningGuardrail() + guardrail.scan_only_tool_results = True + original_read_file = { + "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": [original_read_file], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"], ( + "two returned tools sharing a name must not both be forwarded to the provider" + ) + assert data["tools"][1]["function"]["parameters"]["properties"] == {"first": {"type": "string"}} + + @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?" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 65d6e33588f..76a695ce3fd 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3670,3 +3670,40 @@ async def test_moderation_hook_honors_the_mcp_event_type(mode, call_type, should "the scan must be logged under the event it actually ran for, so guardrail logs, " "OTel spans, and Langfuse metadata do not misclassify MCP enforcement as an LLM call" ) + + +class TestScanOnlyToolResultsWithLatestRoleFilter: + @pytest.mark.asyncio + async def test_warns_and_skips_when_scoped_payload_has_no_user_message(self): + """scan_only_tool_results hands Bedrock a tool-role-only payload, but + experimental_use_latest_role_message_only scans only the latest user + message: the silent no-op must warn.""" + guardrail = BedrockGuardrail( + guardrail_name="bedrock-latest-role-scoped", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + experimental_use_latest_role_message_only=True, + ) + guardrail.scan_only_tool_results = True + inputs = { + "texts": ["TOOL-RESULT"], + "structured_messages": [{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}], + } + + with ( + patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api, + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.warning" + ) as mock_warning, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"litellm_call_id": "test-call-id"}, + input_type="request", + ) + + mock_api.assert_not_called() + assert result["texts"] == ["TOOL-RESULT"] + warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args) + assert "scan_only_tool_results" in warning_text diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 431a7aa6f02..2f0fd51539d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -1696,6 +1696,34 @@ async def test_apply_guardrail_allow(self, handler): request_data=request_data, guardrail_name=handler.guardrail_name ) + @pytest.mark.asyncio + async def test_apply_guardrail_warns_when_tool_results_scope_leaves_nothing_scannable(self, handler): + """scan_only_tool_results hands PANW a tool-role-only payload, but PANW's role + filter only scans user/system/developer rows: the silent no-op must warn.""" + handler.scan_only_tool_results = True + inputs: GenericGuardrailAPIInputs = { + "texts": ["TOOL-RESULT"], + "structured_messages": [{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}], + } + request_data = {"litellm_call_id": "test-call-id", "model": "gpt-4"} + + with ( + patch.object(handler, "_call_panw_api", new_callable=AsyncMock) as mock_api, + patch( + "litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs.panw_prisma_airs.verbose_proxy_logger.warning" + ) as mock_warning, + ): + result = await handler.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + mock_api.assert_not_called() + assert result["texts"] == ["TOOL-RESULT"] + warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args) + assert "scan_only_tool_results" in warning_text + @pytest.mark.asyncio async def test_apply_guardrail_block(self, handler): """Test block action raises HTTPException(400).""" diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 6bd109f0f95..729dbce6b9a 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -558,3 +558,102 @@ def test_reinitialized_judge_guardrail_uses_lazy_router_provider(): finally: for cb_list, snapshot in zip(lists, snapshots): cb_list[:] = snapshot + + +class TestScanOnlyToolResultsInitRefusal: + """A guardrail whose role filtering never scans tool results must be rejected at + initialization when configured with scan_only_tool_results, instead of booting a + proxy that silently scans nothing on every request.""" + + def _initialize(self, name: str, params: dict): + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + return InMemoryGuardrailHandler().initialize_guardrail( + guardrail={"guardrail_name": name, "litellm_params": params}, + ) + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + def test_panw_prisma_airs_with_scan_only_tool_results_is_rejected(self): + with pytest.raises(ValueError, match="never scans tool results"): + self._initialize( + "panw-scan-only-combo", + { + "guardrail": "panw_prisma_airs", + "mode": "pre_call", + "api_key": "test-key", + "profile_name": "test-profile", + "scan_only_tool_results": True, + }, + ) + + def test_bedrock_latest_role_with_scan_only_tool_results_is_rejected(self): + with pytest.raises(ValueError, match="never scans tool results"): + self._initialize( + "bedrock-latest-role-scan-only-combo", + { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "gr-1", + "guardrailVersion": "1", + "experimental_use_latest_role_message_only": True, + "scan_only_tool_results": True, + }, + ) + + def test_bedrock_without_latest_role_accepts_scan_only_tool_results(self): + result = self._initialize( + "bedrock-scan-only-ok", + { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "gr-1", + "guardrailVersion": "1", + "scan_only_tool_results": True, + }, + ) + assert result is not None + + def test_prompt_security_default_tool_filtering_rejects_scan_only_tool_results(self, monkeypatch): + monkeypatch.delenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", raising=False) + with pytest.raises(ValueError, match="never scans tool results"): + self._initialize( + "prompt-security-scan-only-combo", + { + "guardrail": "prompt_security", + "mode": "pre_call", + "api_key": "test-key", + "api_base": "https://ps.example.com", + "scan_only_tool_results": True, + }, + ) + + def test_prompt_security_check_tool_results_accepts_scan_only_tool_results(self, monkeypatch): + monkeypatch.setenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", "true") + result = self._initialize( + "prompt-security-scan-only-ok", + { + "guardrail": "prompt_security", + "mode": "pre_call", + "api_key": "test-key", + "api_base": "https://ps.example.com", + "scan_only_tool_results": True, + }, + ) + assert result is not None + + def test_skip_tool_message_with_scan_only_tool_results_is_rejected(self): + with pytest.raises(ValueError, match="skip_tool_message_in_guardrail are enabled together"): + self._initialize( + "bedrock-skip-tool-scan-only-combo", + { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "gr-1", + "guardrailVersion": "1", + "skip_tool_message_in_guardrail": True, + "scan_only_tool_results": True, + }, + ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e26ce54ede7..8064e63f1aa 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23343 + "limit": 23332 }, "LIT002": { "limit": 27213 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1093 + "limit": 1091 }, "LIT007": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16802 + "limit": 16792 }, "LIT011": { "limit": 5602