Comfy-Org/ComfyUI_frontend
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI_frontend.git synced 2026-04-29 02:32:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: extract shouldUseAssetBrowser(), fix missing inputNameForBrowser, sanitize logs (#8867)

Browse Source 
## Summary

Extract duplicated asset-browser eligibility guard into
`shouldUseAssetBrowser()`, fix a missing parameter bug, and sanitize a
log statement.

## Changes

- **What**: 
- DRY: Extract the repeated 3-condition guard (`isCloud &&
isUsingAssetAPI && isAssetBrowserEligible`) into
`assetService.shouldUseAssetBrowser()`, used by `widgetInputs.ts` and
`useComboWidget.ts`
- Bug fix: `createAssetBrowserWidget()` in `useComboWidget.ts` was
missing the `inputNameForBrowser` parameter, which could show wrong
assets for nodes with multiple model inputs
- Security: `createAssetWidget.ts` no longer logs the full raw asset
object on validation failure
- `WidgetSelect.vue` keeps an inline guard because it has a special
`widget.type === "asset"` fallback that must stay gated behind
`isUsingAssetAPI`

## Review Focus

- `shouldUseAssetBrowser()` correctly combines the three conditions that
were previously duplicated
- `WidgetSelect.vue` preserves exact behavioral equivalence (a widget
can have `type === "asset"` when the setting is off if a user toggles
the setting after creating asset widgets)

Fixes #8744

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-8867-refactor-extract-shouldUseAssetBrowser-fix-missing-inputNameForBrowser-sanitize-logs-3076d73d3650818cabdcd76a351dac31)
by [Unito](https://www.unito.io)
This commit is contained in:
Christian Byrne
2026-02-20 02:23:29 -08:00
committed by GitHub
parent 0792d26f77
commit c809ac5a43
9 changed files with 146 additions and 503 deletions
Download Patch File Download Diff File Expand all files Collapse all files

465
src/platform/assets/services/assetService.test.ts
View File

@@ -1,423 +1,106 @@
import { beforeEach, describe, expect, it, vi } from 'vitest'
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
import { assetService } from '@/platform/assets/services/assetService'
import { api } from '@/scripts/api'
vi.mock('@/scripts/api', () => ({
api: {
fetchApi: vi.fn(),
addEventListener: vi.fn(),
removeEventListener: vi.fn()
const mockDistributionState = vi.hoisted(() => ({ isCloud: false }))
const mockSettingStoreGet = vi.hoisted(() => vi.fn(() => false))
vi.mock('@/platform/distribution/types', () => ({
get isCloud() {
return mockDistributionState.isCloud
}
}))
const mockGetCategoryForNodeType = vi.fn()
vi.mock('@/stores/modelToNodeStore', () => ({
useModelToNodeStore: vi.fn(() => ({
getRegisteredNodeTypes: vi.fn(() => ({
CheckpointLoaderSimple: 'ckpt_name',
LoraLoader: 'lora_name',
VAELoader: 'vae_name',
TestNode: ''
})),
getCategoryForNodeType: mockGetCategoryForNodeType,
modelToNodeMap: {
checkpoints: [{ nodeDef: { name: 'CheckpointLoaderSimple' } }],
loras: [{ nodeDef: { name: 'LoraLoader' } }],
vae: [{ nodeDef: { name: 'VAELoader' } }]
}
vi.mock('@/platform/settings/settingStore', () => ({
useSettingStore: vi.fn(() => ({
get: mockSettingStoreGet
}))
}))
// Helper to create API-compliant test assets
function createTestAsset(overrides: Partial<AssetItem> = {}) {
vi.mock('@/stores/modelToNodeStore', () => {
const registeredNodeTypes: Record<string, string> = {
CheckpointLoaderSimple: 'ckpt_name',
LoraLoader: 'lora_name'
}
return {
id: 'test-uuid',
name: 'test-model.safetensors',
asset_hash: 'blake3:test123',
size: 123456,
mime_type: 'application/octet-stream',
tags: ['models', 'checkpoints'],
created_at: '2024-01-01T00:00:00Z',
updated_at: '2024-01-01T00:00:00Z',
last_access_time: '2024-01-01T00:00:00Z',
...overrides
useModelToNodeStore: vi.fn(() => ({
getRegisteredNodeTypes: () => registeredNodeTypes,
getCategoryForNodeType: vi.fn()
}))
}
}
})
// Test data constants
const MOCK_ASSETS = {
checkpoints: createTestAsset({
id: 'uuid-1',
name: 'model1.safetensors',
tags: ['models', 'checkpoints']
}),
loras: createTestAsset({
id: 'uuid-2',
name: 'model2.safetensors',
tags: ['models', 'loras']
}),
vae: createTestAsset({
id: 'uuid-3',
name: 'vae1.safetensors',
tags: ['models', 'vae']
})
} as const
// Helper functions
function mockApiResponse(assets: unknown[], options = {}) {
const response = {
assets,
total: assets.length,
has_more: false,
...options
vi.mock('@/scripts/api', () => ({
api: {
fetchApi: vi.fn()
}
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 })
)
}
vi.mock('@/i18n', () => ({
st: vi.fn((_key: string, fallback: string) => fallback)
}))
describe('assetService', () => {
describe(assetService.shouldUseAssetBrowser, () => {
beforeEach(() => {
vi.resetAllMocks()
vi.spyOn(api, 'fetchApi')
vi.clearAllMocks()
mockDistributionState.isCloud = false
mockSettingStoreGet.mockReturnValue(false)
})
describe('getAssetModelFolders', () => {
it('should extract directory names from asset tags and filter blacklisted ones', async () => {
const assets = [
createTestAsset({
id: 'uuid-1',
name: 'checkpoint1.safetensors',
tags: ['models', 'checkpoints']
}),
createTestAsset({
id: 'uuid-2',
name: 'config.yaml',
tags: ['models', 'configs'] // Blacklisted
}),
createTestAsset({
id: 'uuid-3',
name: 'vae1.safetensors',
tags: ['models', 'vae']
})
]
mockApiResponse(assets)
it('returns false when not on cloud', () => {
mockDistributionState.isCloud = false
mockSettingStoreGet.mockReturnValue(true)
const result = await assetService.getAssetModelFolders()
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=models&limit=500'
)
expect(result).toHaveLength(2)
const folderNames = result.map((f) => f.name)
expect(folderNames).toEqual(['checkpoints', 'vae'])
expect(folderNames).not.toContain('configs')
})
it('should handle empty responses', async () => {
mockApiResponse([])
const emptyResult = await assetService.getAssetModelFolders()
expect(emptyResult).toHaveLength(0)
})
it('should handle network errors', async () => {
vi.mocked(api.fetchApi).mockRejectedValueOnce(new Error('Network error'))
await expect(assetService.getAssetModelFolders()).rejects.toThrow(
'Network error'
)
})
it('should handle HTTP errors', async () => {
mockApiError(500)
await expect(assetService.getAssetModelFolders()).rejects.toThrow(
'Unable to load model folders: Server returned 500. Please try again.'
)
})
expect(
assetService.shouldUseAssetBrowser('CheckpointLoaderSimple', 'ckpt_name')
).toBe(false)
})
describe('getAssetModels', () => {
it('should return filtered models for folder', async () => {
const assets = [
{ ...MOCK_ASSETS.checkpoints, name: 'valid.safetensors' },
{ ...MOCK_ASSETS.loras, name: 'lora.safetensors' }, // Wrong tag
createTestAsset({
id: 'uuid-4',
name: 'missing-model.safetensors',
tags: ['models', 'checkpoints', 'missing'] // Has missing tag
})
]
mockApiResponse(assets)
it('returns false when asset API setting is disabled', () => {
mockDistributionState.isCloud = true
mockSettingStoreGet.mockReturnValue(false)
const result = await assetService.getAssetModels('checkpoints')
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=models%2Ccheckpoints&limit=500'
)
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.'
)
})
expect(
assetService.shouldUseAssetBrowser('CheckpointLoaderSimple', 'ckpt_name')
).toBe(false)
})
describe('isAssetBrowserEligible', () => {
it.for<[string, string, boolean, string]>([
['CheckpointLoaderSimple', 'ckpt_name', true, 'valid inputs'],
['LoraLoader', 'lora_name', true, 'valid inputs'],
['VAELoader', 'vae_name', true, 'valid inputs'],
['CheckpointLoaderSimple', 'type', false, 'other combo widgets'],
['UnknownNode', 'widget', false, 'unregistered types'],
['NotRegistered', 'widget', false, 'unregistered types']
])(
'isAssetBrowserEligible("%s", "%s") should return %s for %s',
([type, name, expected]) => {
expect(assetService.isAssetBrowserEligible(type, name)).toBe(expected)
}
it('returns false when node type is not eligible', () => {
mockDistributionState.isCloud = true
mockSettingStoreGet.mockReturnValue(true)
expect(
assetService.shouldUseAssetBrowser('UnknownNode', 'some_input')
).toBe(false)
})
it('returns true when cloud, setting enabled, and node is eligible', () => {
mockDistributionState.isCloud = true
mockSettingStoreGet.mockReturnValue(true)
expect(
assetService.shouldUseAssetBrowser('CheckpointLoaderSimple', 'ckpt_name')
).toBe(true)
})
it('returns false when nodeType is undefined', () => {
mockDistributionState.isCloud = true
mockSettingStoreGet.mockReturnValue(true)
expect(assetService.shouldUseAssetBrowser(undefined, 'ckpt_name')).toBe(
false
)
})
describe('getAssetsForNodeType', () => {
beforeEach(() => {
mockGetCategoryForNodeType.mockClear()
})
it('returns false when widget name does not match registered input', () => {
mockDistributionState.isCloud = true
mockSettingStoreGet.mockReturnValue(true)
it('should return empty array for unregistered node types', async () => {
mockGetCategoryForNodeType.mockReturnValue(undefined)
const result = await assetService.getAssetsForNodeType('UnknownNode')
expect(mockGetCategoryForNodeType).toHaveBeenCalledWith('UnknownNode')
expect(result).toEqual([])
})
it('should use getCategoryForNodeType for efficient category lookup', async () => {
mockGetCategoryForNodeType.mockReturnValue('checkpoints')
const testAssets = [MOCK_ASSETS.checkpoints]
mockApiResponse(testAssets)
const result = await assetService.getAssetsForNodeType(
'CheckpointLoaderSimple'
expect(
assetService.shouldUseAssetBrowser(
'CheckpointLoaderSimple',
'wrong_input'
)
expect(mockGetCategoryForNodeType).toHaveBeenCalledWith(
'CheckpointLoaderSimple'
)
expect(result).toEqual(testAssets)
// Verify API call includes correct category (comma is URL-encoded by URLSearchParams)
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=models%2Ccheckpoints&limit=500'
)
})
it('should return empty array when no category found', async () => {
mockGetCategoryForNodeType.mockReturnValue(undefined)
const result = await assetService.getAssetsForNodeType('TestNode')
expect(result).toEqual([])
expect(api.fetchApi).not.toHaveBeenCalled()
})
it('should handle API errors gracefully', async () => {
mockGetCategoryForNodeType.mockReturnValue('loras')
mockApiError(500, 'Internal Server Error')
await expect(
assetService.getAssetsForNodeType('LoraLoader')
).rejects.toThrow(
'Unable to load assets for LoraLoader: Server returned 500. Please try again.'
)
})
it('should return all assets without filtering for different categories', async () => {
// Test checkpoints
mockGetCategoryForNodeType.mockReturnValue('checkpoints')
const checkpointAssets = [MOCK_ASSETS.checkpoints]
mockApiResponse(checkpointAssets)
let result = await assetService.getAssetsForNodeType(
'CheckpointLoaderSimple'
)
expect(result).toEqual(checkpointAssets)
// Test loras
mockGetCategoryForNodeType.mockReturnValue('loras')
const loraAssets = [MOCK_ASSETS.loras]
mockApiResponse(loraAssets)
result = await assetService.getAssetsForNodeType('LoraLoader')
expect(result).toEqual(loraAssets)
// Test vae
mockGetCategoryForNodeType.mockReturnValue('vae')
const vaeAssets = [MOCK_ASSETS.vae]
mockApiResponse(vaeAssets)
result = await assetService.getAssetsForNodeType('VAELoader')
expect(result).toEqual(vaeAssets)
})
})
describe('getAssetsByTag', () => {
it('should fetch assets with correct tag query parameter', async () => {
const testAssets = [MOCK_ASSETS.checkpoints, MOCK_ASSETS.loras]
mockApiResponse(testAssets)
const result = await assetService.getAssetsByTag('models')
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=models&limit=500&include_public=true'
)
expect(result).toEqual(testAssets)
})
it('should filter out assets with missing tag', async () => {
const testAssets = [
MOCK_ASSETS.checkpoints,
createTestAsset({
id: 'uuid-missing',
name: 'missing.safetensors',
tags: ['models', 'checkpoints', 'missing']
}),
MOCK_ASSETS.loras
]
mockApiResponse(testAssets)
const result = await assetService.getAssetsByTag('models')
expect(result).toHaveLength(2)
expect(result).toEqual([MOCK_ASSETS.checkpoints, MOCK_ASSETS.loras])
expect(result.some((a) => a.id === 'uuid-missing')).toBe(false)
})
it('should return empty array on API error', async () => {
mockApiError(500)
await expect(assetService.getAssetsByTag('models')).rejects.toThrow(
'Unable to load assets for tag models: Server returned 500. Please try again.'
)
})
it('should return empty array for empty response', async () => {
mockApiResponse([])
const result = await assetService.getAssetsByTag('nonexistent')
expect(result).toEqual([])
})
it('should return AssetItem[] with full metadata', async () => {
const fullAsset = createTestAsset({
id: 'test-full',
name: 'full-model.safetensors',
asset_hash: 'blake3:full123',
size: 999999,
tags: ['models', 'checkpoints'],
user_metadata: { filename: 'models/checkpoints/full-model.safetensors' }
})
mockApiResponse([fullAsset])
const result = await assetService.getAssetsByTag('models')
expect(result).toHaveLength(1)
expect(result[0]).toEqual(fullAsset)
expect(result[0]).toHaveProperty('asset_hash', 'blake3:full123')
expect(result[0]).toHaveProperty('user_metadata')
})
it('should exclude public assets when includePublic is false', async () => {
const testAssets = [MOCK_ASSETS.checkpoints]
mockApiResponse(testAssets)
const result = await assetService.getAssetsByTag('input', false)
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=input&limit=500&include_public=false'
)
expect(result).toEqual(testAssets)
})
it('should include public assets when includePublic is true', async () => {
const testAssets = [MOCK_ASSETS.checkpoints, MOCK_ASSETS.loras]
mockApiResponse(testAssets)
const result = await assetService.getAssetsByTag('models', true)
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=models&limit=500&include_public=true'
)
expect(result).toEqual(testAssets)
})
it('should accept custom limit via options', async () => {
const testAssets = [MOCK_ASSETS.checkpoints]
mockApiResponse(testAssets)
const result = await assetService.getAssetsByTag('input', false, {
limit: 100
})
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=input&limit=100&include_public=false'
)
expect(result).toEqual(testAssets)
})
it('should accept custom offset via options', async () => {
const testAssets = [MOCK_ASSETS.loras]
mockApiResponse(testAssets)
const result = await assetService.getAssetsByTag('models', true, {
offset: 50
})
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=models&limit=500&offset=50&include_public=true'
)
expect(result).toEqual(testAssets)
})
it('should accept both limit and offset via options', async () => {
const testAssets = [MOCK_ASSETS.checkpoints]
mockApiResponse(testAssets)
const result = await assetService.getAssetsByTag('input', false, {
limit: 100,
offset: 25
})
expect(api.fetchApi).toHaveBeenCalledWith(
'/assets?include_tags=input&limit=100&offset=25&include_public=false'
)
expect(result).toEqual(testAssets)
})
).toBe(false)
})
})

27
src/platform/assets/services/assetService.ts
View File

@@ -18,6 +18,8 @@ import type {
ModelFolder,
TagsOperationResult
} from '@/platform/assets/schemas/assetSchema'
import { isCloud } from '@/platform/distribution/types'
import { useSettingStore } from '@/platform/settings/settingStore'
import { api } from '@/scripts/api'
import { useModelToNodeStore } from '@/stores/modelToNodeStore'
@@ -298,6 +300,29 @@ function createAssetService() {
)
}
/**
* Checks if the asset API is enabled (cloud environment + user setting).
*/
function isAssetAPIEnabled(): boolean {
if (!isCloud) return false
return !!useSettingStore().get('Comfy.Assets.UseAssetAPI')
}
/**
* Checks if the asset browser should be used for a given node input.
* Combines the cloud environment check, user setting, and eligibility check.
*
* @param nodeType - The ComfyUI node comfyClass
* @param widgetName - The name of the widget to check
* @returns true if this input should use the asset browser
*/
function shouldUseAssetBrowser(
nodeType: string | undefined,
widgetName: string
): boolean {
return isAssetAPIEnabled() && isAssetBrowserEligible(nodeType, widgetName)
}
/**
* Gets assets for a specific node type by finding the matching category
* and fetching all assets with that category tag
@@ -732,7 +757,9 @@ function createAssetService() {
return {
getAssetModelFolders,
getAssetModels,
isAssetAPIEnabled,
isAssetBrowserEligible,
shouldUseAssetBrowser,
getAssetsForNodeType,
getAssetDetails,
getAssetsByTag,

6
src/platform/assets/utils/createAssetWidget.ts
View File

@@ -1,3 +1,5 @@
import { fromZodError } from 'zod-validation-error'
import { t } from '@/i18n'
import type { LGraphNode } from '@/lib/litegraph/src/litegraph'
import type {
@@ -66,9 +68,7 @@ export function createAssetWidget(
if (!validatedAsset.success) {
console.error(
'Invalid asset item:',
validatedAsset.error.errors,
'Received:',
asset
fromZodError(validatedAsset.error).message
)
toastStore.add({
severity: 'error',