From 19390c112a1da280eb83cc2fed5462a0e7a96135 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 1 Jun 2026 16:24:48 -0700 Subject: [PATCH] Add server-side TypeResolver for prompt-graph type resolution Resolves the concrete io_type of any output/input slot in a prompt by walking the graph, so API-submitted workflows (no frontend) and the execution engine agree on resolved types even when MatchType chains are involved. * New comfy_execution/type_resolver.py: TypeResolver class with output resolution (incl. MatchType template walking, cycle detection, depth cap, AnyType fallback + one-shot warning), input resolution (links and literals), is_output_list / is_input_list helpers, effective slot io_type peeling for dynamic wrappers (Autogrow -> wrapped element type, DynamicSlot -> underlying slot type), and bulk compute_live_input_types. * DynamicPrompt now lazily exposes get_type_resolver() and invalidates the resolver cache on add_ephemeral_node. * get_finalized_class_inputs / parse_class_inputs / DYNAMIC_INPUT_LOOKUP callable signature accept an optional live_input_types dict. Existing Autogrow/DynamicSlot/DynamicCombo expansions accept and ignore it; future per-type dynamic inputs use it as their discriminator. * validate_inputs and get_input_data both build live_input_types via the resolver and pass it through; validate_inputs also uses the resolver to determine received_type for linked inputs so MatchType chains in API workflows validate correctly. * validate_prompt builds one TypeResolver and shares it across all output-node validations to avoid re-walking chains. * tests-unit/execution_test/test_type_resolver.py covers V1 static return types, V1 wildcard warning behavior, MatchType resolution including first-wins, cycle termination, chain walking, input resolution, Autogrow peeling, list info, and cache invalidation. Amp-Thread-ID: https://ampcode.com/threads/T-019e8568-f382-743d-a97f-0de3ff29d501 Co-authored-by: Amp --- comfy_api/latest/_io.py | 48 ++- comfy_execution/graph.py | 15 + comfy_execution/type_resolver.py | 372 +++++++++++++++++ execution.py | 45 ++- .../execution_test/test_type_resolver.py | 377 ++++++++++++++++++ 5 files changed, 835 insertions(+), 22 deletions(-) create mode 100644 comfy_execution/type_resolver.py create mode 100644 tests-unit/execution_test/test_type_resolver.py diff --git a/comfy_api/latest/_io.py b/comfy_api/latest/_io.py index a3aa508ce..d6e92c17b 100644 --- a/comfy_api/latest/_io.py +++ b/comfy_api/latest/_io.py @@ -1090,7 +1090,7 @@ class Autogrow(ComfyTypeI): self.template.validate() @staticmethod - def _expand_schema_for_dynamic(out_dict: dict[str, Any], live_inputs: dict[str, Any], value: tuple[str, dict[str, Any]], input_type: str, curr_prefix: list[str] | None): + def _expand_schema_for_dynamic(out_dict: dict[str, Any], live_inputs: dict[str, Any], value: tuple[str, dict[str, Any]], input_type: str, curr_prefix: list[str] | None, live_input_types: dict[str, str] | None = None): # NOTE: purposely do not include self in out_dict; instead use only the template inputs # need to figure out names based on template type is_names = ("names" in value[1]["template"]) @@ -1139,7 +1139,7 @@ class Autogrow(ComfyTypeI): finalized_prefix = finalize_prefix(curr_prefix) out_dict["dynamic_paths"][finalized_prefix] = finalized_prefix out_dict["dynamic_paths_default_value"][finalized_prefix] = DynamicPathsDefaultValue.EMPTY_DICT - parse_class_inputs(out_dict, live_inputs, new_dict, curr_prefix) + parse_class_inputs(out_dict, live_inputs, new_dict, curr_prefix, live_input_types) @comfytype(io_type="COMFY_DYNAMICCOMBO_V3") class DynamicCombo(ComfyTypeI): @@ -1177,7 +1177,7 @@ class DynamicCombo(ComfyTypeI): input.validate() @staticmethod - def _expand_schema_for_dynamic(out_dict: dict[str, Any], live_inputs: dict[str, Any], value: tuple[str, dict[str, Any]], input_type: str, curr_prefix: list[str] | None): + def _expand_schema_for_dynamic(out_dict: dict[str, Any], live_inputs: dict[str, Any], value: tuple[str, dict[str, Any]], input_type: str, curr_prefix: list[str] | None, live_input_types: dict[str, str] | None = None): finalized_id = finalize_prefix(curr_prefix) if finalized_id in live_inputs: key = live_inputs[finalized_id] @@ -1189,7 +1189,7 @@ class DynamicCombo(ComfyTypeI): selected_option = option break if selected_option is not None: - parse_class_inputs(out_dict, live_inputs, selected_option["inputs"], curr_prefix) + parse_class_inputs(out_dict, live_inputs, selected_option["inputs"], curr_prefix, live_input_types) # add self to inputs out_dict[input_type][finalized_id] = value out_dict["dynamic_paths"][finalized_id] = finalize_prefix(curr_prefix, curr_prefix[-1]) @@ -1232,11 +1232,11 @@ class DynamicSlot(ComfyTypeI): input.validate() @staticmethod - def _expand_schema_for_dynamic(out_dict: dict[str, Any], live_inputs: dict[str, Any], value: tuple[str, dict[str, Any]], input_type: str, curr_prefix: list[str] | None): + def _expand_schema_for_dynamic(out_dict: dict[str, Any], live_inputs: dict[str, Any], value: tuple[str, dict[str, Any]], input_type: str, curr_prefix: list[str] | None, live_input_types: dict[str, str] | None = None): finalized_id = finalize_prefix(curr_prefix) if finalized_id in live_inputs: inputs = value[1]["inputs"] - parse_class_inputs(out_dict, live_inputs, inputs, curr_prefix) + parse_class_inputs(out_dict, live_inputs, inputs, curr_prefix, live_input_types) # add self to inputs out_dict[input_type][finalized_id] = value out_dict["dynamic_paths"][finalized_id] = finalize_prefix(curr_prefix, curr_prefix[-1]) @@ -1357,11 +1357,21 @@ class Range(ComfyTypeIO): }) -DYNAMIC_INPUT_LOOKUP: dict[str, Callable[[dict[str, Any], dict[str, Any], tuple[str, dict[str, Any]], str, list[str] | None], None]] = {} -def register_dynamic_input_func(io_type: str, func: Callable[[dict[str, Any], dict[str, Any], tuple[str, dict[str, Any]], str, list[str] | None], None]): +# Signature: (out_dict, live_inputs, value, input_type, curr_prefix, live_input_types) +# live_input_types is an optional {input_id: resolved_io_type} dict produced +# by comfy_execution.type_resolver.TypeResolver. Existing dynamic-input +# implementations may ignore it; future type-discriminated dynamic inputs +# (e.g. a per-connected-type variant of DynamicCombo) use it as their +# discriminator instead of literal live_inputs values. +_DynamicInputFunc = Callable[ + [dict[str, Any], dict[str, Any], tuple[str, dict[str, Any]], str, list[str] | None, dict[str, str] | None], + None, +] +DYNAMIC_INPUT_LOOKUP: dict[str, _DynamicInputFunc] = {} +def register_dynamic_input_func(io_type: str, func: _DynamicInputFunc): DYNAMIC_INPUT_LOOKUP[io_type] = func -def get_dynamic_input_func(io_type: str) -> Callable[[dict[str, Any], dict[str, Any], tuple[str, dict[str, Any]], str, list[str] | None], None]: +def get_dynamic_input_func(io_type: str) -> _DynamicInputFunc: return DYNAMIC_INPUT_LOOKUP[io_type] def setup_dynamic_input_funcs(): @@ -1709,7 +1719,19 @@ class Schema: ) return info -def get_finalized_class_inputs(d: dict[str, Any], live_inputs: dict[str, Any], include_hidden=False) -> tuple[dict[str, Any], V3Data]: +def get_finalized_class_inputs(d: dict[str, Any], live_inputs: dict[str, Any], include_hidden=False, live_input_types: dict[str, str] | None = None) -> tuple[dict[str, Any], V3Data]: + """Expand a node's V3 schema against a concrete prompt. + + Args: + d: ``INPUT_TYPES()``-shaped dict for the node. + live_inputs: Concrete ``{input_id: value}`` map from the prompt + (values may be links ``[node_id, slot_idx]`` or literals). + include_hidden: When True, retain hidden inputs in the returned dict. + live_input_types: Optional ``{input_id: resolved_io_type}`` map + produced by ``comfy_execution.type_resolver.TypeResolver``. Future + dynamic-input strategies that branch on connected type use this as + their discriminator. Existing dynamic types ignore it. + """ out_dict = { "required": {}, "optional": {}, @@ -1719,7 +1741,7 @@ def get_finalized_class_inputs(d: dict[str, Any], live_inputs: dict[str, Any], i d = d.copy() # ignore hidden for parsing hidden = d.pop("hidden", None) - parse_class_inputs(out_dict, live_inputs, d) + parse_class_inputs(out_dict, live_inputs, d, None, live_input_types) if hidden is not None and include_hidden: out_dict["hidden"] = hidden v3_data = {} @@ -1732,7 +1754,7 @@ def get_finalized_class_inputs(d: dict[str, Any], live_inputs: dict[str, Any], i v3_data["dynamic_paths_default_value"] = dynamic_paths_default_value return out_dict, hidden, v3_data -def parse_class_inputs(out_dict: dict[str, Any], live_inputs: dict[str, Any], curr_dict: dict[str, Any], curr_prefix: list[str] | None=None) -> None: +def parse_class_inputs(out_dict: dict[str, Any], live_inputs: dict[str, Any], curr_dict: dict[str, Any], curr_prefix: list[str] | None=None, live_input_types: dict[str, str] | None = None) -> None: for input_type, inner_d in curr_dict.items(): for id, value in inner_d.items(): io_type = value[0] @@ -1740,7 +1762,7 @@ def parse_class_inputs(out_dict: dict[str, Any], live_inputs: dict[str, Any], cu # dynamic inputs need to be handled with lookup functions dynamic_input_func = get_dynamic_input_func(io_type) new_prefix = handle_prefix(curr_prefix, id) - dynamic_input_func(out_dict, live_inputs, value, input_type, new_prefix) + dynamic_input_func(out_dict, live_inputs, value, input_type, new_prefix, live_input_types) else: # non-dynamic inputs get directly transferred finalized_id = finalize_prefix(curr_prefix, id) diff --git a/comfy_execution/graph.py b/comfy_execution/graph.py index 479ee8a53..ae5b338e7 100644 --- a/comfy_execution/graph.py +++ b/comfy_execution/graph.py @@ -26,6 +26,9 @@ class DynamicPrompt: self.ephemeral_prompt = {} self.ephemeral_parents = {} self.ephemeral_display = {} + # Lazily-built type resolver, scoped to this DynamicPrompt's lifetime. + # Invalidated whenever the graph mutates via add_ephemeral_node. + self._type_resolver = None def get_node(self, node_id): if node_id in self.ephemeral_prompt: @@ -41,6 +44,18 @@ class DynamicPrompt: self.ephemeral_prompt[node_id] = node_info self.ephemeral_parents[node_id] = parent_id self.ephemeral_display[node_id] = display_id + # Conservatively invalidate the entire resolver cache. Selective + # downstream invalidation would require topological info we don't have + # here cheaply; the resolver's cache is small and easy to rebuild. + if self._type_resolver is not None: + self._type_resolver.invalidate() + + def get_type_resolver(self): + """Lazily build and return the per-prompt TypeResolver.""" + if self._type_resolver is None: + from comfy_execution.type_resolver import TypeResolver + self._type_resolver = TypeResolver(self) + return self._type_resolver def get_real_node_id(self, node_id): while node_id in self.ephemeral_parents: diff --git a/comfy_execution/type_resolver.py b/comfy_execution/type_resolver.py new file mode 100644 index 000000000..975d37097 --- /dev/null +++ b/comfy_execution/type_resolver.py @@ -0,0 +1,372 @@ +"""Server-side type resolver for prompt graphs. + +Resolves the concrete io_type of an output slot or input slot by walking the +prompt graph. Handles: + + * Static V1/V3 ``RETURN_TYPES`` (returned as-is). + * V3 ``MatchType.Output`` (resolved by walking inputs that share the same + ``template_id`` until a concrete type is found). + * Cycles and unbounded recursion (terminates at ``AnyType`` with a one-shot + warning). + * Unknown / unresolvable / wildcard outputs (fall back to ``AnyType`` with a + one-shot warning). + +The resolver works against either a raw prompt dict +(``{node_id: {"class_type": str, "inputs": dict}}``) or a +``comfy_execution.graph.DynamicPrompt`` instance. + +All resolved values are plain strings, so the resolver state is trivially +serializable across processes if needed. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from comfy_api.latest import io +from comfy_api.internal import _ComfyNodeInternal + +# Sentinel for "type is unknown / wildcard". Matches AnyType.io_type ("*"). +ANY_TYPE: str = io.AnyType.io_type + +# Hard cap on resolver recursion depth. MatchType chains should never be +# anywhere near this deep; this is a belt-and-suspenders guard against malformed +# graphs and pathological cycles. +MAX_RESOLVE_DEPTH: int = 64 + + +class TypeResolver: + """Resolves concrete io_types for a prompt graph. + + Instantiate once per prompt (or per ``DynamicPrompt``) and reuse; results + are cached. Call :py:meth:`invalidate` (or :py:meth:`invalidate_node`) when + the underlying graph mutates (e.g. when an ephemeral node is added). + """ + + def __init__(self, prompt_source: Any): + """Args: + prompt_source: Either a ``DynamicPrompt`` (anything with + ``get_node(node_id)`` / ``has_node(node_id)``) or a plain + ``dict[node_id, {"class_type", "inputs"}]``. + """ + self._source = prompt_source + self._output_cache: dict[tuple[str, int], str] = {} + self._is_output_list_cache: dict[tuple[str, int], bool] = {} + self._warned: set[tuple[str, Any, str]] = set() + + # ---- prompt access ---------------------------------------------------- + def _has_node(self, node_id: str) -> bool: + if hasattr(self._source, "has_node"): + return self._source.has_node(node_id) + return node_id in self._source + + def _get_node(self, node_id: str) -> dict[str, Any] | None: + try: + if hasattr(self._source, "get_node"): + return self._source.get_node(node_id) + return self._source[node_id] + except Exception: + return None + + @staticmethod + def _get_class_def(class_type: str): + # Local import to avoid a hard import-cycle between nodes.py and + # comfy_execution at module-load time. + import nodes + return nodes.NODE_CLASS_MAPPINGS.get(class_type) + + # ---- cache management ------------------------------------------------- + def invalidate(self) -> None: + """Clear all cached resolutions. Cheap; call after any graph mutation.""" + self._output_cache.clear() + self._is_output_list_cache.clear() + # Intentionally do NOT clear self._warned: those messages are already + # logged and re-warning would just spam the log. + + def invalidate_node(self, node_id: str) -> None: + """Clear cached entries for a single node (e.g. after node-level expand).""" + for key in [k for k in self._output_cache if k[0] == node_id]: + del self._output_cache[key] + for key in [k for k in self._is_output_list_cache if k[0] == node_id]: + del self._is_output_list_cache[key] + + # ---- output resolution ----------------------------------------------- + def resolve_output_type(self, node_id: str, slot_idx: int, + _stack: frozenset[tuple[str, int]] | None = None) -> str: + """Return the resolved io_type string of ``node_id``'s output slot. + + Falls back to ``ANY_TYPE`` on cycle, depth-overflow, unknown class, + out-of-range slot, missing node, or unresolved MatchType template. + """ + cache_key = (node_id, slot_idx) + if cache_key in self._output_cache: + return self._output_cache[cache_key] + + if _stack is None: + _stack = frozenset() + if cache_key in _stack: + self._warn(node_id, slot_idx, "cycle detected during type resolution; defaulting to AnyType") + return ANY_TYPE + if len(_stack) >= MAX_RESOLVE_DEPTH: + self._warn(node_id, slot_idx, f"exceeded MAX_RESOLVE_DEPTH={MAX_RESOLVE_DEPTH}; defaulting to AnyType") + return ANY_TYPE + next_stack = _stack | {cache_key} + + if not self._has_node(node_id): + return ANY_TYPE + + node = self._get_node(node_id) + if node is None: + return ANY_TYPE + + class_type = node.get("class_type") + class_def = self._get_class_def(class_type) if class_type is not None else None + if class_def is None: + return ANY_TYPE + + try: + return_types = class_def.RETURN_TYPES + except Exception: + return ANY_TYPE + if return_types is None or slot_idx < 0 or slot_idx >= len(return_types): + return ANY_TYPE + + declared = return_types[slot_idx] + + # V3 nodes may have MatchType outputs that need to be traced through + # the schema. V1 nodes (and V3 nodes with plain outputs) just use the + # declared RETURN_TYPES string. + resolved = declared + if isinstance(class_def, type) and issubclass(class_def, _ComfyNodeInternal): + schema = getattr(class_def, "SCHEMA", None) + if schema is None: + # Trigger schema computation. RETURN_TYPES would have done this + # already, but be defensive. + try: + schema = class_def.GET_SCHEMA() + except Exception: + schema = None + if schema is not None and slot_idx < len(schema.outputs): + out = schema.outputs[slot_idx] + if isinstance(out, io.MatchType.Output): + resolved = self._resolve_match_template( + node_id, schema, out.template.template_id, next_stack + ) + + # Treat the legacy wildcard literally as AnyType. We warn only when the + # source node's *declared* type was already wildcard, so MatchType-style + # "no upstream connected" cases (which warn elsewhere) don't double-warn. + if isinstance(resolved, str) and resolved == ANY_TYPE and declared == ANY_TYPE: + self._warn( + node_id, slot_idx, + f"node '{class_type}' output slot {slot_idx} is wildcard; defaulting to AnyType", + ) + + if not isinstance(resolved, str): + # Non-string types (e.g., legacy combos passed as list) — bail to AnyType. + self._warn(node_id, slot_idx, + f"node '{class_type}' output slot {slot_idx} has non-string return type {type(resolved).__name__}; defaulting to AnyType") + resolved = ANY_TYPE + + self._output_cache[cache_key] = resolved + return resolved + + def _resolve_match_template(self, node_id: str, schema, template_id: str, + stack: frozenset[tuple[str, int]]) -> str: + """Resolve a MatchType.Output by inspecting the node's MatchType.Inputs + with the same template_id. + + Strategy (per design decision): walk inputs in schema order, pick the + FIRST concrete (non-AnyType) resolution. If none resolve, return + AnyType with a one-shot warning. + """ + node = self._get_node(node_id) + inputs_dict = (node or {}).get("inputs", {}) or {} + any_input_seen = False + for inp in schema.inputs: + if not isinstance(inp, io.MatchType.Input): + continue + if inp.template.template_id != template_id: + continue + any_input_seen = True + val = inputs_dict.get(inp.id) + if val is None: + continue + if isinstance(val, list) and len(val) == 2 and isinstance(val[0], str): + src_node, src_slot = val[0], val[1] + t = self.resolve_output_type(src_node, src_slot, stack) + if t != ANY_TYPE: + return t + # Literal value: a MatchType slot has no concrete declared type, so + # we cannot infer anything useful here. + if not any_input_seen: + # Schema declared a template_id with no Input bearing it. This is a + # node-author bug; warn once. + self._warn(node_id, None, + f"MatchType output template '{template_id}' has no matching Input on the node; defaulting to AnyType") + else: + self._warn(node_id, None, + f"MatchType template '{template_id}' has no bound concrete upstream input; defaulting to AnyType") + return ANY_TYPE + + def is_output_list(self, node_id: str, slot_idx: int) -> bool: + """Whether the source slot is declared as a list output (``OUTPUT_IS_LIST[idx]``).""" + cache_key = (node_id, slot_idx) + if cache_key in self._is_output_list_cache: + return self._is_output_list_cache[cache_key] + result = False + node = self._get_node(node_id) + if node is not None: + class_def = self._get_class_def(node.get("class_type")) + if class_def is not None: + lst = getattr(class_def, "OUTPUT_IS_LIST", None) + if lst is not None and 0 <= slot_idx < len(lst): + result = bool(lst[slot_idx]) + self._is_output_list_cache[cache_key] = result + return result + + # ---- input resolution ------------------------------------------------ + def resolve_input_type(self, node_id: str, input_id: str) -> str: + """Resolve the io_type of the value currently bound to a node's input. + + * If the value is a link, return the resolved type of the source slot. + * If the value is a literal, return the declared slot's effective + io_type (peeling dynamic-input wrappers — e.g. an Autogrow-of-Image + slot resolves to ``IMAGE``, not ``COMFY_AUTOGROW_V3``). + * If the value is missing or the slot is unknown, return ``ANY_TYPE``. + """ + node = self._get_node(node_id) + if node is None: + return ANY_TYPE + inputs = node.get("inputs", {}) or {} + if input_id not in inputs: + return ANY_TYPE + val = inputs[input_id] + if isinstance(val, list) and len(val) == 2 and isinstance(val[0], str): + return self.resolve_output_type(val[0], val[1]) + return self.get_declared_slot_io_type(node_id, input_id) + + def is_input_list(self, node_id: str, input_id: str) -> bool: + """Whether the value bound to ``input_id`` originates from a list output.""" + node = self._get_node(node_id) + if node is None: + return False + val = (node.get("inputs", {}) or {}).get(input_id) + if isinstance(val, list) and len(val) == 2 and isinstance(val[0], str): + return self.is_output_list(val[0], val[1]) + return False + + def get_declared_slot_io_type(self, node_id: str, input_id: str) -> str: + """Return the effective declared io_type of a node's input slot. + + Peels dynamic-input wrappers so that the user-facing element type is + returned: + + * Autogrow → wrapped template input's io_type + * DynamicSlot → underlying slot's io_type + * Anything else → the slot's own io_type + * DynamicCombo / unsupported → ``ANY_TYPE`` (the combo key is itself + dynamic, not a meaningful type for consumers) + """ + node = self._get_node(node_id) + if node is None: + return ANY_TYPE + class_def = self._get_class_def(node.get("class_type")) + if class_def is None: + return ANY_TYPE + + # Prefer V3 schema (carries dynamic-input wrapper info). + if isinstance(class_def, type) and issubclass(class_def, _ComfyNodeInternal): + schema = getattr(class_def, "SCHEMA", None) + if schema is None: + try: + class_def.GET_SCHEMA() + schema = getattr(class_def, "SCHEMA", None) + except Exception: + schema = None + if schema is not None: + # First, try a top-level input id match. + for inp in schema.inputs: + if inp.id == input_id: + return self._effective_io_type(inp) + # Then a nested match (DynamicSlot / DynamicCombo prefix.child). + if "." in input_id: + top, _, _ = input_id.partition(".") + for inp in schema.inputs: + if inp.id != top: + continue + for child in inp.get_all(): + if child is inp: + continue + if child.id == input_id.split(".", 1)[1]: + return self._effective_io_type(child) + # Fall through to V1 dict for hidden inputs etc. + + # V1 fallback: look at INPUT_TYPES() dict. + try: + inputs = class_def.INPUT_TYPES() + except Exception: + return ANY_TYPE + for section in ("required", "optional"): + section_d = inputs.get(section, {}) + if input_id in section_d: + entry = section_d[input_id] + if not entry: + return ANY_TYPE + t = entry[0] + if isinstance(t, str): + return t + if isinstance(t, list): + # legacy combo declared as a list of options. + return io.Combo.io_type + return ANY_TYPE + return ANY_TYPE + + @staticmethod + def _effective_io_type(inp) -> str: + """Return the consumer-facing io_type of a (possibly dynamic) input.""" + # Autogrow wraps a template input — the element type is what matters. + if isinstance(inp, io.Autogrow.Input): + try: + return inp.template.input.get_io_type() + except Exception: + return ANY_TYPE + # DynamicSlot wraps an underlying slot input. + if isinstance(inp, io.DynamicSlot.Input): + try: + return inp.slot.get_io_type() + except Exception: + return ANY_TYPE + # DynamicCombo's "type" is a key value selector, not a connection type. + if isinstance(inp, io.DynamicCombo.Input): + return ANY_TYPE + # Everything else: trust the input's declared io_type. + try: + return inp.get_io_type() + except Exception: + return ANY_TYPE + + # ---- bulk helpers ---------------------------------------------------- + def compute_live_input_types(self, node_id: str) -> dict[str, str]: + """Build the ``{input_id: resolved_io_type}`` map for a node. + + Used by :py:func:`comfy_api.latest._io.get_finalized_class_inputs` so + future dynamic-input expansion strategies (per-type DynamicType, etc.) + can branch on what was actually connected. + """ + node = self._get_node(node_id) + if node is None: + return {} + out: dict[str, str] = {} + for input_id in (node.get("inputs", {}) or {}).keys(): + out[input_id] = self.resolve_input_type(node_id, input_id) + return out + + # ---- diagnostics ----------------------------------------------------- + def _warn(self, node_id: str, slot_idx: Any, msg: str) -> None: + key = (node_id, slot_idx, msg) + if key in self._warned: + return + self._warned.add(key) + logging.warning("TypeResolver: node=%s slot=%s %s", node_id, slot_idx, msg) diff --git a/execution.py b/execution.py index 5246d651c..f7d0924c8 100644 --- a/execution.py +++ b/execution.py @@ -83,8 +83,9 @@ class IsChangedCache: self.is_changed[node_id] = node["is_changed"] return self.is_changed[node_id] - # Intentionally do not use cached outputs here. We only want constants in IS_CHANGED - input_data_all, _, v3_data = get_input_data(node["inputs"], class_def, node_id, None) + # Intentionally do not use cached outputs here. We only want constants in IS_CHANGED. + # Pass dynprompt so the TypeResolver can resolve link types for V3 dynamic schemas. + input_data_all, _, v3_data = get_input_data(node["inputs"], class_def, node_id, None, self.dynprompt) try: is_changed = await _async_map_node_over_list(self.prompt_id, node_id, class_def, input_data_all, is_changed_name, v3_data=v3_data) is_changed = await resolve_map_node_over_list_results(is_changed) @@ -158,7 +159,15 @@ def get_input_data(inputs, class_def, unique_id, execution_list=None, dynprompt= hidden_inputs_v3 = {} valid_inputs = class_def.INPUT_TYPES() if is_v3: - valid_inputs, hidden, v3_data = _io.get_finalized_class_inputs(valid_inputs, inputs) + # Build the type-resolution map for this node so dynamic schemas can + # branch on resolved upstream types (and not only on literal values). + # When no DynamicPrompt is available (e.g. some IsChangedCache paths + # in tests), live_input_types stays None and only literal-driven + # dynamic types continue to work. + live_input_types = None + if dynprompt is not None and hasattr(dynprompt, "get_type_resolver"): + live_input_types = dynprompt.get_type_resolver().compute_live_input_types(unique_id) + valid_inputs, hidden, v3_data = _io.get_finalized_class_inputs(valid_inputs, inputs, live_input_types=live_input_types) input_data_all = {} missing_keys = {} for x in inputs: @@ -821,9 +830,19 @@ class PromptExecutor: self._notify_prompt_lifecycle("end", prompt_id) -async def validate_inputs(prompt_id, prompt, item, validated, visiting=None): +async def validate_inputs(prompt_id, prompt, item, validated, visiting=None, type_resolver=None): + """Validate inputs for a single node, recursing into upstream nodes. + + ``type_resolver`` (a ``comfy_execution.type_resolver.TypeResolver``) is + built once at the top of the recursion and reused so MatchType chains are + only walked once. It also gives V3 dynamic schemas an accurate map of + resolved upstream types for API-submitted workflows. + """ if visiting is None: visiting = [] + if type_resolver is None: + from comfy_execution.type_resolver import TypeResolver + type_resolver = TypeResolver(prompt) unique_id = item if unique_id in validated: @@ -858,7 +877,8 @@ async def validate_inputs(prompt_id, prompt, item, validated, visiting=None): if issubclass(obj_class, _ComfyNodeInternal): obj_class: _io._ComfyNodeBaseInternal class_inputs = obj_class.INPUT_TYPES() - class_inputs, _, v3_data = _io.get_finalized_class_inputs(class_inputs, inputs) + live_input_types = type_resolver.compute_live_input_types(unique_id) + class_inputs, _, v3_data = _io.get_finalized_class_inputs(class_inputs, inputs, live_input_types=live_input_types) validate_function_name = "validate_inputs" validate_function = first_real_override(obj_class, validate_function_name) else: @@ -909,8 +929,11 @@ async def validate_inputs(prompt_id, prompt, item, validated, visiting=None): o_id = val[0] o_class_type = prompt[o_id]['class_type'] - r = nodes.NODE_CLASS_MAPPINGS[o_class_type].RETURN_TYPES - received_type = r[val[1]] + # Resolve the upstream output's effective type through the + # TypeResolver. This walks MatchType/template chains, so an API + # workflow without frontend-injected type metadata still gets the + # same answer the UI does. + received_type = type_resolver.resolve_output_type(o_id, val[1]) received_types[x] = received_type if 'input_types' not in validate_function_inputs and not validate_node_input(received_type, input_type): details = f"{x}, received_type({received_type}) mismatch input_type({input_type})" @@ -930,7 +953,7 @@ async def validate_inputs(prompt_id, prompt, item, validated, visiting=None): try: visiting.append(unique_id) try: - r = await validate_inputs(prompt_id, prompt, o_id, validated, visiting) + r = await validate_inputs(prompt_id, prompt, o_id, validated, visiting, type_resolver) finally: visiting.pop() if r[0] is False: @@ -1155,11 +1178,15 @@ async def validate_prompt(prompt_id, prompt, partial_execution_list: Union[list[ errors = [] node_errors = {} validated = {} + # Share one TypeResolver across all output validations so MatchType chains + # are only walked once per prompt. + from comfy_execution.type_resolver import TypeResolver + type_resolver = TypeResolver(prompt) for o in outputs: valid = False reasons = [] try: - m = await validate_inputs(prompt_id, prompt, o, validated) + m = await validate_inputs(prompt_id, prompt, o, validated, None, type_resolver) valid = m[0] reasons = m[1] except Exception as ex: diff --git a/tests-unit/execution_test/test_type_resolver.py b/tests-unit/execution_test/test_type_resolver.py new file mode 100644 index 000000000..f79384cc0 --- /dev/null +++ b/tests-unit/execution_test/test_type_resolver.py @@ -0,0 +1,377 @@ +"""Unit tests for :mod:`comfy_execution.type_resolver`. + +These tests stand up a small in-memory ``NODE_CLASS_MAPPINGS`` for the test +node classes (V1 and V3) and a fake DynamicPrompt-like dict, then verify the +resolver's behaviour for: + + * Static V1 ``RETURN_TYPES`` resolution. + * V1 wildcard outputs (must yield ``AnyType`` and warn once). + * V3 ``MatchType`` chains resolved via the downstream node's bound inputs. + * ``MatchType`` with no upstream bound (fall back to ``AnyType`` + warn). + * ``MatchType`` cycles (termination at ``AnyType`` + warn, no recursion blow-up). + * Deep chains capped by ``MAX_RESOLVE_DEPTH``. + * Input-type resolution for both literal values and links. + * Effective slot io_type peeling for ``Autogrow`` (returns the wrapped type). + * ``compute_live_input_types`` produces the right shape. + * Cache invalidation. + +The tests deliberately patch ``nodes.NODE_CLASS_MAPPINGS`` so they don't need +the whole ComfyUI bootstrap. +""" + +from __future__ import annotations + +import logging +import sys +import types as _pytypes + +import pytest + + +# --------------------------------------------------------------------------- +# Lightweight V1 test node factory +# --------------------------------------------------------------------------- + +def _v1_node(return_types: tuple[str, ...], input_types_dict: dict | None = None, + output_is_list: tuple[bool, ...] | None = None): + """Build a V1 node class with the given RETURN_TYPES / INPUT_TYPES().""" + if input_types_dict is None: + input_types_dict = {"required": {}} + + class _V1: + RETURN_TYPES = return_types + if output_is_list is not None: + OUTPUT_IS_LIST = output_is_list + + @classmethod + def INPUT_TYPES(cls): + return input_types_dict + + return _V1 + + +# --------------------------------------------------------------------------- +# Fixture: install fake nodes module before importing the resolver +# --------------------------------------------------------------------------- + +@pytest.fixture +def fake_nodes_module(): + """Install a synthetic ``nodes`` module with an empty mappings dict. + + Yields the mappings dict so tests can populate it per case. Cleans up + afterwards. We also have to make sure comfy_execution.type_resolver picks + up our fake module on its local import. + """ + real_nodes = sys.modules.get("nodes") + fake = _pytypes.ModuleType("nodes") + fake.NODE_CLASS_MAPPINGS = {} + sys.modules["nodes"] = fake + try: + yield fake.NODE_CLASS_MAPPINGS + finally: + if real_nodes is not None: + sys.modules["nodes"] = real_nodes + else: + del sys.modules["nodes"] + + +@pytest.fixture +def TypeResolver(fake_nodes_module): + # Late import so it picks up our fake `nodes` module. + from comfy_execution.type_resolver import TypeResolver as TR + return TR + + +# --------------------------------------------------------------------------- +# V1 resolution +# --------------------------------------------------------------------------- + +def test_v1_static_return_types_resolves(fake_nodes_module, TypeResolver): + fake_nodes_module["AddNode"] = _v1_node(("INT",)) + prompt = {"n1": {"class_type": "AddNode", "inputs": {}}} + r = TypeResolver(prompt) + assert r.resolve_output_type("n1", 0) == "INT" + + +def test_v1_wildcard_warns_once_and_returns_any(fake_nodes_module, TypeResolver, caplog): + fake_nodes_module["WildNode"] = _v1_node(("*",)) + prompt = {"n1": {"class_type": "WildNode", "inputs": {}}} + r = TypeResolver(prompt) + with caplog.at_level(logging.WARNING, logger="root"): + assert r.resolve_output_type("n1", 0) == "*" + # second call should still return * but not produce a second warning + assert r.resolve_output_type("n1", 0) == "*" + warnings = [rec for rec in caplog.records if "TypeResolver" in rec.message] + assert len(warnings) == 1, f"expected exactly one warning, got {warnings}" + + +def test_unknown_node_returns_any(fake_nodes_module, TypeResolver): + prompt = {"n1": {"class_type": "NopeNode", "inputs": {}}} + r = TypeResolver(prompt) + assert r.resolve_output_type("n1", 0) == "*" + + +def test_out_of_range_slot_returns_any(fake_nodes_module, TypeResolver): + fake_nodes_module["AddNode"] = _v1_node(("INT",)) + prompt = {"n1": {"class_type": "AddNode", "inputs": {}}} + r = TypeResolver(prompt) + assert r.resolve_output_type("n1", 5) == "*" + + +def test_missing_node_returns_any(fake_nodes_module, TypeResolver): + fake_nodes_module["AddNode"] = _v1_node(("INT",)) + prompt = {"n1": {"class_type": "AddNode", "inputs": {}}} + r = TypeResolver(prompt) + assert r.resolve_output_type("nonexistent", 0) == "*" + + +# --------------------------------------------------------------------------- +# is_output_list / is_input_list +# --------------------------------------------------------------------------- + +def test_is_output_list(fake_nodes_module, TypeResolver): + fake_nodes_module["ListNode"] = _v1_node(("IMAGE", "MASK"), output_is_list=(True, False)) + prompt = {"n1": {"class_type": "ListNode", "inputs": {}}} + r = TypeResolver(prompt) + assert r.is_output_list("n1", 0) is True + assert r.is_output_list("n1", 1) is False + + +def test_is_input_list_follows_link(fake_nodes_module, TypeResolver): + fake_nodes_module["ListNode"] = _v1_node(("IMAGE",), output_is_list=(True,)) + fake_nodes_module["Consumer"] = _v1_node( + ("INT",), + {"required": {"img": ("IMAGE",)}}, + ) + prompt = { + "src": {"class_type": "ListNode", "inputs": {}}, + "dst": {"class_type": "Consumer", "inputs": {"img": ["src", 0]}}, + } + r = TypeResolver(prompt) + assert r.is_input_list("dst", "img") is True + + +# --------------------------------------------------------------------------- +# V3 MatchType resolution +# --------------------------------------------------------------------------- + +def _make_switch_node_class(): + """Build a V3 Switch-like node with MatchType inputs/outputs.""" + from comfy_api.latest import io + + class Switch(io.ComfyNode): + @classmethod + def define_schema(cls): + template = io.MatchType.Template("switch") + return io.Schema( + node_id="TestSwitch", + inputs=[ + io.Boolean.Input("switch"), + io.MatchType.Input("on_false", template=template, optional=True), + io.MatchType.Input("on_true", template=template, optional=True), + ], + outputs=[io.MatchType.Output(template=template)], + ) + + @classmethod + def execute(cls, switch, on_false=None, on_true=None): + return io.NodeOutput(on_true if switch else on_false) + + # Force schema computation so SCHEMA / RETURN_TYPES are populated. + Switch.GET_SCHEMA() + return Switch + + +def test_matchtype_resolves_to_upstream_concrete(fake_nodes_module, TypeResolver): + fake_nodes_module["TestSwitch"] = _make_switch_node_class() + fake_nodes_module["ImageSrc"] = _v1_node(("IMAGE",)) + prompt = { + "img": {"class_type": "ImageSrc", "inputs": {}}, + "sw": { + "class_type": "TestSwitch", + "inputs": {"switch": True, "on_true": ["img", 0]}, + }, + } + r = TypeResolver(prompt) + assert r.resolve_output_type("sw", 0) == "IMAGE" + + +def test_matchtype_first_concrete_wins(fake_nodes_module, TypeResolver): + fake_nodes_module["TestSwitch"] = _make_switch_node_class() + fake_nodes_module["ImageSrc"] = _v1_node(("IMAGE",)) + fake_nodes_module["LatentSrc"] = _v1_node(("LATENT",)) + prompt = { + "img": {"class_type": "ImageSrc", "inputs": {}}, + "lat": {"class_type": "LatentSrc", "inputs": {}}, + "sw": { + "class_type": "TestSwitch", + "inputs": { + "switch": False, + "on_false": ["img", 0], # listed first in schema → wins + "on_true": ["lat", 0], + }, + }, + } + r = TypeResolver(prompt) + assert r.resolve_output_type("sw", 0) == "IMAGE" + + +def test_matchtype_no_bound_input_returns_any(fake_nodes_module, TypeResolver, caplog): + fake_nodes_module["TestSwitch"] = _make_switch_node_class() + prompt = {"sw": {"class_type": "TestSwitch", "inputs": {"switch": True}}} + r = TypeResolver(prompt) + with caplog.at_level(logging.WARNING, logger="root"): + assert r.resolve_output_type("sw", 0) == "*" + assert any("MatchType" in rec.message for rec in caplog.records) + + +def test_matchtype_skips_wildcard_input(fake_nodes_module, TypeResolver): + """If the first matched input resolves to AnyType, the resolver tries the next.""" + fake_nodes_module["TestSwitch"] = _make_switch_node_class() + fake_nodes_module["WildNode"] = _v1_node(("*",)) + fake_nodes_module["ImageSrc"] = _v1_node(("IMAGE",)) + prompt = { + "wild": {"class_type": "WildNode", "inputs": {}}, + "img": {"class_type": "ImageSrc", "inputs": {}}, + "sw": { + "class_type": "TestSwitch", + "inputs": { + "switch": True, + "on_false": ["wild", 0], + "on_true": ["img", 0], + }, + }, + } + r = TypeResolver(prompt) + assert r.resolve_output_type("sw", 0) == "IMAGE" + + +def test_matchtype_cycle_terminates_at_any(fake_nodes_module, TypeResolver): + """Two switches that feed each other must not recurse forever.""" + fake_nodes_module["TestSwitch"] = _make_switch_node_class() + prompt = { + "a": {"class_type": "TestSwitch", "inputs": {"switch": True, "on_true": ["b", 0]}}, + "b": {"class_type": "TestSwitch", "inputs": {"switch": True, "on_true": ["a", 0]}}, + } + r = TypeResolver(prompt) + # Must not raise / recurse forever; both resolve to AnyType. + assert r.resolve_output_type("a", 0) == "*" + assert r.resolve_output_type("b", 0) == "*" + + +def test_matchtype_chain_resolves_through(fake_nodes_module, TypeResolver): + """A → B → C → IMAGE: chain must walk all the way.""" + fake_nodes_module["TestSwitch"] = _make_switch_node_class() + fake_nodes_module["ImageSrc"] = _v1_node(("IMAGE",)) + prompt = { + "src": {"class_type": "ImageSrc", "inputs": {}}, + "a": {"class_type": "TestSwitch", "inputs": {"switch": True, "on_true": ["src", 0]}}, + "b": {"class_type": "TestSwitch", "inputs": {"switch": True, "on_true": ["a", 0]}}, + "c": {"class_type": "TestSwitch", "inputs": {"switch": True, "on_true": ["b", 0]}}, + } + r = TypeResolver(prompt) + assert r.resolve_output_type("c", 0) == "IMAGE" + + +# --------------------------------------------------------------------------- +# Input resolution and effective io_type peeling +# --------------------------------------------------------------------------- + +def test_resolve_input_type_literal_uses_declared(fake_nodes_module, TypeResolver): + fake_nodes_module["Sink"] = _v1_node(("INT",), {"required": {"steps": ("INT",)}}) + prompt = {"n1": {"class_type": "Sink", "inputs": {"steps": 20}}} + r = TypeResolver(prompt) + assert r.resolve_input_type("n1", "steps") == "INT" + + +def test_resolve_input_type_link(fake_nodes_module, TypeResolver): + fake_nodes_module["Src"] = _v1_node(("LATENT",)) + fake_nodes_module["Sink"] = _v1_node(("INT",), {"required": {"x": ("*",)}}) + prompt = { + "src": {"class_type": "Src", "inputs": {}}, + "sink": {"class_type": "Sink", "inputs": {"x": ["src", 0]}}, + } + r = TypeResolver(prompt) + assert r.resolve_input_type("sink", "x") == "LATENT" + + +def test_effective_slot_type_peels_autogrow(fake_nodes_module, TypeResolver): + from comfy_api.latest import io + + class AutogrowImg(io.ComfyNode): + @classmethod + def define_schema(cls): + template = io.Autogrow.TemplatePrefix( + input=io.Image.Input("img"), + prefix="img", + min=1, + ) + return io.Schema( + node_id="AutogrowImg", + inputs=[io.Autogrow.Input("imgs", template=template)], + outputs=[io.Image.Output()], + ) + + @classmethod + def execute(cls, imgs): + return io.NodeOutput(None) + + AutogrowImg.GET_SCHEMA() + fake_nodes_module["AutogrowImg"] = AutogrowImg + prompt = {"n1": {"class_type": "AutogrowImg", "inputs": {}}} + r = TypeResolver(prompt) + # The user-facing element type, not the autogrow wrapper. + assert r.get_declared_slot_io_type("n1", "imgs") == "IMAGE" + + +# --------------------------------------------------------------------------- +# compute_live_input_types +# --------------------------------------------------------------------------- + +def test_compute_live_input_types_mixes_links_and_literals(fake_nodes_module, TypeResolver): + fake_nodes_module["Src"] = _v1_node(("MODEL",)) + fake_nodes_module["Sink"] = _v1_node( + ("INT",), + {"required": {"model": ("MODEL",), "steps": ("INT",)}}, + ) + prompt = { + "src": {"class_type": "Src", "inputs": {}}, + "sink": { + "class_type": "Sink", + "inputs": {"model": ["src", 0], "steps": 20}, + }, + } + r = TypeResolver(prompt) + assert r.compute_live_input_types("sink") == {"model": "MODEL", "steps": "INT"} + + +# --------------------------------------------------------------------------- +# Cache invalidation +# --------------------------------------------------------------------------- + +def test_invalidate_clears_cache(fake_nodes_module, TypeResolver): + fake_nodes_module["Src"] = _v1_node(("IMAGE",)) + prompt = {"n1": {"class_type": "Src", "inputs": {}}} + r = TypeResolver(prompt) + assert r.resolve_output_type("n1", 0) == "IMAGE" + # Mutate the underlying class and invalidate; the resolver must re-read. + fake_nodes_module["Src"] = _v1_node(("LATENT",)) + r.invalidate() + assert r.resolve_output_type("n1", 0) == "LATENT" + + +def test_invalidate_node_only_clears_that_node(fake_nodes_module, TypeResolver): + fake_nodes_module["SrcA"] = _v1_node(("IMAGE",)) + fake_nodes_module["SrcB"] = _v1_node(("LATENT",)) + prompt = { + "a": {"class_type": "SrcA", "inputs": {}}, + "b": {"class_type": "SrcB", "inputs": {}}, + } + r = TypeResolver(prompt) + r.resolve_output_type("a", 0) + r.resolve_output_type("b", 0) + fake_nodes_module["SrcA"] = _v1_node(("MASK",)) + r.invalidate_node("a") + assert r.resolve_output_type("a", 0) == "MASK" + # b's cached result survives even though SrcB was unchanged + assert ("b", 0) in r._output_cache