fix: properly redact toolResult blocks

Author: leotac

src/strands/agent/agent.py

@@ -687,9 +687,9 @@ async def _run_loop(self, messages: Messages, invocation_state: dict[str, Any])
                     and event.chunk.get("redactContent")
                     and event.chunk["redactContent"].get("redactUserContentMessage")
                 ):
-                    self.messages[-1]["content"] = [
-                        {"text": str(event.chunk["redactContent"]["redactUserContentMessage"])}
-                    ]
+                    self.messages[-1]["content"] = self._redact_user_content(
+                        self.messages[-1]["content"], str(event.chunk["redactContent"]["redactUserContentMessage"])
+                    )
                     if self._session_manager:
                         self._session_manager.redact_latest_message(self.messages[-1], self)
                 yield event
@@ -891,3 +891,29 @@ def _append_message(self, message: Message) -> None:
         """Appends a message to the agent's list of messages and invokes the callbacks for the MessageCreatedEvent."""
         self.messages.append(message)
         self.hooks.invoke_callbacks(MessageAddedEvent(agent=self, message=message))
+
+    def _redact_user_content(self, content: list[ContentBlock], redact_message: str) -> list[ContentBlock]:
+        """Redact user content preserving toolResult blocks.
+
+        Args:
+            content: content blocks to be redacted
+            redact_message: redact message to be replaced
+
+        Returns:
+            Redacted content, as follows:
+            - if the message contains at least a toolResult block,
+                all toolResult blocks(s) are kept, redacting only the result content;
+            - otherwise, the entire content of the message is replaced
+                with a single text block with the redact message.
+        """
+        redacted_content = []
+        for block in content:
+            if "toolResult" in block:
+                block["toolResult"]["content"] = [{"text": redact_message}]
+                redacted_content.append(block)
+
+        if not redacted_content:
+            # Text content is added only if no toolResult blocks were found
+            redacted_content = [{"text": redact_message}]
+
+        return redacted_content

tests/strands/agent/test_agent.py

@@ -2061,3 +2061,58 @@ def test_agent_tool_caller_interrupt(user):
     exp_message = r"cannot directly call tool during interrupt"
     with pytest.raises(RuntimeError, match=exp_message):
         agent.tool.test_tool()
+
+
+@pytest.mark.parametrize(
+    "content, expected",
+    [
+        # Single toolResult block - preserves structure, redacts content
+        (
+            [{"toolResult": {"toolUseId": "123", "content": [{"text": "original result"}], "status": "success"}}],
+            [{"toolResult": {"toolUseId": "123", "content": [{"text": "REDACTED"}], "status": "success"}}],
+        ),
+        # Multiple toolResult blocks - preserves all, redacts each content
+        (
+            [
+                {"toolResult": {"toolUseId": "123", "content": [{"text": "result1"}], "status": "success"}},
+                {"toolResult": {"toolUseId": "456", "content": [{"text": "result2"}], "status": "error"}},
+            ],
+            [
+                {"toolResult": {"toolUseId": "123", "content": [{"text": "REDACTED"}], "status": "success"}},
+                {"toolResult": {"toolUseId": "456", "content": [{"text": "REDACTED"}], "status": "error"}},
+            ],
+        ),
+        # Text only content - replaces with single text block
+        (
+            [{"text": "sensitive data"}],
+            [{"text": "REDACTED"}],
+        ),
+        # Mixed content with toolResult - keeps only toolResult blocks
+        # (This should not actually happen, toolResult is never mixed with other content)
+        (
+            [
+                {"text": "some text"},
+                {"toolResult": {"toolUseId": "789", "content": [{"text": "tool output"}], "status": "success"}},
+                {"image": {"format": "png", "source": {"bytes": b"fake_data"}}},
+            ],
+            [{"toolResult": {"toolUseId": "789", "content": [{"text": "REDACTED"}], "status": "success"}}],
+        ),
+        # Empty content - returns single text block
+        (
+            [],
+            [{"text": "REDACTED"}],
+        ),
+    ],
+    ids=[
+        "single_tool_result",
+        "multiple_tool_results",
+        "text_only",
+        "mixed_content_with_tool_result",
+        "empty_content",
+    ],
+)
+def test_redact_user_content(content, expected):
+    """Test _redact_user_content function with various content types."""
+    agent = Agent()
+    result = agent._redact_user_content(content, "REDACTED")
+    assert result == expected

tests_integ/test_bedrock_guardrails.py

@@ -5,7 +5,7 @@
 import boto3
 import pytest
 
-from strands import Agent
+from strands import Agent, tool
 from strands.models.bedrock import BedrockModel
 from strands.session.file_session_manager import FileSessionManager
 
@@ -187,7 +187,7 @@ def test_guardrail_output_intervention_redact_output(bedrock_guardrail, processi
     In async streaming: The buffering is non-blocking. 
     Tokens are streamed while Guardrails processes the buffered content in the background. 
     This means the response may be returned before Guardrails has finished processing.
-    As a result, we cannot guarantee that the REDACT_MESSAGE is in the response
+    As a result, we cannot guarantee that the REDACT_MESSAGE is in the response.
     """
     if processing_mode == "sync":
         assert REDACT_MESSAGE in str(response1)
@@ -203,6 +203,79 @@ def test_guardrail_output_intervention_redact_output(bedrock_guardrail, processi
         )
 
 
+@pytest.mark.parametrize("processing_mode", ["sync", "async"])
+def test_guardrail_intervention_properly_redacts_tool_result(bedrock_guardrail, processing_mode):
+    INPUT_REDACT_MESSAGE = "Input redacted."
+    OUTPUT_REDACT_MESSAGE = "Output redacted."
+    bedrock_model = BedrockModel(
+        guardrail_id=bedrock_guardrail,
+        guardrail_version="DRAFT",
+        guardrail_stream_processing_mode=processing_mode,
+        guardrail_redact_output=True,
+        guardrail_redact_input_message=INPUT_REDACT_MESSAGE,
+        guardrail_redact_output_message=OUTPUT_REDACT_MESSAGE,
+        region_name="us-east-1",
+    )
+
+    @tool
+    def list_users() -> str:
+        "List my users"
+        return """[{"name": "Jerry Merry"}, {"name": "Mr. CACTUS"}]"""
+
+    agent = Agent(
+        model=bedrock_model,
+        system_prompt="You are a helpful assistant.",
+        callback_handler=None,
+        load_tools_from_directory=False,
+        tools=[list_users],
+    )
+
+    response1 = agent("List my users.")
+    response2 = agent("Thank you!")
+
+    """ Message sequence:
+    0 (user): request1
+    1 (assistant): reasoning + tool call
+    2 (user): tool result
+    3 (assistant): response1 -> output guardrail intervenes
+    4 (user): request2
+    5 (assistant): response2
+
+    Guardrail intervened on output in message 3 will cause
+    the redaction of the preceding input (message 2) and message 3.
+    """
+
+    assert response1.stop_reason == "guardrail_intervened"
+
+    if processing_mode == "sync":
+        """ In sync mode the guardrail processing is blocking.
+        The response is already blocked and redacted. """
+
+        assert OUTPUT_REDACT_MESSAGE in str(response1)
+        assert OUTPUT_REDACT_MESSAGE not in str(response2)
+
+    """
+    In async streaming, the buffering is non-blocking,
+    so the response may be returned before Guardrails has finished processing.
+
+    However, in both sync and async, with guardrail_redact_output=True:
+    
+    1. the content should be properly redacted in memory, so that
+    response2 is not blocked by guardrails;
+    """
+    assert response2.stop_reason != "guardrail_intervened"
+
+    """
+    2. the tool result block should be redacted properly, so that the
+    conversation is not corrupted.
+    """
+
+    tool_call = [b for b in agent.messages[1]["content"] if "toolUse" in b][0]["toolUse"]
+    tool_result = [b for b in agent.messages[2]["content"] if "toolResult" in b][0]["toolResult"]
+    assert tool_result["toolUseId"] == tool_call["toolUseId"]
+    assert tool_result["content"][0]["text"] == INPUT_REDACT_MESSAGE
+
+
 def test_guardrail_input_intervention_properly_redacts_in_session(boto_session, bedrock_guardrail, temp_dir):
     bedrock_model = BedrockModel(
         guardrail_id=bedrock_guardrail,