feat: implement progressive pagination for Asset Browser model assets (#8212)

## Summary

Implements progressive pagination for model assets - returns the first
batch immediately while loading remaining batches in the background.

## Changes

### Store (`assetsStore.ts`)
- Adds `ModelPaginationState` tracking (assets Map, offset, hasMore,
loading, error)
- `updateModelsForKey()` returns first batch, then calls
`loadRemainingBatches()` to fetch the rest
- Accessor functions `getAssets(key)`, `isModelLoading(key)` replace
direct Map access

### API (`assetService.ts`)
- Adds `PaginationOptions` interface (`{ limit?, offset? }`)

### Components
- `AssetBrowserModal.vue` uses new accessor API

### Tests
- Updated mocks for new accessor pattern

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-8212-feat-implement-progressive-pagination-for-Asset-Browser-model-assets-2ef6d73d36508157af04d1264780997e)
by [Unito](https://www.unito.io)

---------

Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
Alexander Brown
2026-01-21 17:59:08 -08:00
committed by GitHub
parent f08b0f44ef
commit 482159957e
9 changed files with 441 additions and 171 deletions

View File

@@ -6,6 +6,9 @@ import AssetBrowserModal from '@/platform/assets/components/AssetBrowserModal.vu
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
import { useAssetsStore } from '@/stores/assetsStore'
const mockAssetsByKey = vi.hoisted(() => new Map<string, AssetItem[]>())
const mockLoadingByKey = vi.hoisted(() => new Map<string, boolean>())
vi.mock('@/i18n', () => ({
t: (key: string, params?: Record<string, string>) =>
params ? `${key}:${JSON.stringify(params)}` : key,
@@ -13,13 +16,20 @@ vi.mock('@/i18n', () => ({
}))
vi.mock('@/stores/assetsStore', () => {
const store = {
modelAssetsByNodeType: new Map<string, AssetItem[]>(),
modelLoadingByNodeType: new Map<string, boolean>(),
updateModelsForNodeType: vi.fn(),
updateModelsForTag: vi.fn()
const getAssets = vi.fn((key: string) => mockAssetsByKey.get(key) ?? [])
const isModelLoading = vi.fn(
(key: string) => mockLoadingByKey.get(key) ?? false
)
const updateModelsForNodeType = vi.fn()
const updateModelsForTag = vi.fn()
return {
useAssetsStore: () => ({
getAssets,
isModelLoading,
updateModelsForNodeType,
updateModelsForTag
})
}
return { useAssetsStore: () => store }
})
vi.mock('@/stores/modelToNodeStore', () => ({
@@ -183,12 +193,10 @@ describe('AssetBrowserModal', () => {
})
}
const mockStore = useAssetsStore()
beforeEach(() => {
vi.resetAllMocks()
mockStore.modelAssetsByNodeType.clear()
mockStore.modelLoadingByNodeType.clear()
mockAssetsByKey.clear()
mockLoadingByKey.clear()
})
describe('Integration with useAssetBrowser', () => {
@@ -197,7 +205,7 @@ describe('AssetBrowserModal', () => {
createTestAsset('asset1', 'Model A', 'checkpoints'),
createTestAsset('asset2', 'Model B', 'loras')
]
mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
mockAssetsByKey.set('CheckpointLoaderSimple', assets)
const wrapper = createWrapper({ nodeType: 'CheckpointLoaderSimple' })
await flushPromises()
@@ -214,7 +222,7 @@ describe('AssetBrowserModal', () => {
createTestAsset('c1', 'model.safetensors', 'checkpoints'),
createTestAsset('l1', 'lora.pt', 'loras')
]
mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
mockAssetsByKey.set('CheckpointLoaderSimple', assets)
const wrapper = createWrapper({
nodeType: 'CheckpointLoaderSimple',
@@ -231,17 +239,18 @@ describe('AssetBrowserModal', () => {
describe('Data fetching', () => {
it('triggers store refresh for node type on mount', async () => {
const store = useAssetsStore()
createWrapper({ nodeType: 'CheckpointLoaderSimple' })
await flushPromises()
expect(mockStore.updateModelsForNodeType).toHaveBeenCalledWith(
expect(store.updateModelsForNodeType).toHaveBeenCalledWith(
'CheckpointLoaderSimple'
)
})
it('displays cached assets immediately from store', async () => {
const assets = [createTestAsset('asset1', 'Cached Model', 'checkpoints')]
mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
mockAssetsByKey.set('CheckpointLoaderSimple', assets)
const wrapper = createWrapper({ nodeType: 'CheckpointLoaderSimple' })
@@ -253,15 +262,16 @@ describe('AssetBrowserModal', () => {
})
it('triggers store refresh for asset type (tag) on mount', async () => {
const store = useAssetsStore()
createWrapper({ assetType: 'models' })
await flushPromises()
expect(mockStore.updateModelsForTag).toHaveBeenCalledWith('models')
expect(store.updateModelsForTag).toHaveBeenCalledWith('models')
})
it('uses tag: prefix for cache key when assetType is provided', async () => {
const assets = [createTestAsset('asset1', 'Tagged Model', 'models')]
mockStore.modelAssetsByNodeType.set('tag:models', assets)
mockAssetsByKey.set('tag:models', assets)
const wrapper = createWrapper({ assetType: 'models' })
await flushPromises()
@@ -277,7 +287,7 @@ describe('AssetBrowserModal', () => {
describe('Asset Selection', () => {
it('emits asset-select event when asset is selected', async () => {
const assets = [createTestAsset('asset1', 'Model A', 'checkpoints')]
mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
mockAssetsByKey.set('CheckpointLoaderSimple', assets)
const wrapper = createWrapper({ nodeType: 'CheckpointLoaderSimple' })
await flushPromises()
@@ -290,7 +300,7 @@ describe('AssetBrowserModal', () => {
it('executes onSelect callback when provided', async () => {
const assets = [createTestAsset('asset1', 'Model A', 'checkpoints')]
mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
mockAssetsByKey.set('CheckpointLoaderSimple', assets)
const onSelect = vi.fn()
const wrapper = createWrapper({
@@ -333,7 +343,7 @@ describe('AssetBrowserModal', () => {
createTestAsset('asset1', 'Model A', 'checkpoints'),
createTestAsset('asset2', 'Model B', 'loras')
]
mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
mockAssetsByKey.set('CheckpointLoaderSimple', assets)
const wrapper = createWrapper({
nodeType: 'CheckpointLoaderSimple',
@@ -366,7 +376,7 @@ describe('AssetBrowserModal', () => {
it('passes computed contentTitle to BaseModalLayout when no title prop', async () => {
const assets = [createTestAsset('asset1', 'Model A', 'checkpoints')]
mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
mockAssetsByKey.set('CheckpointLoaderSimple', assets)
const wrapper = createWrapper({ nodeType: 'CheckpointLoaderSimple' })
await flushPromises()

View File

@@ -112,27 +112,21 @@ const cacheKey = computed(() => {
})
// Read directly from store cache - reactive to any store updates
const fetchedAssets = computed(
() => assetStore.modelAssetsByNodeType.get(cacheKey.value) ?? []
)
const fetchedAssets = computed(() => assetStore.getAssets(cacheKey.value))
const isStoreLoading = computed(
() => assetStore.modelLoadingByNodeType.get(cacheKey.value) ?? false
)
const isStoreLoading = computed(() => assetStore.isModelLoading(cacheKey.value))
// Only show loading spinner when loading AND no cached data
const isLoading = computed(
() => isStoreLoading.value && fetchedAssets.value.length === 0
)
async function refreshAssets(): Promise<AssetItem[]> {
async function refreshAssets(): Promise<void> {
if (props.nodeType) {
return await assetStore.updateModelsForNodeType(props.nodeType)
await assetStore.updateModelsForNodeType(props.nodeType)
} else if (props.assetType) {
await assetStore.updateModelsForTag(props.assetType)
}
if (props.assetType) {
return await assetStore.updateModelsForTag(props.assetType)
}
return []
}
// Trigger background refresh on mount

View File

@@ -160,7 +160,7 @@ describe('assetService', () => {
const result = await assetService.getAssetModels('checkpoints')
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=models,checkpoints&limit=500'
'/assets?include_tags=models%2Ccheckpoints&limit=500'
)
expect(result).toEqual([
expect.objectContaining({ name: 'valid.safetensors', pathIndex: 0 })
@@ -231,9 +231,9 @@ describe('assetService', () => {
)
expect(result).toEqual(testAssets)
// Verify API call includes correct category
// Verify API call includes correct category (comma is URL-encoded by URLSearchParams)
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=models,checkpoints&limit=500'
'/assets?include_tags=models%2Ccheckpoints&limit=500'
)
})
@@ -400,7 +400,7 @@ describe('assetService', () => {
})
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=models&limit=500&include_public=true&offset=50'
'/assets?include_tags=models&limit=500&offset=50&include_public=true'
)
expect(result).toEqual(testAssets)
})
@@ -415,7 +415,7 @@ describe('assetService', () => {
})
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=input&limit=100&include_public=false&offset=25'
'/assets?include_tags=input&limit=100&offset=25&include_public=false'
)
expect(result).toEqual(testAssets)
})

View File

@@ -1,6 +1,7 @@
import { fromZodError } from 'zod-validation-error'
import { st } from '@/i18n'
import {
assetItemSchema,
assetResponseSchema,
@@ -17,6 +18,16 @@ import type {
import { api } from '@/scripts/api'
import { useModelToNodeStore } from '@/stores/modelToNodeStore'
export interface PaginationOptions {
limit?: number
offset?: number
}
interface AssetRequestOptions extends PaginationOptions {
includeTags: string[]
includePublic?: boolean
}
/**
* Maps CivitAI validation error codes to localized error messages
*/
@@ -77,9 +88,27 @@ function createAssetService() {
* Handles API response with consistent error handling and Zod validation
*/
async function handleAssetRequest(
url: string,
options: AssetRequestOptions,
context: string
): Promise<AssetResponse> {
const {
includeTags,
limit = DEFAULT_LIMIT,
offset,
includePublic
} = options
const queryParams = new URLSearchParams({
include_tags: includeTags.join(','),
limit: limit.toString()
})
if (offset !== undefined && offset > 0) {
queryParams.set('offset', offset.toString())
}
if (includePublic !== undefined) {
queryParams.set('include_public', includePublic ? 'true' : 'false')
}
const url = `${ASSETS_ENDPOINT}?${queryParams.toString()}`
const res = await api.fetchApi(url)
if (!res.ok) {
throw new Error(
@@ -101,7 +130,7 @@ function createAssetService() {
*/
async function getAssetModelFolders(): Promise<ModelFolder[]> {
const data = await handleAssetRequest(
`${ASSETS_ENDPOINT}?include_tags=${MODELS_TAG}&limit=${DEFAULT_LIMIT}`,
{ includeTags: [MODELS_TAG] },
'model folders'
)
@@ -130,7 +159,7 @@ function createAssetService() {
*/
async function getAssetModels(folder: string): Promise<ModelFile[]> {
const data = await handleAssetRequest(
`${ASSETS_ENDPOINT}?include_tags=${MODELS_TAG},${folder}&limit=${DEFAULT_LIMIT}`,
{ includeTags: [MODELS_TAG, folder] },
`models for ${folder}`
)
@@ -169,9 +198,15 @@ function createAssetService() {
* and fetching all assets with that category tag
*
* @param nodeType - The ComfyUI node type (e.g., 'CheckpointLoaderSimple')
* @param options - Pagination options
* @param options.limit - Maximum number of assets to return (default: 500)
* @param options.offset - Number of assets to skip (default: 0)
* @returns Promise<AssetItem[]> - Full asset objects with preserved metadata
*/
async function getAssetsForNodeType(nodeType: string): Promise<AssetItem[]> {
async function getAssetsForNodeType(
nodeType: string,
{ limit = DEFAULT_LIMIT, offset = 0 }: PaginationOptions = {}
): Promise<AssetItem[]> {
if (!nodeType || typeof nodeType !== 'string') {
return []
}
@@ -186,7 +221,7 @@ function createAssetService() {
// Fetch assets for this category using same API pattern as getAssetModels
const data = await handleAssetRequest(
`${ASSETS_ENDPOINT}?include_tags=${MODELS_TAG},${category}&limit=${DEFAULT_LIMIT}`,
{ includeTags: [MODELS_TAG, category], limit, offset },
`assets for ${nodeType}`
)
@@ -242,23 +277,10 @@ function createAssetService() {
async function getAssetsByTag(
tag: string,
includePublic: boolean = true,
{
limit = DEFAULT_LIMIT,
offset = 0
}: { limit?: number; offset?: number } = {}
{ limit = DEFAULT_LIMIT, offset = 0 }: PaginationOptions = {}
): Promise<AssetItem[]> {
const queryParams = new URLSearchParams({
include_tags: tag,
limit: limit.toString(),
include_public: includePublic ? 'true' : 'false'
})
if (offset > 0) {
queryParams.set('offset', offset.toString())
}
const data = await handleAssetRequest(
`${ASSETS_ENDPOINT}?${queryParams.toString()}`,
{ includeTags: [tag], limit, offset, includePublic },
`assets for tag ${tag}`
)