Fix: Add thread locks to memory storage classes to prevent race conditions

This commit introduces threading locks to `ShortTermMemory`, `MidTermMemory`, and `LongTermMemory` classes to ensure thread safety during file write operations.

Problem:
- In multi-threaded environments (e.g., when running as an MCP server or handling concurrent requests), multiple threads might attempt to write to the JSON storage files simultaneously.
- This race condition can lead to data corruption, lost memory entries, or file truncation.

Solution:
- Added `self.lock = threading.Lock()` to the `__init__` method of each memory class.
- Wrapped all `json.dump` operations within `save()` methods using a `with self.lock:` context manager.
- This ensures that file writes are atomic and serialized, preventing concurrent write conflicts.

Changes:
- Modified `memoryos-pypi/short_term.py`
- Modified `memoryos-pypi/mid_term.py`
- Modified `memoryos-pypi/long_term.py`

Author: OmniJax

memoryos-pypi/long_term.py

@@ -2,6 +2,7 @@
 import numpy as np
 import faiss
 from collections import deque
+import threading
 try:
     from .utils import get_timestamp, get_embedding, normalize_vector, ensure_directory_exists
 except ImportError:
@@ -19,6 +20,7 @@ def __init__(self, file_path, knowledge_capacity=100, embedding_model_name: str
 
         self.embedding_model_name = embedding_model_name
         self.embedding_model_kwargs = embedding_model_kwargs if embedding_model_kwargs is not None else {}
+        self.lock = threading.Lock()
         self.load()
 
     def update_user_profile(self, user_id, new_data, merge=True):
@@ -144,8 +146,9 @@ def save(self):
             "assistant_knowledge": list(self.assistant_knowledge)
         }
         try:
-            with open(self.file_path, "w", encoding="utf-8") as f:
-                json.dump(data, f, ensure_ascii=False, indent=2)
+            with self.lock:
+                with open(self.file_path, "w", encoding="utf-8") as f:
+                    json.dump(data, f, ensure_ascii=False, indent=2)
         except IOError as e:
             print(f"Error saving LongTermMemory to {self.file_path}: {e}")
 

memoryos-pypi/mid_term.py

@@ -3,6 +3,7 @@
 from collections import defaultdict
 import faiss
 import heapq
+import threading
 from datetime import datetime
 
 try:
@@ -46,6 +47,7 @@ def __init__(self, file_path: str, client: OpenAIClient, max_capacity=2000, embe
 
         self.embedding_model_name = embedding_model_name
         self.embedding_model_kwargs = embedding_model_kwargs if embedding_model_kwargs is not None else {}
+        self.lock = threading.Lock()
         self.load()
 
     def get_page_by_id(self, page_id):
@@ -370,8 +372,9 @@ def save(self):
             # "heap_snapshot": self.heap 
         }
         try:
-            with open(self.file_path, "w", encoding="utf-8") as f:
-                json.dump(data_to_save, f, ensure_ascii=False, indent=2)
+            with self.lock:
+                with open(self.file_path, "w", encoding="utf-8") as f:
+                    json.dump(data_to_save, f, ensure_ascii=False, indent=2)
         except IOError as e:
             print(f"Error saving MidTermMemory to {self.file_path}: {e}")
 

memoryos-pypi/short_term.py

@@ -1,5 +1,6 @@
 import json
 from collections import deque
+import threading
 try:
     from .utils import get_timestamp, ensure_directory_exists
 except ImportError:
@@ -11,6 +12,7 @@ def __init__(self, file_path, max_capacity=10):
         self.file_path = file_path
         ensure_directory_exists(self.file_path)
         self.memory = deque(maxlen=max_capacity)
+        self.lock = threading.Lock()
         self.load()
 
     def add_qa_pair(self, qa_pair):
@@ -38,8 +40,9 @@ def pop_oldest(self):
 
     def save(self):
         try:
-            with open(self.file_path, "w", encoding="utf-8") as f:
-                json.dump(list(self.memory), f, ensure_ascii=False, indent=2)
+            with self.lock:
+                with open(self.file_path, "w", encoding="utf-8") as f:
+                    json.dump(list(self.memory), f, ensure_ascii=False, indent=2)
         except IOError as e:
             print(f"Error saving ShortTermMemory to {self.file_path}: {e}")