feat(nodes): make Lora Stack Combiner inputs dynamic

tests/frontend/core/loraStackDynamicInputs.test.js

@@ -0,0 +1,195 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const { APP_MODULE, EXTENSION_MODULE, appMock, registeredExtensions } =
+  vi.hoisted(() => {
+    const registeredExtensions = [];
+    const appMock = {
+      configuringGraph: false,
+      registerExtension: (ext) => registeredExtensions.push(ext),
+    };
+    return {
+      APP_MODULE: new URL("../../../scripts/app.js", import.meta.url).pathname,
+      EXTENSION_MODULE: new URL(
+        "../../../web/comfyui/lora_stack_dynamic_inputs.js",
+        import.meta.url
+      ).pathname,
+      appMock,
+      registeredExtensions,
+    };
+  });
+
+vi.mock(APP_MODULE, () => ({
+  app: appMock,
+}));
+
+describe("Lora Stack Combiner dynamic inputs", () => {
+  let extension;
+
+  beforeEach(async () => {
+    vi.resetModules();
+    registeredExtensions.length = 0;
+    appMock.configuringGraph = false;
+    await import(EXTENSION_MODULE);
+    extension = registeredExtensions.find(
+      (ext) => ext.name === "Comfy.LoraManager.LoraStackCombiner"
+    );
+    expect(extension).toBeDefined();
+  });
+
+  function createNodeType() {
+    const nodeType = { prototype: {} };
+    extension.beforeRegisterNodeDef(
+      nodeType,
+      { name: "Lora Stack Combiner (LoraManager)" },
+      appMock
+    );
+    return nodeType;
+  }
+
+  function createNode(inputs = []) {
+    const node = {
+      comfyClass: "Lora Stack Combiner (LoraManager)",
+      inputs: inputs.map((name) => ({ name, type: "LORA_STACK" })),
+      addInput: vi.fn(function (name, type, opts) {
+        this.inputs.push({ name, type, ...opts });
+      }),
+      removeInput: vi.fn(function (index) {
+        this.inputs.splice(index, 1);
+      }),
+    };
+    return node;
+  }
+
+  function makeLinkInfo() {
+    return { id: 999, origin_id: 1, target_id: 2 };
+  }
+
+  it("adds a third input when the last slot gets connected", () => {
+    const nodeType = createNodeType();
+    const node = createNode(["lora_stack1", "lora_stack2"]);
+    node.onConnectionsChange = nodeType.prototype.onConnectionsChange;
+
+    node.onConnectionsChange(1, 1, true, makeLinkInfo());
+
+    expect(node.inputs.map((input) => input.name)).toEqual([
+      "lora_stack1",
+      "lora_stack2",
+      "lora_stack3",
+    ]);
+  });
+
+  it("does not add an input when a non-last slot gets connected", () => {
+    const nodeType = createNodeType();
+    const node = createNode(["lora_stack1", "lora_stack2", "lora_stack3"]);
+    node.onConnectionsChange = nodeType.prototype.onConnectionsChange;
+
+    node.onConnectionsChange(1, 0, true, makeLinkInfo());
+
+    expect(node.inputs.map((input) => input.name)).toEqual([
+      "lora_stack1",
+      "lora_stack2",
+      "lora_stack3",
+    ]);
+  });
+
+  it("removes a disconnected middle slot and renumbers", () => {
+    // Simulates a real LiteGraph disconnect event: it fires only for slots that
+    // had a link, and input.link has already been cleared before the event fires.
+    const nodeType = createNodeType();
+    const node = createNode(["lora_stack1", "lora_stack2", "lora_stack3"]);
+    node.inputs[0].link = 11;
+    node.inputs[1].link = null; // slot 2 was just disconnected
+    node.inputs[2].link = 13;
+    node.onConnectionsChange = nodeType.prototype.onConnectionsChange;
+
+    node.onConnectionsChange(1, 1, false, makeLinkInfo());
+
+    expect(node.inputs.map((input) => input.name)).toEqual([
+      "lora_stack1",
+      "lora_stack2",
+    ]);
+  });
+
+  it("keeps the last slot when it is disconnected", () => {
+    const nodeType = createNodeType();
+    const node = createNode(["lora_stack1", "lora_stack2", "lora_stack3"]);
+    node.inputs[0].link = 11;
+    node.inputs[1].link = 12;
+    node.inputs[2].link = null; // last slot was just disconnected
+    node.onConnectionsChange = nodeType.prototype.onConnectionsChange;
+
+    node.onConnectionsChange(1, 2, false, makeLinkInfo());
+
+    expect(node.inputs.map((input) => input.name)).toEqual([
+      "lora_stack1",
+      "lora_stack2",
+      "lora_stack3",
+    ]);
+    expect(node.removeInput).not.toHaveBeenCalled();
+  });
+
+  it("keeps at least two inputs when disconnecting", () => {
+    const nodeType = createNodeType();
+    const node = createNode(["lora_stack1", "lora_stack2"]);
+    node.inputs[0].link = 11;
+    node.inputs[1].link = null; // slot 2 was just disconnected
+    node.onConnectionsChange = nodeType.prototype.onConnectionsChange;
+
+    node.onConnectionsChange(1, 1, false, makeLinkInfo());
+
+    expect(node.inputs.map((input) => input.name)).toEqual([
+      "lora_stack1",
+      "lora_stack2",
+    ]);
+    expect(node.removeInput).not.toHaveBeenCalled();
+  });
+
+  it("does nothing while the graph is being configured", () => {
+    appMock.configuringGraph = true;
+    const nodeType = createNodeType();
+    const node = createNode(["lora_stack1", "lora_stack2"]);
+    node.onConnectionsChange = nodeType.prototype.onConnectionsChange;
+
+    node.onConnectionsChange(1, 1, true, makeLinkInfo());
+
+    expect(node.inputs.map((input) => input.name)).toEqual([
+      "lora_stack1",
+      "lora_stack2",
+    ]);
+    expect(node.addInput).not.toHaveBeenCalled();
+  });
+
+  it("leaves legacy lora_stack_a/b inputs untouched", () => {
+    const nodeType = createNodeType();
+    const node = createNode(["lora_stack_a", "lora_stack_b"]);
+    node.onConnectionsChange = nodeType.prototype.onConnectionsChange;
+
+    node.onConnectionsChange(1, 0, true, makeLinkInfo());
+
+    expect(node.inputs.map((input) => input.name)).toEqual([
+      "lora_stack_a",
+      "lora_stack_b",
+    ]);
+    expect(node.addInput).not.toHaveBeenCalled();
+  });
+
+  it("ensures two numbered inputs exist on creation", () => {
+    const node = createNode([]);
+    extension.nodeCreated(node, {});
+
+    expect(node.inputs.map((input) => input.name)).toEqual([
+      "lora_stack1",
+      "lora_stack2",
+    ]);
+  });
+
+  it("does not add numbered inputs to legacy workflows", () => {
+    const node = createNode(["lora_stack_a", "lora_stack_b"]);
+    extension.nodeCreated(node, {});
+
+    expect(node.inputs.map((input) => input.name)).toEqual([
+      "lora_stack_a",
+      "lora_stack_b",
+    ]);
+  });
+});

tests/nodes/test_lora_stack_combiner.py

@@ -1,4 +1,11 @@
-from py.nodes.lora_stack_combiner import LoraStackCombinerLM
+import types
+
+import pytest
+
+from py.nodes.lora_stack_combiner import (
+    LoraStackCombinerLM,
+    _LoraStackOptionalInputs,
+)
 
 
 def test_combine_stacks_preserves_order():
@@ -63,8 +70,95 @@ def test_combine_stacks_returns_other_when_one_unconnected():
     node = LoraStackCombinerLM()
     stack_a = [("folder/a.safetensors", 0.7, 0.6)]
 
-    (combined_stack_a,) = node.combine_stacks(lora_stack_a=stack_a)
-    (combined_stack_b,) = node.combine_stacks(lora_stack_b=stack_a)
+    (combined_stack_a,) = node.combine_stacks(lora_stack1=stack_a)
+    (combined_stack_b,) = node.combine_stacks(lora_stack2=stack_a)
 
     assert combined_stack_a == stack_a
     assert combined_stack_b == stack_a
+
+
+def test_combine_stacks_with_dynamic_third_slot():
+    node = LoraStackCombinerLM()
+    stack_a = [("folder/a.safetensors", 0.7, 0.6)]
+    stack_b = [("folder/b.safetensors", 0.8, 0.8)]
+    stack_c = [("folder/c.safetensors", 1.0, 0.9)]
+
+    (combined_stack,) = node.combine_stacks(
+        lora_stack1=stack_a, lora_stack2=stack_b, lora_stack3=stack_c
+    )
+
+    assert combined_stack == stack_a + stack_b + stack_c
+
+
+def test_combine_stacks_orders_by_slot_number_not_call_order():
+    node = LoraStackCombinerLM()
+    stack_a = [("folder/a.safetensors", 0.7, 0.6)]
+    stack_b = [("folder/b.safetensors", 0.8, 0.8)]
+    stack_c = [("folder/c.safetensors", 1.0, 0.9)]
+
+    (combined_stack,) = node.combine_stacks(
+        lora_stack3=stack_c, lora_stack2=stack_b, lora_stack1=stack_a
+    )
+
+    assert combined_stack == stack_a + stack_b + stack_c
+
+
+def test_combine_stacks_accepts_only_dynamic_slot():
+    node = LoraStackCombinerLM()
+    stack_c = [("folder/c.safetensors", 1.0, 0.9)]
+
+    (combined_stack,) = node.combine_stacks(lora_stack3=stack_c)
+
+    assert combined_stack == stack_c
+
+
+def test_combine_stacks_handles_legacy_input_names():
+    node = LoraStackCombinerLM()
+    stack_a = [("folder/a.safetensors", 0.7, 0.6)]
+    stack_b = [("folder/b.safetensors", 0.8, 0.8)]
+
+    (combined_stack,) = node.combine_stacks(lora_stack_a=stack_a, lora_stack_b=stack_b)
+
+    assert combined_stack == stack_a + stack_b
+
+
+def test_input_types_exposes_two_default_slots():
+    input_types = LoraStackCombinerLM.INPUT_TYPES()
+
+    assert set(input_types["optional"]) == {"lora_stack1", "lora_stack2"}
+    assert input_types["optional"]["lora_stack1"][0] == "LORA_STACK"
+    assert input_types["optional"]["lora_stack2"][0] == "LORA_STACK"
+
+
+def test_input_types_recognizes_dynamic_slots_from_get_input_info(monkeypatch):
+    frames = [None, None, types.SimpleNamespace(function="get_input_info")]
+    monkeypatch.setattr(
+        "py.nodes.lora_stack_combiner.inspect.stack", lambda: frames
+    )
+
+    input_types = LoraStackCombinerLM.INPUT_TYPES()
+    optional = input_types["optional"]
+
+    assert "lora_stack3" in optional
+    assert optional["lora_stack3"][0] == "LORA_STACK"
+    assert "lora_stack25" in optional
+    assert optional["lora_stack25"][0] == "LORA_STACK"
+
+
+def test_lora_stack_optional_inputs_proxy():
+    proxy = _LoraStackOptionalInputs({"lora_stack1": ("LORA_STACK", {})})
+
+    assert "lora_stack1" in proxy
+    assert "lora_stack2" in proxy
+    assert "lora_stack10" in proxy
+    assert "lora_stack_a" in proxy
+    assert "lora_stack" not in proxy
+    assert "lora_stacka" not in proxy
+    assert "lora_stack_1" not in proxy
+    assert "text" not in proxy
+
+    assert proxy["lora_stack1"][0] == "LORA_STACK"
+    assert proxy["lora_stack5"][0] == "LORA_STACK"
+
+    with pytest.raises(KeyError):
+        proxy["not_a_stack"]