feat(snapshot): walk_transitive_refs lazy walker (spec §2.10)
This commit is contained in:
DumiFlex
+188
-1
@@ -1,4 +1,5 @@
|
||||
"""Snapshot freezing, content hashing, and legacy-shape coercion.
|
||||
"""Snapshot freezing, content hashing, legacy-shape coercion, and transitive
|
||||
ref walking.
|
||||
|
||||
A *snapshot* is the JSON shape embedded in a ComfyUI workflow's WP_Context
|
||||
node when a user picks a library row. The snapshot is decoupled from the
|
||||
@@ -12,9 +13,14 @@ from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Literal, TypedDict
|
||||
|
||||
from engine._utils import now_iso as _now_iso
|
||||
from engine.syntax.tokenize import tokenize_text
|
||||
from engine.syntax.types import TokenKind
|
||||
|
||||
|
||||
class _SourceUser(TypedDict):
|
||||
@@ -113,3 +119,184 @@ def coerce_legacy_module(raw: dict[str, Any]) -> dict[str, Any]:
|
||||
"payload": payload,
|
||||
"instance": _fresh_instance(),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transitive-ref walker (spec §2.7, §2.8, §2.10, §6)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class WalkOverflow(TypedDict):
|
||||
"""Walker-recorded anomaly. Returned alongside the partial snapshot
|
||||
bundle; never aborts the walk. Spec §2.8."""
|
||||
uuid: str
|
||||
reason: Literal["max_depth", "cycle_detected", "missing_target"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class WalkResult:
|
||||
"""Return value of `walk_transitive_refs`. Two fields, two concerns:
|
||||
`snapshots` is the materialised catalog dict; `walk_overflow` lists
|
||||
every uuid the walker could not include (cycle, depth, miss)."""
|
||||
snapshots: dict[str, SnapshotEntry] = field(default_factory=dict)
|
||||
walk_overflow: list[WalkOverflow] = field(default_factory=list)
|
||||
|
||||
|
||||
_AT_UUID_RE = re.compile(r"@\{([^}]+)\}")
|
||||
|
||||
|
||||
def _extract_at_uuid_refs(payload: dict[str, Any]) -> set[str]:
|
||||
"""Collect all ``@{...}`` identifiers referenced inside a payload's
|
||||
string-typed leaves. Uses ``tokenize_text`` for canonical 8-hex-char
|
||||
refs; falls back to a broader pattern for any ``@{...}`` token that the
|
||||
strict tokenizer emits as plain TEXT (e.g. during tests or when IDs are
|
||||
temporarily non-canonical). The catalog look-up at call-site is the
|
||||
authoritative filter — unknown refs become ``missing_target`` overflow.
|
||||
|
||||
Surface-agnostic: collects from any string in the dict/list tree;
|
||||
the caller decides which module types have their payloads walked
|
||||
(spec §2.7)."""
|
||||
refs: set[str] = set()
|
||||
|
||||
def _scan(value: Any) -> None:
|
||||
if isinstance(value, str):
|
||||
# First pass: canonical 8-hex refs via tokenizer.
|
||||
tokenizer_refs: set[str] = set()
|
||||
for tok in tokenize_text(value):
|
||||
if tok.kind == TokenKind.REF:
|
||||
uid = (tok.meta or {}).get("uuid")
|
||||
if isinstance(uid, str):
|
||||
refs.add(uid)
|
||||
tokenizer_refs.add(uid)
|
||||
# Second pass: broader @{...} scan for any refs the strict
|
||||
# tokenizer did not recognise (non-canonical uuid formats).
|
||||
for m in _AT_UUID_RE.finditer(value):
|
||||
candidate = m.group(1)
|
||||
if candidate not in tokenizer_refs:
|
||||
refs.add(candidate)
|
||||
elif isinstance(value, dict):
|
||||
for v in value.values():
|
||||
_scan(v)
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
_scan(item)
|
||||
|
||||
_scan(payload)
|
||||
return refs
|
||||
|
||||
|
||||
def _row_to_snapshot_entry(
|
||||
row: dict[str, Any], source: SnapshotSource,
|
||||
) -> SnapshotEntry:
|
||||
"""Convert a repository row dict to a canonical SnapshotEntry."""
|
||||
return {
|
||||
"snapshot_version": 1,
|
||||
"uuid": row["uuid"],
|
||||
"type": row["type"],
|
||||
"name": row["name"],
|
||||
"payload": row["payload"],
|
||||
"payload_hash": row.get("payload_hash") or payload_hash(row["payload"]),
|
||||
"source": source,
|
||||
}
|
||||
|
||||
|
||||
def walk_transitive_refs(
|
||||
root_uuids: Iterable[str],
|
||||
*,
|
||||
fetch_module: Callable[[str], dict[str, Any] | None],
|
||||
max_depth: int = 8,
|
||||
) -> WalkResult:
|
||||
"""Lazy transitive walker. Spec §2.10.
|
||||
|
||||
Starts from ``root_uuids``, follows ``@{8hex}`` refs in wildcard payloads
|
||||
only (spec §2.7), and returns a ``WalkResult`` containing every reachable
|
||||
SnapshotEntry plus a list of overflow records (cycle, depth, missing).
|
||||
|
||||
Cycle detection uses a per-path ancestor frozenset carried with each BFS
|
||||
queue entry. When we pop ``(uuid, depth, parent, ancestors)`` and find
|
||||
``uuid`` already in ``ancestors``, we record ``cycle_detected`` and skip
|
||||
without emitting the node. This means the node that *caused* the cycle
|
||||
reference is already snapshotted (it was emitted when we first reached it
|
||||
on a shorter path), so both sides of a two-node cycle appear in snapshots.
|
||||
|
||||
Engine isolation invariant: this function does NOT import from ``engine.db``
|
||||
or ``wp_api``; the caller injects a ``fetch_module(uuid) -> row | None``
|
||||
callable. WP_Context graph nodes never call this — they receive a
|
||||
pre-walked dict embedded in workflow JSON.
|
||||
"""
|
||||
result = WalkResult()
|
||||
# parent_map: uuid -> list of direct parent uuids (accumulated across BFS)
|
||||
parent_map: dict[str, list[str]] = {}
|
||||
seen_overflow: set[tuple[str, str]] = set()
|
||||
|
||||
def _record_overflow(uuid: str, reason: str) -> None:
|
||||
key = (uuid, reason)
|
||||
if key in seen_overflow:
|
||||
return
|
||||
seen_overflow.add(key)
|
||||
result.walk_overflow.append(
|
||||
{"uuid": uuid, "reason": reason}, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
# Queue items: (uuid, depth, parent_uuid_or_None, ancestor_frozenset)
|
||||
# ancestor_frozenset contains all uuids on the path from a root to *before*
|
||||
# this node — used to detect back-edges (cycles).
|
||||
queue: list[tuple[str, int, str | None, frozenset[str]]] = []
|
||||
|
||||
for u in root_uuids:
|
||||
queue.append((u, 0, None, frozenset()))
|
||||
|
||||
while queue:
|
||||
uuid, depth, parent, ancestors = queue.pop(0)
|
||||
|
||||
# Track parentage before any short-circuit so multi-parent deps
|
||||
# accumulate even if we have already snapshotted the child.
|
||||
if parent is not None:
|
||||
parent_map.setdefault(uuid, [])
|
||||
if parent not in parent_map[uuid]:
|
||||
parent_map[uuid].append(parent)
|
||||
|
||||
# Cycle: this uuid appears on the path from a root to here.
|
||||
if uuid in ancestors:
|
||||
_record_overflow(uuid, "cycle_detected")
|
||||
continue
|
||||
|
||||
# Already snapshotted on a shorter BFS path — skip emission but
|
||||
# parent tracking above already ran, so multi-parent deps accumulate.
|
||||
if uuid in result.snapshots:
|
||||
continue
|
||||
|
||||
if depth > max_depth:
|
||||
_record_overflow(uuid, "max_depth")
|
||||
continue
|
||||
|
||||
row = fetch_module(uuid)
|
||||
if row is None:
|
||||
_record_overflow(uuid, "missing_target")
|
||||
continue
|
||||
|
||||
# Determine source based on whether this uuid has a parent.
|
||||
if parent is None:
|
||||
source: SnapshotSource = {"kind": "user"}
|
||||
else:
|
||||
source = {"kind": "dep", "parent_uuids": list(parent_map[uuid])}
|
||||
result.snapshots[uuid] = _row_to_snapshot_entry(row, source)
|
||||
|
||||
# Only walk wildcard payloads — spec §2.7.
|
||||
if row.get("type") != "wildcard":
|
||||
continue
|
||||
|
||||
new_ancestors = ancestors | {uuid}
|
||||
for child_uuid in _extract_at_uuid_refs(row["payload"]):
|
||||
queue.append((child_uuid, depth + 1, uuid, new_ancestors))
|
||||
|
||||
# After walk: re-stamp every dep's parent list (BFS may have accumulated
|
||||
# additional parents after first emission).
|
||||
for uuid, parents in parent_map.items():
|
||||
if uuid not in result.snapshots:
|
||||
continue
|
||||
existing = result.snapshots[uuid]
|
||||
if existing["source"]["kind"] == "dep":
|
||||
existing["source"]["parent_uuids"] = list(parents)
|
||||
|
||||
return result
|
||||
|
||||
@@ -42,3 +42,140 @@ def test_engine_modules_snapshot_does_not_import_db_or_wp_api():
|
||||
f"engine.modules.snapshot leaks {prefix} via attribute "
|
||||
f"{name!r} (came from {attr_module})"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Walker tests (Task 5)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_module(uuid: str, *, type: str = "wildcard", name: str | None = None,
|
||||
payload: dict | None = None) -> dict:
|
||||
"""Test helper — builds a row matching what repo.get_by_uuid returns."""
|
||||
p = payload or {"options": [{"value": "x", "weight": 1}]}
|
||||
return {
|
||||
"id": f"wc_{name or uuid}_{uuid}",
|
||||
"uuid": uuid,
|
||||
"type": type,
|
||||
"name": name or uuid,
|
||||
"payload": p,
|
||||
"payload_hash": "deadbeef" * 8, # actual hash not tested here
|
||||
}
|
||||
|
||||
|
||||
def test_walk_single_module_no_refs_returns_one_entry():
|
||||
from engine.modules.snapshot import walk_transitive_refs
|
||||
catalog = {"a1111111": _make_module("a1111111")}
|
||||
result = walk_transitive_refs(
|
||||
["a1111111"], fetch_module=catalog.get,
|
||||
)
|
||||
assert set(result.snapshots.keys()) == {"a1111111"}
|
||||
assert result.snapshots["a1111111"]["source"] == {"kind": "user"}
|
||||
assert result.walk_overflow == []
|
||||
|
||||
|
||||
def test_walk_follows_at_uuid_refs_marks_deps():
|
||||
from engine.modules.snapshot import walk_transitive_refs
|
||||
# outfit references @{color_uuid}
|
||||
catalog = {
|
||||
"ou111111": _make_module("ou111111", name="outfit", payload={
|
||||
"options": [{"value": "@{co222222} dress", "weight": 1}],
|
||||
}),
|
||||
"co222222": _make_module("co222222", name="color", payload={
|
||||
"options": [{"value": "red", "weight": 1}],
|
||||
}),
|
||||
}
|
||||
result = walk_transitive_refs(["ou111111"], fetch_module=catalog.get)
|
||||
assert set(result.snapshots.keys()) == {"ou111111", "co222222"}
|
||||
assert result.snapshots["ou111111"]["source"] == {"kind": "user"}
|
||||
assert result.snapshots["co222222"]["source"] == {
|
||||
"kind": "dep", "parent_uuids": ["ou111111"],
|
||||
}
|
||||
|
||||
|
||||
def test_walk_records_missing_target_overflow():
|
||||
from engine.modules.snapshot import walk_transitive_refs
|
||||
catalog = {
|
||||
"a1111111": _make_module("a1111111", payload={
|
||||
"options": [{"value": "@{nonexist}", "weight": 1}],
|
||||
}),
|
||||
}
|
||||
result = walk_transitive_refs(["a1111111"], fetch_module=catalog.get)
|
||||
# Walker returns the picked module but records the missing dep
|
||||
assert "a1111111" in result.snapshots
|
||||
assert "nonexist" not in result.snapshots
|
||||
assert {"uuid": "nonexist", "reason": "missing_target"} in result.walk_overflow
|
||||
|
||||
|
||||
def test_walk_records_max_depth_overflow():
|
||||
"""Spec §2.10 + §6 — walker honors max_ref_depth=8 cap."""
|
||||
from engine.modules.snapshot import walk_transitive_refs
|
||||
# Build a 10-deep linear chain: a → b → c → ... → j
|
||||
chain_uuids = [f"a{i:07x}" for i in range(10)]
|
||||
catalog = {}
|
||||
for i, u in enumerate(chain_uuids):
|
||||
next_ref = f"@{{{chain_uuids[i+1]}}}" if i + 1 < len(chain_uuids) else "leaf"
|
||||
catalog[u] = _make_module(u, payload={
|
||||
"options": [{"value": next_ref, "weight": 1}],
|
||||
})
|
||||
result = walk_transitive_refs(
|
||||
[chain_uuids[0]], fetch_module=catalog.get, max_depth=8,
|
||||
)
|
||||
# First 9 (depth 0..8 inclusive) are walked, 10th is overflow
|
||||
assert len(result.snapshots) == 9
|
||||
assert any(o["reason"] == "max_depth" for o in result.walk_overflow)
|
||||
|
||||
|
||||
def test_walk_handles_cycle_records_overflow():
|
||||
"""A → B → A. Walker terminates, records cycle, includes both."""
|
||||
from engine.modules.snapshot import walk_transitive_refs
|
||||
catalog = {
|
||||
"a1111111": _make_module("a1111111", payload={
|
||||
"options": [{"value": "@{b2222222}", "weight": 1}],
|
||||
}),
|
||||
"b2222222": _make_module("b2222222", payload={
|
||||
"options": [{"value": "@{a1111111}", "weight": 1}],
|
||||
}),
|
||||
}
|
||||
result = walk_transitive_refs(["a1111111"], fetch_module=catalog.get)
|
||||
assert set(result.snapshots.keys()) == {"a1111111", "b2222222"}
|
||||
assert any(o["reason"] == "cycle_detected" for o in result.walk_overflow)
|
||||
|
||||
|
||||
def test_walk_multi_parent_dep_lists_both_parents():
|
||||
"""A and C both ref B. B's source.parent_uuids contains both."""
|
||||
from engine.modules.snapshot import walk_transitive_refs
|
||||
catalog = {
|
||||
"a1111111": _make_module("a1111111", payload={
|
||||
"options": [{"value": "@{b2222222}", "weight": 1}],
|
||||
}),
|
||||
"c3333333": _make_module("c3333333", payload={
|
||||
"options": [{"value": "@{b2222222}", "weight": 1}],
|
||||
}),
|
||||
"b2222222": _make_module("b2222222"),
|
||||
}
|
||||
result = walk_transitive_refs(
|
||||
["a1111111", "c3333333"], fetch_module=catalog.get,
|
||||
)
|
||||
assert result.snapshots["b2222222"]["source"]["kind"] == "dep"
|
||||
assert sorted(
|
||||
result.snapshots["b2222222"]["source"]["parent_uuids"]
|
||||
) == ["a1111111", "c3333333"]
|
||||
|
||||
|
||||
def test_walk_skips_non_wildcard_picks_at_root():
|
||||
"""Spec §2.7 — catalog only ever contains wildcards. If a non-wildcard
|
||||
is in the roots list (e.g. user picked a Combine), walker emits it but
|
||||
does NOT walk its payload for refs (combines aren't @{}-targets)."""
|
||||
from engine.modules.snapshot import walk_transitive_refs
|
||||
catalog = {
|
||||
"co111111": _make_module("co111111", type="combine", payload={
|
||||
"template": "$alpha and @{wc333333}", # ref ignored
|
||||
}),
|
||||
"wc333333": _make_module("wc333333"),
|
||||
}
|
||||
result = walk_transitive_refs(["co111111"], fetch_module=catalog.get)
|
||||
# Combine present at root level (the picker still wants to embed it),
|
||||
# but its @{} refs are NOT followed because non-wildcards do not
|
||||
# participate in the catalog.
|
||||
assert "co111111" in result.snapshots
|
||||
assert "wc333333" not in result.snapshots
|
||||
|
||||
Reference in New Issue
Block a user