aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/services/state_stores/json_store.py

@@ -0,0 +1,67 @@
+import asyncio
+import json
+import os
+import threading
+import time
+
+from aethergraph.contracts.services.state_stores import GraphSnapshot, GraphStateStore, StateEvent
+
+
+class JsonGraphStateStore(GraphStateStore):
+    def __init__(self, root: str):
+        self.root = root
+        os.makedirs(root, exist_ok=True)
+        self._alock = asyncio.Lock()
+        self._tlock = threading.RLock()
+
+    def _run_dir(self, run_id: str) -> str:
+        d = os.path.join(self.root, run_id)
+        os.makedirs(d, exist_ok=True)
+        return d
+
+    async def save_snapshot(self, snap: GraphSnapshot) -> None:
+        d = self._run_dir(snap.run_id)
+        ts = int(time.time())
+        fn = f"snapshot_{snap.rev:08d}_{ts}.json"
+        tmp = os.path.join(d, fn + ".tmp")
+        dst = os.path.join(d, fn)
+        with self._tlock:  # <— thread-safe region
+            with open(tmp, "w", encoding="utf-8") as f:
+                json.dump(snap.__dict__, f, ensure_ascii=False)
+                f.flush()
+                os.fsync(f.fileno())
+            os.replace(tmp, dst)
+
+    async def load_latest_snapshot(self, run_id: str) -> GraphSnapshot | None:
+        d = self._run_dir(run_id)
+        with self._tlock:
+            files = [x for x in os.listdir(d) if x.startswith("snapshot_")]
+            if not files:
+                return None
+            files.sort()
+            with open(os.path.join(d, files[-1]), encoding="utf-8") as f:
+                return GraphSnapshot(**json.load(f))
+
+    async def append_event(self, ev: StateEvent) -> None:
+        p = os.path.join(self._run_dir(ev.run_id), "events.jsonl")
+        line = json.dumps(ev.__dict__, ensure_ascii=False) + "\n"
+        with self._tlock, open(p, "a", encoding="utf-8") as f:
+            f.write(line)
+            f.flush()
+            os.fsync(f.fileno())
+
+    async def load_events_since(self, run_id: str, from_rev: int) -> list[StateEvent]:
+        p = os.path.join(self._run_dir(run_id), "events.jsonl")
+        if not os.path.exists(p):
+            return []
+        out = []
+        with open(p, encoding="utf-8") as f:
+            for line in f:
+                rec = json.loads(line)
+                if rec["rev"] > from_rev:
+                    out.append(StateEvent(**rec))
+        return out
+
+    async def list_run_ids(self, graph_id: str | None = None) -> list[str]:
+        # best-effort: return all directories; filter by graph_id by reading latest snapshot if needed
+        return [d for d in os.listdir(self.root) if os.path.isdir(os.path.join(self.root, d))]

aethergraph/services/state_stores/resume_policy.py

@@ -0,0 +1,119 @@
+# aethergraph/core/persist/resume_policy.py
+from __future__ import annotations
+
+from typing import Any
+
+from aethergraph.contracts.errors.errors import ResumeIncompatibleSnapshot
+
+_JSON_PRIMITIVES = (str, int, float, bool, type(None))
+
+
+def _is_json_like(x: Any) -> bool:
+    if isinstance(x, _JSON_PRIMITIVES):
+        return True
+    if isinstance(x, list):
+        return all(_is_json_like(v) for v in x)
+    if isinstance(x, dict):
+        # Treat any dict that has __aether_ref__ as NOT allowed in strict JSON-only policy
+        if "__aether_ref__" in x:
+            return False
+        return all(isinstance(k, str) and _is_json_like(v) for k, v in x.items())
+    return False
+
+
+def _walk_non_json(obj, path="$"):
+    from collections.abc import Mapping, Sequence
+
+    if isinstance(obj, str | int | float | bool) or obj is None:
+        return
+    if isinstance(obj, Mapping):
+        if "__aether_ref__" in obj:
+            yield path
+            return
+        for k, v in obj.items():
+            yield from _walk_non_json(v, f"{path}.{k}")
+        return
+    if isinstance(obj, Sequence) and not isinstance(obj, str | bytes | bytearray):
+        for i, v in enumerate(obj):
+            yield from _walk_non_json(v, f"{path}[{i}]")
+        return
+    yield path  # non-JSON leaf
+
+
+def assert_snapshot_json_only(
+    run_id: str,
+    snap_json: dict,
+    *,
+    mode: str = "reuse_only",  # "strict" | "reuse_only"
+    ignore_nodes: set[str] | None = None,  # node_ids to skip (e.g., graph output producers)
+) -> None:
+    """
+    - mode="strict": scan ALL nodes; forbid any non-JSON/ref anywhere.
+    - mode="reuse_only": ONLY check nodes whose outputs reuse (status DONE/SKIPPED).
+    - ignore_nodes: always skip these node_ids (e.g., final/sink nodes that feed graph outputs).
+    """
+    ignore_nodes = ignore_nodes or set()
+    reasons: list[str] = []
+
+    state = snap_json.get("state") or snap_json
+    nodes = state.get("nodes", {})
+
+    def _should_check(nid: str, ns: dict) -> bool:
+        if nid in ignore_nodes:
+            return False
+        if mode == "strict":
+            return True
+        st = (ns.get("status") or "").upper()
+        return st in {"DONE", "SKIPPED"}
+
+    for nid, ns in nodes.items():
+        if not _should_check(nid, ns):
+            continue
+        outs = ns.get("outputs")
+        if not outs:
+            continue
+        bad_paths = list(_walk_non_json(outs))
+        if bad_paths:
+            reasons.append(
+                f"node '{nid}' outputs contain non-JSON or refs at: "
+                + ", ".join(bad_paths[:8])
+                + (" ..." if len(bad_paths) > 8 else "")
+            )
+
+    if reasons:
+        raise ResumeIncompatibleSnapshot(run_id, reasons)
+
+
+def output_node_ids_from_graph(graph) -> set[str]:
+    """
+    Collect node_ids that directly feed the declared graph outputs.
+    Expected binding shape:
+      {'sum': {'_type': 'ref', 'from': 'combine_3', 'key': 'sum'}}
+    """
+    try:
+        bindings = graph.io_signature().get("outputs", {}).get("bindings", {}) or {}
+    except Exception:
+        bindings = {}
+
+    out_nodes: set[str] = set()
+
+    def _collect(ref):
+        # Handle canonical case
+        if isinstance(ref, dict) and ref.get("_type") == "ref":
+            src = ref.get("from")
+            if isinstance(src, str) and src:
+                out_nodes.add(src)
+            return
+
+        # Be forgiving if future formats appear (list of refs, nested, etc.)
+        if isinstance(ref, dict):
+            for v in ref.values():
+                _collect(v)
+        elif isinstance(ref, list | tuple):
+            for v in ref:
+                _collect(v)
+
+    for v in bindings.values():
+        _collect(v)
+
+    return out_nodes

aethergraph/services/state_stores/serialize.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import asdict
+from typing import Any
+
+from .externalize import externalize_to_artifact
+
+_JSON_SCALARS = (str, int, float, bool, type(None))
+
+
+def is_json_pure(obj: Any) -> bool:
+    if isinstance(obj, _JSON_SCALARS):
+        return True
+    if isinstance(obj, Mapping):
+        # disallow references in "pure" mode
+        if "__aether_ref__" in obj or obj.get("__externalized__") is True:
+            return False
+        for k, v in obj.items():
+            if not isinstance(k, str):
+                return False
+            if not is_json_pure(v):
+                return False
+        return True
+    if isinstance(obj, Sequence) and not isinstance(obj, str | bytes | bytearray):
+        return all(is_json_pure(v) for v in obj)
+    return False
+
+
+def jsonish_or_ref(
+    obj: Any,
+    *,
+    mk_ref: callable[[Any], dict[str, Any]] | None = None,
+) -> tuple[Any, dict[str, Any] | None]:
+    """
+    Returns (jsonish, ref) where:
+      - if obj is JSON-pure → (obj, None)
+      - else → ({"__aether_ref__": <...>}, ref_meta)
+    mk_ref(obj) must return a dict with at least {"__aether_ref__": "<uri-or-id>"}.
+    """
+    if is_json_pure(obj):
+        return obj, None
+
+    if mk_ref is None:
+        # Default: opaque marker (will block resume; artifacts still saved by caller if desired)
+        return {"__aether_ref__": "opaque:nonjson"}, {"__aether_ref__": "opaque:nonjson"}
+
+    ref = mk_ref(obj) or {"__aether_ref__": "opaque:nonjson"}
+    return {
+        "__aether_ref__": ref.get("__aether_ref__", "opaque:nonjson"),
+        **{k: v for k, v in ref.items() if k != "__aether_ref__"},
+    }, ref
+
+
+def map_jsonish_or_ref(
+    payload: Any,
+    *,
+    mk_ref: callable[[Any], dict[str, Any]] | None = None,
+) -> tuple[Any, bool]:
+    """
+    Walk nested structures. Returns (jsonish_payload, had_refs).
+    Any non-JSON leaf becomes a {"__aether_ref__": ...} marker via mk_ref.
+    """
+    if is_json_pure(payload):
+        return payload, False
+
+    # dict
+    if isinstance(payload, Mapping):
+        out = {}
+        had_ref = False
+        for k, v in payload.items():
+            if not isinstance(k, str):
+                # stringify non-string keys to keep snapshot JSON-safe
+                k = str(k)
+            vv, r = map_jsonish_or_ref(v, mk_ref=mk_ref)
+            out[k] = vv
+            had_ref = had_ref or r
+        return out, had_ref
+
+    # list/tuple
+    if isinstance(payload, Sequence) and not isinstance(payload, str | bytes | bytearray):
+        out = []
+        had_ref = False
+        for v in payload:
+            vv, r = map_jsonish_or_ref(v, mk_ref=mk_ref)
+            out.append(vv)
+            had_ref = had_ref or r
+        return out, had_ref
+
+    # leaf non-JSON → ref
+    jsonish, _ref_meta = jsonish_or_ref(payload, mk_ref=mk_ref)
+    return jsonish, True
+
+
+async def _externalize_leaf_to_artifact(
+    obj: Any,
+    *,
+    run_id: str,
+    graph_id: str,
+    node_id: str,
+    tool_name: str | None,
+    tool_version: str | None,
+    artifacts,
+):
+    ref = await externalize_to_artifact(
+        obj,
+        run_id=run_id,
+        graph_id=graph_id,
+        node_id=node_id,
+        tool_name=tool_name,
+        tool_version=tool_version,
+        artifacts=artifacts,
+    )
+    return ref
+
+
+async def _jsonish_outputs_with_refs(
+    *,
+    outputs: dict[str, Any] | None,
+    run_id: str,
+    graph_id: str,
+    node_id: str,
+    tool_name: str | None,
+    tool_version: str | None,
+    artifacts,  # AsyncArtifactStore or None
+    allow_externalize: bool,  # toggle
+) -> dict[str, Any] | None:
+    if outputs is None:
+        return None
+
+    def mk_ref(obj):
+        # If we can't (or shouldn't) externalize, mark opaque
+        if not allow_externalize or artifacts is None:
+            return {"__aether_ref__": "opaque:nonjson"}
+        # We'll externalize synchronously at call site; placeholder here—replaced below
+        return {"__aether_ref__": "pending:externalize"}
+
+    # First pass: mark structure with refs
+    jsonish, had_refs = map_jsonish_or_ref(outputs, mk_ref=mk_ref)
+
+    if not had_refs or not allow_externalize or artifacts is None:
+        # Nothing to externalize (or disabled) → return as-is
+        return jsonish
+
+    # Second pass: walk and replace "pending:externalize" leaves with real URIs
+    async def _resolve_refs(x):
+        if isinstance(x, dict):
+            if x.get("__aether_ref__") == "pending:externalize" and len(x) == 1:
+                # We need original object to externalize; this simple walker can’t see it anymore.
+                # Approach: re-walk original outputs in parallel to find leaves that caused refs.
+                # For simplicity & performance, do a single-pass direct externalization below instead.
+                return x
+            return {k: await _resolve_refs(v) for k, v in x.items()}
+        if isinstance(x, list):
+            return [await _resolve_refs(v) for v in x]
+        return x
+
+    # Simpler approach: do a second real pass that externalizes by diffing original leaves.
+    # Implement a focused externalizer that walks original outputs again.
+    async def _externalize_in_place(orig):
+        # returns jsonish w/ real refs
+        from collections.abc import Mapping, Sequence
+
+        if isinstance(orig, str | int | float | bool | type(None)):
+            return orig
+        if isinstance(orig, Mapping):
+            out = {}
+            for k, v in orig.items():
+                out[str(k)] = await _externalize_in_place(v)
+            return out
+        if isinstance(orig, Sequence) and not isinstance(orig, str | bytes | bytearray):
+            return [await _externalize_in_place(v) for v in orig]
+
+        # leaf non-JSON → actual artifact ref
+        ref = await _externalize_leaf_to_artifact(
+            orig,
+            run_id=run_id,
+            graph_id=graph_id,
+            node_id=node_id,
+            tool_name=tool_name,
+            tool_version=tool_version,
+            artifacts=artifacts,
+        )
+        return {
+            "__aether_ref__": ref["__aether_ref__"],
+            **{k: v for k, v in ref.items() if k != "__aether_ref__"},
+        }
+
+    return await _externalize_in_place(outputs)
+
+
+async def state_to_json_safe(
+    state_obj,
+    *,
+    run_id: str,
+    graph_id: str,
+    artifacts=None,
+    allow_externalize: bool = False,  # Do not externalize by default until fixing artifacts writer
+    include_wait_spec: bool = True,
+) -> dict[str, Any]:
+    """
+    Convert TaskGraphState to a JSON-safe dict.
+    - JSON outputs inlined
+    - non-JSON leaves become {"__aether_ref__": "..."} and (optionally) are externalized to artifacts
+    """
+    nodes_block = {}
+    for nid, ns in state_obj.nodes.items():
+        status = getattr(ns, "status", None)
+        status_name = getattr(status, "name", status)  # Enum.name or string
+        tool_name = getattr(ns, "tool_name", None) or getattr(
+            getattr(ns, "spec", None), "tool_name", None
+        )
+        tool_version = getattr(ns, "tool_version", None) or getattr(
+            getattr(ns, "spec", None), "tool_version", None
+        )
+
+        outputs_json = await _jsonish_outputs_with_refs(
+            outputs=getattr(ns, "outputs", None),
+            run_id=run_id,
+            graph_id=graph_id,
+            node_id=nid,
+            tool_name=tool_name,
+            tool_version=tool_version,
+            artifacts=artifacts,
+            allow_externalize=allow_externalize,
+        )
+
+        entry = {
+            "status": status_name,
+            "outputs": outputs_json,
+            "error": getattr(ns, "error", None),
+            "attempts": getattr(ns, "attempts", 0),
+            "next_wakeup_at": getattr(ns, "next_wakeup_at", None),
+            "wait_token": getattr(ns, "wait_token", None),
+        }
+        if include_wait_spec:
+            ws = getattr(ns, "wait_spec", None)
+            if ws:
+                # ensure JSON-safe wait_spec (it should be strings/lists/dicts already)
+                entry["wait_spec"] = ws
+        nodes_block[nid] = entry
+
+    return {
+        "run_id": getattr(state_obj, "run_id", run_id),
+        "rev": getattr(state_obj, "rev", None),
+        "patches": [asdict(p) for p in getattr(state_obj, "patches", [])],
+        "_bound_inputs": getattr(state_obj, "_bound_inputs", None),
+        "nodes": nodes_block,
+    }

aethergraph/services/state_stores/utils.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from dataclasses import asdict, is_dataclass
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+from aethergraph.contracts.services.state_stores import GraphSnapshot
+
+from .serialize import state_to_json_safe
+
+
+async def snapshot_from_graph(
+    run_id: str,
+    graph_id: str,
+    rev: int,
+    spec_hash: str,
+    state_obj,
+    *,
+    artifacts=None,  # AsyncArtifactStore or None
+    allow_externalize: bool = False,
+    include_wait_spec: bool = True,
+):
+    json_state = await state_to_json_safe(
+        state_obj,
+        run_id=run_id,
+        graph_id=graph_id,
+        artifacts=artifacts,
+        allow_externalize=allow_externalize,
+        include_wait_spec=include_wait_spec,
+    )
+    return GraphSnapshot(
+        run_id=run_id,
+        graph_id=graph_id,
+        rev=rev,
+        created_at=datetime.utcnow().timestamp(),
+        spec_hash=spec_hash,
+        state=json_state,
+    )
+
+
+def _status_to_str(s) -> str:
+    if s is None:
+        return "PENDING"
+    if isinstance(s, Enum):
+        return s.name
+    # already a string or something printable
+    return str(s)
+
+
+def _enum_name_or_str(x):
+    try:
+        return x.name  # Enum
+    except AttributeError:
+        return str(x)
+
+
+def _state_to_json(state_obj) -> dict[str, Any]:
+    return {
+        "run_id": getattr(state_obj, "run_id", None),
+        "rev": getattr(state_obj, "rev", None),
+        "patches": [asdict(p) if is_dataclass(p) else p for p in getattr(state_obj, "patches", [])],
+        "_bound_inputs": getattr(state_obj, "_bound_inputs", None),
+        "nodes": {
+            nid: {
+                "status": _enum_name_or_str(getattr(ns, "status", "PENDING")),
+                "outputs": getattr(ns, "outputs", None),
+                "error": getattr(ns, "error", None),
+                "attempts": getattr(ns, "attempts", 0),
+                "next_wakeup_at": getattr(ns, "next_wakeup_at", None),
+                "wait_token": getattr(ns, "wait_token", None),
+                # NEW: safe subset of wait_spec (avoid inline payload & tokens)
+                "wait_spec": _sanitize_wait_spec(getattr(ns, "wait_spec", None)),
+            }
+            for nid, ns in getattr(state_obj, "nodes", {}).items()
+        },
+    }
+
+
+def _sanitize_wait_spec(ws):
+    if not ws:
+        return None
+    # Drop volatile/sensitive fields if present
+    return {
+        "kind": ws.get("kind"),
+        "channel": ws.get("channel"),
+        "prompt": ws.get("prompt"),
+        "options": ws.get("options"),
+        "meta": ws.get("meta") or {},
+        # DO NOT store: token / inline_payload / resume_schema with secrets
+    }

aethergraph/services/state_stores/validate.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+
+class ResumptionNotSupported(Exception):
+    """Raised when a snapshot contains non-JSON-pure outputs and resume is disabled."""
+
+    pass
+
+
+_JSON_SCALARS = (str, int, float, bool, type(None))
+
+
+def _is_json_pure(obj: Any) -> bool:
+    """
+    Strict check: only JSON scalars, lists/tuples, and dicts with string keys.
+    No custom objects, no bytes, no externalization markers.
+    """
+    if isinstance(obj, _JSON_SCALARS):
+        return True
+
+    if isinstance(obj, Mapping):
+        # Disallow future externalization markers proactively
+        if "__aether_ref__" in obj or obj.get("__externalized__") is True:
+            return False
+        # JSON requires string keys
+        for k, v in obj.items():
+            if not isinstance(k, str):
+                return False
+            if not _is_json_pure(v):
+                return False
+        return True
+
+    if isinstance(obj, Sequence) and not isinstance(obj, (str | bytes | bytearray)):
+        return all(_is_json_pure(v) for v in obj)
+
+    return False
+
+
+def assert_snapshot_json_pure(
+    snapshot_state: dict, *, run_id: str, graph_id: str, allow_non_json: bool = False
+) -> None:
+    """
+    Validate the *serialized* (dict) snapshot state produced by snapshot_from_graph(...).
+    If any node outputs are not strictly JSON-pure, raise ResumptionNotSupported
+    (unless allow_non_json=True).
+    """
+    if allow_non_json:
+        return
+
+    if not isinstance(snapshot_state, dict):
+        raise ResumptionNotSupported(
+            f"Resume blocked: snapshot state is not a dict (run_id={run_id}, graph_id={graph_id})."
+        )
+
+    nodes = snapshot_state.get("nodes", {})
+    bad_nodes = []
+
+    for nid, ns in nodes.items():
+        if not isinstance(ns, dict):
+            bad_nodes.append(nid)
+            continue
+        outs = ns.get("outputs", None)
+        if outs is None:
+            continue  # output not produced yet is fine
+        if not _is_json_pure(outs):
+            bad_nodes.append(nid)
+
+    if bad_nodes:
+        listed = ", ".join(bad_nodes)
+        raise ResumptionNotSupported(
+            "Resume blocked: snapshot contains non-JSON outputs; "
+            f"nodes=({listed}) run_id={run_id} graph_id={graph_id}. "
+            "Ensure tools return JSON-friendly outputs (dict/list/str/float/int/bool/null), "
+            "or disable strict resume checks once tested."
+        )

aethergraph/services/tracing/noop.py

@@ -0,0 +1,18 @@
+# services/tracing/noop.py
+import contextlib
+import logging
+import time
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("aethergraph.tracing.noop")
+
+
+class NoopTracer:
+    @contextlib.contextmanager
+    def span(self, name: str, **attrs):
+        t = time.time()
+        yield
+        dt = (time.time() - t) * 1000  # milliseconds
+        # optionally log duration
+        logger.info(f"Span '{name}' took {dt} ms", **attrs)
+        return