From 4bfc8e9e338ce1d6cf1e7b21437f36ba923e561b Mon Sep 17 00:00:00 2001 From: Christian Byrne Date: Thu, 27 Mar 2025 08:49:05 -0700 Subject: [PATCH] [Manager] Fetch lists of node packs in single request (#3250) --- .../content/manager/ManagerDialogContent.vue | 84 +++++++++---------- src/composables/nodePack/useNodePacks.ts | 69 ++++----------- src/services/comfyRegistryService.ts | 4 + src/stores/comfyRegistryStore.ts | 53 +++++++++++- .../tests/store/comfyRegistryStore.test.ts | 71 +++++++++++++++- 5 files changed, 181 insertions(+), 100 deletions(-) diff --git a/src/components/dialog/content/manager/ManagerDialogContent.vue b/src/components/dialog/content/manager/ManagerDialogContent.vue index 0c707d727..539de24ae 100644 --- a/src/components/dialog/content/manager/ManagerDialogContent.vue +++ b/src/components/dialog/content/manager/ManagerDialogContent.vue @@ -188,59 +188,52 @@ const { isLoading: isLoadingWorkflow } = useWorkflowPacks() -const getInstalledResults = () => { - if (isEmptySearch.value) { - startFetchInstalled() - return installedPacks.value - } else { - return filterInstalledPack(searchResults.value) - } -} - -const getInWorkflowResults = () => { - if (isEmptySearch.value) { - startFetchWorkflowPacks() - return workflowPacks.value - } else { - return filterWorkflowPack(searchResults.value) - } -} - const filterMissingPacks = (packs: components['schemas']['Node'][]) => packs.filter((pack) => !comfyManagerStore.isPackInstalled(pack.id)) -const setMissingPacks = () => { - displayPacks.value = filterMissingPacks(workflowPacks.value) -} +const isInstalledTab = computed( + () => selectedTab.value?.id === ManagerTab.Installed +) +const isMissingTab = computed( + () => selectedTab.value?.id === ManagerTab.Missing +) +const isWorkflowTab = computed( + () => selectedTab.value?.id === ManagerTab.Workflow +) +const isAllTab = computed(() => selectedTab.value?.id === ManagerTab.All) -const getMissingPacks = () => { - if (isEmptySearch.value) { - startFetchWorkflowPacks() - whenever(() => workflowPacks.value.length, setMissingPacks, { - immediate: true, - once: true - }) - return filterMissingPacks(workflowPacks.value) +watch([isInstalledTab, installedPacks], () => { + if (!isInstalledTab.value) return + + if (!isEmptySearch.value) { + displayPacks.value = filterInstalledPack(searchResults.value) + } else if (!installedPacks.value.length) { + startFetchInstalled() } else { - return filterMissingPacks(filterWorkflowPack(searchResults.value)) + displayPacks.value = installedPacks.value } -} +}) -const onTabChange = () => { - switch (selectedTab.value?.id) { - case ManagerTab.Installed: - displayPacks.value = getInstalledResults() - break - case ManagerTab.Workflow: - displayPacks.value = getInWorkflowResults() - break - case ManagerTab.Missing: - displayPacks.value = getMissingPacks() - break - default: - displayPacks.value = searchResults.value +watch([isMissingTab, isWorkflowTab, workflowPacks], () => { + if (!isWorkflowTab.value && !isMissingTab.value) return + + if (!isEmptySearch.value) { + displayPacks.value = isMissingTab.value + ? filterMissingPacks(filterWorkflowPack(searchResults.value)) + : filterWorkflowPack(searchResults.value) + } else if (!workflowPacks.value.length) { + startFetchWorkflowPacks() + } else { + displayPacks.value = isMissingTab.value + ? filterMissingPacks(workflowPacks.value) + : workflowPacks.value } -} +}) + +watch([isAllTab, searchResults], () => { + if (!isAllTab.value) return + displayPacks.value = searchResults.value +}) const onResultsChange = () => { switch (selectedTab.value?.id) { @@ -260,7 +253,6 @@ const onResultsChange = () => { } } -whenever(selectedTab, onTabChange) watch(searchResults, onResultsChange, { flush: 'pre' }) watch(() => comfyManagerStore.installedPacksIds, onResultsChange) diff --git a/src/composables/nodePack/useNodePacks.ts b/src/composables/nodePack/useNodePacks.ts index 1ce87a736..e4fab9d69 100644 --- a/src/composables/nodePack/useNodePacks.ts +++ b/src/composables/nodePack/useNodePacks.ts @@ -1,14 +1,8 @@ -import { useAsyncState } from '@vueuse/core' -import { chunk } from 'lodash' -import { Ref, computed, isRef, ref } from 'vue' +import { get, useAsyncState } from '@vueuse/core' +import { Ref } from 'vue' import { useComfyRegistryStore } from '@/stores/comfyRegistryStore' import { UseNodePacksOptions } from '@/types/comfyManagerTypes' -import { components } from '@/types/comfyRegistryTypes' - -const DEFAULT_MAX_CONCURRENT = 6 - -type NodePack = components['schemas']['Node'] /** * Handles fetching node packs from the registry given a list of node pack IDs @@ -17,54 +11,25 @@ export const useNodePacks = ( packsIds: string[] | Ref, options: UseNodePacksOptions = {} ) => { - const { immediate = false, maxConcurrent = DEFAULT_MAX_CONCURRENT } = options - const { getPackById } = useComfyRegistryStore() + const { immediate = false } = options + const { getPacksByIds } = useComfyRegistryStore() - const nodePacks = ref([]) - const processedIds = ref>(new Set()) + const fetchPacks = () => getPacksByIds.call(get(packsIds).filter(Boolean)) - const queuedPackIds = isRef(packsIds) ? packsIds : ref(packsIds) - const remainingIds = computed(() => - queuedPackIds.value?.filter((id) => !processedIds.value.has(id)) - ) - const chunks = computed(() => - remainingIds.value?.length ? chunk(remainingIds.value, maxConcurrent) : [] - ) - - const fetchPack = (id: Parameters[0]) => - id ? getPackById.call(id) : null - - const toRequestBatch = async (ids: string[]) => - Promise.all(ids.map(fetchPack)) - - const isValidResponse = (response: NodePack | null) => response !== null - - const fetchPacks = async () => { - for (const chunk of chunks.value) { - const resolvedChunk = await toRequestBatch(chunk) - chunk.forEach((id) => processedIds.value.add(id)) - if (!resolvedChunk) continue - nodePacks.value.push(...resolvedChunk.filter(isValidResponse)) - } - } - - const { isReady, isLoading, error, execute } = useAsyncState( - fetchPacks, - null, - { - immediate - } - ) - - const clear = () => { - queuedPackIds.value = [] - isReady.value = false - isLoading.value = false - } + const { + isReady, + isLoading, + error, + execute, + state: nodePacks + } = useAsyncState(fetchPacks, [], { + immediate + }) const cleanup = () => { - getPackById.cancel() - clear() + getPacksByIds.cancel() + isReady.value = false + isLoading.value = false } return { diff --git a/src/services/comfyRegistryService.ts b/src/services/comfyRegistryService.ts index df120108c..7a55959fc 100644 --- a/src/services/comfyRegistryService.ts +++ b/src/services/comfyRegistryService.ts @@ -11,6 +11,10 @@ const registryApiClient = axios.create({ baseURL: API_BASE_URL, headers: { 'Content-Type': 'application/json' + }, + paramsSerializer: { + // Disables PHP-style notation (e.g. param[]=value) in favor of repeated params (e.g. param=value1¶m=value2) + indexes: null } }) diff --git a/src/stores/comfyRegistryStore.ts b/src/stores/comfyRegistryStore.ts index 4a2800cc9..12eb929c3 100644 --- a/src/stores/comfyRegistryStore.ts +++ b/src/stores/comfyRegistryStore.ts @@ -1,3 +1,5 @@ +import QuickLRU from '@alloc/quick-lru' +import { partition } from 'lodash' import { defineStore } from 'pinia' import { useCachedRequest } from '@/composables/useCachedRequest' @@ -5,7 +7,7 @@ import { useComfyRegistryService } from '@/services/comfyRegistryService' import type { components, operations } from '@/types/comfyRegistryTypes' const PACK_LIST_CACHE_SIZE = 20 -const PACK_BY_ID_CACHE_SIZE = 50 +const PACK_BY_ID_CACHE_SIZE = 64 type NodePack = components['schemas']['Node'] type ListPacksParams = operations['listAllNodes']['parameters']['query'] @@ -14,12 +16,21 @@ type ListPacksResult = type ComfyNode = components['schemas']['ComfyNode'] type GetPackByIdPath = operations['getNode']['parameters']['path']['nodeId'] +const isNodePack = (pack: NodePack | undefined): pack is NodePack => { + return pack !== undefined && 'id' in pack +} + /** * Store for managing remote custom nodes */ export const useComfyRegistryStore = defineStore('comfyRegistry', () => { const registryService = useComfyRegistryService() + let getPacksByIdController: AbortController | null = null + const getPacksByIdCache = new QuickLRU({ + maxSize: PACK_BY_ID_CACHE_SIZE + }) + /** * Get a list of all node packs from the registry */ @@ -39,6 +50,41 @@ export const useComfyRegistryStore = defineStore('comfyRegistry', () => { { maxSize: PACK_BY_ID_CACHE_SIZE } ) + /** + * Get a list of packs by their IDs from the registry + */ + const getPacksByIds = async (ids: NodePack['id'][]): Promise => { + const [cachedPacksIds, uncachedPacksIds] = partition(ids, (id) => + getPacksByIdCache.has(id) + ) + + const resolvedPacks = cachedPacksIds + .map((id) => getPacksByIdCache.get(id)) + .filter(isNodePack) + + if (uncachedPacksIds.length) { + getPacksByIdController = new AbortController() + const uncachedPacks = await registryService.listAllPacks( + { + node_id: uncachedPacksIds.filter( + (id): id is string => id !== undefined + ) + }, + getPacksByIdController.signal + ) + + const { nodes = [] } = uncachedPacks ?? {} + nodes.forEach((pack) => { + if (pack?.id) { + getPacksByIdCache.set(pack.id, pack) + resolvedPacks.push(pack) + } + }) + } + + return resolvedPacks + } + /** * Get the node definitions for a pack */ @@ -63,11 +109,16 @@ export const useComfyRegistryStore = defineStore('comfyRegistry', () => { getNodeDefs.cancel() listAllPacks.cancel() getPackById.cancel() + getPacksByIdController?.abort() } return { listAllPacks, getPackById, + getPacksByIds: { + call: getPacksByIds, + cancel: () => getPacksByIdController?.abort() + }, getNodeDefs, clearCache, diff --git a/tests-ui/tests/store/comfyRegistryStore.test.ts b/tests-ui/tests/store/comfyRegistryStore.test.ts index 64f2eb9a5..44db6dfeb 100644 --- a/tests-ui/tests/store/comfyRegistryStore.test.ts +++ b/tests-ui/tests/store/comfyRegistryStore.test.ts @@ -26,6 +26,38 @@ const mockNodePack: components['schemas']['Node'] = { } } +const mockNodePack2: components['schemas']['Node'] = { + id: 'test-pack-id-2', + name: 'Test Pack 2', + description: 'A second test node pack', + downloads: 1000, + publisher: { + id: 'test-publisher', + name: 'Test Publisher' + }, + latest_version: { + id: 'test-version', + version: '1.0.0', + createdAt: '2023-01-01T00:00:00Z' + } +} + +const mockNodePack3: components['schemas']['Node'] = { + id: 'test-pack-id-3', + name: 'Test Pack 3', + description: 'A third test node pack', + downloads: 1000, + publisher: { + id: 'test-publisher', + name: 'Test Publisher' + }, + latest_version: { + id: 'test-version', + version: '1.0.0', + createdAt: '2023-01-01T00:00:00Z' + } +} + const mockListResult: operations['listAllNodes']['responses'][200]['content']['application/json'] = { nodes: [mockNodePack], @@ -48,7 +80,32 @@ describe('useComfyRegistryStore', () => { mockRegistryService = { isLoading: ref(false), error: ref(null), - listAllPacks: vi.fn().mockResolvedValue(mockListResult), + listAllPacks: vi.fn().mockImplementation((params) => { + // If node_id is provided, return specific nodes + if (params.node_id) { + return Promise.resolve({ + nodes: params.node_id + .map((id: string) => { + switch (id) { + case 'test-pack-id': + return mockNodePack + case 'test-pack-id-2': + return mockNodePack2 + case 'test-pack-id-3': + return mockNodePack3 + default: + return null + } + }) + .filter(Boolean), + total: params.node_id.length, + page: 1, + limit: 10 + }) + } + // Otherwise return paginated results + return Promise.resolve(mockListResult) + }), getPackById: vi.fn().mockResolvedValue(mockNodePack) } @@ -117,4 +174,16 @@ describe('useComfyRegistryStore', () => { expect(result).toBeNull() }) + + it('should fetch packs by IDs', async () => { + const store = useComfyRegistryStore() + const packIds = ['test-pack-id', 'test-pack-id-2', 'test-pack-id-3'] + const result = await store.getPacksByIds.call(packIds) + + expect(result).toEqual([mockNodePack, mockNodePack2, mockNodePack3]) + expect(mockRegistryService.listAllPacks).toHaveBeenCalledWith( + { node_id: packIds }, + expect.any(Object) // abort signal + ) + }) })