agent-tracegraph 0.2.0__py3-none-any.whl

tracegraph/sqlite_snapshot.py

@@ -0,0 +1,65 @@
+"""Analyze checkpoints on a private SQLite backup, without initializing the source.
+
+Reading a checkpoint database is the one place this tool touches a file format that can
+*execute*. LangGraph's msgpack deserializer revives objects by importing a module and calling
+a name, both taken from the stored payload; by default an unrecognized target only logs a
+warning and is then called. A crafted checkpoint row naming ``os.system`` therefore runs a
+command on the machine doing the analysis — which is the opposite of what a read-only
+forensic tool should do. So this module hands the saver a serializer with an **empty**
+allowlist: framework-registered types still revive, and anything else comes back as inert
+data instead of being invoked.
+"""
+
+from contextlib import closing
+from pathlib import Path
+import sqlite3
+import tempfile
+import time
+
+from tracegraph.adapters.langgraph_checkpoint import LangGraphCheckpointAdapter
+from tracegraph.model import RawTrace
+
+
+def _inert_serde():
+    """A checkpoint deserializer that revives data but never invokes an arbitrary callable.
+
+    ``allowed_msgpack_modules=()`` keeps the framework's own registered types working — dates,
+    sets, `Send` packets, the values the adapter actually reads — while an unregistered
+    ``(module, name)`` pair is refused and handed back as its raw arguments. LangGraph offers
+    the same lockdown through ``LANGGRAPH_STRICT_MSGPACK``, but an environment variable is not
+    a guarantee a library can make about itself.
+    """
+    from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+    return JsonPlusSerializer(allowed_msgpack_modules=())
+
+
+def ingest_snapshot(source: Path, thread: str, *, error_channel: str, timeout: float = 30) -> RawTrace:
+    from langgraph.checkpoint.sqlite import SqliteSaver
+
+    deadline = time.monotonic() + timeout
+
+    def progress(status: int, remaining: int, total: int) -> None:
+        if time.monotonic() >= deadline:
+            raise ValueError("SQLite snapshot deadline exceeded")
+
+    try:
+        with tempfile.TemporaryDirectory(prefix="tracegraph-snapshot-") as directory:
+            snapshot = Path(directory) / "checkpoints.sqlite"
+            src = sqlite3.connect(source.resolve().as_uri() + "?mode=ro", timeout=timeout, uri=True)
+            try:
+                dest = sqlite3.connect(snapshot)
+                try:
+                    src.backup(dest, pages=256, progress=progress, sleep=0.01)
+                finally:
+                    dest.close()
+            finally:
+                src.close()
+            # Not `from_conn_string`: it does not take a serializer, and the default one
+            # will call whatever the payload names. `check_same_thread=False` mirrors what
+            # that helper does.
+            with closing(sqlite3.connect(str(snapshot), check_same_thread=False)) as conn:
+                saver = SqliteSaver(conn, serde=_inert_serde())
+                return LangGraphCheckpointAdapter(saver, error_channel=error_channel).ingest(thread)
+    except sqlite3.Error as exc:
+        raise ValueError("cannot read SQLite checkpoint snapshot") from exc

tracegraph/store/__init__.py

@@ -0,0 +1,22 @@
+"""Pluggable query backends. The pure-Python ``InMemoryStore`` is the default.
+
+The optional :class:`LadybugStore` lives behind the ``[cypher]`` extra and is exported here
+**only** when ``ladybug`` is importable, so a default install never pays the import cost (or
+errors on a missing optional dep).
+"""
+
+from tracegraph.store.base import GraphStore
+from tracegraph.store.in_memory import InMemoryStore
+
+__all__ = ["GraphStore", "InMemoryStore"]
+
+try:
+    from tracegraph.store.ladybug import LadybugStore  # noqa: F401
+except ModuleNotFoundError as exc:
+    # Swallow ONLY "the optional dep itself is missing"; anything else (a broken ladybug
+    # install, a typo in our own imports, a removed internal module) must surface as a
+    # real error, not silently hide LadybugStore from the public API.
+    if exc.name != "ladybug":
+        raise
+else:
+    __all__.append("LadybugStore")

tracegraph/store/base.py

@@ -0,0 +1,43 @@
+"""The pluggable query backend interface — deliberately small.
+
+The store does **not** expose a ``query(cypher)`` method: that would leak Cypher
+upward and make the pure-Python path a second-class citizen. Instead, analyses are
+expressed as *operations* the store implements; each backend supplies its own
+implementation (pure-Python traversal here; compiled Cypher in the optional LadybugDB
+backend for pattern matching).
+
+Pattern matching intentionally stays out of the minimal ``GraphStore`` protocol. It is a
+separate backend capability (pure-Python functions for the default path; ``LadybugStore`` adds
+``find_matches``) so RCA/diff callers do not depend on a Cypher-capable backend.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Protocol, runtime_checkable
+
+from tracegraph.model import Edge, NormalizedTrace, Step
+
+
+@runtime_checkable
+class GraphStore(Protocol):
+    def init_schema(self) -> None: ...
+
+    def upsert_nodes(self, steps: list[Step]) -> None: ...
+
+    def upsert_edges(self, edges: list[Edge]) -> None:
+        """Store ALL edges, including every raw CAUSED_BY (not just the tree parent)."""
+        ...
+
+    def ancestors(self, step_id: str) -> list[Step]:
+        """Raw CAUSED_BY traversal from ``step_id`` toward root causes (the RCA primitive)."""
+        ...
+
+    def export_artifact(self, path: str | Path) -> None: ...
+
+    @classmethod
+    def load_artifact(cls, path: str | Path) -> "GraphStore": ...
+
+    def trace(self) -> NormalizedTrace:
+        """Return the loaded trace as a NormalizedTrace (for rendering/diffing)."""
+        ...

tracegraph/store/in_memory.py

@@ -0,0 +1,94 @@
+"""Pure-Python, in-process graph store — the default backend.
+
+Zero dependencies beyond the model, no server, no optional deps. It holds the
+normalized trace in memory and answers the raw-graph traversal that ``explain``/RCA
+need. The portable artifact remains the source of record; this store is rebuilt from
+it on load.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from pathlib import Path
+
+from tracegraph import artifact
+from tracegraph.model import Edge, EdgeType, NormalizedTrace, RawTrace, Step, Trace
+from tracegraph.normalize import normalize, validate_normalized
+
+
+class InMemoryStore:
+    def __init__(self) -> None:
+        self._trace: Trace | None = None
+        self._steps: dict[str, Step] = {}
+        self._edges: list[Edge] = []
+        # adjacency for raw causal edges: effect -> [causes]
+        self._caused_by: dict[str, list[str]] = {}
+
+    # --- construction ---
+
+    @classmethod
+    def from_trace(cls, nt: NormalizedTrace) -> "InMemoryStore":
+        """Load an already-normalized trace, validating both edge layers at the boundary."""
+        validate_normalized(nt)
+        store = cls()
+        store.init_schema()
+        store._trace = nt.trace.model_copy(deep=True)
+        store.upsert_nodes(nt.steps)
+        store.upsert_edges(nt.edges)
+        return store
+
+    @classmethod
+    def from_raw(cls, raw: RawTrace) -> "InMemoryStore":
+        """Normalize a raw adapter output and load it. The intended ingestion path."""
+        return cls.from_trace(normalize(raw))
+
+    @classmethod
+    def load_artifact(cls, path: str | Path) -> "InMemoryStore":
+        return cls.from_trace(artifact.load(path))
+
+    # --- GraphStore protocol ---
+
+    def init_schema(self) -> None:  # nothing to do for an in-memory store
+        pass
+
+    def upsert_nodes(self, steps: list[Step]) -> None:
+        for s in steps:
+            self._steps[s.step_id] = s.model_copy(deep=True)
+
+    def upsert_edges(self, edges: list[Edge]) -> None:
+        self._edges.extend(e.model_copy(deep=True) for e in edges)
+        for e in edges:
+            if e.type is EdgeType.CAUSED_BY:
+                self._caused_by.setdefault(e.src, []).append(e.dst)
+
+    def ancestors(self, step_id: str) -> list[Step]:
+        """All raw causal ancestors of ``step_id``, nearest cause first (BFS over CAUSED_BY)."""
+        if step_id not in self._steps:
+            raise KeyError(f"unknown step {step_id!r}")
+        out: list[Step] = []
+        seen: set[str] = {step_id}
+        queue: deque[str] = deque(self._caused_by.get(step_id, []))
+        while queue:
+            cur = queue.popleft()
+            if cur in seen:
+                continue
+            seen.add(cur)
+            if cur not in self._steps:
+                # A CAUSED_BY edge pointing at an unknown step means a corrupted raw
+                # graph. explain/RCA must surface that, not silently drop the cause.
+                raise KeyError(f"CAUSED_BY edge points to unknown step {cur!r}")
+            out.append(self._steps[cur].model_copy(deep=True))
+            queue.extend(self._caused_by.get(cur, []))
+        return out
+
+    def export_artifact(self, path: str | Path) -> None:
+        artifact.save_atomic(self.trace(), path)
+
+    def trace(self) -> NormalizedTrace:
+        if self._trace is None:
+            raise RuntimeError("store has no trace loaded")
+        return NormalizedTrace(
+            trace=self._trace,
+            steps=list(self._steps.values()),
+            edges=list(self._edges),
+        ).model_copy(deep=True)