AbstractRuntime 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

abstractruntime/storage/in_memory.py

@@ -58,6 +58,47 @@ class InMemoryRunStore(RunStore):
 
         return results[:limit]
 
+    def list_run_index(
+        self,
+        *,
+        status: Optional[RunStatus] = None,
+        workflow_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        root_only: bool = False,
+        limit: int = 100,
+    ) -> List[Dict[str, Any]]:
+        lim = max(1, int(limit or 100))
+        out: List[Dict[str, Any]] = []
+
+        for run in self._runs.values():
+            if status is not None and run.status != status:
+                continue
+            if workflow_id is not None and run.workflow_id != workflow_id:
+                continue
+            if session_id is not None and str(run.session_id or "").strip() != str(session_id or "").strip():
+                continue
+            if bool(root_only) and str(run.parent_run_id or "").strip():
+                continue
+
+            waiting = run.waiting
+            out.append(
+                {
+                    "run_id": str(run.run_id),
+                    "workflow_id": str(run.workflow_id),
+                    "status": str(getattr(run.status, "value", run.status)),
+                    "wait_reason": str(getattr(getattr(waiting, "reason", None), "value", waiting.reason)) if waiting is not None else None,
+                    "wait_until": str(getattr(waiting, "until", None)) if waiting is not None else None,
+                    "parent_run_id": str(run.parent_run_id) if run.parent_run_id else None,
+                    "actor_id": str(run.actor_id) if run.actor_id else None,
+                    "session_id": str(run.session_id) if run.session_id else None,
+                    "created_at": str(run.created_at) if run.created_at else None,
+                    "updated_at": str(run.updated_at) if run.updated_at else None,
+                }
+            )
+
+        out.sort(key=lambda r: str(r.get("updated_at") or ""), reverse=True)
+        return out[:lim]
+
     def list_due_wait_until(
         self,
         *,
@@ -116,4 +157,3 @@ class InMemoryLedgerStore(LedgerStore):
     def list(self, run_id: str) -> List[Dict[str, Any]]:
         return list(self._records.get(run_id, []))
 
-

abstractruntime/storage/json_files.py

@@ -10,6 +10,7 @@ This is meant as a straightforward MVP backend.
 from __future__ import annotations
 
 import json
+import threading
 import uuid
 from dataclasses import asdict
 from pathlib import Path
@@ -25,16 +26,80 @@ class JsonFileRunStore(RunStore):
     Implements both RunStore (ABC) and QueryableRunStore (Protocol).
 
     Query operations scan all run_*.json files, which is acceptable for MVP
-    but may need indexing for large deployments.
+    but needs lightweight indexing for interactive workloads (e.g. WS tick loops)
+    once the run directory grows.
     """
 
     def __init__(self, base_dir: str | Path):
         self._base = Path(base_dir)
         self._base.mkdir(parents=True, exist_ok=True)
+        self._index_lock = threading.Lock()
+        self._children_index: Optional[Dict[str, set[str]]] = None
+        self._run_parent_index: Dict[str, Optional[str]] = {}
+        self._run_cache_lock = threading.Lock()
+        # run_id -> (mtime_ns, RunState)
+        self._run_cache: Dict[str, tuple[int, RunState]] = {}
 
     def _path(self, run_id: str) -> Path:
         return self._base / f"run_{run_id}.json"
 
+    def _run_id_from_path(self, p: Path) -> str:
+        name = str(getattr(p, "name", "") or "")
+        if not name.startswith("run_") or not name.endswith(".json"):
+            return ""
+        return name[len("run_") : -len(".json")]
+
+    def _ensure_children_index(self) -> None:
+        if self._children_index is not None:
+            return
+        with self._index_lock:
+            if self._children_index is not None:
+                return
+
+            children: Dict[str, set[str]] = {}
+            run_parent: Dict[str, Optional[str]] = {}
+
+            for run in self._iter_all_runs():
+                parent = run.parent_run_id
+                run_parent[run.run_id] = parent
+                if isinstance(parent, str) and parent:
+                    children.setdefault(parent, set()).add(run.run_id)
+
+            self._children_index = children
+            self._run_parent_index = run_parent
+
+    def _drop_from_children_index(self, run_id: str) -> None:
+        with self._index_lock:
+            if self._children_index is None:
+                return
+            parent = self._run_parent_index.pop(run_id, None)
+            if isinstance(parent, str) and parent:
+                siblings = self._children_index.get(parent)
+                if siblings is not None:
+                    siblings.discard(run_id)
+                    if not siblings:
+                        self._children_index.pop(parent, None)
+
+    def _update_children_index_on_save(self, run: RunState) -> None:
+        run_id = run.run_id
+        new_parent = run.parent_run_id
+
+        with self._index_lock:
+            if self._children_index is None:
+                return
+
+            old_parent = self._run_parent_index.get(run_id)
+            if isinstance(old_parent, str) and old_parent and old_parent != new_parent:
+                siblings = self._children_index.get(old_parent)
+                if siblings is not None:
+                    siblings.discard(run_id)
+                    if not siblings:
+                        self._children_index.pop(old_parent, None)
+
+            self._run_parent_index[run_id] = new_parent
+            if isinstance(new_parent, str) and new_parent:
+                self._children_index.setdefault(new_parent, set()).add(run_id)
+
     def save(self, run: RunState) -> None:
         p = self._path(run.run_id)
         # Atomic write to prevent corrupted/partial JSON when multiple threads/processes
@@ -51,6 +116,15 @@ class JsonFileRunStore(RunStore):
                     tmp.unlink()
             except Exception:
                 pass
+        self._update_children_index_on_save(run)
+        try:
+            st = p.stat()
+            mtime_ns = int(getattr(st, "st_mtime_ns", 0) or 0)
+        except Exception:
+            mtime_ns = 0
+        if mtime_ns > 0:
+            with self._run_cache_lock:
+                self._run_cache[str(run.run_id)] = (mtime_ns, run)
 
     def load(self, run_id: str) -> Optional[RunState]:
         p = self._path(run_id)
@@ -60,6 +134,17 @@ class JsonFileRunStore(RunStore):
 
     def _load_from_path(self, p: Path) -> Optional[RunState]:
         """Load a RunState from a file path."""
+        rid_hint = self._run_id_from_path(p)
+        try:
+            st = p.stat()
+            mtime_ns = int(getattr(st, "st_mtime_ns", 0) or 0)
+        except Exception:
+            mtime_ns = 0
+        if rid_hint and mtime_ns > 0:
+            with self._run_cache_lock:
+                cached = self._run_cache.get(rid_hint)
+            if cached is not None and int(cached[0]) == int(mtime_ns):
+                return cached[1]
         try:
             with p.open("r", encoding="utf-8") as f:
                 data = json.load(f)
@@ -89,7 +174,7 @@ class JsonFileRunStore(RunStore):
                 details=raw_waiting.get("details"),
             )
 
-        return RunState(
+        run = RunState(
             run_id=data["run_id"],
             workflow_id=data["workflow_id"],
             status=status,
@@ -104,6 +189,11 @@ class JsonFileRunStore(RunStore):
             session_id=data.get("session_id"),
             parent_run_id=data.get("parent_run_id"),
         )
+        rid = str(getattr(run, "run_id", "") or "").strip() or rid_hint
+        if rid and mtime_ns > 0:
+            with self._run_cache_lock:
+                self._run_cache[rid] = (mtime_ns, run)
+        return run
 
     def _iter_all_runs(self) -> List[RunState]:
         """Iterate over all stored runs."""
@@ -124,11 +214,28 @@ class JsonFileRunStore(RunStore):
         workflow_id: Optional[str] = None,
         limit: int = 100,
     ) -> List[RunState]:
-        """List runs matching the given filters."""
-        results: List[RunState] = []
+        """List runs matching the given filters.
+
+        Performance note:
+        - We order by run file mtime (close to updated_at) and stop once we have `limit` matches.
+        - This avoids parsing every historical run JSON file on large runtimes.
+        """
+        lim = max(1, int(limit or 100))
+        ranked: list[tuple[int, Path]] = []
+        for p in self._base.glob("run_*.json"):
+            try:
+                st = p.stat()
+                mtime_ns = int(getattr(st, "st_mtime_ns", 0) or 0)
+            except Exception:
+                continue
+            ranked.append((mtime_ns, p))
+        ranked.sort(key=lambda x: x[0], reverse=True)
 
-        for run in self._iter_all_runs():
-            # Apply filters
+        results: List[RunState] = []
+        for _mtime_ns, p in ranked:
+            run = self._load_from_path(p)
+            if run is None:
+                continue
             if status is not None and run.status != status:
                 continue
             if workflow_id is not None and run.workflow_id != workflow_id:
@@ -136,13 +243,70 @@ class JsonFileRunStore(RunStore):
             if wait_reason is not None:
                 if run.waiting is None or run.waiting.reason != wait_reason:
                     continue
-
             results.append(run)
+            if len(results) >= lim:
+                break
 
-        # Sort by updated_at descending (most recent first)
         results.sort(key=lambda r: r.updated_at or "", reverse=True)
+        return results[:lim]
 
-        return results[:limit]
+    def list_run_index(
+        self,
+        *,
+        status: Optional[RunStatus] = None,
+        workflow_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        root_only: bool = False,
+        limit: int = 100,
+    ) -> List[Dict[str, Any]]:
+        """List lightweight run index rows without depending on full RunState consumers."""
+        lim = max(1, int(limit or 100))
+        ranked: list[tuple[int, Path]] = []
+        for p in self._base.glob("run_*.json"):
+            try:
+                st = p.stat()
+                mtime_ns = int(getattr(st, "st_mtime_ns", 0) or 0)
+            except Exception:
+                continue
+            ranked.append((mtime_ns, p))
+        ranked.sort(key=lambda x: x[0], reverse=True)
+
+        out: List[Dict[str, Any]] = []
+        sid = str(session_id or "").strip() if session_id is not None else None
+
+        for _mtime_ns, p in ranked:
+            run = self._load_from_path(p)
+            if run is None:
+                continue
+            if status is not None and run.status != status:
+                continue
+            if workflow_id is not None and run.workflow_id != workflow_id:
+                continue
+            if sid is not None and str(run.session_id or "").strip() != sid:
+                continue
+            if bool(root_only) and str(run.parent_run_id or "").strip():
+                continue
+
+            waiting = run.waiting
+            out.append(
+                {
+                    "run_id": str(run.run_id),
+                    "workflow_id": str(run.workflow_id),
+                    "status": str(getattr(run.status, "value", run.status)),
+                    "wait_reason": str(getattr(getattr(waiting, "reason", None), "value", waiting.reason)) if waiting is not None else None,
+                    "wait_until": str(getattr(waiting, "until", None)) if waiting is not None else None,
+                    "parent_run_id": str(run.parent_run_id) if run.parent_run_id else None,
+                    "actor_id": str(run.actor_id) if run.actor_id else None,
+                    "session_id": str(run.session_id) if run.session_id else None,
+                    "created_at": str(run.created_at) if run.created_at else None,
+                    "updated_at": str(run.updated_at) if run.updated_at else None,
+                }
+            )
+            if len(out) >= lim:
+                break
+
+        out.sort(key=lambda r: str(r.get("updated_at") or ""), reverse=True)
+        return out[:lim]
 
     def list_due_wait_until(
         self,
@@ -180,10 +344,15 @@ class JsonFileRunStore(RunStore):
         status: Optional[RunStatus] = None,
     ) -> List[RunState]:
         """List child runs of a parent."""
-        results: List[RunState] = []
+        self._ensure_children_index()
+        with self._index_lock:
+            child_ids = list((self._children_index or {}).get(parent_run_id, set()))
 
-        for run in self._iter_all_runs():
-            if run.parent_run_id != parent_run_id:
+        results: List[RunState] = []
+        for run_id in sorted(child_ids):
+            run = self.load(run_id)
+            if run is None:
+                self._drop_from_children_index(run_id)
                 continue
             if status is not None and run.status != status:
                 continue
@@ -219,3 +388,17 @@ class JsonlLedgerStore(LedgerStore):
                 out.append(json.loads(line))
         return out
 
+    def count(self, run_id: str) -> int:
+        """Return the number of ledger records for run_id (fast path).
+
+        This avoids JSON parsing when only a count is needed (e.g. UI dropdowns).
+        """
+        p = self._path(run_id)
+        if not p.exists():
+            return 0
+        n = 0
+        with p.open("r", encoding="utf-8") as f:
+            for line in f:
+                if line.strip():
+                    n += 1
+        return n

abstractruntime/storage/observable.py

@@ -74,6 +74,44 @@ class ObservableLedgerStore(LedgerStore):
     def list(self, run_id: str) -> List[LedgerRecordDict]:
         return self._inner.list(run_id)
 
+    def count(self, run_id: str) -> int:
+        """Best-effort record count for run_id.
+
+        When the inner LedgerStore implements a faster `count()` method (e.g. JsonlLedgerStore),
+        delegate to it. Otherwise, fall back to `len(list(run_id))`.
+        """
+        inner = self._inner
+        try:
+            count_fn = getattr(inner, "count", None)
+            if callable(count_fn):
+                return int(count_fn(run_id))
+        except Exception:
+            pass
+        try:
+            return len(inner.list(run_id))
+        except Exception:
+            return 0
+
+    def count_many(self, run_ids: List[str]) -> Dict[str, int]:  # type: ignore[override]
+        fn = getattr(self._inner, "count_many", None)
+        if callable(fn):
+            try:
+                out = fn(run_ids)
+                return out if isinstance(out, dict) else {}
+            except Exception:
+                return {}
+        return {str(r or "").strip(): self.count(str(r or "").strip()) for r in (run_ids or []) if str(r or "").strip()}
+
+    def metrics_many(self, run_ids: List[str]) -> Dict[str, Dict[str, int]]:  # type: ignore[override]
+        fn = getattr(self._inner, "metrics_many", None)
+        if callable(fn):
+            try:
+                out = fn(run_ids)
+                return out if isinstance(out, dict) else {}
+            except Exception:
+                return {}
+        return {}
+
     def subscribe(
         self,
         callback: LedgerSubscriber,
@@ -96,4 +134,3 @@ class ObservableLedgerStore(LedgerStore):
         """Clear all subscribers (test utility)."""
         with self._lock:
             self._subscribers.clear()
-