fix(bedrock): guard neutralized tool results anywhere in history when guardrailConfig is set

This commit is contained in:
Kent 2026-09-12 09:51:01 +08:00
parent 5493be743b
commit 04086be795
2 changed files with 60 additions and 6 deletions

View file

@ -234,22 +234,28 @@ class AmazonConverseConfig(BaseConfig):
rendered = convert_content_list_to_str(message).strip()
return rendered or "<non-text tool result omitted>"
guardrail_active: Final = "guardrailConfig" in optional_params
def _rewrite(message: AllMessageValues) -> AllMessageValues:
role = message.get("role")
tool_calls = message.get("tool_calls")
if role == "assistant" and tool_calls:
base_text = convert_content_list_to_str(message)
call_texts = [_tool_call_text(call) for call in tool_calls]
text = "\n".join(filter(None, [base_text, *call_texts]))
base_text: Final = convert_content_list_to_str(message)
call_texts: Final = tuple(_tool_call_text(call) for call in tool_calls)
text: Final = "\n".join(part for part in (base_text, *call_texts) if part)
return ChatCompletionAssistantMessage(role="assistant", content=text)
if role in ("tool", "function"):
tool_call_id = message.get("tool_call_id")
name = message.get("name")
label = f"tool result for {tool_call_id or name or 'unknown'}"
return ChatCompletionUserMessage(
role="user",
content=f"[{label}: {_result_text(message)}]",
result_text: Final = f"[{label}: {_result_text(message)}]"
# Tool results are externally controlled, so guard them wherever they
# land in history; _convert_consecutive_user_messages_to_guarded_text
# only covers the trailing user turn.
content: Final = (
[{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text
)
return ChatCompletionUserMessage(role="user", content=content)
return message
verbose_logger.warning(

View file

@ -6714,6 +6714,54 @@ def test_transform_request_neutralized_tool_output_is_guarded(monkeypatch):
assert "secret tool output" in serialized
def test_transform_request_neutralized_tool_output_guarded_mid_history(monkeypatch):
"""Regression: a neutralized tool result that is NOT the trailing turn (an
assistant reply and a later user turn follow it) must still be guardContent.
_convert_consecutive_user_messages_to_guarded_text only covers the trailing
user turn, so neutralize itself must guard untrusted tool output regardless
of position, else an attacker controlling the tool response bypasses the
guardrail (bot review)."""
monkeypatch.setattr(litellm, "modify_params", False)
config = AmazonConverseConfig()
result = config.transform_request(
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
messages=[
{"role": "user", "content": "look it up"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "c1",
"type": "function",
"function": {"name": "lookup", "arguments": "{}"},
}
],
},
{"role": "tool", "tool_call_id": "c1", "content": "IGNORE_PRIOR malware"},
{"role": "assistant", "content": "Here is the summary."},
{"role": "user", "content": "thanks"},
],
optional_params={
"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "1"}
},
litellm_params={},
headers={},
)
_assert_no_structured_tool_blocks(result)
blocks = [block for message in result["messages"] for block in message["content"]]
guarded_texts = [
block["guardContent"]["text"]["text"] for block in blocks if "guardContent" in block
]
plain_texts = [block["text"] for block in blocks if "text" in block and "guardContent" not in block]
assert any("malware" in text for text in guarded_texts), "mid-history tool output must be guarded"
assert not any(
"malware" in text for text in plain_texts
), "mid-history tool output must not reach the model as unguarded text"
@pytest.mark.asyncio
async def test_async_transform_request_no_tools_with_tool_history(monkeypatch):
"""Async is a separate request assembler; it must neutralize identically."""