feat: switch to a default behavior to recover from max tokens reached

Author: dbschmigelski

src/strands/agent/agent.py

@@ -37,7 +37,7 @@
 from ..tools.registry import ToolRegistry
 from ..tools.watcher import ToolWatcher
 from ..types.content import ContentBlock, Message, Messages
-from ..types.exceptions import ContextWindowOverflowException, MaxTokensReachedException
+from ..types.exceptions import ContextWindowOverflowException
 from ..types.tools import ToolResult, ToolUse
 from ..types.traces import AttributeValue
 from .agent_result import AgentResult
@@ -582,21 +582,18 @@ async def _execute_event_loop_cycle(self, invocation_state: dict[str, Any]) -> A
             )
             async for event in events:
                 yield event
-            return
+
         except ContextWindowOverflowException as e:
             # Try reducing the context size and retrying
-            self.conversation_manager.reduce_context(agent=self, e=e)
-        except MaxTokensReachedException as e:
-            # Recover conversation state after token limit exceeded, then continue with next cycle
-            await self.conversation_manager.handle_token_limit_reached(agent=self, e=e)
+            self.conversation_manager.reduce_context(self, e=e)
 
-        # Sync agent after handling exception to keep conversation_manager_state up to date in the session
-        if self._session_manager:
-            self._session_manager.sync_agent(self)
+            # Sync agent after reduce_context to keep conversation_manager_state up to date in the session
+            if self._session_manager:
+                self._session_manager.sync_agent(self)
 
-        events = self._execute_event_loop_cycle(invocation_state)
-        async for event in events:
-            yield event
+            events = self._execute_event_loop_cycle(invocation_state)
+            async for event in events:
+                yield event
 
     def _record_tool_execution(
         self,

src/strands/agent/conversation_manager/__init__.py

@@ -15,14 +15,12 @@
 
 from .conversation_manager import ConversationManager
 from .null_conversation_manager import NullConversationManager
-from .recover_tool_use_on_max_tokens_reached import recover_tool_use_on_max_tokens_reached
 from .sliding_window_conversation_manager import SlidingWindowConversationManager
 from .summarizing_conversation_manager import SummarizingConversationManager
 
 __all__ = [
     "ConversationManager",
     "NullConversationManager",
-    "recover_tool_use_on_max_tokens_reached",
     "SlidingWindowConversationManager",
     "SummarizingConversationManager",
 ]

src/strands/agent/conversation_manager/conversation_manager.py

@@ -4,7 +4,6 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 from ...types.content import Message
-from ...types.exceptions import MaxTokensReachedException
 
 if TYPE_CHECKING:
     from ...agent.agent import Agent
@@ -87,17 +86,3 @@ def reduce_context(self, agent: "Agent", e: Optional[Exception] = None, **kwargs
             **kwargs: Additional keyword arguments for future extensibility.
         """
         pass
-
-    async def handle_token_limit_reached(self, agent: "Agent", e: MaxTokensReachedException, **kwargs: Any) -> None:
-        """Called when MaxTokensReachedException is thrown to recover conversation state.
-
-        This method should implement recovery strategies when the token limit is exceeded and the message array
-        may be in a broken state.
-
-        Args:
-            agent: The agent whose conversation state will be recovered.
-                This list is modified in-place.
-            e: The MaxTokensReachedException that triggered the recovery.
-            **kwargs: Additional keyword arguments for future extensibility.
-        """
-        raise e

src/strands/agent/conversation_manager/recover_tool_use_on_max_tokens_reached.py

@@ -7,9 +7,8 @@
     from ...agent.agent import Agent
 
 from ...types.content import Messages
-from ...types.exceptions import ContextWindowOverflowException, MaxTokensReachedException
+from ...types.exceptions import ContextWindowOverflowException
 from .conversation_manager import ConversationManager
-from .recover_tool_use_on_max_tokens_reached import recover_tool_use_on_max_tokens_reached
 
 logger = logging.getLogger(__name__)
 
@@ -113,16 +112,6 @@ def reduce_context(self, agent: "Agent", e: Optional[Exception] = None, **kwargs
         # Overwrite message history
         messages[:] = messages[trim_index:]
 
-    async def handle_token_limit_reached(self, agent: "Agent", e: MaxTokensReachedException, **kwargs: Any) -> None:
-        """Apply sliding window strategy for token limit recovery.
-
-        Args:
-            agent: The agent whose conversation state will be recovered.
-            e: The MaxTokensReachedException that triggered the recovery.
-            **kwargs: Additional keyword arguments for future extensibility.
-        """
-        await recover_tool_use_on_max_tokens_reached(agent, e)
-
     def _truncate_tool_results(self, messages: Messages, msg_idx: int) -> bool:
         """Truncate tool results in a message to reduce context size.
 

src/strands/agent/conversation_manager/sliding_window_conversation_manager.py

@@ -6,9 +6,8 @@
 from typing_extensions import override
 
 from ...types.content import Message
-from ...types.exceptions import ContextWindowOverflowException, MaxTokensReachedException
+from ...types.exceptions import ContextWindowOverflowException
 from .conversation_manager import ConversationManager
-from .recover_tool_use_on_max_tokens_reached import recover_tool_use_on_max_tokens_reached
 
 if TYPE_CHECKING:
     from ..agent import Agent
@@ -167,16 +166,6 @@ def reduce_context(self, agent: "Agent", e: Optional[Exception] = None, **kwargs
             logger.error("Summarization failed: %s", summarization_error)
             raise summarization_error from e
 
-    async def handle_token_limit_reached(self, agent: "Agent", e: MaxTokensReachedException, **kwargs: Any) -> None:
-        """Apply summarization strategy for token limit recovery.
-
-        Args:
-            agent: The agent whose conversation state will be recovered.
-            e: The MaxTokensReachedException that triggered the recovery.
-            **kwargs: Additional keyword arguments for future extensibility.
-        """
-        await recover_tool_use_on_max_tokens_reached(agent, e)
-
     def _generate_summary(self, messages: List[Message], agent: "Agent") -> Message:
         """Generate a summary of the provided messages.
 

src/strands/agent/conversation_manager/summarizing_conversation_manager.py

@@ -0,0 +1,76 @@
+"""Message recovery utilities for handling max token limit scenarios.
+
+This module provides functionality to recover and clean up incomplete messages that occur
+when model responses are truncated due to maximum token limits being reached. It specifically
+handles cases where tool use blocks are incomplete or malformed due to truncation.
+"""
+
+import logging
+
+from ..types.content import ContentBlock, Message
+from ..types.tools import ToolUse
+
+logger = logging.getLogger(__name__)
+
+
+def recover_message_on_max_tokens_reached(message: Message) -> Message:
+    """Recover and clean up incomplete 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:
+
+    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)
+    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.
+
+    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.
+        The returned message maintains the same role as the input message.
+
+    Example:
+        If a message contains an incomplete tool use like:
+        ```
+        {"toolUse": {"name": "calculator"}}  # missing input and toolUseId
+        ```
+
+        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")
+
+    valid_content: list[ContentBlock] = []
+    for content in message["content"] or []:
+        tool_use: ToolUse | None = content.get("toolUse")
+        if not tool_use:
+            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>"
+        logger.warning("tool_name=<%s> | replacing with error message due to max_tokens truncation.", display_name)
+
+        valid_content.append(
+            {
+                "text": f"The selected tool {display_name}'s tool use was incomplete due "
+                f"to maximum token limits being reached."
+            }
+        )
+
+    return {"content": valid_content, "role": message["role"]}

src/strands/event_loop/_recover_message_on_max_tokens_reached.py

@@ -36,6 +36,7 @@
 )
 from ..types.streaming import Metrics, StopReason
 from ..types.tools import ToolChoice, ToolChoiceAuto, ToolConfig, ToolGenerator, ToolResult, ToolUse
+from ._recover_message_on_max_tokens_reached import recover_message_on_max_tokens_reached
 from .streaming import stream_messages
 
 if TYPE_CHECKING:
@@ -156,6 +157,9 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
                     )
                 )
 
+                if stop_reason == "max_tokens":
+                    message = recover_message_on_max_tokens_reached(message)
+
                 if model_invoke_span:
                     tracer.end_model_invoke_span(model_invoke_span, message, usage, stop_reason)
                 break  # Success! Break out of retry loop
@@ -192,6 +196,19 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
                     raise e
 
     try:
+        # Add message in trace and mark the end of the stream messages trace
+        stream_trace.add_message(message)
+        stream_trace.end()
+
+        # Add the response message to the conversation
+        agent.messages.append(message)
+        agent.hooks.invoke_callbacks(MessageAddedEvent(agent=agent, message=message))
+        yield {"callback": {"message": message}}
+
+        # Update metrics
+        agent.event_loop_metrics.update_usage(usage)
+        agent.event_loop_metrics.update_metrics(metrics)
+
         if stop_reason == "max_tokens":
             """
             Handle max_tokens limit reached by the model.
@@ -205,21 +222,8 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
                     "Agent has reached an unrecoverable state due to max_tokens limit. "
                     "For more information see: "
                     "https://strandsagents.com/latest/user-guide/concepts/agents/agent-loop/#maxtokensreachedexception"
-                ),
-                incomplete_message=message,
+                )
             )
-        # Add message in trace and mark the end of the stream messages trace
-        stream_trace.add_message(message)
-        stream_trace.end()
-
-        # Add the response message to the conversation
-        agent.messages.append(message)
-        agent.hooks.invoke_callbacks(MessageAddedEvent(agent=agent, message=message))
-        yield {"callback": {"message": message}}
-
-        # Update metrics
-        agent.event_loop_metrics.update_usage(usage)
-        agent.event_loop_metrics.update_metrics(metrics)
 
         # If the model is requesting to use tools
         if stop_reason == "tool_use":

src/strands/event_loop/event_loop.py

@@ -2,8 +2,6 @@
 
 from typing import Any
 
-from strands.types.content import Message
-
 
 class EventLoopException(Exception):
     """Exception raised by the event loop."""
@@ -28,14 +26,12 @@ class MaxTokensReachedException(Exception):
     the complexity of the response, or when the model naturally reaches its configured output limit during generation.
     """
 
-    def __init__(self, message: str, incomplete_message: Message):
+    def __init__(self, message: str):
         """Initialize the exception with an error message and the incomplete message object.
 
         Args:
             message: The error message describing the token limit issue
-            incomplete_message: The valid Message object with incomplete content due to token limits
         """
-        self.incomplete_message = incomplete_message
         super().__init__(message)