[Manager] Fetch lists of node packs in single request (#3250)

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)
 

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<string[]>,
   options: UseNodePacksOptions = {}
 ) => {
-  const { immediate = false, maxConcurrent = DEFAULT_MAX_CONCURRENT } = options
-  const { getPackById } = useComfyRegistryStore()
+  const { immediate = false } = options
+  const { getPacksByIds } = useComfyRegistryStore()
 
-  const nodePacks = ref<NodePack[]>([])
-  const processedIds = ref<Set<string>>(new Set())
+  const fetchPacks = () => getPacksByIds.call(get(packsIds).filter(Boolean))
 
-  const queuedPackIds = isRef(packsIds) ? packsIds : ref<string[]>(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<typeof getPackById.call>[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 {

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&param=value2)
+    indexes: null
   }
 })
 

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<NodePack['id'], NodePack>({
+    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<NodePack[]> => {
+    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,

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
+    )
+  })
 })