Bootstrap tamq local tmux agent message queue

Assistant: codex
Assistant-Model: gpt-5.6-sol
Assistant-Session: 01a02f34-ead4-7a60-85cd-d7f08e22fa0e

src/tamq/store.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+import json
+import sqlite3
+import time
+from pathlib import Path
+from typing import Iterable
+from uuid import uuid4
+
+SCHEMA_VERSION = 1
+
+
+class Store:
+    def __init__(self, path: Path):
+        path.parent.mkdir(parents=True, exist_ok=True)
+        self.db = sqlite3.connect(path, timeout=5)
+        self.db.row_factory = sqlite3.Row
+        self.db.execute("PRAGMA journal_mode=WAL")
+        self.db.execute("PRAGMA foreign_keys=ON")
+        self.db.execute("PRAGMA busy_timeout=5000")
+        self.db.executescript("""
+        CREATE TABLE IF NOT EXISTS metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL);
+        CREATE TABLE IF NOT EXISTS messages (
+          message_id TEXT PRIMARY KEY,
+          sender_repo TEXT NOT NULL,
+          target_repo TEXT NOT NULL,
+          body TEXT NOT NULL,
+          created_at REAL NOT NULL,
+          state TEXT NOT NULL,
+          endpoint_id TEXT,
+          provenance TEXT,
+          injected_at REAL,
+          acknowledged_at REAL
+        );
+        CREATE INDEX IF NOT EXISTS messages_target_state ON messages(target_repo, state);
+        CREATE TABLE IF NOT EXISTS endpoints (
+          endpoint_id TEXT PRIMARY KEY,
+          pid INTEGER NOT NULL,
+          session TEXT NOT NULL,
+          repos TEXT NOT NULL,
+          connected_at REAL NOT NULL,
+          disconnected_at REAL
+        );
+        CREATE TABLE IF NOT EXISTS leases (
+          message_id TEXT PRIMARY KEY REFERENCES messages(message_id) ON DELETE CASCADE,
+          lease_id TEXT NOT NULL,
+          endpoint_id TEXT NOT NULL,
+          acquired_at REAL NOT NULL,
+          expires_at REAL NOT NULL
+        );
+        """)
+        self.db.execute("INSERT OR IGNORE INTO metadata(key,value) VALUES('schema_version',?)", (str(SCHEMA_VERSION),))
+        self.db.commit()
+
+    def close(self) -> None:
+        self.db.close()
+
+    def register_endpoint(self, endpoint_id: str, pid: int, session: str, repos: list[str]) -> None:
+        import json
+        self.db.execute(
+            "INSERT INTO endpoints(endpoint_id,pid,session,repos,connected_at,disconnected_at) VALUES(?,?,?,?,strftime('%s','now'),NULL) "
+            "ON CONFLICT(endpoint_id) DO UPDATE SET pid=excluded.pid, session=excluded.session, repos=excluded.repos, connected_at=excluded.connected_at, disconnected_at=NULL",
+            (endpoint_id, pid, session, json.dumps(repos)),
+        )
+        self.db.commit()
+
+    def disconnect_endpoint(self, endpoint_id: str) -> None:
+        self.db.execute("UPDATE endpoints SET disconnected_at=strftime('%s','now') WHERE endpoint_id=?", (endpoint_id,))
+        self.db.commit()
+
+    def disconnect_all(self) -> None:
+        self.db.execute("UPDATE endpoints SET disconnected_at=strftime('%s','now') WHERE disconnected_at IS NULL")
+        self.db.commit()
+
+    def endpoints(self) -> list[sqlite3.Row]:
+        return list(self.db.execute("SELECT * FROM endpoints ORDER BY endpoint_id"))
+
+    def endpoint(self, endpoint_id: str) -> sqlite3.Row | None:
+        row = self.db.execute("SELECT * FROM endpoints WHERE endpoint_id=? AND disconnected_at IS NULL", (endpoint_id,)).fetchone()
+        if row is not None:
+            return row
+        if endpoint_id.startswith("tmux-amq-"):
+            try:
+                pid = int(endpoint_id.removeprefix("tmux-amq-").split("-", 1)[0])
+            except ValueError:
+                return None
+            return self.db.execute("SELECT * FROM endpoints WHERE pid=? AND disconnected_at IS NULL ORDER BY connected_at DESC LIMIT 1", (pid,)).fetchone()
+        return None
+
+    def history_stats(self) -> tuple[int, float | None]:
+        size = 0
+        try:
+            size = self.db.execute("SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()").fetchone()[0]
+        except sqlite3.DatabaseError:
+            pass
+        oldest = self.db.execute("SELECT MIN(created_at) FROM messages").fetchone()[0]
+        return size, oldest
+
+    def add(self, sender: str, target: str, body: str, *, endpoint: str | None = None, provenance: str | None = None) -> str:
+        if len(body.encode("utf-8")) > 8192:
+            raise ValueError("message body exceeds 8 KiB limit")
+        message_id = f"m-{uuid4()}"
+        self.db.execute(
+            "INSERT INTO messages VALUES(?,?,?,?,?,?,?,?,?,?)",
+            (message_id, sender, target, body, time.time(), "pending", endpoint, provenance, None, None),
+        )
+        self.db.commit()
+        return message_id
+
+    def list(self, target: str | None = None, state: str | None = None) -> list[sqlite3.Row]:
+        clauses, values = [], []
+        if target:
+            clauses.append("target_repo=?"); values.append(target)
+        if state:
+            clauses.append("state=?"); values.append(state)
+        where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
+        return list(self.db.execute(f"SELECT * FROM messages{where} ORDER BY created_at", values))
+
+    def set_state(self, message_id: str, state: str) -> None:
+        column = {"injected": "injected_at", "acknowledged": "acknowledged_at"}.get(state)
+        if column:
+            self.db.execute(f"UPDATE messages SET state=?, {column}=? WHERE message_id=?", (state, time.time(), message_id))
+        else:
+            self.db.execute("UPDATE messages SET state=? WHERE message_id=?", (state, message_id))
+        self.db.commit()
+
+    def acknowledge(self, message_id: str) -> bool:
+        row = self.db.execute("SELECT 1 FROM messages WHERE message_id=?", (message_id,)).fetchone()
+        if row is None:
+            return False
+        self.set_state(message_id, "acknowledged")
+        return True
+
+    def claim(self, message_id: str, endpoint_id: str, ttl: float = 30.0) -> str | None:
+        now = time.time()
+        lease_id = f"lease-{uuid4()}"
+        with self.db:
+            self.db.execute("DELETE FROM leases WHERE expires_at < ?", (now,))
+            try:
+                self.db.execute("INSERT INTO leases VALUES(?,?,?,?,?)", (message_id, lease_id, endpoint_id, now, now + ttl))
+            except sqlite3.IntegrityError:
+                return None
+        return lease_id
+
+    def renew(self, message_id: str, lease_id: str, ttl: float = 30.0) -> bool:
+        with self.db:
+            result = self.db.execute("UPDATE leases SET expires_at=? WHERE message_id=? AND lease_id=?", (time.time() + ttl, message_id, lease_id))
+        return result.rowcount == 1
+
+    def release(self, message_id: str, lease_id: str, state: str = "injected") -> bool:
+        with self.db:
+            result = self.db.execute("DELETE FROM leases WHERE message_id=? AND lease_id=?", (message_id, lease_id))
+            if result.rowcount:
+                self.set_state(message_id, state)
+        return result.rowcount == 1
+
+    def export(self, rows: Iterable[sqlite3.Row], path: Path) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("w", encoding="utf-8") as stream:
+            for row in rows:
+                stream.write(json.dumps(dict(row), sort_keys=True) + "\n")
+
+    def purge(self, before: float | None = None, max_bytes: int | None = None) -> int:
+        rows = self.list()
+        ids = [r["message_id"] for r in rows if before is not None and r["created_at"] < before]
+        if max_bytes is not None and self.db.execute("SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()").fetchone()[0] > max_bytes:
+            ids += [r["message_id"] for r in rows if r["message_id"] not in ids][:max(0, len(rows) // 2)]
+        if ids:
+            self.db.executemany("DELETE FROM messages WHERE message_id=?", ((i,) for i in set(ids)))
+            self.db.commit()
+        return len(set(ids))