diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index be5632dc..242b2fdc 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/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,6 +601,9 @@ 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, @@ -601,13 +611,14 @@ def create_message( 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") diff --git a/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py b/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py index 11f505a0..3f48e3a3 100644 --- a/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py +++ b/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py @@ -599,6 +599,25 @@ def test_update_message_updates_latest_agent_message(self, session_manager, mock assert session_manager._latest_agent_message["test-agent-123"].message_id == "new_event_456" + def test_redact_persisted_blob_message(self, session_manager, mock_memory_client): + """A blob event must retain its event ID so redaction can replace it.""" + session_manager._latest_agent_message = {} + session_manager.session_repository = session_manager + mock_memory_client.gmdp_client.create_event.return_value = {"event": {"eventId": "blob_event_123"}} + mock_memory_client.create_event.return_value = {"eventId": "redacted_event_456"} + agent = Mock(agent_id="test-agent-123") + message = {"role": "user", "content": [{"text": "x" * (CONVERSATIONAL_MAX_SIZE + 100)}]} + + session_manager.append_message(message, agent) + assert session_manager._latest_agent_message[agent.agent_id].message_id == "blob_event_123" + + message["content"] = [{"text": "[User input redacted.]"}] + session_manager.redact_latest_message(message, agent) + + assert "[User input redacted.]" in mock_memory_client.create_event.call_args.kwargs["messages"][0][0] + assert mock_memory_client.gmdp_client.delete_event.call_args.kwargs["eventId"] == "blob_event_123" + assert session_manager._latest_agent_message[agent.agent_id].message_id == "redacted_event_456" + def test_update_message_wrong_session(self, session_manager): """Test updating a message with wrong session ID.""" message = SessionMessage(message={"role": "user", "content": [{"text": "Hello"}]}, message_id=1) @@ -606,16 +625,18 @@ def test_update_message_wrong_session(self, session_manager): with pytest.raises(SessionException, match="Session ID mismatch"): session_manager.update_message("wrong-session-id", "test-agent-123", message) - def test_update_message_no_message_id(self, session_manager): - """Test updating a message with no message_id (not yet persisted) skips gracefully.""" + def test_update_message_no_message_id(self, session_manager, caplog): + """A missing event ID and buffer entry must be visible when redaction is skipped.""" message = SessionMessage( message={"role": "user", "content": [{"text": "redacted"}]}, message_id=None, created_at="2024-01-01T12:00:00Z", ) - # Should not raise - just skips since message isn't persisted and buffer is empty - session_manager.update_message("test-session-456", "test-agent-123", message) + with caplog.at_level(logging.WARNING): + session_manager.update_message("test-session-456", "test-agent-123", message) + + assert "Message has no event ID and was not found in buffer" in caplog.text def test_update_message_create_fails(self, session_manager, mock_memory_client): """Test update_message raises SessionException when create fails and does not delete.""" @@ -1572,6 +1593,198 @@ def test_update_buffered_message(self, batching_session_manager, mock_memory_cli mock_memory_client.create_event.assert_not_called() mock_memory_client.gmdp_client.delete_event.assert_not_called() + def test_redact_buffered_blob_question(self, batching_session_manager, mock_memory_client): + """A redacted conversational message must replace its buffered blob.""" + batching_session_manager._latest_agent_message = {} + batching_session_manager.session_repository = batching_session_manager + agent = Mock(agent_id="test-agent-123") + question = {"role": "user", "content": [{"text": "x" * (CONVERSATIONAL_MAX_SIZE + 100)}]} + + batching_session_manager.append_message(question, agent) + assert batching_session_manager._message_buffer[0].is_blob + + question["content"] = [{"text": "[User input redacted.]"}] + batching_session_manager.redact_latest_message(question, agent) + batching_session_manager._flush_messages_only() + + payload = mock_memory_client.gmdp_client.create_event.call_args.kwargs["payload"] + assert len(payload) == 1 + stored = json.loads(payload[0]["conversational"]["content"]["text"]) + assert stored["message"]["content"] == [{"text": "[User input redacted.]"}] + + def test_redact_buffered_blob_tool_result(self, batching_session_manager, mock_memory_client): + """Redacting a blob tool result must leave the earlier question intact.""" + batching_session_manager._latest_agent_message = {} + batching_session_manager.session_repository = batching_session_manager + agent = Mock(agent_id="test-agent-123") + question = {"role": "user", "content": [{"text": "Summarise the report."}]} + tool_use = { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "get_report", "input": {}}}], + } + tool_result = { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "t1", "content": [{"text": "x" * (CONVERSATIONAL_MAX_SIZE + 100)}]}} + ], + } + for message in (question, tool_use, tool_result): + batching_session_manager.append_message(message, agent) + assert batching_session_manager._message_buffer[-1].is_blob + + tool_result["content"][0]["toolResult"]["content"] = [{"text": "[User input redacted.]"}] + batching_session_manager.redact_latest_message(tool_result, agent) + batching_session_manager._flush_messages_only() + + payload = mock_memory_client.gmdp_client.create_event.call_args.kwargs["payload"] + stored = [json.loads(item["conversational"]["content"]["text"])["message"] for item in payload] + assert stored[0]["content"] == question["content"] + assert stored[1]["content"] == tool_use["content"] + assert stored[2]["content"] == tool_result["content"] + + @pytest.mark.parametrize("flush_trigger,batch_size", [("batch_size", 2), ("interval", 10)]) + def test_latest_redaction_survives_early_flush(self, flush_trigger, batch_size, mock_memory_client): + """Early flushes may send older messages but must retain the redaction target.""" + config = AgentCoreMemoryConfig( + memory_id="test-memory-123", + session_id="test-session-456", + actor_id="test-actor-789", + batch_size=batch_size, + ) + manager = _create_session_manager(config, mock_memory_client) + manager._latest_agent_message = {} + manager.session_repository = manager + agent = Mock(agent_id="test-agent-123") + question = {"role": "user", "content": [{"text": "Summarise the report."}]} + tool_use = { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "get_report", "input": {}}}], + } + tool_result = { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "t1", "content": [{"text": "x" * (CONVERSATIONAL_MAX_SIZE + 100)}]}} + ], + } + + for message in (question, tool_use, tool_result): + manager.append_message(message, agent) + if flush_trigger == "interval": + manager._interval_flush_callback() + + assert mock_memory_client.gmdp_client.create_event.call_count == 1 + early_payload = mock_memory_client.gmdp_client.create_event.call_args.kwargs["payload"] + assert len(early_payload) == 2 + assert manager.pending_message_count() == 1 + assert manager._message_buffer[0].protected + + tool_result["content"][0]["toolResult"]["content"] = [{"text": "[User input redacted.]"}] + manager.redact_latest_message(tool_result, agent) + manager.close() + + final_payload = mock_memory_client.gmdp_client.create_event.call_args.kwargs["payload"] + assert len(final_payload) == 1 + stored = json.loads(final_payload[0]["conversational"]["content"]["text"]) + assert stored["message"]["content"] == tool_result["content"] + assert manager.pending_message_count() == 0 + + def test_agent_state_auto_flush_defers_latest_message(self, mock_memory_client): + """An agent-state threshold must not flush a message awaiting redaction.""" + config = AgentCoreMemoryConfig( + memory_id="test-memory-123", + session_id="test-session-456", + actor_id="test-actor-789", + batch_size=2, + ) + manager = _create_session_manager(config, mock_memory_client) + manager._latest_agent_message = {} + manager.session_repository = manager + agent = Mock(agent_id="test-agent-123") + question = {"role": "user", "content": [{"text": "x" * (CONVERSATIONAL_MAX_SIZE + 100)}]} + manager.append_message(question, agent) + + for i in range(2): + state = SessionAgent(agent_id=agent.agent_id, state={"turn": i}, conversation_manager_state={}) + manager.create_agent(config.session_id, state) + + assert manager.pending_message_count() == 1 + assert manager.pending_agent_state_count() == 2 + mock_memory_client.gmdp_client.create_event.assert_not_called() + + question["content"] = [{"text": "[User input redacted.]"}] + manager.redact_latest_message(question, agent) + manager.close() + + first_payload = mock_memory_client.gmdp_client.create_event.call_args_list[0].kwargs["payload"] + stored = json.loads(first_payload[0]["conversational"]["content"]["text"]) + assert stored["message"]["content"] == question["content"] + assert manager.pending_agent_state_count() == 0 + + def test_deferred_agent_states_respect_event_payload_limit(self, mock_memory_client): + """A long invocation sends deferred states in events of at most 100 items.""" + config = AgentCoreMemoryConfig( + memory_id="test-memory-123", + session_id="test-session-456", + actor_id="test-actor-789", + batch_size=2, + ) + manager = _create_session_manager(config, mock_memory_client) + manager._latest_agent_message = {} + manager.session_repository = manager + agent = Mock(agent_id="test-agent-123") + question = {"role": "user", "content": [{"text": "x" * (CONVERSATIONAL_MAX_SIZE + 100)}]} + manager.append_message(question, agent) + + for turn in range(101): + manager.create_agent( + config.session_id, + SessionAgent(agent_id=agent.agent_id, state={"turn": turn}, conversation_manager_state={}), + ) + + mock_memory_client.gmdp_client.create_event.assert_not_called() + question["content"] = [{"text": "[User input redacted.]"}] + manager.redact_latest_message(question, agent) + manager.close() + + calls = mock_memory_client.gmdp_client.create_event.call_args_list + assert [len(call.kwargs["payload"]) for call in calls] == [1, 100, 1] + stored_question = json.loads(calls[0].kwargs["payload"][0]["conversational"]["content"]["text"]) + assert stored_question["message"]["content"] == question["content"] + stored_states = [ + json.loads(item["blob"])["state"]["turn"] for call in calls[1:] for item in call.kwargs["payload"] + ] + assert stored_states == list(range(101)) + assert manager.pending_agent_state_count() == 0 + + # read_agent fetches the newest event and its last payload item. + mock_memory_client.list_events.return_value = [{"payload": calls[-1].kwargs["payload"]}] + restored_agent = manager.read_agent(config.session_id, agent.agent_id) + assert restored_agent.state == {"turn": 100} + + def test_failed_early_flush_preserves_latest_redaction_target(self, batching_session_manager, mock_memory_client): + """Restoring a failed flush must keep older messages before the protected latest one.""" + batching_session_manager._latest_agent_message = {} + batching_session_manager.session_repository = batching_session_manager + agent = Mock(agent_id="test-agent-123") + earlier = {"role": "user", "content": [{"text": "Earlier question"}]} + latest = {"role": "user", "content": [{"text": "x" * (CONVERSATIONAL_MAX_SIZE + 100)}]} + batching_session_manager.append_message(earlier, agent) + batching_session_manager.append_message(latest, agent) + + mock_memory_client.gmdp_client.create_event.side_effect = Exception("temporary failure") + batching_session_manager._interval_flush_callback() + assert batching_session_manager.pending_message_count() == 2 + assert [buf.protected for buf in batching_session_manager._message_buffer] == [False, True] + + mock_memory_client.gmdp_client.create_event.side_effect = None + latest["content"] = [{"text": "[User input redacted.]"}] + batching_session_manager.redact_latest_message(latest, agent) + batching_session_manager.close() + + payload = mock_memory_client.gmdp_client.create_event.call_args.kwargs["payload"] + stored = [json.loads(item["conversational"]["content"]["text"])["message"] for item in payload] + assert [message["content"] for message in stored] == [earlier["content"], latest["content"]] + def test_buffer_auto_flushes_at_batch_size(self, batching_session_manager, mock_memory_client): """Test buffer automatically flushes when reaching batch_size.""" mock_memory_client.gmdp_client.create_event.return_value = {"eventId": "event_123"} @@ -3935,6 +4148,30 @@ def fail_after_concurrent_append(**kwargs): assert [a.state for _, a in batching_session_manager._agent_state_buffer] == [{"v": 1}, {"v": 2}] + def test_flush_agent_states_failure_retries_only_unsent_chunk(self, batching_session_manager, mock_memory_client): + """A failed later chunk must not replay states already persisted.""" + with batching_session_manager._agent_state_lock: + batching_session_manager._agent_state_buffer = [ + ( + "test-session-456", + SessionAgent(agent_id="agent-1", state={"turn": turn}, conversation_manager_state={}), + ) + for turn in range(101) + ] + mock_memory_client.gmdp_client.create_event.side_effect = [ + {"eventId": "first-chunk"}, + Exception("second chunk failed"), + ] + + with pytest.raises(SessionException, match="Failed to flush agent states"): + batching_session_manager._flush_agent_states_only() + + assert [agent.state["turn"] for _, agent in batching_session_manager._agent_state_buffer] == [100] + mock_memory_client.gmdp_client.create_event.side_effect = None + batching_session_manager._flush_agent_states_only() + assert len(mock_memory_client.gmdp_client.create_event.call_args.kwargs["payload"]) == 1 + assert batching_session_manager.pending_agent_state_count() == 0 + def test_agent_state_auto_flush_persists_buffered_messages_first( self, batching_session_manager, mock_memory_client ):