diff --git a/src/platform/assets/composables/media/assetMappers.test.ts b/src/platform/assets/composables/media/assetMappers.test.ts deleted file mode 100644 index 4bea3580b6..0000000000 --- a/src/platform/assets/composables/media/assetMappers.test.ts +++ /dev/null @@ -1,55 +0,0 @@ -import { describe, expect, it, vi } from 'vitest' - -import { mapInputFileToAssetItem } from './assetMappers' - -vi.mock('@/scripts/api', () => ({ - api: { - apiURL: (path: string) => `/api${path}` - } -})) - -vi.mock('@/platform/distribution/cloudPreviewUtil', () => ({ - appendCloudResParam: vi.fn() -})) - -describe('mapInputFileToAssetItem', () => { - it('preserves a clean filename', () => { - const asset = mapInputFileToAssetItem('photo.png', 0, 'input') - - expect(asset.name).toBe('photo.png') - expect(asset.id).toBe('input-0-photo.png') - expect(asset.preview_url).toBe('/api/view?filename=photo.png&type=input') - }) - - it.for([ - ['photo.png [input]', 'photo.png'], - ['photo.png [output]', 'photo.png'], - ['photo.png [temp]', 'photo.png'], - ['clip.mp4[input]', 'clip.mp4'], - ['MyFile.WEBP [Input]', 'MyFile.WEBP'] - ])( - 'strips ComfyUI directory annotation: %s -> %s', - ([input, expectedName]) => { - const asset = mapInputFileToAssetItem(input, 1, 'input') - - expect(asset.name).toBe(expectedName) - expect(asset.id).toBe(`input-1-${expectedName}`) - expect(asset.preview_url).toBe( - `/api/view?filename=${encodeURIComponent(expectedName)}&type=input` - ) - } - ) - - it('leaves non-annotation brackets in the filename intact', () => { - const asset = mapInputFileToAssetItem('my [draft] image.png', 0, 'input') - - expect(asset.name).toBe('my [draft] image.png') - }) - - it('uses the directory passed in for the type query param', () => { - const asset = mapInputFileToAssetItem('clip.mp4 [output]', 0, 'output') - - expect(asset.preview_url).toBe('/api/view?filename=clip.mp4&type=output') - expect(asset.tags).toEqual(['output']) - }) -}) diff --git a/src/platform/assets/composables/media/assetMappers.ts b/src/platform/assets/composables/media/assetMappers.ts index 67544a0e87..d394264a3a 100644 --- a/src/platform/assets/composables/media/assetMappers.ts +++ b/src/platform/assets/composables/media/assetMappers.ts @@ -1,8 +1,6 @@ import type { AssetItem } from '@/platform/assets/schemas/assetSchema' import type { OutputAssetMetadata } from '@/platform/assets/schemas/assetMetadataSchema' import type { AssetContext } from '@/platform/assets/schemas/mediaAssetSchema' -import { appendCloudResParam } from '@/platform/distribution/cloudPreviewUtil' -import { api } from '@/scripts/api' import type { ResultItemImpl, TaskItemImpl } from '@/stores/queueStore' /** @@ -50,42 +48,3 @@ export function mapTaskOutputToAssetItem( user_metadata: metadata } } - -/** - * Strips ComfyUI's trailing directory-type annotation (e.g. ` [input]`, - * ` [output]`, `[temp]`) from a filename returned by the OSS internal - * `/internal/files/{type}` endpoint. The annotation is part of the wire - * format LoadImage-style widgets expect, but for the assets sidebar we - * want the canonical on-disk filename so type detection / titles work. - */ -function stripDirectoryAnnotation(filename: string): string { - return filename.replace(/\s*\[(?:input|output|temp)\]\s*$/i, '') -} - -/** - * Maps input directory file to AssetItem format - * @param filename The filename - * @param index File index for unique ID - * @param directory The directory type - * @returns AssetItem formatted object - */ -export function mapInputFileToAssetItem( - filename: string, - index: number, - directory: 'input' | 'output' = 'input' -): AssetItem { - const cleanName = stripDirectoryAnnotation(filename) - const params = new URLSearchParams({ filename: cleanName, type: directory }) - const preview_url = api.apiURL(`/view?${params}`) - appendCloudResParam(params, cleanName) - - return { - id: `${directory}-${index}-${cleanName}`, - name: cleanName, - size: 0, - created_at: new Date().toISOString(), - tags: [directory], - thumbnail_url: api.apiURL(`/view?${params}`), - preview_url - } -} diff --git a/src/stores/assetsStore.test.ts b/src/stores/assetsStore.test.ts index af822532f8..88eb88d988 100644 --- a/src/stores/assetsStore.test.ts +++ b/src/stores/assetsStore.test.ts @@ -36,12 +36,8 @@ vi.mock('@/platform/assets/services/assetService', () => ({ OUTPUT_TAG: 'output' })) -// Mock distribution type - hoisted so it can be changed per test -const mockIsCloud = vi.hoisted(() => ({ value: false })) vi.mock('@/platform/distribution/types', () => ({ - get isCloud() { - return mockIsCloud.value - } + isCloud: true })) // Mock modelToNodeStore with proper node providers and category lookups @@ -155,14 +151,6 @@ vi.mock('@/stores/queueStore', () => ({ // Mock asset mappers - add unique timestamps vi.mock('@/platform/assets/composables/media/assetMappers', () => ({ - mapInputFileToAssetItem: vi.fn((name, index, type) => ({ - id: `${type}-${index}`, - name, - size: 0, - created_at: new Date(Date.now() - index * 1000).toISOString(), - tags: [type], - preview_url: `http://test.com/${name}` - })), mapTaskOutputToAssetItem: vi.fn((task, output) => { const index = parseInt(task.jobId.split('_')[1]) || 0 return { @@ -770,17 +758,12 @@ describe('assetsStore - Refactored (Option A)', () => { }) }) -describe('assetsStore - Model Assets Cache (Cloud)', () => { +describe('assetsStore - Model Assets Cache', () => { beforeEach(() => { setActivePinia(createTestingPinia({ stubActions: false })) - mockIsCloud.value = true vi.clearAllMocks() }) - afterEach(() => { - mockIsCloud.value = false - }) - const createMockAsset = (id: string, tags: string[] = ['models']) => ({ id, name: `asset-${id}`, @@ -1454,25 +1437,20 @@ describe('assetsStore - Deletion State and Input Mapping', () => { describe('getInputName', () => { it('resolves a hashed filename to the human-readable name when the input asset is in the cache', async () => { - mockIsCloud.value = true - try { - setActivePinia(createTestingPinia({ stubActions: false })) - const store = useAssetsStore() + setActivePinia(createTestingPinia({ stubActions: false })) + const store = useAssetsStore() - vi.mocked(assetService.getAssetsByTag).mockResolvedValueOnce([ - { - id: 'input-1', - name: 'cute-puppy.png', - asset_hash: 'abc123def.png', - tags: ['input'] - } - ]) - await store.updateInputs() + vi.mocked(assetService.getAssetsByTag).mockResolvedValueOnce([ + { + id: 'input-1', + name: 'cute-puppy.png', + asset_hash: 'abc123def.png', + tags: ['input'] + } + ]) + await store.updateInputs() - expect(store.getInputName('abc123def.png')).toBe('cute-puppy.png') - } finally { - mockIsCloud.value = false - } + expect(store.getInputName('abc123def.png')).toBe('cute-puppy.png') }) it('falls back to the original filename when the input asset is not cached', () => { @@ -1481,27 +1459,22 @@ describe('assetsStore - Deletion State and Input Mapping', () => { }) }) - describe('updateInputs cloud routing', () => { - it('reads from assetService.getAssetsByTag with limit 100 when isCloud is true', async () => { - mockIsCloud.value = true - try { - setActivePinia(createTestingPinia({ stubActions: false })) - const store = useAssetsStore() + describe('updateInputs', () => { + it('reads from assetService.getAssetsByTag with limit 100', async () => { + setActivePinia(createTestingPinia({ stubActions: false })) + const store = useAssetsStore() - vi.mocked(assetService.getAssetsByTag).mockResolvedValueOnce([]) - await store.updateInputs() + vi.mocked(assetService.getAssetsByTag).mockResolvedValueOnce([]) + await store.updateInputs() - expect(vi.mocked(assetService.getAssetsByTag)).toHaveBeenCalledWith( - 'input', - false, - { limit: 100 } - ) - expect( - assetService.invalidateInputAssetsIncludingPublic - ).toHaveBeenCalledOnce() - } finally { - mockIsCloud.value = false - } + expect(vi.mocked(assetService.getAssetsByTag)).toHaveBeenCalledWith( + 'input', + false, + { limit: 100 } + ) + expect( + assetService.invalidateInputAssetsIncludingPublic + ).toHaveBeenCalledOnce() }) }) }) diff --git a/src/stores/assetsStore.ts b/src/stores/assetsStore.ts index 645df3c8fc..0c391b4ace 100644 --- a/src/stores/assetsStore.ts +++ b/src/stores/assetsStore.ts @@ -2,10 +2,7 @@ import { useAsyncState, whenever } from '@vueuse/core' import { difference } from 'es-toolkit' import { defineStore } from 'pinia' import { computed, reactive, ref, shallowReactive } from 'vue' -import { - mapInputFileToAssetItem, - mapTaskOutputToAssetItem -} from '@/platform/assets/composables/media/assetMappers' +import { mapTaskOutputToAssetItem } from '@/platform/assets/composables/media/assetMappers' import type { AssetItem, TagsOperationResult @@ -16,7 +13,6 @@ import { assetService } from '@/platform/assets/services/assetService' import type { PaginationOptions } from '@/platform/assets/services/assetService' -import { isCloud } from '@/platform/distribution/types' import type { JobListItem } from '@/platform/remote/comfyui/jobs/jobTypes' import { api } from '@/scripts/api' @@ -26,30 +22,7 @@ import { useModelToNodeStore } from './modelToNodeStore' const INPUT_LIMIT = 100 -/** - * Fetch input files from the internal API (OSS version) - */ -async function fetchInputFilesFromAPI(): Promise { - const response = await fetch(api.internalURL('/files/input'), { - headers: { - 'Comfy-User': api.user - } - }) - - if (!response.ok) { - throw new Error('Failed to fetch input files') - } - - const filenames: string[] = await response.json() - return filenames.map((name, index) => - mapInputFileToAssetItem(name, index, 'input') - ) -} - -/** - * Fetch input files from cloud service - */ -async function fetchInputFilesFromCloud(): Promise { +async function fetchInputFiles(): Promise { return await assetService.getAssetsByTag(INPUT_TAG, false, { limit: INPUT_LIMIT }) @@ -123,10 +96,6 @@ export const useAssetsStore = defineStore('assets', () => { const loadedIds = shallowReactive(new Set()) - const fetchInputFiles = isCloud - ? fetchInputFilesFromCloud - : fetchInputFilesFromAPI - const { state: inputAssets, isLoading: inputLoading, @@ -385,418 +354,400 @@ export const useAssetsStore = defineStore('assets', () => { * Cloud-only feature - empty Maps in desktop builds */ const getModelState = () => { - if (isCloud) { - const modelStateByCategory = ref(new Map()) + const modelStateByCategory = ref(new Map()) - const assetsArrayCache = new Map< - string, - { source: Map; array: AssetItem[] } - >() + const assetsArrayCache = new Map< + string, + { source: Map; array: AssetItem[] } + >() - const pendingRequestByCategory = new Map() - const pendingPromiseByCategory = new Map>() + const pendingRequestByCategory = new Map() + const pendingPromiseByCategory = new Map>() - function createState( - existingAssets?: Map - ): ModelPaginationState { - const assets = new Map(existingAssets) - return reactive({ - assets, - offset: 0, - hasMore: true, - isLoading: true - }) + function createState( + existingAssets?: Map + ): ModelPaginationState { + const assets = new Map(existingAssets) + return reactive({ + assets, + offset: 0, + hasMore: true, + isLoading: true + }) + } + + function isStale(category: string, state: ModelPaginationState): boolean { + const committed = modelStateByCategory.value.get(category) + const pending = pendingRequestByCategory.get(category) + return committed !== state && pending !== state + } + + const EMPTY_ASSETS: AssetItem[] = [] + + /** + * Resolve a key to a category. Handles both nodeType and tag:xxx formats. + * @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models') + * @returns The category or undefined if not resolvable + */ + function resolveCategory(key: string): string | undefined { + if (key.startsWith('tag:')) { + return key + } + return modelToNodeStore.getCategoryForNodeType(key) + } + + /** + * Get assets by nodeType or tag key. + * Translates nodeType to category internally for cache lookup. + * @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models') + */ + function getAssets(key: string): AssetItem[] { + const category = resolveCategory(key) + if (!category) return EMPTY_ASSETS + + const state = modelStateByCategory.value.get(category) + const assetsMap = state?.assets + if (!assetsMap) return EMPTY_ASSETS + + const cached = assetsArrayCache.get(category) + if (cached && cached.source === assetsMap) { + return cached.array } - function isStale(category: string, state: ModelPaginationState): boolean { - const committed = modelStateByCategory.value.get(category) - const pending = pendingRequestByCategory.get(category) - return committed !== state && pending !== state + const array = Array.from(assetsMap.values()) + assetsArrayCache.set(category, { source: assetsMap, array }) + return array + } + + function isLoading(key: string): boolean { + const category = resolveCategory(key) + if (!category) return false + return modelStateByCategory.value.get(category)?.isLoading ?? false + } + + function getError(key: string): Error | undefined { + const category = resolveCategory(key) + if (!category) return undefined + return modelStateByCategory.value.get(category)?.error + } + + function hasMore(key: string): boolean { + const category = resolveCategory(key) + if (!category) return false + return modelStateByCategory.value.get(category)?.hasMore ?? false + } + + function hasAssetKey(key: string): boolean { + const category = resolveCategory(key) + if (!category) return false + return modelStateByCategory.value.has(category) + } + + /** + * Check if a category exists in the cache. + * Checks both direct category keys and tag-prefixed keys. + * @param category The category to check (e.g., 'checkpoints', 'loras') + */ + function hasCategory(category: string): boolean { + return ( + modelStateByCategory.value.has(category) || + modelStateByCategory.value.has(`tag:${category}`) + ) + } + + /** + * Internal helper to fetch and cache assets for a category. + * Loads first batch immediately, then progressively loads remaining batches. + * Keeps existing data visible until new data is successfully fetched. + * + * Concurrent calls for the same category are short-circuited: if a request + * is already in progress (tracked via pendingRequestByCategory), subsequent + * calls return immediately to avoid redundant work. + */ + async function updateModelsForCategory( + category: string, + fetcher: (options: PaginationOptions) => Promise + ): Promise { + if (pendingPromiseByCategory.has(category)) { + return pendingPromiseByCategory.get(category)! } - const EMPTY_ASSETS: AssetItem[] = [] + const existingState = modelStateByCategory.value.get(category) + const state = createState(existingState?.assets) - /** - * Resolve a key to a category. Handles both nodeType and tag:xxx formats. - * @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models') - * @returns The category or undefined if not resolvable - */ - function resolveCategory(key: string): string | undefined { - if (key.startsWith('tag:')) { - return key - } - return modelToNodeStore.getCategoryForNodeType(key) + const seenIds = new Set() + + const hasExistingData = modelStateByCategory.value.has(category) + if (hasExistingData) { + pendingRequestByCategory.set(category, state) + } else { + // Also track in pending map for initial loads to prevent concurrent calls + pendingRequestByCategory.set(category, state) + modelStateByCategory.value.set(category, state) } - /** - * Get assets by nodeType or tag key. - * Translates nodeType to category internally for cache lookup. - * @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models') - */ - function getAssets(key: string): AssetItem[] { - const category = resolveCategory(key) - if (!category) return EMPTY_ASSETS + async function loadBatches(): Promise { + while (state.hasMore) { + try { + const newAssets = await fetcher({ + limit: MODEL_BATCH_SIZE, + offset: state.offset + }) - const state = modelStateByCategory.value.get(category) - const assetsMap = state?.assets - if (!assetsMap) return EMPTY_ASSETS + if (isStale(category, state)) return - const cached = assetsArrayCache.get(category) - if (cached && cached.source === assetsMap) { - return cached.array - } - - const array = Array.from(assetsMap.values()) - assetsArrayCache.set(category, { source: assetsMap, array }) - return array - } - - function isLoading(key: string): boolean { - const category = resolveCategory(key) - if (!category) return false - return modelStateByCategory.value.get(category)?.isLoading ?? false - } - - function getError(key: string): Error | undefined { - const category = resolveCategory(key) - if (!category) return undefined - return modelStateByCategory.value.get(category)?.error - } - - function hasMore(key: string): boolean { - const category = resolveCategory(key) - if (!category) return false - return modelStateByCategory.value.get(category)?.hasMore ?? false - } - - function hasAssetKey(key: string): boolean { - const category = resolveCategory(key) - if (!category) return false - return modelStateByCategory.value.has(category) - } - - /** - * Check if a category exists in the cache. - * Checks both direct category keys and tag-prefixed keys. - * @param category The category to check (e.g., 'checkpoints', 'loras') - */ - function hasCategory(category: string): boolean { - return ( - modelStateByCategory.value.has(category) || - modelStateByCategory.value.has(`tag:${category}`) - ) - } - - /** - * Internal helper to fetch and cache assets for a category. - * Loads first batch immediately, then progressively loads remaining batches. - * Keeps existing data visible until new data is successfully fetched. - * - * Concurrent calls for the same category are short-circuited: if a request - * is already in progress (tracked via pendingRequestByCategory), subsequent - * calls return immediately to avoid redundant work. - */ - async function updateModelsForCategory( - category: string, - fetcher: (options: PaginationOptions) => Promise - ): Promise { - if (pendingPromiseByCategory.has(category)) { - return pendingPromiseByCategory.get(category)! - } - - const existingState = modelStateByCategory.value.get(category) - const state = createState(existingState?.assets) - - const seenIds = new Set() - - const hasExistingData = modelStateByCategory.value.has(category) - if (hasExistingData) { - pendingRequestByCategory.set(category, state) - } else { - // Also track in pending map for initial loads to prevent concurrent calls - pendingRequestByCategory.set(category, state) - modelStateByCategory.value.set(category, state) - } - - async function loadBatches(): Promise { - while (state.hasMore) { - try { - const newAssets = await fetcher({ - limit: MODEL_BATCH_SIZE, - offset: state.offset - }) - - if (isStale(category, state)) return - - const isFirstBatch = state.offset === 0 - if (isFirstBatch) { - assetsArrayCache.delete(category) - if (hasExistingData) { - pendingRequestByCategory.delete(category) - modelStateByCategory.value.set(category, state) - } + const isFirstBatch = state.offset === 0 + if (isFirstBatch) { + assetsArrayCache.delete(category) + if (hasExistingData) { + pendingRequestByCategory.delete(category) + modelStateByCategory.value.set(category, state) } - - // Merge new assets into existing map and track seen IDs - for (const asset of newAssets) { - seenIds.add(asset.id) - state.assets.set(asset.id, asset) - } - state.assets = new Map(state.assets) - - state.offset += newAssets.length - state.hasMore = newAssets.length === MODEL_BATCH_SIZE - - if (isFirstBatch) { - state.isLoading = false - } - - if (state.hasMore) { - await new Promise((resolve) => setTimeout(resolve, 50)) - } - } catch (err) { - if (isStale(category, state)) return - console.error(`Error loading batch for ${category}:`, err) - - state.error = err instanceof Error ? err : new Error(String(err)) - state.hasMore = false - state.isLoading = false - pendingRequestByCategory.delete(category) - - return } - } - const staleIds = [...state.assets.keys()].filter( - (id) => !seenIds.has(id) - ) - for (const id of staleIds) { - state.assets.delete(id) + // Merge new assets into existing map and track seen IDs + for (const asset of newAssets) { + seenIds.add(asset.id) + state.assets.set(asset.id, asset) + } + state.assets = new Map(state.assets) + + state.offset += newAssets.length + state.hasMore = newAssets.length === MODEL_BATCH_SIZE + + if (isFirstBatch) { + state.isLoading = false + } + + if (state.hasMore) { + await new Promise((resolve) => setTimeout(resolve, 50)) + } + } catch (err) { + if (isStale(category, state)) return + console.error(`Error loading batch for ${category}:`, err) + + state.error = err instanceof Error ? err : new Error(String(err)) + state.hasMore = false + state.isLoading = false + pendingRequestByCategory.delete(category) + + return } - assetsArrayCache.delete(category) - pendingRequestByCategory.delete(category) } - const promise = loadBatches().finally(() => { - pendingPromiseByCategory.delete(category) - }) - pendingPromiseByCategory.set(category, promise) - await promise - } - - /** - * Fetch and cache model assets for a specific node type. - * Translates nodeType to category internally - multiple node types - * sharing the same category will share the same cache entry. - * @param nodeType The node type to fetch assets for (e.g., 'CheckpointLoaderSimple') - */ - async function updateModelsForNodeType(nodeType: string): Promise { - const category = modelToNodeStore.getCategoryForNodeType(nodeType) - if (!category) return - - // Use category as cache key but fetch using nodeType for API compatibility - await updateModelsForCategory(category, (opts) => - assetService.getAssetsForNodeType(nodeType, opts) + const staleIds = [...state.assets.keys()].filter( + (id) => !seenIds.has(id) ) - } - - /** - * Fetch and cache model assets for a specific tag - * @param tag The tag to fetch assets for (e.g., 'models') - */ - async function updateModelsForTag(tag: string): Promise { - const category = `tag:${tag}` - await updateModelsForCategory(category, (opts) => - assetService.getAssetsByTag(tag, true, opts) - ) - } - - /** - * Invalidate the cache for a specific category. - * Forces a refetch on next access. - * @param category The category to invalidate (e.g., 'checkpoints', 'loras') - */ - function invalidateCategory(category: string): void { - modelStateByCategory.value.delete(category) + for (const id of staleIds) { + state.assets.delete(id) + } assetsArrayCache.delete(category) pendingRequestByCategory.delete(category) + } + + const promise = loadBatches().finally(() => { pendingPromiseByCategory.delete(category) - } + }) + pendingPromiseByCategory.set(category, promise) + await promise + } - /** - * Optimistically update an asset in the cache - * @param assetId The asset ID to update - * @param updates Partial asset data to merge - * @param cacheKey Optional cache key to target (nodeType or 'tag:xxx') - */ - function updateAssetInCache( - assetId: string, - updates: Partial, - cacheKey?: string - ) { - const category = cacheKey ? resolveCategory(cacheKey) : undefined - if (cacheKey && !category) return + /** + * Fetch and cache model assets for a specific node type. + * Translates nodeType to category internally - multiple node types + * sharing the same category will share the same cache entry. + * @param nodeType The node type to fetch assets for (e.g., 'CheckpointLoaderSimple') + */ + async function updateModelsForNodeType(nodeType: string): Promise { + const category = modelToNodeStore.getCategoryForNodeType(nodeType) + if (!category) return - const categoriesToCheck = category - ? [category] - : Array.from(modelStateByCategory.value.keys()) + // Use category as cache key but fetch using nodeType for API compatibility + await updateModelsForCategory(category, (opts) => + assetService.getAssetsForNodeType(nodeType, opts) + ) + } - for (const cat of categoriesToCheck) { - const state = modelStateByCategory.value.get(cat) - if (!state?.assets) continue + /** + * Fetch and cache model assets for a specific tag + * @param tag The tag to fetch assets for (e.g., 'models') + */ + async function updateModelsForTag(tag: string): Promise { + const category = `tag:${tag}` + await updateModelsForCategory(category, (opts) => + assetService.getAssetsByTag(tag, true, opts) + ) + } - const existingAsset = state.assets.get(assetId) - if (existingAsset) { - const updatedAsset = { ...existingAsset, ...updates } - state.assets.set(assetId, updatedAsset) - assetsArrayCache.delete(cat) - if (cacheKey) return - } + /** + * Invalidate the cache for a specific category. + * Forces a refetch on next access. + * @param category The category to invalidate (e.g., 'checkpoints', 'loras') + */ + function invalidateCategory(category: string): void { + modelStateByCategory.value.delete(category) + assetsArrayCache.delete(category) + pendingRequestByCategory.delete(category) + pendingPromiseByCategory.delete(category) + } + + /** + * Optimistically update an asset in the cache + * @param assetId The asset ID to update + * @param updates Partial asset data to merge + * @param cacheKey Optional cache key to target (nodeType or 'tag:xxx') + */ + function updateAssetInCache( + assetId: string, + updates: Partial, + cacheKey?: string + ) { + const category = cacheKey ? resolveCategory(cacheKey) : undefined + if (cacheKey && !category) return + + const categoriesToCheck = category + ? [category] + : Array.from(modelStateByCategory.value.keys()) + + for (const cat of categoriesToCheck) { + const state = modelStateByCategory.value.get(cat) + if (!state?.assets) continue + + const existingAsset = state.assets.get(assetId) + if (existingAsset) { + const updatedAsset = { ...existingAsset, ...updates } + state.assets.set(assetId, updatedAsset) + assetsArrayCache.delete(cat) + if (cacheKey) return } } - - /** - * Update asset metadata with optimistic cache update - * @param asset The asset to update - * @param userMetadata The user_metadata to save - * @param cacheKey Optional cache key to target for optimistic update - */ - async function updateAssetMetadata( - asset: AssetItem, - userMetadata: Record, - cacheKey?: string - ) { - const originalMetadata = asset.user_metadata - updateAssetInCache(asset.id, { user_metadata: userMetadata }, cacheKey) - - try { - const updatedAsset = await assetService.updateAsset(asset.id, { - user_metadata: userMetadata - }) - updateAssetInCache(asset.id, updatedAsset, cacheKey) - } catch (error) { - console.error('Failed to update asset metadata:', error) - updateAssetInCache( - asset.id, - { user_metadata: originalMetadata }, - cacheKey - ) - } - } - - /** - * Update asset tags using add/remove endpoints - * @param asset The asset to update (used to read current tags) - * @param newTags The desired tags array - * @param cacheKey Optional cache key to target for optimistic update - */ - async function updateAssetTags( - asset: AssetItem, - newTags: string[], - cacheKey?: string - ) { - const originalTags = asset.tags - const tagsToAdd = difference(newTags, originalTags) - const tagsToRemove = difference(originalTags, newTags) - - if (tagsToAdd.length === 0 && tagsToRemove.length === 0) return - - updateAssetInCache(asset.id, { tags: newTags }, cacheKey) - - let removedTagsOnServer: string[] = [] - try { - let removeResult: TagsOperationResult | undefined - if (tagsToRemove.length > 0) { - removeResult = await assetService.removeAssetTags( - asset.id, - tagsToRemove - ) - removedTagsOnServer = removeResult.removed ?? tagsToRemove - } - - const addResult = - tagsToAdd.length > 0 - ? await assetService.addAssetTags(asset.id, tagsToAdd) - : undefined - - const finalTags = (addResult ?? removeResult)?.total_tags - if (finalTags) { - updateAssetInCache(asset.id, { tags: finalTags }, cacheKey) - } - } catch (error) { - console.error('Failed to update asset tags:', error) - updateAssetInCache(asset.id, { tags: originalTags }, cacheKey) - - if (removedTagsOnServer.length > 0) { - try { - await assetService.addAssetTags(asset.id, removedTagsOnServer) - } catch (compensationError) { - console.error( - 'Failed to restore tags after partial failure; invalidating cache to force refetch:', - compensationError - ) - const categoriesToInvalidate = new Set() - const resolved = cacheKey ? resolveCategory(cacheKey) : undefined - if (resolved) { - categoriesToInvalidate.add(resolved) - } - for (const [ - category, - state - ] of modelStateByCategory.value.entries()) { - if (state.assets?.has(asset.id)) { - categoriesToInvalidate.add(category) - } - } - for (const category of categoriesToInvalidate) { - invalidateCategory(category) - } - } - } - } - } - - /** - * Invalidate model caches for a given category (e.g., 'checkpoints', 'loras') - * Clears the category cache and tag-based caches so next access triggers refetch - * @param category The model category to invalidate (e.g., 'checkpoints') - */ - function invalidateModelsForCategory(category: string): void { - invalidateCategory(category) - invalidateCategory(`tag:${category}`) - invalidateCategory('tag:models') - } - - return { - getAssets, - isLoading, - getError, - hasMore, - hasAssetKey, - hasCategory, - updateModelsForNodeType, - updateModelsForTag, - invalidateCategory, - updateAssetMetadata, - updateAssetTags, - invalidateModelsForCategory - } } - const emptyAssets: AssetItem[] = [] + /** + * Update asset metadata with optimistic cache update + * @param asset The asset to update + * @param userMetadata The user_metadata to save + * @param cacheKey Optional cache key to target for optimistic update + */ + async function updateAssetMetadata( + asset: AssetItem, + userMetadata: Record, + cacheKey?: string + ) { + const originalMetadata = asset.user_metadata + updateAssetInCache(asset.id, { user_metadata: userMetadata }, cacheKey) + + try { + const updatedAsset = await assetService.updateAsset(asset.id, { + user_metadata: userMetadata + }) + updateAssetInCache(asset.id, updatedAsset, cacheKey) + } catch (error) { + console.error('Failed to update asset metadata:', error) + updateAssetInCache( + asset.id, + { user_metadata: originalMetadata }, + cacheKey + ) + } + } + + /** + * Update asset tags using add/remove endpoints + * @param asset The asset to update (used to read current tags) + * @param newTags The desired tags array + * @param cacheKey Optional cache key to target for optimistic update + */ + async function updateAssetTags( + asset: AssetItem, + newTags: string[], + cacheKey?: string + ) { + const originalTags = asset.tags + const tagsToAdd = difference(newTags, originalTags) + const tagsToRemove = difference(originalTags, newTags) + + if (tagsToAdd.length === 0 && tagsToRemove.length === 0) return + + updateAssetInCache(asset.id, { tags: newTags }, cacheKey) + + let removedTagsOnServer: string[] = [] + try { + let removeResult: TagsOperationResult | undefined + if (tagsToRemove.length > 0) { + removeResult = await assetService.removeAssetTags( + asset.id, + tagsToRemove + ) + removedTagsOnServer = removeResult.removed ?? tagsToRemove + } + + const addResult = + tagsToAdd.length > 0 + ? await assetService.addAssetTags(asset.id, tagsToAdd) + : undefined + + const finalTags = (addResult ?? removeResult)?.total_tags + if (finalTags) { + updateAssetInCache(asset.id, { tags: finalTags }, cacheKey) + } + } catch (error) { + console.error('Failed to update asset tags:', error) + updateAssetInCache(asset.id, { tags: originalTags }, cacheKey) + + if (removedTagsOnServer.length > 0) { + try { + await assetService.addAssetTags(asset.id, removedTagsOnServer) + } catch (compensationError) { + console.error( + 'Failed to restore tags after partial failure; invalidating cache to force refetch:', + compensationError + ) + const categoriesToInvalidate = new Set() + const resolved = cacheKey ? resolveCategory(cacheKey) : undefined + if (resolved) { + categoriesToInvalidate.add(resolved) + } + for (const [ + category, + state + ] of modelStateByCategory.value.entries()) { + if (state.assets?.has(asset.id)) { + categoriesToInvalidate.add(category) + } + } + for (const category of categoriesToInvalidate) { + invalidateCategory(category) + } + } + } + } + } + + /** + * Invalidate model caches for a given category (e.g., 'checkpoints', 'loras') + * Clears the category cache and tag-based caches so next access triggers refetch + * @param category The model category to invalidate (e.g., 'checkpoints') + */ + function invalidateModelsForCategory(category: string): void { + invalidateCategory(category) + invalidateCategory(`tag:${category}`) + invalidateCategory('tag:models') + } + return { - getAssets: () => emptyAssets, - isLoading: () => false, - getError: () => undefined, - hasMore: () => false, - hasAssetKey: () => false, - hasCategory: () => false, - updateModelsForNodeType: async () => {}, - invalidateCategory: () => {}, - updateModelsForTag: async () => {}, - updateAssetMetadata: async () => {}, - updateAssetTags: async () => {}, - invalidateModelsForCategory: () => {} + getAssets, + isLoading, + getError, + hasMore, + hasAssetKey, + hasCategory, + updateModelsForNodeType, + updateModelsForTag, + invalidateCategory, + updateAssetMetadata, + updateAssetTags, + invalidateModelsForCategory } }