From beb7d2916d36ab4efc0387936fccaecded88003e Mon Sep 17 00:00:00 2001 From: Namraa Patel Date: Wed, 28 Oct 2026 18:37:31 +0530 Subject: [PATCH] Fix Bedrock parallel tool result handling. --- .../agent_framework_bedrock/_chat_client.py | 12 +- .../bedrock/tests/test_bedrock_client.py | 182 ++++++++++++++++++ 2 files changed, 189 insertions(+), 5 deletions(-) diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index b9f125f7bc..d03137e2eb 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -554,9 +554,12 @@ def _prepare_bedrock_messages( ) elif message.role == "tool": content_blocks = self._align_tool_results_with_pending(content_blocks, pending_tool_use_ids) - pending_tool_use_ids.clear() if not content_blocks: continue + # Merge consecutive tool messages into a single user message + if conversation and conversation[-1]["role"] == "user": + conversation[-1]["content"].extend(content_blocks) + continue else: pending_tool_use_ids.clear() @@ -576,7 +579,6 @@ def _align_tool_results_with_pending( ] aligned_blocks: list[dict[str, Any]] = [] - pending = deque(pending_tool_use_ids) for block in content_blocks: if not isinstance(block, MutableMapping): aligned_blocks.append(block) @@ -585,18 +587,18 @@ def _align_tool_results_with_pending( if not tool_result: aligned_blocks.append(block) continue - if not pending: + if not pending_tool_use_ids: logger.debug("Dropping extra tool result block due to missing pending tool uses: %s", block) continue tool_use_id = tool_result.get("toolUseId") if tool_use_id: try: - pending.remove(tool_use_id) + pending_tool_use_ids.remove(tool_use_id) except ValueError: logger.debug("Tool result references unknown toolUseId '%s'. Dropping block.", tool_use_id) continue else: - tool_result["toolUseId"] = pending.popleft() + tool_result["toolUseId"] = pending_tool_use_ids.popleft() aligned_blocks.append(block) return aligned_blocks diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index fbe23f4847..9f753e8835 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -1025,6 +1025,188 @@ def test_align_tool_results_handles_pending_edge_cases() -> None: assert unmatched == [] +def test_prepare_bedrock_messages_preserves_separate_parallel_tool_results() -> None: + """Parallel tool results arriving as separate tool messages must not be lost.""" + client = _make_client() + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c1", name="tool_a", arguments='{"arg": "value1"}'), + Content.from_function_call(call_id="c2", name="tool_b", arguments='{"arg": "value2"}'), + ], + ), + Message(role="tool", contents=[Content.from_function_result(call_id="c1", result="result1")]), + Message(role="tool", contents=[Content.from_function_result(call_id="c2", result="result2")]), + ] + + _, conversation = client._prepare_bedrock_messages(messages) + + # Should have assistant message with toolUse blocks and one user message with both toolResult blocks + assert len(conversation) == 2 + assert conversation[0]["role"] == "assistant" + assert len(conversation[0]["content"]) == 2 + assert conversation[0]["content"][0]["toolUse"]["toolUseId"] == "c1" + assert conversation[0]["content"][1]["toolUse"]["toolUseId"] == "c2" + + assert conversation[1]["role"] == "user" + assert len(conversation[1]["content"]) == 2 + assert conversation[1]["content"][0]["toolResult"]["toolUseId"] == "c1" + assert conversation[1]["content"][1]["toolResult"]["toolUseId"] == "c2" + + +def test_prepare_bedrock_messages_rejects_duplicate_and_unknown_tool_results() -> None: + """Duplicate tool results must be deduplicated; unknown IDs must be rejected.""" + client = _make_client() + messages = [ + Message( + role="assistant", + contents=[Content.from_function_call(call_id="c1", name="tool_a", arguments="{}")], + ), + Message(role="tool", contents=[Content.from_function_result(call_id="c1", result="result1")]), + Message(role="tool", contents=[Content.from_function_result(call_id="c1", result="result1")]), + Message(role="tool", contents=[Content.from_function_result(call_id="unknown-id", result="result2")]), + ] + + _, conversation = client._prepare_bedrock_messages(messages) + + # Should have assistant message and one user message with a single toolResult (c1) + assert len(conversation) == 2 + assert conversation[0]["role"] == "assistant" + assert conversation[1]["role"] == "user" + assert len(conversation[1]["content"]) == 1 + assert conversation[1]["content"][0]["toolResult"]["toolUseId"] == "c1" + + +def test_prepare_bedrock_messages_handles_single_tool_call() -> None: + """Single tool calls must continue working correctly.""" + client = _make_client() + messages = [ + Message( + role="assistant", + contents=[Content.from_function_call(call_id="c1", name="tool_a", arguments="{}")], + ), + Message(role="tool", contents=[Content.from_function_result(call_id="c1", result="result1")]), + ] + + _, conversation = client._prepare_bedrock_messages(messages) + + assert len(conversation) == 2 + assert conversation[0]["role"] == "assistant" + assert conversation[1]["role"] == "user" + assert len(conversation[1]["content"]) == 1 + assert conversation[1]["content"][0]["toolResult"]["toolUseId"] == "c1" + + +def test_prepare_bedrock_messages_handles_parallel_results_in_one_message() -> None: + """Parallel tool results in a single tool message must continue working.""" + client = _make_client() + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c1", name="tool_a", arguments="{}"), + Content.from_function_call(call_id="c2", name="tool_b", arguments="{}"), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result(call_id="c1", result="result1"), + Content.from_function_result(call_id="c2", result="result2"), + ], + ), + ] + + _, conversation = client._prepare_bedrock_messages(messages) + + assert len(conversation) == 2 + assert conversation[0]["role"] == "assistant" + assert conversation[1]["role"] == "user" + assert len(conversation[1]["content"]) == 2 + assert conversation[1]["content"][0]["toolResult"]["toolUseId"] == "c1" + assert conversation[1]["content"][1]["toolResult"]["toolUseId"] == "c2" + + +def test_prepare_bedrock_messages_handles_three_parallel_separate_results() -> None: + """Three or more parallel tool calls with separate results must all be preserved.""" + client = _make_client() + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c1", name="tool_a", arguments="{}"), + Content.from_function_call(call_id="c2", name="tool_b", arguments="{}"), + Content.from_function_call(call_id="c3", name="tool_c", arguments="{}"), + ], + ), + Message(role="tool", contents=[Content.from_function_result(call_id="c1", result="result1")]), + Message(role="tool", contents=[Content.from_function_result(call_id="c2", result="result2")]), + Message(role="tool", contents=[Content.from_function_result(call_id="c3", result="result3")]), + ] + + _, conversation = client._prepare_bedrock_messages(messages) + + assert len(conversation) == 2 + assert conversation[0]["role"] == "assistant" + assert conversation[1]["role"] == "user" + assert len(conversation[1]["content"]) == 3 + assert conversation[1]["content"][0]["toolResult"]["toolUseId"] == "c1" + assert conversation[1]["content"][1]["toolResult"]["toolUseId"] == "c2" + assert conversation[1]["content"][2]["toolResult"]["toolUseId"] == "c3" + + +def test_prepare_bedrock_messages_does_not_merge_unrelated_messages() -> None: + """Tool results followed by another assistant message must not merge incorrectly.""" + client = _make_client() + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c1", name="tool_a", arguments="{}"), + Content.from_function_call(call_id="c2", name="tool_b", arguments="{}"), + ], + ), + Message(role="tool", contents=[Content.from_function_result(call_id="c1", result="result1")]), + Message(role="tool", contents=[Content.from_function_result(call_id="c2", result="result2")]), + Message(role="assistant", contents=[Content.from_text(text="next response")]), + ] + + _, conversation = client._prepare_bedrock_messages(messages) + + # Should have: assistant (toolUse), user (toolResult), assistant (text) + assert len(conversation) == 3 + assert conversation[0]["role"] == "assistant" + assert conversation[1]["role"] == "user" + assert conversation[2]["role"] == "assistant" + assert len(conversation[1]["content"]) == 2 + + +def test_prepare_bedrock_messages_does_not_invent_missing_results() -> None: + """Missing tool results must not be fabricated.""" + client = _make_client() + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c1", name="tool_a", arguments="{}"), + Content.from_function_call(call_id="c2", name="tool_b", arguments="{}"), + ], + ), + Message(role="tool", contents=[Content.from_function_result(call_id="c1", result="result1")]), + # c2 result is missing + ] + + _, conversation = client._prepare_bedrock_messages(messages) + + assert len(conversation) == 2 + assert conversation[0]["role"] == "assistant" + assert conversation[1]["role"] == "user" + # Only c1 result should be present + assert len(conversation[1]["content"]) == 1 + assert conversation[1]["content"][0]["toolResult"]["toolUseId"] == "c1" + + def test_convert_content_to_bedrock_block_handles_errors_and_missing_items() -> None: """Function result conversion should serialize items, rich content warnings, and fallback results.""" client = _make_client()