refactor: centralize in-memory store locking

src/agentmsgs/stores/memory.py

@@ -4,10 +4,20 @@ import datetime
 import threading
 import uuid
 from collections import defaultdict
+from functools import wraps
 
 from agentmsgs.core.types import Agent, Message, Thread, ThreadInfo, ThreadReadState
 
 
+def _lock(method):
+    @wraps(method)
+    def wrapped(self, *args, **kwargs):
+        with self._lock:
+            return method(self, *args, **kwargs)
+
+    return wrapped
+
+
 class InMemoryStore:
     def __init__(self) -> None:
         self._agents: dict[uuid.UUID, Agent] = {}
@@ -17,14 +27,14 @@ class InMemoryStore:
         self._cursors: dict[tuple[uuid.UUID, uuid.UUID], int] = {}
         self._lock = threading.RLock()
 
+    @_lock
     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)
 
+    @_lock
     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)
@@ -32,19 +42,19 @@ class InMemoryStore:
         self._name_idx[name] = a.id
         return a
 
+    @_lock
     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)
 
+    @_lock
     def get_agent_by_id(self, id: uuid.UUID) -> Agent | None:
-        with self._lock:
         return self._agents.get(id)
 
+    @_lock
     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)
@@ -56,12 +66,12 @@ class InMemoryStore:
                         created_at=t.created_at,
                     )
 
+    @_lock
     def list_agents(self) -> list[Agent]:
-        with self._lock:
         return list(self._agents.values())
 
+    @_lock
     def create_thread(self, participants: set[Agent]) -> Thread:
-        with self._lock:
         if len(participants) < 2:
             raise ValueError("need >=2 participants")
         t = Thread(
@@ -72,12 +82,12 @@ class InMemoryStore:
         self._threads[t.id] = t
         return t
 
+    @_lock
     def get_thread(self, id: uuid.UUID) -> Thread | None:
-        with self._lock:
         return self._threads.get(id)
 
+    @_lock
     def get_thread_info(self, id: uuid.UUID) -> ThreadInfo | None:
-        with self._lock:
         thread = self.get_thread(id)
         if thread is None:
             return None
@@ -88,10 +98,10 @@ class InMemoryStore:
             latest_message=messages[-1] if messages else None,
         )
 
+    @_lock
     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, [])
@@ -107,13 +117,13 @@ class InMemoryStore:
             ),
         )
 
+    @_lock
     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]
 
+    @_lock
     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
@@ -121,8 +131,8 @@ class InMemoryStore:
         self._threads[thread_id] = nt
         return nt
 
+    @_lock
     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
@@ -130,19 +140,18 @@ class InMemoryStore:
         self._threads[thread_id] = nt
         return nt
 
+    @_lock
     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
         for k in list(self._cursors.keys()):
             if k[0] == id:
                 self._cursors.pop(k, None)
 
+    @_lock
     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:
@@ -158,14 +167,14 @@ class InMemoryStore:
         self._msgs[thread_id].append(m)
         return m
 
+    @_lock
     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]
 
+    @_lock
     def get_cursor(self, thread_id: uuid.UUID, agent: Agent) -> int:
-        with self._lock:
         return self._cursors.get((thread_id, agent.id), 0)
 
+    @_lock
     def set_cursor(self, thread_id: uuid.UUID, agent: Agent, seq: int) -> None:
-        with self._lock:
         self._cursors[(thread_id, agent.id)] = seq