defensive lock

Author: fcenedes

agent-memory-client/agent_memory_client/models.py

@@ -6,6 +6,7 @@
 """
 
 import logging
+import threading
 from datetime import datetime, timedelta, timezone
 from enum import Enum
 from typing import Any, ClassVar, Literal
@@ -68,6 +69,7 @@ class MemoryMessage(BaseModel):
     # Track message IDs that have been warned (in-memory, per-process)
     # Used to rate-limit deprecation warnings
     _warned_message_ids: ClassVar[set[str]] = set()
+    _warned_message_ids_lock: ClassVar[threading.Lock] = threading.Lock()
     _max_warned_ids: ClassVar[int] = 10000  # Prevent unbounded growth
 
     # Default tolerance for future timestamp validation (5 minutes)
@@ -106,15 +108,20 @@ def validate_created_at(cls, data: Any) -> Any:
         created_at_provided = "created_at" in data and data["created_at"] is not None
 
         if not created_at_provided:
-            # Rate-limit warnings by message ID
+            # Rate-limit warnings by message ID (thread-safe)
             msg_id = data.get("id", "unknown")
 
-            if msg_id not in cls._warned_message_ids:
-                # Prevent unbounded memory growth
-                if len(cls._warned_message_ids) >= cls._max_warned_ids:
-                    cls._warned_message_ids.clear()
-                cls._warned_message_ids.add(msg_id)
-
+            with cls._warned_message_ids_lock:
+                if msg_id not in cls._warned_message_ids:
+                    # Prevent unbounded memory growth
+                    if len(cls._warned_message_ids) >= cls._max_warned_ids:
+                        cls._warned_message_ids.clear()
+                    cls._warned_message_ids.add(msg_id)
+                    should_warn = True
+                else:
+                    should_warn = False
+
+            if should_warn:
                 logger.warning(
                     "MemoryMessage created without explicit created_at timestamp. "
                     "This will become required in a future version. "

agent_memory_server/models.py

@@ -1,4 +1,5 @@
 import logging
+import threading
 from collections.abc import Callable
 from datetime import UTC, datetime, timedelta
 from enum import Enum
@@ -85,8 +86,12 @@ class MemoryMessage(BaseModel):
     # Track message IDs that have been warned (in-memory, per-worker)
     # Used to rate-limit deprecation warnings
     _warned_message_ids: ClassVar[set[str]] = set()
+    _warned_message_ids_lock: ClassVar[threading.Lock] = threading.Lock()
     _max_warned_ids: ClassVar[int] = 10000  # Prevent unbounded growth
 
+    # Thread-local storage for passing created_at_provided flag from validator to model_post_init
+    _created_at_provided_thread_local: ClassVar[threading.local] = threading.local()
+
     role: str
     content: str
     id: str = Field(
@@ -109,11 +114,15 @@ class MemoryMessage(BaseModel):
     # Used for deprecation header in API responses
     _created_at_was_provided: bool = PrivateAttr(default=False)
 
-    def __init__(self, **data):
-        # Check if created_at was provided before calling super().__init__
-        created_at_provided = "created_at" in data and data["created_at"] is not None
-        super().__init__(**data)
-        self._created_at_was_provided = created_at_provided
+    def model_post_init(self, __context: Any) -> None:
+        """Set _created_at_was_provided from thread-local storage after model is constructed."""
+        # Retrieve the flag from thread-local storage (set by validator)
+        self._created_at_was_provided = getattr(
+            self._created_at_provided_thread_local, "value", False
+        )
+        # Clean up thread-local storage
+        if hasattr(self._created_at_provided_thread_local, "value"):
+            del self._created_at_provided_thread_local.value
 
     @model_validator(mode="before")
     @classmethod
@@ -130,22 +139,30 @@ def validate_created_at(cls, data: Any) -> Any:
 
         created_at_provided = "created_at" in data and data["created_at"] is not None
 
+        # Store in thread-local for model_post_init to pick up
+        cls._created_at_provided_thread_local.value = created_at_provided
+
         if not created_at_provided:
             # Handle missing created_at
             if settings.require_message_timestamps:
                 raise ValueError(
                     "created_at is required for messages. "
                     "Please provide the timestamp when the message was created."
                 )
-            # Rate-limit warnings by message ID
+            # Rate-limit warnings by message ID (thread-safe)
             msg_id = data.get("id", "unknown")
 
-            if msg_id not in cls._warned_message_ids:
-                # Prevent unbounded memory growth
-                if len(cls._warned_message_ids) >= cls._max_warned_ids:
-                    cls._warned_message_ids.clear()
-                cls._warned_message_ids.add(msg_id)
+            with cls._warned_message_ids_lock:
+                if msg_id not in cls._warned_message_ids:
+                    # Prevent unbounded memory growth
+                    if len(cls._warned_message_ids) >= cls._max_warned_ids:
+                        cls._warned_message_ids.clear()
+                    cls._warned_message_ids.add(msg_id)
+                    should_warn = True
+                else:
+                    should_warn = False
 
+            if should_warn:
                 logger.warning(
                     "MemoryMessage created without explicit created_at timestamp. "
                     "This will become required in a future version. "

tests/test_models.py

@@ -256,7 +256,8 @@ def test_message_with_future_timestamp_rejected(self):
                 created_at=future_time,
             )
 
-        assert "cannot be in the future" in str(exc_info.value)
+        assert "cannot be more than" in str(exc_info.value)
+        assert "seconds in the future" in str(exc_info.value)
 
     def test_message_with_near_future_timestamp_allowed(self):
         """Test that timestamps within tolerance are allowed"""