Skip to content

fix: implement update_message() for guardrail redaction support#388

Merged
jariy17 merged 7 commits intomainfrom
fix/update-message-guardrail-redaction
Apr 21, 2026
Merged

fix: implement update_message() for guardrail redaction support#388
jariy17 merged 7 commits intomainfrom
fix/update-message-guardrail-redaction

Conversation

@notgitika
Copy link
Copy Markdown
Contributor

Summary

  • Implement update_message() in AgentCoreMemorySessionManager so that Strands' built-in guardrail redaction (redact_latest_message()) works out of the box
  • Since AgentCore Memory events are immutable, updates are performed via create-then-delete (same pattern as update_agent())
  • Handle the batch_size > 1 case by replacing messages in the send buffer before they are flushed

Context

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 messages
  • test_update_message_wrong_session — session ID mismatch raises SessionException
  • test_update_message_no_message_id — graceful skip when message has no event ID
  • test_update_message_create_fails — raises SessionException on create failure
  • test_update_message_delete_fails — raises SessionException on delete failure
  • test_update_buffered_message — in-buffer replacement when batch_size > 1
  • Full test suite passes (140/140)

@notgitika notgitika requested a review from a team April 2, 2026 05:56
@github-actions
Copy link
Copy Markdown
Contributor

github-actions Bot commented Apr 2, 2026

✅ No Breaking Changes Detected

No public API breaking changes found in this PR.

Comment thread src/bedrock_agentcore/memory/integrations/strands/session_manager.py Outdated
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.
@notgitika notgitika force-pushed the fix/update-message-guardrail-redaction branch from 828699f to a6df5f4 Compare April 11, 2026 03:08
Copy link
Copy Markdown
Contributor

@jariy17 jariy17 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Detailed Review: PR #388update_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)

  1. When batch_size > 1 and 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 since new_event_id would be None.
  2. If create_message() returns None (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_message being updated after successful update
  • No test for rollback success path (only double-failure is accidentally tested)
  • test_update_buffered_message checks buffer count but never verifies content actually changed
  • test_update_message_create_fails doesn't assert delete_event was NOT called
  • No test for PersistenceMode.NONE or batch_size > 1 with 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
@jariy17 jariy17 merged commit 9ba4512 into main Apr 21, 2026
32 checks passed
@jariy17 jariy17 deleted the fix/update-message-guardrail-redaction branch April 21, 2026 16:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants