fix(registry): handle compound subgraph node IDs, add proactive node push from graph hooks

- 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
This commit is contained in:
Will Miao
2026-07-13 18:02:26 +08:00
parent 9d85c2a44a
commit 2018722cc8
4 changed files with 228 additions and 35 deletions

View File

@@ -573,12 +573,18 @@ class NodeRegistry:
tab_nodes[nd["unique_id"]] = nd tab_nodes[nd["unique_id"]] = nd
async with self._lock: async with self._lock:
prev_count = len(self._tab_nodes.get(sid, {}))
self._tab_nodes[sid] = tab_nodes self._tab_nodes[sid] = tab_nodes
self._waiting_clients.discard(sid) self._waiting_clients.discard(sid)
if not self._waiting_clients: if not self._waiting_clients:
self._ready.set() 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: 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.""" """Set the list of client IDs we expect to hear from during the next refresh cycle."""
@@ -601,10 +607,17 @@ class NodeRegistry:
longer connected.""" longer connected."""
async with self._lock: async with self._lock:
# Garbage-collect stale entries (disconnected tabs) # Garbage-collect stale entries (disconnected tabs)
stale_sids = []
if active_sids is not None: if active_sids is not None:
for sid in list(self._tab_nodes): for sid in list(self._tab_nodes):
if sid not in active_sids: if sid not in active_sids:
stale_sids.append(sid)
del self._tab_nodes[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] = {} merged: dict[str, dict] = {}
tab_info: dict[str, dict] = {} tab_info: dict[str, dict] = {}
@@ -3116,6 +3129,8 @@ class NodeRegistryHandler:
self._node_registry = node_registry self._node_registry = node_registry
self._prompt_server = prompt_server self._prompt_server = prompt_server
self._standalone_mode = standalone_mode 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: async def register_nodes(self, request: web.Request) -> web.Response:
try: try:
@@ -3162,7 +3177,12 @@ class NodeRegistryHandler:
) )
graph_name = node.get("graph_name") graph_name = node.get("graph_name")
try: 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): except (TypeError, ValueError):
return web.json_response( return web.json_response(
{ {
@@ -3203,39 +3223,95 @@ class NodeRegistryHandler:
status=503, 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()) 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( registry_info = await self._node_registry.get_merged_registry(
active_sids=current_sids 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: if registry_info["node_count"] == 0:
logger.warning("No nodes registered after refresh") logger.warning("No nodes registered after refresh")

View File

@@ -552,6 +552,8 @@ async function fetchWorkflowRegistry() {
if (!registryData.success) { if (!registryData.success) {
if (registryData.error === 'Standalone Mode Active') { if (registryData.error === 'Standalone Mode Active') {
showToast('toast.general.cannotInteractStandalone', {}, 'warning'); showToast('toast.general.cannotInteractStandalone', {}, 'warning');
} else if (registryData.error === 'Empty Registry') {
showToast('uiHelpers.workflow.noSupportedNodes', {}, 'warning');
} else { } else {
showToast('toast.general.failedWorkflowInfo', {}, 'error'); showToast('toast.general.failedWorkflowInfo', {}, 'error');
} }

View File

@@ -728,6 +728,54 @@ async def test_register_nodes_includes_capabilities():
assert stored_node["widget_names"] == ["ckpt_name"] 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 @pytest.mark.asyncio
async def test_update_node_widget_sends_payload(): async def test_update_node_widget_sends_payload():
send_calls: list[tuple[str, dict]] = [] send_calls: list[tuple[str, dict]] = []

View File

@@ -1,8 +1,10 @@
import { app } from "../../scripts/app.js"; import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.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"; import { ensureLmStyles } from "./lm_styles_loader.js";
const DEBOUNCE_DELAY = 500;
const LORA_NODE_CLASSES = new Set([ const LORA_NODE_CLASSES = new Set([
"Lora Loader (LoraManager)", "Lora Loader (LoraManager)",
"Lora Stacker (LoraManager)", "Lora Stacker (LoraManager)",
@@ -79,6 +81,7 @@ app.registerExtension({
setup() { setup() {
ensureLmStyles(); ensureLmStyles();
this._log("extension initialized, clientId=%s", api.clientId ?? api.initialClientId ?? "(pending)");
api.addEventListener("lora_registry_refresh", () => { api.addEventListener("lora_registry_refresh", () => {
this.refreshRegistry(); this.refreshRegistry();
@@ -88,10 +91,64 @@ app.registerExtension({
this.applyWidgetUpdate(event?.detail ?? {}); this.applyWidgetUpdate(event?.detail ?? {});
}); });
// React to marker changes from the Node Marker extension
window.addEventListener("lm_marker_changed", () => { window.addEventListener("lm_marker_changed", () => {
this.refreshRegistry(); 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() { async refreshRegistry() {
@@ -115,7 +172,6 @@ app.registerExtension({
const hasTextWidget = TEXT_CAPABLE_CLASSES.has(node.comfyClass); const hasTextWidget = TEXT_CAPABLE_CLASSES.has(node.comfyClass);
const markerRole = node.properties?.lm_marker_role ?? null; const markerRole = node.properties?.lm_marker_role ?? null;
// Skip nodes with no relevant capability UNLESS they are marked
if (!supportsLora && !hasTargetWidget && !hasTextWidget && !markerRole) { if (!supportsLora && !hasTargetWidget && !hasTextWidget && !markerRole) {
continue; 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", { const response = await fetch("/api/lm/register-nodes", {
method: "POST", method: "POST",
headers: { headers: {
@@ -153,7 +220,7 @@ app.registerExtension({
}, },
body: JSON.stringify({ body: JSON.stringify({
nodes: workflowNodes, nodes: workflowNodes,
client_id: api.clientId ?? api.initialClientId ?? "", client_id: clientId,
}), }),
}); });