Skip to content

fix: prevent assistant-first tool-use history after context reduction #1750

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
KKamJi98 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: prevent assistant-first tool-use history after context reduction #1750

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
32 changes: 27 additions & 5 deletions src/strands/agent/conversation_manager/sliding_window_conversation_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,10 +162,15 @@ def reduce_context(self, agent: "Agent", e: Exception | None = None, **kwargs: A
logger.debug(
"message_index=<%s> | found message with tool results at index", last_message_idx_with_tool_results
)
results_truncated = self._truncate_tool_results(messages, last_message_idx_with_tool_results)
if results_truncated:
logger.debug("message_index=<%s> | tool results truncated", last_message_idx_with_tool_results)
return
for message_index in range(last_message_idx_with_tool_results, -1, -1):
if not any("toolResult" in content for content in messages[message_index]["content"]):
continue

results_truncated = self._truncate_tool_results(messages, message_index)
if results_truncated:
logger.debug("message_index=<%s> | tool results truncated", message_index)
self._ensure_starts_with_user_turn(messages, e)
return

# Try to trim index id when tool result cannot be truncated anymore
# If the number of messages is less than the window_size, then we default to 2, otherwise, trim to window size
Expand Down Expand Up @@ -195,6 +200,23 @@ def reduce_context(self, agent: "Agent", e: Exception | None = None, **kwargs: A

# Overwrite message history
messages[:] = messages[trim_index:]
self._ensure_starts_with_user_turn(messages, e)

def _ensure_starts_with_user_turn(self, messages: Messages, e: Exception | None = None) -> None:
"""Ensure reduced history does not start with assistant tool calls."""
removed_non_user_messages = 0
while (
messages
and messages[0]["role"] != "user"
and any("toolUse" in content for content in messages[0]["content"])
):
messages.pop(0)
removed_non_user_messages += 1

self.removed_message_count += removed_non_user_messages

if not messages:
raise ContextWindowOverflowException("Unable to trim conversation context!") from e

def _truncate_tool_results(self, messages: Messages, msg_idx: int) -> bool:
"""Truncate tool results in a message to reduce context size.
Expand Down Expand Up @@ -227,7 +249,7 @@ def _truncate_tool_results(self, messages: Messages, msg_idx: int) -> bool:
and tool_result_content_text == tool_result_too_large_message
):
logger.info("ToolResult has already been updated, skipping overwrite")
return False
continue
# Update status to error with informative message
message["content"][i]["toolResult"]["status"] = "error"
message["content"][i]["toolResult"]["content"] = [{"text": tool_result_too_large_message}]
Expand Down
56 changes: 49 additions & 7 deletions tests/strands/agent/test_conversation_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ def conversation_manager(request):
{"role": "assistant", "content": [{"text": "Second response"}]},
],
),
# 7 - Message count above max window size - Preserve tool use/tool result pairs
# 7 - Message count above max window size - Keep latest user turn as first message
(
{"window_size": 2},
[
Expand All @@ -115,11 +115,10 @@ def conversation_manager(request):
{"role": "user", "content": [{"toolResult": {"toolUseId": "456", "content": [], "status": "success"}}]},
],
[
{"role": "assistant", "content": [{"toolUse": {"toolUseId": "123", "name": "tool1", "input": {}}}]},
{"role": "user", "content": [{"toolResult": {"toolUseId": "456", "content": [], "status": "success"}}]},
],
),
# 8 - Test sliding window behavior - preserve tool use/result pairs across cut boundary
# 8 - Test sliding window behavior - avoid assistant-first reduced history
(
{"window_size": 3},
[
Expand All @@ -129,12 +128,11 @@ def conversation_manager(request):
{"role": "assistant", "content": [{"text": "Response after tool use"}]},
],
[
{"role": "assistant", "content": [{"toolUse": {"toolUseId": "123", "name": "tool1", "input": {}}}]},
{"role": "user", "content": [{"toolResult": {"toolUseId": "123", "content": [], "status": "success"}}]},
{"role": "assistant", "content": [{"text": "Response after tool use"}]},
],
),
# 9 - Test sliding window with multiple tool pairs that need preservation
# 9 - Test sliding window with multiple tool pairs while keeping user as first turn
(
{"window_size": 4},
[
Expand All @@ -146,7 +144,6 @@ def conversation_manager(request):
{"role": "assistant", "content": [{"text": "Final response"}]},
],
[
{"role": "assistant", "content": [{"toolUse": {"toolUseId": "456", "name": "tool2", "input": {}}}]},
{"role": "user", "content": [{"toolResult": {"toolUseId": "456", "content": [], "status": "success"}}]},
{"role": "assistant", "content": [{"text": "Final response"}]},
],
Expand Down Expand Up @@ -192,7 +189,6 @@ def test_sliding_window_conversation_manager_with_tool_results_truncated():
manager.reduce_context(test_agent)

expected_messages = [
{"role": "assistant", "content": [{"toolUse": {"toolUseId": "456", "name": "tool1", "input": {}}}]},
{
"role": "user",
"content": [
Expand All @@ -210,6 +206,52 @@ def test_sliding_window_conversation_manager_with_tool_results_truncated():
assert messages == expected_messages


def test_sliding_window_conversation_manager_scans_older_tool_results_when_latest_is_already_truncated():
manager = SlidingWindowConversationManager(1)
messages = [
{"role": "assistant", "content": [{"toolUse": {"toolUseId": "old", "name": "tool1", "input": {}}}]},
{
"role": "user",
"content": [
{"toolResult": {"toolUseId": "old", "content": [{"text": "large input"}], "status": "success"}}
],
},
{"role": "assistant", "content": [{"toolUse": {"toolUseId": "new", "name": "tool2", "input": {}}}]},
{
"role": "user",
"content": [
{
"toolResult": {
"toolUseId": "new",
"content": [{"text": "The tool result was too large!"}],
"status": "error",
}
}
],
},
]
test_agent = Agent(messages=messages)

manager.reduce_context(test_agent)

assert messages[0]["role"] == "user"
assert messages[0]["content"][0]["toolResult"]["toolUseId"] == "old"
assert messages[0]["content"][0]["toolResult"]["content"] == [{"text": "The tool result was too large!"}]
assert messages[0]["content"][0]["toolResult"]["status"] == "error"


def test_sliding_window_conversation_manager_raises_when_reduced_history_has_no_user_messages():
manager = SlidingWindowConversationManager(1, should_truncate_results=False)
messages = [
{"role": "user", "content": [{"text": "First message"}]},
{"role": "assistant", "content": [{"toolUse": {"toolUseId": "123", "name": "tool1", "input": {}}}]},
]
test_agent = Agent(messages=messages)

with pytest.raises(ContextWindowOverflowException, match="Unable to trim conversation context!"):
manager.apply_management(test_agent)


def test_null_conversation_manager_reduce_context_raises_context_window_overflow_exception():
"""Test that NullConversationManager doesn't modify messages."""
manager = NullConversationManager()
Expand Down