From 6a0e7fbeb45850594a94f0f734e162ede1f57c56 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sat, 23 May 2026 20:46:29 +0800 Subject: [PATCH] fix: normalize restored tool results --- .../session/repository_session_manager.py | 50 +++++++++++------ .../test_repository_session_manager.py | 53 +++++++++++++++++++ 2 files changed, 88 insertions(+), 15 deletions(-) diff --git a/strands-py/src/strands/session/repository_session_manager.py b/strands-py/src/strands/session/repository_session_manager.py index c1032a85ea..30f742736e 100644 --- a/strands-py/src/strands/session/repository_session_manager.py +++ b/strands-py/src/strands/session/repository_session_manager.py @@ -281,26 +281,46 @@ def _fix_broken_tool_use(self, messages: list[Message]) -> list[Message]: content["toolUse"]["toolUseId"] for content in message["content"] if "toolUse" in content ] - # 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 + # Keep only the toolResults that belong to this assistant turn. + 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 = [ + tool_use_id for tool_use_id in tool_use_ids if tool_use_id not in tool_result_ids ] + if missing_tool_use_ids or removed_tool_results: + if missing_tool_use_ids: + logger.warning( + "Session message history has an orphaned toolUse with no toolResult. " + "Adding toolResult content blocks to create valid conversation." + ) + if removed_tool_results: + logger.warning( + "Session message history has duplicate or unrelated toolResult blocks. " + "Removing extra toolResult content blocks to create valid conversation." + ) - missing_tool_use_ids = list(set(tool_use_ids) - set(tool_result_ids)) - # If there are missing tool use ids, that means the messages history is broken - if missing_tool_use_ids: - logger.warning( - "Session message history has an orphaned toolUse with no toolResult. " - "Adding toolResult content blocks to create valid conversation." - ) # 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/strands-py/tests/strands/session/test_repository_session_manager.py b/strands-py/tests/strands/session/test_repository_session_manager.py index 1d50481132..f63611e10e 100644 --- a/strands-py/tests/strands/session/test_repository_session_manager.py +++ b/strands-py/tests/strands/session/test_repository_session_manager.py @@ -370,6 +370,59 @@ 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): + """Keep valid toolResults, drop stale duplicates, and fill missing ones.""" + 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_drops_stale_results_with_no_missing_ids(session_manager): + """Drop stale toolResults even when every toolUse already has a valid result.""" + messages = [ + { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "complete-123", "name": "test_tool", "input": {"input": "test1"}}}, + ], + }, + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "complete-123", "status": "success", "content": [{"text": "ok"}]}}, + {"toolResult": {"toolUseId": "stale-789", "status": "error", "content": [{"text": "old"}]}}, + ], + }, + ] + + 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"] + assert tool_results[0]["content"] == [{"text": "ok"}] + + def test_fix_broken_tool_use_handles_multiple_orphaned_tools(existing_session_manager): """Test fixing multiple orphaned toolUse messages."""