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.
2026-08-23 18:25:22 +00:00 · 2026-06-12 07:30:31 +00:00
parent 5aab7f5b9d
commit d8a1656e8b
2 changed files with 271 additions and 36 deletions
@@ -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
@@ -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()