From 2018722cc882cd91573f92ca81e0441f667a340c Mon Sep 17 00:00:00 2001 From: Will Miao Date: Mon, 13 Jul 2026 18:02:26 +0800 Subject: [PATCH] fix(registry): handle compound subgraph node IDs, add proactive node push from graph hooks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Handle compound node IDs (e.g. "252:0") from expanded group subgraphs to fix 400 Bad Request on workflows with group nodes - Frontend proactively pushes node data via afterConfigureGraph and LiteGraph hooks (onNodeAdded/onNodeRemoved/graphChanged), eliminating WebSocket round-trip latency for most "Send to Workflow" operations - Add content-fingerprint dedup to skip duplicate register-nodes POSTs - Fast-path cache returns immediately when tabs are registered (including 0-node registrations), avoiding unnecessary WS refresh cycles - Distinguish "Empty Registry" from other errors in standalone UI toast - Reduce WS refresh timeout 2s→0.5s, add cooldown and lock to prevent concurrent refresh storms - All [LM:Registry] logs at DEBUG level --- py/routes/handlers/misc_handlers.py | 138 +++++++++++++++++++++------- static/js/utils/uiHelpers.js | 2 + tests/routes/test_misc_routes.py | 48 ++++++++++ web/comfyui/workflow_registry.js | 75 ++++++++++++++- 4 files changed, 228 insertions(+), 35 deletions(-) diff --git a/py/routes/handlers/misc_handlers.py b/py/routes/handlers/misc_handlers.py index f270b24e..3325e4d4 100644 --- a/py/routes/handlers/misc_handlers.py +++ b/py/routes/handlers/misc_handlers.py @@ -573,12 +573,18 @@ class NodeRegistry: tab_nodes[nd["unique_id"]] = nd async with self._lock: + prev_count = len(self._tab_nodes.get(sid, {})) self._tab_nodes[sid] = tab_nodes self._waiting_clients.discard(sid) if not self._waiting_clients: self._ready.set() + total_tabs = len(self._tab_nodes) - logger.debug("Registered %s nodes from client %s", len(nodes), sid) + if len(nodes) != prev_count or len(nodes) > 0: + logger.debug( + "[LM:Registry] stored %s nodes (was %s) for client %s (total tabs: %s)", + len(nodes), prev_count, sid, total_tabs, + ) def prepare_for_refresh(self, active_sids: list[str]) -> None: """Set the list of client IDs we expect to hear from during the next refresh cycle.""" @@ -601,10 +607,17 @@ class NodeRegistry: longer connected.""" async with self._lock: # Garbage-collect stale entries (disconnected tabs) + stale_sids = [] if active_sids is not None: for sid in list(self._tab_nodes): if sid not in active_sids: + stale_sids.append(sid) del self._tab_nodes[sid] + if stale_sids: + logger.debug( + "[LM:Registry] GC pruned %s disconnected tabs: %s", + len(stale_sids), stale_sids, + ) merged: dict[str, dict] = {} tab_info: dict[str, dict] = {} @@ -3116,6 +3129,8 @@ class NodeRegistryHandler: self._node_registry = node_registry self._prompt_server = prompt_server self._standalone_mode = standalone_mode + self._refresh_lock = asyncio.Lock() + self._last_slow_path_ts: float = 0.0 async def register_nodes(self, request: web.Request) -> web.Response: try: @@ -3162,7 +3177,12 @@ class NodeRegistryHandler: ) graph_name = node.get("graph_name") try: - node["node_id"] = int(node_id) + # Handle compound node IDs from expanded group subgraphs, + # e.g. "252:0" → 0 (parent scope is already in graph_id) + if isinstance(node_id, str) and ":" in node_id: + node["node_id"] = int(node_id.rsplit(":", 1)[-1]) + else: + node["node_id"] = int(node_id) except (TypeError, ValueError): return web.json_response( { @@ -3203,39 +3223,95 @@ class NodeRegistryHandler: status=503, ) - # Snapshot of currently-connected ComfyUI tabs - active_sids = list(self._prompt_server.instance.sockets.keys()) - self._node_registry.prepare_for_refresh(active_sids) - - try: - self._prompt_server.instance.send_sync("lora_registry_refresh", {}) - logger.debug( - "Sent registry refresh request (expecting %s clients)", len(active_sids) - ) - except Exception as exc: - logger.error("Failed to send registry refresh message: %s", exc) - return web.json_response( - { - "success": False, - "error": "Communication Error", - "message": f"Failed to communicate with ComfyUI frontend: {exc}", - }, - status=500, - ) - - if not await self._node_registry.wait_for_all(timeout=2.0): - logger.warning( - "Registry refresh timeout after 2s (%s/%s clients responded)", - len(active_sids) - self._node_registry.pending_client_count, - len(active_sids), - ) - - # Re-read current sockets after the wait: a tab may have connected - # while we were waiting, and we don't want to garbage-collect it. current_sids = set(self._prompt_server.instance.sockets.keys()) + + # Fast path: if the frontend has already pushed node data (via + # afterConfigureGraph / graphChanged hooks), return it immediately + # without triggering a WebSocket round-trip. registry_info = await self._node_registry.get_merged_registry( active_sids=current_sids ) + if registry_info["tab_count"] > 0: + logger.debug( + "[LM:Registry] fast path: %s nodes across %s tabs %s", + registry_info["node_count"], + registry_info["tab_count"], + dict(registry_info.get("tabs", {})), + ) + return web.json_response({"success": True, "data": registry_info}) + + # Slow path: registry is empty — trigger refresh via WebSocket. + # Serialize with an async lock so concurrent callers don't all + # trigger separate WS refresh cycles. The second caller will + # re-check the fast path and (usually) find populated data. + async with self._refresh_lock: + # Re-check after acquiring the lock — another concurrent call + # may have populated the cache while we were waiting. + registry_info = await self._node_registry.get_merged_registry( + active_sids=current_sids + ) + if registry_info["tab_count"] > 0: + logger.debug( + "[LM:Registry] fast path after lock wait: %s nodes across %s tabs", + registry_info["node_count"], + registry_info["tab_count"], + ) + return web.json_response({"success": True, "data": registry_info}) + + # Cooldown: if the slow path ran recently (< 2 s) and + # returned empty, skip another WS round-trip. + elapsed = time.monotonic() - self._last_slow_path_ts + if elapsed < 2.0: + logger.debug( + "[LM:Registry] slow path cooldown (%.1fs since last refresh), returning empty", + elapsed, + ) + return web.json_response( + { + "success": False, + "error": "Empty Registry", + "message": "No workflow nodes found — ensure ComfyUI is open and the extension is loaded.", + }, + status=408, + ) + + logger.debug( + "[LM:Registry] slow path: cache empty, triggering WS refresh (%s connected tabs: %s)", + len(current_sids), list(current_sids)[:5], + ) + active_sids = list(current_sids) + self._node_registry.prepare_for_refresh(active_sids) + + try: + self._prompt_server.instance.send_sync("lora_registry_refresh", {}) + logger.debug( + "Sent registry refresh request (expecting %s clients)", len(active_sids) + ) + except Exception as exc: + logger.error("Failed to send registry refresh message: %s", exc) + return web.json_response( + { + "success": False, + "error": "Communication Error", + "message": f"Failed to communicate with ComfyUI frontend: {exc}", + }, + status=500, + ) + + if not await self._node_registry.wait_for_all(timeout=0.5): + logger.warning( + "Registry refresh timeout after 0.5s (%s/%s clients responded)", + len(active_sids) - self._node_registry.pending_client_count, + len(active_sids), + ) + + # Re-read current sockets after the wait: a tab may have connected + # while we were waiting, and we don't want to garbage-collect it. + current_sids = set(self._prompt_server.instance.sockets.keys()) + registry_info = await self._node_registry.get_merged_registry( + active_sids=current_sids + ) + self._last_slow_path_ts = time.monotonic() if registry_info["node_count"] == 0: logger.warning("No nodes registered after refresh") diff --git a/static/js/utils/uiHelpers.js b/static/js/utils/uiHelpers.js index a628027d..4350b9e7 100644 --- a/static/js/utils/uiHelpers.js +++ b/static/js/utils/uiHelpers.js @@ -552,6 +552,8 @@ async function fetchWorkflowRegistry() { if (!registryData.success) { if (registryData.error === 'Standalone Mode Active') { showToast('toast.general.cannotInteractStandalone', {}, 'warning'); + } else if (registryData.error === 'Empty Registry') { + showToast('uiHelpers.workflow.noSupportedNodes', {}, 'warning'); } else { showToast('toast.general.failedWorkflowInfo', {}, 'error'); } diff --git a/tests/routes/test_misc_routes.py b/tests/routes/test_misc_routes.py index 3e3a8031..b6a7566d 100644 --- a/tests/routes/test_misc_routes.py +++ b/tests/routes/test_misc_routes.py @@ -728,6 +728,54 @@ async def test_register_nodes_includes_capabilities(): assert stored_node["widget_names"] == ["ckpt_name"] +@pytest.mark.asyncio +async def test_register_nodes_accepts_compound_node_ids(): + """Subgraph nodes from expanded group nodes have compound IDs like '252:0'.""" + node_registry = NodeRegistry() + handler = NodeRegistryHandler( + node_registry=node_registry, + prompt_server=FakePromptServer, + standalone_mode=False, + ) + + request = FakeRequest( + json_data={ + "nodes": [ + { + "node_id": "252:0", + "graph_id": "252", + "type": "CheckpointLoaderSimple", + "title": "Checkpoint Loader (subgraph)", + }, + { + "node_id": "252:1", + "graph_id": "252", + "type": "CLIPLoader", + "title": "CLIP Loader (subgraph)", + }, + ], + "client_id": "test-client-1", + } + ) + + response = await handler.register_nodes(request) + payload = json.loads(response.text) + + assert response.status == 200 + assert payload["success"] is True + assert "2 nodes registered" in payload["message"] + + registry = await node_registry.get_merged_registry() + assert registry["node_count"] == 2 + + nodes_map = registry["nodes"] + assert "252:0" in nodes_map + assert "252:1" in nodes_map + assert nodes_map["252:0"]["id"] == 0 + assert nodes_map["252:0"]["graph_id"] == "252" + assert nodes_map["252:1"]["id"] == 1 + + @pytest.mark.asyncio async def test_update_node_widget_sends_payload(): send_calls: list[tuple[str, dict]] = [] diff --git a/web/comfyui/workflow_registry.js b/web/comfyui/workflow_registry.js index 29e968a7..947b0993 100644 --- a/web/comfyui/workflow_registry.js +++ b/web/comfyui/workflow_registry.js @@ -1,8 +1,10 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; -import { getAllGraphNodes, getNodeReference, getNodeFromGraph } from "./utils.js"; +import { getAllGraphNodes, getNodeReference, getNodeFromGraph, chainCallback } from "./utils.js"; import { ensureLmStyles } from "./lm_styles_loader.js"; +const DEBOUNCE_DELAY = 500; + const LORA_NODE_CLASSES = new Set([ "Lora Loader (LoraManager)", "Lora Stacker (LoraManager)", @@ -79,6 +81,7 @@ app.registerExtension({ setup() { ensureLmStyles(); + this._log("extension initialized, clientId=%s", api.clientId ?? api.initialClientId ?? "(pending)"); api.addEventListener("lora_registry_refresh", () => { this.refreshRegistry(); @@ -88,10 +91,64 @@ app.registerExtension({ this.applyWidgetUpdate(event?.detail ?? {}); }); - // React to marker changes from the Node Marker extension window.addEventListener("lm_marker_changed", () => { this.refreshRegistry(); }); + + this._hookGraphChanges(); + }, + + async afterConfigureGraph(_missingNodeTypes, _app) { + this._log("afterConfigureGraph: workflow loaded (%s missing types)", _missingNodeTypes?.length ?? 0); + await this.refreshRegistry(); + }, + + _hookGraphChanges() { + const graph = app.graph; + if (!graph) { + this._log("app.graph not available, skipping proactive hooks"); + return; + } + + let hooksInstalled = 0; + + const scheduleRefresh = (source) => { + if (this._debounceTimer != null) { + clearTimeout(this._debounceTimer); + } + this._debounceTimer = setTimeout(() => { + this._debounceTimer = null; + this.refreshRegistry(); + }, DEBOUNCE_DELAY); + }; + + try { + chainCallback(graph, "onNodeAdded", () => scheduleRefresh("onNodeAdded")); + chainCallback(graph, "onNodeRemoved", () => scheduleRefresh("onNodeRemoved")); + hooksInstalled += 2; + } catch (e) { + this._log("failed to chain LiteGraph hooks: %s", e.message); + } + + if (typeof api.addEventListener === "function") { + try { + api.addEventListener("graphChanged", () => scheduleRefresh("graphChanged")); + hooksInstalled += 1; + } catch (_e) { + // graphChanged may not be available on older ComfyUI versions + } + } + + this._log("%s proactive hooks installed on graph", hooksInstalled); + }, + + _log(format, ...args) { + const ts = new Date().toISOString().slice(11, 23); + let msg = format; + for (const arg of args) { + msg = msg.replace(/%s/g, String(arg)); + } + console.debug(`[LM:Registry ${ts}] ${msg}`); }, async refreshRegistry() { @@ -115,7 +172,6 @@ app.registerExtension({ const hasTextWidget = TEXT_CAPABLE_CLASSES.has(node.comfyClass); const markerRole = node.properties?.lm_marker_role ?? null; - // Skip nodes with no relevant capability UNLESS they are marked if (!supportsLora && !hasTargetWidget && !hasTextWidget && !markerRole) { continue; } @@ -146,6 +202,17 @@ app.registerExtension({ }); } + const clientId = api.clientId ?? api.initialClientId ?? ""; + + // Content-based dedup: skip POST if identical to last sent payload. + const fingerprint = JSON.stringify( + workflowNodes.map(n => `${n.graph_id}:${n.node_id}|${n.marker_role ?? ""}|${n.mode ?? 0}`).sort() + ); + if (fingerprint === this._lastFingerprint) { + return; + } + this._lastFingerprint = fingerprint; + const response = await fetch("/api/lm/register-nodes", { method: "POST", headers: { @@ -153,7 +220,7 @@ app.registerExtension({ }, body: JSON.stringify({ nodes: workflowNodes, - client_id: api.clientId ?? api.initialClientId ?? "", + client_id: clientId, }), });