feat(ui): show live scan progress and ETA for cache refresh

Broadcast typed scan_progress messages over /ws/fetch-progress from the manual refresh/rebuild paths of ModelScanner and RecipeScanner, and render percent, processed/total, current file name and an EMA-smoothed ETA in the loading overlay. Hardcoded refresh strings move to i18n (common.scanProgress); WS connection failure falls back to the previous static loading behavior.
2026-09-21 03:01:27 -03:00 · 2026-09-03 11:38:27 +08:00
parent da71985c3e
commit 14da8a6f17
19 changed files with 1668 additions and 32 deletions
@@ -0,0 +1,348 @@
+import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
+
+const {
+  BASE_MODEL_API_MODULE,
+  STATE_MODULE,
+  UI_HELPERS_MODULE,
+  I18N_MODULE,
+  STORAGE_MODULE,
+  API_CONFIG_MODULE,
+  API_FACTORY_MODULE,
+  SIDEBAR_MANAGER_MODULE,
+} = vi.hoisted(() => ({
+  BASE_MODEL_API_MODULE: new URL('../../../static/js/api/baseModelApi.js', import.meta.url).pathname,
+  STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
+  UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
+  I18N_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
+  STORAGE_MODULE: new URL('../../../static/js/utils/storageHelpers.js', import.meta.url).pathname,
+  API_CONFIG_MODULE: new URL('../../../static/js/api/apiConfig.js', import.meta.url).pathname,
+  API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
+  SIDEBAR_MANAGER_MODULE: new URL('../../../static/js/components/SidebarManager.js', import.meta.url).pathname,
+}));
+
+const showToastMock = vi.fn();
+const showMock = vi.fn();
+const showCancelButtonMock = vi.fn();
+const hideMock = vi.fn();
+const restoreProgressBarMock = vi.fn();
+const setProgressMock = vi.fn();
+const setStatusMock = vi.fn();
+const resetAndReloadMock = vi.fn();
+
+vi.mock(STATE_MODULE, () => ({
+  state: {
+    loadingManager: {
+      show: showMock,
+      showCancelButton: showCancelButtonMock,
+      hide: hideMock,
+      restoreProgressBar: restoreProgressBarMock,
+      setProgress: setProgressMock,
+      setStatus: setStatusMock,
+    },
+  },
+  getCurrentPageState: vi.fn(() => ({})),
+}));
+
+vi.mock(UI_HELPERS_MODULE, () => ({
+  showToast: showToastMock,
+}));
+
+vi.mock(I18N_MODULE, () => ({
+  translate: vi.fn((key, params, fallback) => {
+    if (fallback) {
+      return Object.entries(params || {}).reduce(
+        (text, [name, value]) => text.replaceAll(`{${name}}`, value),
+        fallback
+      );
+    }
+    return key;
+  }),
+}));
+
+vi.mock(STORAGE_MODULE, () => ({
+  getStorageItem: vi.fn(),
+  getSessionItem: vi.fn(),
+  removeSessionItem: vi.fn(),
+  saveMapToStorage: vi.fn(),
+}));
+
+vi.mock(API_CONFIG_MODULE, () => ({
+  getCompleteApiConfig: vi.fn(() => ({
+    endpoints: { scan: '/api/lm/loras/scan' },
+    config: { displayName: 'LoRA', singularName: 'lora' },
+  })),
+  getCurrentModelType: vi.fn(() => 'loras'),
+  isValidModelType: vi.fn(() => true),
+  DOWNLOAD_ENDPOINTS: {},
+  HF_ENDPOINTS: {},
+  WS_ENDPOINTS: { fetchProgress: '/ws/fetch-progress' },
+}));
+
+vi.mock(API_FACTORY_MODULE, () => ({
+  resetAndReload: resetAndReloadMock,
+}));
+
+vi.mock(SIDEBAR_MANAGER_MODULE, () => ({
+  sidebarManager: { refresh: vi.fn() },
+}));
+
+class FakeWebSocket {
+  static instances = [];
+  static failNextConnection = false;
+
+  constructor(url) {
+    this.url = url;
+    this.onopen = null;
+    this.onerror = null;
+    this.onmessage = null;
+    this.close = vi.fn();
+    FakeWebSocket.instances.push(this);
+    const shouldFail = FakeWebSocket.failNextConnection;
+    FakeWebSocket.failNextConnection = false;
+    queueMicrotask(() => {
+      if (shouldFail) {
+        this.onerror?.(new Error('connection refused'));
+      } else {
+        this.onopen?.();
+      }
+    });
+  }
+
+  emit(data) {
+    this.onmessage?.({ data: JSON.stringify(data) });
+  }
+}
+
+async function createClient() {
+  const { BaseModelApiClient } = await import(BASE_MODEL_API_MODULE);
+  class TestClient extends BaseModelApiClient {}
+  return new TestClient('loras');
+}
+
+async function flushMicrotasks() {
+  await new Promise((resolve) => setTimeout(resolve, 0));
+}
+
+describe('BaseModelApiClient.refreshModels scan progress', () => {
+  beforeEach(() => {
+    showToastMock.mockReset();
+    showMock.mockReset();
+    showCancelButtonMock.mockReset();
+    hideMock.mockReset();
+    restoreProgressBarMock.mockReset();
+    setProgressMock.mockReset();
+    setStatusMock.mockReset();
+    resetAndReloadMock.mockReset();
+    FakeWebSocket.instances = [];
+    FakeWebSocket.failNextConnection = false;
+    vi.stubGlobal('WebSocket', FakeWebSocket);
+  });
+
+  afterEach(() => {
+    delete global.fetch;
+    vi.unstubAllGlobals();
+    vi.restoreAllMocks();
+  });
+
+  function mockFetchPending() {
+    let resolveFetch;
+    global.fetch = vi.fn(() => new Promise((resolve) => { resolveFetch = resolve; }));
+    return {
+      resolveOk: (payload = { status: 'success' }) =>
+        resolveFetch({ ok: true, json: async () => payload }),
+    };
+  }
+
+  async function startRefresh(client, fullRebuild = false) {
+    const promise = client.refreshModels(fullRebuild);
+    await vi.waitFor(() => {
+      expect(FakeWebSocket.instances.length).toBe(1);
+    });
+    await flushMicrotasks();
+    const socket = FakeWebSocket.instances[0];
+    await vi.waitFor(() => {
+      expect(socket.onmessage).toBeTruthy();
+    });
+    return { promise, socket };
+  }
+
+  it('shows scan progress updates from the WebSocket channel', async () => {
+    const fetchControl = mockFetchPending();
+    const client = await createClient();
+    const { promise, socket } = await startRefresh(client);
+
+    expect(socket.url).toBe(`ws://${window.location.host}/ws/fetch-progress`);
+
+    socket.emit({
+      type: 'scan_progress',
+      status: 'started',
+      stage: 'scan_folders',
+      model_type: 'lora',
+      pageType: 'loras',
+      full_rebuild: false,
+      progress: 0,
+    });
+    socket.emit({
+      type: 'scan_progress',
+      status: 'processing',
+      stage: 'process_models',
+      model_type: 'lora',
+      pageType: 'loras',
+      full_rebuild: false,
+      progress: 50,
+      processed: 5,
+      total: 10,
+      current_name: 'style.safetensors',
+    });
+
+    expect(setProgressMock).toHaveBeenCalledWith(0);
+    expect(setProgressMock).toHaveBeenCalledWith(50);
+    const lastStatus = setStatusMock.mock.calls.at(-1)[0];
+    expect(lastStatus).toContain('(5/10)');
+    expect(lastStatus).toContain('style.safetensors');
+    // First ETA sample only anchors the timer
+    expect(lastStatus).toContain('Estimating time...');
+
+    fetchControl.resolveOk();
+    await promise;
+
+    expect(resetAndReloadMock).toHaveBeenCalledWith(true);
+    expect(showToastMock).toHaveBeenCalledWith(
+      'toast.api.refreshComplete',
+      { action: 'Refresh' },
+      'success'
+    );
+    expect(socket.close).toHaveBeenCalled();
+    expect(hideMock).toHaveBeenCalled();
+  });
+
+  it('ignores messages for other types or other model types', async () => {
+    const fetchControl = mockFetchPending();
+    const client = await createClient();
+    const { promise, socket } = await startRefresh(client);
+
+    socket.emit({
+      type: 'scan_progress',
+      status: 'processing',
+      stage: 'process_models',
+      model_type: 'checkpoint',
+      progress: 33,
+      processed: 1,
+      total: 3,
+    });
+    socket.emit({
+      type: 'example_images_progress',
+      status: 'running',
+      model_type: 'lora',
+      progress: 66,
+      processed: 2,
+      total: 3,
+    });
+
+    expect(setProgressMock).not.toHaveBeenCalled();
+    expect(setStatusMock).not.toHaveBeenCalled();
+
+    fetchControl.resolveOk();
+    await promise;
+  });
+
+  it('falls back to plain loading when the WebSocket connection fails', async () => {
+    FakeWebSocket.failNextConnection = true;
+    global.fetch = vi.fn().mockResolvedValue({
+      ok: true,
+      json: async () => ({ status: 'success' }),
+    });
+
+    const client = await createClient();
+    await client.refreshModels(true);
+
+    expect(global.fetch).toHaveBeenCalled();
+    const [url] = global.fetch.mock.calls[0];
+    expect(url.searchParams.get('full_rebuild')).toBe('true');
+    expect(showMock).toHaveBeenCalledWith('Full rebuild LoRAs...', 0);
+    expect(showToastMock).toHaveBeenCalledWith(
+      'toast.api.refreshComplete',
+      { action: 'Full rebuild' },
+      'success'
+    );
+  });
+
+  it('computes an ETA with EMA smoothing once enough samples arrive', async () => {
+    const fetchControl = mockFetchPending();
+    let now = 1000;
+    vi.spyOn(Date, 'now').mockImplementation(() => now);
+
+    const client = await createClient();
+    const { promise, socket } = await startRefresh(client);
+
+    const emitProcessing = (processed, total) => socket.emit({
+      type: 'scan_progress',
+      status: 'processing',
+      stage: 'process_models',
+      model_type: 'lora',
+      progress: Math.floor((processed / total) * 100),
+      processed,
+      total,
+    });
+
+    // First sample anchors the timer
+    emitProcessing(1, 10);
+    expect(setStatusMock.mock.calls.at(-1)[0]).toContain('Estimating time...');
+
+    // 100s elapsed for 2 files -> 50s per file -> 400s remaining -> ~7 min
+    now = 101000;
+    emitProcessing(2, 10);
+    expect(setStatusMock.mock.calls.at(-1)[0]).toContain('~7 min remaining');
+
+    // 110s elapsed for 4 files -> EMA = 50000*0.7 + 27500*0.3 = 43250ms/file
+    // remaining 6 files -> 259.5s -> ~4 min
+    now = 111000;
+    emitProcessing(4, 10);
+    expect(setStatusMock.mock.calls.at(-1)[0]).toContain('~4 min remaining');
+
+    fetchControl.resolveOk();
+    await promise;
+  });
+
+  it('shows the cancelled toast when the server reports cancellation', async () => {
+    const fetchControl = mockFetchPending();
+    const client = await createClient();
+    const { promise } = await startRefresh(client);
+
+    fetchControl.resolveOk({ status: 'cancelled' });
+    await promise;
+
+    expect(showToastMock).toHaveBeenCalledWith('toast.api.operationCancelled', {}, 'info');
+    expect(resetAndReloadMock).not.toHaveBeenCalled();
+  });
+});
+
+describe('createScanEtaTracker / formatScanRemainingTime', () => {
+  it('estimates remaining time from EMA of per-file cost', async () => {
+    const { createScanEtaTracker } = await import(BASE_MODEL_API_MODULE);
+    let now = 0;
+    vi.spyOn(Date, 'now').mockImplementation(() => now);
+
+    const tracker = createScanEtaTracker();
+    expect(tracker.update(1, 10)).toBe('Estimating time...');
+
+    now = 60000; // 60s for 3 files -> 20s/file -> 7 * 20s = 140s -> ~2 min
+    expect(tracker.update(3, 10)).toBe('~2 min remaining');
+
+    now = 61000; // tiny delta keeps EMA near 20s/file
+    expect(tracker.update(4, 10)).toBe('~2 min remaining');
+
+    // Done: no ETA
+    expect(tracker.update(10, 10)).toBeNull();
+    expect(tracker.update(0, 0)).toBeNull();
+
+    vi.restoreAllMocks();
+  });
+
+  it('formats hours and sub-minute remainders', async () => {
+    const { formatScanRemainingTime } = await import(BASE_MODEL_API_MODULE);
+    expect(formatScanRemainingTime(30000)).toBe('Less than a minute remaining');
+    expect(formatScanRemainingTime(5 * 60000)).toBe('~5 min remaining');
+    expect(formatScanRemainingTime(3600000 + 30 * 60000)).toBe('~1 hr 30 min remaining');
+  });
+});
@@ -0,0 +1,285 @@
+import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
+
+const showToastMock = vi.hoisted(() => vi.fn());
+const loadingManagerMock = vi.hoisted(() => ({
+  show: vi.fn(),
+  hide: vi.fn(),
+  restoreProgressBar: vi.fn(),
+  setProgress: vi.fn(),
+  setStatus: vi.fn(),
+}));
+const virtualScrollerMock = vi.hoisted(() => ({
+  refreshWithData: vi.fn(),
+}));
+const getCurrentPageStateMock = vi.hoisted(() => vi.fn());
+const etaUpdateMock = vi.hoisted(() => vi.fn(() => 'ETA soon'));
+
+vi.mock('../../../static/js/components/RecipeCard.js', () => ({
+  RecipeCard: vi.fn(() => ({ element: document.createElement('div') })),
+}));
+
+vi.mock('../../../static/js/state/index.js', () => ({
+  state: {
+    loadingManager: loadingManagerMock,
+    virtualScroller: virtualScrollerMock,
+  },
+  getCurrentPageState: getCurrentPageStateMock,
+}));
+
+vi.mock('../../../static/js/utils/uiHelpers.js', () => ({
+  showToast: showToastMock,
+}));
+
+vi.mock('../../../static/js/utils/i18nHelpers.js', () => ({
+  translate: vi.fn((key, params, fallback) => {
+    if (fallback) {
+      return Object.entries(params || {}).reduce(
+        (text, [name, value]) => text.replaceAll(`{${name}}`, value),
+        fallback
+      );
+    }
+    return key;
+  }),
+}));
+
+vi.mock('../../../static/js/utils/infiniteScroll.js', () => ({
+  captureScrollPosition: vi.fn(),
+  restoreScrollPosition: vi.fn(),
+}));
+
+vi.mock('../../../static/js/api/apiConfig.js', () => ({
+  WS_ENDPOINTS: { fetchProgress: '/ws/fetch-progress' },
+}));
+
+vi.mock('../../../static/js/utils/scanEtaUtils.js', () => ({
+  createScanEtaTracker: () => ({ update: etaUpdateMock }),
+}));
+
+import { refreshRecipes } from '../../../static/js/api/recipeApi.js';
+
+class FakeWebSocket {
+  static instances = [];
+  static failNextConnection = false;
+
+  constructor(url) {
+    this.url = url;
+    this.onopen = null;
+    this.onerror = null;
+    this.onmessage = null;
+    this.close = vi.fn();
+    FakeWebSocket.instances.push(this);
+    const shouldFail = FakeWebSocket.failNextConnection;
+    FakeWebSocket.failNextConnection = false;
+    queueMicrotask(() => {
+      if (shouldFail) {
+        this.onerror?.(new Error('connection refused'));
+      } else {
+        this.onopen?.();
+      }
+    });
+  }
+
+  emit(data) {
+    this.onmessage?.({ data: JSON.stringify(data) });
+  }
+}
+
+async function flushMicrotasks() {
+  await new Promise((resolve) => setTimeout(resolve, 0));
+}
+
+describe('refreshRecipes scan progress', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    getCurrentPageStateMock.mockReturnValue({
+      pageSize: 50,
+      currentPage: 1,
+      hasMore: true,
+      isLoading: false,
+      sortBy: 'date:desc',
+      showFavoritesOnly: false,
+      activeFolder: null,
+      searchOptions: { recursive: true },
+      customFilter: { active: false },
+      filters: {},
+    });
+    FakeWebSocket.instances = [];
+    FakeWebSocket.failNextConnection = false;
+    vi.stubGlobal('WebSocket', FakeWebSocket);
+  });
+
+  afterEach(() => {
+    delete global.fetch;
+    vi.unstubAllGlobals();
+  });
+
+  function mockFetchPendingScan() {
+    let resolveScan;
+    global.fetch = vi.fn((input) => {
+      const url = String(input);
+      if (url.includes('/scan')) {
+        return new Promise((resolve) => { resolveScan = resolve; });
+      }
+      // Recipe list reload after the scan completes
+      return Promise.resolve({
+        ok: true,
+        json: async () => ({ items: [], total: 0, total_pages: 0 }),
+      });
+    });
+    return {
+      resolveOk: (payload = { status: 'success' }) =>
+        resolveScan({ ok: true, json: async () => payload }),
+      resolveNotOk: () =>
+        resolveScan({ ok: false, status: 500, statusText: 'Server Error' }),
+    };
+  }
+
+  async function startRefresh(fullRebuild = true) {
+    const promise = refreshRecipes(fullRebuild);
+    await vi.waitFor(() => {
+      expect(FakeWebSocket.instances.length).toBe(1);
+    });
+    await flushMicrotasks();
+    const socket = FakeWebSocket.instances[0];
+    await vi.waitFor(() => {
+      expect(socket.onmessage).toBeTruthy();
+    });
+    return { promise, socket };
+  }
+
+  it('shows scan progress updates from the WebSocket channel', async () => {
+    const fetchControl = mockFetchPendingScan();
+    const { promise, socket } = await startRefresh();
+
+    expect(socket.url).toBe(`ws://${window.location.host}/ws/fetch-progress`);
+
+    socket.emit({
+      type: 'scan_progress',
+      status: 'started',
+      stage: 'scan_folders',
+      model_type: 'recipe',
+      pageType: 'recipes',
+      full_rebuild: true,
+      progress: 0,
+    });
+    socket.emit({
+      type: 'scan_progress',
+      status: 'processing',
+      stage: 'process_models',
+      model_type: 'recipe',
+      pageType: 'recipes',
+      full_rebuild: true,
+      progress: 50,
+      processed: 5,
+      total: 10,
+      current_name: 'style.recipe.json',
+    });
+
+    expect(loadingManagerMock.setProgress).toHaveBeenCalledWith(0);
+    expect(loadingManagerMock.setProgress).toHaveBeenCalledWith(50);
+    const lastStatus = loadingManagerMock.setStatus.mock.calls.at(-1)[0];
+    expect(lastStatus).toContain('(5/10)');
+    expect(lastStatus).toContain('style.recipe.json');
+    expect(lastStatus).toContain('ETA soon');
+    expect(etaUpdateMock).toHaveBeenCalledWith(5, 10);
+
+    fetchControl.resolveOk();
+    await promise;
+
+    expect(showToastMock).toHaveBeenCalledWith(
+      'toast.api.refreshComplete',
+      { action: 'Full rebuild' },
+      'success'
+    );
+    expect(socket.close).toHaveBeenCalled();
+    expect(loadingManagerMock.hide).toHaveBeenCalled();
+  });
+
+  it('ignores messages for other types or other model types', async () => {
+    const fetchControl = mockFetchPendingScan();
+    const { promise, socket } = await startRefresh();
+
+    socket.emit({
+      type: 'scan_progress',
+      status: 'processing',
+      stage: 'process_models',
+      model_type: 'lora',
+      progress: 33,
+      processed: 1,
+      total: 3,
+    });
+    socket.emit({
+      type: 'example_images_progress',
+      status: 'running',
+      model_type: 'recipe',
+      progress: 66,
+      processed: 2,
+      total: 3,
+    });
+
+    expect(loadingManagerMock.setProgress).not.toHaveBeenCalled();
+    expect(loadingManagerMock.setStatus).not.toHaveBeenCalled();
+
+    fetchControl.resolveOk();
+    await promise;
+  });
+
+  it('falls back to plain loading when the WebSocket connection fails', async () => {
+    FakeWebSocket.failNextConnection = true;
+    global.fetch = vi.fn((input) => {
+      const url = String(input);
+      if (url.includes('/scan')) {
+        return Promise.resolve({
+          ok: true,
+          json: async () => ({ status: 'success' }),
+        });
+      }
+      return Promise.resolve({
+        ok: true,
+        json: async () => ({ items: [], total: 0, total_pages: 0 }),
+      });
+    });
+
+    await refreshRecipes(false);
+
+    expect(global.fetch).toHaveBeenCalled();
+    const [url] = global.fetch.mock.calls[0];
+    expect(url.searchParams.get('full_rebuild')).toBe('false');
+    expect(loadingManagerMock.show).toHaveBeenCalledWith('Refreshing Recipes...', 0);
+    expect(showToastMock).toHaveBeenCalledWith(
+      'toast.api.refreshComplete',
+      { action: 'Refresh' },
+      'success'
+    );
+  });
+
+  it('shows the cancelled toast when the server reports cancellation', async () => {
+    const fetchControl = mockFetchPendingScan();
+    const { promise } = await startRefresh();
+
+    fetchControl.resolveOk({ status: 'cancelled' });
+    await promise;
+
+    expect(showToastMock).toHaveBeenCalledWith('toast.api.operationCancelled', {}, 'info');
+    expect(showToastMock).not.toHaveBeenCalledWith(
+      'toast.api.refreshComplete',
+      expect.anything(),
+      expect.anything()
+    );
+  });
+
+  it('reports refresh failures through the error toast', async () => {
+    const fetchControl = mockFetchPendingScan();
+    const { promise } = await startRefresh();
+
+    fetchControl.resolveNotOk();
+    await promise;
+
+    expect(showToastMock).toHaveBeenCalledWith(
+      'toast.api.refreshFailed',
+      { action: 'rebuild', type: 'recipe' },
+      'error'
+    );
+    expect(loadingManagerMock.hide).toHaveBeenCalled();
+  });
+});