Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 29 additions & 3 deletions src/strands/agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -752,9 +752,9 @@ async def _run_loop(
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
Expand Down Expand Up @@ -969,3 +969,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
56 changes: 56 additions & 0 deletions tests/strands/agent/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -2160,3 +2160,59 @@ def shell(command: str):

# And that it continued to the LLM call
assert agent.messages[-1] == {"content": [{"text": "I invoked a tool!"}], "role": "assistant"}



@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
77 changes: 75 additions & 2 deletions tests_integ/test_bedrock_guardrails.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand All @@ -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,
Expand Down