fix: synchronize in-memory store access

src/agentmsgs/stores/memory.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import datetime
+import threading
 import uuid
 from collections import defaultdict
 
@@ -14,13 +15,16 @@ class InMemoryStore:
         self._threads: dict[uuid.UUID, Thread] = {}
         self._msgs: dict[uuid.UUID, list[Message]] = defaultdict(list)
         self._cursors: dict[tuple[uuid.UUID, uuid.UUID], int] = {}
+        self._lock = threading.RLock()
 
     def get_or_create_agent(self, name: str) -> Agent:
+        with self._lock:
             if name in self._name_idx:
                 return self._agents[self._name_idx[name]]
             return self.create_agent(name)
 
     def create_agent(self, name: str) -> Agent:
+        with self._lock:
             if name in self._name_idx:
                 raise ValueError(f"agent name taken: {name}")
             a = Agent(id=uuid.uuid4(), name=name)
@@ -29,15 +33,18 @@ class InMemoryStore:
             return a
 
     def get_agent_by_name(self, name: str) -> Agent | None:
+        with self._lock:
             uid = self._name_idx.get(name)
             if uid is None:
                 return None
             return self._agents.get(uid)
 
     def get_agent_by_id(self, id: uuid.UUID) -> Agent | None:
+        with self._lock:
             return self._agents.get(id)
 
     def delete_agent(self, id: uuid.UUID) -> None:
+        with self._lock:
             a = self._agents.pop(id, None)
             if a is not None:
                 self._name_idx.pop(a.name, None)
@@ -50,9 +57,11 @@ class InMemoryStore:
                         )
 
     def list_agents(self) -> list[Agent]:
+        with self._lock:
             return list(self._agents.values())
 
     def create_thread(self, participants: set[Agent]) -> Thread:
+        with self._lock:
             if len(participants) < 2:
                 raise ValueError("need >=2 participants")
             t = Thread(
@@ -64,9 +73,11 @@ class InMemoryStore:
             return t
 
     def get_thread(self, id: uuid.UUID) -> Thread | None:
+        with self._lock:
             return self._threads.get(id)
 
     def get_thread_info(self, id: uuid.UUID) -> ThreadInfo | None:
+        with self._lock:
             thread = self.get_thread(id)
             if thread is None:
                 return None
@@ -80,6 +91,7 @@ class InMemoryStore:
     def get_thread_read_state(
         self, id: uuid.UUID, agent: Agent
     ) -> ThreadReadState | None:
+        with self._lock:
             if self.get_thread(id) is None:
                 return None
             messages = self._msgs.get(id, [])
@@ -96,10 +108,12 @@ class InMemoryStore:
             )
 
     def find_threads(self, containing: set[Agent]) -> list[Thread]:
+        with self._lock:
             c = set(containing)
             return [t for t in self._threads.values() if c <= t.participants]
 
     def add_participant(self, thread_id: uuid.UUID, agent: Agent) -> Thread:
+        with self._lock:
             t = self._threads[thread_id]
             nt = Thread(
                 id=t.id, participants=t.participants | {agent}, created_at=t.created_at
@@ -108,6 +122,7 @@ class InMemoryStore:
             return nt
 
     def remove_participant(self, thread_id: uuid.UUID, agent: Agent) -> Thread:
+        with self._lock:
             t = self._threads[thread_id]
             nt = Thread(
                 id=t.id, participants=t.participants - {agent}, created_at=t.created_at
@@ -116,6 +131,7 @@ class InMemoryStore:
             return nt
 
     def delete_thread(self, id: uuid.UUID) -> None:
+        with self._lock:
             self._threads.pop(id, None)
             self._msgs.pop(id, None)
             # also delete cursors for that thread
@@ -126,6 +142,7 @@ class InMemoryStore:
     def append_message(
         self, thread_id: uuid.UUID, sender: Agent, content: str
     ) -> Message:
+        with self._lock:
             if thread_id not in self._threads:
                 raise KeyError(thread_id)
             if sender not in self._threads[thread_id].participants:
@@ -142,10 +159,13 @@ class InMemoryStore:
             return m
 
     def list_messages(self, thread_id: uuid.UUID, after_seq: int = 0) -> list[Message]:
+        with self._lock:
             return [m for m in self._msgs.get(thread_id, []) if m.seq > after_seq]
 
     def get_cursor(self, thread_id: uuid.UUID, agent: Agent) -> int:
+        with self._lock:
             return self._cursors.get((thread_id, agent.id), 0)
 
     def set_cursor(self, thread_id: uuid.UUID, agent: Agent, seq: int) -> None:
+        with self._lock:
             self._cursors[(thread_id, agent.id)] = seq

tests/test_store_memory.py

@@ -1,4 +1,7 @@
+import threading
+import time
 import uuid
+from concurrent.futures import ThreadPoolExecutor
 
 from agentmsgs.stores.memory import InMemoryStore
 
@@ -86,3 +89,57 @@ def test_memory_thread_info_for_empty_thread():
     assert state.unread_from_others_count == 0
     assert s.get_thread_info(uuid.uuid4()) is None
     assert s.get_thread_read_state(uuid.uuid4(), a) is None
+
+
+def test_memory_concurrent_get_or_create_agent_is_idempotent():
+    store = InMemoryStore()
+    start = threading.Barrier(2)
+
+    class SlowNameIndex(dict):
+        def __contains__(self, key):
+            present = super().__contains__(key)
+            if not present:
+                time.sleep(0.01)
+            return present
+
+    store._name_idx = SlowNameIndex()
+
+    def create_agent():
+        start.wait()
+        return store.get_or_create_agent("Alice")
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        agents = list(executor.map(lambda _: create_agent(), range(2)))
+
+    assert len({agent.id for agent in agents}) == 1
+    assert len(store.list_agents()) == 1
+
+
+def test_memory_concurrent_append_assigns_unique_sequences():
+    store = InMemoryStore()
+    alice = store.get_or_create_agent("Alice")
+    bob = store.get_or_create_agent("Bob")
+    thread = store.create_thread({alice, bob})
+    gate = threading.Barrier(2, timeout=0.2)
+    local = threading.local()
+
+    class SlowMessageList(list):
+        def __len__(self):
+            length = super().__len__()
+            if not hasattr(local, "waited"):
+                try:
+                    gate.wait()
+                except threading.BrokenBarrierError:
+                    pass
+                local.waited = True
+            return length
+
+    store._msgs[thread.id] = SlowMessageList()
+
+    def append_message(index):
+        return store.append_message(thread.id, alice, str(index))
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        messages = list(executor.map(append_message, range(2)))
+
+    assert sorted(message.seq for message in messages) == [1, 2]