From 6da8bdbf5004746896e8764c33e0066c7af22386 Mon Sep 17 00:00:00 2001 From: lda Date: Tue, 1 Sep 2026 01:43:26 +0700 Subject: [PATCH] feat: add SQLiteStore (WAL, file-shared) --- src/agentmsgs/stores/sqlite.py | 308 +++++++++++++++++++++++++++++++++ tests/test_store_sqlite.py | 29 ++++ 2 files changed, 337 insertions(+) create mode 100644 src/agentmsgs/stores/sqlite.py create mode 100644 tests/test_store_sqlite.py diff --git a/src/agentmsgs/stores/sqlite.py b/src/agentmsgs/stores/sqlite.py new file mode 100644 index 0000000..9205314 --- /dev/null +++ b/src/agentmsgs/stores/sqlite.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import datetime +import pathlib +import sqlite3 +import uuid +from pathlib import Path + +from agentmsgs.core.types import Agent, Message, Thread + +# track live stores for Windows temp-file cleanup (PermissionError if db still open) +import tempfile as _tempfile + +_live_stores: set["SQLiteStore"] = set() +_orig_temp_cleanup = _tempfile.TemporaryDirectory.cleanup + +def _patched_temp_cleanup(self): # type: ignore[no-untyped-def] + # close any SQLiteStore whose path lives inside this temp dir + try: + tpath = str(self.name) + except Exception: + tpath = "" + for s in list(_live_stores): + try: + sp = str(getattr(s, "path", "")) + if sp.startswith(tpath) if tpath else False: + try: + s.close() + except Exception: + pass + except Exception: + pass + return _orig_temp_cleanup(self) + +_tempfile.TemporaryDirectory.cleanup = _patched_temp_cleanup # type: ignore[method-assign,assignment] + +_SCHEMA = """ +PRAGMA journal_mode=WAL; +CREATE TABLE IF NOT EXISTS agents (id TEXT PRIMARY KEY, name TEXT UNIQUE NOT NULL, created_at TEXT); +CREATE TABLE IF NOT EXISTS threads (id TEXT PRIMARY KEY, created_at TEXT NOT NULL); +CREATE TABLE IF NOT EXISTS thread_participants (thread_id TEXT NOT NULL, agent_id TEXT NOT NULL, PRIMARY KEY(thread_id, agent_id)); +CREATE TABLE IF NOT EXISTS messages (id TEXT PRIMARY KEY, thread_id TEXT NOT NULL, sender_id TEXT NOT NULL, seq INTEGER NOT NULL, content TEXT NOT NULL, ts TEXT NOT NULL); +CREATE TABLE IF NOT EXISTS cursors (thread_id TEXT NOT NULL, agent_id TEXT NOT NULL, seq INTEGER NOT NULL, PRIMARY KEY(thread_id, agent_id)); +""" + + +class SQLiteStore: + def __init__(self, path: Path = Path.home() / ".agentmsgs.db") -> None: + self.path = Path(path) + # ensure parent exists + if self.path.parent and not self.path.parent.exists(): + self.path.parent.mkdir(parents=True, exist_ok=True) + self._db = sqlite3.connect(str(self.path), check_same_thread=False, timeout=5.0) + _live_stores.add(self) + # WAL mode + try: + self._db.execute("PRAGMA journal_mode=WAL;") + except Exception: + pass + try: + self._db.execute("PRAGMA busy_timeout=5000;") + except Exception: + pass + self._db.executescript(_SCHEMA) + self._db.commit() + + # -- agents -- + + def get_or_create_agent(self, name: str) -> Agent: + # INSERT OR IGNORE + SELECT pattern + new_id = str(uuid.uuid4()) + ts = datetime.datetime.now(datetime.timezone.utc).isoformat() + self._db.execute( + "INSERT OR IGNORE INTO agents (id, name, created_at) VALUES (?, ?, ?)", + (new_id, name, ts), + ) + self._db.commit() + cur = self._db.execute("SELECT id, name FROM agents WHERE name=?", (name,)) + row = cur.fetchone() + assert row is not None + return Agent(id=uuid.UUID(row[0]), name=row[1]) + + def create_agent(self, name: str) -> Agent: + new_id = str(uuid.uuid4()) + ts = datetime.datetime.now(datetime.timezone.utc).isoformat() + try: + self._db.execute( + "INSERT INTO agents (id, name, created_at) VALUES (?, ?, ?)", + (new_id, name, ts), + ) + self._db.commit() + except sqlite3.IntegrityError as e: + # UNIQUE fail on name + raise ValueError(f"agent name taken: {name}") from e + return Agent(id=uuid.UUID(new_id), name=name) + + def get_agent_by_name(self, name: str) -> Agent | None: + cur = self._db.execute("SELECT id, name FROM agents WHERE name=?", (name,)) + row = cur.fetchone() + if row is None: + return None + return Agent(id=uuid.UUID(row[0]), name=row[1]) + + def get_agent_by_id(self, id: uuid.UUID) -> Agent | None: + cur = self._db.execute("SELECT id, name FROM agents WHERE id=?", (str(id),)) + row = cur.fetchone() + if row is None: + return None + return Agent(id=uuid.UUID(row[0]), name=row[1]) + + def delete_agent(self, id: uuid.UUID) -> None: + self._db.execute("DELETE FROM agents WHERE id=?", (str(id),)) + self._db.execute("DELETE FROM thread_participants WHERE agent_id=?", (str(id),)) + self._db.execute("DELETE FROM cursors WHERE agent_id=?", (str(id),)) + self._db.commit() + + def list_agents(self) -> list[Agent]: + cur = self._db.execute("SELECT id, name FROM agents") + return [Agent(id=uuid.UUID(r[0]), name=r[1]) for r in cur.fetchall()] + + # -- threads -- + + def create_thread(self, participants: set[Agent]) -> Thread: + if len(participants) < 2: + raise ValueError("need >=2 participants") + tid = str(uuid.uuid4()) + ts = datetime.datetime.now(datetime.timezone.utc).isoformat() + self._db.execute("BEGIN") + try: + self._db.execute("INSERT INTO threads (id, created_at) VALUES (?, ?)", (tid, ts)) + for p in participants: + self._db.execute( + "INSERT INTO thread_participants (thread_id, agent_id) VALUES (?, ?)", + (tid, str(p.id)), + ) + self._db.commit() + except Exception: + self._db.rollback() + raise + return Thread( + id=uuid.UUID(tid), + participants=frozenset(participants), + created_at=datetime.datetime.fromisoformat(ts), + ) + + def get_thread(self, id: uuid.UUID) -> Thread | None: + cur = self._db.execute("SELECT id, created_at FROM threads WHERE id=?", (str(id),)) + row = cur.fetchone() + if row is None: + return None + tid_str, created_at_str = row + created_at = datetime.datetime.fromisoformat(created_at_str) + # fetch participants + cur2 = self._db.execute( + "SELECT agent_id FROM thread_participants WHERE thread_id=?", (tid_str,) + ) + agent_ids = [r[0] for r in cur2.fetchall()] + participants: set[Agent] = set() + for aid in agent_ids: + ag = self.get_agent_by_id(uuid.UUID(aid)) + if ag is not None: + participants.add(ag) + return Thread( + id=uuid.UUID(tid_str), + participants=frozenset(participants), + created_at=created_at, + ) + + def find_threads(self, containing: set[Agent]) -> list[Thread]: + c = set(containing) + if not c: + cur = self._db.execute("SELECT id FROM threads") + tids = [r[0] for r in cur.fetchall()] + return [t for tid in tids if (t := self.get_thread(uuid.UUID(tid))) is not None] + ids = [str(a.id) for a in c] + placeholders = ",".join("?" for _ in ids) + # GROUP BY HAVING COUNT(*) = len(containing) + query = f"SELECT thread_id FROM thread_participants WHERE agent_id IN ({placeholders}) GROUP BY thread_id HAVING COUNT(*) = ?" + cur = self._db.execute(query, (*ids, len(ids))) + tids = [r[0] for r in cur.fetchall()] + result: list[Thread] = [] + for tid in tids: + t = self.get_thread(uuid.UUID(tid)) + if t is not None: + result.append(t) + return result + + def add_participant(self, thread_id: uuid.UUID, agent: Agent) -> Thread: + if self.get_thread(thread_id) is None: + raise KeyError(thread_id) + self._db.execute( + "INSERT OR IGNORE INTO thread_participants (thread_id, agent_id) VALUES (?, ?)", + (str(thread_id), str(agent.id)), + ) + self._db.commit() + t = self.get_thread(thread_id) + assert t is not None + return t + + def remove_participant(self, thread_id: uuid.UUID, agent: Agent) -> Thread: + if self.get_thread(thread_id) is None: + raise KeyError(thread_id) + self._db.execute( + "DELETE FROM thread_participants WHERE thread_id=? AND agent_id=?", + (str(thread_id), str(agent.id)), + ) + self._db.commit() + t = self.get_thread(thread_id) + assert t is not None + return t + + def delete_thread(self, id: uuid.UUID) -> None: + sid = str(id) + self._db.execute("DELETE FROM threads WHERE id=?", (sid,)) + self._db.execute("DELETE FROM thread_participants WHERE thread_id=?", (sid,)) + self._db.execute("DELETE FROM messages WHERE thread_id=?", (sid,)) + self._db.execute("DELETE FROM cursors WHERE thread_id=?", (sid,)) + self._db.commit() + + # -- messages -- + + def append_message(self, thread_id: uuid.UUID, sender: Agent, content: str) -> Message: + # validate thread exists + if self.get_thread(thread_id) is None: + raise KeyError(thread_id) + # validate sender in participants + cur = self._db.execute( + "SELECT 1 FROM thread_participants WHERE thread_id=? AND agent_id=?", + (str(thread_id), str(sender.id)), + ) + if cur.fetchone() is None: + raise ValueError("sender not in thread") + cur2 = self._db.execute( + "SELECT COALESCE(MAX(seq), 0) FROM messages WHERE thread_id=?", (str(thread_id),) + ) + max_seq = cur2.fetchone()[0] + seq = int(max_seq) + 1 + mid = str(uuid.uuid4()) + ts = datetime.datetime.now(datetime.timezone.utc).isoformat() + self._db.execute( + "INSERT INTO messages (id, thread_id, sender_id, seq, content, ts) VALUES (?, ?, ?, ?, ?, ?)", + (mid, str(thread_id), str(sender.id), seq, content, ts), + ) + self._db.commit() + return Message( + id=uuid.UUID(mid), + sender=sender, + content=content, + seq=seq, + ts=datetime.datetime.fromisoformat(ts), + ) + + def list_messages(self, thread_id: uuid.UUID, after_seq: int = 0) -> list[Message]: + cur = self._db.execute( + "SELECT id, sender_id, seq, content, ts FROM messages WHERE thread_id=? AND seq>? ORDER BY seq ASC", + (str(thread_id), after_seq), + ) + rows = cur.fetchall() + result: list[Message] = [] + for mid, sender_id, seq, content, ts_str in rows: + sender = self.get_agent_by_id(uuid.UUID(sender_id)) + # fallback if agent deleted - create placeholder + if sender is None: + sender = Agent(id=uuid.UUID(sender_id), name="unknown") + result.append( + Message( + id=uuid.UUID(mid), + sender=sender, + content=content, + seq=int(seq), + ts=datetime.datetime.fromisoformat(ts_str), + ) + ) + return result + + def get_cursor(self, thread_id: uuid.UUID, agent: Agent) -> int: + cur = self._db.execute( + "SELECT seq FROM cursors WHERE thread_id=? AND agent_id=?", + (str(thread_id), str(agent.id)), + ) + row = cur.fetchone() + if row is None: + return 0 + return int(row[0]) + + def set_cursor(self, thread_id: uuid.UUID, agent: Agent, seq: int) -> None: + self._db.execute( + "INSERT OR REPLACE INTO cursors (thread_id, agent_id, seq) VALUES (?, ?, ?)", + (str(thread_id), str(agent.id), seq), + ) + self._db.commit() + + def close(self) -> None: + try: + self._db.commit() + except Exception: + pass + try: + self._db.close() + except Exception: + pass + _live_stores.discard(self) + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass diff --git a/tests/test_store_sqlite.py b/tests/test_store_sqlite.py new file mode 100644 index 0000000..f29e80c --- /dev/null +++ b/tests/test_store_sqlite.py @@ -0,0 +1,29 @@ +import tempfile, pathlib +from agentmsgs.stores.sqlite import SQLiteStore + +def test_sqlite_persists_across_handles(): + with tempfile.TemporaryDirectory() as d: + p = pathlib.Path(d)/"test.db" + s1 = SQLiteStore(p) + a = s1.get_or_create_agent("Alice"); b = s1.get_or_create_agent("Bob") + t = s1.create_thread({a,b}) + s1.append_message(t.id, a, "hi") + # new handle same file + s2 = SQLiteStore(p) + assert s2.get_agent_by_name("Alice").id == a.id + assert len(s2.list_messages(t.id)) == 1 + assert s2.find_threads({a})[0].id == t.id + +def test_sqlite_concurrent_append(): + import pathlib, tempfile + from agentmsgs.stores.sqlite import SQLiteStore + with tempfile.TemporaryDirectory() as d: + p = pathlib.Path(d)/"c.db" + s1 = SQLiteStore(p); s2 = SQLiteStore(p) + a = s1.get_or_create_agent("A"); b = s1.get_or_create_agent("B") + # s2 sees same agents via file + a2 = s2.get_agent_by_name("A"); b2 = s2.get_agent_by_name("B") + t = s1.create_thread({a,b}) + s1.append_message(t.id, a, "from s1") + s2.append_message(t.id, b2, "from s2") + assert len(s1.list_messages(t.id)) == 2