diff --git a/src/constants/coreSettings.ts b/src/constants/coreSettings.ts index 8f023abbe..3d7859f3b 100644 --- a/src/constants/coreSettings.ts +++ b/src/constants/coreSettings.ts @@ -972,5 +972,14 @@ export const CORE_SETTINGS: SettingParams[] = [ defaultValue: false, experimental: true, versionAdded: '1.27.1' + }, + { + id: 'Comfy.Assets.UseAssetAPI', + name: 'Use Asset API for model library', + type: 'boolean', + tooltip: + 'Use new asset API instead of experiment endpoints for model browsing', + defaultValue: false, + experimental: true } ] diff --git a/src/schemas/apiSchema.ts b/src/schemas/apiSchema.ts index b4283d231..3def57252 100644 --- a/src/schemas/apiSchema.ts +++ b/src/schemas/apiSchema.ts @@ -467,6 +467,7 @@ const zSettings = z.object({ 'Comfy.Minimap.RenderErrorState': z.boolean(), 'Comfy.Canvas.NavigationMode': z.string(), 'Comfy.VueNodes.Enabled': z.boolean(), + 'Comfy.Assets.UseAssetAPI': z.boolean(), 'Comfy-Desktop.AutoUpdate': z.boolean(), 'Comfy-Desktop.SendStatistics': z.boolean(), 'Comfy-Desktop.WindowStyle': z.string(), diff --git a/src/scripts/api.ts b/src/scripts/api.ts index 857626674..9a6a0b7d3 100644 --- a/src/scripts/api.ts +++ b/src/scripts/api.ts @@ -682,7 +682,8 @@ export class ComfyApi extends EventTarget { } const folderBlacklist = ['configs', 'custom_nodes'] return (await res.json()).filter( - (folder: string) => !folderBlacklist.includes(folder) + (folder: { name: string; folders: string[] }) => + !folderBlacklist.includes(folder.name) ) } @@ -1019,7 +1020,13 @@ export class ComfyApi extends EventTarget { } async getFolderPaths(): Promise> { - return (await axios.get(this.internalURL('/folder_paths'))).data + const response = await axios + .get(this.internalURL('/folder_paths')) + .catch(() => null) + if (!response) { + return {} // Fallback: no filesystem paths known when API unavailable + } + return response.data } /* Frees memory by unloading models and optionally freeing execution cache diff --git a/src/services/assetService.ts b/src/services/assetService.ts new file mode 100644 index 000000000..90a4d8320 --- /dev/null +++ b/src/services/assetService.ts @@ -0,0 +1,151 @@ +import { api } from '@/scripts/api' + +const ASSETS_ENDPOINT = '/assets' +const MODELS_TAG = 'models' +const MISSING_TAG = 'missing' + +// Types for asset API responses +interface AssetResponse { + assets?: Asset[] + total?: number + has_more?: boolean +} + +interface Asset { + id: string + name: string + tags: string[] + size: number + created_at?: string +} + +/** + * Type guard for validating asset structure + */ +function isValidAsset(asset: unknown): asset is Asset { + return ( + asset !== null && + typeof asset === 'object' && + 'id' in asset && + 'name' in asset && + 'tags' in asset && + Array.isArray((asset as Asset).tags) + ) +} + +/** + * Creates predicate for filtering assets by folder and excluding missing ones + */ +function createAssetFolderFilter(folder?: string) { + return (asset: unknown): asset is Asset => { + if (!isValidAsset(asset) || asset.tags.includes(MISSING_TAG)) { + return false + } + if (folder && !asset.tags.includes(folder)) { + return false + } + return true + } +} + +/** + * Creates predicate for filtering folder assets (requires name) + */ +function createFolderAssetFilter(folder: string) { + return (asset: unknown): asset is Asset => { + if (!isValidAsset(asset) || !asset.name) { + return false + } + return asset.tags.includes(folder) && !asset.tags.includes(MISSING_TAG) + } +} + +/** + * Private service for asset-related network requests + * Not exposed globally - used internally by ComfyApi + */ +function createAssetService() { + /** + * Handles API response with consistent error handling + */ + async function handleAssetRequest( + url: string, + context: string + ): Promise { + const res = await api.fetchApi(url) + if (!res.ok) { + throw new Error( + `Unable to load ${context}: Server returned ${res.status}. Please try again.` + ) + } + return await res.json() + } + /** + * Gets a list of model folder keys from the asset API + * + * Logic: + * 1. Extract directory names directly from asset tags + * 2. Filter out blacklisted directories + * 3. Return alphabetically sorted directories with assets + * + * @returns The list of model folder keys + */ + async function getAssetModelFolders(): Promise< + { name: string; folders: string[] }[] + > { + const data = await handleAssetRequest( + `${ASSETS_ENDPOINT}?include_tags=${MODELS_TAG}`, + 'model folders' + ) + + // Blacklist directories we don't want to show + const blacklistedDirectories = ['configs'] + + // Extract directory names from assets that actually exist, exclude missing assets + const discoveredFolders = new Set() + if (data?.assets) { + const directoryTags = data.assets + .filter(createAssetFolderFilter()) + .flatMap((asset) => asset.tags) + .filter( + (tag) => tag !== MODELS_TAG && !blacklistedDirectories.includes(tag) + ) + + for (const tag of directoryTags) { + discoveredFolders.add(tag) + } + } + + // Return only discovered folders in alphabetical order + const sortedFolders = Array.from(discoveredFolders).sort() + return sortedFolders.map((name) => ({ name, folders: [] })) + } + + /** + * Gets a list of models in the specified folder from the asset API + * @param folder The folder to list models from, such as 'checkpoints' + * @returns The list of model filenames within the specified folder + */ + async function getAssetModels( + folder: string + ): Promise<{ name: string; pathIndex: number }[]> { + const data = await handleAssetRequest( + `${ASSETS_ENDPOINT}?include_tags=${MODELS_TAG},${folder}`, + `models for ${folder}` + ) + + return data?.assets + ? data.assets.filter(createFolderAssetFilter(folder)).map((asset) => ({ + name: asset.name, + pathIndex: 0 + })) + : [] + } + + return { + getAssetModelFolders, + getAssetModels + } +} + +export const assetService = createAssetService() diff --git a/src/stores/modelStore.ts b/src/stores/modelStore.ts index dc03c6145..f8dc5441d 100644 --- a/src/stores/modelStore.ts +++ b/src/stores/modelStore.ts @@ -2,6 +2,8 @@ import { defineStore } from 'pinia' import { computed, ref } from 'vue' import { api } from '@/scripts/api' +import { assetService } from '@/services/assetService' +import { useSettingStore } from '@/stores/settingStore' /** (Internal helper) finds a value in a metadata object from any of a list of keys. */ function _findInMetadata(metadata: any, ...keys: string[]): string | null { @@ -153,7 +155,12 @@ export class ModelFolder { models: Record = {} state: ResourceState = ResourceState.Uninitialized - constructor(public directory: string) {} + constructor( + public directory: string, + private getModelsFunc: ( + folder: string + ) => Promise<{ name: string; pathIndex: number }[]> + ) {} get key(): string { return this.directory + '/' @@ -167,7 +174,7 @@ export class ModelFolder { return this } this.state = ResourceState.Loading - const models = await api.getModels(this.directory) + const models = await this.getModelsFunc(this.directory) for (const model of models) { this.models[`${model.pathIndex}/${model.name}`] = new ComfyModelDef( model.name, @@ -182,6 +189,7 @@ export class ModelFolder { /** Model store handler, wraps individual per-folder model stores */ export const useModelStore = defineStore('models', () => { + const settingStore = useSettingStore() const modelFolderNames = ref([]) const modelFolderByName = ref>({}) const modelFolders = computed(() => @@ -197,11 +205,22 @@ export const useModelStore = defineStore('models', () => { * Loads the model folders from the server */ async function loadModelFolders() { - const resData = await api.getModelFolders() + const useAssetAPI: boolean = settingStore.get('Comfy.Assets.UseAssetAPI') + + const resData = useAssetAPI + ? await assetService.getAssetModelFolders() + : await api.getModelFolders() modelFolderNames.value = resData.map((folder) => folder.name) modelFolderByName.value = {} for (const folderName of modelFolderNames.value) { - modelFolderByName.value[folderName] = new ModelFolder(folderName) + const getModelsFunc = useAssetAPI + ? (folder: string) => assetService.getAssetModels(folder) + : (folder: string) => api.getModels(folder) + + modelFolderByName.value[folderName] = new ModelFolder( + folderName, + getModelsFunc + ) } } diff --git a/tests-ui/tests/api.folderPaths.test.ts b/tests-ui/tests/api.folderPaths.test.ts new file mode 100644 index 000000000..2b0bf7fd1 --- /dev/null +++ b/tests-ui/tests/api.folderPaths.test.ts @@ -0,0 +1,30 @@ +import axios from 'axios' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { api } from '@/scripts/api' + +vi.mock('axios') + +describe('getFolderPaths', () => { + beforeEach(() => { + vi.resetAllMocks() + }) + + it('returns legacy API response when available', async () => { + const mockResponse = { checkpoints: ['/test/checkpoints'] } + vi.mocked(axios.get).mockResolvedValueOnce({ data: mockResponse }) + + const result = await api.getFolderPaths() + + expect(result).toEqual(mockResponse) + }) + + it('returns empty object when legacy API unavailable (dynamic discovery)', async () => { + vi.mocked(axios.get).mockRejectedValueOnce(new Error()) + + const result = await api.getFolderPaths() + + // With dynamic discovery, we don't pre-generate directories when API is unavailable + expect(result).toEqual({}) + }) +}) diff --git a/tests-ui/tests/services/assetService.test.ts b/tests-ui/tests/services/assetService.test.ts new file mode 100644 index 000000000..d7c4a673b --- /dev/null +++ b/tests-ui/tests/services/assetService.test.ts @@ -0,0 +1,150 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { api } from '@/scripts/api' +import { assetService } from '@/services/assetService' + +// Test data constants +const MOCK_ASSETS = { + checkpoints: { + id: 'uuid-1', + name: 'model1.safetensors', + tags: ['models', 'checkpoints'], + size: 123456 + }, + loras: { + id: 'uuid-2', + name: 'model2.safetensors', + tags: ['models', 'loras'], + size: 654321 + }, + vae: { + id: 'uuid-3', + name: 'vae1.safetensors', + tags: ['models', 'vae'], + size: 789012 + } +} as const + +// Helper functions +function mockApiResponse(assets: any[], options = {}) { + const response = { + assets, + total: assets.length, + has_more: false, + ...options + } + vi.mocked(api.fetchApi).mockResolvedValueOnce(Response.json(response)) + return response +} + +function mockApiError(status: number, statusText = 'Error') { + vi.mocked(api.fetchApi).mockResolvedValueOnce( + new Response(null, { status, statusText }) + ) +} + +describe('assetService', () => { + beforeEach(() => { + vi.resetAllMocks() + vi.spyOn(api, 'fetchApi') + }) + + describe('getAssetModelFolders', () => { + it('should extract directory names from asset tags and filter blacklisted ones', async () => { + const assets = [ + { + id: 'uuid-1', + name: 'checkpoint1.safetensors', + tags: ['models', 'checkpoints'], + size: 123456 + }, + { + id: 'uuid-2', + name: 'config.yaml', + tags: ['models', 'configs'], // Blacklisted + size: 654321 + }, + { + id: 'uuid-3', + name: 'vae1.safetensors', + tags: ['models', 'vae'], + size: 789012 + } + ] + mockApiResponse(assets) + + const result = await assetService.getAssetModelFolders() + + expect(api.fetchApi).toHaveBeenCalledWith('/assets?include_tags=models') + expect(result).toHaveLength(2) + + const folderNames = result.map((f) => f.name) + expect(folderNames).toEqual(['checkpoints', 'vae']) + expect(folderNames).not.toContain('configs') + }) + + it('should handle errors and empty responses', async () => { + // Empty response + mockApiResponse([]) + const emptyResult = await assetService.getAssetModelFolders() + expect(emptyResult).toHaveLength(0) + + // Network error + vi.mocked(api.fetchApi).mockRejectedValueOnce(new Error('Network error')) + await expect(assetService.getAssetModelFolders()).rejects.toThrow( + 'Network error' + ) + + // HTTP error + mockApiError(500) + await expect(assetService.getAssetModelFolders()).rejects.toThrow( + 'Unable to load model folders: Server returned 500. Please try again.' + ) + }) + }) + + describe('getAssetModels', () => { + it('should return filtered models for folder', async () => { + const assets = [ + { ...MOCK_ASSETS.checkpoints, name: 'valid.safetensors' }, + { ...MOCK_ASSETS.checkpoints, name: undefined }, // Invalid name + { ...MOCK_ASSETS.loras, name: 'lora.safetensors' }, // Wrong tag + { + id: 'uuid-4', + name: 'missing-model.safetensors', + tags: ['models', 'checkpoints', 'missing'], // Has missing tag + size: 654321 + } + ] + mockApiResponse(assets) + + const result = await assetService.getAssetModels('checkpoints') + + expect(api.fetchApi).toHaveBeenCalledWith( + '/assets?include_tags=models,checkpoints' + ) + expect(result).toEqual([ + expect.objectContaining({ name: 'valid.safetensors', pathIndex: 0 }) + ]) + }) + + it('should handle errors and empty responses', async () => { + // Empty response + mockApiResponse([]) + const emptyResult = await assetService.getAssetModels('nonexistent') + expect(emptyResult).toEqual([]) + + // Network error + vi.mocked(api.fetchApi).mockRejectedValueOnce(new Error('Network error')) + await expect(assetService.getAssetModels('checkpoints')).rejects.toThrow( + 'Network error' + ) + + // HTTP error + mockApiError(404) + await expect(assetService.getAssetModels('checkpoints')).rejects.toThrow( + 'Unable to load models for checkpoints: Server returned 404. Please try again.' + ) + }) + }) +}) diff --git a/tests-ui/tests/store/modelStore.test.ts b/tests-ui/tests/store/modelStore.test.ts index 6964b04e8..26fa8be65 100644 --- a/tests-ui/tests/store/modelStore.test.ts +++ b/tests-ui/tests/store/modelStore.test.ts @@ -2,7 +2,9 @@ import { createPinia, setActivePinia } from 'pinia' import { beforeEach, describe, expect, it, vi } from 'vitest' import { api } from '@/scripts/api' +import { assetService } from '@/services/assetService' import { useModelStore } from '@/stores/modelStore' +import { useSettingStore } from '@/stores/settingStore' // Mock the api vi.mock('@/scripts/api', () => ({ @@ -13,7 +15,32 @@ vi.mock('@/scripts/api', () => ({ } })) -function enableMocks() { +// Mock the assetService +vi.mock('@/services/assetService', () => ({ + assetService: { + getAssetModelFolders: vi.fn(), + getAssetModels: vi.fn() + } +})) + +// Mock the settingStore +vi.mock('@/stores/settingStore', () => ({ + useSettingStore: vi.fn() +})) + +function enableMocks(useAssetAPI = false) { + // Mock settingStore to return the useAssetAPI setting + const mockSettingStore = { + get: vi.fn().mockImplementation((key: string) => { + if (key === 'Comfy.Assets.UseAssetAPI') { + return useAssetAPI + } + return false + }) + } + vi.mocked(useSettingStore).mockReturnValue(mockSettingStore as any) + + // Mock experimental API - returns objects with name and folders properties vi.mocked(api.getModels).mockResolvedValue([ { name: 'sdxl.safetensors', pathIndex: 0 }, { name: 'sdv15.safetensors', pathIndex: 0 }, @@ -23,6 +50,18 @@ function enableMocks() { { name: 'checkpoints', folders: ['/path/to/checkpoints'] }, { name: 'vae', folders: ['/path/to/vae'] } ]) + + // Mock asset API - also returns objects with name and folders properties + vi.mocked(assetService.getAssetModelFolders).mockResolvedValue([ + { name: 'checkpoints', folders: ['/path/to/checkpoints'] }, + { name: 'vae', folders: ['/path/to/vae'] } + ]) + vi.mocked(assetService.getAssetModels).mockResolvedValue([ + { name: 'sdxl.safetensors', pathIndex: 0 }, + { name: 'sdv15.safetensors', pathIndex: 0 }, + { name: 'noinfo.safetensors', pathIndex: 0 } + ]) + vi.mocked(api.viewMetadata).mockImplementation((_, model) => { if (model === 'noinfo.safetensors') { return Promise.resolve({}) @@ -46,26 +85,25 @@ describe('useModelStore', () => { beforeEach(async () => { setActivePinia(createPinia()) - store = useModelStore() vi.resetAllMocks() }) it('should load models', async () => { enableMocks() + store = useModelStore() await store.loadModelFolders() const folderStore = await store.getLoadedModelFolder('checkpoints') - expect(folderStore).not.toBeNull() - if (!folderStore) return - expect(Object.keys(folderStore.models).length).toBe(3) + expect(folderStore).toBeDefined() + expect(Object.keys(folderStore!.models)).toHaveLength(3) }) it('should load model metadata', async () => { enableMocks() + store = useModelStore() await store.loadModelFolders() const folderStore = await store.getLoadedModelFolder('checkpoints') - expect(folderStore).not.toBeNull() - if (!folderStore) return - const model = folderStore.models['0/sdxl.safetensors'] + expect(folderStore).toBeDefined() + const model = folderStore!.models['0/sdxl.safetensors'] await model.load() expect(model.title).toBe('Title of sdxl.safetensors') expect(model.architecture_id).toBe('stable-diffusion-xl-base-v1') @@ -79,11 +117,11 @@ describe('useModelStore', () => { it('should handle no metadata', async () => { enableMocks() + store = useModelStore() await store.loadModelFolders() const folderStore = await store.getLoadedModelFolder('checkpoints') - expect(folderStore).not.toBeNull() - if (!folderStore) return - const model = folderStore.models['0/noinfo.safetensors'] + expect(folderStore).toBeDefined() + const model = folderStore!.models['0/noinfo.safetensors'] await model.load() expect(model.file_name).toBe('noinfo.safetensors') expect(model.title).toBe('noinfo') @@ -95,6 +133,7 @@ describe('useModelStore', () => { it('should cache model information', async () => { enableMocks() + store = useModelStore() await store.loadModelFolders() expect(api.getModels).toHaveBeenCalledTimes(0) await store.getLoadedModelFolder('checkpoints') @@ -102,4 +141,36 @@ describe('useModelStore', () => { await store.getLoadedModelFolder('checkpoints') expect(api.getModels).toHaveBeenCalledTimes(1) }) + + describe('API switching functionality', () => { + it('should use experimental API for complete workflow when UseAssetAPI setting is false', async () => { + enableMocks(false) // useAssetAPI = false + store = useModelStore() + await store.loadModelFolders() + const folderStore = await store.getLoadedModelFolder('checkpoints') + + // Both APIs return objects with .name property, modelStore extracts folder.name in both cases + expect(api.getModelFolders).toHaveBeenCalledTimes(1) + expect(api.getModels).toHaveBeenCalledWith('checkpoints') + expect(assetService.getAssetModelFolders).toHaveBeenCalledTimes(0) + expect(assetService.getAssetModels).toHaveBeenCalledTimes(0) + expect(folderStore).toBeDefined() + expect(Object.keys(folderStore!.models)).toHaveLength(3) + }) + + it('should use asset API for complete workflow when UseAssetAPI setting is true', async () => { + enableMocks(true) // useAssetAPI = true + store = useModelStore() + await store.loadModelFolders() + const folderStore = await store.getLoadedModelFolder('checkpoints') + + // Both APIs return objects with .name property, modelStore extracts folder.name in both cases + expect(assetService.getAssetModelFolders).toHaveBeenCalledTimes(1) + expect(assetService.getAssetModels).toHaveBeenCalledWith('checkpoints') + expect(api.getModelFolders).toHaveBeenCalledTimes(0) + expect(api.getModels).toHaveBeenCalledTimes(0) + expect(folderStore).toBeDefined() + expect(Object.keys(folderStore!.models)).toHaveLength(3) + }) + }) })