fix(guardrails): scan tool-result-only messages under experimental_use_latest_role_message_only

_select_messages_for_apply_guardrail used _count_message_texts, which only
recognizes a top-level "text" key, to decide whether the latest user message
had scannable content. A message whose only content is a tool_result block
passed _count_message_texts(...) == 0 and the whole INPUT scan was skipped
via skip_scan=True, bypassing Bedrock's ApplyGuardrail API entirely for that
message under experimental_use_latest_role_message_only.

Use get_content_items_for_message (the tool-result-aware extractor already
used elsewhere in this fix) instead, so a tool-result-only message is
correctly recognized as having scannable content.

Addresses veria-ai review comment on #33092.
This commit is contained in:
bricelalu 2026-07-15 22:07:34 +02:00
parent 5817d7cdc9
commit 598eac84bb
No known key found for this signature in database
GPG key ID: 72D3579A5CD0D821
2 changed files with 59 additions and 1 deletions

View file

@ -532,7 +532,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
return ApplyGuardrailMessageSelection(None, None, True, skip_scan=True)
selected_message = structured_messages[latest_user_index]
if self._count_message_texts(selected_message) == 0:
# _count_message_texts only recognizes a top-level "text" key and does not know
# about tool_result-nested text, so a message whose only content is a tool_result
# block would otherwise look empty here and skip_scan, bypassing the guardrail
# entirely. get_content_items_for_message is the tool-result-aware extractor.
if not self.get_content_items_for_message(message=selected_message):
verbose_proxy_logger.debug(
"Bedrock Guardrail: latest user message has no text content, skipping INPUT scan"
)

View file

@ -3684,3 +3684,57 @@ def test_locate_message_texts_slice_returns_none_on_tool_result_count_mismatch()
)
assert result is None
def test_select_messages_for_apply_guardrail_does_not_skip_tool_result_only_message():
"""veria-ai review (PR #33092): _select_messages_for_apply_guardrail used
_count_message_texts -- which only recognizes a top-level "text" key -- to
decide whether the latest user message has scannable content. A message
whose *only* content is a tool_result block passed
_count_message_texts(...) == 0, so the whole INPUT scan was skipped via
skip_scan=True: a complete guardrail bypass under
experimental_use_latest_role_message_only, since Bedrock's ApplyGuardrail
API is never even called (worse than a masking-misalignment bug -- BLOCKED
detection doesn't happen either). Must use the tool-result-aware
extractor (get_content_items_for_message) instead."""
guardrail = BedrockGuardrail(
guardrailIdentifier="test-guardrail",
guardrailVersion="DRAFT",
experimental_use_latest_role_message_only=True,
)
tool_result_only_message = _string_tool_result_message()
selection = guardrail._select_messages_for_apply_guardrail(
texts=[],
inputs={"texts": [], "structured_messages": [tool_result_only_message]},
request_data={},
input_type="request",
)
assert selection.skip_scan is False
assert selection.filtered_messages == [tool_result_only_message]
@pytest.mark.asyncio
async def test_apply_guardrail_scans_tool_result_only_message_under_latest_role_only():
"""End-to-end companion to the unit test above: apply_guardrail must
actually call Bedrock's ApplyGuardrail API with the tool_result-only
message's content, not silently skip_scan and let it through unscanned."""
guardrail = BedrockGuardrail(
guardrailIdentifier="test-guardrail",
guardrailVersion="DRAFT",
experimental_use_latest_role_message_only=True,
)
tool_result_only_message = _string_tool_result_message()
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
await guardrail.apply_guardrail(
inputs={"texts": [], "structured_messages": [tool_result_only_message]},
request_data={"messages": [tool_result_only_message]},
input_type="request",
)
mock_api.assert_called_once()
assert mock_api.call_args.kwargs["messages"] == [tool_result_only_message]