diff --git a/src/strands/session/repository_session_manager.py b/src/strands/session/repository_session_manager.py index c1032a85e..a7cf6bb92 100644 --- a/src/strands/session/repository_session_manager.py +++ b/src/strands/session/repository_session_manager.py @@ -282,15 +282,29 @@ def _fix_broken_tool_use(self, messages: list[Message]) -> list[Message]: ] # Check if there are more messages after the current toolUse message - tool_result_ids = [ - content["toolResult"]["toolUseId"] - for content in messages[index + 1]["content"] - if "toolResult" in content - ] + next_message = messages[index + 1] + tool_use_id_set = set(tool_use_ids) + tool_result_ids = [] + filtered_content = [] + removed_tool_results = False + + for content in next_message["content"]: + if "toolResult" not in content: + filtered_content.append(content) + continue + + tool_use_id = content["toolResult"]["toolUseId"] + if tool_use_id in tool_use_id_set and tool_use_id not in tool_result_ids: + tool_result_ids.append(tool_use_id) + filtered_content.append(content) + else: + removed_tool_results = True - missing_tool_use_ids = list(set(tool_use_ids) - set(tool_result_ids)) + missing_tool_use_ids = [ + tool_use_id for tool_use_id in tool_use_ids if tool_use_id not in tool_result_ids + ] # If there are missing tool use ids, that means the messages history is broken - if missing_tool_use_ids: + if missing_tool_use_ids or removed_tool_results: logger.warning( "Session message history has an orphaned toolUse with no toolResult. " "Adding toolResult content blocks to create valid conversation." @@ -298,9 +312,9 @@ def _fix_broken_tool_use(self, messages: list[Message]) -> list[Message]: # Create the missing toolResult content blocks missing_content_blocks = generate_missing_tool_result_content(missing_tool_use_ids) - if tool_result_ids: - # If there were any toolResult ids, that means only some of the content blocks are missing - messages[index + 1]["content"].extend(missing_content_blocks) + if tool_result_ids or removed_tool_results: + # Keep only toolResults that match the previous assistant toolUse blocks. + next_message["content"] = filtered_content + missing_content_blocks else: # The message following the toolUse was not a toolResult, so lets insert it messages.insert(index + 1, {"role": "user", "content": missing_content_blocks}) diff --git a/tests/strands/session/test_repository_session_manager.py b/tests/strands/session/test_repository_session_manager.py index 1d5048113..ea190901f 100644 --- a/tests/strands/session/test_repository_session_manager.py +++ b/tests/strands/session/test_repository_session_manager.py @@ -370,6 +370,33 @@ def test_fix_broken_tool_use_extends_partial_tool_results(existing_session_manag assert missing_result["toolResult"]["content"][0]["text"] == "Tool was interrupted." +def test_fix_broken_tool_use_drops_extra_tool_results(session_manager): + messages = [ + { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "complete-123", "name": "test_tool", "input": {"input": "test1"}}}, + {"toolUse": {"toolUseId": "missing-456", "name": "test_tool", "input": {"input": "test2"}}}, + ], + }, + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "complete-123", "status": "success", "content": [{"text": "ok"}]}}, + {"toolResult": {"toolUseId": "stale-789", "status": "error", "content": [{"text": "old"}]}}, + {"toolResult": {"toolUseId": "complete-123", "status": "success", "content": [{"text": "dup"}]}}, + ], + }, + ] + + fixed_messages = session_manager._fix_broken_tool_use(messages) + + tool_results = [content["toolResult"] for content in fixed_messages[1]["content"]] + assert [tool_result["toolUseId"] for tool_result in tool_results] == ["complete-123", "missing-456"] + assert tool_results[0]["content"] == [{"text": "ok"}] + assert tool_results[1]["status"] == "error" + + def test_fix_broken_tool_use_handles_multiple_orphaned_tools(existing_session_manager): """Test fixing multiple orphaned toolUse messages."""