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,138 +15,157 @@ 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:
-        if name in self._name_idx:
-            return self._agents[self._name_idx[name]]
-        return self.create_agent(name)
+        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:
-        if name in self._name_idx:
-            raise ValueError(f"agent name taken: {name}")
-        a = Agent(id=uuid.uuid4(), name=name)
-        self._agents[a.id] = a
-        self._name_idx[name] = a.id
-        return a
+        with self._lock:
+            if name in self._name_idx:
+                raise ValueError(f"agent name taken: {name}")
+            a = Agent(id=uuid.uuid4(), name=name)
+            self._agents[a.id] = a
+            self._name_idx[name] = a.id
+            return a
 
     def get_agent_by_name(self, name: str) -> Agent | None:
-        uid = self._name_idx.get(name)
-        if uid is None:
-            return None
-        return self._agents.get(uid)
+        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:
-        return self._agents.get(id)
+        with self._lock:
+            return self._agents.get(id)
 
     def delete_agent(self, id: uuid.UUID) -> None:
-        a = self._agents.pop(id, None)
-        if a is not None:
-            self._name_idx.pop(a.name, None)
-            for tid, t in list(self._threads.items()):
-                if a in t.participants:
-                    self._threads[tid] = Thread(
-                        id=t.id,
-                        participants=t.participants - {a},
-                        created_at=t.created_at,
-                    )
+        with self._lock:
+            a = self._agents.pop(id, None)
+            if a is not None:
+                self._name_idx.pop(a.name, None)
+                for tid, t in list(self._threads.items()):
+                    if a in t.participants:
+                        self._threads[tid] = Thread(
+                            id=t.id,
+                            participants=t.participants - {a},
+                            created_at=t.created_at,
+                        )
 
     def list_agents(self) -> list[Agent]:
-        return list(self._agents.values())
+        with self._lock:
+            return list(self._agents.values())
 
     def create_thread(self, participants: set[Agent]) -> Thread:
-        if len(participants) < 2:
-            raise ValueError("need >=2 participants")
-        t = Thread(
-            id=uuid.uuid4(),
-            participants=frozenset(participants),
-            created_at=datetime.datetime.now(datetime.UTC),
-        )
-        self._threads[t.id] = t
-        return t
+        with self._lock:
+            if len(participants) < 2:
+                raise ValueError("need >=2 participants")
+            t = Thread(
+                id=uuid.uuid4(),
+                participants=frozenset(participants),
+                created_at=datetime.datetime.now(datetime.UTC),
+            )
+            self._threads[t.id] = t
+            return t
 
     def get_thread(self, id: uuid.UUID) -> Thread | None:
-        return self._threads.get(id)
+        with self._lock:
+            return self._threads.get(id)
 
     def get_thread_info(self, id: uuid.UUID) -> ThreadInfo | None:
-        thread = self.get_thread(id)
-        if thread is None:
-            return None
-        messages = self._msgs.get(id, [])
-        return ThreadInfo(
-            thread=thread,
-            message_count=len(messages),
-            latest_message=messages[-1] if messages else None,
-        )
+        with self._lock:
+            thread = self.get_thread(id)
+            if thread is None:
+                return None
+            messages = self._msgs.get(id, [])
+            return ThreadInfo(
+                thread=thread,
+                message_count=len(messages),
+                latest_message=messages[-1] if messages else None,
+            )
 
     def get_thread_read_state(
         self, id: uuid.UUID, agent: Agent
     ) -> ThreadReadState | None:
-        if self.get_thread(id) is None:
-            return None
-        messages = self._msgs.get(id, [])
-        cursor_seq = self.get_cursor(id, agent)
-        unread = [message for message in messages if message.seq > cursor_seq]
-        return ThreadReadState(
-            agent=agent,
-            cursor_seq=cursor_seq,
-            latest_seq=messages[-1].seq if messages else 0,
-            unread_count=len(unread),
-            unread_from_others_count=sum(
-                message.sender.id != agent.id for message in unread
-            ),
-        )
+        with self._lock:
+            if self.get_thread(id) is None:
+                return None
+            messages = self._msgs.get(id, [])
+            cursor_seq = self.get_cursor(id, agent)
+            unread = [message for message in messages if message.seq > cursor_seq]
+            return ThreadReadState(
+                agent=agent,
+                cursor_seq=cursor_seq,
+                latest_seq=messages[-1].seq if messages else 0,
+                unread_count=len(unread),
+                unread_from_others_count=sum(
+                    message.sender.id != agent.id for message in unread
+                ),
+            )
 
     def find_threads(self, containing: set[Agent]) -> list[Thread]:
-        c = set(containing)
-        return [t for t in self._threads.values() if c <= t.participants]
+        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:
-        t = self._threads[thread_id]
-        nt = Thread(
-            id=t.id, participants=t.participants | {agent}, created_at=t.created_at
-        )
-        self._threads[thread_id] = nt
-        return nt
+        with self._lock:
+            t = self._threads[thread_id]
+            nt = Thread(
+                id=t.id, participants=t.participants | {agent}, created_at=t.created_at
+            )
+            self._threads[thread_id] = nt
+            return nt
 
     def remove_participant(self, thread_id: uuid.UUID, agent: Agent) -> Thread:
-        t = self._threads[thread_id]
-        nt = Thread(
-            id=t.id, participants=t.participants - {agent}, created_at=t.created_at
-        )
-        self._threads[thread_id] = nt
-        return nt
+        with self._lock:
+            t = self._threads[thread_id]
+            nt = Thread(
+                id=t.id, participants=t.participants - {agent}, created_at=t.created_at
+            )
+            self._threads[thread_id] = nt
+            return nt
 
     def delete_thread(self, id: uuid.UUID) -> None:
-        self._threads.pop(id, None)
-        self._msgs.pop(id, None)
-        # also delete cursors for that thread
-        for k in list(self._cursors.keys()):
-            if k[0] == id:
-                self._cursors.pop(k, None)
+        with self._lock:
+            self._threads.pop(id, None)
+            self._msgs.pop(id, None)
+            # also delete cursors for that thread
+            for k in list(self._cursors.keys()):
+                if k[0] == id:
+                    self._cursors.pop(k, None)
 
     def append_message(
         self, thread_id: uuid.UUID, sender: Agent, content: str
     ) -> Message:
-        if thread_id not in self._threads:
-            raise KeyError(thread_id)
-        if sender not in self._threads[thread_id].participants:
-            raise ValueError("sender not in thread")
-        seq = len(self._msgs[thread_id]) + 1
-        m = Message(
-            id=uuid.uuid4(),
-            sender=sender,
-            content=content,
-            seq=seq,
-            ts=datetime.datetime.now(datetime.UTC),
-        )
-        self._msgs[thread_id].append(m)
-        return m
+        with self._lock:
+            if thread_id not in self._threads:
+                raise KeyError(thread_id)
+            if sender not in self._threads[thread_id].participants:
+                raise ValueError("sender not in thread")
+            seq = len(self._msgs[thread_id]) + 1
+            m = Message(
+                id=uuid.uuid4(),
+                sender=sender,
+                content=content,
+                seq=seq,
+                ts=datetime.datetime.now(datetime.UTC),
+            )
+            self._msgs[thread_id].append(m)
+            return m
 
     def list_messages(self, thread_id: uuid.UUID, after_seq: int = 0) -> list[Message]:
-        return [m for m in self._msgs.get(thread_id, []) if m.seq > after_seq]
+        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:
-        return self._cursors.get((thread_id, agent.id), 0)
+        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:
-        self._cursors[(thread_id, agent.id)] = seq
+        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]