Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 35 additions & 15 deletions strands-py/src/strands/session/repository_session_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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})
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Issue: Missing docstring. The neighboring tests (test_fix_broken_tool_use_extends_partial_tool_results, test_fix_broken_tool_use_handles_multiple_orphaned_tools) all have docstrings describing what they test. This one should too for consistency.

Suggestion: Add a brief docstring:

def test_fix_broken_tool_use_drops_extra_tool_results(session_manager):
    """Test that stale and duplicate toolResult blocks are removed."""

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here, docstring would be great

"""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"
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Issue: Consider adding a test for the edge case where removed_tool_results=True but missing_tool_use_ids is empty (i.e., all toolUse IDs have a matching result, but there are also stale/duplicate results to drop). This scenario exercises the if missing_tool_use_ids or removed_tool_results path differently — generate_missing_tool_result_content returns [] and the content is just filtered_content.

Suggestion:

def test_fix_broken_tool_use_drops_stale_results_with_no_missing_ids(session_manager):
    """Test that stale toolResults are removed even when all toolUse IDs are satisfied."""
    messages = [
        {
            "role": "assistant",
            "content": [
                {"toolUse": {"toolUseId": "id-1", "name": "t", "input": {}}},
            ],
        },
        {
            "role": "user",
            "content": [
                {"toolResult": {"toolUseId": "id-1", "status": "success", "content": [{"text": "ok"}]}},
                {"toolResult": {"toolUseId": "stale-999", "status": "error", "content": [{"text": "old"}]}},
            ],
        },
    ]
    fixed = session_manager._fix_broken_tool_use(messages)
    results = [c["toolResult"] for c in fixed[1]["content"]]
    assert len(results) == 1
    assert results[0]["toolUseId"] == "id-1"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with the bot again! I think we need this test for sure



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."""

Expand Down