fix(memory): preserve guardrail redactions

src/bedrock_agentcore/memory/integrations/strands/session_manager.py

@@ -58,19 +58,24 @@
 
 # Maximum metadata key-value pairs per event (API limit)
 MAX_METADATA_KEYS = 15
+MAX_EVENT_PAYLOAD_ITEMS = 100
 
 # Reserved internal metadata keys that users cannot override
 RESERVED_METADATA_KEYS = frozenset({STATE_TYPE_KEY, AGENT_ID_KEY})
 
 
 class BufferedMessage(NamedTuple):
-    """A pre-processed message waiting to be flushed to AgentCore Memory."""
+    """A pre-processed message waiting to be flushed to AgentCore Memory.
+
+    A protected message is Strands' latest message, which it may still redact.
+    """
 
     session_id: str
     messages: list[tuple[str, str]]
     is_blob: bool
     timestamp: datetime
     metadata: Optional[Dict[str, MetadataValue]] = None
+    protected: bool = False
 
 
 class StateType(Enum):
@@ -403,7 +408,7 @@ def create_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: A
             # Flush outside the lock to prevent deadlock. Flush messages too, not just agent states: agent state
             # depends on the messages (removed_message_count indexes into them), so they must be persisted together.
             if should_flush:
-                self._flush_messages()
+                self._flush_messages(include_protected=False)
 
             logger.info(
                 "Buffered agent creation: %s in session: %s",
@@ -540,7 +545,9 @@ def create_message(
     ) -> Optional[dict[str, Any]]:
         """Create a new message in AgentCore Memory.
 
-        If batch_size > 1, the message is buffered and sent when the buffer reaches batch_size.
+        If batch_size > 1, the message is buffered and sendable messages are flushed
+        when they reach batch_size. Strands' latest message stays buffered until
+        invocation completion so a guardrail can still redact it.
         Use _flush_messages() or close() to send any remaining buffered messages.
 
         Args:
@@ -594,20 +601,24 @@ def create_message(
             # Buffer the pre-processed message
             should_flush = False
             with self._message_lock:
+                if kwargs.get("_protect_latest"):
+                    # Strands can still redact its latest message. Older messages can be flushed.
+                    self._message_buffer = [buf._replace(protected=False) for buf in self._message_buffer]
                 self._message_buffer.append(
                     BufferedMessage(
                         session_id=session_id,
                         messages=messages,
                         is_blob=is_blob,
                         timestamp=monotonic_timestamp,
                         metadata=merged_metadata,
+                        protected=bool(kwargs.get("_protect_latest")),
                     )
                 )
-                should_flush = len(self._message_buffer) >= self.config.batch_size
+                should_flush = sum(not buf.protected for buf in self._message_buffer) >= self.config.batch_size
 
             # Flush only messages outside the lock to prevent deadlock
             if should_flush:
-                self._flush_messages_only()
+                self._flush_messages_only(include_protected=False)
 
             return {}  # No eventId yet
 
@@ -632,7 +643,7 @@ def create_message(
                 }
                 if merged_metadata:
                     create_event_kwargs["metadata"] = merged_metadata
-                event = self.memory_client.gmdp_client.create_event(**create_event_kwargs)
+                event = self.memory_client.gmdp_client.create_event(**create_event_kwargs)["event"]
             logger.debug("Created event: %s for message: %s", event.get("eventId"), session_message.message_id)
             return event
         except Exception as e:
@@ -689,7 +700,7 @@ def update_message(self, session_id: str, agent_id: str, session_message: Sessio
             if self._update_buffered_message(session_message):
                 logger.debug("Updated buffered message (not yet persisted to AgentCore Memory)")
                 return
-            logger.debug("Message has no event ID and was not found in buffer - skipping update")
+            logger.warning("Message has no event ID and was not found in buffer - skipping update")
             return
 
         # Create a new event with the updated message content
@@ -836,12 +847,17 @@ def _filter_restored_tool_context(self, messages: list[SessionMessage]) -> list[
     def append_message(self, message: Message, agent: "LocalAgent", **kwargs: Any) -> None:
         """Append a message to the agent's session using AgentCore's eventId as message_id.
 
+        With batching, hold this message through early flushes because Strands may
+        still request redaction of its latest message.
+
         Args:
             message: Message to add to the agent in the session
             agent: Agent to append the message to
             **kwargs: Additional keyword arguments for future extensibility.
         """
-        created_message = self.create_message(self.session_id, agent.agent_id, SessionMessage.from_message(message, 0))
+        created_message = self.create_message(
+            self.session_id, agent.agent_id, SessionMessage.from_message(message, 0), _protect_latest=True
+        )
         if created_message is None:
             return
         session_message = SessionMessage.from_message(message, created_message.get("eventId"))
@@ -1034,25 +1050,27 @@ def _update_buffered_message(self, session_message: SessionMessage) -> bool:
         is_blob = self.converter.exceeds_conversational_limit(updated_messages[0])
 
         with self._message_lock:
-            # Search from the end (most recent) to find the message to update
-            for i in range(len(self._message_buffer) - 1, -1, -1):
+            # Strands redacts its latest message. Prefer the protected entry when present;
+            # direct create_message() callers can still use the role-based fallback.
+            protected_indices = [i for i, buf in enumerate(self._message_buffer) if buf.protected]
+            indices = reversed(protected_indices) if protected_indices else range(len(self._message_buffer) - 1, -1, -1)
+            for i in indices:
                 buf = self._message_buffer[i]
                 if buf.session_id == self.config.session_id and buf.messages:
                     # Match by role - the most recent message with the same role
-                    existing_role = buf.messages[0][1] if not buf.is_blob else None
-                    new_role = updated_messages[0][1] if not is_blob else None
-                    if existing_role == new_role:
+                    if buf.messages[0][1] == updated_messages[0][1]:
                         self._message_buffer[i] = BufferedMessage(
                             session_id=buf.session_id,
                             messages=updated_messages,
                             is_blob=is_blob,
                             timestamp=buf.timestamp,
                             metadata=buf.metadata,
+                            protected=buf.protected,
                         )
                         return True
         return False
 
-    def _flush_messages_only(self) -> list[dict[str, Any]]:
+    def _flush_messages_only(self, include_protected: bool = True) -> list[dict[str, Any]]:
         """Flush only buffered messages to AgentCore Memory.
 
         Call this method to send any remaining buffered messages when batch_size > 1.
@@ -1061,6 +1079,8 @@ def _flush_messages_only(self) -> list[dict[str, Any]]:
         session are combined into a single create_event() call to reduce API calls.
         Messages that exceed the conversational payload limit are sent as blob events individually
         as they require a different API path.
+        Early flushes pass include_protected=False to keep Strands' latest message
+        available for redaction.
 
         Returns:
             list[dict[str, Any]]: List of created event responses from AgentCore Memory.
@@ -1072,8 +1092,12 @@ def _flush_messages_only(self) -> list[dict[str, Any]]:
             return []
 
         with self._message_lock:
-            messages_to_send = list(self._message_buffer)
-            self._message_buffer.clear()
+            if include_protected:
+                messages_to_send = list(self._message_buffer)
+                self._message_buffer.clear()
+            else:
+                messages_to_send = [buf for buf in self._message_buffer if not buf.protected]
+                self._message_buffer = [buf for buf in self._message_buffer if buf.protected]
 
         if not messages_to_send:
             return []
@@ -1129,7 +1153,7 @@ def _flush_messages_only(self) -> list[dict[str, Any]]:
         except Exception as e:
             # Restore messages to buffer so they aren't lost
             with self._message_lock:
-                self._message_buffer.extend(messages_to_send)
+                self._message_buffer[:0] = messages_to_send
             logger.error("Failed to flush messages to AgentCore Memory: %s", e)
             raise SessionException(f"Failed to flush messages: {e}") from e
 
@@ -1140,8 +1164,9 @@ def _flush_agent_states_only(self) -> list[dict[str, Any]]:
         """Flush only buffered agent states to AgentCore Memory.
 
         Call this method to send any remaining agent state when batch_size > 1.
-        Agent states are grouped by agent_id and sent as separate events so that
-        each event carries the correct agentId metadata for read_agent() lookups.
+        Agent states are grouped by agent_id and sent in events of at most 100
+        payload items so each event carries the correct agentId metadata for
+        read_agent() lookups.
 
         Returns:
             list[dict[str, Any]]: List of created event responses from AgentCore Memory.
@@ -1160,46 +1185,57 @@ def _flush_agent_states_only(self) -> list[dict[str, Any]]:
             return []
 
         results = []
+        sent_indices: set[int] = set()
         try:
             # Group agent states by agent_id
-            agent_groups: dict[str, list[dict]] = {}
-            for _session_id, session_agent in agent_states_to_send:
+            agent_groups: dict[str, list[tuple[int, SessionAgent]]] = {}
+            for index, (_session_id, session_agent) in enumerate(agent_states_to_send):
                 agent_id = session_agent.agent_id
                 if agent_id not in agent_groups:
                     agent_groups[agent_id] = []
-                agent_groups[agent_id].append({"blob": json.dumps(session_agent.to_dict())})
-
-            # Send one event per agent_id with correct metadata
-            for agent_id, payloads in agent_groups.items():
-                event = self.memory_client.gmdp_client.create_event(
-                    memoryId=self.config.memory_id,
-                    actorId=self.config.actor_id,
-                    sessionId=self.config.session_id,
-                    payload=payloads,
-                    eventTimestamp=self._get_monotonic_timestamp(),
-                    metadata={
-                        STATE_TYPE_KEY: {"stringValue": StateType.AGENT.value},
-                        AGENT_ID_KEY: {"stringValue": agent_id},
-                    },
-                )
-                results.append(event)
-                logger.debug("Flushed %d agent states for agent %s: %s", len(payloads), agent_id, event.get("eventId"))
+                agent_groups[agent_id].append((index, session_agent))
+
+            # Send each agent's states in order, respecting the CreateEvent payload limit.
+            for agent_id, states in agent_groups.items():
+                for start in range(0, len(states), MAX_EVENT_PAYLOAD_ITEMS):
+                    chunk = states[start : start + MAX_EVENT_PAYLOAD_ITEMS]
+                    payloads = [{"blob": json.dumps(agent.to_dict())} for _, agent in chunk]
+                    event = self.memory_client.gmdp_client.create_event(
+                        memoryId=self.config.memory_id,
+                        actorId=self.config.actor_id,
+                        sessionId=self.config.session_id,
+                        payload=payloads,
+                        eventTimestamp=self._get_monotonic_timestamp(),
+                        metadata={
+                            STATE_TYPE_KEY: {"stringValue": StateType.AGENT.value},
+                            AGENT_ID_KEY: {"stringValue": agent_id},
+                        },
+                    )
+                    results.append(event)
+                    sent_indices.update(index for index, _ in chunk)
+                    logger.debug(
+                        "Flushed %d agent states for agent %s: %s", len(payloads), agent_id, event.get("eventId")
+                    )
 
         except Exception as e:
-            # Prepend so they aren't lost: states buffered during the send are newer, and read_agent restores the last
+            # Prepend only unsent states. States buffered during the send are newer,
+            # and read_agent restores the last state.
             with self._agent_state_lock:
-                self._agent_state_buffer[:0] = agent_states_to_send
+                self._agent_state_buffer[:0] = [
+                    state for index, state in enumerate(agent_states_to_send) if index not in sent_indices
+                ]
             logger.error("Failed to flush agent states to AgentCore Memory: %s", e)
             raise SessionException(f"Failed to flush agent states: {e}") from e
 
         logger.info("Flushed %d agent state events to AgentCore Memory", len(results))
         return results
 
-    def _flush_messages(self) -> list[dict[str, Any]]:
+    def _flush_messages(self, include_protected: bool = True) -> list[dict[str, Any]]:
         """Flush all buffered messages and agent state to AgentCore Memory.
 
         Call this method to send any remaining buffered messages and agent state messages.
         This is automatically called when the session is complete (via close() or context manager).
+        Early flushes retain Strands' latest message and defer agent state until it is sent.
 
         Returns:
             list[dict[str, Any]]: List of created event responses from AgentCore Memory.
@@ -1208,8 +1244,11 @@ def _flush_messages(self) -> list[dict[str, Any]]:
             SessionException: If any message or agent state creation fails.
         """
         results = []
-        results.extend(self._flush_messages_only())
-        results.extend(self._flush_agent_states_only())
+        results.extend(self._flush_messages_only(include_protected=include_protected))
+        with self._message_lock:
+            has_protected_message = any(buf.protected for buf in self._message_buffer)
+        if not has_protected_message:
+            results.extend(self._flush_agent_states_only())
         return results
 
     def pending_message_count(self) -> int:
@@ -1310,7 +1349,7 @@ def _interval_flush_callback(self) -> None:
                     pending_messages,
                     pending_agent_states,
                 )
-                self._flush_messages()
+                self._flush_messages(include_protected=False)
             else:
                 logger.debug("Interval flush skipped: buffers are empty")