fix: all tool uses now must be replaced

Author: dbschmigelski

src/strands/event_loop/_recover_message_on_max_tokens_reached.py

@@ -14,39 +14,44 @@
 
 
 def recover_message_on_max_tokens_reached(message: Message) -> Message:
-    """Recover and clean up incomplete messages when max token limits are reached.
+    """Recover and clean up messages when max token limits are reached.
 
-    When a model response is truncated due to maximum token limits, tool use blocks may be
-    incomplete or malformed. This function inspects the message content and:
+    When a model response is truncated due to maximum token limits, all tool use blocks
+    should be replaced with informative error messages since they may be incomplete or
+    unreliable. This function inspects the message content and:
 
-    1. Identifies incomplete tool use blocks (missing name, input, or toolUseId)
-    2. Replaces incomplete tool uses with informative error messages
-    3. Preserves all valid content blocks (text and complete tool uses)
+    1. Identifies all tool use blocks (regardless of validity)
+    2. Replaces all tool uses with informative error messages
+    3. Preserves all non-tool content blocks (text, images, etc.)
     4. Returns a cleaned message suitable for conversation history
 
     This recovery mechanism ensures that the conversation can continue gracefully even when
-    model responses are truncated, providing clear feedback about what happened.
+    model responses are truncated, providing clear feedback about what happened and preventing
+    potentially incomplete or corrupted tool executions.
+
+    TODO: after https://github.com/strands-agents/sdk-python/issues/561 is completed, only the verifiable
+        invalid tool_use content blocks need to be replaced.
 
     Args:
         message: The potentially incomplete message from the model that was truncated
                 due to max token limits.
 
     Returns:
-        A cleaned Message with incomplete tool uses replaced by explanatory text content.
+        A cleaned Message with all tool uses replaced by explanatory text content.
         The returned message maintains the same role as the input message.
 
     Example:
-        If a message contains an incomplete tool use like:
+        If a message contains any tool use (complete or incomplete):
         ```
-        {"toolUse": {"name": "calculator"}}  # missing input and toolUseId
+        {"toolUse": {"name": "calculator", "input": {"expression": "2+2"}, "toolUseId": "123"}}
         ```
 
         It will be replaced with:
         ```
         {"text": "The selected tool calculator's tool use was incomplete due to maximum token limits being reached."}
         ```
     """
-    logger.info("handling max_tokens stop reason - inspecting incomplete message for invalid tool uses")
+    logger.info("handling max_tokens stop reason - replacing all tool uses with error messages")
 
     valid_content: list[ContentBlock] = []
     for content in message["content"] or []:
@@ -55,15 +60,8 @@ def recover_message_on_max_tokens_reached(message: Message) -> Message:
             valid_content.append(content)
             continue
 
-        # Check if tool use is incomplete (missing or empty required fields)
-        tool_name = tool_use.get("name")
-        if tool_name and tool_use.get("input") and tool_use.get("toolUseId"):
-            # As far as we can tell, tool use is valid if this condition is true
-            valid_content.append(content)
-            continue
-
-        # Tool use is incomplete due to max_tokens truncation
-        display_name = tool_name if tool_name else "<unknown>"
+        # Replace all tool uses with error messages when max_tokens is reached
+        display_name = tool_use.get("name", "<unknown>")
         logger.warning("tool_name=<%s> | replacing with error message due to max_tokens truncation.", display_name)
 
         valid_content.append(

tests/strands/event_loop/test_recover_message_on_max_tokens_reached.py

@@ -95,7 +95,7 @@ def test_recover_message_on_max_tokens_reached_with_missing_tool_use_id():
 
 
 def test_recover_message_on_max_tokens_reached_with_valid_tool_use():
-    """Test that valid tool uses are preserved unchanged."""
+    """Test that even valid tool uses are replaced with error messages."""
     complete_message: Message = {
         "role": "assistant",
         "content": [
@@ -106,13 +106,15 @@ def test_recover_message_on_max_tokens_reached_with_valid_tool_use():
 
     result = recover_message_on_max_tokens_reached(complete_message)
 
-    # Should preserve the message exactly as-is
+    # Should replace even valid tool uses with error messages
     assert result["role"] == "assistant"
     assert len(result["content"]) == 2
     assert result["content"][0] == {"text": "I'll help you with that."}
-    assert result["content"][1] == {
-        "toolUse": {"name": "calculator", "input": {"expression": "2+2"}, "toolUseId": "123"}
-    }
+
+    # Valid tool use should also be replaced with error message
+    assert "text" in result["content"][1]
+    assert "calculator" in result["content"][1]["text"]
+    assert "incomplete due to maximum token limits" in result["content"][1]["text"]
 
 
 def test_recover_message_on_max_tokens_reached_with_empty_content():