diff --git a/src/agentmsgs/stores/memory.py b/src/agentmsgs/stores/memory.py index 773558a..380e395 100644 --- a/src/agentmsgs/stores/memory.py +++ b/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 diff --git a/tests/test_store_memory.py b/tests/test_store_memory.py index b2b40be..2ce8e90 100644 --- a/tests/test_store_memory.py +++ b/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]