feat: override the timestamp serializer of the _data_plane_client to support millisecond granularity

Author: Jonathan James

src/bedrock_agentcore/memory/session.py

@@ -110,47 +110,98 @@ def __init__(
         Raises:
             ValueError: If region_name parameter conflicts with boto3_session region.
         """
+        # Initialize core attributes
+        self._memory_id = memory_id
+
+        # Setup session and validate region consistency
         session = boto3_session if boto3_session else boto3.Session()
+        self.region_name = self._validate_and_resolve_region(region_name, session)
+
+        # Configure and create boto3 client
+        client_config = self._build_client_config(boto_client_config)
+        self._data_plane_client = session.client(
+            "bedrock-agentcore", region_name=self.region_name, config=client_config
+        )
+
+        # Configure timestamp serialization for microsecond precision
+        self._configure_timestamp_serialization()
+
+        # Define allowed data plane methods
+        self._ALLOWED_DATA_PLANE_METHODS = {
+            "retrieve_memory_records",
+            "get_memory_record",
+            "delete_memory_record",
+            "list_memory_records",
+            "create_event",
+            "get_event",
+            "delete_event",
+            "list_events",
+        }
+
+    def _validate_and_resolve_region(self, region_name: Optional[str], session: boto3.Session) -> str:
+        """Validate region consistency and resolve the final region to use.
+
+        Args:
+            region_name: Explicitly provided region name
+            session: Boto3 session instance
+
+        Returns:
+            The resolved region name to use
+
+        Raises:
+            ValueError: If region_name conflicts with session region
+        """
         session_region = session.region_name
 
         # Validate region consistency if both are provided
-        if region_name and boto3_session and session_region and region_name != session_region:
+        if region_name and session_region and isinstance(session_region, str) and region_name != session_region:
             raise ValueError(
                 f"Region mismatch: provided region_name '{region_name}' does not match "
                 f"boto3_session region '{session_region}'. Please ensure both "
                 f"parameters specify the same region or omit the region_name parameter "
                 f"to use the session's region."
             )
 
-        # Configure boto3 client with merged configuration
+        return region_name or session_region
+
+    def _build_client_config(self, boto_client_config: Optional[BotocoreConfig]) -> BotocoreConfig:
+        """Build the final boto3 client configuration with SDK user agent.
+
+        Args:
+            boto_client_config: Optional user-provided client configuration
+
+        Returns:
+            Final client configuration with SDK user agent
+        """
+        sdk_user_agent = "bedrock-agentcore-sdk"
+
         if boto_client_config:
             existing_user_agent = getattr(boto_client_config, "user_agent_extra", None)
             if existing_user_agent:
-                new_user_agent = f"{existing_user_agent} bedrock-agentcore-sdk"
+                new_user_agent = f"{existing_user_agent} {sdk_user_agent}"
             else:
-                new_user_agent = "bedrock-agentcore-sdk"
-            client_config = boto_client_config.merge(BotocoreConfig(user_agent_extra=new_user_agent))
+                new_user_agent = sdk_user_agent
+            return boto_client_config.merge(BotocoreConfig(user_agent_extra=new_user_agent))
         else:
-            client_config = BotocoreConfig(user_agent_extra="bedrock-agentcore-sdk")
+            return BotocoreConfig(user_agent_extra=sdk_user_agent)
 
-        # Use provided region_name or fall back to session region
-        self.region_name = region_name or session_region
-        self._memory_id = memory_id
-        self._data_plane_client = session.client(
-            "bedrock-agentcore", region_name=self.region_name, config=client_config
-        )
+    def _configure_timestamp_serialization(self) -> None:
+        """Configure the boto3 client to serialize timestamps with microsecond precision.
 
-        # AgentCore Memory data plane methods
-        self._ALLOWED_DATA_PLANE_METHODS = {
-            "retrieve_memory_records",
-            "get_memory_record",
-            "delete_memory_record",
-            "list_memory_records",
-            "create_event",
-            "get_event",
-            "delete_event",
-            "list_events",
-        }
+        This method overrides the default timestamp serialization to preserve microsecond
+        precision when sending datetime objects to the AgentCore Memory service.
+        """
+        original_serialize_timestamp = self._data_plane_client._serializer._serializer._serialize_type_timestamp
+
+        def serialize_timestamp_with_microseconds(serialized, value, shape, name):
+            if isinstance(value, datetime):
+                serialized[name] = value.timestamp()  # Float with microseconds
+            else:
+                original_serialize_timestamp(serialized, value, shape, name)
+
+        self._data_plane_client._serializer._serializer._serialize_type_timestamp = (
+            serialize_timestamp_with_microseconds
+        )
 
     def __getattr__(self, name: str):
         """Dynamically forward method calls to the appropriate boto3 client.