feat(versions): add file-variant badge and hide download button for in-library versions (#1058)

tests/frontend/components/modelVersionsTab.downloadButton.test.js

@@ -0,0 +1,285 @@
+import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
+
+const {
+  MODEL_VERSIONS_MODULE,
+  API_FACTORY_MODULE,
+  DOWNLOAD_MANAGER_MODULE,
+  UI_HELPERS_MODULE,
+  STATE_MODULE,
+  I18N_HELPERS_MODULE,
+  UTILS_MODULE,
+} = vi.hoisted(() => ({
+  MODEL_VERSIONS_MODULE: new URL('../../../static/js/components/shared/ModelVersionsTab.js', import.meta.url).pathname,
+  API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
+  DOWNLOAD_MANAGER_MODULE: new URL('../../../static/js/managers/DownloadManager.js', import.meta.url).pathname,
+  UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
+  STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
+  I18N_HELPERS_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
+  UTILS_MODULE: new URL('../../../static/js/components/shared/utils.js', import.meta.url).pathname,
+}));
+
+const downloadVersionWithDefaults = vi.fn();
+const openFileSelectionForVersion = vi.fn();
+
+vi.mock(DOWNLOAD_MANAGER_MODULE, () => ({
+  downloadManager: {
+    downloadVersionWithDefaults,
+    openFileSelectionForVersion,
+  },
+}));
+
+vi.mock(UI_HELPERS_MODULE, () => ({
+  showToast: vi.fn(),
+  openCivitaiUrl: vi.fn(),
+}));
+
+const stateMock = {
+  global: {
+    settings: {
+      autoplay_on_hover: false,
+      version_grouping: 'any',
+    },
+  },
+};
+vi.mock(STATE_MODULE, () => ({
+  state: stateMock,
+}));
+
+vi.mock(I18N_HELPERS_MODULE, () => ({
+  translate: vi.fn((_, __, fallback) => fallback ?? ''),
+}));
+
+vi.mock(UTILS_MODULE, () => ({
+  formatFileSize: vi.fn(() => '1 MB'),
+}));
+
+vi.mock(API_FACTORY_MODULE, () => ({
+  getModelApiClient: vi.fn(),
+}));
+
+function buildRecord(versions) {
+  return {
+    success: true,
+    record: {
+      shouldIgnore: false,
+      inLibraryVersionIds: versions.filter(v => v.isInLibrary).map(v => v.versionId),
+      versions,
+    },
+  };
+}
+
+async function renderVersions(record) {
+  const { initVersionsTab } = await import(MODEL_VERSIONS_MODULE);
+  const controller = initVersionsTab({
+    modalId: 'model-versions-modal',
+    modelType: 'loras',
+    modelId: 123,
+    currentVersionId: null,
+  });
+  await controller.load();
+}
+
+function downloadButtonFor(versionId) {
+  return document.querySelector(
+    `.model-version-row[data-version-id="${versionId}"] [data-version-action="download"]`
+  );
+}
+
+function filesBadgeFor(versionId) {
+  return document.querySelector(
+    `.model-version-row[data-version-id="${versionId}"] [data-version-files]`
+  );
+}
+
+describe('ModelVersionsTab download button visibility', () => {
+  let getModelApiClient;
+  let fetchModelUpdateVersions;
+
+  beforeEach(async () => {
+    vi.resetModules();
+    downloadVersionWithDefaults.mockReset();
+    downloadVersionWithDefaults.mockResolvedValue(true);
+    openFileSelectionForVersion.mockReset();
+    openFileSelectionForVersion.mockResolvedValue(undefined);
+    document.body.innerHTML = `
+      <div id="model-versions-modal">
+        <div id="versions-tab">
+          <div class="model-versions-tab"></div>
+        </div>
+      </div>
+    `;
+    ({ getModelApiClient } = await import(API_FACTORY_MODULE));
+    fetchModelUpdateVersions = vi.fn();
+    getModelApiClient.mockReturnValue({
+      fetchModelUpdateVersions,
+      fetchModelRoots: vi.fn(),
+      setModelUpdateIgnore: vi.fn(),
+      setVersionUpdateIgnore: vi.fn(),
+      deleteModel: vi.fn(),
+    });
+  });
+
+  afterEach(() => {
+    document.body.innerHTML = '';
+  });
+
+  it('hides the download button for a single-file in-library version', async () => {
+    fetchModelUpdateVersions.mockResolvedValue(buildRecord([
+      {
+        versionId: 10,
+        name: 'v1.0',
+        baseModel: 'Illustrious',
+        isInLibrary: true,
+        shouldIgnore: false,
+        filePath: '/models/loras/file.safetensors',
+        fileCount: 1,
+      },
+    ]));
+
+    await renderVersions();
+
+    expect(downloadButtonFor(10)).toBeFalsy();
+    expect(filesBadgeFor(10)).toBeFalsy();
+    // The delete affordance must remain for in-library versions.
+    expect(document.querySelector(
+      '.model-version-row[data-version-id="10"] [data-version-action="delete"]'
+    )).toBeTruthy();
+  });
+
+  it('shows the files badge instead of a download button for a multi-file in-library version', async () => {
+    fetchModelUpdateVersions.mockResolvedValue(buildRecord([
+      {
+        versionId: 10,
+        name: 'v1.0',
+        baseModel: 'Illustrious',
+        isInLibrary: true,
+        shouldIgnore: false,
+        filePath: '/models/loras/file.safetensors',
+        fileCount: 3,
+      },
+    ]));
+
+    await renderVersions();
+
+    expect(downloadButtonFor(10)).toBeFalsy();
+    const badge = filesBadgeFor(10);
+    expect(badge).toBeTruthy();
+    expect(badge.textContent).toContain('3 files');
+    expect(badge.getAttribute('title')).toBe('Choose which files to download');
+
+    badge.click();
+    await new Promise(resolve => setTimeout(resolve, 0));
+
+    expect(openFileSelectionForVersion).toHaveBeenCalledWith('loras', 123, 10);
+  });
+
+  it('hides the download button when fileCount is unknown for an in-library version', async () => {
+    fetchModelUpdateVersions.mockResolvedValue(buildRecord([
+      {
+        versionId: 10,
+        name: 'v1.0',
+        baseModel: 'Illustrious',
+        isInLibrary: true,
+        shouldIgnore: false,
+        filePath: '/models/loras/file.safetensors',
+      },
+    ]));
+
+    await renderVersions();
+
+    expect(downloadButtonFor(10)).toBeFalsy();
+    expect(filesBadgeFor(10)).toBeFalsy();
+  });
+
+  it('shows the download button for versions not in the library', async () => {
+    fetchModelUpdateVersions.mockResolvedValue(buildRecord([
+      {
+        versionId: 11,
+        name: 'v1.1',
+        baseModel: 'Illustrious',
+        isInLibrary: false,
+        shouldIgnore: false,
+        fileCount: 1,
+      },
+      {
+        versionId: 12,
+        name: 'v1.2',
+        baseModel: 'Illustrious',
+        isInLibrary: false,
+        shouldIgnore: false,
+      },
+    ]));
+
+    await renderVersions();
+
+    expect(downloadButtonFor(11)).toBeTruthy();
+    expect(downloadButtonFor(12)).toBeTruthy();
+    // Single-file and unknown-count versions get no files badge.
+    expect(filesBadgeFor(11)).toBeFalsy();
+    expect(filesBadgeFor(12)).toBeFalsy();
+  });
+
+  it('keeps the default-file download button and offers the files badge for a multi-file version not in the library', async () => {
+    fetchModelUpdateVersions.mockResolvedValue(buildRecord([
+      {
+        versionId: 11,
+        name: 'v1.1',
+        baseModel: 'Illustrious',
+        isInLibrary: false,
+        shouldIgnore: false,
+        fileCount: 2,
+      },
+    ]));
+
+    await renderVersions();
+
+    // The Download button stays bound to the default (primary) file.
+    const button = downloadButtonFor(11);
+    expect(button).toBeTruthy();
+    expect(button.getAttribute('title')).toBe('Download this version');
+
+    button.click();
+    await new Promise(resolve => setTimeout(resolve, 0));
+
+    expect(downloadVersionWithDefaults).toHaveBeenCalledWith(
+      'loras', 123, 11,
+      expect.objectContaining({ versionName: 'v1.1' })
+    );
+    expect(openFileSelectionForVersion).not.toHaveBeenCalled();
+
+    // The badge is the advanced entry into the file-selection step.
+    const badge = filesBadgeFor(11);
+    expect(badge).toBeTruthy();
+    expect(badge.textContent).toContain('2 files');
+
+    badge.click();
+    await new Promise(resolve => setTimeout(resolve, 0));
+
+    expect(openFileSelectionForVersion).toHaveBeenCalledWith('loras', 123, 11);
+  });
+
+  it('keeps the direct default download for a single-file version not in the library', async () => {
+    fetchModelUpdateVersions.mockResolvedValue(buildRecord([
+      {
+        versionId: 11,
+        name: 'v1.1',
+        baseModel: 'Illustrious',
+        isInLibrary: false,
+        shouldIgnore: false,
+        fileCount: 1,
+      },
+    ]));
+
+    await renderVersions();
+
+    expect(filesBadgeFor(11)).toBeFalsy();
+    downloadButtonFor(11).click();
+    await new Promise(resolve => setTimeout(resolve, 0));
+
+    expect(downloadVersionWithDefaults).toHaveBeenCalledWith(
+      'loras', 123, 11,
+      expect.objectContaining({ versionName: 'v1.1' })
+    );
+    expect(openFileSelectionForVersion).not.toHaveBeenCalled();
+  });
+});

tests/routes/test_model_update_handler.py

@@ -693,3 +693,21 @@ async def test_enrich_early_access_details_skips_permanent_paid(monkeypatch):
     enriched_map = {v.version_id: v for v in enriched.versions}
     assert enriched_map[200].early_access_ends_at == "2099-01-01T00:00:00.000Z"
     assert enriched_map[100].early_access_ends_at is None
+
+
+def test_serialize_version_includes_file_count():
+    version = ModelVersionRecord(
+        version_id=11, name="v11", base_model=None, released_at=None, size_bytes=None,
+        preview_url=None, is_in_library=True, should_ignore=False, file_count=2,
+    )
+    serialized = ModelUpdateHandler._serialize_version(version, None)
+    assert serialized["fileCount"] == 2
+
+
+def test_serialize_version_file_count_defaults_to_none():
+    version = ModelVersionRecord(
+        version_id=12, name="v12", base_model=None, released_at=None, size_bytes=None,
+        preview_url=None, is_in_library=False, should_ignore=False,
+    )
+    serialized = ModelUpdateHandler._serialize_version(version, None)
+    assert serialized["fileCount"] is None

tests/services/test_model_update_service.py

@@ -798,3 +798,126 @@ def test_build_record_from_remote_preserves_paid_fields(tmp_path):
     rebuilt = record.versions[0]
     assert rebuilt.paid_access == '{"permanent": true, "endsAt": null}'
     assert rebuilt.is_paid is True
+
+
+def test_extract_file_count_counts_weight_files(tmp_path):
+    """file_count counts only weight-type files; a missing files array stays
+    None (unknown) so the UI can distinguish it from "no weight files"."""
+    db_path = tmp_path / "updates.sqlite"
+    service = ModelUpdateService(str(db_path))
+
+    response = {
+        "modelVersions": [
+            {
+                "id": 42,
+                "files": [
+                    {"sizeKB": 100, "type": "Model", "primary": True},
+                    {"sizeKB": 10, "type": "Training Data"},
+                    {"sizeKB": 50, "type": "Pruned Model"},
+                ],
+                "images": [],
+            },
+            {"id": 43, "images": []},
+            {"id": 44, "files": [], "images": []},
+        ]
+    }
+
+    versions = service._extract_versions(response)
+    assert versions is not None
+    assert versions[0].file_count == 2
+    assert versions[1].file_count is None
+    assert versions[2].file_count == 0
+
+
+@pytest.mark.asyncio
+async def test_refresh_persists_file_count(tmp_path):
+    db_path = tmp_path / "updates.sqlite"
+    service = ModelUpdateService(str(db_path), ttl_seconds=3600)
+    raw_data = [{"civitai": {"modelId": 1, "id": 11}}]
+    scanner = DummyScanner(raw_data)
+    provider = DummyProvider(
+        {
+            "modelVersions": [
+                {
+                    "id": 11,
+                    "name": "v1",
+                    "baseModel": "SD15",
+                    "files": [
+                        {"sizeKB": 1024, "type": "Model", "primary": True},
+                        {"sizeKB": 2048, "type": "Model"},
+                        {"sizeKB": 128, "type": "Training Data"},
+                    ],
+                    "images": [],
+                }
+            ]
+        }
+    )
+
+    await service.refresh_for_model_type("lora", scanner, provider)
+    record = await service.get_record("lora", 1)
+
+    assert record is not None
+    assert record.versions[0].file_count == 2
+
+
+def test_build_record_from_remote_preserves_file_count(tmp_path):
+    """A remote payload without files data must not clobber the previously
+    persisted file_count; a populated payload wins."""
+    db_path = tmp_path / "updates.sqlite"
+    service = ModelUpdateService(str(db_path))
+
+    existing = make_record(
+        ModelVersionRecord(
+            version_id=7,
+            name="v7",
+            base_model=None,
+            released_at=None,
+            size_bytes=None,
+            preview_url=None,
+            is_in_library=True,
+            should_ignore=False,
+            file_count=3,
+        )
+    )
+    remote_without_count = ModelVersionRecord(
+        version_id=7,
+        name="v7",
+        base_model=None,
+        released_at=None,
+        size_bytes=None,
+        preview_url=None,
+        is_in_library=False,
+        should_ignore=False,
+        file_count=None,
+    )
+
+    record = service._build_record_from_remote(
+        model_type="lora",
+        model_id=999,
+        local_versions=[7],
+        remote_versions=[remote_without_count],
+        existing=existing,
+        timestamp=1.0,
+    )
+    assert record.versions[0].file_count == 3
+
+    remote_with_count = ModelVersionRecord(
+        version_id=7,
+        name="v7",
+        base_model=None,
+        released_at=None,
+        size_bytes=None,
+        preview_url=None,
+        is_in_library=False,
+        should_ignore=False,
+        file_count=1,
+    )
+    record = service._build_record_from_remote(
+        model_type="lora",
+        model_id=999,
+        local_versions=[7],
+        remote_versions=[remote_with_count],
+        existing=existing,
+        timestamp=2.0,
+    )
+    assert record.versions[0].file_count == 1