From d79784f40685a61844f29378985c1a9538816083 Mon Sep 17 00:00:00 2001 From: Patrick Gray Date: Tue, 15 Sep 2026 22:17:19 -0400 Subject: [PATCH 1/3] fix(memory): use current Strands bidi session hooks --- pyproject.toml | 4 +-- .../integrations/strands/session_manager.py | 30 +++++-------------- .../test_agentcore_memory_session_manager.py | 19 ++++-------- uv.lock | 10 +++---- 4 files changed, 20 insertions(+), 43 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3b07516c..f15bbdf8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -149,7 +149,7 @@ dev = [ "ruff>=0.12.0", "websockets>=14.1", "wheel>=0.45.1", - "strands-agents>=1.46.0", + "strands-agents>=1.56.0", "strands-agents-evals>=1.0.3,<2.0.0", "deepeval>=3.5.0,<5.0.0", "autoevals>=0.3.0,<1.0.0", @@ -169,7 +169,7 @@ a2a = ["a2a-sdk[http-server]>=0.3,<0.4"] a2a-v1 = ["a2a-sdk[http-server]>=1.0.1,<2.0"] ag-ui = ["ag-ui-protocol>=0.1.10"] strands-agents = [ - "strands-agents>=1.46.0", + "strands-agents>=1.56.0", "mcp>=1.23.0,<2.0.0", ] langgraph = [ diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index f930b548..0bda3c0a 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -11,11 +11,7 @@ import boto3 from botocore.config import Config as BotocoreConfig -from strands.experimental.hooks.events import ( - BidiAfterInvocationEvent, - BidiAgentInitializedEvent, - BidiMessageAddedEvent, -) +from strands.experimental.bidi.hooks import BidiAgentStopEvent from strands.experimental.hooks.multiagent.events import ( AfterMultiAgentInvocationEvent, AfterNodeCallEvent, @@ -45,7 +41,7 @@ from .converters import MemoryConverter if TYPE_CHECKING: - from strands.agent.agent import Agent + from strands.types.agent import LocalAgent logger = logging.getLogger(__name__) @@ -829,7 +825,7 @@ def _filter_restored_tool_context(self, messages: list[SessionMessage]) -> list[ # region RepositorySessionManager overrides @override - def append_message(self, message: Message, agent: "Agent", **kwargs: Any) -> None: + 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. Args: @@ -951,9 +947,8 @@ def register_hooks(self, registry: HookRegistry, **kwargs) -> None: return # Async mode: register async callbacks that offload the existing sync - # methods to a worker thread via asyncio.to_thread. AgentInitializedEvent - # and BidiAgentInitializedEvent must stay sync (Strands disallows async - # callbacks for AgentInitializedEvent — see strands/hooks/registry.py:227). + # methods to a worker thread via asyncio.to_thread. Initialization + # callbacks must stay synchronous. logger.warning( "AgentCoreMemorySessionManager async_mode=True: the agent must be invoked " "via the async path (e.g. agent.stream_async(...) or agent.invoke_async(...)). " @@ -990,21 +985,10 @@ async def _on_message_added_persist(event: MessageAddedEvent) -> None: 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)) - # 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_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(BidiMessageAddedEvent, _on_bidi_message_added) - registry.add_callback(BidiAfterInvocationEvent, _offload(self.sync_bidi_agent, lambda e: e.agent)) + registry.add_callback(BidiAgentStopEvent, _offload(self.sync_agent, lambda e: e.agent)) @override - def initialize(self, agent: "Agent", **kwargs: Any) -> None: + def initialize(self, agent: "LocalAgent", **kwargs: Any) -> None: if self.has_existing_agent: logger.warning( "An Agent already exists in session %s. We currently support one agent per session.", self.session_id 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 b225cfa9..f3403334 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 @@ -11,17 +11,13 @@ 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.bidi.hooks import BidiAgentStopEvent from strands.experimental.hooks.multiagent.events import ( AfterMultiAgentInvocationEvent, AfterNodeCallEvent, MultiAgentInitializedEvent, ) -from strands.hooks import AfterInvocationEvent, MessageAddedEvent +from strands.hooks import AfterInvocationEvent, AgentInitializedEvent, MessageAddedEvent from strands.hooks.registry import HookRegistry from strands.types.exceptions import SessionException from strands.types.session import Session, SessionAgent, SessionMessage, SessionType @@ -3749,16 +3745,13 @@ def test_async_mode_registers_bidi_agent_callbacks(self, 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 = list(registry.get_callbacks_for(BidiAgentInitializedEvent(agent=Mock()))) - assert init_callbacks, "No callbacks registered for BidiAgentInitializedEvent" + init_callbacks = list(registry.get_callbacks_for(AgentInitializedEvent(agent=Mock()))) + assert init_callbacks 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 in ( - BidiMessageAddedEvent(agent=Mock(), message={"role": "user", "content": [{"text": "x"}]}), - BidiAfterInvocationEvent(agent=Mock()), + MessageAddedEvent(agent=Mock(), message={"role": "user", "content": [{"text": "x"}]}), + BidiAgentStopEvent(agent=Mock()), ): callbacks = list(registry.get_callbacks_for(event)) assert callbacks, f"No callbacks registered for {type(event).__name__}" diff --git a/uv.lock b/uv.lock index 6f398f71..240f746f 100644 --- a/uv.lock +++ b/uv.lock @@ -606,7 +606,7 @@ requires-dist = [ { name = "ragas", marker = "extra == 'ragas'", specifier = ">=0.4.3,<1.0.0" }, { name = "requests", marker = "extra == 'datasets'", specifier = ">=2.31.0" }, { name = "starlette", specifier = ">=0.46.2" }, - { name = "strands-agents", marker = "extra == 'strands-agents'", specifier = ">=1.46.0" }, + { name = "strands-agents", marker = "extra == 'strands-agents'", specifier = ">=1.56.0" }, { name = "strands-agents-evals", marker = "extra == 'autoevals'", specifier = ">=1.0.3,<2.0.0" }, { name = "strands-agents-evals", marker = "extra == 'deepeval'", specifier = ">=1.0.3,<2.0.0" }, { name = "strands-agents-evals", marker = "extra == 'ragas'", specifier = ">=1.0.3,<2.0.0" }, @@ -641,7 +641,7 @@ dev = [ { name = "pytest-rerunfailures", specifier = ">=15.0" }, { name = "ragas", specifier = ">=0.4.3,<1.0.0" }, { name = "ruff", specifier = ">=0.12.0" }, - { name = "strands-agents", specifier = ">=1.46.0" }, + { name = "strands-agents", specifier = ">=1.56.0" }, { name = "strands-agents-evals", specifier = ">=1.0.3,<2.0.0" }, { name = "websockets", specifier = ">=14.1" }, { name = "wheel", specifier = ">=0.45.1" }, @@ -5059,7 +5059,7 @@ wheels = [ [[package]] name = "strands-agents" -version = "1.50.2" +version = "1.56.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "boto3" }, @@ -5076,9 +5076,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "watchdog" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d0/7a/b5e20f5ee859c71e451e2e06de8767aab73d850018167c436e3a96446caf/strands_agents-1.50.2.tar.gz", hash = "sha256:7b45c41593e41f22383962a233b6616360b8c3686ae8b902059b6a6fab3deacf", size = 1224067, upload-time = "2026-07-27T20:38:48.238Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4c/2f/c980715e56a2017280ba597aea83e77ec37a94eeaaabc47114e354a6846e/strands_agents-1.56.0.tar.gz", hash = "sha256:d0a4587065f9b9f4890455ce947474067e5deaadf140d1b2ef5153a67ebf80a3", size = 1662954, upload-time = "2026-09-15T19:54:39.56Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d8/d6/2fd27147f53c45e546c0c542c6ea4a24b93c1f3908dcbad164624e5ddfa6/strands_agents-1.50.2-py3-none-any.whl", hash = "sha256:39c6b755e579e0b631ea01af78a0a201019320670b87adcf0b5c8dd1e5a2cef8", size = 638411, upload-time = "2026-07-27T20:38:46.12Z" }, + { url = "https://files.pythonhosted.org/packages/89/07/0d8955760ed673c3112af65d2a026e79bbcb3c869b3eadc896820b64edd7/strands_agents-1.56.0-py3-none-any.whl", hash = "sha256:8026a2fde7ca3d2f2760ea57003dd7163761bfba74972762468b745b4f7bfd08", size = 838144, upload-time = "2026-09-15T19:54:37.448Z" }, ] [[package]] From 7ce133152e0443e810e00d193f46fb7a9876d166 Mon Sep 17 00:00:00 2001 From: Patrick Gray Date: Wed, 16 Sep 2026 09:34:24 -0400 Subject: [PATCH 2/3] fix(memory): flush session batches after final state sync --- .../integrations/strands/session_manager.py | 20 ++++++----- .../test_agentcore_memory_session_manager.py | 34 +++++++++++++++++-- 2 files changed, 43 insertions(+), 11 deletions(-) diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index 0bda3c0a..ef0af6da 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -920,8 +920,7 @@ def register_hooks(self, registry: HookRegistry, **kwargs) -> None: """Register additional hooks. In sync mode (the default), delegates to the base class and adds the - retrieve_customer_context + batching callbacks synchronously, preserving - existing behavior exactly. + retrieve_customer_context + batching callbacks synchronously. In async mode, registers async callbacks that wrap every per-turn boto3-backed operation (append_message, sync_agent, buffer flushes, @@ -938,12 +937,13 @@ def register_hooks(self, registry: HookRegistry, **kwargs) -> None: **kwargs: Additional keyword arguments. """ if not self.config.async_mode: - RepositorySessionManager.register_hooks(self, registry, **kwargs) - registry.add_callback(MessageAddedEvent, lambda event: self.retrieve_customer_context(event)) - - # Only register AfterInvocationEvent hook when batching is enabled if self.config.batch_size > 1: + # Completion callbacks run in reverse order, so register flushes before state syncs. registry.add_callback(AfterInvocationEvent, lambda event: self._flush_messages()) + registry.add_callback(BidiAgentStopEvent, lambda event: self._flush_messages()) + + RepositorySessionManager.register_hooks(self, registry, **kwargs) + registry.add_callback(MessageAddedEvent, lambda event: self.retrieve_customer_context(event)) return # Async mode: register async callbacks that offload the existing sync @@ -967,6 +967,11 @@ async def _callback(event): return _callback + if self.config.batch_size > 1: + # Completion callbacks run in reverse order, so register flushes before state syncs. + registry.add_callback(AfterInvocationEvent, _offload(self._flush_messages)) + registry.add_callback(BidiAgentStopEvent, _offload(self._flush_messages)) + registry.add_callback(AgentInitializedEvent, lambda event: self.initialize(event.agent)) async def _on_message_added_persist(event: MessageAddedEvent) -> None: @@ -977,9 +982,6 @@ async def _on_message_added_persist(event: MessageAddedEvent) -> None: 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: - registry.add_callback(AfterInvocationEvent, _offload(self._flush_messages)) - # Register multi-agent callbacks so async-mode parity matches sync-mode 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)) 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 f3403334..fb9bc7b9 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 @@ -2,6 +2,7 @@ import asyncio import inspect +import json import logging import time from datetime import datetime, timedelta, timezone @@ -2675,8 +2676,8 @@ def test_retrieve_customer_context_default_context_tag(self, mock_memory_client) assert "" in content[0]["text"] -class TestAfterInvocationHook: - """Test AfterInvocationEvent hook integration.""" +class TestSessionHooks: + """Test session lifecycle hook integration.""" def test_after_invocation_hook_registered(self, batching_session_manager): """Test that AfterInvocationEvent hook is registered when batching is enabled.""" @@ -2741,6 +2742,35 @@ def spy_add_callback(event_type, callback): ] assert len(flush_callbacks) == 0 + @pytest.mark.parametrize("async_mode", [False, True]) + @pytest.mark.parametrize("event_type", [AfterInvocationEvent, BidiAgentStopEvent]) + async def test_completion_flushes_messages_and_final_state( + self, batching_config, mock_memory_client, async_mode, event_type + ): + """Completion flushes a partial batch, including the final state update.""" + batching_config.async_mode = async_mode + manager = _create_session_manager(batching_config, mock_memory_client) + manager.session_repository = manager + agent = Mock(agent_id="test-agent") + agent.state.get.return_value = {} + manager.create_agent(manager.session_id, SessionAgent.from_agent(agent)) + manager.create_message( + manager.session_id, + agent.agent_id, + SessionMessage(message={"role": "user", "content": [{"text": "Hello"}]}, message_id=0), + ) + manager.memory_client.gmdp_client.create_event.assert_not_called() + + agent.state.get.return_value = {"status": "stopped"} + registry = HookRegistry() + manager.register_hooks(registry) + await registry.invoke_callbacks_async(event_type(agent=agent)) + + assert manager.pending_message_count() == 0 + assert manager.pending_agent_state_count() == 0 + state_payloads = manager.memory_client.gmdp_client.create_event.call_args.kwargs["payload"] + assert [json.loads(payload["blob"])["state"] for payload in state_payloads] == [{}, {"status": "stopped"}] + class TestIntervalFlush: """Test interval-based flush mechanism for long-running agents.""" From 250a901308fd87212a0eb39fc7b1ca0915944647 Mon Sep 17 00:00:00 2001 From: Patrick Gray Date: Wed, 16 Sep 2026 09:51:57 -0400 Subject: [PATCH 3/3] fix(memory): skip automatic retrieval for bidi agents --- .../memory/integrations/strands/config.py | 3 +- .../integrations/strands/session_manager.py | 6 +++- .../test_agentcore_memory_session_manager.py | 28 +++++++++++++++++++ 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/src/bedrock_agentcore/memory/integrations/strands/config.py b/src/bedrock_agentcore/memory/integrations/strands/config.py index a42e824e..8f7f4117 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/config.py +++ b/src/bedrock_agentcore/memory/integrations/strands/config.py @@ -47,7 +47,8 @@ class AgentCoreMemoryConfig(BaseModel): memory_id: Required Bedrock AgentCore Memory ID session_id: Required unique ID for the session actor_id: Required unique ID for the agent instance/user - retrieval_config: Optional dictionary mapping namespaces to retrieval configurations + retrieval_config: Optional dictionary mapping namespaces to retrieval configurations. + Automatic context retrieval applies to Agent only. batch_size: Number of messages to batch before sending to AgentCore Memory. Default of 1 means immediate sending (no batching). Max 100. flush_interval_seconds: Optional interval in seconds for automatic buffer flushing. diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index ef0af6da..d6d20ab7 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -11,6 +11,7 @@ import boto3 from botocore.config import Config as BotocoreConfig +from strands.experimental.bidi import BidiAgent from strands.experimental.bidi.hooks import BidiAgentStopEvent from strands.experimental.hooks.multiagent.events import ( AfterMultiAgentInvocationEvent, @@ -840,11 +841,14 @@ def append_message(self, message: Message, agent: "LocalAgent", **kwargs: Any) - self._latest_agent_message[agent.agent_id] = session_message def retrieve_customer_context(self, event: MessageAddedEvent) -> None: - """Retrieve customer LTM context before processing support query. + """Retrieve customer LTM context for regular Agent invocations. Args: event (MessageAddedEvent): The message added event containing the agent and message data. """ + if isinstance(event.agent, BidiAgent): + return None + messages = event.agent.messages if not messages or messages[-1].get("role") != "user": return None 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 fb9bc7b9..35d3e8b9 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 @@ -12,6 +12,7 @@ from botocore.config import Config as BotocoreConfig from botocore.exceptions import ClientError from strands.agent.agent import Agent +from strands.experimental.bidi import BidiAgent from strands.experimental.bidi.hooks import BidiAgentStopEvent from strands.experimental.hooks.multiagent.events import ( AfterMultiAgentInvocationEvent, @@ -2679,6 +2680,33 @@ def test_retrieve_customer_context_default_context_tag(self, mock_memory_client) class TestSessionHooks: """Test session lifecycle hook integration.""" + @pytest.mark.parametrize("async_mode", [False, True]) + async def test_bidi_message_persists_without_retrieval( + self, agentcore_config_with_retrieval, mock_memory_client, async_mode + ): + """Bidi messages are persisted without retrieving or injecting context.""" + agentcore_config_with_retrieval.async_mode = async_mode + manager = _create_session_manager(agentcore_config_with_retrieval, mock_memory_client) + manager.session_repository = Mock() + manager._latest_agent_message = {} + agent = Mock( + spec=BidiAgent, + agent_id="test-agent", + messages=[{"role": "user", "content": [{"text": "Hello"}]}], + state=Mock(), + ) + agent.state.get.return_value = {} + mock_memory_client.retrieve_memories.return_value = [{"content": {"text": "User prefers blue"}, "score": 1.0}] + registry = HookRegistry() + manager.register_hooks(registry) + + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=agent.messages[0])) + + mock_memory_client.retrieve_memories.assert_not_called() + assert agent.messages == [{"role": "user", "content": [{"text": "Hello"}]}] + mock_memory_client.create_event.assert_called_once() + manager.session_repository.update_agent.assert_called_once() + def test_after_invocation_hook_registered(self, batching_session_manager): """Test that AfterInvocationEvent hook is registered when batching is enabled.""" registry = HookRegistry()