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."""