diff --git a/pyproject.toml b/pyproject.toml index b24e9149..3cfe362c 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.56.0", + "strands-agents>=1.57.1", "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.56.0", + "strands-agents>=1.57.1", "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 242b2fdc..290e5f69 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -4,10 +4,11 @@ import json import logging import threading +from collections import deque 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 from botocore.config import Config as BotocoreConfig @@ -18,7 +19,7 @@ AfterNodeCallEvent, MultiAgentInitializedEvent, ) -from strands.hooks import AfterInvocationEvent, MessageAddedEvent +from strands.hooks import AfterInvocationEvent, MessageAddedEvent, MessageUpdatedEvent from strands.hooks.events import AgentInitializedEvent from strands.hooks.registry import HookRegistry from strands.session.repository_session_manager import RepositorySessionManager @@ -64,6 +65,19 @@ 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 _HeldMessage(NamedTuple): + message: Message + agent: "LocalAgent" + ready: bool + + class BufferedMessage(NamedTuple): """A pre-processed message waiting to be flushed to AgentCore Memory. @@ -171,6 +185,9 @@ def __init__( self._agent_state_buffer: list[tuple[str, SessionAgent]] = [] self._agent_state_lock = threading.Lock() + self._held_messages: deque[_HeldMessage] = deque() + self._held_lock = threading.Lock() + # Cache for agent created_at timestamps to avoid fetching on every update self._agent_created_at_cache: dict[str, datetime] = {} @@ -855,6 +872,14 @@ def append_message(self, message: Message, agent: "LocalAgent", **kwargs: Any) - agent: Agent to append the message to **kwargs: Additional keyword arguments for future extensibility. """ + ready = _bidi_status(message) != "pending" + with self._held_lock: + if ready and not self._held_messages: + self._write_message(message, agent) + else: + self._held_messages.append(_HeldMessage(message, agent, ready)) + + def _write_message(self, message: Message, agent: "LocalAgent") -> None: created_message = self.create_message( self.session_id, agent.agent_id, SessionMessage.from_message(message, 0), _protect_latest=True ) @@ -863,6 +888,54 @@ def append_message(self, message: Message, agent: "LocalAgent", **kwargs: Any) - session_message = SessionMessage.from_message(message, created_message.get("eventId")) self._latest_agent_message[agent.agent_id] = session_message + def _complete_message(self, event: MessageUpdatedEvent) -> None: + status = _bidi_status(event.message) + if status == "pending": + return + with self._held_lock: + index = next( + ( + i + for i, held in enumerate(self._held_messages) + if not held.ready and held.message.get("tracking_id") == event.tracking_id + ), + None, + ) + if index is None: + logger.debug("tracking_id=<%s> | no waiting message for transcript update, ignoring", event.tracking_id) + return + if status == "complete": + self._held_messages[index] = _HeldMessage(event.message, event.agent, True) + elif status == "incomplete": + del self._held_messages[index] + self._write_ready_messages() + + def _write_ready_messages(self) -> None: + first_error: Optional[Exception] = None + while self._held_messages and self._held_messages[0].ready: + held = self._held_messages.popleft() + try: + self._write_message(held.message, held.agent) + except Exception as error: + if first_error is not None: + logger.error("Failed to write held message: %s", error) + else: + first_error = error + if first_error is not None: + raise first_error + + def _release_held_messages(self) -> None: + with self._held_lock: + self._held_messages = deque(held for held in self._held_messages if held.ready) + self._write_ready_messages() + + def _release_and_flush(self) -> None: + try: + self._release_held_messages() + finally: + if self.config.batch_size > 1: + self._flush_messages() + def retrieve_customer_context(self, event: MessageAddedEvent) -> None: """Retrieve customer LTM context for regular Agent invocations. @@ -967,10 +1040,11 @@ def register_hooks(self, registry: HookRegistry, **kwargs) -> None: 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()) + registry.add_callback(BidiAgentStopEvent, lambda event: self._release_and_flush()) RepositorySessionManager.register_hooks(self, registry, **kwargs) registry.add_callback(MessageAddedEvent, lambda event: self.retrieve_customer_context(event)) + registry.add_callback(MessageUpdatedEvent, self._complete_message) return # Async mode: register async callbacks that offload the existing sync @@ -997,7 +1071,7 @@ async def _callback(event): 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(BidiAgentStopEvent, _offload(self._release_and_flush)) registry.add_callback(AgentInitializedEvent, lambda event: self.initialize(event.agent)) @@ -1006,6 +1080,7 @@ 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) + registry.add_callback(MessageUpdatedEvent, _offload(self._complete_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)) @@ -1277,7 +1352,10 @@ def close(self) -> None: manager protocol (with statement) for automatic cleanup. """ self._stop_flush_timer() - self._flush_messages() + try: + self._release_held_messages() + finally: + self._flush_messages() def __enter__(self) -> "AgentCoreMemorySessionManager": """Enter the context manager. 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 3f48e3a3..943e8200 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 @@ -19,7 +19,7 @@ AfterNodeCallEvent, MultiAgentInitializedEvent, ) -from strands.hooks import AfterInvocationEvent, AgentInitializedEvent, MessageAddedEvent +from strands.hooks import AfterInvocationEvent, AgentInitializedEvent, MessageAddedEvent, MessageUpdatedEvent from strands.hooks.registry import HookRegistry from strands.types.exceptions import SessionException from strands.types.session import Session, SessionAgent, SessionMessage, SessionType @@ -86,6 +86,7 @@ def _create_session_manager(config, mock_memory_client): manager = AgentCoreMemorySessionManager(config) manager.session_id = config.session_id manager.session = Session(session_id=config.session_id, session_type=SessionType.AGENT) + manager._held_records = {} return manager @@ -2998,7 +2999,7 @@ def test_after_invocation_hook_flushes_buffer(self, batching_session_manager, mo batching_session_manager.register_hooks(registry) # Create mock event with mock agent - mock_agent = Mock() + mock_agent = Mock(messages=[]) mock_event = AfterInvocationEvent(agent=mock_agent) registry.invoke_callbacks(mock_event) @@ -3038,7 +3039,7 @@ async def test_completion_flushes_messages_and_final_state( 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 = Mock(agent_id="test-agent", messages=[]) agent.state.get.return_value = {} manager.create_agent(manager.session_id, SessionAgent.from_agent(agent)) manager.create_message( @@ -3059,6 +3060,358 @@ 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, tracking_id=None): + content = [] if text is None else [{"text": text}] + return { + "role": role, + "content": content, + "tracking_id": tracking_id or f"{role}-1", + "metadata": {"custom": {"bidi": {"kind": "transcript", "status": status}}}, + } + + +def _stored_texts(mock_memory_client): + return [ + (role.upper(), json.loads(payload)["message"]["content"][0]) + for call in mock_memory_client.create_event.call_args_list + for payload, role in call.kwargs["messages"] + ] + + +class _KeepEmptyConverter(AgentCoreMemoryConverter): + @staticmethod + def message_to_payload(session_message): + session_dict = session_message.to_dict() + return [(json.dumps(session_dict), session_dict["message"]["role"])] + + +class TestBidiTranscripts: + """BidiAgent transcripts are stored once they are complete, in history order.""" + + @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 manager, registry, Mock(spec=BidiAgent, agent_id="test-agent") + + async def test_pending_placeholder_is_not_stored_by_converter_keeping_empty_content( + self, bidi_manager, mock_memory_client + ): + manager, registry, agent = bidi_manager + manager.converter = _KeepEmptyConverter + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async(BidiAgentStopEvent(agent=agent)) + + 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): + manager, 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.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): + manager, 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)) + + assert _stored_texts(mock_memory_client) == [(memory_role, {"text": "I need a hotel in London."})] + + @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): + manager, registry, agent = bidi_manager + update = _bidi_message("assistant", text, status) + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("assistant", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, update["tracking_id"], update)) + + mock_memory_client.create_event.assert_not_called() + + async def test_update_without_bidi_status_is_not_stored(self, bidi_manager, mock_memory_client): + manager, 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() + + async def test_reply_completed_before_question_is_stored_after_it(self, bidi_manager, mock_memory_client): + manager, registry, agent = bidi_manager + question = _bidi_message("user", "Any rooms in London?", "complete") + reply = _bidi_message("assistant", "Yes, two rooms.", "complete") + follow_up = {"role": "user", "content": [{"text": "Book one."}]} + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("assistant", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, reply["tracking_id"], reply)) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=follow_up)) + mock_memory_client.create_event.assert_not_called() + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, question["tracking_id"], question)) + + assert _stored_texts(mock_memory_client) == [ + ("USER", {"text": "Any rooms in London?"}), + ("ASSISTANT", {"text": "Yes, two rooms."}), + ("USER", {"text": "Book one."}), + ] + + async def test_incomplete_transcript_releases_messages_after_it(self, bidi_manager, mock_memory_client): + manager, registry, agent = bidi_manager + lost = _bidi_message("user", "[Transcript unavailable.]", "incomplete") + reply = _bidi_message("assistant", "Yes, two rooms.", "complete") + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("assistant", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, reply["tracking_id"], reply)) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, lost["tracking_id"], lost)) + + assert _stored_texts(mock_memory_client) == [("ASSISTANT", {"text": "Yes, two rooms."})] + + async def test_update_without_bidi_status_keeps_transcript_waiting(self, bidi_manager, mock_memory_client): + manager, registry, agent = bidi_manager + redacted = {"role": "user", "content": [{"text": "redacted"}], "tracking_id": "user-1"} + question = _bidi_message("user", "Any rooms in London?", "complete") + follow_up = {"role": "assistant", "content": [{"text": "Yes, two rooms."}]} + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, "user-1", redacted)) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=follow_up)) + mock_memory_client.create_event.assert_not_called() + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, question["tracking_id"], question)) + + assert _stored_texts(mock_memory_client) == [ + ("USER", {"text": "Any rooms in London?"}), + ("ASSISTANT", {"text": "Yes, two rooms."}), + ] + + @pytest.mark.parametrize("async_mode", [False, True], ids=["sync", "async"]) + async def test_batched_messages_keep_history_order(self, agentcore_config, mock_memory_client, async_mode): + agentcore_config.async_mode = async_mode + agentcore_config.batch_size = 10 + manager = _create_session_manager(agentcore_config, mock_memory_client) + manager._latest_agent_message = {} + manager.sync_agent = Mock() + registry = HookRegistry() + manager.register_hooks(registry) + agent = Mock(spec=BidiAgent, agent_id="test-agent") + question = _bidi_message("user", "Any rooms in London?", "complete") + reply = _bidi_message("assistant", "Yes, two rooms.", "complete") + follow_up = {"role": "user", "content": [{"text": "Book one."}]} + confirmation = {"role": "assistant", "content": [{"text": "Booked."}]} + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("assistant", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, reply["tracking_id"], reply)) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=follow_up)) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, question["tracking_id"], question)) + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending", tracking_id="user-2")) + ) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=confirmation)) + assert manager.pending_message_count() == 3 + await registry.invoke_callbacks_async(BidiAgentStopEvent(agent=agent)) + + payload = mock_memory_client.gmdp_client.create_event.call_args.kwargs["payload"] + assert [json.loads(item["conversational"]["content"]["text"])["message"]["content"][0] for item in payload] == [ + {"text": "Any rooms in London?"}, + {"text": "Yes, two rooms."}, + {"text": "Book one."}, + {"text": "Booked."}, + ] + + @pytest.mark.parametrize("finish", ["stop", "close"]) + async def test_finishing_writes_messages_held_behind_open_transcript( + self, bidi_manager, mock_memory_client, finish + ): + manager, registry, agent = bidi_manager + reply = _bidi_message("assistant", "Yes, two rooms.", "complete") + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("assistant", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, reply["tracking_id"], reply)) + if finish == "stop": + await registry.invoke_callbacks_async(BidiAgentStopEvent(agent=agent)) + else: + manager.close() + + assert _stored_texts(mock_memory_client) == [("ASSISTANT", {"text": "Yes, two rooms."})] + + async def test_failed_write_does_not_block_messages_after_it(self, bidi_manager, mock_memory_client): + manager, registry, agent = bidi_manager + question = _bidi_message("user", "Any rooms in London?", "complete") + reply = _bidi_message("assistant", "Yes, two rooms.", "complete") + follow_up = {"role": "user", "content": [{"text": "Book one."}]} + mock_memory_client.create_event.side_effect = [ClientError({"Error": {"Code": "500"}}, "CreateEvent"), {}, {}] + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("assistant", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, reply["tracking_id"], reply)) + with pytest.raises(SessionException): + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, question["tracking_id"], question)) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=follow_up)) + + failed, *stored = _stored_texts(mock_memory_client) + assert failed == ("USER", {"text": "Any rooms in London?"}) + assert stored == [("ASSISTANT", {"text": "Yes, two rooms."}), ("USER", {"text": "Book one."})] + + async def test_complete_without_waiting_message_is_ignored(self, bidi_manager, mock_memory_client): + manager, registry, agent = bidi_manager + final = _bidi_message("assistant", "Late reply.", "complete", tracking_id="late") + + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, "late", final)) + manager.close() + + mock_memory_client.create_event.assert_not_called() + + async def test_repeated_complete_is_stored_once(self, bidi_manager, mock_memory_client): + manager, registry, agent = bidi_manager + final = _bidi_message("assistant", "Yes, two rooms.", "complete") + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("assistant", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, final["tracking_id"], final)) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, final["tracking_id"], final)) + + assert _stored_texts(mock_memory_client) == [("ASSISTANT", {"text": "Yes, two rooms."})] + + async def test_repeated_complete_while_waiting_is_stored_once(self, bidi_manager, mock_memory_client): + manager, registry, agent = bidi_manager + question = _bidi_message("user", "Any rooms in London?", "complete") + reply = _bidi_message("assistant", "Yes, two rooms.", "complete") + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("assistant", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, reply["tracking_id"], reply)) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, reply["tracking_id"], reply)) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, question["tracking_id"], question)) + + assert _stored_texts(mock_memory_client) == [ + ("USER", {"text": "Any rooms in London?"}), + ("ASSISTANT", {"text": "Yes, two rooms."}), + ] + + async def test_incomplete_transcript_in_middle_is_skipped(self, bidi_manager, mock_memory_client): + manager, registry, agent = bidi_manager + question = _bidi_message("user", "Any rooms in London?", "complete") + cut_off = _bidi_message("assistant", "[Transcript unavailable.]", "incomplete") + follow_up = {"role": "user", "content": [{"text": "Book one."}]} + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("assistant", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=follow_up)) + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, cut_off["tracking_id"], cut_off)) + mock_memory_client.create_event.assert_not_called() + await registry.invoke_callbacks_async(MessageUpdatedEvent(agent, question["tracking_id"], question)) + + assert _stored_texts(mock_memory_client) == [ + ("USER", {"text": "Any rooms in London?"}), + ("USER", {"text": "Book one."}), + ] + assert not manager._held_messages + + async def test_close_writes_and_flushes_after_failed_held_write(self, bidi_manager, mock_memory_client): + manager, registry, agent = bidi_manager + reply = _bidi_message("assistant", "Yes, two rooms.", "complete") + follow_up = {"role": "user", "content": [{"text": "Book one."}]} + mock_memory_client.create_event.side_effect = [ClientError({"Error": {"Code": "500"}}, "CreateEvent"), {}] + manager._flush_messages = Mock(wraps=manager._flush_messages) + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=reply)) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=follow_up)) + with pytest.raises(SessionException): + manager.close() + + assert _stored_texts(mock_memory_client) == [ + ("ASSISTANT", {"text": "Yes, two rooms."}), + ("USER", {"text": "Book one."}), + ] + assert not manager._held_messages + manager._flush_messages.assert_called_once() + + +@pytest.mark.parametrize("async_mode", [False, True]) +async def test_stop_syncs_and_flushes_after_failed_held_write(agentcore_config, mock_memory_client, async_mode): + agentcore_config.async_mode = async_mode + agentcore_config.batch_size = 10 + manager = _create_session_manager(agentcore_config, mock_memory_client) + manager._latest_agent_message = {} + manager.sync_agent = Mock() + registry = HookRegistry() + manager.register_hooks(registry) + agent = Mock(spec=BidiAgent, agent_id="test-agent") + reply = _bidi_message("assistant", "Yes, two rooms.", "complete") + follow_up = {"role": "user", "content": [{"text": "Book one."}]} + + await registry.invoke_callbacks_async( + MessageAddedEvent(agent=agent, message=_bidi_message("user", None, "pending")) + ) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=reply)) + await registry.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=follow_up)) + manager.sync_agent.reset_mock() + manager.create_message = Mock(side_effect=[SessionException("write failed"), {"eventId": "e1"}]) + manager._flush_messages = Mock() + with pytest.raises(SessionException): + await registry.invoke_callbacks_async(BidiAgentStopEvent(agent=agent)) + + manager.sync_agent.assert_called_once_with(agent) + assert manager.create_message.call_count == 2 + manager._flush_messages.assert_called_once() + assert not manager._held_messages + + class TestIntervalFlush: """Test interval-based flush mechanism for long-running agents.""" diff --git a/uv.lock b/uv.lock index 2af3e035..46b5ef60 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.56.0" }, + { name = "strands-agents", marker = "extra == 'strands-agents'", specifier = ">=1.57.1" }, { 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.56.0" }, + { name = "strands-agents", specifier = ">=1.57.1" }, { name = "strands-agents-evals", specifier = ">=1.0.3,<2.0.0" }, { name = "websockets", specifier = ">=14.1" }, { name = "wheel", specifier = ">=0.45.1" }, @@ -5043,7 +5043,7 @@ wheels = [ [[package]] name = "strands-agents" -version = "1.56.0" +version = "1.57.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "boto3" }, @@ -5060,9 +5060,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "watchdog" }, ] -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" } +sdist = { url = "https://files.pythonhosted.org/packages/6d/2c/061ba1354b8586e8bd6edb42b2f8de3c6a4db83488b6b28d82cbc88a9185/strands_agents-1.57.1.tar.gz", hash = "sha256:109ea0a0f7da82a3b32e4c02cef12f53d46637d3e1e0fb8a285c5fa7116eb1d8", size = 1709312, upload-time = "2026-09-25T20:12:35.457Z" } wheels = [ - { 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" }, + { url = "https://files.pythonhosted.org/packages/8c/1b/2b7e12fafee0e6ab6ce9fa16a0ad850bb570d8e2bd6326f5cca7c2cd75a7/strands_agents-1.57.1-py3-none-any.whl", hash = "sha256:797ea3de2ba4a66aaf7e804cba8308e3d582e2412c59978433d3216973d02a9d", size = 858212, upload-time = "2026-09-25T20:12:33.361Z" }, ] [[package]]