feat: Update SlidingWindowConversationManager (strands-agents#120)

Author: Unshure

src/strands/agent/conversation_manager/sliding_window_conversation_manager.py

@@ -1,12 +1,10 @@
 """Sliding window conversation history management."""
 
-import json
 import logging
-from typing import List, Optional, cast
+from typing import Optional
 
-from ...types.content import ContentBlock, Message, Messages
+from ...types.content import Message, Messages
 from ...types.exceptions import ContextWindowOverflowException
-from ...types.tools import ToolResult
 from .conversation_manager import ConversationManager
 
 logger = logging.getLogger(__name__)
@@ -110,8 +108,9 @@ def _remove_dangling_messages(self, messages: Messages) -> None:
     def reduce_context(self, messages: Messages, e: Optional[Exception] = None) -> None:
         """Trim the oldest messages to reduce the conversation context size.
 
-        The method handles special cases where tool results need to be converted to regular content blocks to maintain
-        conversation coherence after trimming.
+        The method handles special cases where trimming the messages leads to:
+         - toolResult with no corresponding toolUse
+         - toolUse with no corresponding toolResult
 
         Args:
             messages: The messages to reduce.
@@ -126,52 +125,24 @@ def reduce_context(self, messages: Messages, e: Optional[Exception] = None) -> N
         # If the number of messages is less than the window_size, then we default to 2, otherwise, trim to window size
         trim_index = 2 if len(messages) <= self.window_size else len(messages) - self.window_size
 
-        # Throw if we cannot trim any messages from the conversation
-        if trim_index >= len(messages):
-            raise ContextWindowOverflowException("Unable to trim conversation context!") from e
-
-        # If the message at the cut index has ToolResultContent, then we map that to ContentBlock. This gets around the
-        # limitation of needing ToolUse and ToolResults to be paired.
-        if any("toolResult" in content for content in messages[trim_index]["content"]):
-            if len(messages[trim_index]["content"]) == 1:
-                messages[trim_index]["content"] = self._map_tool_result_content(
-                    cast(ToolResult, messages[trim_index]["content"][0]["toolResult"])
+        # Find the next valid trim_index
+        while trim_index < len(messages):
+            if (
+                # Oldest message cannot be a toolResult because it needs a toolUse preceding it
+                any("toolResult" in content for content in messages[trim_index]["content"])
+                or (
+                    # Oldest message can be a toolUse only if a toolResult immediately follows it.
+                    any("toolUse" in content for content in messages[trim_index]["content"])
+                    and trim_index + 1 < len(messages)
+                    and not any("toolResult" in content for content in messages[trim_index + 1]["content"])
                 )
-
-            # If there is more content than just one ToolResultContent, then we cannot cut at this index.
+            ):
+                trim_index += 1
             else:
-                raise ContextWindowOverflowException("Unable to trim conversation context!") from e
+                break
+        else:
+            # If we didn't find a valid trim_index, then we throw
+            raise ContextWindowOverflowException("Unable to trim conversation context!") from e
 
         # Overwrite message history
         messages[:] = messages[trim_index:]
-
-    def _map_tool_result_content(self, tool_result: ToolResult) -> List[ContentBlock]:
-        """Convert a ToolResult to a list of standard ContentBlocks.
-
-        This method transforms tool result content into standard content blocks that can be preserved when trimming the
-        conversation history.
-
-        Args:
-            tool_result: The ToolResult to convert.
-
-        Returns:
-            A list of content blocks representing the tool result.
-        """
-        contents = []
-        text_content = "Tool Result Status: " + tool_result["status"] if tool_result["status"] else ""
-
-        for tool_result_content in tool_result["content"]:
-            if "text" in tool_result_content:
-                text_content = "\nTool Result Text Content: " + tool_result_content["text"] + f"\n{text_content}"
-            elif "json" in tool_result_content:
-                text_content = (
-                    "\nTool Result JSON Content: " + json.dumps(tool_result_content["json"]) + f"\n{text_content}"
-                )
-            elif "image" in tool_result_content:
-                contents.append(ContentBlock(image=tool_result_content["image"]))
-            elif "document" in tool_result_content:
-                contents.append(ContentBlock(document=tool_result_content["document"]))
-            else:
-                logger.warning("unsupported content type")
-        contents.append(ContentBlock(text=text_content))
-        return contents

tests/strands/agent/test_agent.py

@@ -459,7 +459,7 @@ def test_agent__call__always_sliding_window_conversation_manager_doesnt_infinite
     with pytest.raises(ContextWindowOverflowException):
         agent("Test!")
 
-    assert conversation_manager_spy.reduce_context.call_count == 251
+    assert conversation_manager_spy.reduce_context.call_count > 0
     assert conversation_manager_spy.apply_management.call_count == 1
 
 

tests/strands/agent/test_conversation_manager.py

@@ -111,41 +111,7 @@ def conversation_manager(request):
                 {"role": "assistant", "content": [{"text": "Second response"}]},
             ],
         ),
-        # 7 - Message count above max window size - Remove dangling tool uses and tool results
-        (
-            {"window_size": 1},
-            [
-                {"role": "user", "content": [{"text": "First message"}]},
-                {"role": "assistant", "content": [{"toolUse": {"toolUseId": "321", "name": "tool1", "input": {}}}]},
-                {
-                    "role": "user",
-                    "content": [
-                        {"toolResult": {"toolUseId": "123", "content": [{"text": "Hello!"}], "status": "success"}}
-                    ],
-                },
-            ],
-            [
-                {
-                    "role": "user",
-                    "content": [{"text": "\nTool Result Text Content: Hello!\nTool Result Status: success"}],
-                },
-            ],
-        ),
-        # 8 - Message count above max window size - Remove multiple tool use/tool result pairs
-        (
-            {"window_size": 1},
-            [
-                {"role": "user", "content": [{"toolResult": {"toolUseId": "123", "content": [], "status": "success"}}]},
-                {"role": "assistant", "content": [{"toolUse": {"toolUseId": "123", "name": "tool1", "input": {}}}]},
-                {"role": "user", "content": [{"toolResult": {"toolUseId": "456", "content": [], "status": "success"}}]},
-                {"role": "assistant", "content": [{"toolUse": {"toolUseId": "456", "name": "tool1", "input": {}}}]},
-                {"role": "user", "content": [{"toolResult": {"toolUseId": "789", "content": [], "status": "success"}}]},
-            ],
-            [
-                {"role": "user", "content": [{"text": "Tool Result Status: success"}]},
-            ],
-        ),
-        # 9 - Message count above max window size - Preserve tool use/tool result pairs
+        # 7 - Message count above max window size - Preserve tool use/tool result pairs
         (
             {"window_size": 2},
             [
@@ -158,7 +124,7 @@ def conversation_manager(request):
                 {"role": "user", "content": [{"toolResult": {"toolUseId": "456", "content": [], "status": "success"}}]},
             ],
         ),
-        # 10 - Test sliding window behavior - preserve tool use/result pairs across cut boundary
+        # 8 - Test sliding window behavior - preserve tool use/result pairs across cut boundary
         (
             {"window_size": 3},
             [
@@ -173,7 +139,7 @@ def conversation_manager(request):
                 {"role": "assistant", "content": [{"text": "Response after tool use"}]},
             ],
         ),
-        # 11 - Test sliding window with multiple tool pairs that need preservation
+        # 9 - Test sliding window with multiple tool pairs that need preservation
         (
             {"window_size": 4},
             [
@@ -185,7 +151,6 @@ def conversation_manager(request):
                 {"role": "assistant", "content": [{"text": "Final response"}]},
             ],
             [
-                {"role": "user", "content": [{"text": "Tool Result Status: success"}]},
                 {"role": "assistant", "content": [{"toolUse": {"toolUseId": "456", "name": "tool2", "input": {}}}]},
                 {"role": "user", "content": [{"toolResult": {"toolUseId": "456", "content": [], "status": "success"}}]},
                 {"role": "assistant", "content": [{"text": "Final response"}]},
@@ -200,6 +165,20 @@ def test_apply_management(conversation_manager, messages, expected_messages):
     assert messages == expected_messages
 
 
+def test_sliding_window_conversation_manager_with_untrimmable_history_raises_context_window_overflow_exception():
+    manager = strands.agent.conversation_manager.SlidingWindowConversationManager(1)
+    messages = [
+        {"role": "assistant", "content": [{"toolUse": {"toolUseId": "456", "name": "tool1", "input": {}}}]},
+        {"role": "user", "content": [{"toolResult": {"toolUseId": "789", "content": [], "status": "success"}}]},
+    ]
+    original_messages = messages.copy()
+
+    with pytest.raises(ContextWindowOverflowException):
+        manager.apply_management(messages)
+
+    assert messages == original_messages
+
+
 def test_null_conversation_manager_reduce_context_raises_context_window_overflow_exception():
     """Test that NullConversationManager doesn't modify messages."""
     manager = strands.agent.conversation_manager.NullConversationManager()