fix(bedrock): include assistant response in guardrail_last_turn_only context

Author: aianch

src/strands/models/bedrock.py

@@ -311,14 +311,14 @@ def _format_request(
     def _get_last_turn_messages(self, messages: Messages) -> Messages:
         """Get the last turn messages for guardrail evaluation.
 
-        Returns the latest user message and the previous assistant message (if it exists).
+        Returns the latest user message and the assistant's response (if it exists).
         This reduces the conversation context sent to guardrails when guardrail_last_turn_only is True.
 
         Args:
             messages: Full conversation messages.
 
         Returns:
-            Messages containing only the last turn (user + previous assistant if exists).
+            Messages containing only the last turn (user + assistant response if exists).
         """
         if not messages:
             return []
@@ -334,13 +334,12 @@ def _get_last_turn_messages(self, messages: Messages) -> Messages:
             # No user message found, return empty
             return []
 
-        # Include the previous assistant message if it exists
-        result_messages: Messages = []
-        if last_user_index > 0 and messages[last_user_index - 1]["role"] == "assistant":
-            result_messages.append(messages[last_user_index - 1])
+        # Start with the last user message
+        result_messages: Messages = [messages[last_user_index]]
 
-        # Add the last user message
-        result_messages.append(messages[last_user_index])
+        # Include the assistant's response if it exists (the message after the user message)
+        if last_user_index < len(messages) - 1 and messages[last_user_index + 1]["role"] == "assistant":
+            result_messages.append(messages[last_user_index + 1])
 
         return result_messages
 

tests/strands/models/test_bedrock.py

@@ -2057,29 +2057,43 @@ def test_get_last_turn_messages(model):
     # Test empty messages
     assert model._get_last_turn_messages([]) == []
 
-    # Test single user message
+    # Test single user message (no assistant response yet)
     messages = [{"role": "user", "content": [{"text": "Hello"}]}]
     result = model._get_last_turn_messages(messages)
     assert len(result) == 1
     assert result[0]["role"] == "user"
 
-    # Test user-assistant pair
+    # Test user-assistant pair (user message with assistant response)
     messages = [
         {"role": "user", "content": [{"text": "Hello"}]},
         {"role": "assistant", "content": [{"text": "Hi"}]},
         {"role": "user", "content": [{"text": "How are you?"}]},
+        {"role": "assistant", "content": [{"text": "I'm doing well"}]},
     ]
     result = model._get_last_turn_messages(messages)
     assert len(result) == 2
-    assert result[0]["role"] == "assistant"
-    assert result[1]["role"] == "user"
-    assert result[1]["content"][0]["text"] == "How are you?"
+    assert result[0]["role"] == "user"
+    assert result[0]["content"][0]["text"] == "How are you?"
+    assert result[1]["role"] == "assistant"
+    assert result[1]["content"][0]["text"] == "I'm doing well"
+
+    # Test last user message without assistant response
+    messages = [
+        {"role": "user", "content": [{"text": "Hello"}]},
+        {"role": "assistant", "content": [{"text": "Hi"}]},
+        {"role": "user", "content": [{"text": "How are you?"}]},
+    ]
+    result = model._get_last_turn_messages(messages)
+    assert len(result) == 1
+    assert result[0]["role"] == "user"
+    assert result[0]["content"][0]["text"] == "How are you?"
 
 
 def test_format_request_with_guardrail_last_turn_only(model, model_id):
     """Test _format_request uses filtered messages when guardrail_last_turn_only=True."""
     model.update_config(guardrail_id="test-guardrail", guardrail_version="DRAFT", guardrail_last_turn_only=True)
 
+    # Test with last user message only (no assistant response yet)
     messages = [
         {"role": "user", "content": [{"text": "First message"}]},
         {"role": "assistant", "content": [{"text": "First response"}]},
@@ -2088,12 +2102,29 @@ def test_format_request_with_guardrail_last_turn_only(model, model_id):
 
     request = model._format_request(messages)
 
-    # Should only include the last turn (assistant + user)
+    # Should only include the last user message (no assistant response after it yet)
     formatted_messages = request["messages"]
-    assert len(formatted_messages) == 2
-    assert formatted_messages[0]["role"] == "assistant"
-    assert formatted_messages[1]["role"] == "user"
-    assert formatted_messages[1]["content"][0]["text"] == "Latest message"
+    assert len(formatted_messages) == 1
+    assert formatted_messages[0]["role"] == "user"
+    assert formatted_messages[0]["content"][0]["text"] == "Latest message"
+
+    # Test with last user message + assistant response
+    messages_with_response = [
+        {"role": "user", "content": [{"text": "First message"}]},
+        {"role": "assistant", "content": [{"text": "First response"}]},
+        {"role": "user", "content": [{"text": "How are you?"}]},
+        {"role": "assistant", "content": [{"text": "I'm good"}]},
+    ]
+
+    request2 = model._format_request(messages_with_response)
+
+    # Should include last user + assistant response
+    formatted_messages2 = request2["messages"]
+    assert len(formatted_messages2) == 2
+    assert formatted_messages2[0]["role"] == "user"
+    assert formatted_messages2[0]["content"][0]["text"] == "How are you?"
+    assert formatted_messages2[1]["role"] == "assistant"
+    assert formatted_messages2[1]["content"][0]["text"] == "I'm good"
 
 
 @pytest.mark.asyncio

tests_integ/test_bedrock_guardrails.py

@@ -343,6 +343,39 @@ def test_guardrail_last_turn_only_recovery_scenario(boto_session, bedrock_guardr
     assert len(agent.messages) == 4  # 2 user + 2 assistant messages
 
 
+def test_guardrail_last_turn_only_output_intervention(boto_session, bedrock_guardrail):
+    """Test that guardrail_last_turn_only works with OUTPUT guardrails.
+
+    This tests that when the assistant tries to output blocked content,
+    the OUTPUT guardrail intervenes, even with guardrail_last_turn_only=True.
+    Then verifies that subsequent normal responses work correctly.
+    """
+    bedrock_model = BedrockModel(
+        guardrail_id=bedrock_guardrail,
+        guardrail_version="DRAFT",
+        guardrail_last_turn_only=True,
+        guardrail_stream_processing_mode="sync",
+        boto_session=boto_session,
+    )
+
+    agent = Agent(
+        model=bedrock_model,
+        system_prompt="When asked to say the word, say CACTUS. Otherwise respond normally.",
+        callback_handler=None,
+        load_tools_from_directory=False,
+    )
+
+    # First turn - assistant tries to output "CACTUS", should be blocked by OUTPUT guardrail
+    response1 = agent("Say the word.")
+    assert response1.stop_reason == "guardrail_intervened"
+    assert BLOCKED_OUTPUT in str(response1)
+
+    # Second turn - normal question should work fine
+    response2 = agent("What is 2+2?")
+    assert response2.stop_reason != "guardrail_intervened"
+    assert BLOCKED_OUTPUT not in str(response2)
+
+
 def test_guardrail_input_intervention_properly_redacts_in_session(boto_session, bedrock_guardrail, temp_dir):
     bedrock_model = BedrockModel(
         guardrail_id=bedrock_guardrail,