fix(bedrock): apply Converse guardrail rewrites texts-first and keep main's write-backs

Map tool results onto their text and json parts only, accept a structured row the flat texts already reproduce, let a flat-text mask win over a structured leaf, ignore texts that do not line up with the request when structured rows come back, keep tool call arguments out of the rows under skip_tool_message_in_guardrail and send no tool definitions to tool-only guardrails

litellm/llms/bedrock/passthrough/guardrail_translation/handler.py

@@ -2,6 +2,7 @@ import copy
 import json
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
+from functools import partial, reduce
 from itertools import chain, groupby
 from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias
 
@@ -37,8 +38,8 @@ if TYPE_CHECKING:
     from litellm.llms.pass_through.guardrail_translation.handler import (
         PassThroughEndpointHandler,
     )
-    from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
     from litellm.proxy.utils import ProxyLogging
+    from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
 
 _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
 _EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream"
@@ -168,7 +169,7 @@ def _extract_converse_texts(
 
 
 class _ConverseModel(BaseModel):
-    model_config = ConfigDict(extra="ignore", frozen=True, populate_by_name=True)
+    model_config = ConfigDict(extra="ignore", frozen=True)
 
 
 class _ConverseGuardText(_ConverseModel):
@@ -263,7 +264,11 @@ class _ConvertedMessage:
 
 _IndexedBlock: TypeAlias = tuple[_BlockRef, _ConverseContentBlock]
 _TEXT_PARTS: Final = TypeAdapter(tuple[ChatCompletionTextObject, ...])
+_OBJECT: Final = TypeAdapter(Mapping[str, object])
+_ITEMS: Final = TypeAdapter(Sequence[object])
+_PARSED_JSON: Final = TypeAdapter(object)
 _TextEdit: TypeAlias = tuple[_BlockRef, object]
+_JsonPair: TypeAlias = tuple[tuple[str | int, ...], object, object]
 
 
 def _indexed_blocks(message_index: int | None, blocks: Sequence[_ConverseContentBlock]) -> tuple[_IndexedBlock, ...]:
@@ -282,15 +287,17 @@ def _refs(text_blocks: Sequence[tuple[_BlockRef, str]]) -> tuple[_BlockRef, ...]
     return tuple(ref for ref, _ in text_blocks)
 
 
+def _scannable_tool_parts(tool_result: _ConverseToolResult) -> tuple[tuple[int, _ConverseToolResultContent], ...]:
+    return tuple(
+        (index, part) for index, part in enumerate(tool_result.content) if part.text or part.json_value is not None
+    )
+
+
 def _converse_tool_message(tool_result: _ConverseToolResult) -> ChatCompletionToolMessage:
     return ChatCompletionToolMessage(
         role="tool",
         tool_call_id=tool_result.tool_use_id,
-        content="\n".join(
-            part.text or json.dumps(part.json_value)
-            for part in tool_result.content
-            if part.text or part.json_value is not None
-        ),
+        content="\n".join(part.text or json.dumps(part.json_value) for _, part in _scannable_tool_parts(tool_result)),
     )
 
 
@@ -328,10 +335,14 @@ def _converse_user_messages(blocks: Sequence[_IndexedBlock]) -> tuple[_Converted
     )
 
 
-def _converse_assistant_messages(blocks: Sequence[_IndexedBlock]) -> tuple[_ConvertedMessage, ...]:
+def _converse_assistant_messages(
+    blocks: Sequence[_IndexedBlock], include_tool_calls: bool
+) -> tuple[_ConvertedMessage, ...]:
     text_blocks: Final = _text_blocks(blocks)
     content: Final = list(_text_content(text_blocks))
-    tool_calls: Final = [_converse_tool_call(block.tool_use) for _, block in blocks if block.tool_use]
+    tool_calls: Final = [
+        _converse_tool_call(block.tool_use) for _, block in blocks if block.tool_use and include_tool_calls
+    ]
     if tool_calls:
         message: Final = ChatCompletionAssistantMessage(
             role="assistant", content=content or None, tool_calls=tool_calls
@@ -348,16 +359,20 @@ def _converse_assistant_messages(blocks: Sequence[_IndexedBlock]) -> tuple[_Conv
     return (_ConvertedMessage(ChatCompletionAssistantMessage(role="assistant", content=content), _refs(text_blocks)),)
 
 
-def _converse_message(message_index: int, message: _ConverseMessage) -> tuple[_ConvertedMessage, ...]:
+def _converse_message(
+    message_index: int, message: _ConverseMessage, include_tool_calls: bool
+) -> tuple[_ConvertedMessage, ...]:
     blocks: Final = _indexed_blocks(message_index, message.content)
     if message.role == "assistant":
-        return _converse_assistant_messages(blocks)
+        return _converse_assistant_messages(blocks, include_tool_calls)
     return _converse_user_messages(blocks)
 
 
-def _converse_structured_messages(request: _ConverseRequest) -> tuple[_ConvertedMessage, ...]:
+def _converse_structured_messages(request: _ConverseRequest, include_tool_calls: bool) -> tuple[_ConvertedMessage, ...]:
     converted: Final = tuple(
-        chain.from_iterable(_converse_message(index, message) for index, message in enumerate(request.messages))
+        chain.from_iterable(
+            _converse_message(index, message, include_tool_calls) for index, message in enumerate(request.messages)
+        )
     )
     system_blocks: Final = _text_blocks(_indexed_blocks(None, request.system))
     if not system_blocks:
@@ -366,143 +381,269 @@ def _converse_structured_messages(request: _ConverseRequest) -> tuple[_Converted
     return (_ConvertedMessage(system, _refs(system_blocks)), *converted)
 
 
+def _without(mapping: Mapping[str, object], key: str) -> tuple[tuple[str, object], ...]:
+    return tuple(sorted(((name, value) for name, value in mapping.items() if name != key), key=lambda item: item[0]))
+
+
 def _without_content(message: Mapping[str, object]) -> tuple[tuple[str, object], ...]:
-    return tuple(sorted(((key, value) for key, value in message.items() if key != "content"), key=lambda item: item[0]))
+    return _without(message, "content")
+
+
+def _fields(message: AllMessageValues) -> Mapping[str, object]:
+    return message
+
+
+def _as_object(value: object) -> Mapping[str, object] | None:
+    try:
+        return _OBJECT.validate_python(value)
+    except ValidationError:
+        return None
 
 
 def _returned_text_parts(message: Mapping[str, object]) -> tuple[str, ...] | None:
-    content = message.get("content")
+    content: Final = message.get("content")
     if isinstance(content, str):
         return (content,)
     try:
-        parts: Final = _TEXT_PARTS.validate_python(message.get("content"))
+        parts: Final = _TEXT_PARTS.validate_python(content)
     except ValidationError:
         return None
     return tuple(part["text"] for part in parts)
 
 
+def _merged_edits(children: Sequence[tuple[_TextEdit, ...] | None]) -> tuple[_TextEdit, ...] | None:
+    if any(child is None for child in children):
+        return None
+    return tuple(chain.from_iterable(child for child in children if child is not None))
+
+
+def _json_children(path: tuple[str | int, ...], original: object, returned: object) -> tuple[_JsonPair, ...] | None:
+    original_object: Final = _as_object(original)
+    returned_object: Final = _as_object(returned)
+    if original_object is not None and returned_object is not None:
+        if original_object.keys() != returned_object.keys():
+            return None
+        return tuple(((*path, key), value, returned_object[key]) for key, value in original_object.items())
+    original_items: Final = _as_items(original)
+    returned_items: Final = _as_items(returned)
+    if original_items is None or returned_items is None or len(original_items) != len(returned_items):
+        return None
+    return tuple(((*path, index), value, returned_items[index]) for index, value in enumerate(original_items))
+
+
 def _json_edits(
     ref: _BlockRef, path: tuple[str | int, ...], original: object, returned: object
 ) -> tuple[_TextEdit, ...] | None:
     """Only write changed leaves, preserving independent flat-text redactions."""
-    if original == returned:
-        return ()
-    if isinstance(original, dict) and isinstance(returned, dict) and original.keys() == returned.keys():
-        children = [_json_edits(ref, (*path, key), value, returned[key]) for key, value in original.items()]
-    elif isinstance(original, list) and isinstance(returned, list) and len(original) == len(returned):
-        children = [_json_edits(ref, (*path, index), value, returned[index]) for index, value in enumerate(original)]
-    elif isinstance(original, str) and isinstance(returned, str):
-        return ((_BlockRef(ref.message_index, ref.block_index, path), returned),)
-    else:
+    pending: Final[list[_JsonPair]] = [(path, original, returned)]  # mutable-ok: work stack instead of recursion
+    edits: Final[list[_TextEdit]] = []  # mutable-ok: collects the leaves the work stack visits
+    while pending:
+        leaf_path, before, after = pending.pop()
+        if before == after:
+            continue
+        if isinstance(before, str) and isinstance(after, str):
+            edits.append((_BlockRef(ref.message_index, ref.block_index, leaf_path), after))
+            continue
+        if (children := _json_children(leaf_path, before, after)) is None:
+            return None
+        pending.extend(children)
+    return tuple(edits)
+
+
+def _as_items(value: object) -> Sequence[object] | None:
+    if not isinstance(value, list):
         return None
-    if any(child is None for child in children):
+    try:
+        return _ITEMS.validate_python(value)
+    except ValidationError:
         return None
-    return tuple(edit for child in children if child is not None for edit in child)
+
+
+def _tool_part_edits(
+    ref: _BlockRef, index: int, part: _ConverseToolResultContent, text: str
+) -> tuple[_TextEdit, ...] | None:
+    if part.text:
+        return ((_BlockRef(ref.message_index, ref.block_index, ("toolResult", "content", index, "text")), text),)
+    try:
+        value: Final = _PARSED_JSON.validate_json(text)
+    except ValidationError:
+        return None
+    return _json_edits(ref, ("toolResult", "content", index, "json"), part.json_value, value)
 
 
 def _tool_result_edits(converted: _ConvertedMessage, returned: Mapping[str, object]) -> tuple[_TextEdit, ...] | None:
     if converted.tool_result is None or _without_content(returned) != _without_content(converted.message):
         return None
     ref, result = converted.tool_result
-    texts = _returned_text_parts(returned)
-    # A joined string cannot be split safely across multiple source blocks.
-    if texts is None or len(texts) != len(result.content):
+    texts: Final = _returned_text_parts(returned)
+    parts: Final = _scannable_tool_parts(result)
+    if texts is None or len(texts) != len(parts):
         return None
-    edits: list[_TextEdit] = []
-    for index, (part, text) in enumerate(zip(result.content, texts)):
-        if part.text:
-            if text != part.text:
-                edits.append(
-                    (_BlockRef(ref.message_index, ref.block_index, ("toolResult", "content", index, "text")), text)
-                )
-        elif part.json_value is not None:
-            try:
-                value = json.loads(text)
-            except ValueError:
-                return None
-            changes = _json_edits(ref, ("toolResult", "content", index, "json"), part.json_value, value)
-            if changes is None:
-                return None
-            edits.extend(changes)
-        else:
-            return None
-    return tuple(edits)
+    return _merged_edits(tuple(_tool_part_edits(ref, index, part, text) for (index, part), text in zip(parts, texts)))
+
+
+def _call_function(call: object) -> Mapping[str, object] | None:
+    call_object: Final = _as_object(call)
+    return _as_object(call_object.get("function")) if call_object is not None else None
+
+
+def _tool_call_edits(
+    tool_use: tuple[_BlockRef, _ConverseToolUse], before: object, after: object
+) -> tuple[_TextEdit, ...] | None:
+    if before == after:
+        return ()
+    before_call: Final = _as_object(before)
+    after_call: Final = _as_object(after)
+    before_function: Final = _call_function(before)
+    after_function: Final = _call_function(after)
+    if before_call is None or after_call is None or before_function is None or after_function is None:
+        return None
+    arguments: Final = after_function.get("arguments")
+    if (
+        not isinstance(arguments, str)
+        or _without(after_call, "function") != _without(before_call, "function")
+        or _without(after_function, "arguments") != _without(before_function, "arguments")
+    ):
+        return None
+    try:
+        value: Final = _PARSED_JSON.validate_json(arguments)
+    except ValidationError:
+        return None
+    ref, tool = tool_use
+    return _json_edits(ref, ("toolUse", "input"), tool.tool_input, value)
+
+
+def _tool_calls_edits(
+    tool_uses: Sequence[tuple[_BlockRef, _ConverseToolUse]], before: object, after: object
+) -> tuple[_TextEdit, ...] | None:
+    before_calls: Final = _as_items(before)
+    after_calls: Final = _as_items(after)
+    if before_calls is None or after_calls is None or len(before_calls) != len(after_calls):
+        return None
+    return _merged_edits(
+        tuple(
+            _tool_call_edits(tool_use, original, returned)
+            for tool_use, original, returned in zip(tool_uses, before_calls, after_calls)
+        )
+    )
+
+
+def _content_edits(converted: _ConvertedMessage, returned: Mapping[str, object]) -> tuple[_TextEdit, ...] | None:
+    original: Final = _fields(converted.message)
+    if returned.get("content") == original.get("content"):
+        return ()
+    texts: Final = _returned_text_parts(returned)
+    before_texts: Final = _returned_text_parts(original)
+    if texts is None or before_texts is None or len(texts) != len(converted.text_blocks):
+        return None
+    return tuple((ref, text) for ref, before, text in zip(converted.text_blocks, before_texts, texts) if text != before)
 
 
 def _message_text_edits(converted: _ConvertedMessage, returned: Mapping[str, object]) -> tuple[_TextEdit, ...] | None:
-    if not isinstance(returned, Mapping):
-        return None
     if returned == converted.message:
         return ()
     if converted.tool_result is not None:
         return _tool_result_edits(converted, returned)
-    edits: list[_TextEdit] = []
-    original = converted.message
-    comparable = dict(returned)
-    if converted.tool_uses:
-        calls = comparable.get("tool_calls")
-        original_calls = original.get("tool_calls")
-        if not isinstance(calls, list) or not isinstance(original_calls, list) or len(calls) != len(original_calls):
-            return None
-        for (ref, tool), before, after in zip(converted.tool_uses, original_calls, calls):
-            if before == after:
-                continue
-            if not isinstance(after, dict) or not isinstance(after.get("function"), dict):
-                return None
-            expected = copy.deepcopy(before)
-            expected["function"]["arguments"] = after["function"].get("arguments")
-            if expected != after or not isinstance(after["function"].get("arguments"), str):
-                return None
-            try:
-                arguments = json.loads(after["function"]["arguments"])
-            except ValueError:
-                return None
-            changes = _json_edits(ref, ("toolUse", "input"), tool.tool_input, arguments)
-            if changes is None:
-                return None
-            edits.extend(changes)
-        comparable["tool_calls"] = original_calls
-    if _without_content(comparable) != _without_content(original):
+    original_calls: Final = _fields(converted.message).get("tool_calls")
+    call_edits: Final = (
+        _tool_calls_edits(converted.tool_uses, original_calls, returned.get("tool_calls"))
+        if converted.tool_uses
+        else ()
+    )
+    comparable: Final = {**returned, "tool_calls": original_calls} if converted.tool_uses else returned
+    if call_edits is None or _without_content(comparable) != _without_content(converted.message):
         return None
-    if comparable.get("content") != original.get("content"):
-        texts = _returned_text_parts(comparable)
-        before_texts = _returned_text_parts(original)
-        if texts is None or before_texts is None or len(texts) != len(converted.text_blocks):
-            return None
-        edits.extend(
-            (ref, text) for ref, before, text in zip(converted.text_blocks, before_texts, texts) if text != before
-        )
-    return tuple(edits)
+    return _merged_edits((call_edits, _content_edits(converted, comparable)))
+
+
+def _leaf(body: Mapping[str, Any], ref: _BlockRef) -> _StringHolder:
+    blocks: Final = body["system"] if ref.message_index is None else body["messages"][ref.message_index]["content"]
+    block: Final = blocks[ref.block_index]
+    if not ref.path:
+        return (block, "text") if block.get("text") else (block["guardContent"]["text"], "text")
+    return reduce(lambda node, key: node[key], ref.path[:-1], block), ref.path[-1]
+
+
+def _leaf_identity(holder: _StringHolder) -> tuple[int, str | int]:
+    container, key = holder
+    return id(container), key
+
+
+def _changed_leaves(
+    holders: Sequence[_StringHolder], texts: Sequence[str], guardrailed_texts: Sequence[str]
+) -> frozenset[tuple[int, str | int]]:
+    return frozenset(
+        _leaf_identity(holder) for holder, before, after in zip(holders, texts, guardrailed_texts) if after != before
+    )
+
+
+def _texts_applied_view(
+    body: Mapping[str, Any], texts: Sequence[str], skip_system: bool, skip_tool: bool, scan_only_tool_results: bool
+) -> tuple[_ConvertedMessage, ...]:
+    rewritten: Final = copy.deepcopy(dict(body))
+    _, holders = _extract_converse_texts(rewritten, skip_system, skip_tool, scan_only_tool_results)
+    _write_back_texts(texts, holders)
+    return _converse_structured_messages(_ConverseRequest.model_validate(rewritten), not skip_tool)
+
+
+def _same_content(expected: Mapping[str, object], returned: Mapping[str, object]) -> bool:
+    if expected.get("content") == returned.get("content"):
+        return True
+    expected_parts: Final = _returned_text_parts(expected)
+    returned_parts: Final = _returned_text_parts(returned)
+    return (
+        expected_parts is not None
+        and returned_parts is not None
+        and "\n".join(expected_parts) == "\n".join(returned_parts)
+    )
+
+
+def _reproduces(rewritten: _ConvertedMessage, returned: Mapping[str, object] | None) -> bool:
+    """Whether the flat-text rewrite already yields the row the guardrail returned."""
+    return (
+        returned is not None
+        and _without_content(returned) == _without_content(rewritten.message)
+        and _same_content(rewritten.message, returned)
+    )
 
 
 def _structured_text_edits(
+    leaf: Callable[[_BlockRef], _StringHolder],
     converted: Sequence[_ConvertedMessage],
     scoped_indices: Sequence[int],
-    returned: Sequence[Mapping[str, object]],
-) -> tuple[_TextEdit, ...] | None:
+    returned: Sequence[object],
+    flat_changes: frozenset[tuple[int, str | int]],
+    flat_view: Callable[[], tuple[_ConvertedMessage, ...]],
+) -> tuple[tuple[_StringHolder, object], ...] | None:
+    """
+    Structured edits for the leaves the flat texts left unchanged.
+
+    A row whose rewrite cannot be mapped back onto its Converse blocks is accepted only when
+    the flat-text rewrite already produces it.
+    """
     if len(returned) != len(scoped_indices):
         return None
-    per_message: Final = tuple(
-        _message_text_edits(converted[index], message) for index, message in zip(scoped_indices, returned)
+    rows: Final = tuple(_as_object(message) for message in returned)
+    proposed: Final = tuple(
+        _message_text_edits(converted[index], row) if row is not None else None
+        for index, row in zip(scoped_indices, rows)
     )
-    if any(edits is None for edits in per_message):
-        return None
-    return tuple(chain.from_iterable(edits for edits in per_message if edits is not None))
+    unmapped: Final = tuple((index, row) for index, row, edits in zip(scoped_indices, rows, proposed) if edits is None)
+    if unmapped:
+        view: Final = flat_view()
+        if len(view) != len(converted):
+            return None
+        if not all(_reproduces(view[index], row) for index, row in unmapped):
+            return None
+    resolved: Final = tuple(
+        (leaf(ref), text) for ref, text in chain.from_iterable(edits for edits in proposed if edits is not None)
+    )
+    return tuple((holder, text) for holder, text in resolved if _leaf_identity(holder) not in flat_changes)
 
 
-def _write_back_structured_texts(body: Mapping[str, Any], edits: Sequence[_TextEdit]) -> None:
-    for ref, text in edits:
-        blocks = body["system"] if ref.message_index is None else body["messages"][ref.message_index]["content"]
-        block = blocks[ref.block_index]
-        if ref.path:
-            target = block
-            for key in ref.path[:-1]:
-                target = target[key]
-            target[ref.path[-1]] = text
-        elif "text" in block:
-            block["text"] = text
-        else:
-            block["guardContent"]["text"]["text"] = text
+def _write_back_structured_texts(edits: Sequence[tuple[_StringHolder, object]]) -> None:
+    for (container, key), text in edits:
+        container[key] = text
 
 
 def _converse_tools(request: _ConverseRequest) -> tuple[ChatCompletionToolParam, ...]:
@@ -555,8 +696,8 @@ def _extract_converse_output_texts(
 
 
 def _write_back_texts(
-    guardrailed_texts: list[str],
-    holders: list[_StringHolder],
+    guardrailed_texts: Sequence[str],
+    holders: Sequence[_StringHolder],
 ) -> None:
     if len(guardrailed_texts) < len(holders):
         verbose_proxy_logger.warning(
@@ -827,7 +968,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
         scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
         texts, holders = _extract_converse_texts(body, skip_system, skip_tool, scan_only_tool_results)
         converse_request: Final = _parse_converse_request(lambda: _ConverseRequest.model_validate(body))
-        converted: Final = _converse_structured_messages(converse_request) if converse_request else ()
+        converted: Final = _converse_structured_messages(converse_request, not skip_tool) if converse_request else ()
         structured_messages: Final = tuple(item.message for item in converted)
         scoped_message_indices: Final = scoped_structured_message_indices(
             structured_messages,
@@ -839,12 +980,12 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
         if not texts and not scoped_message_indices:
             return data
 
-        inputs: Final = GenericGuardrailAPIInputs(texts=texts)
+        inputs: Final = GenericGuardrailAPIInputs(texts=list(texts))
         if scoped_message_indices:
             inputs["structured_messages"] = copy.deepcopy(
                 [structured_messages[index] for index in scoped_message_indices]
             )
-        tools: Final = _converse_tools(converse_request) if converse_request else ()
+        tools: Final = _converse_tools(converse_request) if converse_request and not scan_only_tool_results else ()
         if tools:
             inputs["tools"] = list(tools)
         model: Final = data.get("model")
@@ -859,16 +1000,33 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
         )
 
         guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
-        edits: tuple[_TextEdit, ...] = ()
         guardrailed_messages: Final = guardrailed_inputs.get("structured_messages")
-        if guardrailed_messages is not None:
-            proposed_edits = _structured_text_edits(converted, scoped_message_indices, guardrailed_messages)
-            if proposed_edits is None:
-                raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
-            edits = proposed_edits
-        if guardrailed_texts:
-            _write_back_texts(guardrailed_texts, holders)
-        _write_back_structured_texts(body, edits)
+        sent_messages: Final = [structured_messages[index] for index in scoped_message_indices]
+        rewritten_rows: Final = (
+            guardrailed_messages
+            if sent_messages and isinstance(guardrailed_messages, list) and guardrailed_messages != sent_messages
+            else None
+        )
+        applied_texts: Final = (
+            texts if rewritten_rows is not None and len(guardrailed_texts) != len(texts) else guardrailed_texts
+        )
+        edits: Final = (
+            ()
+            if rewritten_rows is None
+            else _structured_text_edits(
+                partial(_leaf, body),
+                converted,
+                scoped_message_indices,
+                rewritten_rows,
+                _changed_leaves(holders, texts, applied_texts),
+                partial(_texts_applied_view, body, applied_texts, skip_system, skip_tool, scan_only_tool_results),
+            )
+        )
+        if edits is None:
+            raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
+        if applied_texts:
+            _write_back_texts(applied_texts, holders)
+        _write_back_structured_texts(edits)
 
         return data
 

tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py

@@ -0,0 +1,451 @@
+from __future__ import annotations
+
+import json
+import re
+from collections.abc import Callable, Iterator, Mapping
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+import yaml
+from pydantic import JsonValue, TypeAdapter
+
+from tests.integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, list_value, object_value
+from tests.integration._support.process import owned_proxy
+from tests.integration._support.wire import Reply, Request, Wire, wire_server
+
+_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
+_MODEL: Final = "converse-guarded"
+_CARD: Final = "4111 1111 1111 1111"
+_CARD_PATTERN: Final = re.compile(r"\b(?:\d{4} ){3}\d{4}\b")
+_MASK: Final = "****"
+_FLAGGED: Final = "flagged-by-moderation"
+_GUARD_PATH: Final = "/beta/litellm_basic_guardrail_api"
+_PNG: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
+_JSON: Final = TypeAdapter(JsonValue)
+_UNAPPLIABLE: Final = "rewrote the request in a way this endpoint cannot apply"
+_GENERIC_MODES: Final = (
+    "masksm",
+    "masksm_flat",
+    "smonly",
+    "sm_partial",
+    "echo_sm",
+    "mask",
+    "block",
+    "rowcount",
+)
+
+
+@dataclass(frozen=True, slots=True)
+class _Rig:
+    gateway: Gateway
+    vendor: Wire
+    bedrock: Wire
+    keys: Mapping[str, str]
+
+
+def _masked(value: JsonValue) -> JsonValue:
+    return _JSON.validate_json(_CARD_PATTERN.sub(_MASK, json.dumps(value)))
+
+
+def _flattened(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+    content: Final = row.get("content")
+    if not isinstance(content, list):
+        return dict(row)
+    parts: Final = tuple(object_value(part).get("text") for part in content)
+    return {**row, "content": "\n".join(part for part in parts if isinstance(part, str))}
+
+
+def _normalised_first_part(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+    content: Final = row.get("content")
+    if not isinstance(content, list) or not content:
+        return dict(row)
+    first: Final = object_value(content[0])
+    text: Final = first.get("text")
+    if not isinstance(text, str):
+        return dict(row)
+    return {**row, "content": [{**first, "text": text.strip() + " "}, *content[1:]]}
+
+
+def _generic_verdict(mode: str, body: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+    texts: Final = list_value(body.get("texts") or [])
+    rows: Final = tuple(object_value(row) for row in list_value(body.get("structured_messages") or []))
+    masked_texts: Final = _masked(list(texts))
+    masked_rows: Final = list_value(_masked([dict(row) for row in rows]))
+    if _CARD_PATTERN.search(json.dumps(body)) is None:
+        return {"action": "NONE"}
+    verdicts: Final[Mapping[str, Callable[[], dict[str, JsonValue]]]] = {
+        "masksm": lambda: {"action": "GUARDRAIL_INTERVENED", "texts": masked_texts, "structured_messages": masked_rows},
+        "masksm_flat": lambda: {
+            "action": "GUARDRAIL_INTERVENED",
+            "texts": masked_texts,
+            "structured_messages": [_flattened(object_value(row)) for row in masked_rows],
+        },
+        "smonly": lambda: {"action": "GUARDRAIL_INTERVENED", "structured_messages": masked_rows},
+        "sm_partial": lambda: {
+            "action": "GUARDRAIL_INTERVENED",
+            "texts": masked_texts,
+            "structured_messages": [_normalised_first_part(row) for row in rows],
+        },
+        "echo_sm": lambda: {
+            "action": "GUARDRAIL_INTERVENED",
+            "texts": masked_texts,
+            "structured_messages": [dict(row) for row in rows],
+        },
+        "mask": lambda: {"action": "GUARDRAIL_INTERVENED", "texts": masked_texts},
+        "block": lambda: {"action": "BLOCKED", "blocked_reason": "card number"},
+        "rowcount": lambda: {
+            "action": "GUARDRAIL_INTERVENED",
+            "texts": masked_texts,
+            "structured_messages": masked_rows[:-1],
+        },
+    }
+    return verdicts[mode]()
+
+
+def _vendor(request: Request) -> Reply:
+    body: Final = JSON_OBJECT.validate_json(request.body)
+    if request.target == "/moderation/moderations":
+        flagged: Final = _FLAGGED in str(body.get("input"))
+        result: Final = {
+            "flagged": flagged,
+            "categories": {"violence": flagged},
+            "category_scores": {"violence": 0.99 if flagged else 0.0},
+            "category_applied_input_types": {"violence": ["text"]},
+        }
+        return Reply(
+            body=json.dumps({"id": "modr-scripted", "model": "omni-moderation-latest", "results": [result]}).encode()
+        )
+    if request.target == "/promptguard/api/v1/guard":
+        messages: Final = body.get("messages")
+        if _CARD_PATTERN.search(json.dumps(messages)) is None:
+            return Reply(body=json.dumps({"decision": "allow"}).encode())
+        return Reply(body=json.dumps({"decision": "redact", "redacted_messages": _masked(messages)}).encode())
+    mode: Final = request.target.removeprefix("/gga/").removesuffix(_GUARD_PATH)
+    return Reply(body=json.dumps(_generic_verdict(mode, body)).encode())
+
+
+def _bedrock(request: Request) -> Reply:
+    assert request.target == f"/model/{_MODEL_ID}/converse", request.target
+    return Reply(
+        body=json.dumps(
+            {
+                "output": {"message": {"role": "assistant", "content": [{"text": "guarded reply"}]}},
+                "stopReason": "end_turn",
+                "usage": {"inputTokens": 7, "outputTokens": 2, "totalTokens": 9},
+                "metrics": {"latencyMs": 1},
+            }
+        ).encode()
+    )
+
+
+def _guardrail(name: str, params: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+    return {"guardrail_name": name, "litellm_params": {"mode": "pre_call", "default_on": False, **params}}
+
+
+def _config(directory: Path, vendor: Wire, bedrock: Wire) -> Path:
+    configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    configuration["model_list"] = [
+        *configuration["model_list"],
+        {
+            "model_name": _MODEL,
+            "litellm_params": {
+                "model": f"bedrock/{_MODEL_ID}",
+                "api_base": bedrock.url,
+                "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
+                "aws_secret_access_key": "scripted-secret",
+                "aws_region_name": "us-east-1",
+            },
+        },
+    ]
+    generic: Final = {"guardrail": "generic_guardrail_api", "api_key": "synthetic-guardrail-key"}
+    configuration["guardrails"] = [
+        *(_guardrail(f"gga-{mode}", {**generic, "api_base": f"{vendor.url}/gga/{mode}"}) for mode in _GENERIC_MODES),
+        _guardrail(
+            "gga-block-skip-tool",
+            {**generic, "api_base": f"{vendor.url}/gga/block", "skip_tool_message_in_guardrail": True},
+        ),
+        _guardrail(
+            "gga-mask-tool-only",
+            {**generic, "api_base": f"{vendor.url}/gga/mask", "scan_only_tool_results": True},
+        ),
+        _guardrail(
+            "promptguard",
+            {"guardrail": "promptguard", "api_base": f"{vendor.url}/promptguard", "api_key": "synthetic-pg-key"},
+        ),
+        _guardrail(
+            "moderation",
+            {"guardrail": "openai_moderation", "api_base": f"{vendor.url}/moderation", "api_key": "synthetic-key"},
+        ),
+    ]
+    path: Final = directory / "converse-guardrails.yaml"
+    path.write_text(yaml.safe_dump(configuration))
+    return path
+
+
+@pytest.fixture(scope="module")
+def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
+    directory: Final = tmp_path_factory.mktemp("converse-guardrail-rewrites")
+    with (
+        gateway_from_environment() as gateway,
+        wire_server(_vendor) as vendor,
+        wire_server(_bedrock) as bedrock,
+        owned_proxy(gateway, directory, {}, config=_config(directory, vendor, bedrock), workers=2) as owned,
+        owned.scenario() as scenario,
+    ):
+        names: Final = (
+            *(f"gga-{mode}" for mode in _GENERIC_MODES),
+            "gga-block-skip-tool",
+            "gga-mask-tool-only",
+            "promptguard",
+            "moderation",
+        )
+        yield _Rig(owned, vendor, bedrock, {name: scenario.key(guardrails=[name]) for name in names})
+
+
+@pytest.fixture(autouse=True)
+def _drained(rig: _Rig) -> Iterator[None]:
+    rig.vendor.drain()
+    rig.bedrock.drain()
+    yield
+    rig.vendor.drain()
+    rig.bedrock.drain()
+
+
+def _converse(rig: _Rig, guardrail: str, body: Mapping[str, JsonValue]) -> httpx.Response:
+    return rig.gateway.request("POST", f"/bedrock/model/{_MODEL}/converse", body, key=rig.keys[guardrail])
+
+
+def _forwarded(rig: _Rig, response: httpx.Response) -> dict[str, JsonValue]:
+    assert response.status_code == 200, response.text
+    sent: Final = rig.bedrock.drain()
+    assert len(sent) == 1, response.text
+    return JSON_OBJECT.validate_json(sent[0].body)
+
+
+def _refused(rig: _Rig, response: httpx.Response) -> str:
+    assert response.status_code >= 400, response.text
+    assert rig.bedrock.drain() == (), response.text
+    return response.text
+
+
+def _vendor_bodies(rig: _Rig) -> tuple[dict[str, JsonValue], ...]:
+    return tuple(JSON_OBJECT.validate_json(request.body) for request in rig.vendor.drain())
+
+
+def _user(*blocks: JsonValue) -> dict[str, JsonValue]:
+    return {"role": "user", "content": list(blocks)}
+
+
+def _assistant(*blocks: JsonValue) -> dict[str, JsonValue]:
+    return {"role": "assistant", "content": list(blocks)}
+
+
+def _tool_use(card: str) -> dict[str, JsonValue]:
+    return {"toolUse": {"toolUseId": "tooluse_lookup", "name": "lookup_card", "input": {"card": card, "limit": 2}}}
+
+
+def _tool_result(*content: JsonValue) -> dict[str, JsonValue]:
+    return {"toolResult": {"toolUseId": "tooluse_lookup", "content": list(content), "status": "success"}}
+
+
+def test_multi_block_user_message_with_flattened_vendor_reply_reaches_bedrock_masked(rig: _Rig) -> None:
+    body: Final = {"messages": [_user({"text": "Order details below."}, {"text": f"Card {_CARD} please."})]}
+    sent: Final = _forwarded(rig, _converse(rig, "gga-masksm_flat", body))
+    assert sent["messages"] == [_user({"text": "Order details below."}, {"text": f"Card {_MASK} please."})]
+
+
+def test_multi_block_system_prompt_with_flattened_vendor_reply_reaches_bedrock_masked(rig: _Rig) -> None:
+    body: Final = {
+        "system": [{"text": "You are a billing assistant."}, {"text": f"Card on file {_CARD}."}],
+        "messages": [_user({"text": "What card is on file?"})],
+    }
+    sent: Final = _forwarded(rig, _converse(rig, "gga-masksm_flat", body))
+    assert sent["system"] == [{"text": "You are a billing assistant."}, {"text": f"Card on file {_MASK}."}]
+    assert sent["messages"] == body["messages"]
+
+
+def test_tool_result_with_text_and_json_parts_reaches_bedrock_masked(rig: _Rig) -> None:
+    body: Final = {
+        "messages": [
+            _user({"text": "Look up my card."}),
+            _assistant(_tool_use("last four 1111")),
+            _user(_tool_result({"text": f"found {_CARD}"}, {"json": {"card": _CARD, "status": "active"}})),
+        ]
+    }
+    sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body))
+    assert sent["messages"] == [
+        _user({"text": "Look up my card."}),
+        _assistant(_tool_use("last four 1111")),
+        _user(_tool_result({"text": f"found {_MASK}"}, {"json": {"card": _MASK, "status": "active"}})),
+    ]
+
+
+def test_tool_result_with_text_and_image_parts_reaches_bedrock_masked_with_image_intact(rig: _Rig) -> None:
+    image: Final = {"image": {"format": "png", "source": {"bytes": _PNG}}}
+    body: Final = {
+        "messages": [
+            _user({"text": "Read the screenshot."}),
+            _assistant(_tool_use("screenshot")),
+            _user(_tool_result({"text": f"screenshot shows {_CARD}"}, image)),
+        ]
+    }
+    sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body))
+    assert sent["messages"] == [
+        _user({"text": "Read the screenshot."}),
+        _assistant(_tool_use("screenshot")),
+        _user(_tool_result({"text": f"screenshot shows {_MASK}"}, image)),
+    ]
+
+
+def test_flat_text_mask_wins_over_an_unmasked_structured_leaf(rig: _Rig) -> None:
+    body: Final = {"messages": [_user({"text": f"Card {_CARD} on file."}, {"text": "Thanks."})]}
+    sent: Final = _forwarded(rig, _converse(rig, "gga-sm_partial", body))
+    assert sent["messages"] == [_user({"text": f"Card {_MASK} on file."}, {"text": "Thanks."})]
+
+
+def test_promptguard_redaction_keeps_every_turn_in_place(rig: _Rig) -> None:
+    body: Final = {
+        "system": [{"text": "You are a support bot."}],
+        "messages": [
+            _user({"text": "My name is Ann."}),
+            _assistant({"text": "Hi Ann, how can I help?"}),
+            _user({"text": f"Repeat exactly: card {_CARD} ok"}),
+        ],
+    }
+    sent: Final = _forwarded(rig, _converse(rig, "promptguard", body))
+    assert sent["system"] == [{"text": "You are a support bot."}]
+    assert sent["messages"] == [
+        _user({"text": "My name is Ann."}),
+        _assistant({"text": "Hi Ann, how can I help?"}),
+        _user({"text": f"Repeat exactly: card {_MASK} ok"}),
+    ]
+
+
+def test_moderation_blocks_a_flagged_latest_user_turn_before_bedrock(rig: _Rig) -> None:
+    body: Final = {
+        "system": [{"text": "You are a support bot."}],
+        "messages": [
+            _user({"text": "Say hi."}),
+            _assistant({"text": "Hi."}),
+            _user({"text": f"please {_FLAGGED} now"}),
+        ],
+    }
+    refused: Final = _refused(rig, _converse(rig, "moderation", body))
+    assert "Violated OpenAI moderation policy" in refused
+
+
+def test_skip_tool_messages_keeps_tool_call_arguments_away_from_the_guardrail(rig: _Rig) -> None:
+    body: Final = {
+        "messages": [
+            _user({"text": "Check my card."}),
+            _assistant(_tool_use(_CARD)),
+            _user(_tool_result({"text": "card is active"})),
+        ]
+    }
+    sent: Final = _forwarded(rig, _converse(rig, "gga-block-skip-tool", body))
+    assert sent["messages"] == body["messages"]
+    assert all(_CARD not in json.dumps(scanned) for scanned in _vendor_bodies(rig))
+
+
+def test_tool_only_guardrail_receives_no_tool_definitions(rig: _Rig) -> None:
+    body: Final = {
+        "messages": [
+            _user({"text": f"My card is {_CARD}."}),
+            _assistant(_tool_use("lookup")),
+            _user(_tool_result({"text": f"found {_CARD}"})),
+        ],
+        "toolConfig": {
+            "tools": [
+                {
+                    "toolSpec": {
+                        "name": "lookup_card",
+                        "description": "Look up a card",
+                        "inputSchema": {"json": {"type": "object", "properties": {"card": {"type": "string"}}}},
+                    }
+                }
+            ]
+        },
+    }
+    sent: Final = _forwarded(rig, _converse(rig, "gga-mask-tool-only", body))
+    scanned: Final = _vendor_bodies(rig)
+    assert len(scanned) == 1
+    assert scanned[0].get("tools") is None, scanned[0]
+    assert scanned[0]["texts"] == [f"found {_CARD}"]
+    assert sent["messages"] == [
+        _user({"text": f"My card is {_CARD}."}),
+        _assistant(_tool_use("lookup")),
+        _user(_tool_result({"text": f"found {_MASK}"})),
+    ]
+
+
+def test_structured_only_vendor_mask_reaches_bedrock(rig: _Rig) -> None:
+    body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]}
+    sent: Final = _forwarded(rig, _converse(rig, "gga-smonly", body))
+    assert sent["messages"] == [_user({"text": f"Card {_MASK} please."})]
+
+
+def test_guard_content_turn_reaches_bedrock_masked(rig: _Rig) -> None:
+    body: Final = {"messages": [_user({"text": "Check this."}, {"guardContent": {"text": {"text": f"Card {_CARD}."}}})]}
+    sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body))
+    assert sent["messages"] == [_user({"text": "Check this."}, {"guardContent": {"text": {"text": f"Card {_MASK}."}}})]
+
+
+def test_tool_call_arguments_in_history_reach_bedrock_masked(rig: _Rig) -> None:
+    body: Final = {
+        "messages": [
+            _user({"text": "Check my card."}),
+            _assistant(_tool_use(_CARD)),
+            _user(_tool_result({"text": "card is active"})),
+        ]
+    }
+    sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body))
+    assert sent["messages"] == [
+        _user({"text": "Check my card."}),
+        _assistant(_tool_use(_MASK)),
+        _user(_tool_result({"text": "card is active"})),
+    ]
+
+
+def test_echoed_structured_rows_keep_the_text_mask(rig: _Rig) -> None:
+    body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]}
+    sent: Final = _forwarded(rig, _converse(rig, "gga-echo_sm", body))
+    assert sent["messages"] == [_user({"text": f"Card {_MASK} please."})]
+
+
+def test_single_block_string_reply_reaches_bedrock_masked(rig: _Rig) -> None:
+    body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]}
+    sent: Final = _forwarded(rig, _converse(rig, "gga-masksm_flat", body))
+    assert sent["messages"] == [_user({"text": f"Card {_MASK} please."})]
+
+
+def test_missing_structured_row_is_refused_before_bedrock(rig: _Rig) -> None:
+    body: Final = {"messages": [_user({"text": "Hello."}), _assistant({"text": "Hi."}), _user({"text": f"{_CARD}"})]}
+    assert _UNAPPLIABLE in _refused(rig, _converse(rig, "gga-rowcount", body))
+
+
+def test_blocked_verdict_never_reaches_bedrock(rig: _Rig) -> None:
+    body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]}
+    _refused(rig, _converse(rig, "gga-block", body))
+
+
+def test_concurrent_requests_keep_their_own_rewrites(rig: _Rig) -> None:
+    cards: Final = tuple(f"4111 1111 1111 {index:04d}" for index in range(24))
+    bodies: Final = tuple(
+        {"messages": [_user({"text": f"request {index}"}, {"text": f"card {card}"})]}
+        for index, card in enumerate(cards)
+    )
+    with ThreadPoolExecutor(max_workers=12) as pool:
+        responses: Final = tuple(pool.map(lambda body: _converse(rig, "gga-masksm_flat", body), bodies))
+    assert [response.status_code for response in responses] == [200] * len(bodies)
+    sent: Final = sorted(
+        (JSON_OBJECT.validate_json(request.body)["messages"] for request in rig.bedrock.drain()), key=json.dumps
+    )
+    expected: Final = sorted(
+        ([_user({"text": f"request {index}"}, {"text": f"card {_MASK}"})] for index in range(len(cards))),
+        key=json.dumps,
+    )
+    assert sent == expected

tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py

@@ -1650,6 +1650,9 @@ class TestConverseRewriteRegressions:
                     },
                 ],
                 "additionalModelRequestFields": {"prompt": "extra"},
+                "toolConfig": {
+                    "tools": [{"toolSpec": {"name": "search", "inputSchema": {"json": {"type": "object"}}}}]
+                },
             },
         }
         guard = _make_guardrail({})
@@ -1658,6 +1661,7 @@ class TestConverseRewriteRegressions:
         inputs = guard.apply_guardrail.call_args.kwargs["inputs"]
         assert inputs["texts"] == ["result"]
         assert inputs["structured_messages"] == [{"role": "tool", "tool_call_id": "t1", "content": "result"}]
+        assert "tools" not in inputs
 
     @pytest.mark.asyncio
     async def test_ambiguous_string_over_multiple_blocks_fails_closed_without_mutation(self):
@@ -1781,6 +1785,303 @@ class TestConverseRewriteRegressions:
         assert result == data
         guard.apply_guardrail.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_tool_result_image_part_is_skipped_when_mapping_the_text_part(self):
+        image = {"image": {"format": "png", "source": {"bytes": "iVBORw0KGgo="}}}
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [{"toolResult": {"toolUseId": "t1", "content": [{"text": "card 4111"}, image]}}],
+                    }
+                ]
+            },
+        }
+        guard = _make_guardrail(
+            {"structured_messages": [{"role": "tool", "tool_call_id": "t1", "content": "card [CARD]"}]}
+        )
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert result["data"]["messages"][0]["content"] == [
+            {"toolResult": {"toolUseId": "t1", "content": [{"text": "card [CARD]"}, image]}}
+        ]
+
+    @pytest.mark.asyncio
+    async def test_joined_tool_result_row_is_accepted_when_the_flat_texts_reproduce_it(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "toolResult": {
+                                    "toolUseId": "t1",
+                                    "content": [{"text": "found 4111"}, {"json": {"card": "4111"}}],
+                                }
+                            }
+                        ],
+                    }
+                ]
+            },
+        }
+        guard = _make_guardrail(
+            {
+                "texts": ["found [CARD]", "[CARD]"],
+                "structured_messages": [
+                    {"role": "tool", "tool_call_id": "t1", "content": 'found [CARD]\n{"card": "[CARD]"}'}
+                ],
+            }
+        )
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert result["data"]["messages"][0]["content"] == [
+            {"toolResult": {"toolUseId": "t1", "content": [{"text": "found [CARD]"}, {"json": {"card": "[CARD]"}}]}}
+        ]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("role", ["system", "user"])
+    async def test_flattened_multi_block_row_is_accepted_when_the_flat_texts_reproduce_it(self, role):
+        blocks = [{"text": "Order details."}, {"text": "card 4111"}]
+        body = (
+            {"system": blocks, "messages": [{"role": "user", "content": [{"text": "hi"}]}]}
+            if role == "system"
+            else {"messages": [{"role": "user", "content": blocks}]}
+        )
+        data = {"endpoint": "model/test/converse", "data": body}
+        rows = [{"role": role, "content": "Order details.\ncard [CARD]"}]
+        texts = ["Order details.", "card [CARD]"]
+        if role == "system":
+            rows.append({"role": "user", "content": [{"type": "text", "text": "hi"}]})
+            texts.append("hi")
+        guard = _make_guardrail({"texts": texts, "structured_messages": rows})
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        written = result["data"]["system"] if role == "system" else result["data"]["messages"][0]["content"]
+        assert written == [{"text": "Order details."}, {"text": "card [CARD]"}]
+
+    @pytest.mark.asyncio
+    async def test_flat_text_mask_wins_over_a_changed_but_unmasked_structured_leaf(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {"messages": [{"role": "user", "content": [{"text": "card 4111 on file."}, {"text": "Alice"}]}]},
+        }
+        guard = _make_guardrail(
+            {
+                "texts": ["card [CARD] on file.", "Alice"],
+                "structured_messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": "card 4111 on file. "},
+                            {"type": "text", "text": "[NAME]"},
+                        ],
+                    }
+                ],
+            }
+        )
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert result["data"]["messages"][0]["content"] == [{"text": "card [CARD] on file."}, {"text": "[NAME]"}]
+
+    @pytest.mark.asyncio
+    async def test_texts_that_do_not_line_up_with_the_request_are_ignored_next_to_structured_rows(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {
+                "system": [{"text": "You are a support bot."}],
+                "messages": [
+                    {"role": "user", "content": [{"text": "My name is Ann."}]},
+                    {"role": "assistant", "content": [{"text": "Hi Ann."}]},
+                    {"role": "user", "content": [{"text": "card 4111 ok"}]},
+                ],
+            },
+        }
+        guard = _make_guardrail(
+            {
+                "texts": ["My name is Ann.", "card [CARD] ok"],
+                "structured_messages": [
+                    {"role": "system", "content": [{"type": "text", "text": "You are a support bot."}]},
+                    {"role": "user", "content": [{"type": "text", "text": "My name is Ann."}]},
+                    {"role": "assistant", "content": [{"type": "text", "text": "Hi Ann."}]},
+                    {"role": "user", "content": [{"type": "text", "text": "card [CARD] ok"}]},
+                ],
+            }
+        )
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert result["data"]["system"] == [{"text": "You are a support bot."}]
+        assert result["data"]["messages"] == [
+            {"role": "user", "content": [{"text": "My name is Ann."}]},
+            {"role": "assistant", "content": [{"text": "Hi Ann."}]},
+            {"role": "user", "content": [{"text": "card [CARD] ok"}]},
+        ]
+
+    @pytest.mark.asyncio
+    async def test_structured_messages_that_are_not_rows_fall_back_to_texts(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {"messages": [{"role": "user", "content": [{"text": "4111"}]}]},
+        }
+        guard = _make_guardrail({"texts": ["[CARD]"], "structured_messages": {"evaluation": {"action": "Redact"}}})
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert result["data"]["messages"][0]["content"] == [{"text": "[CARD]"}]
+
+    @pytest.mark.asyncio
+    async def test_skip_tool_keeps_tool_call_arguments_out_of_structured_rows(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {
+                "messages": [
+                    {
+                        "role": "assistant",
+                        "content": [
+                            {"text": "Checking."},
+                            {"toolUse": {"toolUseId": "t1", "name": "pay", "input": {"card": "4111"}}},
+                        ],
+                    },
+                    {"role": "user", "content": [{"text": "thanks"}]},
+                ]
+            },
+        }
+        guard = _make_guardrail({})
+        guard.skip_tool_message_in_guardrail = True
+        await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert guard.apply_guardrail.call_args.kwargs["inputs"]["structured_messages"] == [
+            {"role": "assistant", "content": [{"type": "text", "text": "Checking."}]},
+            {"role": "user", "content": [{"type": "text", "text": "thanks"}]},
+        ]
+
+    @pytest.mark.asyncio
+    async def test_unchanged_structured_rows_keep_the_flat_text_write_back(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {"messages": [{"role": "user", "content": [{"text": "4111"}, {"text": "Alice"}]}]},
+        }
+        guard = _make_guardrail(
+            {
+                "texts": ["[CARD]"],
+                "structured_messages": [
+                    {"role": "user", "content": [{"type": "text", "text": "4111"}, {"type": "text", "text": "Alice"}]}
+                ],
+            }
+        )
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert result["data"]["messages"][0]["content"] == [{"text": "[CARD]"}, {"text": "Alice"}]
+
+    @pytest.mark.asyncio
+    async def test_rows_returned_when_none_were_sent_are_ignored(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {"messages": [{"role": "user", "content": [{"text": "4111"}, "not a block"]}]},
+        }
+        guard = _make_guardrail({"texts": ["[CARD]"], "structured_messages": [{"role": "user", "content": "[CARD]"}]})
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert "structured_messages" not in guard.apply_guardrail.call_args.kwargs["inputs"]
+        assert result["data"]["messages"][0]["content"] == [{"text": "[CARD]"}, "not a block"]
+
+    @pytest.mark.asyncio
+    async def test_missing_structured_row_fails_closed_without_mutation(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {
+                "messages": [
+                    {"role": "user", "content": [{"text": "hello"}]},
+                    {"role": "assistant", "content": [{"text": "hi"}]},
+                    {"role": "user", "content": [{"text": "4111"}]},
+                ]
+            },
+        }
+        original = copy.deepcopy(data)
+        guard = _make_guardrail(
+            {
+                "texts": ["hello", "hi", "[CARD]"],
+                "structured_messages": [
+                    {"role": "user", "content": [{"type": "text", "text": "hello"}]},
+                    {"role": "assistant", "content": [{"type": "text", "text": "hi"}]},
+                ],
+            }
+        )
+        with pytest.raises(UnappliableRequestRewrite):
+            await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert data == original
+
+    @pytest.mark.asyncio
+    async def test_flattened_row_after_a_turn_the_texts_emptied_fails_closed(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {
+                "messages": [
+                    {"role": "user", "content": [{"text": "hello"}]},
+                    {"role": "assistant", "content": [{"text": "hi"}]},
+                    {"role": "user", "content": [{"text": "card"}, {"text": "4111"}]},
+                ]
+            },
+        }
+        guard = _make_guardrail(
+            {
+                "texts": ["hello", "", "card", "[CARD]"],
+                "structured_messages": [
+                    {"role": "user", "content": [{"type": "text", "text": "hello"}]},
+                    {"role": "assistant", "content": [{"type": "text", "text": ""}]},
+                    {"role": "user", "content": "card\n[CARD]"},
+                ],
+            }
+        )
+        with pytest.raises(UnappliableRequestRewrite):
+            await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+
+    @pytest.mark.asyncio
+    async def test_joined_tool_result_row_the_flat_texts_do_not_reproduce_fails_closed(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"toolResult": {"toolUseId": "t1", "content": [{"text": "found 4111"}, {"text": "ok"}]}}
+                        ],
+                    }
+                ]
+            },
+        }
+        original = copy.deepcopy(data)
+        guard = _make_guardrail(
+            {"structured_messages": [{"role": "tool", "tool_call_id": "t1", "content": "found [CARD]\nok"}]}
+        )
+        with pytest.raises(UnappliableRequestRewrite):
+            await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert data == original
+
+    @pytest.mark.asyncio
+    async def test_flat_mask_written_into_the_sent_texts_list_still_wins(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {"messages": [{"role": "user", "content": [{"text": "ssn 123 email a@b.c"}]}]},
+        }
+
+        def mask_in_place(inputs: dict, **_: object) -> dict:
+            inputs["texts"][0] = "ssn [SSN] email a@b.c"
+            rows = copy.deepcopy(inputs["structured_messages"])
+            rows[0]["content"][0]["text"] = "ssn 123 email [EMAIL]"
+            return {**inputs, "structured_messages": rows}
+
+        guard = _make_guardrail({})
+        guard.apply_guardrail = AsyncMock(side_effect=lambda inputs, **kwargs: mask_in_place(inputs))
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert result["data"]["messages"][0]["content"] == [{"text": "ssn [SSN] email a@b.c"}]
+
+    @pytest.mark.asyncio
+    async def test_snake_case_block_keys_are_not_read_as_converse_fields(self):
+        data = {
+            "endpoint": "model/test/converse",
+            "data": {"messages": [{"role": "user", "content": [{"guard_content": {"text": {"text": "4111"}}}]}]},
+        }
+        original = copy.deepcopy(data)
+        guard = _make_guardrail({"structured_messages": [{"role": "user", "content": "[CARD]"}]})
+        result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
+        assert result == original
+        guard.apply_guardrail.assert_not_called()
+
 
 async def _run_post_call_guardrail_on_stream(
     stream_bytes: bytes, guardrail: CustomGuardrail, monkeypatch: pytest.MonkeyPatch