fix: copy and clear messages within same lock acquisition

Author: nborges-aws

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

@@ -11,6 +11,11 @@
 
 import boto3
 from botocore.config import Config as BotocoreConfig
+from strands.experimental.hooks.events import (
+    BidiAfterInvocationEvent,
+    BidiAgentInitializedEvent,
+    BidiMessageAddedEvent,
+)
 from strands.experimental.hooks.multiagent.events import (
     AfterMultiAgentInvocationEvent,
     AfterNodeCallEvent,
@@ -942,50 +947,56 @@ def register_hooks(self, registry: HookRegistry, **kwargs) -> None:
 
         # Async mode: register async callbacks that offload the existing sync
         # methods to a worker thread via asyncio.to_thread. AgentInitializedEvent
-        # must stay sync (Strands disallows async callbacks on this event; see
-        # strands/hooks/registry.py:174).
+        # and BidiAgentInitializedEvent must stay sync (Strands disallows async
+        # callbacks for AgentInitializedEvent — see strands/hooks/registry.py:227).
         logger.warning(
             "AgentCoreMemorySessionManager async_mode=True: the agent must be invoked "
             "via the async path (e.g. agent.stream_async(...) or agent.invoke_async(...)). "
             "Sync invocation will raise RuntimeError from Strands' hook registry."
         )
 
+        def _offload(method, *event_args):
+            """Build an async callback that offloads `method(*[a(event) for a in event_args])` to a thread.
+
+            Each entry in `event_args` is a callable that extracts an argument from the event;
+            pass none for a zero-arg method.
+            """
+
+            async def _callback(event):
+                await asyncio.to_thread(method, *(extract(event) for extract in event_args))
+
+            return _callback
+
         registry.add_callback(AgentInitializedEvent, lambda event: self.initialize(event.agent))
 
         async def _on_message_added_persist(event: MessageAddedEvent) -> None:
             await asyncio.to_thread(self.append_message, event.message, event.agent)
             await asyncio.to_thread(self.sync_agent, event.agent)
 
-        async def _on_message_added_retrieve(event: MessageAddedEvent) -> None:
-            await asyncio.to_thread(self.retrieve_customer_context, event)
-
-        async def _on_after_invocation_sync(event: AfterInvocationEvent) -> None:
-            await asyncio.to_thread(self.sync_agent, event.agent)
-
         registry.add_callback(MessageAddedEvent, _on_message_added_persist)
-        registry.add_callback(AfterInvocationEvent, _on_after_invocation_sync)
-        registry.add_callback(MessageAddedEvent, _on_message_added_retrieve)
+        registry.add_callback(AfterInvocationEvent, _offload(self.sync_agent, lambda e: e.agent))
+        registry.add_callback(MessageAddedEvent, _offload(self.retrieve_customer_context, lambda e: e))
 
         if self.config.batch_size > 1:
-
-            async def _on_after_invocation_flush(event: AfterInvocationEvent) -> None:
-                await asyncio.to_thread(self._flush_messages)
-
-            registry.add_callback(AfterInvocationEvent, _on_after_invocation_flush)
+            registry.add_callback(AfterInvocationEvent, _offload(self._flush_messages))
 
         # Register multi-agent callbacks so async-mode parity matches sync-mode
-        async def _on_multi_agent_initialized(event: MultiAgentInitializedEvent) -> None:
-            await asyncio.to_thread(self.initialize_multi_agent, event.source)
+        registry.add_callback(MultiAgentInitializedEvent, _offload(self.initialize_multi_agent, lambda e: e.source))
+        registry.add_callback(AfterNodeCallEvent, _offload(self.sync_multi_agent, lambda e: e.source))
+        registry.add_callback(AfterMultiAgentInvocationEvent, _offload(self.sync_multi_agent, lambda e: e.source))
 
-        async def _on_after_node_call(event: AfterNodeCallEvent) -> None:
-            await asyncio.to_thread(self.sync_multi_agent, event.source)
+        # Register BidiAgent callbacks so async-mode parity matches sync-mode.
+        # BidiAgentInitializedEvent dispatches through invoke_callbacks (sync),
+        # so its callback must stay sync; the other two dispatch through
+        # invoke_callbacks_async, so async wrappers are safe.
+        registry.add_callback(BidiAgentInitializedEvent, lambda event: self.initialize_bidi_agent(event.agent))
 
-        async def _on_after_multi_agent_invocation(event: AfterMultiAgentInvocationEvent) -> None:
-            await asyncio.to_thread(self.sync_multi_agent, event.source)
+        async def _on_bidi_message_added(event: BidiMessageAddedEvent) -> None:
+            await asyncio.to_thread(self.append_bidi_message, event.message, event.agent)
+            await asyncio.to_thread(self.sync_bidi_agent, event.agent)
 
-        registry.add_callback(MultiAgentInitializedEvent, _on_multi_agent_initialized)
-        registry.add_callback(AfterNodeCallEvent, _on_after_node_call)
-        registry.add_callback(AfterMultiAgentInvocationEvent, _on_after_multi_agent_invocation)
+        registry.add_callback(BidiMessageAddedEvent, _on_bidi_message_added)
+        registry.add_callback(BidiAfterInvocationEvent, _offload(self.sync_bidi_agent, lambda e: e.agent))
 
     @override
     def initialize(self, agent: "Agent", **kwargs: Any) -> None:
@@ -1141,6 +1152,7 @@ def _flush_agent_states_only(self) -> list[dict[str, Any]]:
 
         with self._agent_state_lock:
             agent_states_to_send = list(self._agent_state_buffer)
+            self._agent_state_buffer.clear()
 
         if not agent_states_to_send:
             return []
@@ -1171,11 +1183,10 @@ def _flush_agent_states_only(self) -> list[dict[str, Any]]:
                 results.append(event)
                 logger.debug("Flushed %d agent states for agent %s: %s", len(payloads), agent_id, event.get("eventId"))
 
-            # Clear agent state buffer only after ALL events succeed
-            with self._agent_state_lock:
-                self._agent_state_buffer.clear()
-
         except Exception as e:
+            # Restore agent states to buffer so they aren't lost
+            with self._agent_state_lock:
+                self._agent_state_buffer.extend(agent_states_to_send)
             logger.error("Failed to flush agent states to AgentCore Memory: %s", e)
             raise SessionException(f"Failed to flush agent states: {e}") from e
 

tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py

@@ -11,6 +11,11 @@
 from botocore.config import Config as BotocoreConfig
 from botocore.exceptions import ClientError
 from strands.agent.agent import Agent
+from strands.experimental.hooks.events import (
+    BidiAfterInvocationEvent,
+    BidiAgentInitializedEvent,
+    BidiMessageAddedEvent,
+)
 from strands.experimental.hooks.multiagent.events import (
     AfterMultiAgentInvocationEvent,
     AfterNodeCallEvent,
@@ -3707,3 +3712,76 @@ def test_async_mode_logs_sync_invocation_warning(self, mock_memory_client, caplo
             manager.register_hooks(registry)
 
         assert any("async_mode=True" in rec.message and "stream_async" in rec.message for rec in caplog.records)
+
+    def test_async_mode_registers_bidi_agent_callbacks(self, mock_memory_client):
+        """async_mode=True: BidiAgent events get callbacks; init stays sync, others are async."""
+        config = AgentCoreMemoryConfig(memory_id="m", session_id="s", actor_id="a", async_mode=True)
+        manager = _create_session_manager(config, mock_memory_client)
+        registry = HookRegistry()
+        manager.register_hooks(registry)
+
+        # BidiAgentInitializedEvent dispatches via the sync hook path, so its callback must NOT be a coroutine.
+        init_callbacks = registry._registered_callbacks.get(BidiAgentInitializedEvent, [])
+        assert init_callbacks, "No callbacks registered for BidiAgentInitializedEvent"
+        assert not any(asyncio.iscoroutinefunction(cb) for cb in init_callbacks)
+
+        # BidiMessageAddedEvent and BidiAfterInvocationEvent dispatch via invoke_callbacks_async,
+        # so their callbacks should be async to keep the event loop unblocked.
+        for event_type in (BidiMessageAddedEvent, BidiAfterInvocationEvent):
+            callbacks = registry._registered_callbacks.get(event_type, [])
+            assert callbacks, f"No callbacks registered for {event_type.__name__}"
+            assert all(asyncio.iscoroutinefunction(cb) for cb in callbacks)
+
+
+class TestFlushAgentStatesRaceCondition:
+    """Tests for the copy-and-clear-under-one-lock fix in _flush_agent_states_only."""
+
+    def test_flush_agent_states_does_not_drop_concurrent_appends(self, batching_session_manager, mock_memory_client):
+        """States appended during the network I/O window must survive the flush."""
+        states_appended_during_flush = []
+
+        # Pre-populate the buffer with one state to force a flush.
+        initial_agent = SessionAgent(
+            agent_id="agent-1",
+            state={"description": "initial"},
+            conversation_manager_state={},
+        )
+        batching_session_manager.create_agent("test-session-456", initial_agent)
+        assert batching_session_manager.pending_agent_state_count() == 1
+
+        # Simulate a concurrent create_agent during the boto3 call. The mock fires
+        # while the buffer is being flushed — i.e. between the copy and any clear —
+        # so a second append must NOT be lost.
+        def create_event_and_append_concurrently(**kwargs):
+            new_agent = SessionAgent(
+                agent_id="agent-2",
+                state={"description": "appended-mid-flush"},
+                conversation_manager_state={},
+            )
+            batching_session_manager.create_agent("test-session-456", new_agent)
+            states_appended_during_flush.append(new_agent)
+            return {"eventId": "event_during_flush"}
+
+        mock_memory_client.gmdp_client.create_event.side_effect = create_event_and_append_concurrently
+
+        batching_session_manager._flush_agent_states_only()
+
+        # The state appended during the flush must still be in the buffer afterwards.
+        assert states_appended_during_flush, "Test setup error: concurrent append did not run"
+        assert batching_session_manager.pending_agent_state_count() == 1, (
+            "State appended during flush was dropped — copy/clear is not atomic"
+        )
+
+    def test_flush_agent_states_failure_restores_buffer(self, batching_session_manager, mock_memory_client):
+        """A failed flush must restore the originally-buffered states (no data loss)."""
+        mock_memory_client.gmdp_client.create_event.side_effect = Exception("API Error")
+
+        agent = SessionAgent(agent_id="agent-1", state={"description": "v1"}, conversation_manager_state={})
+        batching_session_manager.create_agent("test-session-456", agent)
+        assert batching_session_manager.pending_agent_state_count() == 1
+
+        with pytest.raises(SessionException):
+            batching_session_manager._flush_agent_states_only()
+
+        # State must be back in the buffer for retry.
+        assert batching_session_manager.pending_agent_state_count() == 1