From 2ff1d3d042755f0943d367997b849cdbf04e5cdb Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 28 Apr 2026 03:03:41 -0700 Subject: [PATCH] fix: make NodeReplaceManager.register() idempotent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the same (old_node_id, new_node_id) pair is registered more than once in the same process, the duplicate is now silently ignored (with a debug log). Previously every call appended a new entry, so reloading a custom node — e.g. via ComfyUI-Manager hot-reload — would accumulate stale duplicate replacements that surfaced through GET /node_replacements. Multiple distinct alternatives for the same old_node_id (different new_node_ids) are still preserved, matching the existing 'list of alternatives' design used by apply_replacements(). Adds unit tests in tests-unit/app_test/node_replace_manager_test.py covering normal registration, multi-alternative support, dedupe, and first-write-wins ordering. Co-authored-by: Amp Amp-Thread-ID: https://ampcode.com/threads/T-019dd37c-4751-72ef-9927-3182b5825db0 --- app/node_replace_manager.py | 20 ++++- .../app_test/node_replace_manager_test.py | 76 +++++++++++++++++++ 2 files changed, 94 insertions(+), 2 deletions(-) create mode 100644 tests-unit/app_test/node_replace_manager_test.py diff --git a/app/node_replace_manager.py b/app/node_replace_manager.py index d9aab5b22..72e8ac2b1 100644 --- a/app/node_replace_manager.py +++ b/app/node_replace_manager.py @@ -1,5 +1,7 @@ from __future__ import annotations +import logging + from aiohttp import web from typing import TYPE_CHECKING, TypedDict @@ -31,8 +33,22 @@ class NodeReplaceManager: self._replacements: dict[str, list[NodeReplace]] = {} def register(self, node_replace: NodeReplace): - """Register a node replacement mapping.""" - self._replacements.setdefault(node_replace.old_node_id, []).append(node_replace) + """Register a node replacement mapping. + + Idempotent: if a replacement with the same (old_node_id, new_node_id) + is already registered, the duplicate is ignored. This prevents stale + entries from accumulating when custom nodes are reloaded in the same + process (e.g. via ComfyUI-Manager). + """ + existing = self._replacements.setdefault(node_replace.old_node_id, []) + for entry in existing: + if entry.new_node_id == node_replace.new_node_id: + logging.debug( + "Node replacement %s -> %s already registered, ignoring duplicate.", + node_replace.old_node_id, node_replace.new_node_id, + ) + return + existing.append(node_replace) def get_replacement(self, old_node_id: str) -> list[NodeReplace] | None: """Get replacements for an old node ID.""" diff --git a/tests-unit/app_test/node_replace_manager_test.py b/tests-unit/app_test/node_replace_manager_test.py new file mode 100644 index 000000000..cf1be3034 --- /dev/null +++ b/tests-unit/app_test/node_replace_manager_test.py @@ -0,0 +1,76 @@ +"""Tests for NodeReplaceManager registration behavior.""" +import sys +import types + +# Stub `nodes` to avoid the heavy torch + ComfyUI import chain — register() +# only touches its own dict, so it doesn't actually need NODE_CLASS_MAPPINGS. +if "nodes" not in sys.modules: + _fake_nodes = types.ModuleType("nodes") + _fake_nodes.NODE_CLASS_MAPPINGS = {} + sys.modules["nodes"] = _fake_nodes + +from app.node_replace_manager import NodeReplaceManager # noqa: E402 + + +class FakeNodeReplace: + """Lightweight stand-in for comfy_api.latest._io.NodeReplace.""" + def __init__(self, new_node_id, old_node_id, old_widget_ids=None, + input_mapping=None, output_mapping=None): + self.new_node_id = new_node_id + self.old_node_id = old_node_id + self.old_widget_ids = old_widget_ids + self.input_mapping = input_mapping + self.output_mapping = output_mapping + + +def test_register_adds_replacement(): + manager = NodeReplaceManager() + manager.register(FakeNodeReplace(new_node_id="NewNode", old_node_id="OldNode")) + assert manager.has_replacement("OldNode") + assert len(manager.get_replacement("OldNode")) == 1 + + +def test_register_allows_multiple_alternatives_for_same_old_node(): + """Different new_node_ids for the same old_node_id should all be kept.""" + manager = NodeReplaceManager() + manager.register(FakeNodeReplace(new_node_id="AltA", old_node_id="OldNode")) + manager.register(FakeNodeReplace(new_node_id="AltB", old_node_id="OldNode")) + replacements = manager.get_replacement("OldNode") + assert len(replacements) == 2 + assert {r.new_node_id for r in replacements} == {"AltA", "AltB"} + + +def test_register_is_idempotent_for_duplicate_pair(): + """Re-registering the same (old_node_id, new_node_id) should be a no-op.""" + manager = NodeReplaceManager() + manager.register(FakeNodeReplace(new_node_id="NewNode", old_node_id="OldNode")) + manager.register(FakeNodeReplace(new_node_id="NewNode", old_node_id="OldNode")) + manager.register(FakeNodeReplace(new_node_id="NewNode", old_node_id="OldNode")) + assert len(manager.get_replacement("OldNode")) == 1 + + +def test_register_idempotent_preserves_first_registration(): + """First registration wins; later duplicates with different mappings are ignored.""" + manager = NodeReplaceManager() + first = FakeNodeReplace( + new_node_id="NewNode", old_node_id="OldNode", + input_mapping=[{"new_id": "a", "old_id": "x"}], + ) + second = FakeNodeReplace( + new_node_id="NewNode", old_node_id="OldNode", + input_mapping=[{"new_id": "b", "old_id": "y"}], + ) + manager.register(first) + manager.register(second) + replacements = manager.get_replacement("OldNode") + assert len(replacements) == 1 + assert replacements[0] is first + + +def test_register_dedupe_does_not_affect_other_old_nodes(): + manager = NodeReplaceManager() + manager.register(FakeNodeReplace(new_node_id="NewA", old_node_id="OldA")) + manager.register(FakeNodeReplace(new_node_id="NewA", old_node_id="OldA")) + manager.register(FakeNodeReplace(new_node_id="NewB", old_node_id="OldB")) + assert len(manager.get_replacement("OldA")) == 1 + assert len(manager.get_replacement("OldB")) == 1