fix: implement update_message() for guardrail redaction support

[notgitika]

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)

[github-actions]

✅ No Breaking Changes Detected

No public API breaking changes found in this PR.

[jariy17]

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)

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

[jariy17]

P1 — Potential data loss when create_message returns None or {}

Two scenarios:

1. batch_size > 1: create_message() returns {} (buffered, no eventId). The code then proceeds to delete_event() on the old message. The new content is sitting in an unflushed buffer. If the process crashes before flush, the message is permanently lost. Rollback is also impossible since new_event.get("eventId") returns None.

2. Empty converter payload: If create_message() returns None (converter produces empty payload), the code proceeds to delete the old event with no replacement.

Suggestion: Guard against both cases:

new_event = self.create_message(session_id, agent_id, updated_message)
if not new_event or not new_event.get("eventId"):
    logger.warning("create_message did not return an eventId — skipping delete of old event %s", old_message_id)
    return

[jariy17]

P2 — Race condition between buffer update and flush

_flush_messages_only copies the buffer (list(self._message_buffer)) then later clears it. If _update_buffered_message runs between the copy and the clear:

• The flush sends the old (unredacted) content from the copy
• The buffer now contains the new (redacted) content
• After flush succeeds, the buffer is cleared — redaction silently lost

The fix would be to hold _message_lock around the entire flush-copy-and-clear operation, or to check if the flush path also needs to coordinate with updates.

[jariy17]

Test gap — content of created event not verified

mock_memory_client.create_event.assert_called_once() verifies the call happened but never inspects what content was passed. The replacement event could contain garbage and this test would still pass.

Suggest adding:

create_kwargs = mock_memory_client.create_event.call_args.kwargs
# or inspect the payload to verify it contains "redacted"

[jariy17]

Test gap — _latest_agent_message update not tested

The source code explicitly updates self._latest_agent_message[agent_id] when the old message_id matches. This is critical for guardrail correctness — subsequent turns must reference the new eventId.

Suggest adding a test that:

1. Pre-populates session_manager._latest_agent_message["test-agent-123"] with a SessionMessage using message_id="old_event_123"
2. Calls update_message
3. Asserts session_manager._latest_agent_message["test-agent-123"].message_id == "new_event_456"

[jariy17]

Test gap — rollback not verified

Setting delete_event.side_effect = Exception(...) makes ALL calls fail — including the rollback attempt. This accidentally tests only the double-failure path.

Missing tests:

1. Rollback succeeds: Use side_effect=[Exception("Delete failed"), None] so the first call (delete old) fails but the second (rollback new) succeeds. Assert delete_event called twice.
2. Current test: Should assert delete_event was called at least twice (once for old, once for rollback) and that the rollback targeted new_event_456.

[jariy17]

Test gap — create_fails should verify no delete attempted

When create_event raises, the code should bail out before calling delete_event. Add:

mock_memory_client.gmdp_client.delete_event.assert_not_called()