From faef3a09dcb674b288e53c75dee9522c5ead7ee5 Mon Sep 17 00:00:00 2001 From: Abhiram Shaji Date: Tue, 27 Oct 2026 20:42:07 -0400 Subject: [PATCH] fix(memory): store completed BidiAgent transcripts BidiAgent adds each spoken turn to history as an empty placeholder (MessageAddedEvent, bidi status "pending") and later replaces it with the finished text through MessageUpdatedEvent (status "complete"). The session manager only persisted on MessageAddedEvent, so the placeholder was dropped as empty and the finished text was never written. Voice agents therefore stored no user or assistant turns and long term memory strategies had nothing to extract. Persist MessageUpdatedEvent messages whose bidi status is "complete" through append_message, the same path text messages use, in both sync and async mode. Pending placeholders are skipped explicitly so a custom converter cannot store them, and "incomplete" updates are ignored. Each transcript completes once, so it is written exactly once. The hook is only registered when the installed strands-agents provides MessageUpdatedEvent. Closes #716 --- .../integrations/strands/session_manager.py | 22 ++++- .../test_agentcore_memory_session_manager.py | 83 +++++++++++++++++++ 2 files changed, 104 insertions(+), 1 deletion(-) diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index be5632dc..ee1127a7 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -7,9 +7,10 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timedelta, timezone from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Optional +from typing import TYPE_CHECKING, Any, Dict, Mapping, NamedTuple, Optional import boto3 +import strands.hooks from botocore.config import Config as BotocoreConfig from strands.experimental.bidi import BidiAgent from strands.experimental.bidi.hooks import BidiAgentStopEvent @@ -46,6 +47,8 @@ logger = logging.getLogger(__name__) +MessageUpdatedEvent = getattr(strands.hooks, "MessageUpdatedEvent", None) + MAX_FETCH_ALL_RESULTS = 10000 # Legacy prefixes for backwards compatibility with old events @@ -63,6 +66,13 @@ RESERVED_METADATA_KEYS = frozenset({STATE_TYPE_KEY, AGENT_ID_KEY}) +def _bidi_status(message: Mapping[str, Any]) -> Optional[str]: + metadata = message.get("metadata") or {} + bidi = (metadata.get("custom") or {}).get("bidi") or {} + status: Optional[str] = bidi.get("status") + return status + + class BufferedMessage(NamedTuple): """A pre-processed message waiting to be flushed to AgentCore Memory.""" @@ -841,12 +851,18 @@ def append_message(self, message: Message, agent: "LocalAgent", **kwargs: Any) - agent: Agent to append the message to **kwargs: Additional keyword arguments for future extensibility. """ + if _bidi_status(message) == "pending": + return created_message = self.create_message(self.session_id, agent.agent_id, SessionMessage.from_message(message, 0)) if created_message is None: return session_message = SessionMessage.from_message(message, created_message.get("eventId")) self._latest_agent_message[agent.agent_id] = session_message + def _append_completed_message(self, event: Any) -> None: + if _bidi_status(event.message) == "complete": + self.append_message(event.message, event.agent) + def retrieve_customer_context(self, event: MessageAddedEvent) -> None: """Retrieve customer LTM context for regular Agent invocations. @@ -955,6 +971,8 @@ def register_hooks(self, registry: HookRegistry, **kwargs) -> None: RepositorySessionManager.register_hooks(self, registry, **kwargs) registry.add_callback(MessageAddedEvent, lambda event: self.retrieve_customer_context(event)) + if MessageUpdatedEvent is not None: + registry.add_callback(MessageUpdatedEvent, self._append_completed_message) return # Async mode: register async callbacks that offload the existing sync @@ -990,6 +1008,8 @@ async def _on_message_added_persist(event: MessageAddedEvent) -> None: await asyncio.to_thread(self.sync_agent, event.agent) registry.add_callback(MessageAddedEvent, _on_message_added_persist) + if MessageUpdatedEvent is not None: + registry.add_callback(MessageUpdatedEvent, _offload(self._append_completed_message, lambda e: e)) registry.add_callback(AfterInvocationEvent, _offload(self.sync_agent, lambda e: e.agent)) registry.add_callback(MessageAddedEvent, _offload(self.retrieve_customer_context, lambda e: e)) 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..1d4348e1 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,6 +11,7 @@ import pytest from botocore.config import Config as BotocoreConfig from botocore.exceptions import ClientError +from strands import hooks as strands_hooks from strands.agent.agent import Agent from strands.experimental.bidi import BidiAgent from strands.experimental.bidi.hooks import BidiAgentStopEvent @@ -34,6 +35,8 @@ BufferedMessage, ) +MessageUpdatedEvent = getattr(strands_hooks, "MessageUpdatedEvent", None) + @pytest.fixture def agentcore_config(): @@ -2846,6 +2849,86 @@ async def test_completion_flushes_messages_and_final_state( assert [json.loads(payload["blob"])["state"] for payload in state_payloads] == [{}, {"status": "stopped"}] +def _bidi_message(role, text, status): + content = [] if text is None else [{"text": text}] + return { + "role": role, + "content": content, + "tracking_id": f"{role}-1", + "metadata": {"custom": {"bidi": {"kind": "transcript", "status": status}}}, + } + + +class TestBidiTranscripts: + """BidiAgent transcripts are stored once they are complete.""" + + @pytest.fixture(params=[False, True], ids=["sync", "async"]) + def bidi_manager(self, request, agentcore_config, mock_memory_client): + agentcore_config.async_mode = request.param + manager = _create_session_manager(agentcore_config, mock_memory_client) + manager._latest_agent_message = {} + manager.sync_agent = Mock() + registry = HookRegistry() + manager.register_hooks(registry) + return registry, Mock(spec=BidiAgent, agent_id="test-agent") + + async def test_pending_placeholder_is_not_stored(self, bidi_manager, mock_memory_client): + registry, agent = bidi_manager + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + + mock_memory_client.create_event.assert_not_called() + + @pytest.mark.parametrize("metadata", [None, {}, {"custom": None}, {"custom": {"bidi": {"kind": "tool_result"}}}]) + async def test_message_without_bidi_status_is_stored(self, bidi_manager, mock_memory_client, metadata): + registry, agent = bidi_manager + message = {"role": "user", "content": [{"text": "Hello"}], "metadata": metadata} + + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=message)) + + mock_memory_client.create_event.assert_called_once() + + @pytest.mark.skipif(MessageUpdatedEvent is None, reason="strands-agents has no MessageUpdatedEvent") + @pytest.mark.parametrize(("role", "memory_role"), [("user", "USER"), ("assistant", "ASSISTANT")]) + async def test_complete_transcript_is_stored_once(self, bidi_manager, mock_memory_client, role, memory_role): + registry, agent = bidi_manager + final = _bidi_message(role, "I need a hotel in London.", "complete") + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message(role, None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, final["tracking_id"], final)) + + mock_memory_client.create_event.assert_called_once() + [(payload, stored_role)] = mock_memory_client.create_event.call_args.kwargs["messages"] + assert stored_role.upper() == memory_role + assert json.loads(payload)["message"]["content"] == [{"text": "I need a hotel in London."}] + + @pytest.mark.skipif(MessageUpdatedEvent is None, reason="strands-agents has no MessageUpdatedEvent") + @pytest.mark.parametrize( + ("text", "status"), + [("[Transcript unavailable.]", "incomplete"), ("", "complete"), (" ", "complete"), ("partial", "pending")], + ) + async def test_unfinished_or_empty_transcript_is_not_stored(self, bidi_manager, mock_memory_client, text, status): + registry, agent = bidi_manager + update = _bidi_message("assistant", text, status) + + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, update["tracking_id"], update)) + + mock_memory_client.create_event.assert_not_called() + + @pytest.mark.skipif(MessageUpdatedEvent is None, reason="strands-agents has no MessageUpdatedEvent") + async def test_update_without_bidi_status_is_not_stored(self, bidi_manager, mock_memory_client): + registry, agent = bidi_manager + message = {"role": "assistant", "content": [{"text": "redacted"}], "tracking_id": "x"} + + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, "x", message)) + + mock_memory_client.create_event.assert_not_called() + + class TestIntervalFlush: """Test interval-based flush mechanism for long-running agents."""