From d8a1656e8bdacbd93e6a54d409a7d36b67074b2c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 07:30:31 +0000 Subject: [PATCH] fix(bedrock): scan non-text converse blocks for passthrough guardrails Key/team guardrails on bedrock converse passthrough only saw top-level text blocks, so a caller could hide prompt content in toolUse.input or toolResult.content[].json and have it forwarded to Bedrock without the configured guardrail inspecting it, bypassing blocking guardrails by default. Walk those arbitrary-JSON subtrees and write masked values back in place. Extend the non-streaming converse response path to the equivalent model-output fields (toolUse.input, reasoningContent text and citationsContent text) while leaving structural values such as reasoning signatures and citation sources untouched. --- .../guardrail_translation/handler.py | 117 ++++++++--- .../guardrail_translation/test_handler.py | 190 +++++++++++++++++- 2 files changed, 271 insertions(+), 36 deletions(-) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index f86bd9a7a3..8c59790daa 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation @@ -35,32 +35,60 @@ def _generic_passthrough_handler() -> BaseTranslation: return PassThroughEndpointHandler() +_StringHolder = Tuple[Union[dict, list], Union[str, int]] + + +def _collect_strings(node: Any, holders: List[_StringHolder]) -> None: + """ + Record a (container, key) holder for every non-empty string value nested + under an arbitrary JSON node, so prompt content a caller hides in fields + like ``toolUse.input`` or ``toolResult.content[].json`` is still scanned + and can be written back in place. + """ + if isinstance(node, dict): + for key, value in node.items(): + if isinstance(value, str): + if value: + holders.append((node, key)) + else: + _collect_strings(value, holders) + elif isinstance(node, list): + for index, value in enumerate(node): + if isinstance(value, str): + if value: + holders.append((node, index)) + else: + _collect_strings(value, holders) + + +def _collect_block_text(block: dict, holders: List[_StringHolder]) -> None: + text = block.get("text") + if isinstance(text, str) and text: + holders.append((block, "text")) + + def _extract_converse_texts( body: dict, skip_system: bool, skip_tool: bool, -) -> Tuple[List[str], List[dict]]: +) -> Tuple[List[str], List[_StringHolder]]: """ Walk a Bedrock Converse request body and collect text content. - Returns (texts, holders) where each holder is the dict that owns the - extracted ``text`` key, so write-back mutates it in place. Tool result - text lives under ``toolResult.content[].text`` rather than the top-level - block and is scanned too unless tool blocks are skipped. + Returns (texts, holders) where each holder is the (container, key) pair + that owns the extracted string, so write-back mutates it in place. Besides + top-level ``text`` blocks this scans the arbitrary-JSON fields a caller can + hide prompt content in -- ``toolUse.input`` and + ``toolResult.content[].json`` (alongside ``toolResult.content[].text``) -- + so a blocking guardrail sees them before the request reaches Bedrock. Tool + blocks are skipped entirely when tool messages are excluded. """ - texts: List[str] = [] - holders: List[dict] = [] - - def _collect(block: dict) -> None: - text = block.get("text") - if text: - texts.append(text) - holders.append(block) + holders: List[_StringHolder] = [] if not skip_system: for block in body.get("system") or []: if isinstance(block, dict): - _collect(block) + _collect_block_text(block, holders) for message in body.get("messages") or []: if not isinstance(message, dict): @@ -70,24 +98,62 @@ def _extract_converse_texts( continue if skip_tool and ("toolUse" in block or "toolResult" in block): continue - _collect(block) + _collect_block_text(block, holders) + tool_use = block.get("toolUse") + if isinstance(tool_use, dict): + _collect_strings(tool_use.get("input"), holders) tool_result = block.get("toolResult") if isinstance(tool_result, dict): for inner in tool_result.get("content") or []: if isinstance(inner, dict): - _collect(inner) + _collect_block_text(inner, holders) + _collect_strings(inner.get("json"), holders) + texts = [container[key] for container, key in holders] + return texts, holders + + +def _extract_converse_output_texts( + content_blocks: List[Any], +) -> Tuple[List[str], List[_StringHolder]]: + """ + Collect user-visible text from Bedrock Converse output content blocks. + + Covers ``text`` blocks plus the other content-bearing fields a model can + emit -- ``toolUse.input``, ``reasoningContent.reasoningText.text`` and + ``citationsContent.content[].text`` -- while leaving structural values such + as reasoning signatures and citation sources untouched. + """ + holders: List[_StringHolder] = [] + for block in content_blocks: + if not isinstance(block, dict): + continue + _collect_block_text(block, holders) + tool_use = block.get("toolUse") + if isinstance(tool_use, dict): + _collect_strings(tool_use.get("input"), holders) + reasoning = block.get("reasoningContent") + if isinstance(reasoning, dict): + reasoning_text = reasoning.get("reasoningText") + if isinstance(reasoning_text, dict): + _collect_block_text(reasoning_text, holders) + citations = block.get("citationsContent") + if isinstance(citations, dict): + for cited in citations.get("content") or []: + if isinstance(cited, dict): + _collect_block_text(cited, holders) + texts = [container[key] for container, key in holders] return texts, holders def _write_back_texts( guardrailed_texts: List[str], - holders: List[dict], + holders: List[_StringHolder], ) -> None: - for idx, holder in enumerate(holders): + for idx, (container, key) in enumerate(holders): if idx >= len(guardrailed_texts): break - holder["text"] = guardrailed_texts[idx] + container[key] = guardrailed_texts[idx] class BedrockPassthroughGuardrailHandler(BaseTranslation): @@ -308,12 +374,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): if not isinstance(content_blocks, list): return response - texts: List[str] = [] - text_indices: List[int] = [] - for i, block in enumerate(content_blocks): - if isinstance(block, dict) and "text" in block: - texts.append(block["text"]) - text_indices.append(i) + texts, holders = _extract_converse_output_texts(content_blocks) if not texts: return response @@ -345,8 +406,6 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): ) guardrailed_texts = guardrailed_inputs.get("texts", []) - for list_pos, block_idx in enumerate(text_indices): - if list_pos < len(guardrailed_texts): - content_blocks[block_idx]["text"] = guardrailed_texts[list_pos] + _write_back_texts(guardrailed_texts, holders) return response diff --git a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 51c16ac8fd..8e214d7ae1 100644 --- a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -74,8 +74,8 @@ class TestExtractConverseTexts: } texts, holders = _extract_converse_texts(body, skip_system=False, skip_tool=False) assert texts == ["sys text", "user text"] - assert holders[0] is body["system"][0] - assert holders[1] is body["messages"][0]["content"][0] + assert holders[0] == (body["system"][0], "text") + assert holders[1] == (body["messages"][0]["content"][0], "text") def test_skip_system(self): body = { @@ -84,7 +84,7 @@ class TestExtractConverseTexts: } texts, holders = _extract_converse_texts(body, skip_system=True, skip_tool=False) assert texts == ["user text"] - assert holders == [body["messages"][0]["content"][0]] + assert holders == [(body["messages"][0]["content"][0], "text")] def test_skip_tool_blocks(self): body = { @@ -107,7 +107,7 @@ class TestExtractConverseTexts: texts, _ = _extract_converse_texts(body, skip_system=False, skip_tool=True) assert texts == ["hello"] - def test_extracts_nested_tool_result_text(self): + def test_extracts_nested_tool_result_text_and_json(self): body = { "messages": [ { @@ -119,7 +119,7 @@ class TestExtractConverseTexts: "toolUseId": "1", "content": [ {"text": "blocked tool text"}, - {"json": {"k": "v"}}, + {"json": {"k": "blocked json value"}}, ], } }, @@ -128,8 +128,32 @@ class TestExtractConverseTexts: ] } texts, holders = _extract_converse_texts(body, skip_system=False, skip_tool=False) - assert texts == ["hello", "blocked tool text"] - assert holders[1] is body["messages"][0]["content"][1]["toolResult"]["content"][0] + assert texts == ["hello", "blocked tool text", "blocked json value"] + tool_content = body["messages"][0]["content"][1]["toolResult"]["content"] + assert holders[1] == (tool_content[0], "text") + assert holders[2] == (tool_content[1]["json"], "k") + + def test_extracts_tool_use_input_strings(self): + body = { + "messages": [ + { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "1", + "name": "lookup", + "input": {"query": "blocked input value", "limit": 5}, + } + } + ], + } + ] + } + texts, holders = _extract_converse_texts(body, skip_system=False, skip_tool=False) + assert texts == ["blocked input value"] + tool_use_input = body["messages"][0]["content"][0]["toolUse"]["input"] + assert holders[0] == (tool_use_input, "query") def test_non_text_content_blocks_ignored(self): body = { @@ -273,6 +297,72 @@ class TestBedrockPassthroughGuardrailHandlerInput: tool_result = result["data"]["messages"][0]["content"][2]["toolResult"] assert tool_result["content"][0]["text"] == "[REDACTED]" + @pytest.mark.asyncio + async def test_tool_result_json_scanned_and_masked(self): + """A caller can hide blocked text under toolResult.content[].json; the + guardrail must still see it and write the masked value back in place.""" + handler = BedrockPassthroughGuardrailHandler() + data = _converse_data() + data["data"]["messages"][0]["content"].append( + { + "toolResult": { + "toolUseId": "t1", + "content": [{"json": {"note": "SSN 123-45-6789"}}], + } + } + ) + guardrail = _make_guardrail( + {"texts": ["You are helpful.", "Hello world", "[REDACTED]"]} + ) + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] + assert "SSN 123-45-6789" in sent_texts + tool_result = result["data"]["messages"][0]["content"][2]["toolResult"] + assert tool_result["content"][0]["json"]["note"] == "[REDACTED]" + + @pytest.mark.asyncio + async def test_tool_use_input_scanned_and_masked(self): + """Blocked text hidden in toolUse.input must be scanned and masked.""" + handler = BedrockPassthroughGuardrailHandler() + data = _converse_data() + data["data"]["messages"][0]["content"][1]["toolUse"]["input"] = { + "query": "email john@example.com" + } + guardrail = _make_guardrail( + {"texts": ["You are helpful.", "Hello world", "[REDACTED]"]} + ) + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] + assert "email john@example.com" in sent_texts + tool_use = result["data"]["messages"][0]["content"][1]["toolUse"] + assert tool_use["input"]["query"] == "[REDACTED]" + + @pytest.mark.asyncio + async def test_tool_use_input_blocking_propagates(self): + """A blocking guardrail must reject content hidden in toolUse.input.""" + handler = BedrockPassthroughGuardrailHandler() + data = _converse_data() + data["data"]["messages"][0]["content"][1]["toolUse"]["input"] = { + "query": "blocked content" + } + guardrail = MagicMock() + guardrail.guardrail_name = "block-guard" + guardrail.skip_system_message_in_guardrail = False + guardrail.skip_tool_message_in_guardrail = False + guardrail.apply_guardrail = AsyncMock( + side_effect=HTTPException(status_code=400, detail="Blocked") + ) + + with pytest.raises(HTTPException): + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] + assert "blocked content" in sent_texts + @pytest.mark.asyncio async def test_non_converse_endpoint_scans_full_payload(self): """Invoke routes must not bypass guardrails: the full request payload is @@ -369,6 +459,92 @@ class TestBedrockPassthroughGuardrailHandlerOutput: assert result["output"]["message"]["content"][0]["text"] == "[MASKED]" assert result["stopReason"] == "end_turn" + @pytest.mark.asyncio + async def test_response_reasoning_and_tooluse_extracted_and_masked(self): + """Model output hidden in reasoningContent.reasoningText.text and + toolUse.input must be scanned and masked, but the reasoning signature + must be left untouched.""" + handler = BedrockPassthroughGuardrailHandler() + response = { + "output": { + "message": { + "role": "assistant", + "content": [ + {"text": "visible"}, + { + "reasoningContent": { + "reasoningText": { + "text": "thinking about john@example.com", + "signature": "sig-do-not-touch", + } + } + }, + { + "toolUse": { + "toolUseId": "1", + "name": "lookup", + "input": {"q": "ssn 123-45-6789"}, + } + }, + ], + } + }, + "stopReason": "end_turn", + } + guardrail = _make_guardrail( + {"texts": ["[V]", "[REASON]", "[INPUT]"]} + ) + + result = await handler.process_output_response( + response=response, guardrail_to_apply=guardrail + ) + + sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] + assert "thinking about john@example.com" in sent_texts + assert "ssn 123-45-6789" in sent_texts + blocks = result["output"]["message"]["content"] + assert blocks[0]["text"] == "[V]" + reasoning_text = blocks[1]["reasoningContent"]["reasoningText"] + assert reasoning_text["text"] == "[REASON]" + assert reasoning_text["signature"] == "sig-do-not-touch" + assert blocks[2]["toolUse"]["input"]["q"] == "[INPUT]" + + @pytest.mark.asyncio + async def test_response_citations_content_extracted_and_masked(self): + """citationsContent.content[].text is grounded answer text and must be + scanned, while citation sources/titles are left untouched.""" + handler = BedrockPassthroughGuardrailHandler() + response = { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "citationsContent": { + "content": [{"text": "Contact john@example.com"}], + "citations": [ + {"source": "https://example.com", "title": "Example"} + ], + } + } + ], + } + }, + "stopReason": "end_turn", + } + guardrail = _make_guardrail({"texts": ["[CITED]"]}) + + result = await handler.process_output_response( + response=response, guardrail_to_apply=guardrail + ) + + sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] + assert sent_texts == ["Contact john@example.com"] + citations = result["output"]["message"]["content"][0]["citationsContent"] + assert citations["content"][0]["text"] == "[CITED]" + assert citations["citations"][0]["source"] == "https://example.com" + assert citations["citations"][0]["title"] == "Example" + @pytest.mark.asyncio async def test_non_dict_response_returned_unchanged(self): handler = BedrockPassthroughGuardrailHandler()