aethergraph 0.1.0a1__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

aethergraph/storage/continuation_store/kvdoc_cont.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+from dataclasses import asdict
+from datetime import datetime, timezone
+import hashlib
+import hmac
+import os
+from typing import Any
+
+from aethergraph.contracts.services.continuations import AsyncContinuationStore
+from aethergraph.contracts.services.kv import AsyncKV
+from aethergraph.contracts.storage.doc_store import DocStore
+from aethergraph.contracts.storage.event_log import EventLog
+from aethergraph.services.continuations.continuation import Continuation, Correlator
+
+
+class KVDocContinuationStore(AsyncContinuationStore):
+    """
+    Continuation store backed by:
+      - DocStore: main continuation document (one per (run_id, node_id))
+      - AsyncKV:  token and correlator indices
+      - EventLog: (optional) audit trail of continuation events
+    """
+
+    def __init__(
+        self,
+        *,
+        doc_store: DocStore,
+        kv: AsyncKV,
+        event_log: EventLog | None = None,
+        secret: bytes,
+        namespace: str = "cont",
+    ):
+        self._docs = doc_store
+        self._kv = kv
+        self._log = event_log
+        self._secret = secret
+        self._ns = namespace.rstrip("/")  # namespace prefix for KV keys
+
+    # ---------- key helpers ----------
+    def _cont_id(self, run_id: str, node_id: str) -> str:
+        # one doc per continuation
+        return f"{self._ns}/runs/{run_id}/nodes/{node_id}"
+
+    def _token_key(self, token: str) -> str:
+        return f"{self._ns}:token:{token}"
+
+    def _corr_key(self, corr: Correlator) -> str:
+        scheme, channel, thread, message = corr.key()
+        return f"{self._ns}:corr:{scheme}:{channel}:{thread}:{message}"
+
+    # ---------- token helpers ----------
+    def _hmac(self, *parts: str) -> str:
+        h = hmac.new(self._secret, digestmod=hashlib.sha256)
+        for p in parts:
+            h.update(p.encode("utf-8"))
+        return h.hexdigest()
+
+    async def mint_token(self, run_id: str, node_id: str, attempts: int) -> str:
+        token = self._hmac(run_id, node_id, str(attempts), os.urandom(8).hex())
+        return token
+
+    # ---------- main methods ----------
+    async def save(self, cont: Continuation) -> None:
+        payload = cont.to_dict() if hasattr(cont, "to_dict") else asdict(cont)
+
+        # Normalize datetime fields to ISO format
+        for k in ("deadline", "next_wakeup_at", "created_at"):
+            v = payload.get(k)
+            if isinstance(v, datetime):
+                payload[k] = v.astimezone(timezone.utc).isoformat()
+
+        doc_id = self._cont_id(cont.run_id, cont.node_id)
+        await self._docs.put(doc_id, payload)
+
+        # token -> (run_id, node_id)
+        await self._kv.set(
+            self._token_key(cont.token),
+            {"run_id": cont.run_id, "node_id": cont.node_id},
+        )
+
+        if self._log is not None:
+            evt = {
+                "scope_id": cont.run_id,
+                "kind": "continuation.save",
+                "ts": datetime.now(timezone.utc).isoformat(),
+                "tags": [cont.channel or "", cont.kind or ""],
+                "payload": payload,
+            }
+            await self._log.append(evt)
+
+    async def _doc_to_cont(self, data: dict[str, Any] | None) -> Continuation | None:
+        if data is None:
+            return None
+
+        for k in ("deadline", "next_wakeup_at", "created_at"):
+            v = data.get(k)
+            if isinstance(v, str):
+                data[k] = datetime.fromisoformat(v)
+
+        data["closed"] = bool(data.get("closed", False))
+        return Continuation(**data)
+
+    async def get(self, run_id: str, node_id: str) -> Continuation | None:
+        doc_id = self._cont_id(run_id, node_id)
+        data = await self._docs.get(doc_id)
+        return await self._doc_to_cont(data)
+
+    async def list_cont_by_run(self, run_id: str) -> list[Continuation]:
+        prefix = f"{self._ns}/runs/{run_id}/nodes/"
+        ids = await self._docs.list()
+        out: list[Continuation] = []
+        for doc_id in ids:
+            if not doc_id.startswith(prefix):
+                continue
+            data = await self._docs.get(doc_id)
+            cont = await self._doc_to_cont(data)
+            if cont is not None:
+                out.append(cont)
+        return out
+
+    async def delete(self, run_id: str, node_id: str) -> None:
+        cont = await self.get(run_id, node_id)
+        doc_id = self._cont_id(run_id, node_id)
+        await self._docs.delete(doc_id)
+        if cont:
+            # best-effort delete of token mapping
+            await self._kv.delete(self._token_key(cont.token))
+
+    # ---------- token methods ----------
+    async def get_by_token(self, token: str) -> Continuation | None:
+        ref = await self._kv.get(self._token_key(token), default=None)
+        if not ref:
+            return None
+        run_id = ref.get("run_id")
+        node_id = ref.get("node_id")
+        if not run_id or not node_id:
+            return None
+        return await self.get(run_id, node_id)
+
+    async def mark_closed(self, token: str) -> None:
+        cont = await self.get_by_token(token)
+        if not cont:
+            return
+        if not cont.closed:
+            cont.closed = True
+            await self.save(cont)
+
+    async def verify_token(self, run_id: str, node_id: str, token: str) -> bool:
+        cont = await self.get(run_id, node_id)
+        return bool(cont and hmac.compare_digest(token, cont.token))
+
+    # ---------- correlator methods ----------
+    async def bind_correlator(self, *, token: str, corr: str) -> None:
+        key = self._corr_key(corr)
+        toks: list[str] = await self._kv.get(key, default=[]) or []
+        if token not in toks:
+            toks.append(token)
+            await self._kv.set(key, toks)
+
+        # TODO: consider reverse mapping: token -> correlators for easier cleanup
+        # for now we don't need that
+
+    async def find_by_correlator(self, *, corr: Correlator) -> list[Continuation] | None:
+        """
+        Find all continuations matching the correlator. Following the method:
+        - get all tokens for the correlator
+        - fetch each continuation by token
+        - filter out closed or expired continuations
+        - return the first valid continuation found (or None)
+        """
+        key = self._corr_key(corr)
+        toks: list[str] = await self._kv.get(key, default=[]) or []
+        from datetime import datetime as dt, timezone as tz
+
+        for tok in reversed(toks):
+            cont = await self.get_by_token(tok)
+            if not cont or cont.closed:
+                continue
+            if cont.deadline and dt.now(tz.utc) > cont.deadline.astimezone(tz.utc):
+                continue
+            return cont
+        return None
+
+    # ---------- scans ----------
+    async def last_open(self, *, channel: str, kind: str) -> Continuation | None:
+        """
+        Slow scan (like FS version) – OK for dev scale.
+        We scan all docs and pick the most recent created_at with matching channel/kind.
+        NOTE: this is only for debugging and development purposes. Do not use in production!
+        """
+        ids = await self._docs.list()
+        best: Continuation | None = None
+        best_ts: float | None = None
+
+        for doc_id in ids:
+            if not doc_id.startswith(f"{self._ns}/runs/"):
+                continue
+            data = await self._docs.get(doc_id)
+            cont = await self._doc_to_cont(data)
+            if not cont or cont.closed:
+                continue
+            if cont.channel != channel or cont.kind != kind:
+                continue
+            created = cont.created_at or datetime.min.replace(tzinf=timezone.utc)
+            ts = created.timestamp()
+            if best_ts is None or ts > best_ts:
+                best = cont
+                best_ts = ts
+        return best
+
+    async def list_waits(self) -> list[dict[str, Any]]:
+        """
+        Return all continuations as dicts (similar to FS version).
+        Caller can filter for waits if needed.
+        """
+        ids = await self._docs.list()
+        out: list[dict[str, Any]] = []
+        for doc_id in ids:
+            if not doc_id.startswith(f"{self._ns}/runs/"):
+                continue
+            data = await self._docs.get(doc_id)
+            if data:
+                out.append(data)
+        return out
+
+    async def clear(self) -> None:
+        """
+        Best-effort clear of continuation documents.
+
+        NOTE: KV indexes may remain unless we have a scan-capable KV implementation.
+        They are harmless: get_by_token() will just return None if the doc is gone.
+        """
+        # 1) DocStore cleanup
+        ids = await self._docs.list()
+        for doc_id in ids:
+            if doc_id.startswith(f"{self._ns}/runs/"):
+                await self._docs.delete(doc_id)
+
+        # 2) Best-effort KV cleanup if scan_keys is available
+        # TODO: implement KV indexes if we have scan capability
+        scan = getattr(self._kv, "scan_keys", None)
+        if callable(scan):
+            token_prefix = f"{self._ns}:token:"
+            corr_prefix = f"{self._ns}:corr:"
+
+            for pfx in (token_prefix, corr_prefix):
+                try:
+                    keys: list[str] = await scan(pfx)  # type: ignore[call-arg]
+                except Exception:
+                    # If a backend throws here, we still consider clear() successful
+                    continue
+                for k in keys:
+                    try:
+                        await self._kv.delete(k)
+                    except Exception:
+                        # swallow individual key errors; we're best-effort here
+                        continue
+
+    async def alias_for(self, token: str) -> str | None:
+        return token[:24]

aethergraph/storage/docstore/fs_doc.py

@@ -0,0 +1,63 @@
+import asyncio
+import json
+import os
+from pathlib import Path
+import threading
+from typing import Any
+
+from aethergraph.contracts.storage.doc_store import DocStore
+
+
+class FSDocStore(DocStore):
+    def __init__(self, root: str):
+        self.root = Path(root)
+        self.root.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+
+    def _path_for(self, doc_id: str) -> Path:
+        p = self.root / f"{doc_id}.json"
+        p.parent.mkdir(parents=True, exist_ok=True)
+        return p
+
+    async def put(self, doc_id: str, doc: dict[str, Any]) -> None:
+        path = self._path_for(doc_id)
+
+        def _write():
+            tmp = path.with_suffix(path.suffix + ".tmp")
+            with self._lock, tmp.open("w", encoding="utf-8") as f:
+                json.dump(doc, f, ensure_ascii=False, indent=2)
+            os.replace(tmp, path)
+
+        await asyncio.to_thread(_write)
+
+    async def get(self, doc_id: str) -> dict[str, Any] | None:
+        path = self._path_for(doc_id)
+
+        def _read():
+            if not path.exists():
+                return None
+            with self._lock, path.open("r", encoding="utf-8") as f:
+                return json.load(f)
+
+        return await asyncio.to_thread(_read)
+
+    async def delete(self, doc_id):
+        path = self._path_for(doc_id)
+
+        def _delete():
+            if path.exists():
+                with self._lock:
+                    path.unlink()
+
+        await asyncio.to_thread(_delete)
+
+    async def list(self) -> list[str]:
+        def _list():
+            out = []
+            for p in self.root.rglob("*.json"):
+                rel = p.relative_to(self.root)
+                doc_id = str(rel.with_suffix("").as_posix())
+                out.append(doc_id)
+            return out
+
+        return await asyncio.to_thread(_list)

aethergraph/storage/docstore/sqlite_doc.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from aethergraph.contracts.storage.doc_store import DocStore
+
+from .sqlite_doc_sync import SQLiteDocStoreSync
+
+
+class SqliteDocStore(DocStore):
+    """
+    Async DocStore implemented on top of SQLiteDocStoreSync via asyncio.to_thread.
+
+    Safe to use from multiple threads (sidecar + main loop) due to RLock in sync core.
+    """
+
+    def __init__(self, path: str):
+        self._sync = SQLiteDocStoreSync(path)
+
+    async def put(self, doc_id: str, doc: dict[str, Any]) -> None:
+        await asyncio.to_thread(self._sync.put, doc_id, doc)
+
+    async def get(self, doc_id: str) -> dict[str, Any] | None:
+        return await asyncio.to_thread(self._sync.get, doc_id)
+
+    async def delete(self, doc_id: str) -> None:
+        await asyncio.to_thread(self._sync.delete, doc_id)
+
+    async def list(self) -> list[str]:
+        return await asyncio.to_thread(self._sync.list)

aethergraph/storage/docstore/sqlite_doc_sync.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+import sqlite3
+import threading
+import time
+from typing import Any
+
+"""
+This is not used in the main codebase; only used by async wrapper SqliteDocStore.
+"""
+
+
+class SQLiteDocStoreSync:
+    """
+    Durable document store on SQLite.
+
+    - Single connection per instance.
+    - Thread-safe via RLock.
+    - Values are JSON-serialized dicts.
+    """
+
+    def __init__(self, path: str):
+        path_obj = Path(path)
+        path_obj.parent.mkdir(parents=True, exist_ok=True)
+
+        self._db = sqlite3.connect(
+            str(path_obj),
+            check_same_thread=False,  # allow multi-thread access (guarded by RLock)
+            isolation_level=None,  # autocommit
+        )
+        self._db.execute("PRAGMA journal_mode=WAL;")
+        self._db.execute("PRAGMA synchronous=NORMAL;")
+        self._db.execute(
+            """
+            CREATE TABLE IF NOT EXISTS docs (
+                doc_id     TEXT PRIMARY KEY,
+                data_json  TEXT NOT NULL,
+                updated_at REAL NOT NULL
+            )
+            """
+        )
+        self._lock = threading.RLock()
+
+    def put(self, doc_id: str, doc: dict[str, Any]) -> None:
+        payload = json.dumps(doc, ensure_ascii=False)
+        now = time.time()
+        # TEMP: tiny backoff to avoid rare SQLite stalls under continuations.
+        # This is a hacky workaround.
+        # It happens when following conditions align:
+        # 1) continuation store using sqlite doc to save
+        # 2) the continuation is created under if/else or for loop nodes
+        # NOTE: this bug only appears in SQLite and not in other DBs.
+        # Remove once we move continuations to Postgres or refactor SQLite usage.
+        time.sleep(0.01)
+        with self._lock:
+            try:
+                self._db.execute(
+                    """
+                    INSERT INTO docs (doc_id, data_json, updated_at)
+                    VALUES (?, ?, ?)
+                    ON CONFLICT(doc_id) DO UPDATE SET
+                        data_json = excluded.data_json,
+                        updated_at = excluded.updated_at
+                    """,
+                    (doc_id, payload, now),
+                )
+            except sqlite3.Error as e:
+                print("🍓 SQLiteDocStoreSync ERROR during put:", doc_id, repr(e))
+                raise
+
+    def get(self, doc_id: str) -> dict[str, Any] | None:
+        with self._lock:
+            row = self._db.execute(
+                "SELECT data_json FROM docs WHERE doc_id = ?",
+                (doc_id,),
+            ).fetchone()
+        if not row:
+            return None
+        return json.loads(row[0])
+
+    def delete(self, doc_id: str) -> None:
+        with self._lock:
+            self._db.execute("DELETE FROM docs WHERE doc_id = ?", (doc_id,))
+
+    def list(self) -> list[str]:
+        with self._lock:
+            rows = self._db.execute("SELECT doc_id FROM docs").fetchall()
+        return [r[0] for r in rows]

aethergraph/storage/eventlog/fs_event.py

@@ -0,0 +1,136 @@
+import asyncio
+from datetime import datetime, timezone
+import json
+from pathlib import Path
+import threading
+import time
+
+from aethergraph.contracts.storage.event_log import EventLog
+
+
+def _to_ts_float(v) -> float | None:
+    """
+    Normalize event ts field to a float UNIX timestamp.
+
+    Supports:
+      - float / int already
+      - ISO 8601 string, e.g. '2025-11-27T19:48:09.758687+00:00'
+      - ISO with 'Z' suffix, e.g. '2025-11-27T19:48:09Z'
+    """
+    if v is None:
+        return None
+    if isinstance(v, int | float):
+        return float(v)
+    if isinstance(v, str):
+        try:
+            s = v.replace("Z", "+00:00") if v.endswith("Z") else v
+            dt = datetime.fromisoformat(s)
+            if dt.tzinfo is None:
+                dt = dt.replace(tzinfo=timezone.utc)
+            return dt.timestamp()
+        except Exception:
+            return None
+    if isinstance(v, datetime):
+        if v.tzinfo is None:
+            v = v.replace(tzinfo=timezone.utc)
+        return v.timestamp()
+    return None
+
+
+class FSEventLog(EventLog):
+    def __init__(self, root: str):
+        self.root = Path(root)
+        self.root.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+        self._log_path = self.root / "events.jsonl"
+
+    async def append(self, evt: dict) -> None:
+        def _write():
+            self._log_path.parent.mkdir(parents=True, exist_ok=True)
+            row = evt.copy()
+
+            # Normalize ts to a float UNIX timestamp
+            ts = _to_ts_float(row.get("ts"))
+            if ts is None:
+                ts = time.time()
+            row["ts"] = ts
+
+            with self._lock, self._log_path.open("a", encoding="utf-8") as f:
+                f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+        await asyncio.to_thread(_write)
+
+    async def query(
+        self,
+        *,
+        scope_id: str | None = None,
+        since: datetime | None = None,
+        until: datetime | None = None,
+        kinds: list[str] | None = None,
+        limit: int | None = None,
+        tags: list[str] | None = None,
+        offset: int = 0,
+        user_id: str | None = None,
+        org_id: str | None = None,
+    ) -> list[dict]:
+        """
+        FSEventLog reads the single events.jsonl file linearly, applies
+        all filters (scope_id, time window, kinds, tags, tenant) in Python,
+        and then slices via offset + limit.
+
+        This is fine for dev/demo / low event volumes. For production,
+        prefer SQLiteEventLog or a DB-backed implementation.
+        """
+        if not self._log_path.exists():
+            return []
+
+        def _read() -> list[dict]:
+            out: list[dict] = []
+            t_min = since.timestamp() if since else None
+            t_max = until.timestamp() if until else None
+
+            # If we want to early-break, we need enough rows to cover offset+limit.
+            needed = None
+            if limit is not None:
+                needed = (offset or 0) + limit
+
+            with self._lock, self._log_path.open("r", encoding="utf-8") as f:
+                for line in f:
+                    if not line.strip():
+                        continue
+                    row = json.loads(line)
+
+                    ts_val = _to_ts_float(row.get("ts"))
+
+                    if t_min is not None and ts_val is not None and ts_val < t_min:
+                        continue
+                    if t_max is not None and ts_val is not None and ts_val > t_max:
+                        continue
+                    if scope_id is not None and row.get("scope_id") != scope_id:
+                        continue
+                    if kinds is not None and row.get("kind") not in kinds:
+                        continue
+                    if tags is not None:
+                        row_tags = set(row.get("tags", []))
+                        if not row_tags.issuperset(tags):
+                            continue
+                    if user_id is not None and row.get("user_id") != user_id:
+                        continue
+                    if org_id is not None and row.get("org_id") != org_id:
+                        continue
+
+                    out.append(row)
+
+                    # Only break early when we've collected enough to satisfy offset+limit
+                    if needed is not None and len(out) >= needed:
+                        break
+
+            # Apply offset/limit on the filtered rows
+            if offset > 0:
+                out = out[offset:]
+            if limit is not None:
+                out = out[:limit]
+
+            return out
+
+        return await asyncio.to_thread(_read)

aethergraph/storage/eventlog/sqlite_event.py

@@ -0,0 +1,47 @@
+# storage/events/sqlite_event_log.py
+from __future__ import annotations
+
+import asyncio
+from datetime import datetime
+
+from aethergraph.contracts.storage.event_log import EventLog
+
+from .sqlite_event_sync import SQLiteEventLogSync
+
+
+class SqliteEventLog(EventLog):
+    """
+    Async EventLog wrapper around SQLiteEventLogSync via asyncio.to_thread.
+    """
+
+    def __init__(self, path: str):
+        self._sync = SQLiteEventLogSync(path)
+
+    async def append(self, evt: dict) -> None:
+        await asyncio.to_thread(self._sync.append, evt)
+
+    async def query(
+        self,
+        *,
+        scope_id: str | None = None,
+        since: datetime | None = None,
+        until: datetime | None = None,
+        kinds: list[str] | None = None,
+        limit: int | None = None,
+        tags: list[str] | None = None,
+        offset: int = 0,
+        user_id: str | None = None,
+        org_id: str | None = None,
+    ) -> list[dict]:
+        return await asyncio.to_thread(
+            self._sync.query,
+            scope_id=scope_id,
+            since=since,
+            until=until,
+            kinds=kinds,
+            limit=limit,
+            tags=tags,
+            offset=offset,
+            user_id=user_id,
+            org_id=org_id,
+        )