feat(nodes): flag missing local models at queue and load time (#1057)

tests/frontend/components/lorasWidgetAvailability.test.js

@@ -0,0 +1,228 @@
+import { describe, it, expect, beforeEach, vi } from 'vitest';
+
+const {
+  APP_MODULE,
+  API_MODULE,
+  UTILS_MODULE,
+} = vi.hoisted(() => ({
+  APP_MODULE: new URL('../../../scripts/app.js', import.meta.url).pathname,
+  API_MODULE: new URL('../../../scripts/api.js', import.meta.url).pathname,
+  UTILS_MODULE: new URL('../../../web/comfyui/loras_widget_utils.js', import.meta.url).pathname,
+}));
+
+vi.mock(APP_MODULE, () => ({
+  app: { graph: {} },
+}));
+
+const { fetchApiMock } = vi.hoisted(() => ({ fetchApiMock: vi.fn() }));
+vi.mock(API_MODULE, () => ({
+  api: { fetchApi: fetchApiMock },
+}));
+
+import {
+  normalizeLoraNameKey,
+  buildAvailableLoraSet,
+  isLoraNameAvailable,
+  getAvailableLoras,
+  getAvailableLorasSync,
+  resetAvailableLorasCache,
+  onLibraryChanged,
+  handleLibraryChangeMessage,
+} from '../../../web/comfyui/loras_widget_utils.js';
+
+describe('normalizeLoraNameKey', () => {
+  it('normalizes backslashes to forward slashes', () => {
+    expect(normalizeLoraNameKey('sub\\folder\\lora.safetensors')).toBe(
+      'sub/folder/lora'
+    );
+  });
+
+  it('strips known model extensions case-insensitively', () => {
+    expect(normalizeLoraNameKey('lora.safetensors')).toBe('lora');
+    expect(normalizeLoraNameKey('lora.CKPT')).toBe('lora');
+    expect(normalizeLoraNameKey('lora.pt')).toBe('lora');
+    expect(normalizeLoraNameKey('lora.bin')).toBe('lora');
+  });
+
+  it('keeps extensions the backend does not strip', () => {
+    // Matches backend _strip_lora_extension: .gguf is not a LoRA extension.
+    expect(normalizeLoraNameKey('lora.gguf')).toBe('lora.gguf');
+  });
+
+  it('leaves extension-less names unchanged', () => {
+    expect(normalizeLoraNameKey('lora')).toBe('lora');
+  });
+});
+
+describe('buildAvailableLoraSet', () => {
+  it('registers both path and basename forms without extension', () => {
+    const set = buildAvailableLoraSet(['sub/a.safetensors', 'b.ckpt']);
+    expect(set.has('sub/a')).toBe(true);
+    expect(set.has('a')).toBe(true);
+    expect(set.has('b')).toBe(true);
+  });
+
+  it('ignores empty entries', () => {
+    const set = buildAvailableLoraSet([null, '', 'sub/c.safetensors']);
+    expect(set.has('sub/c')).toBe(true);
+    expect(set.has('')).toBe(false);
+  });
+});
+
+describe('isLoraNameAvailable', () => {
+  const set = buildAvailableLoraSet(['sub/a.safetensors', 'b.ckpt']);
+
+  it('treats everything as available while the set is not loaded', () => {
+    expect(isLoraNameAvailable('anything.safetensors', null)).toBe(true);
+  });
+
+  it('matches by basename with or without extension', () => {
+    expect(isLoraNameAvailable('a', set)).toBe(true);
+    expect(isLoraNameAvailable('a.safetensors', set)).toBe(true);
+    expect(isLoraNameAvailable('b.ckpt', set)).toBe(true);
+  });
+
+  it('matches by full folder path', () => {
+    expect(isLoraNameAvailable('sub/a.safetensors', set)).toBe(true);
+  });
+
+  it('reports names not in the library', () => {
+    expect(isLoraNameAvailable('missing.safetensors', set)).toBe(false);
+  });
+
+  it('falls back to the basename for folder-qualified names', () => {
+    // Mirror of the backend basename fallback: a folder prefix that does not
+    // match a stored path still resolves when the basename exists.
+    expect(isLoraNameAvailable('sub/b.ckpt', set)).toBe(true);
+    expect(isLoraNameAvailable('any/folder/a.safetensors', set)).toBe(true);
+    expect(isLoraNameAvailable('sub/missing.safetensors', set)).toBe(false);
+  });
+
+  it('treats absolute paths as available without verification', () => {
+    expect(isLoraNameAvailable('/abs/path/x.safetensors', set)).toBe(true);
+    expect(isLoraNameAvailable('C:/abs/path/x.safetensors', set)).toBe(true);
+  });
+});
+
+describe('getAvailableLoras caching', () => {
+  beforeEach(() => {
+    fetchApiMock.mockReset();
+    resetAvailableLorasCache();
+  });
+
+  it('fetches the cycler list and builds the availability set', async () => {
+    fetchApiMock.mockResolvedValue({
+      ok: true,
+      json: async () => ({
+        success: true,
+        loras: [{ file_name: 'sub/a.safetensors' }, { file_name: 'b.ckpt' }],
+      }),
+    });
+
+    const set = await getAvailableLoras();
+    expect(fetchApiMock).toHaveBeenCalledWith('/lm/loras/cycler-list', {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: '{}',
+    });
+    expect(set.has('sub/a')).toBe(true);
+    expect(set.has('b')).toBe(true);
+    expect(getAvailableLorasSync().has('sub/a')).toBe(true);
+  });
+
+  it('shares a single in-flight request between concurrent callers', async () => {
+    fetchApiMock.mockResolvedValue({
+      ok: true,
+      json: async () => ({ success: true, loras: [] }),
+    });
+
+    await Promise.all([getAvailableLoras(), getAvailableLoras()]);
+    expect(fetchApiMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('does not refetch while the cache is fresh', async () => {
+    fetchApiMock.mockResolvedValue({
+      ok: true,
+      json: async () => ({ success: true, loras: [] }),
+    });
+
+    await getAvailableLoras();
+    await getAvailableLoras();
+    expect(fetchApiMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('resolves to null and does not cache on fetch failure', async () => {
+    fetchApiMock.mockRejectedValue(new Error('network down'));
+
+    const result = await getAvailableLoras();
+    expect(result).toBeNull();
+    expect(getAvailableLorasSync()).toBeNull();
+  });
+
+  it('resolves to null on non-ok response', async () => {
+    fetchApiMock.mockResolvedValue({ ok: false });
+
+    const result = await getAvailableLoras();
+    expect(result).toBeNull();
+  });
+});
+
+describe('library change invalidation', () => {
+  beforeEach(() => {
+    fetchApiMock.mockReset();
+    resetAvailableLorasCache();
+  });
+
+  it('invalidates the availability cache and notifies listeners on models_changed', async () => {
+    fetchApiMock.mockResolvedValue({
+      ok: true,
+      json: async () => ({
+        success: true,
+        loras: [{ file_name: 'old.safetensors' }],
+      }),
+    });
+
+    const listener = vi.fn();
+    const unsubscribe = onLibraryChanged(listener);
+    try {
+      await getAvailableLoras();
+      expect(getAvailableLorasSync().has('old')).toBe(true);
+
+      // Simulate a deletion in the Lora Manager UI: the cache is dropped and
+      // listeners are notified so widgets re-render with fresh data.
+      handleLibraryChangeMessage({ type: 'models_changed' });
+      expect(listener).toHaveBeenCalledTimes(1);
+      expect(getAvailableLorasSync()).toBeNull();
+
+      fetchApiMock.mockResolvedValue({
+        ok: true,
+        json: async () => ({ success: true, loras: [] }),
+      });
+      await getAvailableLoras();
+      expect(getAvailableLorasSync().has('old')).toBe(false);
+    } finally {
+      unsubscribe();
+    }
+  });
+
+  it('ignores unrelated messages', () => {
+    const listener = vi.fn();
+    const unsubscribe = onLibraryChanged(listener);
+    try {
+      handleLibraryChangeMessage({ type: 'download_progress' });
+      handleLibraryChangeMessage({ type: 'init_progress' });
+      handleLibraryChangeMessage(null);
+      expect(listener).not.toHaveBeenCalled();
+    } finally {
+      unsubscribe();
+    }
+  });
+
+  it('unsubscribes listeners', () => {
+    const listener = vi.fn();
+    const unsubscribe = onLibraryChanged(listener);
+    unsubscribe();
+    handleLibraryChangeMessage({ type: 'models_changed' });
+    expect(listener).not.toHaveBeenCalled();
+  });
+});

tests/nodes/test_checkpoint_name_filtering.py

@@ -0,0 +1,85 @@
+"""Tests for the checkpoint/unet combo-name existence filtering.
+
+Deleted files must drop out of the combo list so ComfyUI flags the node at
+queue time via "value not in list" instead of failing at execution time.
+"""
+
+import pytest
+
+from py.nodes.checkpoint_loader import CheckpointLoaderLM
+from py.nodes.unet_loader import UNETLoaderLM
+
+
+class _FakeCache:
+    def __init__(self, raw_data):
+        self.raw_data = raw_data
+
+
+class _FakeScanner:
+    def __init__(self, raw_data, model_roots):
+        self._raw_data = raw_data
+        self._model_roots = model_roots
+
+    async def get_cached_data(self, force_refresh=False):
+        return _FakeCache(self._raw_data)
+
+    def get_model_roots(self):
+        return self._model_roots
+
+
+@pytest.fixture
+def checkpoint_library(tmp_path, monkeypatch):
+    from py.services.service_registry import ServiceRegistry
+
+    existing = tmp_path / "keep.safetensors"
+    existing.write_bytes(b"x")
+    deleted = tmp_path / "deleted.safetensors"  # referenced but never created
+
+    raw_data = [
+        {"sub_type": "checkpoint", "file_path": str(existing)},
+        {"sub_type": "checkpoint", "file_path": str(deleted)},
+        # Wrong type must stay excluded by the sub_type filter.
+        {"sub_type": "diffusion_model", "file_path": str(existing)},
+    ]
+    async def _fake_scanner():
+        return _FakeScanner(raw_data, [str(tmp_path)])
+
+    monkeypatch.setattr(
+        ServiceRegistry, "get_checkpoint_scanner", _fake_scanner
+    )
+    return tmp_path
+
+
+def test_checkpoint_names_drop_deleted_files(checkpoint_library):
+    assert CheckpointLoaderLM._get_checkpoint_names() == ["keep.safetensors"]
+
+
+def test_unet_names_drop_deleted_files(tmp_path, monkeypatch):
+    from py.services.service_registry import ServiceRegistry
+
+    existing = tmp_path / "keep.safetensors"
+    existing.write_bytes(b"x")
+    deleted = tmp_path / "deleted.safetensors"
+
+    raw_data = [
+        {"sub_type": "diffusion_model", "file_path": str(existing)},
+        {"sub_type": "diffusion_model", "file_path": str(deleted)},
+        {"sub_type": "checkpoint", "file_path": str(existing)},
+    ]
+    async def _fake_scanner():
+        return _FakeScanner(raw_data, [str(tmp_path)])
+
+    monkeypatch.setattr(
+        ServiceRegistry, "get_checkpoint_scanner", _fake_scanner
+    )
+    assert UNETLoaderLM._get_unet_names() == ["keep.safetensors"]
+
+
+def test_checkpoint_names_empty_when_scanner_fails(tmp_path, monkeypatch):
+    from py.services.service_registry import ServiceRegistry
+
+    def _boom():
+        raise RuntimeError("scanner not available")
+
+    monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _boom)
+    assert CheckpointLoaderLM._get_checkpoint_names() == []

tests/nodes/test_lora_validation.py

@@ -0,0 +1,252 @@
+"""Tests for queue-time LoRA validation helpers (validate_lora_entries)."""
+
+import pytest
+
+from py.nodes.utils import _find_missing_loras, validate_lora_entries
+
+
+class _FakeCache:
+    def __init__(self, raw_data):
+        self.raw_data = raw_data
+
+
+class _FakeScanner:
+    def __init__(self, raw_data):
+        self._raw_data = raw_data
+        # Non-None cache marks the scanner as initialized; None means the
+        # real scanner has not hydrated yet (validation must stay lenient).
+        self._cache = object()
+        self._is_initializing = False
+
+    async def get_cached_data(self, force_refresh=False):
+        return _FakeCache(self._raw_data)
+
+
+@pytest.fixture
+def lora_library(tmp_path):
+    """Create a fake on-disk LoRA library plus a matching scanner cache."""
+    existing_a = tmp_path / "a.safetensors"
+    existing_b = tmp_path / "sub" / "b.safetensors"
+    existing_a.write_bytes(b"a")
+    existing_b.parent.mkdir()
+    existing_b.write_bytes(b"b")
+
+    # "gone" is referenced by the cache but the file was deleted afterwards,
+    # simulating a stale scanner cache.
+    gone = tmp_path / "gone.safetensors"
+
+    raw_data = [
+        {"file_name": "a.safetensors", "folder": "", "file_path": str(existing_a)},
+        {
+            "file_name": "b.safetensors",
+            "folder": "sub",
+            "file_path": str(existing_b),
+        },
+        {"file_name": "gone.safetensors", "folder": "", "file_path": str(gone)},
+    ]
+    return raw_data
+
+
+@pytest.fixture
+def mock_lora_scanner(lora_library, monkeypatch):
+    from py.services.service_registry import ServiceRegistry
+
+    async def _fake_scanner():
+        return _FakeScanner(lora_library)
+
+    monkeypatch.setattr(
+        ServiceRegistry, "get_lora_scanner", _fake_scanner
+    )
+
+
+def _entries(*names):
+    return [
+        {"active": True, "name": name, "strength": 1.0, "clipStrength": 1.0}
+        for name in names
+    ]
+
+
+def test_validate_flat_name_with_extension(mock_lora_scanner):
+    assert validate_lora_entries({"loras": _entries("a.safetensors")}) is None
+
+
+def test_validate_flat_name_without_extension(mock_lora_scanner):
+    assert validate_lora_entries({"loras": _entries("a")}) is None
+
+
+def test_validate_subfolder_name(mock_lora_scanner):
+    assert validate_lora_entries({"loras": _entries("sub/b.safetensors")}) is None
+    assert validate_lora_entries({"loras": _entries("sub/b")}) is None
+
+
+def test_validate_stale_cache_entry_reported(mock_lora_scanner):
+    result = validate_lora_entries({"loras": _entries("gone.safetensors")})
+    assert result is not None
+    assert "gone.safetensors" in result
+
+
+def test_validate_unknown_name_reported(mock_lora_scanner):
+    result = validate_lora_entries({"loras": _entries("nope.safetensors")})
+    assert result is not None
+    assert "nope.safetensors" in result
+
+
+def test_validate_multiple_missing_all_listed(mock_lora_scanner):
+    result = validate_lora_entries(
+        {"loras": _entries("gone.safetensors", "nope.safetensors")}
+    )
+    assert result is not None
+    assert "gone.safetensors" in result
+    assert "nope.safetensors" in result
+
+
+def test_validate_inactive_entries_ignored(mock_lora_scanner):
+    kwargs = {"loras": [{"active": False, "name": "gone.safetensors", "strength": 1.0}]}
+    assert validate_lora_entries(kwargs) is None
+
+
+def test_validate_value_wrapper_format(mock_lora_scanner):
+    kwargs = {"loras": {"__value__": _entries("a.safetensors")}}
+    assert validate_lora_entries(kwargs) is None
+    kwargs_missing = {"loras": {"__value__": _entries("gone.safetensors")}}
+    assert validate_lora_entries(kwargs_missing) is not None
+
+
+def test_validate_legacy_basename_fallback(mock_lora_scanner):
+    # A name with a folder that only matches by basename resolves at runtime
+    # via get_lora_info_absolute's fallback, so it must not be flagged.
+    assert validate_lora_entries({"loras": _entries("other/b.safetensors")}) is None
+
+
+def test_validate_existing_absolute_path_ok(mock_lora_scanner, lora_library):
+    existing_a = lora_library[0]["file_path"]
+    assert validate_lora_entries({"loras": _entries(existing_a)}) is None
+
+
+def test_validate_missing_absolute_path_reported(mock_lora_scanner):
+    result = validate_lora_entries(
+        {"loras": _entries("/nonexistent/path/x.safetensors")}
+    )
+    # Legacy syntax format normalizes to the basename, matching execution.
+    assert result is not None
+    assert "x.safetensors" in result
+
+
+def test_validate_empty_or_missing_loras_is_valid(mock_lora_scanner):
+    assert validate_lora_entries({}) is None
+    assert validate_lora_entries({"loras": []}) is None
+
+
+def test_validate_scanner_failure_is_lenient(monkeypatch):
+    from py.services.service_registry import ServiceRegistry
+
+    def _boom():
+        raise RuntimeError("scanner not initialized")
+
+    monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _boom)
+    assert validate_lora_entries({"loras": _entries("a.safetensors")}) is None
+
+
+def test_find_missing_loras_empty(mock_lora_scanner):
+    assert _find_missing_loras([]) == []
+
+
+def test_validate_inputs_only_reports_loras_input(mock_lora_scanner):
+    """VALIDATE_INPUTS must fail for the single 'loras' input only — ComfyUI
+    creates one custom_validation_failed error per declared input, so an
+    explicit loras parameter keeps the report to a single input error."""
+    from py.nodes.lora_loader import LoraLoaderLM
+    from py.nodes.lora_stacker import LoraStackerLM
+    from py.nodes.create_hook_lora import CreateHookLoraLM
+    from py.nodes.wanvideo_lora_select import WanVideoLoraSelectLM
+    from py.nodes.lora_randomizer import LoraRandomizerLM
+
+    import inspect
+
+    for cls in (
+        LoraLoaderLM,
+        LoraStackerLM,
+        CreateHookLoraLM,
+        WanVideoLoraSelectLM,
+        LoraRandomizerLM,
+    ):
+        spec = inspect.getfullargspec(getattr(cls, "VALIDATE_INPUTS"))
+        assert spec.varkw is None, f"{cls.__name__} VALIDATE_INPUTS must not accept **kwargs"
+        assert "loras" in spec.args, f"{cls.__name__} VALIDATE_INPUTS must declare loras"
+
+    result = LoraLoaderLM.VALIDATE_INPUTS(
+        loras=[
+            {"active": True, "name": "gone.safetensors", "strength": 1.0},
+        ]
+    )
+    assert result is not None
+    assert "gone.safetensors" in result
+
+    assert LoraLoaderLM.VALIDATE_INPUTS(loras=None) is True
+    assert LoraLoaderLM.VALIDATE_INPUTS() is True
+    assert LoraLoaderLM.VALIDATE_INPUTS(loras=[]) is True
+
+
+def test_validate_lenient_when_scanner_not_initialized(monkeypatch):
+    """An unhydrated scanner (empty cache) must not reject every LoRA."""
+    from py.services.service_registry import ServiceRegistry
+
+    class _UninitializedScanner(_FakeScanner):
+        def __init__(self, raw_data):
+            super().__init__(raw_data)
+            self._cache = None
+
+    async def _fake_scanner():
+        return _UninitializedScanner([])
+
+    monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _fake_scanner)
+    assert validate_lora_entries({"loras": _entries("a.safetensors")}) is None
+
+
+def test_validate_lenient_while_scanner_initializing(monkeypatch):
+    from py.services.service_registry import ServiceRegistry
+
+    class _InitializingScanner(_FakeScanner):
+        def __init__(self, raw_data):
+            super().__init__(raw_data)
+            self._is_initializing = True
+
+    async def _fake_scanner():
+        return _InitializingScanner([])
+
+    monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _fake_scanner)
+    assert validate_lora_entries({"loras": _entries("a.safetensors")}) is None
+
+
+def test_find_missing_basename_fallback_prefers_folder_prefix(
+    tmp_path, monkeypatch
+):
+    """Folder-prefix candidates win over the first basename match, mirroring
+    get_lora_info_absolute's fallback ordering."""
+    from py.services.service_registry import ServiceRegistry
+
+    root_a = tmp_path / "root_a"
+    root_b = tmp_path / "root_b"
+    root_a.mkdir()
+    root_b.mkdir()
+    live_a = root_a / "x.safetensors"
+    live_a.write_bytes(b"x")
+    stale_b = root_b / "x.safetensors"  # same basename, file deleted
+
+    raw_data = [
+        {"file_name": "x.safetensors", "folder": "root_a", "file_path": str(live_a)},
+        {"file_name": "x.safetensors", "folder": "root_b", "file_path": str(stale_b)},
+    ]
+
+    async def _fake_scanner():
+        return _FakeScanner(raw_data)
+
+    monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _fake_scanner)
+    # Exact path matches are authoritative.
+    assert _find_missing_loras(["root_b/x"]) == ["root_b/x"]
+    # Folder-prefix fallback beats the first basename match: root_a is live.
+    assert _find_missing_loras(["root_a/deep/x"]) == []
+    # Prefix match wins over the first candidate: root_b is stale.
+    assert _find_missing_loras(["root_b/deep/x"]) == ["root_b/deep/x"]
+    # No prefix match: falls back to the first basename candidate (root_a).
+    assert _find_missing_loras(["other/x"]) == []

tests/routes/test_models_changed_broadcast.py

@@ -0,0 +1,40 @@
+"""Tests for the models_changed WebSocket broadcast helper."""
+
+import asyncio
+
+import pytest
+
+from py.routes.handlers import model_handlers
+
+
+@pytest.mark.asyncio
+async def test_broadcast_models_changed_payload(monkeypatch):
+    """The helper must broadcast the models_changed event to all clients."""
+    sent = []
+
+    async def fake_broadcast(data):
+        sent.append(data)
+
+    monkeypatch.setattr(
+        "py.services.websocket_manager.ws_manager.broadcast", fake_broadcast
+    )
+    # Give the create_task a chance to run.
+    for _ in range(10):
+        await asyncio.sleep(0)
+
+    model_handlers._broadcast_models_changed()
+    await asyncio.sleep(0)
+    assert sent == [{"type": "models_changed"}]
+
+
+@pytest.mark.asyncio
+async def test_broadcast_models_changed_survives_broadcast_failure(monkeypatch):
+    """A failing broadcast must not raise out of the mutation handler."""
+    async def fake_broadcast(data):
+        raise RuntimeError("socket closed")
+
+    monkeypatch.setattr(
+        "py.services.websocket_manager.ws_manager.broadcast", fake_broadcast
+    )
+    model_handlers._broadcast_models_changed()
+    await asyncio.sleep(0)