fix: implement update_message() for guardrail redaction support#388
fix: implement update_message() for guardrail redaction support#388
Conversation
✅ No Breaking Changes DetectedNo public API breaking changes found in this PR. |
If create_message succeeds but delete_event fails, attempt to roll back the newly created event to avoid leaving duplicate messages. Addresses review comment about partial failure handling.
…essage Prevents stale eventId references by updating the tracked latest message with the new eventId after a successful create+delete replacement.
828699f to
a6df5f4
Compare
jariy17
left a comment
There was a problem hiding this comment.
Detailed Review: PR #388 — update_message() for guardrail redaction
Overall this is a well-motivated change that addresses a real gap (guardrail-redacted messages being persisted unredacted). The create-then-delete approach is reasonable given immutable events. However, there are several edge cases and test gaps that should be addressed before merge.
Summary of Findings
P1 — Potential Data Loss (2 issues)
- When
batch_size > 1and a persisted message is updated,create_message()returns{}(buffered). The old event is then deleted immediately. If the process crashes before the buffer flushes, the message is lost with no recovery. Rollback is also impossible sincenew_event_idwould beNone. - If
create_message()returnsNone(e.g. converter produces empty payload from redacted content), the code proceeds to delete the old event with no replacement created.
P2 — Race Condition (1 issue)
3. _update_buffered_message holds _message_lock while replacing the entry, but _flush_messages_only copies the buffer before clearing it. A flush concurrent with an update could send the old (unredacted) content, then the updated buffer entry is cleared — silently losing the redaction.
P2 — Fragile Buffer Matching (1 issue)
4. _update_buffered_message matches by role only. Multiple buffered messages with the same role, or two blob messages (None == None), could cause the wrong message to be replaced.
P3 — Code Quality (3 issues)
5. read_message docstring says it's "primarily used internally by update_message" but update_message never calls read_message.
6. getattr(self, "_latest_agent_message", None) is unnecessary — attribute is always initialized in parent __init__.
7. PR description says "same pattern as update_agent()" but update_agent does NOT delete the old event.
Test Gaps (see inline comments on test file)
- No test for
_latest_agent_messagebeing updated after successful update - No test for rollback success path (only double-failure is accidentally tested)
test_update_buffered_messagechecks buffer count but never verifies content actually changedtest_update_message_create_failsdoesn't assertdelete_eventwas NOT called- No test for
PersistenceMode.NONEorbatch_size > 1with persisted messages - No test for buffer role mismatch or multi-message-same-role scenarios
…ition, improve tests - Guard against create_message returning None/empty eventId before deleting old event (P1) - Fix flush race condition by atomically clearing buffer with snapshot; restore on failure (P2) - Fix read_message docstring, simplify rollback logic (P3) - Add test for _latest_agent_message update after replacement - Add test for rollback success path (not just double-failure) - Verify created event content and assert no delete on create failure - Verify buffered message content actually changes after update
Summary
update_message()inAgentCoreMemorySessionManagerso that Strands' built-in guardrail redaction (redact_latest_message()) works out of the boxupdate_agent())batch_size > 1case by replacing messages in the send buffer before they are flushedContext
When using Bedrock Guardrails with AgentCore Memory as the Strands session store,
update_message()was a no-op. This meant guardrail-blocked user messages were persisted unredacted, creating a permanent dead-end conversation — on subsequent turns or reconnect, the guardrail would block again on the persisted offending message.Test plan
test_update_message— verifies new event created + old event deleted for persisted messagestest_update_message_wrong_session— session ID mismatch raisesSessionExceptiontest_update_message_no_message_id— graceful skip when message has no event IDtest_update_message_create_fails— raisesSessionExceptionon create failuretest_update_message_delete_fails— raisesSessionExceptionon delete failuretest_update_buffered_message— in-buffer replacement whenbatch_size > 1