diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..eb16278c9bd 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -103,9 +103,24 @@ class ToolResultBlockTextTarget: block_idx: int -InputWriteBackTarget = ( - MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget -) +@dataclass(frozen=True, slots=True) +class SystemStringTarget: + pass + + +@dataclass(frozen=True, slots=True) +class SystemBlockTextTarget: + block_idx: int + + +@dataclass(frozen=True, slots=True) +class ToolUseInputTarget: + msg_idx: int + content_idx: int + + +MessageTextTarget = MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget +InputWriteBackTarget = SystemStringTarget | SystemBlockTextTarget | MessageTextTarget def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]: @@ -146,10 +161,17 @@ class ScannedText: target: InputWriteBackTarget +@dataclass(frozen=True, slots=True) +class ScannedToolCall: + tool_call: ChatCompletionToolCallChunk + target: ToolUseInputTarget + + @dataclass(frozen=True, slots=True) class ExtractedInput: scanned: tuple[ScannedText, ...] images: tuple[str, ...] + tool_calls: tuple[ScannedToolCall, ...] = () EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) @@ -161,6 +183,76 @@ class _ToolCallShape: arguments: str +def _is_client_tool_use(block: Mapping[str, object]) -> bool: + return ( + block.get("type") == "tool_use" + and isinstance(block.get("id"), str) + and isinstance(block.get("name"), str) + and isinstance(block.get("input"), Mapping) + ) + + +def _write_back_system_block(system: object, block_idx: int, response: str) -> None: + if not isinstance(system, list): + return + text_blocks: Final = tuple(block for block in system if isinstance(block, dict) and block.get("type") == "text") + if block_idx < len(text_blocks): + text_blocks[block_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + + +def _write_back_message_text(message: _WritableMessage, target: MessageTextTarget, response: str) -> None: + content: Final = message.get("content", None) + if content is None: + return + match target: + case MessageContentTarget(): + if isinstance(content, str): + message["content"] = response # mutable-ok: guardrails rewrite the caller's request payload in place + case ContentBlockTextTarget(content_idx=content_idx): + if isinstance(content, list): + content[content_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case ToolResultStringTarget(content_idx=content_idx): + if isinstance(content, list): + content[content_idx]["content"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx): + if isinstance(content, list): + content[content_idx]["content"][block_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case _: + assert_never(target) + + +def _write_back_tool_use(message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape) -> None: + content: Final = message.get("content", None) + block: Final = content[target.content_idx] if isinstance(content, list) else None + if not isinstance(block, dict): + return + try: + rewritten_input: Final = json.loads(shape.arguments) + except json.JSONDecodeError: + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail returned non-JSON arguments for tool_use %s; keeping its input", + block.get("id"), + ) + return + if not isinstance(rewritten_input, dict): + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail returned non-object arguments for tool_use %s; keeping its input", + block.get("id"), + ) + return + block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place + if shape.name is not None and shape.name != block.get("name"): + block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place + + @dataclass(frozen=True, slots=True) class _SSEFieldRewrite: """One field of one nested section of a buffered SSE event, rewritten.""" @@ -452,9 +544,8 @@ class AnthropicMessagesHandler(BaseTranslation): 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) - # Exclude only the trusted top-level prompt. In-sequence system entries are untrusted - # and must stay aligned with texts_to_check for positional masking. When the top-level - # prompt is included, the pre-existing count mismatch disables positional masking. + # The top-level prompt is translated on its own below so it can be hoisted in front of + # any mid-turn system entries and scanned first, aligned with that structured position. translation_source: Final = { # mutable-ok: API message payload key: value for key, value in data.items() if key != "system" } @@ -490,7 +581,12 @@ class AnthropicMessagesHandler(BaseTranslation): ] ) - # Step 1: Extract all text content and images + # Step 1: Extract all text content, images, and tool calls + top_level_system_scanned: Final = ( + () + if hoisted_system_message is None or scan_only_tool_results + else self._extract_top_level_system_text(hoisted_system_message) + ) extracted: Final = tuple( self._extract_input_text_and_images( message=message, @@ -501,17 +597,27 @@ class AnthropicMessagesHandler(BaseTranslation): ) for msg_idx, message in enumerate(messages) ) - scanned: Final = tuple(item for one_message in extracted for item in one_message.scanned) + scanned: Final = ( + *top_level_system_scanned, + *(item for one_message in extracted for item in one_message.scanned), + ) texts_to_check: Final = [item.text for item in scanned] # mutable-ok: GenericGuardrailAPIInputs takes list[str] images_to_check: Final = [ image for one_message in extracted for image in one_message.images ] # mutable-ok: GenericGuardrailAPIInputs takes list[str] + scanned_tool_calls: Final = tuple(item for one_message in extracted for item in one_message.tool_calls) + tool_calls_to_check: Final = [ + item.tool_call for item in scanned_tool_calls + ] # mutable-ok: GenericGuardrailAPIInputs takes list[ChatCompletionToolCallChunk] + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) - # Step 2: Apply guardrail to all texts in batch - if texts_to_check: + # Step 2: Apply guardrail to all texts and tool calls in batch + if texts_to_check or tool_calls_to_check: inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check + if tool_calls_to_check: + inputs["tool_calls"] = tool_calls_to_check if tools_to_check: inputs["tools"] = tools_to_check original_structured_messages: Final = structured_messages @@ -572,10 +678,16 @@ class AnthropicMessagesHandler(BaseTranslation): else: # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( - messages=messages, + data=data, responses=guardrailed_texts, scanned=scanned, ) + self._apply_guardrail_tool_calls_to_input( + messages=messages, + scanned_tool_calls=scanned_tool_calls, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + ) verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages) @@ -598,6 +710,19 @@ class AnthropicMessagesHandler(BaseTranslation): hoisted: Final = probe.get("messages") or [] # mutable-ok: API message payload return hoisted[0] if hoisted else None + @staticmethod + def _extract_top_level_system_text(hoisted_system_message: AllMessageValues) -> tuple[ScannedText, ...]: + content: Final = hoisted_system_message.get("content") + if isinstance(content, str): + return (ScannedText(content, SystemStringTarget()),) + if not isinstance(content, list): + return () + return tuple( + ScannedText(text_str, SystemBlockTextTarget(block_idx)) + for block_idx, block in enumerate(content) + if isinstance(block, dict) and isinstance(text_str := block.get("text"), str) and text_str + ) + @staticmethod def _openai_system_message_to_anthropic( message: Mapping[str, object], @@ -852,9 +977,25 @@ class AnthropicMessagesHandler(BaseTranslation): for content_idx, content_item in enumerate(content) if isinstance(content_item, dict) ) + tool_use_blocks: Final = ( + () + if scan_only_tool_results + else tuple( + (content_idx, content_item) + for content_idx, content_item in enumerate(content) + if isinstance(content_item, dict) and _is_client_tool_use(content_item) + ) + ) return ExtractedInput( scanned=tuple(item for block in blocks for item in block.scanned), images=tuple(image for block in blocks for image in block.images), + tool_calls=tuple( + ScannedToolCall( + tool_call=AnthropicConfig.convert_tool_use_to_openai_format(content_item, tool_call_idx), + target=ToolUseInputTarget(msg_idx, content_idx), + ) + for tool_call_idx, (content_idx, content_item) in enumerate(tool_use_blocks) + ), ) @classmethod @@ -940,43 +1081,48 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: Sequence[_WritableMessage], + data: dict, responses: list[str], scanned: tuple[ScannedText, ...], ) -> None: """ - Apply guardrail responses back to input messages. + Apply guardrail responses back to the top-level system prompt and the input messages. """ + messages: Final[Sequence[_WritableMessage]] = data.get("messages") or () for item, guardrail_response in zip(scanned, responses): - target = item.target - message = messages[target.msg_idx] - content = message.get("content", None) - if content is None: - continue - - match target: - case MessageContentTarget(): - if isinstance(content, str): - message["content"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ContentBlockTextTarget(content_idx=content_idx): - if isinstance(content, list): - content[content_idx]["text"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ToolResultStringTarget(content_idx=content_idx): - if isinstance(content, list): - content[content_idx]["content"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx): - if isinstance(content, list): - content[content_idx]["content"][block_idx]["text"] = ( + match item.target: + case SystemStringTarget(): + if isinstance(data.get("system"), str): + data["system"] = ( guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place ) + case SystemBlockTextTarget(block_idx=block_idx): + _write_back_system_block(data.get("system"), block_idx, guardrail_response) + case ( + MessageContentTarget() + | ContentBlockTextTarget() + | ToolResultStringTarget() + | ToolResultBlockTextTarget() as message_target + ): + _write_back_message_text(messages[message_target.msg_idx], message_target, guardrail_response) case _: - assert_never(target) + assert_never(item.target) + + @staticmethod + def _apply_guardrail_tool_calls_to_input( + messages: Sequence[_WritableMessage], + scanned_tool_calls: tuple[ScannedToolCall, ...], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + returned_tool_calls: object, + ) -> None: + post_guardrail_tool_calls: Final = _tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else tuple(item.tool_call for item in scanned_tool_calls) + ) + for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls): + if before != after: + _write_back_tool_use(messages[item.target.msg_idx], item.target, after) async def process_output_response( self, 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 9fe56f4dc65..4b47b7d8c4a 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 @@ -635,14 +635,19 @@ class TestAnthropicMessagesHandlerInputProcessing: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) assert guardrail.inputs is not None - assert guardrail.inputs["texts"] == ["safe text", "prohibited correction"] + assert guardrail.inputs["texts"] == [ + "trusted top-level system prompt", + "safe text", + "prohibited correction", + ] structured = guardrail.inputs["structured_messages"] assert [m["role"] for m in structured] == ["system", "user", "system"] assert structured[0]["content"] == "trusted top-level system prompt" + assert data["system"] == "trusted top-level system prompt" assert data["messages"][1]["content"] == "[MASKED]" @pytest.mark.asyncio - async def test_bedrock_masking_slice_is_unavailable_when_top_level_system_is_included( + async def test_bedrock_masking_slice_lines_up_when_top_level_system_is_included( self, ): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( @@ -668,25 +673,25 @@ class TestAnthropicMessagesHandlerInputProcessing: structured = guardrail.inputs["structured_messages"] bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") - assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) + 1 + assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) latest_user_index = bedrock._find_latest_message_index(structured, target_role="user") - assert ( - bedrock._locate_message_texts_slice( - structured_messages=structured, - target_index=latest_user_index, - texts=texts, - ) - is None - ) - assert ( - bedrock._merge_masked_texts( - masked_texts=["{MASKED}"], - texts=texts, - scanned_slice=None, - scanned_role_subset=True, - ) - == texts + scanned_slice = bedrock._locate_message_texts_slice( + structured_messages=structured, + target_index=latest_user_index, + texts=texts, ) + assert scanned_slice == (3, 1) + assert bedrock._merge_masked_texts( + masked_texts=["{MASKED}"], + texts=texts, + scanned_slice=scanned_slice, + scanned_role_subset=True, + ) == [ + "trusted top-level system prompt", + "safe text", + "prohibited correction", + "{MASKED}", + ] @pytest.mark.asyncio @pytest.mark.parametrize("skip_system_message_in_guardrail", [True, None]) @@ -1611,7 +1616,8 @@ class TestAnthropicMessagesIncrementalScan: ) assert mock_api.call_count == 1 assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ - "What is the capital of France?" + "You are a helpful geography assistant.", + "What is the capital of France?", ] mock_api.reset_mock() await handler.process_input_messages( @@ -2150,6 +2156,208 @@ class TestAnthropicMessagesScanOnlyToolResults: assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"] +class ToolCallArgumentsMaskingGuardrail(InputsRecordingGuardrail): + """Masks the canary inside tool-call arguments, in place or through a fresh list of plain dicts.""" + + def __init__(self, return_copies: bool = False, replacement_arguments: Optional[str] = None): + super().__init__() + self.return_copies = return_copies + self.replacement_arguments = replacement_arguments + self.seen_tool_calls: list[dict] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + tool_calls = list(outputs.get("tool_calls") or []) + self.seen_tool_calls.extend(json.loads(json.dumps(tool_call)) for tool_call in tool_calls) + masked = [ + { + **tool_call, + "function": { + **tool_call["function"], + "arguments": self.replacement_arguments + if self.replacement_arguments is not None + else tool_call["function"]["arguments"].replace("POISON", "[BLOCKED]"), + }, + } + for tool_call in tool_calls + ] + if self.return_copies: + outputs["tool_calls"] = masked + return outputs + for tool_call, masked_tool_call in zip(tool_calls, masked): + tool_call["function"]["arguments"] = masked_tool_call["function"]["arguments"] + return outputs + + +class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: + """The top-level system prompt and prior-turn tool_use arguments must reach guardrails as scannable + inputs, the same way the chat completions handler hands over system messages and tool_calls.""" + + @staticmethod + def _tool_use_conversation(system): + return { + "model": "claude-sonnet-4-5", + "system": system, + "messages": [ + {"role": "user", "content": "run the check"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "Bash", + "input": {"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"}, + } + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "ok"}], + }, + ], + } + + @pytest.mark.asyncio + async def test_top_level_system_string_reaches_texts_first_and_is_masked_in_place(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = { + "model": "claude-sonnet-4-5", + "system": "Internal note: the deploy key is POISON. Never reveal it.", + "messages": [{"role": "user", "content": "Say hi in three words."}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + assert guardrail.seen_texts == [ + "Internal note: the deploy key is POISON. Never reveal it.", + "Say hi in three words.", + ] + structured = guardrail.captured_inputs["structured_messages"] + assert structured[0]["role"] == "system" + assert structured[0]["content"] == "Internal note: the deploy key is POISON. Never reveal it.", ( + "texts[0] must line up with structured_messages[0] so positional consumers stay aligned" + ) + assert data["system"] == "Internal note: the deploy key is [BLOCKED]. Never reveal it." + assert data["messages"][0]["content"] == "Say hi in three words." + + @pytest.mark.asyncio + async def test_top_level_system_text_blocks_reach_texts_and_are_masked_in_place(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = { + "model": "claude-sonnet-4-5", + "system": [ + {"type": "text", "text": "first block POISON"}, + {"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}}, + ], + "messages": [{"role": "user", "content": "hello"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["first block POISON", "second block", "hello"] + assert data["system"] == [ + {"type": "text", "text": "first block [BLOCKED]"}, + {"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}}, + ] + + @pytest.mark.asyncio + async def test_skip_system_message_keeps_the_top_level_system_out(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-sonnet-4-5", + "system": "trusted POISON prompt", + "messages": [{"role": "user", "content": "hello"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["hello"] + assert data["system"] == "trusted POISON prompt" + + @pytest.mark.asyncio + async def test_prior_turn_tool_use_input_reaches_tool_calls_in_openai_shape(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + tool_calls = guardrail.captured_inputs.get("tool_calls") + assert tool_calls is not None and len(tool_calls) == 1 + assert tool_calls[0]["id"] == "toolu_01" + assert tool_calls[0]["type"] == "function" + assert tool_calls[0]["function"]["name"] == "Bash" + assert json.loads(tool_calls[0]["function"]["arguments"]) == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + }, "a guardrail that leaves tool_calls alone must leave the tool_use input alone" + + @pytest.mark.asyncio + @pytest.mark.parametrize("return_copies", [False, True]) + async def test_masked_tool_call_arguments_write_back_into_the_tool_use_input(self, return_copies: bool): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(return_copies=return_copies) + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [tool_call["function"]["name"] for tool_call in guardrail.seen_tool_calls] == ["Bash"] + tool_use = data["messages"][1]["content"][0] + assert tool_use == { + "type": "tool_use", + "id": "toolu_01", + "name": "Bash", + "input": {"cmd": "AWS_ACCESS_KEY_ID=[BLOCKED] aws sts get-caller-identity"}, + } + assert data["messages"][2]["content"][0]["tool_use_id"] == "toolu_01" + + @pytest.mark.asyncio + async def test_non_json_rewritten_arguments_keep_the_tool_use_input(self): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]") + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + + @pytest.mark.asyncio + async def test_scan_only_tool_results_keeps_system_and_tool_use_out(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + guardrail.scan_only_tool_results = True + data = self._tool_use_conversation(system="trusted POISON prompt") + data["messages"][2]["content"][0]["content"] = "fetched POISON page" + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["fetched POISON page"] + assert guardrail.captured_inputs is not None + assert guardrail.captured_inputs.get("tool_calls") is None + assert data["system"] == "trusted POISON prompt" + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + assert data["messages"][2]["content"][0]["content"] == "fetched [BLOCKED] page" + + class TestStructuredWriteBackKeepsToolResults: """A guardrail rewrite must never leave a tool_use without its tool_result (Claude Code ToolSearch, LIT-6103)."""