fix: override end_turn stop reason when streaming response contains toolUse blocks (#1827)

Co-authored-by: atian8179 <atian8179@users.noreply.github.com>
Co-authored-by: Nicholas Clegg <ncclegg@amazon.com>

Author: atian8179

src/strands/event_loop/streaming.py

@@ -324,16 +324,31 @@ def handle_content_block_stop(state: dict[str, Any]) -> dict[str, Any]:
     return state
 
 
-def handle_message_stop(event: MessageStopEvent) -> StopReason:
+def handle_message_stop(event: MessageStopEvent, content: list[dict[str, Any]]) -> StopReason:
     """Handles the end of a message by returning the stop reason.
 
+    Some models return "end_turn" even when tool calls are present, which prevents the event loop from processing
+    those tool calls. This function overrides to "tool_use" so tool execution proceeds correctly.
+
     Args:
         event: Stop event.
+        content: The message content blocks accumulated during streaming.
 
     Returns:
         The reason for stopping the stream.
     """
-    return event["stopReason"]
+    stop_reason = event["stopReason"]
+
+    if stop_reason == "end_turn" and any("toolUse" in item for item in content):
+        logger.warning(
+            "original_stop_reason=<%s>, new_stop_reason=<%s> | "
+            "overriding stop reason due to toolUse blocks in response",
+            "end_turn",
+            "tool_use",
+        )
+        stop_reason = "tool_use"
+
+    return stop_reason
 
 
 def handle_redact_content(event: RedactContentEvent, state: dict[str, Any]) -> None:
@@ -427,7 +442,7 @@ async def process_stream(
         elif "contentBlockStop" in chunk:
             state = handle_content_block_stop(state)
         elif "messageStop" in chunk:
-            stop_reason = handle_message_stop(chunk["messageStop"])
+            stop_reason = handle_message_stop(chunk["messageStop"], state["message"].get("content", []))
         elif "metadata" in chunk:
             time_to_first_byte_ms = (
                 int(1000 * (first_byte_time - start_time)) if (start_time and first_byte_time) else None

src/strands/models/bedrock.py

@@ -823,8 +823,6 @@ def _stream(
             logger.debug("got response from model")
             if streaming:
                 response = self.client.converse_stream(**request)
-                # Track tool use events to fix stopReason for streaming responses
-                has_tool_use = False
                 for chunk in response["stream"]:
                     if (
                         "metadata" in chunk
@@ -836,24 +834,7 @@ def _stream(
                             for event in self._generate_redaction_events():
                                 callback(event)
 
-                    # Track if we see tool use events
-                    if "contentBlockStart" in chunk and chunk["contentBlockStart"].get("start", {}).get("toolUse"):
-                        has_tool_use = True
-
-                    # Fix stopReason for streaming responses that contain tool use
-                    if (
-                        has_tool_use
-                        and "messageStop" in chunk
-                        and (message_stop := chunk["messageStop"]).get("stopReason") == "end_turn"
-                    ):
-                        # Create corrected chunk with tool_use stopReason
-                        modified_chunk = chunk.copy()
-                        modified_chunk["messageStop"] = message_stop.copy()
-                        modified_chunk["messageStop"]["stopReason"] = "tool_use"
-                        logger.warning("Override stop reason from end_turn to tool_use")
-                        callback(modified_chunk)
-                    else:
-                        callback(chunk)
+                    callback(chunk)
 
             else:
                 response = self.client.converse(**request)
@@ -992,17 +973,9 @@ def _convert_non_streaming_to_streaming(self, response: dict[str, Any]) -> Itera
             yield {"contentBlockStop": {}}
 
         # Yield messageStop event
-        # Fix stopReason for models that return end_turn when they should return tool_use on non-streaming side
-        current_stop_reason = response["stopReason"]
-        if current_stop_reason == "end_turn":
-            message_content = response["output"]["message"]["content"]
-            if any("toolUse" in content for content in message_content):
-                current_stop_reason = "tool_use"
-                logger.warning("Override stop reason from end_turn to tool_use")
-
         yield {
             "messageStop": {
-                "stopReason": current_stop_reason,
+                "stopReason": response["stopReason"],
                 "additionalModelResponseFields": response.get("additionalModelResponseFields"),
             }
         }

tests/strands/event_loop/test_streaming.py

@@ -530,12 +530,30 @@ def test_handle_content_block_stop(state, exp_updated_state):
 def test_handle_message_stop():
     event: MessageStopEvent = {"stopReason": "end_turn"}
 
-    tru_reason = strands.event_loop.streaming.handle_message_stop(event)
+    tru_reason = strands.event_loop.streaming.handle_message_stop(event, [])
     exp_reason = "end_turn"
 
     assert tru_reason == exp_reason
 
 
+def test_handle_message_stop_overrides_end_turn_when_tool_use_present():
+    event: MessageStopEvent = {"stopReason": "end_turn"}
+    content = [{"toolUse": {"toolUseId": "t1", "name": "myTool", "input": {}}}]
+
+    tru_reason = strands.event_loop.streaming.handle_message_stop(event, content)
+
+    assert tru_reason == "tool_use"
+
+
+def test_handle_message_stop_keeps_tool_use_unchanged():
+    event: MessageStopEvent = {"stopReason": "tool_use"}
+    content = [{"toolUse": {"toolUseId": "t1", "name": "myTool", "input": {}}}]
+
+    tru_reason = strands.event_loop.streaming.handle_message_stop(event, content)
+
+    assert tru_reason == "tool_use"
+
+
 def test_extract_usage_metrics():
     event = {
         "usage": {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
@@ -1334,3 +1352,68 @@ async def test_stream_messages_normalizes_messages(agenerator, alist):
         {"content": [{"toolUse": {"name": "INVALID_TOOL_NAME"}}], "role": "assistant"},
         {"content": [{"toolUse": {"name": "INVALID_TOOL_NAME"}}], "role": "assistant"},
     ]
+
+
+@pytest.mark.asyncio
+async def test_process_stream_overrides_end_turn_when_tool_use_present(agenerator, alist):
+    response = [
+        {"messageStart": {"role": "assistant"}},
+        {"contentBlockStart": {"contentBlockIndex": 0, "start": {"toolUse": {"toolUseId": "t1", "name": "myTool"}}}},
+        {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"key": "val"}'}}, "contentBlockIndex": 0}},
+        {"contentBlockStop": {"contentBlockIndex": 0}},
+        {"messageStop": {"stopReason": "end_turn"}},
+        {
+            "metadata": {
+                "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30},
+                "metrics": {"latencyMs": 100},
+            }
+        },
+    ]
+
+    stream = strands.event_loop.streaming.process_stream(agenerator(response))
+    last_event = cast(ModelStopReason, (await alist(stream))[-1])
+
+    assert last_event["stop"][0] == "tool_use"
+
+
+@pytest.mark.asyncio
+async def test_process_stream_keeps_end_turn_when_no_tool_use(agenerator, alist):
+    response = [
+        {"messageStart": {"role": "assistant"}},
+        {"contentBlockDelta": {"delta": {"text": "Hello!"}, "contentBlockIndex": 0}},
+        {"contentBlockStop": {"contentBlockIndex": 0}},
+        {"messageStop": {"stopReason": "end_turn"}},
+        {
+            "metadata": {
+                "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30},
+                "metrics": {"latencyMs": 100},
+            }
+        },
+    ]
+
+    stream = strands.event_loop.streaming.process_stream(agenerator(response))
+    last_event = cast(ModelStopReason, (await alist(stream))[-1])
+
+    assert last_event["stop"][0] == "end_turn"
+
+
+@pytest.mark.asyncio
+async def test_process_stream_keeps_tool_use_stop_reason_unchanged(agenerator, alist):
+    response = [
+        {"messageStart": {"role": "assistant"}},
+        {"contentBlockStart": {"contentBlockIndex": 0, "start": {"toolUse": {"toolUseId": "t1", "name": "myTool"}}}},
+        {"contentBlockDelta": {"delta": {"toolUse": {"input": "{}"}}, "contentBlockIndex": 0}},
+        {"contentBlockStop": {"contentBlockIndex": 0}},
+        {"messageStop": {"stopReason": "tool_use"}},
+        {
+            "metadata": {
+                "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30},
+                "metrics": {"latencyMs": 100},
+            }
+        },
+    ]
+
+    stream = strands.event_loop.streaming.process_stream(agenerator(response))
+    last_event = cast(ModelStopReason, (await alist(stream))[-1])
+
+    assert last_event["stop"][0] == "tool_use"

tests/strands/models/test_bedrock.py

@@ -1565,53 +1565,6 @@ async def test_stream_logging(bedrock_client, model, messages, caplog, alist):
     assert "finished streaming response from model" in log_text
 
 
-@pytest.mark.asyncio
-async def test_stream_stop_reason_override_streaming(bedrock_client, model, messages, alist):
-    """Test that stopReason is overridden from end_turn to tool_use in streaming mode when tool use is detected."""
-    bedrock_client.converse_stream.return_value = {
-        "stream": [
-            {"messageStart": {"role": "assistant"}},
-            {"contentBlockStart": {"start": {"toolUse": {"toolUseId": "123", "name": "test_tool"}}}},
-            {"contentBlockDelta": {"delta": {"test": {"input": '{"param": "value"}'}}}},
-            {"contentBlockStop": {}},
-            {"messageStop": {"stopReason": "end_turn"}},
-        ]
-    }
-
-    response = model.stream(messages)
-    events = await alist(response)
-
-    # Find the messageStop event
-    message_stop_event = next(event for event in events if "messageStop" in event)
-
-    # Verify stopReason was overridden to tool_use
-    assert message_stop_event["messageStop"]["stopReason"] == "tool_use"
-
-
-@pytest.mark.asyncio
-async def test_stream_stop_reason_override_non_streaming(bedrock_client, alist, messages):
-    """Test that stopReason is overridden from end_turn to tool_use in non-streaming mode when tool use is detected."""
-    bedrock_client.converse.return_value = {
-        "output": {
-            "message": {
-                "role": "assistant",
-                "content": [{"toolUse": {"toolUseId": "123", "name": "test_tool", "input": {"param": "value"}}}],
-            }
-        },
-        "stopReason": "end_turn",
-    }
-
-    model = BedrockModel(model_id="test-model", streaming=False)
-    response = model.stream(messages)
-    events = await alist(response)
-
-    # Find the messageStop event
-    message_stop_event = next(event for event in events if "messageStop" in event)
-
-    # Verify stopReason was overridden to tool_use
-    assert message_stop_event["messageStop"]["stopReason"] == "tool_use"
-
-
 def test_format_request_cleans_tool_result_content_blocks(model, model_id):
     messages = [
         {