fix(guardrails): split a newline-joined Converse message back into its parts

A vendor that returns a multi-part message as one newline-joined string now has its masks written back to each part, guardContent included, when no part sent already holds a newline.
This commit is contained in:
yucheng-berri 2026-10-17 15:31:53 -07:00
parent 20c7fb359d
commit a499c28029
2 changed files with 54 additions and 8 deletions

View file

@ -166,12 +166,25 @@ def _carries_no_text(content: object) -> bool:
return False
def _content_texts(row: ConverseRow, returned: Mapping[str, object]) -> tuple[str, ...] | None:
content: Final = returned.get("content")
sent: Final = _text_parts(row.message)
if (
isinstance(content, str)
and len(row.content) > 1
and sent is not None
and all("\n" not in text for text in sent)
):
return tuple(content.split("\n"))
return _text_parts(returned)
def _content_edits(row: ConverseRow, original: object, returned: Mapping[str, object]) -> tuple[Edit, ...] | None:
if returned.get("content") == original:
return ()
if not row.content and _carries_no_text(returned.get("content")):
return ()
texts: Final = _text_parts(returned)
texts: Final = _content_texts(row, returned)
if texts is None or len(texts) != len(row.content):
return None
return _merged(tuple(_slot_edits(slot, text) for slot, text in zip(row.content, texts)))

View file

@ -1741,14 +1741,14 @@ class TestConverseRowEditWriteBack:
id="role-changed",
),
pytest.param(
[
{
"role": "user",
"content": [{"text": "card 4111"}, {"cachePoint": {"type": "default"}}, {"text": "thanks"}],
}
],
[{"role": "user", "content": [{"text": "card\n4111"}, {"text": "thanks"}]}],
[{"role": "user", "content": "card\n[CARD] thanks"}],
id="blocks-that-hold-newlines-joined-with-a-space",
),
pytest.param(
[{"role": "user", "content": [{"text": "card 4111"}, {"text": "thanks"}, {"text": "bye"}]}],
[{"role": "user", "content": "card [CARD]\nthanks"}],
id="blocks-merged-around-a-cache-point",
id="newline-joined-blocks-with-one-missing",
),
pytest.param([_pay_turn(_CARD_TOOL_CALL)], [_pay_row(content="added text")], id="text-added-to-a-call"),
pytest.param([_pay_turn(_CARD_TOOL_CALL)], [_pay_row(call_id="t2")], id="tool-call-id-changed"),
@ -1789,6 +1789,39 @@ class TestConverseRowEditWriteBack:
assert data["data"] == original
@pytest.mark.asyncio
@pytest.mark.parametrize(
("last_block", "masked_last_block"),
[
({"text": "thanks 4111"}, {"text": "thanks [CARD]"}),
({"guardContent": {"text": {"text": "thanks 4111"}}}, {"guardContent": {"text": {"text": "thanks [CARD]"}}}),
],
ids=["text", "guard-content"],
)
async def test_row_joined_with_newlines_is_split_back_into_its_blocks(self, last_block, masked_last_block):
cache_point = {"cachePoint": {"type": "default"}}
data = _converse_request(
{"messages": [{"role": "user", "content": [{"text": "card 4111"}, cache_point, last_block]}]}
)
result = await BedrockPassthroughGuardrailHandler().process_input_messages(
data, _replying({"structured_messages": [{"role": "user", "content": "card [CARD]\nthanks [CARD]"}]})
)
assert result["data"]["messages"] == [
{"role": "user", "content": [{"text": "card [CARD]"}, cache_point, masked_last_block]}
]
@pytest.mark.asyncio
async def test_single_block_row_keeps_a_newline_the_guardrail_wrote(self):
data = _converse_request({"messages": [{"role": "user", "content": [{"text": "card 4111"}]}]})
result = await BedrockPassthroughGuardrailHandler().process_input_messages(
data, _replying({"structured_messages": [{"role": "user", "content": "card\n[CARD]"}]})
)
assert result["data"]["messages"] == [{"role": "user", "content": [{"text": "card\n[CARD]"}]}]
@pytest.mark.asyncio
@pytest.mark.parametrize("content", [None, "", []], ids=["none", "empty-string", "empty-list"])
async def test_tool_call_argument_mask_in_a_row_is_written_back(self, content):