refactor: centralize in-memory store locking
This commit is contained in:
1 file changed
+122
-113
+122
-113
@@ -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,155 +27,154 @@ 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)
|
||||
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)
|
||||
self._agents[a.id] = a
|
||||
self._name_idx[name] = a.id
|
||||
return a
|
||||
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
|
||||
|
||||
@_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)
|
||||
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)
|
||||
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)
|
||||
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,
|
||||
)
|
||||
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,
|
||||
)
|
||||
|
||||
@_lock
|
||||
def list_agents(self) -> list[Agent]:
|
||||
with self._lock:
|
||||
return list(self._agents.values())
|
||||
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(
|
||||
id=uuid.uuid4(),
|
||||
participants=frozenset(participants),
|
||||
created_at=datetime.datetime.now(datetime.UTC),
|
||||
)
|
||||
self._threads[t.id] = t
|
||||
return t
|
||||
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
|
||||
|
||||
@_lock
|
||||
def get_thread(self, id: uuid.UUID) -> Thread | None:
|
||||
with self._lock:
|
||||
return self._threads.get(id)
|
||||
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
|
||||
messages = self._msgs.get(id, [])
|
||||
return ThreadInfo(
|
||||
thread=thread,
|
||||
message_count=len(messages),
|
||||
latest_message=messages[-1] if messages else 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,
|
||||
)
|
||||
|
||||
@_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, [])
|
||||
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
|
||||
),
|
||||
)
|
||||
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
|
||||
),
|
||||
)
|
||||
|
||||
@_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]
|
||||
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
|
||||
)
|
||||
self._threads[thread_id] = nt
|
||||
return nt
|
||||
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
|
||||
|
||||
@_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
|
||||
)
|
||||
self._threads[thread_id] = nt
|
||||
return nt
|
||||
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
|
||||
|
||||
@_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)
|
||||
self._threads.pop(id, None)
|
||||
self._msgs.pop(id, None)
|
||||
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:
|
||||
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
|
||||
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
|
||||
|
||||
@_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]
|
||||
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)
|
||||
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
|
||||
self._cursors[(thread_id, agent.id)] = seq
|
||||
Reference in new issue
Block a user