diff --git a/src/talemate/game/engine/nodes/tools/__init__.py b/src/talemate/game/engine/nodes/tools/__init__.py new file mode 100644 index 00000000..456343a3 --- /dev/null +++ b/src/talemate/game/engine/nodes/tools/__init__.py @@ -0,0 +1,138 @@ +""" +Static analysis, mutation, and layout tools for Talemate node graphs. + +This toolset is primarily aimed at helping **LLM-driven agents** author +and modify Talemate node-graph JSON files safely. Hand-editing these +graphs is error-prone for language models — socket names, property +keys, registry paths, and edge wiring are all easy to typo and produce +silently-broken files. The tools here exist so an agent can: + +* **Analyze** existing graphs (``analysis.py``, exposed via the + ``python -m talemate.game.engine.nodes.tools`` CLI) to understand + structure before modifying anything — ``summary``, ``stages``, + ``node``, ``rtrace``, ``check-registries``, etc. +* **Mutate** graphs programmatically (``writer.py``) via + ``GraphWriter``, which validates socket names and property keys + against the live node-class registry at write time and raises + ``UnknownSocketError`` / ``UnknownPropertyError`` / ``CycleError`` / + etc. rather than letting typos slip through into the JSON. +* **Lay out** graphs deterministically (``layout.py``) so the agent + never has to pick ``(x, y)`` coordinates — a stage-stratified + algorithm handles positioning with predecessor-aware row placement + and height estimation. +* **Load** graph JSON with scene-module auto-registration + (``loader.py``) so sibling scene modules are visible to the writer + and analysis layer without manual registry setup. + +A human working directly in the Talemate visual editor does not need +any of this — the editor handles socket wiring, layout, and validation +natively. These tools exist specifically for the out-of-editor, +machine-authored path. +""" + +from .analysis import ( + ConsumersResult, + EdgeInfo, + FeedsResult, + GraphSummary, + NodeDetails, + NodeListEntry, + RegistryCheckResult, + StageMapResult, + TraceNode, + check_registries, + consumers, + ensure_registry_loaded, + feeds, + get_node, + list_edges, + list_nodes, + resolve_node_id, + stage_map, + summarize, + trace_backward, + trace_forward, +) +from .layout import LayoutError, LayoutOptions, layout_graph +from .loader import GraphLoadError, load_graph, register_scene_modules +from .writer import ( + AlreadyConnectedError, + CycleError, + DynamicInputError, + GraphWriter, + GroupColor, + GroupError, + NodeMetadata, + NodeNotFoundError, + NotConnectedError, + UnknownPropertyError, + UnknownRegistryError, + UnknownSocketError, + WriterError, + add_dynamic_input, + add_group, + add_node, + clear_metadata_cache, + connect, + disconnect, + get_node_metadata, + remove_dynamic_input, + remove_node, +) + +__all__ = [ + # analysis models + "GraphSummary", + "NodeListEntry", + "NodeDetails", + "EdgeInfo", + "FeedsResult", + "ConsumersResult", + "TraceNode", + "StageMapResult", + "RegistryCheckResult", + # analysis functions + "summarize", + "list_nodes", + "get_node", + "list_edges", + "feeds", + "consumers", + "trace_forward", + "trace_backward", + "stage_map", + "check_registries", + "resolve_node_id", + "ensure_registry_loaded", + # loader + "GraphLoadError", + "load_graph", + "register_scene_modules", + # writer + "GraphWriter", + "NodeMetadata", + "GroupColor", + "get_node_metadata", + "clear_metadata_cache", + "add_node", + "remove_node", + "connect", + "disconnect", + "add_dynamic_input", + "remove_dynamic_input", + "add_group", + "WriterError", + "UnknownRegistryError", + "UnknownSocketError", + "UnknownPropertyError", + "AlreadyConnectedError", + "NotConnectedError", + "CycleError", + "DynamicInputError", + "GroupError", + "NodeNotFoundError", + # layout + "layout_graph", + "LayoutError", + "LayoutOptions", +] diff --git a/src/talemate/game/engine/nodes/tools/__main__.py b/src/talemate/game/engine/nodes/tools/__main__.py new file mode 100644 index 00000000..1e4dc20c --- /dev/null +++ b/src/talemate/game/engine/nodes/tools/__main__.py @@ -0,0 +1,8 @@ +"""Allow ``python -m talemate.game.engine.nodes.tools`` to invoke the CLI.""" + +import sys + +from .cli import main + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/talemate/game/engine/nodes/tools/analysis.py b/src/talemate/game/engine/nodes/tools/analysis.py new file mode 100644 index 00000000..67039002 --- /dev/null +++ b/src/talemate/game/engine/nodes/tools/analysis.py @@ -0,0 +1,974 @@ +""" +Read-only static analysis of Talemate node graph JSON dicts. + +Every public function takes an already-loaded ``graph: dict`` (the parsed +contents of a graph .json file) and returns a pydantic model. There is no +file I/O, no formatting, and no runtime instantiation of node classes from +the graph itself - we only consult the live ``NODES`` registry to look up +static socket layouts where available. + +This module is intended as the analysis layer for both a developer CLI +(``tools/cli.py``) and, eventually, a UI surface. Keep it pure. +""" + +from __future__ import annotations + +from collections import Counter, defaultdict +from typing import Any + +import pydantic +import structlog + +log = structlog.get_logger("talemate.game.engine.nodes.tools.analysis") + +__all__ = [ + # registry loader + "ensure_registry_loaded", + # models + "GraphSummary", + "NodeListEntry", + "SocketRef", + "NodeInputInfo", + "NodeOutputInfo", + "NodeDetails", + "EdgeInfo", + "FeedsResult", + "ConsumersResult", + "TraceNode", + "StageInfo", + "StageChain", + "StageMapResult", + "RegistryProblem", + "RegistryCheckResult", + # functions + "summarize", + "list_nodes", + "get_node", + "list_edges", + "feeds", + "consumers", + "trace_forward", + "trace_backward", + "stage_map", + "check_registries", + # helpers + "resolve_node_id", + "split_socket_path", +] + + +# --------------------------------------------------------------------------- +# Registry loader (lazy, side-effecty - only run when actually needed) +# --------------------------------------------------------------------------- + +_REGISTRY_LOADED = False + + +def ensure_registry_loaded() -> None: + """Import all shipped Talemate node definitions exactly once. + + This is lazy because importing the registry has filesystem side effects + (walks ``SEARCH_PATHS``) and we want analysis functions that don't need + the registry to remain cheap. + """ + + global _REGISTRY_LOADED + if _REGISTRY_LOADED: + return + # First import the Python modules that register native node classes + # via @register decorators (core/MakeBool, validation/*, etc). + import talemate.game.engine.nodes.load_definitions # noqa: F401 + from talemate.game.engine.nodes.registry import import_initial_node_definitions + + # Then walk SEARCH_PATHS to register JSON-defined module nodes. + import_initial_node_definitions() + _REGISTRY_LOADED = True + + +def _registry_dict() -> dict[str, Any]: + ensure_registry_loaded() + from talemate.game.engine.nodes.registry import NODES + + return NODES + + +def _base_types_dict() -> dict[str, Any]: + ensure_registry_loaded() + from talemate.game.engine.nodes.base_types import BASE_TYPES + + return BASE_TYPES + + +# --------------------------------------------------------------------------- +# Pydantic models +# --------------------------------------------------------------------------- + + +class GraphSummary(pydantic.BaseModel): + title: str | None + id: str | None + registry: str | None + base_type: str | None + extends: str | None + node_count: int + nodes_by_category: dict[str, int] + stage_node_count: int + input_count: int + output_count: int + module_property_count: int + group_titles: list[str] + + +class NodeListEntry(pydantic.BaseModel): + short_id: str + id: str + registry: str | None + title: str + + +class SocketRef(pydantic.BaseModel): + node_id: str + short_id: str + socket: str + + @property + def full(self) -> str: + return f"{self.node_id}.{self.socket}" + + +class NodeInputInfo(pydantic.BaseModel): + name: str + connected: bool + source: SocketRef | None = None + + +class NodeOutputInfo(pydantic.BaseModel): + name: str + consumers: list[SocketRef] = pydantic.Field(default_factory=list) + + +class NodeDetails(pydantic.BaseModel): + id: str + short_id: str + title: str + registry: str | None + base_type: str | None + properties: dict[str, Any] + x: int = 0 + y: int = 0 + width: int = 0 + height: int = 0 + inputs: list[NodeInputInfo] + outputs: list[NodeOutputInfo] + registered: bool + + +class EdgeInfo(pydantic.BaseModel): + source_node: str + source_short: str + source_socket: str + target_node: str + target_short: str + target_socket: str + + +class FeedsResult(pydantic.BaseModel): + target_node: str + target_short: str + target_socket: str + source: SocketRef | None = None + source_node_title: str | None = None + source_node_registry: str | None = None + + +class ConsumersResult(pydantic.BaseModel): + source_node: str + source_short: str + source_socket: str + consumers: list[EdgeInfo] = pydantic.Field(default_factory=list) + + +class TraceNode(pydantic.BaseModel): + node_id: str + short_id: str + title: str + registry: str | None + via_socket: str | None = None # the socket on *this* node we walked through + children: list["TraceNode"] = pydantic.Field(default_factory=list) + cycle: bool = False + truncated: bool = False # depth limit hit + + +class StageInfo(pydantic.BaseModel): + node_id: str + short_id: str + title: str + stage: int + chain_index: int + + +class StageChain(pydantic.BaseModel): + index: int + node_ids: list[str] + stage_node_ids: list[str] + stages: list[int] + + +class StageMapResult(pydantic.BaseModel): + stage_nodes: list[StageInfo] + chains: list[StageChain] + unstaged_node_ids: list[str] + + +class RegistryProblem(pydantic.BaseModel): + node_id: str + short_id: str + title: str + registry: str | None + base_type: str | None + kind: str # "registry" | "base_type" + + +class RegistryCheckResult(pydantic.BaseModel): + unknown_registries: list[RegistryProblem] + unknown_base_types: list[RegistryProblem] + checked_registry_count: int + checked_base_type_count: int + + @property + def has_problems(self) -> bool: + return bool(self.unknown_registries or self.unknown_base_types) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _short(node_id: str) -> str: + return node_id[:8] + + +def _nodes(graph: dict) -> dict[str, dict]: + return graph.get("nodes") or {} + + +def _edges(graph: dict) -> dict[str, list[str]]: + return graph.get("edges") or {} + + +def split_socket_path(path: str) -> tuple[str, str]: + """Split a ``.`` path on the FIRST dot. + + Node ids (full UUIDs or short prefixes) and socket names in this + codebase never contain dots, so the first dot cleanly separates + the node-id half from the socket-name half. + """ + + if "." not in path: + raise ValueError( + f"Socket path '{path}' must be of the form '.'" + ) + node_part, socket = path.split(".", 1) + return node_part, socket + + +def resolve_node_id(graph: dict, prefix: str) -> str: + """Resolve a possibly-truncated node id to a full node id. + + Accepts a full UUID or any unique prefix. Raises ``ValueError`` on + miss or ambiguity. + """ + + nodes = _nodes(graph) + if not prefix: + raise ValueError("Empty node id prefix") + + if prefix in nodes: + return prefix + + matches = [nid for nid in nodes if nid.startswith(prefix)] + if not matches: + raise ValueError( + f"No node matches prefix '{prefix}'. " + f"(Graph contains {len(nodes)} nodes; try `list-nodes` to find one.)" + ) + if len(matches) > 1: + joined = ", ".join(_short(m) for m in matches[:6]) + more = "" if len(matches) <= 6 else f" (+{len(matches) - 6} more)" + raise ValueError( + f"Ambiguous node prefix '{prefix}' matches {len(matches)} nodes: {joined}{more}" + ) + return matches[0] + + +def _make_socket_ref(node_id: str, socket: str) -> SocketRef: + return SocketRef(node_id=node_id, short_id=_short(node_id), socket=socket) + + +def _registered_node_class(registry: str | None) -> Any | None: + """Return the live class registered under ``registry`` or ``None``.""" + + if not registry: + return None + return _registry_dict().get(registry) + + +def _instance_for_class(node_cls: Any) -> Any | None: + """Try to instantiate a node class to read its socket layout. + + This is a best-effort, side-effect-free probe: registered classes set + up their sockets in ``setup()``, and that ``setup`` runs in + ``__init__``. If anything goes wrong (constructor expects positional + args, etc.), return ``None`` and let the caller fall back. + """ + + if node_cls is None: + return None + try: + return node_cls() + except Exception: + return None + + +def _static_socket_names(registry: str | None) -> tuple[list[str] | None, list[str] | None]: + """Return ``(input_names, output_names)`` from the registered class. + + Each element is ``None`` when we couldn't determine it. + """ + + cls = _registered_node_class(registry) + instance = _instance_for_class(cls) + if instance is None: + return (None, None) + + input_names: list[str] | None + output_names: list[str] | None + try: + input_names = [s.name for s in instance.inputs] + except Exception: + input_names = None + try: + output_names = [s.name for s in instance.outputs] + except Exception: + output_names = None + return input_names, output_names + + +def _node_input_names(graph: dict, node_id: str, node: dict) -> list[str]: + """Best-effort list of input socket names for a node. + + Order of preference: + 1. registered class static inputs (+ dynamic_inputs from JSON) + 2. dynamic_inputs from JSON alone + 3. inferred from edges where this node appears as a target + """ + + static_inputs, _ = _static_socket_names(node.get("registry")) + dyn = [d.get("name") for d in (node.get("dynamic_inputs") or []) if d.get("name")] + inferred = sorted(_inferred_input_names_from_edges(graph, node_id)) + + if static_inputs is not None: + # Combine static + dynamic, then drop duplicates while preserving order. + seen: set[str] = set() + combined: list[str] = [] + for name in [*static_inputs, *dyn, *inferred]: + if name not in seen: + seen.add(name) + combined.append(name) + return combined + + if dyn or inferred: + seen = set() + combined = [] + for name in [*dyn, *inferred]: + if name not in seen: + seen.add(name) + combined.append(name) + return combined + + return [] + + +def _node_output_names(graph: dict, node_id: str, node: dict) -> list[str]: + """Best-effort list of output socket names for a node.""" + + _, static_outputs = _static_socket_names(node.get("registry")) + inferred = sorted(_inferred_output_names_from_edges(graph, node_id)) + + if static_outputs is not None: + seen: set[str] = set() + combined: list[str] = [] + for name in [*static_outputs, *inferred]: + if name not in seen: + seen.add(name) + combined.append(name) + return combined + + return inferred + + +def _inferred_input_names_from_edges(graph: dict, node_id: str) -> set[str]: + names: set[str] = set() + for _src, targets in _edges(graph).items(): + for tgt in targets: + try: + tgt_node, tgt_socket = split_socket_path(tgt) + except ValueError: + continue + if tgt_node == node_id: + names.add(tgt_socket) + return names + + +def _inferred_output_names_from_edges(graph: dict, node_id: str) -> set[str]: + names: set[str] = set() + for src, _targets in _edges(graph).items(): + try: + src_node, src_socket = split_socket_path(src) + except ValueError: + continue + if src_node == node_id: + names.add(src_socket) + return names + + +def _input_source_map(graph: dict) -> dict[tuple[str, str], tuple[str, str]]: + """Map ``(target_node, target_socket) -> (source_node, source_socket)``. + + The graph format treats edges as ``source -> [targets]``, so this + inverts the relation. If a target socket somehow appears with multiple + sources we keep the last one we see (the runtime treats inputs as + single-source). + """ + + out: dict[tuple[str, str], tuple[str, str]] = {} + for src, targets in _edges(graph).items(): + try: + src_node, src_socket = split_socket_path(src) + except ValueError: + continue + for tgt in targets: + try: + tgt_node, tgt_socket = split_socket_path(tgt) + except ValueError: + continue + out[(tgt_node, tgt_socket)] = (src_node, src_socket) + return out + + +# --------------------------------------------------------------------------- +# summarize +# --------------------------------------------------------------------------- + + +def summarize(graph: dict) -> GraphSummary: + nodes = _nodes(graph) + + categories: Counter[str] = Counter() + stage_count = 0 + for node in nodes.values(): + registry = node.get("registry") or "" + if "/" in registry: + category = registry.split("/", 1)[0] + else: + category = registry or "" + categories[category] += 1 + if registry == "core/Stage": + stage_count += 1 + + groups = graph.get("groups") or [] + group_titles = [g.get("title", "") for g in groups] + + return GraphSummary( + title=graph.get("title"), + id=graph.get("id"), + registry=graph.get("registry"), + base_type=graph.get("base_type"), + extends=graph.get("extends"), + node_count=len(nodes), + nodes_by_category=dict(sorted(categories.items())), + stage_node_count=stage_count, + input_count=len(graph.get("inputs") or []), + output_count=len(graph.get("outputs") or []), + module_property_count=len(graph.get("module_properties") or {}), + group_titles=group_titles, + ) + + +# --------------------------------------------------------------------------- +# list_nodes +# --------------------------------------------------------------------------- + + +def list_nodes( + graph: dict, + *, + registry_pattern: str | None = None, + title_pattern: str | None = None, +) -> list[NodeListEntry]: + nodes = _nodes(graph) + rp = registry_pattern.lower() if registry_pattern else None + tp = title_pattern.lower() if title_pattern else None + + out: list[NodeListEntry] = [] + for node_id, node in nodes.items(): + registry = node.get("registry") or "" + title = node.get("title") or "" + if rp and rp not in registry.lower(): + continue + if tp and tp not in title.lower(): + continue + out.append( + NodeListEntry( + short_id=_short(node_id), + id=node_id, + registry=registry or None, + title=title, + ) + ) + out.sort(key=lambda e: (e.registry or "", e.title)) + return out + + +# --------------------------------------------------------------------------- +# get_node +# --------------------------------------------------------------------------- + + +def get_node(graph: dict, node_id: str) -> NodeDetails: + full_id = resolve_node_id(graph, node_id) + node = _nodes(graph)[full_id] + + input_source_map = _input_source_map(graph) + input_names = _node_input_names(graph, full_id, node) + output_names = _node_output_names(graph, full_id, node) + + inputs: list[NodeInputInfo] = [] + for name in input_names: + src = input_source_map.get((full_id, name)) + if src is not None: + src_node, src_socket = src + inputs.append( + NodeInputInfo( + name=name, + connected=True, + source=_make_socket_ref(src_node, src_socket), + ) + ) + else: + inputs.append(NodeInputInfo(name=name, connected=False)) + + outputs: list[NodeOutputInfo] = [] + edges = _edges(graph) + for name in output_names: + full_socket = f"{full_id}.{name}" + targets = edges.get(full_socket, []) + consumer_refs: list[SocketRef] = [] + for tgt in targets: + try: + tgt_node, tgt_socket = split_socket_path(tgt) + except ValueError: + continue + consumer_refs.append(_make_socket_ref(tgt_node, tgt_socket)) + outputs.append(NodeOutputInfo(name=name, consumers=consumer_refs)) + + registered = _registered_node_class(node.get("registry")) is not None + + return NodeDetails( + id=full_id, + short_id=_short(full_id), + title=node.get("title", ""), + registry=node.get("registry"), + base_type=node.get("base_type"), + properties=dict(node.get("properties") or {}), + x=int(node.get("x", 0) or 0), + y=int(node.get("y", 0) or 0), + width=int(node.get("width", 0) or 0), + height=int(node.get("height", 0) or 0), + inputs=inputs, + outputs=outputs, + registered=registered, + ) + + +# --------------------------------------------------------------------------- +# list_edges +# --------------------------------------------------------------------------- + + +def list_edges( + graph: dict, + *, + from_node: str | None = None, + to_node: str | None = None, +) -> list[EdgeInfo]: + from_full = resolve_node_id(graph, from_node) if from_node else None + to_full = resolve_node_id(graph, to_node) if to_node else None + + out: list[EdgeInfo] = [] + for src, targets in _edges(graph).items(): + try: + src_node, src_socket = split_socket_path(src) + except ValueError: + continue + if from_full and src_node != from_full: + continue + for tgt in targets: + try: + tgt_node, tgt_socket = split_socket_path(tgt) + except ValueError: + continue + if to_full and tgt_node != to_full: + continue + out.append( + EdgeInfo( + source_node=src_node, + source_short=_short(src_node), + source_socket=src_socket, + target_node=tgt_node, + target_short=_short(tgt_node), + target_socket=tgt_socket, + ) + ) + out.sort(key=lambda e: (e.source_short, e.source_socket, e.target_short, e.target_socket)) + return out + + +# --------------------------------------------------------------------------- +# feeds / consumers +# --------------------------------------------------------------------------- + + +def feeds(graph: dict, target: str) -> FeedsResult: + """Return the source feeding ``target`` (a ``.``). + + For an input socket: returns the upstream output that connects to it. + """ + + node_part, socket = split_socket_path(target) + full_id = resolve_node_id(graph, node_part) + src = _input_source_map(graph).get((full_id, socket)) + + source_ref: SocketRef | None = None + src_title: str | None = None + src_registry: str | None = None + if src is not None: + src_node_id, src_socket = src + source_ref = _make_socket_ref(src_node_id, src_socket) + src_node = _nodes(graph).get(src_node_id, {}) + src_title = src_node.get("title") + src_registry = src_node.get("registry") + + return FeedsResult( + target_node=full_id, + target_short=_short(full_id), + target_socket=socket, + source=source_ref, + source_node_title=src_title, + source_node_registry=src_registry, + ) + + +def consumers(graph: dict, source: str) -> ConsumersResult: + """Return all targets reading ``source`` (a ``.``).""" + + node_part, socket = split_socket_path(source) + full_id = resolve_node_id(graph, node_part) + full_key = f"{full_id}.{socket}" + edges = _edges(graph).get(full_key, []) + + consumer_edges: list[EdgeInfo] = [] + for tgt in edges: + try: + tgt_node, tgt_socket = split_socket_path(tgt) + except ValueError: + continue + consumer_edges.append( + EdgeInfo( + source_node=full_id, + source_short=_short(full_id), + source_socket=socket, + target_node=tgt_node, + target_short=_short(tgt_node), + target_socket=tgt_socket, + ) + ) + return ConsumersResult( + source_node=full_id, + source_short=_short(full_id), + source_socket=socket, + consumers=consumer_edges, + ) + + +# --------------------------------------------------------------------------- +# trace_forward / trace_backward +# --------------------------------------------------------------------------- + + +def _make_trace_leaf(graph: dict, node_id: str, via: str | None, *, cycle: bool, truncated: bool) -> TraceNode: + node = _nodes(graph).get(node_id, {}) + return TraceNode( + node_id=node_id, + short_id=_short(node_id), + title=node.get("title", ""), + registry=node.get("registry"), + via_socket=via, + cycle=cycle, + truncated=truncated, + ) + + +def _trace( + graph: dict, + start_id: str, + *, + depth: int, + forward: bool, + visited: set[str] | None = None, + via: str | None = None, +) -> TraceNode: + if visited is None: + visited = set() + node = _nodes(graph).get(start_id, {}) + if start_id in visited: + return _make_trace_leaf(graph, start_id, via, cycle=True, truncated=False) + + current = TraceNode( + node_id=start_id, + short_id=_short(start_id), + title=node.get("title", ""), + registry=node.get("registry"), + via_socket=via, + ) + + if depth <= 0: + current.truncated = True + return current + + visited = visited | {start_id} + edges = _edges(graph) + + if forward: + # walk every edge whose source is this node + for src_path, targets in edges.items(): + try: + src_node, src_socket = split_socket_path(src_path) + except ValueError: + continue + if src_node != start_id: + continue + for tgt in targets: + try: + tgt_node, tgt_socket = split_socket_path(tgt) + except ValueError: + continue + child = _trace( + graph, + tgt_node, + depth=depth - 1, + forward=True, + visited=visited, + via=tgt_socket, + ) + current.children.append(child) + else: + # walk every edge whose target is this node + for src_path, targets in edges.items(): + try: + src_node, src_socket = split_socket_path(src_path) + except ValueError: + continue + for tgt in targets: + try: + tgt_node, tgt_socket = split_socket_path(tgt) + except ValueError: + continue + if tgt_node != start_id: + continue + child = _trace( + graph, + src_node, + depth=depth - 1, + forward=False, + visited=visited, + via=src_socket, + ) + current.children.append(child) + + return current + + +def trace_forward(graph: dict, start: str, *, depth: int = 3) -> TraceNode: + full_id = resolve_node_id(graph, start) + return _trace(graph, full_id, depth=depth, forward=True) + + +def trace_backward(graph: dict, end: str, *, depth: int = 3) -> TraceNode: + full_id = resolve_node_id(graph, end) + return _trace(graph, full_id, depth=depth, forward=False) + + +# --------------------------------------------------------------------------- +# stage_map +# --------------------------------------------------------------------------- + + +def _weakly_connected_components(graph: dict) -> list[set[str]]: + """Compute weakly-connected components over the node-level edge graph.""" + + nodes = _nodes(graph) + parent: dict[str, str] = {nid: nid for nid in nodes} + + def find(x: str) -> str: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a: str, b: str) -> None: + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + for src, targets in _edges(graph).items(): + try: + src_node, _ = split_socket_path(src) + except ValueError: + continue + if src_node not in parent: + continue + for tgt in targets: + try: + tgt_node, _ = split_socket_path(tgt) + except ValueError: + continue + if tgt_node not in parent: + continue + union(src_node, tgt_node) + + components: dict[str, set[str]] = defaultdict(set) + for nid in nodes: + components[find(nid)].add(nid) + return list(components.values()) + + +def stage_map(graph: dict) -> StageMapResult: + nodes = _nodes(graph) + components = _weakly_connected_components(graph) + + # index components for stable ordering: smallest stage first, then size + def comp_sort_key(comp: set[str]) -> tuple[int, int]: + stage_vals = [ + int((nodes[nid].get("properties") or {}).get("stage", 0)) + for nid in comp + if nodes[nid].get("registry") == "core/Stage" + ] + min_stage = min(stage_vals) if stage_vals else 10**9 + return (min_stage, -len(comp)) + + components.sort(key=comp_sort_key) + + stage_nodes: list[StageInfo] = [] + chains: list[StageChain] = [] + unstaged: list[str] = [] + + for idx, comp in enumerate(components): + comp_stage_ids = [ + nid for nid in comp if nodes[nid].get("registry") == "core/Stage" + ] + if not comp_stage_ids: + unstaged.extend(sorted(comp)) + continue + + comp_stages: list[int] = [] + for sid in comp_stage_ids: + stage_val = int((nodes[sid].get("properties") or {}).get("stage", 0)) + comp_stages.append(stage_val) + stage_nodes.append( + StageInfo( + node_id=sid, + short_id=_short(sid), + title=nodes[sid].get("title", ""), + stage=stage_val, + chain_index=idx, + ) + ) + chains.append( + StageChain( + index=idx, + node_ids=sorted(comp), + stage_node_ids=sorted(comp_stage_ids), + stages=sorted(comp_stages), + ) + ) + + stage_nodes.sort(key=lambda s: (s.chain_index, s.stage, s.short_id)) + return StageMapResult( + stage_nodes=stage_nodes, + chains=chains, + unstaged_node_ids=sorted(unstaged), + ) + + +# --------------------------------------------------------------------------- +# check_registries +# --------------------------------------------------------------------------- + + +def check_registries(graph: dict) -> RegistryCheckResult: + NODES = _registry_dict() + BASE_TYPES = _base_types_dict() + + unknown_registries: list[RegistryProblem] = [] + unknown_base_types: list[RegistryProblem] = [] + checked_registries = 0 + checked_base_types = 0 + + def _check(node_id: str, node: dict) -> None: + nonlocal checked_registries, checked_base_types + registry = node.get("registry") + base_type = node.get("base_type") + title = node.get("title") or "" + if registry: + checked_registries += 1 + if registry not in NODES: + unknown_registries.append( + RegistryProblem( + node_id=node_id, + short_id=_short(node_id), + title=title, + registry=registry, + base_type=base_type, + kind="registry", + ) + ) + if base_type: + checked_base_types += 1 + if base_type not in BASE_TYPES: + unknown_base_types.append( + RegistryProblem( + node_id=node_id, + short_id=_short(node_id), + title=title, + registry=registry, + base_type=base_type, + kind="base_type", + ) + ) + + # check the top-level graph itself + top_id = graph.get("id") or "" + _check(top_id, graph) + + # check every interior node + for nid, node in _nodes(graph).items(): + _check(nid, node) + + return RegistryCheckResult( + unknown_registries=unknown_registries, + unknown_base_types=unknown_base_types, + checked_registry_count=checked_registries, + checked_base_type_count=checked_base_types, + ) diff --git a/src/talemate/game/engine/nodes/tools/cli.py b/src/talemate/game/engine/nodes/tools/cli.py new file mode 100644 index 00000000..5970d529 --- /dev/null +++ b/src/talemate/game/engine/nodes/tools/cli.py @@ -0,0 +1,398 @@ +""" +Argparse-based CLI wrapping ``analysis.py``. + +Subcommands map 1:1 to analysis functions. Every subcommand supports +``--json`` for machine-readable output. Default output is dense, grep- +friendly text using 8-char short UUID prefixes. +""" + +from __future__ import annotations + +import argparse +import io +import json +import sys +from typing import Any, Callable, TextIO + +import pydantic + +from . import analysis +from .loader import GraphLoadError, load_graph + +__all__ = ["main", "build_parser"] + + +# --------------------------------------------------------------------------- +# Exit codes +# --------------------------------------------------------------------------- + +EXIT_OK = 0 +EXIT_STRUCT_ERROR = 1 +EXIT_REGISTRY_PROBLEMS = 2 + + +# --------------------------------------------------------------------------- +# Formatters +# --------------------------------------------------------------------------- + + +def _fmt_summary(s: analysis.GraphSummary) -> str: + out = io.StringIO() + out.write(f"title: {s.title}\n") + out.write(f"id: {s.id}\n") + out.write(f"registry: {s.registry}\n") + out.write(f"base_type: {s.base_type}\n") + out.write(f"extends: {s.extends}\n") + out.write(f"nodes: {s.node_count}\n") + out.write(f" by category:\n") + for cat, count in s.nodes_by_category.items(): + out.write(f" {cat:<20} {count}\n") + out.write(f"stage nodes: {s.stage_node_count}\n") + out.write( + f"interface: inputs={s.input_count} outputs={s.output_count} " + f"module_properties={s.module_property_count}\n" + ) + if s.group_titles: + out.write("groups:\n") + for t in s.group_titles: + out.write(f" - {t}\n") + return out.getvalue().rstrip() + + +def _fmt_list_nodes(entries: list[analysis.NodeListEntry]) -> str: + if not entries: + return "(no nodes match)" + lines = [] + reg_w = max((len(e.registry or "") for e in entries), default=0) + reg_w = min(reg_w, 50) + for e in entries: + lines.append(f"{e.short_id} {(e.registry or ''):<{reg_w}} {e.title}") + return "\n".join(lines) + + +def _fmt_node_details(d: analysis.NodeDetails) -> str: + out = io.StringIO() + out.write(f"{d.short_id} {d.title}\n") + out.write(f" id: {d.id}\n") + out.write(f" registry: {d.registry} (registered={d.registered})\n") + out.write(f" base_type: {d.base_type}\n") + out.write(f" pos: x={d.x} y={d.y} w={d.width} h={d.height}\n") + if d.properties: + out.write(" properties:\n") + for k, v in d.properties.items(): + v_str = repr(v) + if len(v_str) > 100: + v_str = v_str[:97] + "..." + out.write(f" {k} = {v_str}\n") + out.write(" inputs:\n") + if not d.inputs: + out.write(" (none)\n") + for inp in d.inputs: + if inp.connected and inp.source is not None: + out.write( + f" [x] {inp.name:<20} <- {inp.source.short_id}.{inp.source.socket}\n" + ) + else: + out.write(f" [ ] {inp.name:<20} (unconnected)\n") + out.write(" outputs:\n") + if not d.outputs: + out.write(" (none)\n") + for op in d.outputs: + if not op.consumers: + out.write(f" [ ] {op.name:<20} (no consumers)\n") + else: + out.write(f" [x] {op.name:<20} ->\n") + for c in op.consumers: + out.write(f" {c.short_id}.{c.socket}\n") + return out.getvalue().rstrip() + + +def _fmt_edges(edges: list[analysis.EdgeInfo]) -> str: + if not edges: + return "(no edges)" + lines = [] + for e in edges: + lines.append( + f"{e.source_short}.{e.source_socket} -> " + f"{e.target_short}.{e.target_socket}" + ) + return "\n".join(lines) + + +def _fmt_feeds(r: analysis.FeedsResult) -> str: + out = io.StringIO() + out.write(f"target: {r.target_short}.{r.target_socket}\n") + if r.source is None: + out.write("source: (unconnected)\n") + else: + out.write(f"source: {r.source.short_id}.{r.source.socket}\n") + out.write(f" title: {r.source_node_title}\n") + out.write(f" registry: {r.source_node_registry}\n") + return out.getvalue().rstrip() + + +def _fmt_consumers(r: analysis.ConsumersResult) -> str: + out = io.StringIO() + out.write(f"source: {r.source_short}.{r.source_socket}\n") + if not r.consumers: + out.write("consumers: (none)\n") + else: + out.write(f"consumers ({len(r.consumers)}):\n") + for c in r.consumers: + out.write(f" -> {c.target_short}.{c.target_socket}\n") + return out.getvalue().rstrip() + + +def _fmt_trace(t: analysis.TraceNode, *, indent: int = 0) -> str: + out = io.StringIO() + prefix = " " * indent + via = f" ({t.via_socket})" if t.via_socket else "" + markers = [] + if t.cycle: + markers.append("CYCLE") + if t.truncated: + markers.append("DEPTH") + marker_str = f" [{', '.join(markers)}]" if markers else "" + out.write(f"{prefix}{t.short_id} {t.title} <{t.registry}>{via}{marker_str}\n") + for child in t.children: + out.write(_fmt_trace(child, indent=indent + 1)) + return out.getvalue() + + +def _fmt_stage_map(r: analysis.StageMapResult) -> str: + out = io.StringIO() + out.write(f"chains: {len(r.chains)}\n") + for chain in r.chains: + out.write( + f" chain {chain.index}: nodes={len(chain.node_ids)} " + f"stage_nodes={len(chain.stage_node_ids)} stages={chain.stages}\n" + ) + out.write(f"stage nodes: {len(r.stage_nodes)}\n") + for s in r.stage_nodes: + out.write( + f" chain {s.chain_index} stage {s.stage:<4} {s.short_id} {s.title}\n" + ) + if r.unstaged_node_ids: + out.write( + f"unstaged components: {len(r.unstaged_node_ids)} node(s) " + f"in chains without a stage marker (default priority)\n" + ) + return out.getvalue().rstrip() + + +def _fmt_check_registries(r: analysis.RegistryCheckResult) -> str: + out = io.StringIO() + out.write( + f"checked: {r.checked_registry_count} registry strings, " + f"{r.checked_base_type_count} base_type strings\n" + ) + if not r.unknown_registries and not r.unknown_base_types: + out.write("ok: no unknown registries or base types\n") + return out.getvalue().rstrip() + if r.unknown_registries: + out.write(f"unknown registries ({len(r.unknown_registries)}):\n") + for p in r.unknown_registries: + out.write(f" {p.short_id} {p.registry} ({p.title})\n") + if r.unknown_base_types: + out.write(f"unknown base types ({len(r.unknown_base_types)}):\n") + for p in r.unknown_base_types: + out.write(f" {p.short_id} {p.base_type} ({p.title})\n") + return out.getvalue().rstrip() + + +# --------------------------------------------------------------------------- +# Output dispatch +# --------------------------------------------------------------------------- + + +def _emit( + result: pydantic.BaseModel | list[pydantic.BaseModel], + *, + json_out: bool, + text_formatter: Callable[[Any], str], + out_stream: TextIO, +) -> None: + if json_out: + if isinstance(result, list): + data = [r.model_dump(mode="json") for r in result] + out_stream.write(json.dumps(data, indent=2)) + else: + out_stream.write(result.model_dump_json(indent=2)) + out_stream.write("\n") + return + out_stream.write(text_formatter(result)) + out_stream.write("\n") + + +# --------------------------------------------------------------------------- +# Argparse plumbing +# --------------------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m talemate.game.engine.nodes.tools", + description="Read-only static analysis for Talemate node graph JSON files.", + ) + sub = parser.add_subparsers(dest="cmd", required=True) + + def add_graph_arg(p: argparse.ArgumentParser) -> None: + p.add_argument("graph", help="Path to a node graph JSON file") + p.add_argument( + "--json", dest="json_out", action="store_true", help="Emit JSON output" + ) + + p_summary = sub.add_parser("summary", help="Top-level graph summary") + add_graph_arg(p_summary) + + p_list = sub.add_parser("list-nodes", help="List nodes; filter by registry/title") + add_graph_arg(p_list) + p_list.add_argument("--registry", help="Substring filter on node registry") + p_list.add_argument("--title", help="Substring filter on node title") + + p_node = sub.add_parser("node", help="Show details for a single node") + add_graph_arg(p_node) + p_node.add_argument("node_id", help="Node id (full or unique prefix)") + + p_edges = sub.add_parser("edges", help="List edges, optionally filtered") + add_graph_arg(p_edges) + p_edges.add_argument("--from", dest="from_node", help="Filter by source node") + p_edges.add_argument("--to", dest="to_node", help="Filter by target node") + + p_feeds = sub.add_parser("feeds", help="What feeds a given input socket?") + add_graph_arg(p_feeds) + p_feeds.add_argument("target", help=".") + + p_cons = sub.add_parser("consumers", help="What consumes a given output socket?") + add_graph_arg(p_cons) + p_cons.add_argument("source", help=".") + + p_trace = sub.add_parser("trace", help="Forward trace from a node") + add_graph_arg(p_trace) + p_trace.add_argument("start", help="Node id (full or unique prefix)") + p_trace.add_argument("--depth", type=int, default=3) + + p_rtrace = sub.add_parser("rtrace", help="Backward trace from a node") + add_graph_arg(p_rtrace) + p_rtrace.add_argument("end", help="Node id (full or unique prefix)") + p_rtrace.add_argument("--depth", type=int, default=3) + + p_stages = sub.add_parser("stages", help="Show core/Stage chains") + add_graph_arg(p_stages) + + p_check = sub.add_parser( + "check-registries", help="Verify every registry string is known to the runtime" + ) + add_graph_arg(p_check) + + return parser + + +def _run(args: argparse.Namespace, out_stream: TextIO, err_stream: TextIO) -> int: + # Ensure every Talemate node type is registered before any command + # runs. Loading once up-front gives consistent behavior across + # commands and surfaces loader errors before we start analyzing. + analysis.ensure_registry_loaded() + + try: + graph = load_graph(args.graph) + except GraphLoadError as exc: + err_stream.write(f"error: {exc}\n") + return EXIT_STRUCT_ERROR + + cmd = args.cmd + json_out = args.json_out + + try: + if cmd == "summary": + result = analysis.summarize(graph) + _emit(result, json_out=json_out, text_formatter=_fmt_summary, out_stream=out_stream) + + elif cmd == "list-nodes": + result = analysis.list_nodes( + graph, + registry_pattern=args.registry, + title_pattern=args.title, + ) + _emit( + result, + json_out=json_out, + text_formatter=_fmt_list_nodes, + out_stream=out_stream, + ) + + elif cmd == "node": + result = analysis.get_node(graph, args.node_id) + _emit( + result, + json_out=json_out, + text_formatter=_fmt_node_details, + out_stream=out_stream, + ) + + elif cmd == "edges": + result = analysis.list_edges( + graph, from_node=args.from_node, to_node=args.to_node + ) + _emit(result, json_out=json_out, text_formatter=_fmt_edges, out_stream=out_stream) + + elif cmd == "feeds": + result = analysis.feeds(graph, args.target) + _emit(result, json_out=json_out, text_formatter=_fmt_feeds, out_stream=out_stream) + + elif cmd == "consumers": + result = analysis.consumers(graph, args.source) + _emit( + result, + json_out=json_out, + text_formatter=_fmt_consumers, + out_stream=out_stream, + ) + + elif cmd == "trace": + result = analysis.trace_forward(graph, args.start, depth=args.depth) + _emit(result, json_out=json_out, text_formatter=_fmt_trace, out_stream=out_stream) + + elif cmd == "rtrace": + result = analysis.trace_backward(graph, args.end, depth=args.depth) + _emit(result, json_out=json_out, text_formatter=_fmt_trace, out_stream=out_stream) + + elif cmd == "stages": + result = analysis.stage_map(graph) + _emit( + result, + json_out=json_out, + text_formatter=_fmt_stage_map, + out_stream=out_stream, + ) + + elif cmd == "check-registries": + result = analysis.check_registries(graph) + _emit( + result, + json_out=json_out, + text_formatter=_fmt_check_registries, + out_stream=out_stream, + ) + if result.has_problems: + return EXIT_REGISTRY_PROBLEMS + + else: # pragma: no cover - argparse enforces this + err_stream.write(f"error: unknown command {cmd}\n") + return EXIT_STRUCT_ERROR + + except ValueError as exc: + err_stream.write(f"error: {exc}\n") + return EXIT_STRUCT_ERROR + + return EXIT_OK + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + return _run(args, sys.stdout, sys.stderr) + + +if __name__ == "__main__": # pragma: no cover + sys.exit(main()) diff --git a/src/talemate/game/engine/nodes/tools/layout.py b/src/talemate/game/engine/nodes/tools/layout.py new file mode 100644 index 00000000..251d81a0 --- /dev/null +++ b/src/talemate/game/engine/nodes/tools/layout.py @@ -0,0 +1,803 @@ +""" +Deterministic auto-layout for Talemate node graph JSON dicts. + +The goal is a *readable* placement of freshly added nodes — not aesthetic +perfection. The algorithm is a stage-stratified vertical-band sweep: each +weakly-connected component of the target set is classified by the node +registries it contains, sorted vertically (Input chains on top, Stage +chains in numeric order, passthrough middle chains, Output chains at the +bottom), and then laid out horizontally using a topological depth sweep +with greedy row assignment to avoid overlaps. + +Node heights are NOT read from the writer — the writer intentionally +leaves them unset and layout estimates them from socket / property counts +via ``_estimate_height``. This keeps collision avoidance honest for nodes +that actually render tall. + +Usage:: + + from talemate.game.engine.nodes.tools import layout_graph + + layout_graph(graph, new_node_ids=[id_a, id_b], anchor="right") + +``anchor`` values: + +* ``"right"`` — place the target set to the right of the existing bounding + box, sharing the existing ``min_y``. +* ``"below"`` — place the target set below the existing bounding box, + sharing the existing ``min_x``. +* ``"full"`` — relayout every node in the graph, starting at ``(0, 0)``. + +Notes: + +* Nodes outside the target set are never moved and never resized. +* Heights of nodes inside the target set are always overwritten with an + estimate — the contract is "if you ask layout to position a node, it + also sizes it." +* Groups, comments, and any non-node canvas elements are not touched. +* ``extends`` is not resolved — only the raw JSON of this graph is + considered. +""" + +from __future__ import annotations + +import enum +from collections import defaultdict, deque +from dataclasses import dataclass, field +from itertools import groupby +from typing import Iterable + +import pydantic +import structlog + +from . import writer as writer_mod +from .writer import DEFAULT_NODE_HEIGHT, DEFAULT_NODE_WIDTH + +log = structlog.get_logger("talemate.game.engine.nodes.tools.layout") + +__all__ = [ + "layout_graph", + "LayoutError", + "LayoutOptions", + "WCCKind", +] + +# Default canvas metrics — chosen to match the rough shapes of shipped +# graphs. ``col_width = 360`` gives wires ~150 px of horizontal breathing +# room between a node's right edge (default node width 210) and the +# next column's left edge. +# +# ``DEFAULT_NODE_WIDTH`` / ``DEFAULT_NODE_HEIGHT`` are imported from +# ``writer`` so both modules agree on a single source of truth. The +# height fallback is used only when a node outside the target set is +# missing a ``height`` for bounding-box math; nodes inside the target +# set are always sized by ``_estimate_height``. +DEFAULT_COL_WIDTH = 360 +DEFAULT_RIGHT_GAP = 300 +DEFAULT_BELOW_GAP = 200 + + +class LayoutError(ValueError): + """Raised for malformed layout requests.""" + + +class LayoutOptions(pydantic.BaseModel): + """Tweakable metrics for the layout pass (defaults are fine).""" + + col_width: int = DEFAULT_COL_WIDTH + right_gap: int = DEFAULT_RIGHT_GAP + below_gap: int = DEFAULT_BELOW_GAP + # Vertical gap between stage bands after the super-band merge. + band_gap: int = 120 + # Height-aware packing: minimum vertical gap between two nodes + # stacked in the same column. + min_vertical_gap: int = 50 + # Height estimation formula inputs. + title_bar_height: int = 20 + socket_row_height: int = 22 + padding: int = 0 + min_height: int = 60 + # Extra vertical space for the +/- buttons that DynamicSocketNodeBase + # subclasses render at the bottom of the node. + dynamic_socket_bonus: int = 30 + + +# --------------------------------------------------------------------------- +# WCC classification +# --------------------------------------------------------------------------- + + +class WCCKind(enum.Enum): + """The vertical band a weakly-connected component belongs to.""" + + TOP = "top" # has core/Input, no core/Stage + STAGE = "stage" # has core/Stage + MIDDLE = "middle" # no Input, no Output, no Stage + BOTTOM = "bottom" # has core/Output, no core/Input, no core/Stage + + +@dataclass +class WCCInfo: + """One weakly-connected component and its classification.""" + + nodes: list[str] + kind: WCCKind + stage_value: int | None = None # populated when kind == STAGE + + +# --------------------------------------------------------------------------- +# Graph accessors (private) +# --------------------------------------------------------------------------- + + +def _nodes(graph: dict) -> dict[str, dict]: + return graph.get("nodes") or {} + + +def _edges(graph: dict) -> dict[str, list[str]]: + return graph.get("edges") or {} + + +def _split_edge_node(endpoint: str) -> str | None: + if "." not in endpoint: + return None + return endpoint.split(".", 1)[0] + + +def _node_rect(node: dict) -> tuple[int, int, int, int]: + x = int(node.get("x", 0) or 0) + y = int(node.get("y", 0) or 0) + w = int(node.get("width", 0) or 0) or DEFAULT_NODE_WIDTH + h = int(node.get("height", 0) or 0) or DEFAULT_NODE_HEIGHT + return x, y, w, h + + +def _existing_bounds( + graph: dict, target_set: set[str] +) -> tuple[int, int, int, int]: + """Return (min_x, min_y, max_x, max_y) for nodes NOT in the target set. + + If the "outside" set is empty (e.g. full relayout, or every node is + targeted) the bounds collapse to ``(0, 0, 0, 0)``. + """ + + xs: list[int] = [] + ys: list[int] = [] + maxes_x: list[int] = [] + maxes_y: list[int] = [] + for nid, node in _nodes(graph).items(): + if nid in target_set: + continue + x, y, w, h = _node_rect(node) + xs.append(x) + ys.append(y) + maxes_x.append(x + w) + maxes_y.append(y + h) + if not xs: + return (0, 0, 0, 0) + return (min(xs), min(ys), max(maxes_x), max(maxes_y)) + + +# --------------------------------------------------------------------------- +# Height estimation +# --------------------------------------------------------------------------- + + +def _estimate_height(graph_node: dict, options: LayoutOptions) -> int: + """Compute an estimated render height for a node from its class metadata. + + Formula:: + + title_bar_height + + (max(num_inputs, num_outputs) + num_properties) * socket_row_height + + padding + + dynamic_socket_bonus (if the class is a DynamicSocketNodeBase) + + Properties render as full widget rows **below** the socket area in + LiteGraph, stacking vertically with the socket rows rather than + being a small per-property bonus. Dynamic-socket nodes render + ``+``/``-`` buttons at the bottom for adding/removing inputs; those + take ~30 px and apply whether or not the instance currently has any + dynamic inputs (the buttons are a class capability). + + Floored at ``options.min_height``. If the node's registry is not in + the live ``NODES`` dict, falls back to ``options.min_height``. + + ``num_inputs`` includes any graph-level ``dynamic_inputs`` on the + node, because dynamic inputs also render a socket row. + """ + + registry = graph_node.get("registry") + if not registry: + return options.min_height + try: + meta = writer_mod.get_node_metadata(registry) + except writer_mod.UnknownRegistryError: + return options.min_height + + num_static_inputs = len(meta.inputs) + num_outputs = len(meta.outputs) + num_properties = len(meta.properties) + + dyn = graph_node.get("dynamic_inputs") or [] + num_dynamic_inputs = len(dyn) + num_inputs = num_static_inputs + num_dynamic_inputs + + rows = max(num_inputs, num_outputs) + num_properties + estimated = ( + options.title_bar_height + + rows * options.socket_row_height + + options.padding + ) + if meta.is_dynamic: + estimated += options.dynamic_socket_bonus + return max(estimated, options.min_height) + + +def _apply_estimated_heights( + graph: dict, target_ids: set[str], options: LayoutOptions +) -> None: + """Set ``height`` on each target node to an estimated value. + + Always overwrites the height for nodes in the target set — the + contract is "if you ask layout to position a node, it also sizes it." + Nodes outside the target set are not touched. + """ + + nodes = _nodes(graph) + for nid in target_ids: + node = nodes.get(nid) + if node is None: + continue + node["height"] = _estimate_height(node, options) + + +# --------------------------------------------------------------------------- +# Weakly-connected components over a subset +# --------------------------------------------------------------------------- + + +def _wcc_over_subset( + graph: dict, subset: set[str] +) -> list[list[str]]: + """Compute weakly-connected components restricted to ``subset``. + + Only edges with both endpoints in ``subset`` contribute to adjacency; + edges leaving the subset (bridges to the outside world) don't merge + components for our purposes. A node in ``subset`` with no internal + edges becomes a singleton component. + + Returns a list of components, each a sorted ``list[str]`` of node ids. + The outer list is sorted by the lexicographically-smallest node id of + each component for deterministic iteration order. + """ + + parent: dict[str, str] = {nid: nid for nid in subset} + + def find(x: str) -> str: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a: str, b: str) -> None: + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + for src_key, targets in _edges(graph).items(): + s_node = _split_edge_node(src_key) + if s_node is None or s_node not in subset: + continue + for tgt in targets: + t_node = _split_edge_node(tgt) + if t_node is None or t_node not in subset: + continue + union(s_node, t_node) + + buckets: dict[str, list[str]] = defaultdict(list) + for nid in subset: + buckets[find(nid)].append(nid) + + components = [sorted(comp) for comp in buckets.values()] + components.sort(key=lambda comp: comp[0] if comp else "") + return components + + +# --------------------------------------------------------------------------- +# Component classification +# --------------------------------------------------------------------------- + + +def _classify_wcc(graph: dict, node_ids: list[str]) -> WCCInfo: + """Inspect the registries of the nodes in a component and pick a kind. + + Precedence when multiple flags apply:: + + STAGE > TOP > BOTTOM > MIDDLE + + Rationale: a chain containing both an ``Input`` and a ``Stage`` is + part of that Stage's band, not the top band. A chain with both + ``Output`` and ``Input`` is a TOP (passthrough — the Input anchors + it). A chain with ``Output`` only and no ``Stage`` is BOTTOM. + """ + + nodes = _nodes(graph) + has_input = False + has_output = False + stage_values: list[int] = [] + + for nid in node_ids: + node = nodes.get(nid) or {} + registry = node.get("registry") or "" + if registry == "core/Input": + has_input = True + elif registry == "core/Output": + has_output = True + elif registry == "core/Stage": + stage_prop = (node.get("properties") or {}).get("stage", 0) + try: + stage_values.append(int(stage_prop)) + except (TypeError, ValueError): + stage_values.append(0) + + if stage_values: + return WCCInfo( + nodes=node_ids, + kind=WCCKind.STAGE, + stage_value=min(stage_values), + ) + if has_input: + return WCCInfo(nodes=node_ids, kind=WCCKind.TOP) + if has_output: + return WCCInfo(nodes=node_ids, kind=WCCKind.BOTTOM) + return WCCInfo(nodes=node_ids, kind=WCCKind.MIDDLE) + + +def _band_group_key(info: WCCInfo) -> tuple: + """Grouping key for the band-merging step in ``layout_graph``. + + Two WCCs share a key (and therefore get merged into one super-band) + when they belong to the same vertical category — e.g. all BOTTOM + components form one bottom band, all STAGE components with the same + stage number form one stage band. This eliminates the fragmentation + that arises when a graph has multiple WCCs of the same kind, which + is common when each ``GetState\u2192Output`` emission pair is its own + tiny 2-node WCC. + + The key shape mirrors the first two elements of ``_wcc_sort_key`` so + that consecutive same-keyed entries are guaranteed to be adjacent in + the sorted list (a precondition for ``itertools.groupby``). + """ + + if info.kind is WCCKind.TOP: + return (0, 0) + if info.kind is WCCKind.STAGE: + return (1, info.stage_value) # _classify_wcc guarantees non-None here + if info.kind is WCCKind.MIDDLE: + return (2, 0) + return (3, 0) # BOTTOM + + +def _wcc_sort_key(info: WCCInfo) -> tuple: + """Vertical ordering key: TOP, STAGE(0), STAGE(1), ..., MIDDLE, BOTTOM. + + Within the same band, ties break on the lexicographically-smallest + node id so the order is stable. + """ + + smallest = info.nodes[0] if info.nodes else "" + if info.kind is WCCKind.TOP: + return (0, 0, smallest) + if info.kind is WCCKind.STAGE: + return (1, info.stage_value, smallest) # STAGE kind guarantees non-None + if info.kind is WCCKind.MIDDLE: + return (2, 0, smallest) + return (3, 0, smallest) # BOTTOM + + +# --------------------------------------------------------------------------- +# Topological column sweep for a single band +# --------------------------------------------------------------------------- + + +def _topological_depths( + target_ids: set[str], internal_edges: list[tuple[str, str]] +) -> dict[str, int]: + """Kahn's algorithm; isolated nodes get depth 0. + + A cycle (if any — the writer's ``connect()`` rejects these, but raw + JSON could still have them) short-circuits: remaining nodes fall + back to depth 0 to keep the layout robust. + """ + + adj: dict[str, set[str]] = defaultdict(set) + indeg: dict[str, int] = {nid: 0 for nid in target_ids} + for s, t in internal_edges: + if t not in adj[s]: + adj[s].add(t) + indeg[t] = indeg.get(t, 0) + 1 + + depths: dict[str, int] = {nid: 0 for nid in target_ids} + queue: deque[str] = deque(nid for nid, d in indeg.items() if d == 0) + while queue: + nid = queue.popleft() + for child in adj.get(nid, ()): + candidate = depths[nid] + 1 + if candidate > depths[child]: + depths[child] = candidate + indeg[child] -= 1 + if indeg[child] == 0: + queue.append(child) + + # Any nodes left with indeg > 0 are part of a cycle; leave them at 0 + # rather than dropping them. + return depths + + +def _rects_overlap( + a: tuple[int, int, int, int], b: tuple[int, int, int, int] +) -> bool: + ax, ay, aw, ah = a + bx, by, bw, bh = b + return not (ax + aw <= bx or bx + bw <= ax or ay + ah <= by or by + bh <= ay) + + +@dataclass +class _BandPlacement: + """Internal result of placing a single band.""" + + rects: list[tuple[int, int, int, int]] = field(default_factory=list) + max_y: int = 0 # highest y + h across all placed nodes; 0 for empty + + +def _place_band( + graph: dict, + band_nodes: list[str], + *, + origin_x: int, + origin_y: int, + options: LayoutOptions, +) -> _BandPlacement: + """Place one weakly-connected component as a topological column sweep. + + All nodes in ``band_nodes`` are positioned starting at + ``(origin_x, origin_y)``. + + Packing is height-aware (per-column ``y_cursor``) and + predecessor-aware (Sugiyama-layer barycenter heuristic): for each + column past the first, unplaced nodes are sorted by the average y + of their already-placed predecessors in earlier columns, so a child + tends to land near the row of its parent(s). This eliminates the + diamond wire crossings that a naive stable-id ordering would cause. + """ + + nodes = _nodes(graph) + band_set = set(band_nodes) + + # Internal edges for topological sort (both endpoints in the band). + internal_edges: list[tuple[str, str]] = [] + for src_key, targets in _edges(graph).items(): + s_node = _split_edge_node(src_key) + if s_node is None or s_node not in band_set: + continue + for tgt in targets: + t_node = _split_edge_node(tgt) + if t_node is None or t_node not in band_set: + continue + if s_node == t_node: + continue + internal_edges.append((s_node, t_node)) + + depths = _topological_depths(band_set, internal_edges) + + # Predecessor adjacency (parent set per node, restricted to the band). + predecessors: dict[str, set[str]] = {nid: set() for nid in band_set} + for s, t in internal_edges: + predecessors[t].add(s) + + # Group by column (depth), nodes sorted by id for stable initial order. + by_depth: dict[int, list[str]] = defaultdict(list) + for nid in sorted(band_set): + by_depth[depths.get(nid, 0)].append(nid) + + placement = _BandPlacement() + + # Per-column y cursor: next free y position for that column. + column_cursors: dict[int, int] = {} + # Track placed y positions so predecessor-aware sorting can average + # over already-assigned rows. + placed_y: dict[str, int] = {} + + # Build outgoing-edge index for the entry-point sort key (column 0), + # restricted to edges whose BOTH endpoints live in this band. Any + # edge that leaves the band is irrelevant to the BFS-to-deepest walk + # and filtering it here keeps the inner loop tight. + outgoing: dict[str, list[tuple[str, str]]] = defaultdict(list) + for src_key, targets in _edges(graph).items(): + s_node = _split_edge_node(src_key) + if s_node is None or s_node not in band_set: + continue + for tgt in targets: + t_node, _, t_socket = tgt.partition(".") + if t_node and t_socket and t_node in band_set: + outgoing[s_node].append((t_node, t_socket)) + + def _input_socket_index(target_node_id: str, target_socket: str) -> int: + """Return the index of ``target_socket`` in the target node's input + list (static inputs from the class metadata followed by per-instance + ``dynamic_inputs``). Falls back to a large sentinel when the socket + can't be resolved, so it sorts to the end without crashing. + """ + node = nodes.get(target_node_id) or {} + registry = node.get("registry") or "" + if not registry: + return 999 + try: + meta = writer_mod.get_node_metadata(registry) + except writer_mod.UnknownRegistryError: + return 999 + names = list(meta.inputs) + for d in node.get("dynamic_inputs") or []: + name = d.get("name") if isinstance(d, dict) else None + if name: + names.append(name) + try: + return names.index(target_socket) + except ValueError: + return 999 + + def _entry_point_sort_key(nid: str) -> tuple: + """Sort key for column-0 (no-predecessor) nodes. + + Priority order (highest-priority case wins): + + 1. **``core/Input`` entry point** — sorted by its own ``num`` + property, so module inputs appear in interface-declaration + order (``IN x`` with ``num=1`` sits above ``IN y`` with + ``num=2``). Predecessor-aware placement then aligns each + input's ``SET local.x`` counterpart next to it, producing + horizontal IN-to-SET wires. + 2. **Deepest-descendant is ``core/Output``** — sorted by the + descendant's ``num`` property. Makes the bottom emission band's + ``GET\u2192Output`` pairs land in interface-declaration order. + 3. **General deepest-descendant case** — sorted by + ``(descendant_id, socket_index)``. When entry-point chains + converge on a sink (e.g. four ``IN``\u2192``SET`` chains feeding a + ``Stage 0`` marker), siblings share their first sort component + and differ only on socket order, so predecessor-aware placement + downstream eliminates cross-column wire crossings into the + convergence node. + 4. **No outgoing edges** — fall back to stable id order at the + end of the sort. + + Cases 1 and 2 use string-prefixed tuple components (``"__input__"`` + / ``"__output__"``) to keep them in distinct sort spaces so + cross-kind ordering stays predictable. + """ + + # Case 1: the entry point is itself a core/Input — use its num. + node = nodes.get(nid) or {} + if node.get("registry") == "core/Input": + num_val = (node.get("properties") or {}).get("num", 0) + try: + num_int = int(num_val) + except (TypeError, ValueError): + num_int = 0 + return (0, ("__input__", num_int), nid) + + # BFS forward from nid, tracking depth. ``outgoing`` is already + # pre-filtered to in-band edges, so no per-hop band check. + visited: set[str] = {nid} + bfs: deque[tuple[str, int]] = deque([(nid, 0)]) + candidates: list[tuple[int, str, int]] = [] # (depth, target_id, socket_idx) + + while bfs: + current, depth = bfs.popleft() + for t_node, t_socket in outgoing.get(current, ()): + candidates.append( + (depth + 1, t_node, _input_socket_index(t_node, t_socket)) + ) + if t_node not in visited: + visited.add(t_node) + bfs.append((t_node, depth + 1)) + + if not candidates: + return (1, ("", 0), nid) + + # Pick the deepest reachable (descendant, socket). Tie-break on + # lex-smallest (descendant_id, socket_index) so siblings that + # converge on the same descendant share their first sort component + # and differ only on socket order. + candidates.sort(key=lambda c: (-c[0], c[1], c[2])) + _, desc_id, socket_idx = candidates[0] + + # Case 2: deepest descendant is a core/Output — use its num + # property (interface-declaration order). + desc_node = nodes.get(desc_id) or {} + if desc_node.get("registry") == "core/Output": + num_val = (desc_node.get("properties") or {}).get("num", 0) + try: + num_int = int(num_val) + except (TypeError, ValueError): + num_int = 0 + return (0, ("__output__", num_int), nid) + + # Case 3: general deepest-descendant case. + return (0, (desc_id, socket_idx), nid) + + for depth in sorted(by_depth): + column_x = origin_x + depth * options.col_width + candidates = by_depth[depth] + + if depth == 0: + # First column: no predecessors, so we can't barycenter on + # parents. Sort entry points by where their wires attach on + # their successors instead — see ``_entry_point_sort_key``. + ordered = sorted(candidates, key=_entry_point_sort_key) + else: + # Sort by preferred y = average y of placed predecessors in + # earlier columns. Nodes with no placed predecessor go to + # the end. Secondary key is the stable id so ties are + # deterministic. + def _preferred(nid: str) -> tuple[int, int, str]: + placed_parents = [ + placed_y[p] + for p in predecessors.get(nid, ()) + if p in placed_y + ] + if placed_parents: + avg = sum(placed_parents) // len(placed_parents) + return (0, avg, nid) + return (1, 0, nid) + + ordered = sorted(candidates, key=_preferred) + + for nid in ordered: + node = nodes[nid] + _, _, w, h = _node_rect(node) + + # Preferred y from predecessors (ignored in column 0). + preferred_y: int | None = None + if depth > 0: + placed_parents = [ + placed_y[p] + for p in predecessors.get(nid, ()) + if p in placed_y + ] + if placed_parents: + preferred_y = sum(placed_parents) // len(placed_parents) + + cursor = column_cursors.get(depth, origin_y) + if preferred_y is None: + target_y = cursor + else: + # Never go above the cursor (that would shuffle the + # order of already-placed nodes in this column). + target_y = max(preferred_y, cursor) + + candidate = (column_x, target_y, w, h) + # Height-aware collision check against anything already + # placed in the band (protects against cross-column overlap + # from the wider col_width or adjustment quirks). + while any(_rects_overlap(candidate, r) for r in placement.rects): + target_y += options.min_vertical_gap + candidate = (column_x, target_y, w, h) + + cx, cy, cw, ch = candidate + node["x"] = cx + node["y"] = cy + placement.rects.append(candidate) + placed_y[nid] = cy + column_cursors[depth] = cy + ch + options.min_vertical_gap + if cy + ch > placement.max_y: + placement.max_y = cy + ch + + return placement + + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- + + +def layout_graph( + graph: dict, + *, + new_node_ids: Iterable[str] | None = None, + anchor: str = "right", + options: LayoutOptions | None = None, +) -> None: + """Mutate ``graph`` so the target node set is placed without overlap. + + Parameters + ---------- + graph: + Raw graph dict (as loaded from ``loader.load_graph``). Mutated + in place — nothing is returned. + new_node_ids: + Exact set of nodes to reposition. If ``None``, every node + currently sitting at ``(0, 0)`` is treated as needing placement. + anchor: + ``"right"``, ``"below"`` or ``"full"``. See module docstring. + options: + Optional :class:`LayoutOptions` overriding canvas metrics. + """ + + opts = options or LayoutOptions() + if anchor not in ("right", "below", "full"): + raise LayoutError( + f"Unknown anchor '{anchor}'; expected 'right', 'below', or 'full'." + ) + + nodes = _nodes(graph) + if not nodes: + return + + # Resolve the target set. + if anchor == "full": + target_ids = set(nodes.keys()) + elif new_node_ids is not None: + target_ids = set(new_node_ids) + missing = target_ids - nodes.keys() + if missing: + raise LayoutError( + f"Target set references unknown node ids: {sorted(missing)}" + ) + else: + target_ids = { + nid + for nid, n in nodes.items() + if int(n.get("x", 0) or 0) == 0 and int(n.get("y", 0) or 0) == 0 + } + + if not target_ids: + return + + # Step 1: estimate heights for every target node BEFORE we read any + # rectangle for collision math. + _apply_estimated_heights(graph, target_ids, opts) + + # Step 2: pick the anchor origin based on where everything else is. + min_x, min_y, max_x, max_y = _existing_bounds(graph, target_ids) + if anchor == "right": + anchor_x = max_x + opts.right_gap if max_x else 0 + anchor_y = min_y + elif anchor == "below": + anchor_x = min_x + anchor_y = max_y + opts.below_gap if max_y else 0 + else: # full + anchor_x = 0 + anchor_y = 0 + + # Step 3: compute weakly-connected components of the target set and + # classify each one so we can stack them vertically. + components = _wcc_over_subset(graph, target_ids) + wcc_infos = [_classify_wcc(graph, comp) for comp in components] + wcc_infos.sort(key=_wcc_sort_key) + + # Step 4: place each band, stacking downward. WCCs of the same kind + # (e.g. multiple BOTTOM components, or both halves of a single STAGE) + # are merged into one super-band so they read as one cohesive group + # rather than as multiple fragmented bands separated by ``band_gap``. + # The y-cursor starts at anchor_y and advances by (band_height + + # band_gap) after each merged band. + y_cursor = anchor_y + for _key, group_iter in groupby(wcc_infos, key=_band_group_key): + group = list(group_iter) + # Flatten all nodes from all WCCs in this group into a single + # band. ``_place_band`` builds its own internal topology, so + # passing nodes from multiple disconnected sub-WCCs just means + # each sub-WCC contributes its own column structure to the band. + band_nodes: list[str] = [nid for info in group for nid in info.nodes] + placement = _place_band( + graph, + band_nodes, + origin_x=anchor_x, + origin_y=y_cursor, + options=opts, + ) + if not placement.rects: + continue + band_height = placement.max_y - y_cursor + y_cursor = y_cursor + band_height + opts.band_gap diff --git a/src/talemate/game/engine/nodes/tools/loader.py b/src/talemate/game/engine/nodes/tools/loader.py new file mode 100644 index 00000000..318ca807 --- /dev/null +++ b/src/talemate/game/engine/nodes/tools/loader.py @@ -0,0 +1,227 @@ +""" +File-I/O choke point for the node-graph analysis tools. + +This module also handles **scene-module auto-registration**: when you +load a graph file that lives under ``scenes//nodes/``, every +sibling ``*.json`` in the same directory is registered into the live +``NODES`` dict so that any ``add_node("...")`` or static-analysis call +that references one of those scene-level modules can resolve it. This +removes the footgun where ``ensure_registry_loaded()`` only walks the +shipped ``SEARCH_PATHS`` and silently doesn't see scene modules. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Iterable + +import structlog + +from talemate.game.engine.nodes.registry import NODES, import_node_definition + +from .analysis import ensure_registry_loaded + +__all__ = ["load_graph", "GraphLoadError", "register_scene_modules"] + +log = structlog.get_logger("talemate.game.engine.nodes.tools.loader") + + +class GraphLoadError(ValueError): + """Raised when a graph file is missing or unparseable.""" + + +def load_graph( + path: str | Path, + *, + extra_module_paths: Iterable[str | Path] | None = None, +) -> dict: + """Load and parse a node graph JSON file into a plain dict. + + Raises ``GraphLoadError`` for missing files or malformed JSON. The + returned dict is the raw JSON contents - we deliberately do NOT run + it through any pydantic model so analysis remains side-effect-free. + + **Scene-module auto-registration** runs as a side effect: + + * If ``path`` matches the pattern ``**/scenes//nodes/.json``, + every sibling ``*.json`` in that directory is registered into the + live ``NODES`` dict via ``register_scene_modules``. This means the + writer and analysis layer can introspect scene-level modules that + the standard ``ensure_registry_loaded()`` walk doesn't reach. + * If ``extra_module_paths`` is given, every ``*.json`` under each of + those paths is also registered. Use this for non-standard layouts. + + Auto-registration failures are logged at warn-level and never raise — + the loader's primary job is still to load the requested file. + """ + + p = Path(path) + if not p.exists(): + raise GraphLoadError(f"Graph file not found: {p}") + if not p.is_file(): + raise GraphLoadError(f"Not a regular file: {p}") + try: + with p.open("r", encoding="utf-8") as fh: + data = json.load(fh) + except json.JSONDecodeError as exc: + raise GraphLoadError(f"Failed to parse graph JSON {p}: {exc}") from exc + + if not isinstance(data, dict): + raise GraphLoadError( + f"Graph JSON {p} did not parse to an object (got {type(data).__name__})" + ) + + # Auto-register scene modules when loading from a scene path. This + # fires AFTER parsing the requested file (so a malformed file still + # raises predictably) but BEFORE returning, so downstream callers + # see the populated registry. + _maybe_autoregister_scene_modules(p) + if extra_module_paths: + for extra in extra_module_paths: + register_scene_modules(extra) + + return data + + +def register_scene_modules(directory: str | Path) -> int: + """Walk ``directory`` for ``*.json`` files and register each as a node + type in the live ``NODES`` dict. Returns the number of modules newly + registered. Already-registered registries are skipped. + + Files without a top-level ``registry`` key are silently ignored + (they aren't node modules — they may be scene shells, asset metadata, + or other JSON in the same dir). + + The standard node registry (``ensure_registry_loaded()``) is always + loaded first — scene modules typically reference standard node types + like ``validation/ValidateValueIsSet`` and ``core/MakeBool`` in their + interior, and ``import_node_definition`` validates the full graph at + registration time. Without the standard registry loaded first, every + scene module that uses any standard node would fail to validate. + + A retry loop handles inter-module dependencies inside the directory: + if module A references module B and the alphabetic walk hits A first, + A fails its first attempt because B isn't registered yet, then B + succeeds, then A retries and succeeds. + + Failures after retries are logged but never raised. Useful when you + have node modules in a non-standard location and want them visible + to the writer/analysis layer. + + Why not just call ``registry.import_scene_node_definitions``? + The native function requires a live ``Scene`` object and stores + its classes in a per-scene container (``scene._NODE_DEFINITIONS``) + rather than the global ``NODES`` dict. This helper exists because + the tools package runs in contexts (CLI subprocesses, tests, one- + shot scripts) where there is no active scene, and it needs the + registered classes in the **global** ``NODES`` dict so the writer + and analysis layer can introspect them. Additionally, this + version rolls back half-registered classes between retry attempts + (``@register`` fires synchronously at class definition, and a + failed ``node.model_validate(node_data)`` leaves a broken class + behind) — the native scene loader has that same latent issue but + masks it by wiping its per-scene container on the next load. + Consolidating these via a shared lower-level helper in + ``registry.py`` is a worthwhile refactor but out of scope here. + """ + + d = Path(directory) + if not d.is_dir(): + return 0 + + # Make sure the standard registry is loaded before we try to validate + # scene modules — otherwise references to standard types will fail. + ensure_registry_loaded() + + # Collect all candidate module entries first so we can iterate with + # a retry loop for inter-module dependencies. + candidates: list[tuple[Path, dict, str]] = [] + for json_path in sorted(d.glob("*.json")): + try: + with json_path.open("r", encoding="utf-8") as fh: + module_data = json.load(fh) + except (json.JSONDecodeError, OSError) as exc: + log.debug("scene_module_skipped", path=str(json_path), reason=str(exc)) + continue + if not isinstance(module_data, dict): + continue + registry = module_data.get("registry") + if not registry: + continue + if registry in NODES: + continue + candidates.append((json_path, module_data, registry)) + + registered = 0 + last_errors: dict[str, str] = {} + + # Retry loop: keep trying as long as at least one module made progress + # in the previous pass. ``import_node_definition`` registers the class + # via ``@register`` synchronously and THEN validates by instantiating, + # so a failed validation leaves a half-broken class in NODES. We + # explicitly remove failed registries between attempts so retries get + # a clean slate. + pending = list(candidates) + progress = True + while pending and progress: + progress = False + next_pending: list[tuple[Path, dict, str]] = [] + for entry in pending: + json_path, module_data, registry = entry + try: + import_node_definition(module_data, NODES) + registered += 1 + progress = True + last_errors.pop(registry, None) + except Exception as exc: # noqa: BLE001 — see docstring + # Roll back the half-registered class so the next retry + # attempt re-creates it cleanly. + NODES.pop(registry, None) + last_errors[registry] = str(exc) + next_pending.append(entry) + pending = next_pending + + # Anything left in `pending` failed every attempt — log a single + # summary warning per registry rather than spamming once per attempt. + for json_path, _module_data, registry in pending: + log.warning( + "scene_module_import_failed", + path=str(json_path), + registry=registry, + err=last_errors.get(registry, "unknown"), + ) + + if registered: + # Clear the writer's metadata cache so any newly-registered + # registry gets re-probed on its next ``get_node_metadata`` call. + # Lazy to break the writer -> loader import cycle. + from .writer import clear_metadata_cache + + clear_metadata_cache() + + return registered + + +def _maybe_autoregister_scene_modules(graph_path: Path) -> None: + """If ``graph_path`` looks like ``**/scenes//nodes/*.json``, + register every JSON sibling in the same directory. + + The pattern match is intentionally strict: it requires the immediate + parent directory to be named ``nodes`` AND its grandparent to be a + direct child of a directory named ``scenes``. This avoids accidentally + walking unrelated ``scenes/`` directories elsewhere in the filesystem. + """ + + try: + resolved = graph_path.resolve() + except OSError: + return + parts = resolved.parts + if len(parts) < 4: + return + if parts[-2] != "nodes": + return + if parts[-4] != "scenes": + return + register_scene_modules(resolved.parent) diff --git a/src/talemate/game/engine/nodes/tools/writer.py b/src/talemate/game/engine/nodes/tools/writer.py new file mode 100644 index 00000000..5e61997b --- /dev/null +++ b/src/talemate/game/engine/nodes/tools/writer.py @@ -0,0 +1,986 @@ +""" +Programmatic mutation API for Talemate node graph JSON dicts. + +The functions here take an already-loaded ``graph: dict`` and mutate it in +place — the mirror image of ``analysis.py``, which is strictly read-only. +Pure primitive functions are the core surface; ``GraphWriter`` is an +ergonomic class that wraps them and adds file I/O. + +Note on node sizes: the writer does NOT assign a default ``height`` to new +nodes. Height belongs to the layout pass (``tools.layout.layout_graph``), +which estimates it from the node's socket and property counts so collision +avoidance actually has a realistic rectangle to work with. Callers that +need a pinned height can still pass one explicitly. Width is still set to +a fixed default because every node in the repo uses the same canvas width. + +Node metadata (static inputs, outputs, properties, ``base_type``) is read +from the live ``NODES`` class registry via a lightweight in-process cache +that instantiates each class once. We deliberately avoid +``registry.export_node_definitions()`` for this cache because that helper +calls ``PropertyField.model_dump()``, which in turn calls user-supplied +``generate_choices`` callables — several of those touch runtime-only state +(the websocket handler, the active scene, etc.) and crash outside a live +Talemate session. The instance-based probe here is the same pattern +``analysis.py`` uses in ``_static_socket_names`` and matches the +frontend-facing exporter's own instantiation step. + +Known limitations: + +* The writer does NOT keep the graph's top-level ``inputs`` / ``outputs`` + / ``module_properties`` mirror arrays in sync when interior Input / + Output / ModuleProperty nodes are added or removed. Those arrays are + computed at runtime from the interior nodes, but some shipped graphs + (notably director actions) serialise them anyway. Callers that care + about those arrays should regenerate them manually. +* The writer does NOT resolve ``extends``. It operates on raw JSON. +* The writer does NOT check socket *type* compatibility; many sockets in + the codebase are typed ``any`` and strict type checks would fight real + graphs. It only checks that sockets exist and that fan-in is at most 1. +""" + +from __future__ import annotations + +import json +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable + +import structlog + +from . import analysis +from .analysis import ensure_registry_loaded, resolve_node_id, split_socket_path +from .loader import load_graph + +log = structlog.get_logger("talemate.game.engine.nodes.tools.writer") + +__all__ = [ + # class + "GraphWriter", + # pure primitives + "add_node", + "remove_node", + "connect", + "disconnect", + "add_dynamic_input", + "remove_dynamic_input", + "add_group", + # metadata helpers + "get_node_metadata", + "clear_metadata_cache", + "NodeMetadata", + "GroupColor", + # exceptions + "WriterError", + "UnknownRegistryError", + "UnknownSocketError", + "UnknownPropertyError", + "AlreadyConnectedError", + "NotConnectedError", + "CycleError", + "DynamicInputError", + "NodeNotFoundError", + "GroupError", +] + + +# --------------------------------------------------------------------------- +# Group color palette — matches the LiteGraph presets exposed in the +# frontend's "Create Group from Selection" submenu (see +# ``talemate_frontend/src/utils/recentNodes.js``). Each constant is the +# preset's ``groupcolor`` hex value, which is what gets stored on the +# group's ``color`` field in the graph JSON. Use these by name for +# readability — the agent picks which group gets which color based on +# the cluster's purpose. +# --------------------------------------------------------------------------- + + +class GroupColor: + """Named group-color constants matching the LiteGraph preset palette. + + These are the same colors the frontend's right-click "Create Group + from Selection" menu offers, so auto-generated groups visually match + the convention used by hand-built graphs. + """ + + INPUT = "#88A" # blue — input collection / passthrough + OUTPUT = "#8A8" # green — output emission + PROCESS = "#3f789e" # pale_blue — main computation stages + PREPARE = "#8AA" # cyan — setup / preconditions + VALIDATION = "#b58b2a" # yellow — validation / guards + FUNCTION = "#b06634" # brown — function definitions + SPECIAL = "#a1309b" # purple — helpers / orphans / one-offs + ERROR_HANDLING = "#A88" # red — error-handler chains + UX = "#207e7e" # teal — user-facing UX nodes + + +# Default canvas dimensions used when a node in ``add_group``'s target +# set hasn't been sized by the layout pass yet (rare — call layout first). +DEFAULT_NODE_WIDTH = 210 +DEFAULT_NODE_HEIGHT = 100 + + +# Matches the frontend's group-creation padding constants (see +# ``talemate_frontend/src/utils/groupInteractions.js``). Auto-groups +# should look identical to user-created ones, so we mirror the same +# numbers rather than picking our own. +GROUP_PADDING = 25 +GROUP_TOP_PADDING = 45 # padding + 20 +GROUP_TITLE_HEIGHT = 20 # LiteGraph.NODE_TITLE_HEIGHT default +GROUP_MIN_WIDTH = 140 +GROUP_MIN_HEIGHT = 80 +GROUP_DEFAULT_FONT_SIZE = 24 + + +# --------------------------------------------------------------------------- +# Exceptions +# --------------------------------------------------------------------------- + + +class WriterError(Exception): + """Base class for all writer-layer errors.""" + + +class UnknownRegistryError(WriterError): + """A registry path is not present in the live ``NODES`` dict.""" + + +class UnknownSocketError(WriterError): + """A named socket does not exist on the target node.""" + + +class AlreadyConnectedError(WriterError): + """A target input socket already has a source connection.""" + + +class NotConnectedError(WriterError): + """``disconnect`` was asked to remove an edge that does not exist.""" + + +class CycleError(WriterError): + """The requested ``connect`` would introduce a cycle in the graph.""" + + +class GroupError(WriterError): + """``add_group`` was given an invalid set of node ids or empty input.""" + + +class UnknownPropertyError(WriterError): + """A property key passed to ``add_node`` does not exist on the + node class. Sister to ``UnknownSocketError`` — same kind of typo + catch, but for the ``properties`` dict instead of socket names. + """ + + +class DynamicInputError(WriterError): + """A dynamic-input operation was attempted on an incompatible node.""" + + +class NodeNotFoundError(WriterError): + """A node id was not present in the graph.""" + + +# --------------------------------------------------------------------------- +# Node metadata cache +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class NodeMetadata: + """Static metadata for a registered node class, cached per process. + + ``inputs`` / ``outputs`` are socket names in declaration order (the + ones the class sets up in ``setup()``, before any graph-specific + dynamic inputs). ``properties`` is the tuple of class-declared + property names (from ``instance.properties`` after ``setup()``). + ``base_type`` is the class-level ``_base_type`` / ``base_type`` + computed field. ``is_dynamic`` is True when the class is a subclass + of ``DynamicSocketNodeBase``. + """ + + registry: str + base_type: str + inputs: tuple[str, ...] + outputs: tuple[str, ...] + properties: tuple[str, ...] + is_dynamic: bool + + +_METADATA_CACHE: dict[str, NodeMetadata] = {} + + +def _dynamic_socket_base_cls() -> type: + from talemate.game.engine.nodes.core.dynamic import DynamicSocketNodeBase + + return DynamicSocketNodeBase + + +def get_node_metadata(registry: str) -> NodeMetadata: + """Return cached ``NodeMetadata`` for a registry path. + + Raises ``UnknownRegistryError`` if the registry is not present in + the live ``NODES`` dict. + """ + + cached = _METADATA_CACHE.get(registry) + if cached is not None: + return cached + + ensure_registry_loaded() + from talemate.game.engine.nodes.registry import NODES + + node_cls = NODES.get(registry) + if node_cls is None: + raise UnknownRegistryError( + f"Registry '{registry}' is not registered. " + f"(Did you forget to call ensure_registry_loaded() or import " + f"the module that declares it?)" + ) + + try: + instance = node_cls() + except Exception as exc: # pragma: no cover - defensive + raise UnknownRegistryError( + f"Registry '{registry}' is registered but the class " + f"({node_cls.__name__}) could not be instantiated: {exc}" + ) from exc + + try: + inputs = tuple(s.name for s in instance.inputs) + except Exception: + inputs = tuple() + try: + outputs = tuple(s.name for s in instance.outputs) + except Exception: + outputs = tuple() + try: + properties = tuple((instance.properties or {}).keys()) + except Exception: + properties = tuple() + + base_type = "" + try: + base_type = instance.base_type # computed_field on NodeBase + except Exception: + base_type = getattr(node_cls, "_base_type", "") or "" + + try: + is_dynamic = isinstance(instance, _dynamic_socket_base_cls()) + except Exception: # pragma: no cover - defensive + is_dynamic = False + + meta = NodeMetadata( + registry=registry, + base_type=base_type or "", + inputs=inputs, + outputs=outputs, + properties=properties, + is_dynamic=is_dynamic, + ) + _METADATA_CACHE[registry] = meta + return meta + + +def clear_metadata_cache() -> None: + """Flush the node-metadata cache. + + Called by the loader after registering scene-level modules so that + their socket layouts are re-probed on next access. Also useful from + tests that want to observe fresh probes. + """ + + _METADATA_CACHE.clear() + + +# --------------------------------------------------------------------------- +# Graph accessors +# --------------------------------------------------------------------------- + + +def _nodes(graph: dict) -> dict[str, dict]: + nodes = graph.get("nodes") + if nodes is None: + nodes = {} + graph["nodes"] = nodes + return nodes + + +def _edges(graph: dict) -> dict[str, list[str]]: + edges = graph.get("edges") + if edges is None: + edges = {} + graph["edges"] = edges + return edges + + +def _require_node(graph: dict, node_id: str) -> dict: + nodes = _nodes(graph) + if node_id not in nodes: + raise NodeNotFoundError(f"Node id '{node_id}' is not in the graph") + return nodes[node_id] + + +def _dynamic_input_names(node: dict) -> list[str]: + return [ + d.get("name") + for d in (node.get("dynamic_inputs") or []) + if d.get("name") + ] + + +def _input_names_for_node(node: dict) -> list[str]: + """Inputs = class-declared inputs plus any dynamic_inputs on this node. + + Duplicates are removed while preserving declaration order (class + inputs first, then dynamic inputs). + """ + + registry = node.get("registry") + static_inputs: tuple[str, ...] = tuple() + if registry: + try: + static_inputs = get_node_metadata(registry).inputs + except UnknownRegistryError: + static_inputs = tuple() + + dyn = _dynamic_input_names(node) + seen: set[str] = set() + out: list[str] = [] + for name in list(static_inputs) + dyn: + if name and name not in seen: + seen.add(name) + out.append(name) + return out + + +def _output_names_for_node(node: dict) -> list[str]: + registry = node.get("registry") + if not registry: + return [] + try: + return list(get_node_metadata(registry).outputs) + except UnknownRegistryError: + return [] + + +# --------------------------------------------------------------------------- +# Pure primitives +# --------------------------------------------------------------------------- + + +def add_node( + graph: dict, + registry: str, + *, + title: str | None = None, + properties: dict[str, Any] | None = None, + x: int = 0, + y: int = 0, + width: int = DEFAULT_NODE_WIDTH, + height: int | None = None, +) -> str: + """Add a new node to ``graph`` and return its new UUID. + + The new id is always a freshly generated ``uuid4``. The registry + must be present in the live ``NODES`` dict, otherwise + ``UnknownRegistryError`` is raised. + + Property keys in ``properties`` are validated against the node + class's declared ``Fields`` — an unknown / typo'd key raises + ``UnknownPropertyError`` with a helpful message listing the valid + property names. This is the property-side equivalent of how + ``connect()`` validates socket names. + + ``height`` defaults to ``None`` — the writer intentionally does not + pin a height on new nodes because height should be estimated by + the layout pass from the node's socket / property counts. Pass an + explicit integer to override and pin the height. + """ + + meta = get_node_metadata(registry) + + # Validate property keys against the class metadata. We deliberately + # validate even an empty/None properties arg early in the function + # so the error fires before any partial mutation happens. + if properties: + valid = set(meta.properties) + unknown = [k for k in properties if k not in valid] + if unknown: + unknown_repr = ", ".join(repr(k) for k in unknown) + valid_repr = ", ".join(repr(k) for k in sorted(valid)) or "(none)" + raise UnknownPropertyError( + f"Node {registry!r} has no property/properties " + f"{unknown_repr}. Valid properties: {valid_repr}." + ) + + new_id = str(uuid.uuid4()) + node: dict[str, Any] = { + "title": title if title is not None else registry.split("/")[-1], + "id": new_id, + "properties": dict(properties or {}), + "x": int(x), + "y": int(y), + "width": int(width), + "collapsed": False, + "inherited": False, + "registry": registry, + "base_type": meta.base_type, + } + if height is not None: + node["height"] = int(height) + if meta.is_dynamic: + # Mirror the shape shipped graphs use: an empty list is valid and + # makes future add_dynamic_input calls ergonomic. + node["dynamic_inputs"] = [] + + _nodes(graph)[new_id] = node + return new_id + + +def remove_node(graph: dict, node_id: str) -> None: + """Remove a node and sweep every edge that references it. + + The edge sweep removes: + + * any edge key whose source node is ``node_id`` + * any edge target entry whose target node is ``node_id``, pruning + now-empty target lists (and therefore their edge keys) + """ + + nodes = _nodes(graph) + if node_id not in nodes: + raise NodeNotFoundError(f"Node id '{node_id}' is not in the graph") + + del nodes[node_id] + + edges = _edges(graph) + # Drop entire edge keys whose source is this node + to_drop = [ + src for src in edges if src.split(".", 1)[0] == node_id + ] + for src in to_drop: + del edges[src] + + # For remaining edges, filter out targets pointing at this node. + empty_keys: list[str] = [] + for src, targets in edges.items(): + kept = [ + tgt + for tgt in targets + if tgt.split(".", 1)[0] != node_id + ] + if len(kept) != len(targets): + edges[src] = kept + if not edges[src]: + empty_keys.append(src) + for key in empty_keys: + del edges[key] + + +def _target_has_source(graph: dict, target_node: str, target_socket: str) -> bool: + full_tgt = f"{target_node}.{target_socket}" + for _src, targets in _edges(graph).items(): + if full_tgt in targets: + return True + return False + + +def connect( + graph: dict, + source_id: str, + source_socket: str, + target_id: str, + target_socket: str, +) -> None: + """Connect ``source_id.source_socket`` to ``target_id.target_socket``. + + Validates that: + + * both nodes exist in the graph + * the source output socket exists on the source node's class + * the target input socket exists (class-declared or dynamic) on the + target node + * the target input is not already connected (fan-in is at most 1) + * adding the edge would not introduce a cycle + + On cycle detection the edge is rolled back before ``CycleError`` is + raised. + """ + + source_node = _require_node(graph, source_id) + target_node = _require_node(graph, target_id) + + source_outputs = _output_names_for_node(source_node) + if source_socket not in source_outputs: + raise UnknownSocketError( + f"Source node {source_id} ({source_node.get('registry')}) has no " + f"output socket '{source_socket}'. Known outputs: " + f"{source_outputs or '(none)'}" + ) + + target_inputs = _input_names_for_node(target_node) + if target_socket not in target_inputs: + raise UnknownSocketError( + f"Target node {target_id} ({target_node.get('registry')}) has no " + f"input socket '{target_socket}'. Known inputs: " + f"{target_inputs or '(none)'}" + ) + + if _target_has_source(graph, target_id, target_socket): + raise AlreadyConnectedError( + f"Target input {target_id}.{target_socket} is already connected. " + f"Disconnect the existing source before reassigning." + ) + + edges = _edges(graph) + src_key = f"{source_id}.{source_socket}" + tgt_value = f"{target_id}.{target_socket}" + + targets_list = edges.setdefault(src_key, []) + targets_list.append(tgt_value) + + # Cycle check AFTER adding, so we include the new edge; roll back if + # it would form a cycle. + if _would_cycle_including_new(graph, source_id, target_id): + targets_list.remove(tgt_value) + if not targets_list: + del edges[src_key] + raise CycleError( + f"Connecting {source_id}.{source_socket} -> " + f"{target_id}.{target_socket} would introduce a cycle." + ) + + +def _would_cycle_including_new( + graph: dict, source_node: str, target_node: str +) -> bool: + """Cycle check based on current edge state (edge already inserted). + + With the new edge in place, a cycle exists iff there is a path from + ``target_node`` back to ``source_node`` (counting the new edge as + part of adjacency). Self-loops are cycles. + """ + + if source_node == target_node: + return True + + # Build node-level adjacency from every current edge. Edges are + # ``.`` strings; we only care about the node halves. + adj: dict[str, set[str]] = {} + for src, targets in _edges(graph).items(): + s_node = src.split(".", 1)[0] + for tgt in targets: + t_node = tgt.split(".", 1)[0] + adj.setdefault(s_node, set()).add(t_node) + + # Walk forward from target_node; if we reach source_node we have a + # cycle because source -> ... -> target -> source. + stack = list(adj.get(target_node, ())) + seen: set[str] = set() + while stack: + current = stack.pop() + if current == source_node: + return True + if current in seen: + continue + seen.add(current) + stack.extend(adj.get(current, ())) + return False + + +def disconnect( + graph: dict, + source_id: str, + source_socket: str, + target_id: str, + target_socket: str, +) -> None: + """Remove the edge ``source_id.source_socket -> target_id.target_socket``. + + Raises ``NotConnectedError`` if the edge is not present. + """ + + edges = _edges(graph) + src_key = f"{source_id}.{source_socket}" + tgt_value = f"{target_id}.{target_socket}" + + targets = edges.get(src_key) + if not targets or tgt_value not in targets: + raise NotConnectedError( + f"No edge {source_id}.{source_socket} -> " + f"{target_id}.{target_socket} to remove." + ) + targets.remove(tgt_value) + if not targets: + del edges[src_key] + + +def add_dynamic_input( + graph: dict, + node_id: str, + name: str, + socket_type: str = "*", +) -> None: + """Append a ``dynamic_inputs`` entry on a ``DynamicSocketNodeBase`` node. + + Raises ``DynamicInputError`` if the node's registered class is not a + subclass of ``DynamicSocketNodeBase``, or if a dynamic input with + ``name`` already exists. + + The default ``socket_type`` is ``"*"`` — the LiteGraph **wildcard** + socket type that accepts a connection from any other socket regardless + of type. This is the correct default for collector-style nodes like + ``data/DictCollector``, ``data/ListCollector``, ``data/string/AdvancedFormat``, + and ``data/string/Jinja2Format``, which all want to accept inputs of + any type. **Do NOT use ``"any"`` as the type** — despite the name, + ``"any"`` is a specific named type in LiteGraph, not a wildcard, and + the frontend's socket validation refuses to connect a typed output + (e.g. ``int``, ``str``) to an ``"any"`` input. Only use ``"*"``. + """ + + node = _require_node(graph, node_id) + registry = node.get("registry") + if not registry: + raise DynamicInputError( + f"Node {node_id} has no registry; cannot determine whether it " + f"supports dynamic inputs." + ) + + meta = get_node_metadata(registry) + if not meta.is_dynamic: + raise DynamicInputError( + f"Node {node_id} ({registry}) is not a DynamicSocketNodeBase " + f"subclass; it does not accept dynamic inputs." + ) + + dyn = node.setdefault("dynamic_inputs", []) + if any(d.get("name") == name for d in dyn): + raise DynamicInputError( + f"Dynamic input '{name}' already exists on node {node_id}." + ) + dyn.append({"name": name, "type": socket_type}) + + +def remove_dynamic_input(graph: dict, node_id: str, name: str) -> None: + """Remove a dynamic input entry and sweep any edges using it.""" + + node = _require_node(graph, node_id) + dyn = node.get("dynamic_inputs") or [] + new_dyn = [d for d in dyn if d.get("name") != name] + if len(new_dyn) == len(dyn): + raise DynamicInputError( + f"No dynamic input named '{name}' on node {node_id}." + ) + node["dynamic_inputs"] = new_dyn + + # Sweep any edges targeting this now-removed socket. + edges = _edges(graph) + tgt_value = f"{node_id}.{name}" + empty_keys: list[str] = [] + for src, targets in edges.items(): + if tgt_value in targets: + targets.remove(tgt_value) + if not targets: + empty_keys.append(src) + for key in empty_keys: + del edges[key] + + +# --------------------------------------------------------------------------- +# Group creation +# --------------------------------------------------------------------------- + + +def add_group( + graph: dict, + title: str, + color: str, + node_ids: Iterable[str], + *, + font_size: int = GROUP_DEFAULT_FONT_SIZE, +) -> dict: + """Append a colored, titled group around the given nodes' bounding box. + + Computes the union bounding box of every node in ``node_ids`` (using + each node's ``x``/``y``/``width``/``height``) and then expands it by + the same padding constants the frontend's "Create Group from + Selection" menu uses, so the resulting group is visually identical + to a user-created one. + + Parameters + ---------- + graph: + The graph dict being mutated. + title: + Display title for the group. Use a short human-readable label + (e.g. ``"Input"``, ``"Stage 0"``, ``"Output"``). + color: + The group's border / title color, as a hex string. Use one of + the constants on :class:`GroupColor` to match the frontend's + preset palette. + node_ids: + The nodes the group should encompass. Must contain at least one + valid node id. Short-prefix ids are NOT resolved here — pass + full ids (use ``GraphWriter.add_group`` for prefix support). + font_size: + Title font size. Defaults to 24, which matches every existing + group in shipped graphs. + + Returns + ------- + dict + The newly-created group dict (which has also been appended to + ``graph["groups"]`` in place). Caller can mutate further if + needed. + + Raises + ------ + GroupError + If ``node_ids`` is empty or any id does not exist in the graph. + + Notes + ----- + Call this **after** running ``layout_graph``. Group bounding boxes + are computed from the nodes' actual positions, so positioning the + nodes first is required. + """ + + nodes = _nodes(graph) + ids = list(node_ids) + if not ids: + raise GroupError("add_group requires at least one node id.") + + rects: list[tuple[int, int, int, int]] = [] + for nid in ids: + node = nodes.get(nid) + if node is None: + raise GroupError(f"add_group: unknown node id {nid!r}.") + x = int(node.get("x", 0) or 0) + y = int(node.get("y", 0) or 0) + w = int(node.get("width", DEFAULT_NODE_WIDTH) or DEFAULT_NODE_WIDTH) + h = int(node.get("height", DEFAULT_NODE_HEIGHT) or DEFAULT_NODE_HEIGHT) + rects.append((x, y, w, h)) + + min_x = min(r[0] for r in rects) + min_y = min(r[1] for r in rects) + max_x = max(r[0] + r[2] for r in rects) + max_y = max(r[1] + r[3] for r in rects) + + # Mirror the frontend's group-creation padding (groupInteractions.js). + pos_x = min_x - GROUP_PADDING + pos_y = min_y - GROUP_TOP_PADDING - GROUP_TITLE_HEIGHT + raw_width = (max_x - min_x) + GROUP_PADDING * 2 + raw_height = (max_y - min_y) + GROUP_PADDING + GROUP_TOP_PADDING + GROUP_TITLE_HEIGHT + width = max(GROUP_MIN_WIDTH, raw_width) + height = max(GROUP_MIN_HEIGHT, raw_height) + + group_dict = { + "title": title, + "x": pos_x, + "y": pos_y, + "width": width, + "height": height, + "color": color, + "font_size": font_size, + "inherited": False, + } + + groups = graph.setdefault("groups", []) + groups.append(group_dict) + return group_dict + + +# --------------------------------------------------------------------------- +# GraphWriter ergonomics +# --------------------------------------------------------------------------- + + +class GraphWriter: + """Ergonomic wrapper around the pure mutation primitives. + + Holds a graph dict plus an optional source path for ``save()``. All + mutation methods delegate to the module-level functions above — the + class is sugar, not logic. + """ + + def __init__(self, graph: dict, *, path: str | Path | None = None) -> None: + self.graph = graph + self.path: Path | None = Path(path) if path is not None else None + + # -- construction ------------------------------------------------------ + + @classmethod + def load( + cls, + path: str | Path, + *, + extra_module_paths: Iterable[str | Path] | None = None, + ) -> "GraphWriter": + """Load a graph JSON file from disk and wrap it. + + If ``path`` lives under ``scenes//nodes/``, sibling JSON + modules in that directory are auto-registered into the live + ``NODES`` dict so the writer can introspect them when you + reference one via ``add_node("...")``. Pass ``extra_module_paths`` + to also walk additional directories — useful for non-standard + layouts where related modules live elsewhere. + """ + + graph = load_graph(path, extra_module_paths=extra_module_paths) + return cls(graph, path=path) + + def save(self, path: str | Path | None = None, *, indent: int = 2) -> Path: + """Persist the graph to disk, returning the path written to. + + If ``path`` is omitted the writer uses the path it was loaded + from; if the writer was constructed without a path and no + argument is passed this raises ``WriterError``. + """ + + target = Path(path) if path is not None else self.path + if target is None: + raise WriterError( + "GraphWriter.save() needs an explicit path because this " + "writer was not constructed via GraphWriter.load()." + ) + target.parent.mkdir(parents=True, exist_ok=True) + with target.open("w", encoding="utf-8") as fh: + json.dump(self.graph, fh, indent=indent) + fh.write("\n") + self.path = target + return target + + # -- mutation ---------------------------------------------------------- + + def add_node( + self, + registry: str, + *, + title: str | None = None, + properties: dict[str, Any] | None = None, + x: int = 0, + y: int = 0, + width: int = DEFAULT_NODE_WIDTH, + height: int | None = None, + ) -> str: + return add_node( + self.graph, + registry, + title=title, + properties=properties, + x=x, + y=y, + width=width, + height=height, + ) + + def remove_node(self, node_id: str) -> None: + remove_node(self.graph, self._resolve(node_id)) + + def connect(self, *args: str) -> None: + """Connect two sockets. + + Two call shapes are supported: + + * ``writer.connect(src_id, src_socket, tgt_id, tgt_socket)`` + * ``writer.connect(".", ".")`` + """ + + src_id, src_socket, tgt_id, tgt_socket = self._parse_edge_args(args) + connect( + self.graph, + self._resolve(src_id), + src_socket, + self._resolve(tgt_id), + tgt_socket, + ) + + def disconnect(self, *args: str) -> None: + """Disconnect two sockets. Same call shapes as :meth:`connect`.""" + + src_id, src_socket, tgt_id, tgt_socket = self._parse_edge_args(args) + disconnect( + self.graph, + self._resolve(src_id), + src_socket, + self._resolve(tgt_id), + tgt_socket, + ) + + def add_dynamic_input( + self, node_id: str, name: str, socket_type: str = "*" + ) -> None: + add_dynamic_input( + self.graph, self._resolve(node_id), name, socket_type=socket_type + ) + + def remove_dynamic_input(self, node_id: str, name: str) -> None: + remove_dynamic_input(self.graph, self._resolve(node_id), name) + + def add_group( + self, + title: str, + color: str, + node_ids: Iterable[str], + *, + font_size: int = GROUP_DEFAULT_FONT_SIZE, + ) -> dict: + """Append a colored, titled group around the given nodes. + + Same as the pure :func:`add_group` but resolves short-prefix + node ids transparently. **Call this after layout** — group + bounding boxes are computed from the nodes' actual positions. + + Pick a ``color`` from :class:`GroupColor` (or pass a literal hex + string). The agent decides which color matches the cluster's + purpose; common conventions: + + * ``GroupColor.INPUT`` for input-collection / passthrough bands + * ``GroupColor.OUTPUT`` for output-emission bands + * ``GroupColor.PROCESS`` for main computation stages + * ``GroupColor.VALIDATION`` for guard / validation chains + * ``GroupColor.SPECIAL`` for helpers and one-offs + """ + + resolved = [self._resolve(nid) for nid in node_ids] + return add_group( + self.graph, + title=title, + color=color, + node_ids=resolved, + font_size=font_size, + ) + + # -- lookup helpers ---------------------------------------------------- + + def resolve(self, node_id_or_prefix: str) -> str: + """Resolve a short-prefix node id to a full id.""" + + return self._resolve(node_id_or_prefix) + + def _resolve(self, node_id_or_prefix: str) -> str: + try: + return resolve_node_id(self.graph, node_id_or_prefix) + except ValueError as exc: + raise NodeNotFoundError(str(exc)) from exc + + @staticmethod + def _parse_edge_args( + args: tuple[str, ...], + ) -> tuple[str, str, str, str]: + if len(args) == 4: + return args # type: ignore[return-value] + if len(args) == 2: + src_node, src_socket = split_socket_path(args[0]) + tgt_node, tgt_socket = split_socket_path(args[1]) + return src_node, src_socket, tgt_node, tgt_socket + raise TypeError( + "connect/disconnect expects either 4 string args " + "(src_id, src_socket, tgt_id, tgt_socket) or 2 dotted strings " + "('src_id.socket', 'tgt_id.socket')." + ) + + # -- introspection ----------------------------------------------------- + + def summarize(self) -> analysis.GraphSummary: + return analysis.summarize(self.graph) diff --git a/tests/test_nodegraph_tools.py b/tests/test_nodegraph_tools.py new file mode 100644 index 00000000..7034859a --- /dev/null +++ b/tests/test_nodegraph_tools.py @@ -0,0 +1,382 @@ +""" +Tests for ``talemate.game.engine.nodes.tools`` static analysis CLI/library. +""" + +from __future__ import annotations + +import io +import json +from pathlib import Path + +import pytest + +from talemate.game.engine.nodes.tools import analysis, cli +from talemate.game.engine.nodes.tools.loader import GraphLoadError, load_graph + +REPO_ROOT = Path(__file__).resolve().parent.parent +DIRECTOR_FIXTURE = ( + REPO_ROOT + / "src/talemate/agents/director/modules/director-action-direct-story-arc.json" +) + + +# --------------------------------------------------------------------------- +# Inline fixture graph +# --------------------------------------------------------------------------- + + +def _make_inline_graph() -> dict: + """Hand-built tiny graph used for deterministic unit tests. + + Layout: + + a (core/MakeBool) ---value---> b (core/Stage stage=2) + | + state---> c (core/Watch) + d (mystery/Unknown) (no edges) + """ + + return { + "title": "Inline Test Graph", + "id": "00000000-0000-0000-0000-000000000000", + "registry": "test/InlineGraph", + "base_type": "core/Graph", + "extends": None, + "properties": {}, + "nodes": { + "aaaa1111-aaaa-aaaa-aaaa-aaaaaaaaaaaa": { + "title": "true", + "id": "aaaa1111-aaaa-aaaa-aaaa-aaaaaaaaaaaa", + "properties": {"value": True}, + "registry": "core/MakeBool", + "base_type": "core/Node", + }, + "bbbb2222-bbbb-bbbb-bbbb-bbbbbbbbbbbb": { + "title": "Stage 2", + "id": "bbbb2222-bbbb-bbbb-bbbb-bbbbbbbbbbbb", + "properties": {"stage": 2}, + "registry": "core/Stage", + "base_type": "core/Node", + }, + "cccc3333-cccc-cccc-cccc-cccccccccccc": { + "title": "Watcher", + "id": "cccc3333-cccc-cccc-cccc-cccccccccccc", + "properties": {}, + "registry": "core/Watch", + "base_type": "core/Node", + }, + "dddd4444-dddd-dddd-dddd-dddddddddddd": { + "title": "Mystery Node", + "id": "dddd4444-dddd-dddd-dddd-dddddddddddd", + "properties": {}, + "registry": "mystery/Unknown", + "base_type": "core/Node", + }, + }, + "edges": { + "aaaa1111-aaaa-aaaa-aaaa-aaaaaaaaaaaa.value": [ + "bbbb2222-bbbb-bbbb-bbbb-bbbbbbbbbbbb.state" + ], + "bbbb2222-bbbb-bbbb-bbbb-bbbbbbbbbbbb.state": [ + "cccc3333-cccc-cccc-cccc-cccccccccccc.value" + ], + }, + "groups": [{"title": "Inline Group"}], + "comments": [], + "inputs": [], + "outputs": [], + "module_properties": {}, + } + + +def _make_cycle_graph() -> dict: + """Two-node cycle: A.out -> B.in, B.out -> A.in.""" + + return { + "title": "Cycle", + "id": "ffffffff-ffff-ffff-ffff-ffffffffffff", + "registry": "test/Cycle", + "base_type": "core/Graph", + "nodes": { + "11111111-1111-1111-1111-111111111111": { + "title": "A", + "id": "11111111-1111-1111-1111-111111111111", + "properties": {}, + "registry": "core/MakeBool", + "base_type": "core/Node", + }, + "22222222-2222-2222-2222-222222222222": { + "title": "B", + "id": "22222222-2222-2222-2222-222222222222", + "properties": {}, + "registry": "core/MakeBool", + "base_type": "core/Node", + }, + }, + "edges": { + "11111111-1111-1111-1111-111111111111.out": [ + "22222222-2222-2222-2222-222222222222.in" + ], + "22222222-2222-2222-2222-222222222222.out": [ + "11111111-1111-1111-1111-111111111111.in" + ], + }, + "groups": [], + "comments": [], + "inputs": [], + "outputs": [], + "module_properties": {}, + } + + +# --------------------------------------------------------------------------- +# Inline graph: per-function tests +# --------------------------------------------------------------------------- + + +def test_summarize_inline(): + g = _make_inline_graph() + s = analysis.summarize(g) + + assert s.title == "Inline Test Graph" + assert s.registry == "test/InlineGraph" + assert s.node_count == 4 + assert s.stage_node_count == 1 + assert s.nodes_by_category.get("core") == 3 + assert s.nodes_by_category.get("mystery") == 1 + assert s.group_titles == ["Inline Group"] + + +def test_list_nodes_filter_inline(): + g = _make_inline_graph() + all_nodes = analysis.list_nodes(g) + assert len(all_nodes) == 4 + + only_core = analysis.list_nodes(g, registry_pattern="core/") + assert len(only_core) == 3 + assert all(e.registry and e.registry.startswith("core/") for e in only_core) + + by_title = analysis.list_nodes(g, title_pattern="watch") + assert len(by_title) == 1 + assert by_title[0].title == "Watcher" + + +def test_get_node_inline_with_short_prefix(): + g = _make_inline_graph() + d = analysis.get_node(g, "bbbb2222") + assert d.title == "Stage 2" + assert d.short_id == "bbbb2222" + assert d.registered is True + + # state input is connected + state_in = next(i for i in d.inputs if i.name == "state") + assert state_in.connected + assert state_in.source is not None + assert state_in.source.short_id == "aaaa1111" + + # state output is connected + state_out = next(o for o in d.outputs if o.name == "state") + assert any(c.short_id == "cccc3333" for c in state_out.consumers) + + # mystery node has no static class but should still resolve via edges + mystery = analysis.get_node(g, "dddd4444") + assert mystery.registered is False + + +def test_resolve_short_prefix_unique_ambiguous_missing(): + g = _make_inline_graph() + # add another aaaa-prefixed node to make a collision + g["nodes"]["aaaa9999-zzzz-zzzz-zzzz-zzzzzzzzzzzz"] = { + "title": "Other", + "id": "aaaa9999-zzzz-zzzz-zzzz-zzzzzzzzzzzz", + "properties": {}, + "registry": "core/MakeBool", + "base_type": "core/Node", + } + + # unique-enough prefix still works + assert analysis.resolve_node_id(g, "aaaa1111") == "aaaa1111-aaaa-aaaa-aaaa-aaaaaaaaaaaa" + + # ambiguous + with pytest.raises(ValueError, match="Ambiguous"): + analysis.resolve_node_id(g, "aaaa") + + # missing + with pytest.raises(ValueError, match="No node"): + analysis.resolve_node_id(g, "ffff0000") + + +def test_list_edges_filter(): + g = _make_inline_graph() + all_edges = analysis.list_edges(g) + assert len(all_edges) == 2 + + only_from_a = analysis.list_edges(g, from_node="aaaa1111") + assert len(only_from_a) == 1 + assert only_from_a[0].target_short == "bbbb2222" + + only_to_c = analysis.list_edges(g, to_node="cccc3333") + assert len(only_to_c) == 1 + assert only_to_c[0].source_short == "bbbb2222" + + +def test_feeds_and_consumers(): + g = _make_inline_graph() + + f = analysis.feeds(g, "bbbb2222.state") + assert f.source is not None + assert f.source.short_id == "aaaa1111" + assert f.source_node_registry == "core/MakeBool" + + f_unconnected = analysis.feeds(g, "aaaa1111.in") + assert f_unconnected.source is None + + c = analysis.consumers(g, "bbbb2222.state") + assert len(c.consumers) == 1 + assert c.consumers[0].target_short == "cccc3333" + + +def test_trace_forward_inline(): + g = _make_inline_graph() + t = analysis.trace_forward(g, "aaaa1111", depth=3) + assert t.short_id == "aaaa1111" + assert len(t.children) == 1 + assert t.children[0].short_id == "bbbb2222" + assert t.children[0].children[0].short_id == "cccc3333" + + +def test_trace_backward_inline(): + g = _make_inline_graph() + t = analysis.trace_backward(g, "cccc3333", depth=3) + assert t.short_id == "cccc3333" + assert t.children[0].short_id == "bbbb2222" + assert t.children[0].children[0].short_id == "aaaa1111" + + +def test_trace_cycle_safety(): + g = _make_cycle_graph() + t = analysis.trace_forward(g, "11111111", depth=10) + # walk: A -> B -> A(cycle marker) + assert t.short_id == "11111111" + assert len(t.children) == 1 + b = t.children[0] + assert b.short_id == "22222222" + assert len(b.children) == 1 + cycled = b.children[0] + assert cycled.cycle is True + assert cycled.short_id == "11111111" + # no infinite recursion below the cycle marker + assert cycled.children == [] + + +def test_stage_map_inline(): + g = _make_inline_graph() + sm = analysis.stage_map(g) + # one chain has the stage node, the other (mystery) is unstaged + assert len(sm.stage_nodes) == 1 + assert sm.stage_nodes[0].stage == 2 + assert sm.stage_nodes[0].short_id == "bbbb2222" + assert "dddd4444-dddd-dddd-dddd-dddddddddddd" in sm.unstaged_node_ids + + +def test_check_registries_inline_finds_unknown(): + g = _make_inline_graph() + r = analysis.check_registries(g) + unknowns = {p.registry for p in r.unknown_registries} + # mystery/Unknown is fake, plus the top-level test/InlineGraph + assert "mystery/Unknown" in unknowns + assert "test/InlineGraph" in unknowns + # core/Stage and core/MakeBool are real and should NOT appear + assert "core/Stage" not in unknowns + assert "core/MakeBool" not in unknowns + + +# --------------------------------------------------------------------------- +# Real-graph smoke tests +# --------------------------------------------------------------------------- + + +def test_real_graph_summarize_and_lists(): + assert DIRECTOR_FIXTURE.exists(), DIRECTOR_FIXTURE + g = load_graph(DIRECTOR_FIXTURE) + + s = analysis.summarize(g) + assert s.node_count > 0 + assert s.stage_node_count >= 1 + + nodes = analysis.list_nodes(g) + assert nodes # non-empty + + sm = analysis.stage_map(g) + assert sm.stage_nodes # at least one Stage node + + chk = analysis.check_registries(g) + # Shipped graph should not contain unknown registries. + assert chk.unknown_registries == [], [p.model_dump() for p in chk.unknown_registries] + assert chk.unknown_base_types == [], [p.model_dump() for p in chk.unknown_base_types] + + +def test_loader_missing_file(tmp_path): + with pytest.raises(GraphLoadError): + load_graph(tmp_path / "nope.json") + + +def test_loader_bad_json(tmp_path): + p = tmp_path / "bad.json" + p.write_text("not json", encoding="utf-8") + with pytest.raises(GraphLoadError): + load_graph(p) + + +# --------------------------------------------------------------------------- +# CLI smoke tests +# --------------------------------------------------------------------------- + + +def _run_cli(argv: list[str]) -> tuple[int, str, str]: + parser = cli.build_parser() + args = parser.parse_args(argv) + out = io.StringIO() + err = io.StringIO() + code = cli._run(args, out, err) + return code, out.getvalue(), err.getvalue() + + +def test_cli_summary_text(): + code, out, err = _run_cli(["summary", str(DIRECTOR_FIXTURE)]) + assert code == 0, err + assert "title:" in out + assert "nodes:" in out + assert "stage nodes:" in out + + +def test_cli_summary_json_is_valid(): + code, out, err = _run_cli(["summary", str(DIRECTOR_FIXTURE), "--json"]) + assert code == 0, err + parsed = json.loads(out) + assert parsed["registry"] == "agents/director/chat/directorActionDirectStoryArc" + assert parsed["node_count"] > 0 + + +def test_cli_check_registries_real_graph_returns_zero(): + code, out, err = _run_cli(["check-registries", str(DIRECTOR_FIXTURE)]) + assert code == 0, err + out + assert "ok:" in out + + +def test_cli_check_registries_unknown_exits_2(tmp_path): + g = _make_inline_graph() + p = tmp_path / "g.json" + p.write_text(json.dumps(g), encoding="utf-8") + code, out, err = _run_cli(["check-registries", str(p)]) + assert code == cli.EXIT_REGISTRY_PROBLEMS + assert "mystery/Unknown" in out + + +def test_cli_node_with_short_prefix(): + code, out, err = _run_cli( + ["node", str(DIRECTOR_FIXTURE), "ccb39d43"] + ) + # ccb39d43 is the top-level graph id, not a node id - expect error + assert code == cli.EXIT_STRUCT_ERROR + assert "No node" in err or "matches" in err diff --git a/tests/test_nodegraph_writer.py b/tests/test_nodegraph_writer.py new file mode 100644 index 00000000..0260473a --- /dev/null +++ b/tests/test_nodegraph_writer.py @@ -0,0 +1,1248 @@ +""" +Tests for ``talemate.game.engine.nodes.tools`` mutation + layout APIs. + +See also ``tests/test_nodegraph_tools.py`` for the analysis layer. +""" + +from __future__ import annotations + +import uuid +from pathlib import Path + +import pytest + +from talemate.game.engine.nodes.tools import ( + AlreadyConnectedError, + CycleError, + DynamicInputError, + GraphWriter, + GroupColor, + GroupError, + LayoutError, + LayoutOptions, + NodeNotFoundError, + NotConnectedError, + UnknownPropertyError, + UnknownRegistryError, + UnknownSocketError, + add_group, + analysis, + ensure_registry_loaded, + get_node_metadata, + layout_graph, + load_graph, +) +from talemate.game.engine.nodes.tools import writer as writer_mod +from talemate.game.engine.nodes.tools.layout import ( + WCCKind, + _apply_estimated_heights, + _classify_wcc, + _estimate_height, + _wcc_over_subset, +) + +REPO_ROOT = Path(__file__).resolve().parent.parent +DIRECTOR_FIXTURE = ( + REPO_ROOT + / "src/talemate/agents/director/modules/director-action-direct-story-arc.json" +) +ROLL_DICE_FIXTURE = ( + REPO_ROOT / "scenes/infinity-quest/nodes/roll-dice.json" +) + + +# --------------------------------------------------------------------------- +# Helpers / fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True, scope="module") +def _load_registry(): + ensure_registry_loaded() + yield + + +def _empty_graph() -> dict: + """Minimal valid graph dict for writer tests.""" + + return { + "title": "Test", + "id": str(uuid.uuid4()), + "registry": "test/WriterGraph", + "base_type": "core/Graph", + "extends": None, + "properties": {}, + "nodes": {}, + "edges": {}, + "groups": [], + "comments": [], + "inputs": [], + "outputs": [], + "module_properties": {}, + } + + +def _node_at(graph: dict, nid: str) -> dict: + return graph["nodes"][nid] + + +# --------------------------------------------------------------------------- +# add_node / remove_node +# --------------------------------------------------------------------------- + + +def test_add_node_happy_path(): + g = _empty_graph() + w = GraphWriter(g) + nid = w.add_node( + "data/number/Random", + title="Roll", + properties={"method": "integer"}, + ) + + # id is a valid uuid4 string + uuid.UUID(nid) # raises if invalid + + node = _node_at(g, nid) + assert node["registry"] == "data/number/Random" + assert node["title"] == "Roll" + assert node["properties"] == {"method": "integer"} + assert node["base_type"] == "core/Node" + assert node["x"] == 0 + assert node["y"] == 0 + assert node["width"] == 210 + # Writer intentionally does NOT set a default height; layout owns + # height estimation so collision avoidance has a realistic rect. + assert "height" not in node + + +def test_add_node_respects_explicit_height(): + g = _empty_graph() + w = GraphWriter(g) + nid = w.add_node( + "core/Watch", + title="Pinned", + height=200, + ) + node = _node_at(g, nid) + assert node["height"] == 200 + + +def test_add_node_unknown_registry_raises(): + g = _empty_graph() + w = GraphWriter(g) + with pytest.raises(UnknownRegistryError, match="not registered"): + w.add_node("totally/Fake") + + +def test_add_node_dynamic_class_gets_dynamic_inputs_list(): + g = _empty_graph() + w = GraphWriter(g) + nid = w.add_node("data/string/AdvancedFormat", title="Format") + assert _node_at(g, nid).get("dynamic_inputs") == [] + + +def test_remove_node_prunes_edges(): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch", title="A") + b = w.add_node("core/Watch", title="B") + c = w.add_node("core/Watch", title="C") + + w.connect(a, "value", b, "value") + w.connect(b, "value", c, "value") + + writer_mod.remove_node(g, b) + + assert b not in g["nodes"] + # Every edge that referenced b should be gone + for src, targets in g["edges"].items(): + assert src.split(".", 1)[0] != b + for tgt in targets: + assert tgt.split(".", 1)[0] != b + + +def test_remove_node_missing_raises(): + g = _empty_graph() + w = GraphWriter(g) + with pytest.raises(NodeNotFoundError): + writer_mod.remove_node(g, "no-such-node") + + +# --------------------------------------------------------------------------- +# connect / disconnect +# --------------------------------------------------------------------------- + + +def test_connect_happy_path_creates_edge(): + g = _empty_graph() + w = GraphWriter(g) + get_min = w.add_node( + "state/GetState", + title="GET min", + properties={"name": "min", "scope": "local"}, + ) + roll = w.add_node( + "data/number/Random", + title="Roll", + properties={"method": "integer"}, + ) + + w.connect(get_min, "value", roll, "min") + + edges = g["edges"] + key = f"{get_min}.value" + assert key in edges + assert f"{roll}.min" in edges[key] + + +def test_connect_dotted_shortcut(): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch") + b = w.add_node("core/Watch") + w.connect(f"{a}.value", f"{b}.value") + assert f"{a}.value" in g["edges"] + + +def test_connect_unknown_source_socket(): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch") + b = w.add_node("core/Watch") + with pytest.raises(UnknownSocketError, match="output socket 'nope'"): + w.connect(a, "nope", b, "value") + + +def test_connect_unknown_target_socket(): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch") + b = w.add_node("core/Watch") + with pytest.raises(UnknownSocketError, match="input socket 'nope'"): + w.connect(a, "value", b, "nope") + + +def test_connect_already_connected_raises(): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch") + b = w.add_node("core/Watch") + c = w.add_node("core/Watch") + w.connect(a, "value", b, "value") + with pytest.raises(AlreadyConnectedError): + w.connect(c, "value", b, "value") + + +def test_connect_cycle_rolls_back_edge(): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch") + b = w.add_node("core/Watch") + w.connect(a, "value", b, "value") + + with pytest.raises(CycleError): + w.connect(b, "value", a, "value") + + # The rolled-back edge must not be present. + assert f"{b}.value" not in g["edges"] + # The original edge must still be there. + assert f"{a}.value" in g["edges"] + + +def test_connect_self_loop_is_cycle(): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch") + with pytest.raises(CycleError): + w.connect(a, "value", a, "value") + assert f"{a}.value" not in g["edges"] + + +def test_disconnect_happy_path_and_missing(): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch") + b = w.add_node("core/Watch") + w.connect(a, "value", b, "value") + w.disconnect(a, "value", b, "value") + assert f"{a}.value" not in g["edges"] + + with pytest.raises(NotConnectedError): + w.disconnect(a, "value", b, "value") + + +# --------------------------------------------------------------------------- +# dynamic inputs +# --------------------------------------------------------------------------- + + +def test_add_dynamic_input_on_dynamic_node(): + g = _empty_graph() + w = GraphWriter(g) + fmt = w.add_node("data/string/AdvancedFormat", title="Format") + w.add_dynamic_input(fmt, "item0", "any") + w.add_dynamic_input(fmt, "item1", "str") + dyn = _node_at(g, fmt)["dynamic_inputs"] + assert [d["name"] for d in dyn] == ["item0", "item1"] + assert dyn[1]["type"] == "str" + + +def test_add_dynamic_input_duplicate_raises(): + g = _empty_graph() + w = GraphWriter(g) + fmt = w.add_node("data/string/AdvancedFormat") + w.add_dynamic_input(fmt, "item0") + with pytest.raises(DynamicInputError, match="already exists"): + w.add_dynamic_input(fmt, "item0") + + +def test_add_dynamic_input_on_non_dynamic_raises(): + g = _empty_graph() + w = GraphWriter(g) + # core/Watch is a regular Node, not a DynamicSocketNodeBase subclass. + a = w.add_node("core/Watch") + with pytest.raises(DynamicInputError, match="DynamicSocketNodeBase"): + w.add_dynamic_input(a, "item0") + + +def test_remove_dynamic_input_sweeps_edges(): + g = _empty_graph() + w = GraphWriter(g) + fmt = w.add_node("data/string/AdvancedFormat") + src = w.add_node("core/Watch") + w.add_dynamic_input(fmt, "item0") + w.connect(src, "value", fmt, "item0") + assert f"{src}.value" in g["edges"] + + w.remove_dynamic_input(fmt, "item0") + assert f"{src}.value" not in g["edges"] + assert _node_at(g, fmt)["dynamic_inputs"] == [] + + with pytest.raises(DynamicInputError, match="No dynamic input"): + w.remove_dynamic_input(fmt, "item0") + + +# --------------------------------------------------------------------------- +# round-trip via load / save +# --------------------------------------------------------------------------- + + +def test_graphwriter_load_save_round_trip(tmp_path): + src = tmp_path / "orig.json" + g = _empty_graph() + a = str(uuid.uuid4()) + b = str(uuid.uuid4()) + g["nodes"] = { + a: { + "title": "A", + "id": a, + "properties": {}, + "registry": "core/Watch", + "base_type": "core/Node", + "x": 10, + "y": 10, + "width": 210, + "height": 100, + }, + b: { + "title": "B", + "id": b, + "properties": {}, + "registry": "core/Watch", + "base_type": "core/Node", + "x": 300, + "y": 10, + "width": 210, + "height": 100, + }, + } + g["edges"] = {f"{a}.value": [f"{b}.value"]} + import json + + src.write_text(json.dumps(g), encoding="utf-8") + + w = GraphWriter.load(src) + dest = tmp_path / "out.json" + saved = w.save(dest) + assert saved == dest + + reloaded = load_graph(dest) + summary = analysis.summarize(reloaded) + assert summary.node_count == 2 + + +def test_graphwriter_modify_and_reload(tmp_path): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch") + b = w.add_node("core/Watch") + w.connect(a, "value", b, "value") + + dest = tmp_path / "mod.json" + w.save(dest) + + reloaded = load_graph(dest) + assert analysis.summarize(reloaded).node_count == 2 + assert f"{a}.value" in reloaded["edges"] + + +def test_save_without_path_raises(): + w = GraphWriter(_empty_graph()) + with pytest.raises(writer_mod.WriterError): + w.save() + + +# --------------------------------------------------------------------------- +# layout +# --------------------------------------------------------------------------- + + +def _make_layout_graph_with_existing() -> tuple[dict, list[str]]: + """Graph with one existing node at (50, 20) plus three new nodes at origin.""" + + g = _empty_graph() + existing = str(uuid.uuid4()) + g["nodes"][existing] = { + "title": "Existing", + "id": existing, + "properties": {}, + "registry": "core/Watch", + "base_type": "core/Node", + "x": 50, + "y": 20, + "width": 200, + "height": 100, + } + + w = GraphWriter(g) + a = w.add_node("core/Watch", title="A") + b = w.add_node("core/Watch", title="B") + c = w.add_node("core/Watch", title="C") + w.connect(a, "value", b, "value") + w.connect(b, "value", c, "value") + + return g, [a, b, c] + + +def test_layout_anchor_right_places_offset_from_existing(): + g, target = _make_layout_graph_with_existing() + a, b, c = target + existing_x = 50 + existing_right = existing_x + 200 # 250 + + opts = LayoutOptions() + layout_graph(g, new_node_ids=target, anchor="right", options=opts) + + # All target nodes end up at x >= existing_right + 300 + first_col = existing_right + 300 + assert g["nodes"][a]["x"] == first_col + assert g["nodes"][b]["x"] == first_col + opts.col_width + assert g["nodes"][c]["x"] == first_col + 2 * opts.col_width + # And they share min_y with the existing node. + assert g["nodes"][a]["y"] == 20 + + +def test_layout_anchor_below_places_below_existing(): + g, target = _make_layout_graph_with_existing() + a, b, c = target + + layout_graph(g, new_node_ids=target, anchor="below") + + # Existing node ends at y = 120; target should be at y = 320. + assert g["nodes"][a]["y"] == 320 + assert g["nodes"][a]["x"] == 50 # same min_x as existing + + +def test_layout_anchor_full_starts_at_origin_and_relays_everything(): + g, target = _make_layout_graph_with_existing() + # anchor=full ignores new_node_ids and relays everything. + existing_id = next(nid for nid in g["nodes"] if nid not in target) + + layout_graph(g, anchor="full") + + xs = [n["x"] for n in g["nodes"].values()] + ys = [n["y"] for n in g["nodes"].values()] + assert min(xs) == 0 + assert min(ys) == 0 + # The existing node is now in the target set. + assert g["nodes"][existing_id]["x"] >= 0 + + +def test_layout_topological_ordering(): + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch", title="A") + b = w.add_node("core/Watch", title="B") + c = w.add_node("core/Watch", title="C") + w.connect(a, "value", b, "value") + w.connect(b, "value", c, "value") + + layout_graph(g, new_node_ids=[a, b, c], anchor="full") + + ax = g["nodes"][a]["x"] + bx = g["nodes"][b]["x"] + cx = g["nodes"][c]["x"] + assert ax < bx < cx + + +def test_layout_collision_avoidance_within_column(): + """Two isolated nodes both at depth 0 must land in different rows.""" + + g = _empty_graph() + w = GraphWriter(g) + a = w.add_node("core/Watch") + b = w.add_node("core/Watch") + + layout_graph(g, new_node_ids=[a, b], anchor="full") + + assert g["nodes"][a]["x"] == g["nodes"][b]["x"] + assert g["nodes"][a]["y"] != g["nodes"][b]["y"] + + +def test_layout_isolated_node_gets_column_zero(): + g = _empty_graph() + w = GraphWriter(g) + solo = w.add_node("core/Watch") + + layout_graph(g, new_node_ids=[solo], anchor="full") + + assert g["nodes"][solo]["x"] == 0 + assert g["nodes"][solo]["y"] == 0 + + +def test_layout_does_not_move_nodes_outside_target_set(): + g, target = _make_layout_graph_with_existing() + existing_id = next(nid for nid in g["nodes"] if nid not in target) + orig_x = g["nodes"][existing_id]["x"] + orig_y = g["nodes"][existing_id]["y"] + + layout_graph(g, new_node_ids=target, anchor="right") + + assert g["nodes"][existing_id]["x"] == orig_x + assert g["nodes"][existing_id]["y"] == orig_y + + +def test_layout_unknown_anchor_raises(): + g = _empty_graph() + with pytest.raises(LayoutError): + layout_graph(g, anchor="sideways") + + +def test_layout_unknown_target_raises(): + g = _empty_graph() + w = GraphWriter(g) + w.add_node("core/Watch") # layout needs at least one real node + with pytest.raises(LayoutError): + layout_graph(g, new_node_ids=["no-such-node"], anchor="right") + + +# --------------------------------------------------------------------------- +# combined smoke: writer + layout + analysis +# --------------------------------------------------------------------------- + + +def test_combined_smoke_build_layout_reload(tmp_path): + w = GraphWriter(_empty_graph()) + in_node = w.add_node( + "core/Input", + title="IN value", + properties={"input_type": "int", "input_name": "value"}, + ) + get_state = w.add_node( + "state/GetState", + title="GET local.max", + properties={"name": "max", "scope": "local"}, + ) + roll = w.add_node( + "data/number/Random", + title="Roll", + properties={"method": "integer"}, + ) + out_node = w.add_node( + "core/Output", + title="OUT result", + properties={"output_type": "int", "output_name": "result"}, + ) + + w.connect(in_node, "value", roll, "min") + w.connect(get_state, "value", roll, "max") + w.connect(roll, "result", out_node, "value") + + layout_graph( + w.graph, + new_node_ids=[in_node, get_state, roll, out_node], + anchor="full", + ) + + dest = tmp_path / "combined.json" + w.save(dest) + + reloaded = load_graph(dest) + summary = analysis.summarize(reloaded) + assert summary.node_count == 4 + + # Layout sanity: Random must sit right of both of its inputs. + nodes = reloaded["nodes"] + assert nodes[roll]["x"] > nodes[in_node]["x"] + assert nodes[roll]["x"] > nodes[get_state]["x"] + # And the output must sit right of Random. + assert nodes[out_node]["x"] > nodes[roll]["x"] + + +# --------------------------------------------------------------------------- +# layout v2: height estimation +# --------------------------------------------------------------------------- + + +def _expected_height(registry: str, options: LayoutOptions) -> int: + """Replicate the layout _estimate_height formula for test assertions. + + Mirrors ``layout._estimate_height`` (v2.1 formula): + + title_bar_height + + (max(num_inputs, num_outputs) + num_properties) * socket_row_height + + padding + + dynamic_socket_bonus (iff class is a DynamicSocketNodeBase) + + Floored at ``options.min_height``. + """ + + meta = get_node_metadata(registry) + rows = max(len(meta.inputs), len(meta.outputs)) + len(meta.properties) + estimated = ( + options.title_bar_height + + rows * options.socket_row_height + + options.padding + ) + if meta.is_dynamic: + estimated += options.dynamic_socket_bonus + return max(estimated, options.min_height) + + +def test_layout_estimates_height_for_target_nodes(): + """A target node with no height gets an estimate matching the formula.""" + + g = _empty_graph() + w = GraphWriter(g) + nid = w.add_node("state/SetState") + # Writer must not have set a height. + assert "height" not in g["nodes"][nid] + + opts = LayoutOptions() + layout_graph(g, new_node_ids=[nid], anchor="full", options=opts) + + expected = _expected_height("state/SetState", opts) + assert g["nodes"][nid]["height"] == expected + assert expected > opts.min_height # sanity: SetState has real sockets + + +def test_layout_height_formula_is_deterministic(): + """Direct unit test on _estimate_height to pin the v2.1 formula.""" + + g = _empty_graph() + w = GraphWriter(g) + nid = w.add_node("state/SetState") + opts = LayoutOptions() + + meta = get_node_metadata("state/SetState") + rows = max(len(meta.inputs), len(meta.outputs)) + len(meta.properties) + manual = ( + opts.title_bar_height + + rows * opts.socket_row_height + + opts.padding + ) + if meta.is_dynamic: + manual += opts.dynamic_socket_bonus + manual = max(manual, opts.min_height) + assert _estimate_height(g["nodes"][nid], opts) == manual + + +def test_layout_overwrites_target_height_even_when_preset(): + """Target-set height is always replaced, even if the caller pinned one.""" + + g = _empty_graph() + w = GraphWriter(g) + nid = w.add_node("core/Watch", height=999) + assert g["nodes"][nid]["height"] == 999 # pinned by writer argument + + opts = LayoutOptions() + layout_graph(g, new_node_ids=[nid], anchor="full", options=opts) + + expected = _expected_height("core/Watch", opts) + assert g["nodes"][nid]["height"] == expected + assert g["nodes"][nid]["height"] != 999 + + +def test_layout_does_not_touch_heights_outside_target_set(): + """Non-target nodes keep whatever height they had before layout ran.""" + + g = _empty_graph() + # Hand-build a non-target existing node with a pinned height. + existing_id = str(uuid.uuid4()) + g["nodes"][existing_id] = { + "title": "Existing", + "id": existing_id, + "properties": {}, + "registry": "core/Watch", + "base_type": "core/Node", + "x": 0, + "y": 0, + "width": 210, + "height": 777, + } + + w = GraphWriter(g) + new_id = w.add_node("core/Watch") + + layout_graph(g, new_node_ids=[new_id], anchor="right") + + # Existing node keeps its pinned height + assert g["nodes"][existing_id]["height"] == 777 + # New node got an estimate + assert g["nodes"][new_id]["height"] != 777 + assert g["nodes"][new_id]["height"] == _expected_height( + "core/Watch", LayoutOptions() + ) + + +def test_apply_estimated_heights_skips_unknown_registry(): + """Helper falls back to min_height for a registry not in NODES.""" + + g = _empty_graph() + bogus = str(uuid.uuid4()) + g["nodes"][bogus] = { + "title": "Mystery", + "id": bogus, + "properties": {}, + "registry": "totally/Fake", + "base_type": "core/Node", + "x": 0, + "y": 0, + "width": 210, + } + opts = LayoutOptions() + _apply_estimated_heights(g, {bogus}, opts) + assert g["nodes"][bogus]["height"] == opts.min_height + + +# --------------------------------------------------------------------------- +# layout v2: WCC classification +# --------------------------------------------------------------------------- + + +def test_wcc_classify_input_output_passthrough_is_top(): + """A 2-node Input->Output chain is TOP (passthrough).""" + + g = _empty_graph() + w = GraphWriter(g) + in_node = w.add_node( + "core/Input", + properties={"input_type": "any", "input_name": "state"}, + ) + out_node = w.add_node( + "core/Output", + properties={"output_type": "any", "output_name": "state"}, + ) + w.connect(in_node, "value", out_node, "value") + + components = _wcc_over_subset(g, {in_node, out_node}) + assert len(components) == 1 + info = _classify_wcc(g, components[0]) + assert info.kind is WCCKind.TOP + + +def test_wcc_classify_stage_wins_over_input(): + """A chain containing both Input and Stage classifies as STAGE.""" + + g = _empty_graph() + w = GraphWriter(g) + in_node = w.add_node( + "core/Input", + properties={"input_type": "int", "input_name": "value"}, + ) + stage = w.add_node("core/Stage", properties={"stage": 3}) + w.connect(in_node, "value", stage, "state") + + info = _classify_wcc(g, sorted([in_node, stage])) + assert info.kind is WCCKind.STAGE + assert info.stage_value == 3 + + +def test_wcc_classify_stage_value_is_min_across_multiple_stage_nodes(): + """If a WCC contains multiple Stage nodes, stage_value is the min.""" + + g = _empty_graph() + w = GraphWriter(g) + s_low = w.add_node("core/Stage", properties={"stage": 1}) + s_high = w.add_node("core/Stage", properties={"stage": 5}) + # Link them so they're one WCC. + w.connect(s_low, "state", s_high, "state") + + info = _classify_wcc(g, sorted([s_low, s_high])) + assert info.kind is WCCKind.STAGE + assert info.stage_value == 1 + + +def test_wcc_classify_output_only_is_bottom(): + """A chain with only Output (no Input, no Stage) is BOTTOM.""" + + g = _empty_graph() + w = GraphWriter(g) + out_node = w.add_node( + "core/Output", + properties={"output_type": "int", "output_name": "result"}, + ) + watch = w.add_node("core/Watch") + w.connect(watch, "value", out_node, "value") + + info = _classify_wcc(g, sorted([out_node, watch])) + assert info.kind is WCCKind.BOTTOM + + +def test_wcc_classify_plain_watch_is_middle(): + """A lone Watch (no Input / Output / Stage) is MIDDLE.""" + + g = _empty_graph() + w = GraphWriter(g) + watch = w.add_node("core/Watch") + + info = _classify_wcc(g, [watch]) + assert info.kind is WCCKind.MIDDLE + + +# --------------------------------------------------------------------------- +# layout v2: vertical band ordering +# --------------------------------------------------------------------------- + + +def test_layout_vertical_band_ordering(): + """TOP on top, STAGE(0) above STAGE(1), BOTTOM on the bottom.""" + + g = _empty_graph() + w = GraphWriter(g) + + # TOP component: Input -> Watch (passthrough) + top_in = w.add_node( + "core/Input", + title="IN top", + properties={"input_type": "any", "input_name": "x"}, + ) + top_watch = w.add_node("core/Watch", title="top watch") + w.connect(top_in, "value", top_watch, "value") + + # STAGE(0) component + stage0 = w.add_node( + "core/Stage", title="Stage 0", properties={"stage": 0} + ) + stage0_watch = w.add_node("core/Watch", title="s0 watch") + w.connect(stage0_watch, "value", stage0, "state") + + # STAGE(1) component + stage1 = w.add_node( + "core/Stage", title="Stage 1", properties={"stage": 1} + ) + stage1_watch = w.add_node("core/Watch", title="s1 watch") + w.connect(stage1_watch, "value", stage1, "state") + + # BOTTOM component: Watch -> Output + bottom_watch = w.add_node("core/Watch", title="bottom watch") + bottom_out = w.add_node( + "core/Output", + title="OUT bottom", + properties={"output_type": "any", "output_name": "x"}, + ) + w.connect(bottom_watch, "value", bottom_out, "value") + + opts = LayoutOptions() + layout_graph(g, anchor="full", options=opts) + + def _comp_y(node_ids: list[str]) -> int: + return min(g["nodes"][nid]["y"] for nid in node_ids) + + def _comp_bottom(node_ids: list[str]) -> int: + return max( + g["nodes"][nid]["y"] + g["nodes"][nid]["height"] for nid in node_ids + ) + + top_y = _comp_y([top_in, top_watch]) + stage0_y = _comp_y([stage0, stage0_watch]) + stage1_y = _comp_y([stage1, stage1_watch]) + bottom_y = _comp_y([bottom_watch, bottom_out]) + + # Strict vertical ordering + assert top_y < stage0_y < stage1_y < bottom_y + + # Each band separated from the next by at least band_gap. + assert stage0_y - _comp_bottom([top_in, top_watch]) >= opts.band_gap + assert stage1_y - _comp_bottom([stage0, stage0_watch]) >= opts.band_gap + assert bottom_y - _comp_bottom([stage1, stage1_watch]) >= opts.band_gap + + # Bands don't vertically overlap. + assert _comp_bottom([top_in, top_watch]) < stage0_y + assert _comp_bottom([stage0, stage0_watch]) < stage1_y + assert _comp_bottom([stage1, stage1_watch]) < bottom_y + + +# --------------------------------------------------------------------------- +# layout v2.1: new formula, dynamic bonus, height-aware + predecessor-aware +# packing +# --------------------------------------------------------------------------- + + +def test_layout_v21_height_formula_matches_explicit_math(): + """Spell out the v2.1 height math for state/SetState explicitly.""" + + g = _empty_graph() + w = GraphWriter(g) + nid = w.add_node("state/SetState") + opts = LayoutOptions() + + meta = get_node_metadata("state/SetState") + rows = max(len(meta.inputs), len(meta.outputs)) + len(meta.properties) + expected = ( + opts.title_bar_height + + rows * opts.socket_row_height + + opts.padding + ) + # state/SetState is a regular node, not a DynamicSocketNodeBase. + assert meta.is_dynamic is False + expected = max(expected, opts.min_height) + + assert _estimate_height(g["nodes"][nid], opts) == expected + + +def test_layout_v21_dynamic_socket_bonus_applies_to_dynamic_class(): + """data/string/AdvancedFormat is a DynamicSocketNodeBase subclass.""" + + g = _empty_graph() + w = GraphWriter(g) + dyn_nid = w.add_node("data/string/AdvancedFormat") + + opts = LayoutOptions() + meta = get_node_metadata("data/string/AdvancedFormat") + assert meta.is_dynamic is True + + rows = max(len(meta.inputs), len(meta.outputs)) + len(meta.properties) + base = ( + opts.title_bar_height + + rows * opts.socket_row_height + + opts.padding + ) + expected = max(base + opts.dynamic_socket_bonus, opts.min_height) + assert _estimate_height(g["nodes"][dyn_nid], opts) == expected + + # Matched non-dynamic control: a node without the bonus should be + # exactly ``dynamic_socket_bonus`` shorter for the same row count. + # Compute a plain height using the same formula minus the bonus and + # assert the delta. + without_bonus = max(base, opts.min_height) + assert expected - without_bonus == opts.dynamic_socket_bonus + + +def test_layout_v21_height_aware_column_packing(): + """Two isolated nodes in the same column respect variable heights.""" + + from talemate.game.engine.nodes.tools.layout import _place_band + + opts = LayoutOptions() + g = _empty_graph() + + # Hand-build a tall pre-sized node so we can assert the packer + # clears its actual height rather than the legacy 160 stride. + tall_id = str(uuid.uuid4()) + g["nodes"][tall_id] = { + "title": "Tall", + "id": tall_id, + "properties": {}, + "registry": "core/Watch", + "base_type": "core/Node", + "x": 0, + "y": 0, + "width": 210, + "height": 200, + } + short_id = str(uuid.uuid4()) + g["nodes"][short_id] = { + "title": "Short", + "id": short_id, + "properties": {}, + "registry": "core/Watch", + "base_type": "core/Node", + "x": 0, + "y": 0, + "width": 210, + "height": 100, + } + + # No edges — both nodes land in column 0. Sorted id order + # determines which goes first. + placement = _place_band( + g, + sorted([tall_id, short_id]), + origin_x=0, + origin_y=0, + options=opts, + ) + + y_tall = g["nodes"][tall_id]["y"] + y_short = g["nodes"][short_id]["y"] + assert y_tall != y_short + top_id, bottom_id = ( + (tall_id, short_id) if y_tall < y_short else (short_id, tall_id) + ) + top_y = g["nodes"][top_id]["y"] + top_h = g["nodes"][top_id]["height"] + bottom_y = g["nodes"][bottom_id]["y"] + assert bottom_y >= top_y + top_h + opts.min_vertical_gap + # If the "top" node is the tall one, bottom_y must clear 200 + gap. + if top_id == tall_id: + assert bottom_y >= 200 + opts.min_vertical_gap + # Legacy stride would have put the second node at y=160; make sure + # we're not doing that anymore. + assert bottom_y != 160 + # placement.max_y covers both rects. + assert placement.max_y >= bottom_y + g["nodes"][bottom_id]["height"] + + +def test_layout_v21_predecessor_aware_row_placement_uncrosses_wires(): + """Crossed Input->SetState edges should be un-crossed by the layout pass. + + Fixture:: + + IN a ---> SET local.a + IN b ---> SET local.b + + With the original IDs chosen so that ``a`` and ``b`` would sort in + the "wrong" order in column 1 under pure stable-id ordering, we + should still see ``IN a`` and ``SET local.a`` share a y coordinate + after layout, and likewise for ``b``. + """ + + g = _empty_graph() + w = GraphWriter(g) + # Inputs in column 0 + in_a = w.add_node( + "core/Input", + title="IN a", + properties={"input_type": "any", "input_name": "a"}, + ) + in_b = w.add_node( + "core/Input", + title="IN b", + properties={"input_type": "any", "input_name": "b"}, + ) + # SetStates in column 1 -- connect a -> set_a, b -> set_b, but add + # them to the graph in reverse-dependency order so that stable-id + # sort inside the column puts set_b before set_a. + set_b = w.add_node( + "state/SetState", + title="SET local.b", + properties={"name": "b", "scope": "local"}, + ) + set_a = w.add_node( + "state/SetState", + title="SET local.a", + properties={"name": "a", "scope": "local"}, + ) + w.connect(in_a, "value", set_a, "value") + w.connect(in_b, "value", set_b, "value") + + layout_graph(g, anchor="full") + + # Each SetState must land at the same y as its source Input. + assert g["nodes"][in_a]["y"] == g["nodes"][set_a]["y"], ( + "IN a should line up with SET local.a after predecessor-aware placement" + ) + assert g["nodes"][in_b]["y"] == g["nodes"][set_b]["y"], ( + "IN b should line up with SET local.b after predecessor-aware placement" + ) + + +def test_layout_v21_col_width_default_is_wider(): + """col_width default bumped from 260 to 360 in v2.1.""" + + opts = LayoutOptions() + assert opts.col_width == 360 + + +# --------------------------------------------------------------------------- +# add_group / GroupColor +# --------------------------------------------------------------------------- + + +def _empty_graph() -> dict: + return { + "title": "Test", + "id": str(uuid.uuid4()), + "registry": "test/Test", + "base_type": "core/Graph", + "properties": {}, + "x": 0, "y": 0, "width": 200, "height": 100, + "collapsed": False, + "inherited": False, + "nodes": {}, + "edges": {}, + "groups": [], + "comments": [], + "extends": None, + "inputs": [], + "outputs": [], + "module_properties": {}, + } + + +def test_add_group_creates_group_around_nodes(): + """add_group computes bbox from node positions and produces a group dict.""" + g = GraphWriter(_empty_graph()) + a = g.add_node("core/Input", properties={"input_name": "a", "input_type": "any", "num": 0}) + b = g.add_node("core/Output", properties={"output_name": "a", "output_type": "any", "num": 0}) + g.connect(a, "value", b, "value") + layout_graph(g.graph, anchor="full") + + group = g.add_group("Input", GroupColor.INPUT, [a, b]) + + assert group in g.graph["groups"] + assert group["title"] == "Input" + assert group["color"] == "#88A" + assert group["color"] == GroupColor.INPUT + assert group["font_size"] == 24 + assert group["inherited"] is False + + # Bounding box should contain both placed nodes (with padding). + a_node = g.graph["nodes"][a] + b_node = g.graph["nodes"][b] + assert group["x"] <= min(a_node["x"], b_node["x"]) + assert group["y"] <= min(a_node["y"], b_node["y"]) + assert group["x"] + group["width"] >= max( + a_node["x"] + a_node["width"], b_node["x"] + b_node["width"] + ) + assert group["y"] + group["height"] >= max( + a_node["y"] + a_node["height"], b_node["y"] + b_node["height"] + ) + + +def test_add_group_resolves_short_prefix_ids(): + """GraphWriter.add_group accepts short-prefix node ids like other methods.""" + g = GraphWriter(_empty_graph()) + a = g.add_node("core/Input", properties={"input_name": "a", "input_type": "any", "num": 0}) + layout_graph(g.graph, anchor="full") + group = g.add_group("X", GroupColor.SPECIAL, [a[:8]]) + assert group["title"] == "X" + + +def test_add_group_empty_node_ids_raises(): + g = GraphWriter(_empty_graph()) + with pytest.raises(GroupError): + g.add_group("Empty", GroupColor.INPUT, []) + + +def test_add_group_unknown_node_raises(): + g = GraphWriter(_empty_graph()) + with pytest.raises(GroupError): + add_group(g.graph, "X", GroupColor.INPUT, ["does-not-exist"]) + + +def test_add_group_appends_to_existing_groups(): + """add_group never clobbers user-drawn groups; it appends.""" + g = GraphWriter(_empty_graph()) + g.graph["groups"].append({"title": "User Group", "x": 0, "y": 0, "width": 100, "height": 50}) + a = g.add_node("core/Input", properties={"input_name": "a", "input_type": "any", "num": 0}) + layout_graph(g.graph, anchor="full") + g.add_group("Auto", GroupColor.OUTPUT, [a]) + titles = [grp["title"] for grp in g.graph["groups"]] + assert "User Group" in titles + assert "Auto" in titles + assert len(g.graph["groups"]) == 2 + + +def test_group_color_palette_matches_litegraph_presets(): + """Sanity check the constants. These are the LiteGraph groupcolor hex + values; if any drift, the visual match with the frontend breaks.""" + assert GroupColor.INPUT == "#88A" + assert GroupColor.OUTPUT == "#8A8" + assert GroupColor.PROCESS == "#3f789e" + assert GroupColor.PREPARE == "#8AA" + assert GroupColor.VALIDATION == "#b58b2a" + assert GroupColor.FUNCTION == "#b06634" + assert GroupColor.SPECIAL == "#a1309b" + assert GroupColor.ERROR_HANDLING == "#A88" + assert GroupColor.UX == "#207e7e" + + +# --------------------------------------------------------------------------- +# add_node property validation (UnknownPropertyError) +# --------------------------------------------------------------------------- + + +def test_add_node_accepts_known_property(): + """Sanity check: a property that exists on the class is accepted. + + ``state/SetState`` declares ``name`` and ``scope`` as Fields — passing + either should work without raising. + """ + g = GraphWriter(_empty_graph()) + nid = g.add_node( + "state/SetState", + properties={"name": "foo", "scope": "local"}, + ) + assert g.graph["nodes"][nid]["properties"] == {"name": "foo", "scope": "local"} + + +def test_add_node_unknown_property_raises(): + """A property key not declared on the class raises UnknownPropertyError. + + This is the exact failure mode the agent hit when it set + ``properties={"agent": "summarizer"}`` on ``agents/GetAgent`` instead + of the real property name ``agent_name``. + """ + g = GraphWriter(_empty_graph()) + with pytest.raises(UnknownPropertyError) as excinfo: + g.add_node( + "agents/GetAgent", + properties={"agent": "summarizer"}, # bogus key + ) + msg = str(excinfo.value) + assert "'agent'" in msg + assert "agent_name" in msg # the real property name should appear in the hint + + +def test_add_node_unknown_property_lists_valid_options_in_message(): + """The error message must list the valid property names so the caller + can fix the typo without further lookup.""" + g = GraphWriter(_empty_graph()) + with pytest.raises(UnknownPropertyError) as excinfo: + g.add_node( + "state/SetState", + properties={"naem": "foo"}, # typo for "name" + ) + msg = str(excinfo.value) + assert "'naem'" in msg + assert "'name'" in msg + assert "'scope'" in msg + + +def test_add_node_no_properties_still_works(): + """Calling add_node without a properties arg should NOT trigger + validation (there's nothing to validate).""" + g = GraphWriter(_empty_graph()) + nid = g.add_node("core/Stage") + assert g.graph["nodes"][nid]["properties"] == {} + + +def test_add_node_empty_properties_dict_still_works(): + """Same as above but with an explicit empty dict.""" + g = GraphWriter(_empty_graph()) + nid = g.add_node("core/Stage", properties={}) + assert g.graph["nodes"][nid]["properties"] == {} + + +def test_add_node_validation_runs_before_mutation(): + """When add_node raises UnknownPropertyError, no node should have + been added to the graph (no partial-mutation footprint).""" + g = GraphWriter(_empty_graph()) + before_count = len(g.graph["nodes"]) + with pytest.raises(UnknownPropertyError): + g.add_node("state/SetState", properties={"bogus_key": "value"}) + after_count = len(g.graph["nodes"]) + assert after_count == before_count