diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index 8ff21e28..4677a131 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -98,13 +98,11 @@ class AgentCoreMemorySessionManager(RepositorySessionManager, SessionRepository) - Consistent with existing Strands Session managers (such as: FileSessionManager, S3SessionManager) """ - # Class-level timestamp tracking for monotonic ordering - _timestamp_lock = threading.Lock() - _last_timestamp: Optional[datetime] = None + def _get_monotonic_timestamp(self, desired_timestamp: Optional[datetime] = None) -> datetime: + """Get a monotonically increasing timestamp for this session. - @classmethod - def _get_monotonic_timestamp(cls, desired_timestamp: Optional[datetime] = None) -> datetime: - """Get a monotonically increasing timestamp. + Ties are broken at millisecond granularity, which is the resolution + AgentCore Memory stores and orders ``eventTimestamp`` at. Args: desired_timestamp (Optional[datetime]): The desired timestamp. If None, uses current time. @@ -115,20 +113,12 @@ def _get_monotonic_timestamp(cls, desired_timestamp: Optional[datetime] = None) if desired_timestamp is None: desired_timestamp = datetime.now(timezone.utc) - with cls._timestamp_lock: - if cls._last_timestamp is None: - cls._last_timestamp = desired_timestamp - return desired_timestamp - - # Why the 1 second check? Because Boto3 does NOT support sub 1 second resolution. - if desired_timestamp <= cls._last_timestamp + timedelta(seconds=1): - # Increment by 1 second to ensure ordering - new_timestamp = cls._last_timestamp + timedelta(seconds=1) - else: - new_timestamp = desired_timestamp - - cls._last_timestamp = new_timestamp - return new_timestamp + with self._timestamp_lock: + if self._last_timestamp is not None and desired_timestamp <= self._last_timestamp: + # Break the tie at millisecond granularity (the service's resolution). + desired_timestamp = self._last_timestamp + timedelta(milliseconds=1) + self._last_timestamp = desired_timestamp + return desired_timestamp def __init__( self, @@ -159,6 +149,12 @@ def __init__( session = boto_session or boto3.Session(region_name=region_name) self.has_existing_agent = False + # Instance-scoped monotonic-timestamp state. Per-instance so concurrent + # managers for different sessions in one process do not perturb each + # other's ordering. + self._timestamp_lock = threading.Lock() + self._last_timestamp: Optional[datetime] = None + # Batching support - stores pre-processed messages self._message_buffer: list[BufferedMessage] = [] self._message_lock = threading.Lock() 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 fc93984f..50a04217 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 @@ -4,7 +4,7 @@ import inspect import logging import time -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from unittest.mock import Mock, patch import pytest @@ -3810,3 +3810,80 @@ def test_flush_agent_states_failure_restores_buffer(self, batching_session_manag # State must be back in the buffer for retry. assert batching_session_manager.pending_agent_state_count() == 1 + + +class TestMonotonicTimestamp: + """Tests for monotonic event-timestamp ordering across processes/pods. + + Regression coverage for the multi-process interleave bug: the ordering + counter must be instance-scoped, break ties at millisecond (not second) + granularity, and be seeded from the newest persisted event so a freshly + started process continues after another process's writes. + """ + + def test_state_is_instance_scoped_not_class_level(self, agentcore_config, mock_memory_client): + """Two managers must not share timestamp state (class-level state leaked + across unrelated sessions in the same process).""" + m1 = _create_session_manager(agentcore_config, mock_memory_client) + m2 = _create_session_manager(agentcore_config, mock_memory_client) + + assert m1._last_timestamp is None + assert m2._last_timestamp is None + + base = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + m1._get_monotonic_timestamp(base) + + # m1 advanced; m2 must be untouched. + assert m1._last_timestamp == base + assert m2._last_timestamp is None + # Distinct lock objects, not a shared class attribute. + assert m1._timestamp_lock is not m2._timestamp_lock + + def test_first_timestamp_passes_through_unchanged(self, session_manager): + """With no prior events, the desired timestamp is returned as-is.""" + base = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + assert session_manager._get_monotonic_timestamp(base) == base + + def test_tie_broken_by_one_millisecond_not_one_second(self, session_manager): + """A colliding timestamp is bumped by 1ms — not inflated by 1s like the + old behavior that pushed events seconds into the future.""" + base = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + session_manager._get_monotonic_timestamp(base) + + # Same timestamp again -> must advance by exactly 1ms. + second = session_manager._get_monotonic_timestamp(base) + assert second == base + timedelta(milliseconds=1) + + # Third identical -> another 1ms. + third = session_manager._get_monotonic_timestamp(base) + assert third == base + timedelta(milliseconds=2) + + # A full multi-event turn stays within a few ms of real time, not seconds. + assert third - base < timedelta(seconds=1) + + def test_later_timestamp_is_not_bumped(self, session_manager): + """A desired timestamp clearly after the floor passes through unchanged.""" + base = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + session_manager._get_monotonic_timestamp(base) + + later = base + timedelta(seconds=5) + assert session_manager._get_monotonic_timestamp(later) == later + + def test_none_desired_uses_current_time(self, session_manager): + """Passing None falls back to current UTC time (preserved behavior).""" + before = datetime.now(timezone.utc) + result = session_manager._get_monotonic_timestamp(None) + after = datetime.now(timezone.utc) + assert before <= result <= after + + def test_within_process_burst_stays_ordered_at_ms_resolution(self, session_manager): + """A burst of same-instant events gets strictly increasing 1ms-spaced + timestamps instead of being inflated by whole seconds.""" + base = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + # Five events all requesting the same instant (a same-second burst). + stamps = [session_manager._get_monotonic_timestamp(base) for _ in range(5)] + + assert stamps == sorted(stamps) + assert len(set(stamps)) == len(stamps) # no ties (ambiguous ordering) + # Whole burst stays within a few ms of the requested time, not seconds. + assert stamps[-1] - base == timedelta(milliseconds=4)