diff --git a/src/platform/assets/components/AssetBrowserModal.vue b/src/platform/assets/components/AssetBrowserModal.vue index a45f886ac..cda99726f 100644 --- a/src/platform/assets/components/AssetBrowserModal.vue +++ b/src/platform/assets/components/AssetBrowserModal.vue @@ -69,7 +69,7 @@ const { availableCategories, contentTitle, filteredAssets, - selectAsset + selectAssetWithCallback } = useAssetBrowser(props.assets) // Dialog controls panel visibility via prop @@ -84,13 +84,10 @@ const handleClose = () => { } // Handle asset selection and emit to parent -const handleAssetSelectAndEmit = (asset: AssetDisplayItem) => { - selectAsset(asset) // This logs the selection for dev mode +const handleAssetSelectAndEmit = async (asset: AssetDisplayItem) => { emit('asset-select', asset) // Emit the full asset object - // Call prop callback if provided - if (props.onSelect) { - props.onSelect(asset.name) // Use asset name as the asset path - } + // Use composable for detail fetching and callback execution + await selectAssetWithCallback(asset.id, props.onSelect) } diff --git a/src/platform/assets/composables/useAssetBrowser.ts b/src/platform/assets/composables/useAssetBrowser.ts index 964c81fc7..1cb305fa0 100644 --- a/src/platform/assets/composables/useAssetBrowser.ts +++ b/src/platform/assets/composables/useAssetBrowser.ts @@ -3,6 +3,7 @@ import { computed, ref } from 'vue' import { d, t } from '@/i18n' import type { UUID } from '@/lib/litegraph/src/utils/uuid' import type { AssetItem } from '@/platform/assets/schemas/assetSchema' +import { assetService } from '@/platform/assets/services/assetService' import { getAssetBaseModel, getAssetDescription @@ -162,12 +163,51 @@ export function useAssetBrowser(assets: AssetItem[] = []) { return filtered.map(transformAssetForDisplay) }) - // Actions - function selectAsset(asset: AssetDisplayItem): UUID { + /** + * Asset selection that fetches full details and executes callback with filename + * @param assetId - The asset ID to select and fetch details for + * @param onSelect - Optional callback to execute with the asset filename + */ + async function selectAssetWithCallback( + assetId: string, + onSelect?: (filename: string) => void + ): Promise { + // Always log selection for debugging if (import.meta.env.DEV) { - console.log('Asset selected:', asset.id, asset.name) + console.log('Asset selected:', assetId) + } + + // If no callback provided, just return (no need to fetch details) + if (!onSelect) { + return + } + + try { + // Fetch complete asset details to get user_metadata + const detailAsset = await assetService.getAssetDetails(assetId) + + // Extract filename from user_metadata + const filename = detailAsset.user_metadata?.filename + + // Validate filename exists and is not empty + if (!filename || typeof filename !== 'string' || filename.trim() === '') { + console.error( + 'Invalid asset filename from user_metadata:', + filename || null, + 'for asset:', + assetId + ) + return + } + + // Execute callback with validated filename + onSelect(filename) + } catch (error) { + console.error( + `Failed to fetch asset details for ${assetId}:`, + error + ) } - return asset.id } return { @@ -182,7 +222,6 @@ export function useAssetBrowser(assets: AssetItem[] = []) { filteredAssets, // Actions - selectAsset, - transformAssetForDisplay + selectAssetWithCallback } } diff --git a/src/platform/assets/composables/useAssetBrowserDialog.ts b/src/platform/assets/composables/useAssetBrowserDialog.ts index e5f63eead..327844e05 100644 --- a/src/platform/assets/composables/useAssetBrowserDialog.ts +++ b/src/platform/assets/composables/useAssetBrowserDialog.ts @@ -1,4 +1,6 @@ import AssetBrowserModal from '@/platform/assets/components/AssetBrowserModal.vue' +import type { AssetItem } from '@/platform/assets/schemas/assetSchema' +import { assetService } from '@/platform/assets/services/assetService' import { useDialogStore } from '@/stores/dialogStore' interface AssetBrowserDialogProps { @@ -8,8 +10,11 @@ interface AssetBrowserDialogProps { inputName: string /** Current selected asset value */ currentValue?: string - /** Callback for when an asset is selected */ - onAssetSelected?: (assetPath: string) => void + /** + * Callback for when an asset is selected + * @param {string} filename - The validated filename from user_metadata.filename + */ + onAssetSelected?: (filename: string) => void } export const useAssetBrowserDialog = () => { @@ -20,7 +25,7 @@ export const useAssetBrowserDialog = () => { dialogStore.closeDialog({ key: dialogKey }) } - function show(props: AssetBrowserDialogProps) { + async function show(props: AssetBrowserDialogProps) { const handleAssetSelected = (assetPath: string) => { props.onAssetSelected?.(assetPath) hide() // Auto-close on selection @@ -48,6 +53,14 @@ export const useAssetBrowserDialog = () => { } } + // Fetch assets for the specific node type, fallback to empty array on error + let assets: AssetItem[] = [] + try { + assets = await assetService.getAssetsForNodeType(props.nodeType) + } catch (error) { + console.error('Failed to fetch assets for node type:', props.nodeType, error) + } + dialogStore.showDialog({ key: dialogKey, component: AssetBrowserModal, @@ -55,6 +68,7 @@ export const useAssetBrowserDialog = () => { nodeType: props.nodeType, inputName: props.inputName, currentValue: props.currentValue, + assets, onSelect: handleAssetSelected, onClose: handleClose }, diff --git a/src/platform/assets/schemas/assetSchema.ts b/src/platform/assets/schemas/assetSchema.ts index fab41649a..dcf1db19f 100644 --- a/src/platform/assets/schemas/assetSchema.ts +++ b/src/platform/assets/schemas/assetSchema.ts @@ -6,11 +6,11 @@ const zAsset = z.object({ name: z.string(), asset_hash: z.string(), size: z.number(), - mime_type: z.string(), + mime_type: z.string().nullable(), tags: z.array(z.string()), preview_url: z.string().optional(), created_at: z.string(), - updated_at: z.string(), + updated_at: z.string().optional(), last_access_time: z.string(), user_metadata: z.record(z.unknown()).optional(), // API allows arbitrary key-value pairs preview_id: z.string().nullable().optional() diff --git a/src/platform/assets/services/assetService.ts b/src/platform/assets/services/assetService.ts index 74b20a753..ad1125538 100644 --- a/src/platform/assets/services/assetService.ts +++ b/src/platform/assets/services/assetService.ts @@ -1,6 +1,7 @@ import { fromZodError } from 'zod-validation-error' import { + type AssetItem, type AssetResponse, type ModelFile, type ModelFolder, @@ -127,10 +128,74 @@ function createAssetService() { ) } + /** + * Gets assets for a specific node type by finding the matching category + * and fetching all assets with that category tag + * + * @param nodeType - The ComfyUI node type (e.g., 'CheckpointLoaderSimple') + * @returns Promise - Full asset objects with preserved metadata + */ + async function getAssetsForNodeType(nodeType: string): Promise { + if (!nodeType || typeof nodeType !== 'string') { + return [] + } + + // Find the category for this node type by reverse lookup in modelToNodeMap + const modelToNodeStore = useModelToNodeStore() + const modelToNodeMap = modelToNodeStore.modelToNodeMap + + const category = Object.keys(modelToNodeMap).find(categoryKey => + modelToNodeMap[categoryKey].some(provider => provider.nodeDef.name === nodeType) + ) + + if (!category) { + return [] + } + + // Fetch assets for this category using same API pattern as getAssetModels + const data = await handleAssetRequest( + `${ASSETS_ENDPOINT}?include_tags=${MODELS_TAG},${category}`, + `assets for ${nodeType}` + ) + + // Return full AssetItem[] objects (don't strip like getAssetModels does) + return data?.assets?.filter(asset => + !asset.tags.includes(MISSING_TAG) && asset.tags.includes(category) + ) ?? [] + } + + /** + * Gets complete details for a specific asset by ID + * Calls the detail endpoint which includes user_metadata and all fields + * + * @param id - The asset ID + * @returns Promise - Complete asset object with user_metadata + */ + async function getAssetDetails(id: string): Promise { + const res = await api.fetchApi(`${ASSETS_ENDPOINT}/${id}`) + if (!res.ok) { + throw new Error( + `Unable to load asset details for ${id}: Server returned ${res.status}. Please try again.` + ) + } + const data = await res.json() + + // Validate the single asset response against our schema + const result = assetResponseSchema.safeParse({ assets: [data] }) + if (result.success && result.data.assets?.[0]) { + return result.data.assets[0] + } + + const error = fromZodError(result.error) + throw new Error(`Invalid asset response against zod schema:\n${error}`) + } + return { getAssetModelFolders, getAssetModels, - isAssetBrowserEligible + isAssetBrowserEligible, + getAssetsForNodeType, + getAssetDetails } } diff --git a/src/platform/assets/utils/assetMetadataUtils.ts b/src/platform/assets/utils/assetMetadataUtils.ts index 2d32fa07f..a6ed821e0 100644 --- a/src/platform/assets/utils/assetMetadataUtils.ts +++ b/src/platform/assets/utils/assetMetadataUtils.ts @@ -25,3 +25,38 @@ export function getAssetBaseModel(asset: AssetItem): string | null { ? asset.user_metadata.base_model : null } + +/** + * Safely extracts the ComfyUI-relative filename from user_metadata. + * @param {import('../schemas/assetSchema').AssetItem} asset - The asset item containing user_metadata + * @returns {string | null} ComfyUI-relative path or null if not available + */ +export function getAssetFilename(asset: AssetItem): string | null { + const filename = asset.user_metadata?.filename + + if (typeof filename !== 'string' || !filename.trim()) { + return null + } + + return filename.trim() +} + +/** + * Validates if a filename path is safe for ComfyUI widget usage. + * @param {string} filename - The filename to validate + * @returns {boolean} True if filename is safe for widget usage + */ +export function validateAssetFilename(filename: string): boolean { + if (!filename || typeof filename !== 'string') return false + + const trimmed = filename.trim() + if (!trimmed) return false + + // Reject dangerous patterns but allow forward slashes for subdirectories + // e.g., reject "../../../etc/passwd" but allow "checkpoints/model.safetensors" + if (trimmed.includes('..') || /[<>:"|?*]/.test(trimmed)) { + return false + } + + return true +} diff --git a/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts b/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts index 01c3f3184..6b0d69405 100644 --- a/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts +++ b/src/renderer/extensions/vueNodes/widgets/composables/useComboWidget.ts @@ -4,9 +4,11 @@ import MultiSelectWidget from '@/components/graph/widgets/MultiSelectWidget.vue' import { t } from '@/i18n' import type { LGraphNode } from '@/lib/litegraph/src/litegraph' import type { + IAssetWidget, IBaseWidget, IComboWidget } from '@/lib/litegraph/src/types/widgets' +import { useAssetBrowserDialog } from '@/platform/assets/composables/useAssetBrowserDialog' import { assetService } from '@/platform/assets/services/assetService' import { useSettingStore } from '@/platform/settings/settingStore' import { transformInputSpecV2ToV1 } from '@/schemas/nodeDef/migration' @@ -73,10 +75,21 @@ const addComboWidget = ( const currentValue = getDefaultValue(inputSpec) const displayLabel = currentValue ?? t('widgets.selectModel') - const widget = node.addWidget('asset', inputSpec.name, displayLabel, () => { - console.log( - `Asset Browser would open here for:\nNode: ${node.type}\nWidget: ${inputSpec.name}\nCurrent Value:${currentValue}` - ) + const assetBrowserDialog = useAssetBrowserDialog() + + const widget = node.addWidget('asset', inputSpec.name, displayLabel, async () => { + const assetWidget = widget as IAssetWidget + await assetBrowserDialog.show({ + nodeType: node.comfyClass || '', + inputName: inputSpec.name, + currentValue: assetWidget.value, + onAssetSelected: (filename: string) => { + assetWidget.value = filename + // Must call widget.callback to notify litegraph of value changes + // This ensures proper serialization and triggers any downstream effects + assetWidget.callback?.(assetWidget.value) + } + }) }) return widget diff --git a/tests-ui/platform/assets/components/AssetBrowserModal.test.ts b/tests-ui/platform/assets/components/AssetBrowserModal.test.ts index 4e96dd902..ffb630318 100644 --- a/tests-ui/platform/assets/components/AssetBrowserModal.test.ts +++ b/tests-ui/platform/assets/components/AssetBrowserModal.test.ts @@ -6,6 +6,14 @@ import { nextTick } from 'vue' import AssetBrowserModal from '@/platform/assets/components/AssetBrowserModal.vue' import type { AssetDisplayItem } from '@/platform/assets/composables/useAssetBrowser' import type { AssetItem } from '@/platform/assets/schemas/assetSchema' +import { assetService } from '@/platform/assets/services/assetService' + +// Mock assetService +vi.mock('@/platform/assets/services/assetService', () => ({ + assetService: { + getAssetDetails: vi.fn() + } +})) // Mock external dependencies with minimal functionality needed for business logic tests vi.mock('@/components/input/SearchBox.vue', () => ({ @@ -92,9 +100,18 @@ vi.mock('@/platform/assets/components/AssetGrid.vue', () => ({ vi.mock('vue-i18n', () => ({ useI18n: () => ({ t: (key: string) => key - }) + }), + createI18n: vi.fn() })) +vi.mock('@/i18n', () => ({ + t: (key: string) => key, + d: (date: Date, options?: any) => date.toLocaleDateString() +})) + +// Mock console.error for error handling tests +const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + describe('AssetBrowserModal', () => { const createTestAsset = ( id: string, @@ -255,18 +272,9 @@ describe('AssetBrowserModal', () => { expect(emittedAsset.id).toBe('asset1') }) - it('executes onSelect callback when provided', async () => { - const onSelectSpy = vi.fn() - const assets = [createTestAsset('asset1', 'Test Model', 'checkpoints')] - const wrapper = createWrapper(assets, { onSelect: onSelectSpy }) - - // Click on first asset - await wrapper.find('[data-testid="asset-asset1"]').trigger('click') - - expect(onSelectSpy).toHaveBeenCalledWith('Test Model') - }) }) + describe('Left Panel Conditional Logic', () => { it('hides left panel by default when showLeftPanel prop is undefined', () => { const singleCategoryAssets = [ diff --git a/tests-ui/platform/assets/composables/useAssetBrowser.test.ts b/tests-ui/platform/assets/composables/useAssetBrowser.test.ts index d7d4f74dc..a5d784145 100644 --- a/tests-ui/platform/assets/composables/useAssetBrowser.test.ts +++ b/tests-ui/platform/assets/composables/useAssetBrowser.test.ts @@ -1,10 +1,33 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi, beforeEach } from 'vitest' import { nextTick } from 'vue' import { useAssetBrowser } from '@/platform/assets/composables/useAssetBrowser' import type { AssetItem } from '@/platform/assets/schemas/assetSchema' +import { assetService } from '@/platform/assets/services/assetService' + +vi.mock('@/platform/assets/services/assetService', () => ({ + assetService: { + getAssetDetails: vi.fn() + } +})) + +vi.mock('@/i18n', () => ({ + t: (key: string) => { + const translations: Record = { + 'assetBrowser.allModels': 'All Models', + 'assetBrowser.assets': 'Assets', + 'assetBrowser.unknown': 'unknown' + } + return translations[key] || key + }, + d: (date: Date, options?: any) => date.toLocaleDateString() +})) describe('useAssetBrowser', () => { + beforeEach(() => { + vi.resetAllMocks() + }) + // Test fixtures - minimal data focused on functionality being tested const createApiAsset = (overrides: Partial = {}): AssetItem => ({ id: 'test-id', @@ -26,8 +49,8 @@ describe('useAssetBrowser', () => { user_metadata: { description: 'Test model' } }) - const { transformAssetForDisplay } = useAssetBrowser([apiAsset]) - const result = transformAssetForDisplay(apiAsset) + const { filteredAssets } = useAssetBrowser([apiAsset]) + const result = filteredAssets.value[0] // Get the transformed asset from filteredAssets // Preserves API properties expect(result.id).toBe(apiAsset.id) @@ -49,15 +72,13 @@ describe('useAssetBrowser', () => { user_metadata: undefined }) - const { transformAssetForDisplay } = useAssetBrowser([apiAsset]) - const result = transformAssetForDisplay(apiAsset) + const { filteredAssets } = useAssetBrowser([apiAsset]) + const result = filteredAssets.value[0] expect(result.description).toBe('loras model') }) it('formats various file sizes correctly', () => { - const { transformAssetForDisplay } = useAssetBrowser([]) - const testCases = [ { size: 512, expected: '512 B' }, { size: 1536, expected: '1.5 KB' }, @@ -67,7 +88,8 @@ describe('useAssetBrowser', () => { testCases.forEach(({ size, expected }) => { const asset = createApiAsset({ size }) - const result = transformAssetForDisplay(asset) + const { filteredAssets } = useAssetBrowser([asset]) + const result = filteredAssets.value[0] expect(result.formattedSize).toBe(expected) }) }) @@ -236,18 +258,87 @@ describe('useAssetBrowser', () => { }) }) - describe('Asset Selection', () => { - it('returns selected asset UUID for efficient handling', () => { + + describe('Async Asset Selection with Detail Fetching', () => { + it('should fetch asset details and call onSelect with filename when provided', async () => { + const onSelectSpy = vi.fn() const asset = createApiAsset({ - id: 'test-uuid-123', - name: 'selected_model.safetensors' + id: 'asset-123', + name: 'test-model.safetensors' }) - const { selectAsset, transformAssetForDisplay } = useAssetBrowser([asset]) - const displayAsset = transformAssetForDisplay(asset) - const result = selectAsset(displayAsset) + const detailAsset = createApiAsset({ + id: 'asset-123', + name: 'test-model.safetensors', + user_metadata: { filename: 'checkpoints/test-model.safetensors' } + }) + vi.mocked(assetService.getAssetDetails).mockResolvedValue(detailAsset) - expect(result).toBe('test-uuid-123') + const { selectAssetWithCallback } = useAssetBrowser([asset]) + + await selectAssetWithCallback(asset.id, onSelectSpy) + + expect(assetService.getAssetDetails).toHaveBeenCalledWith('asset-123') + expect(onSelectSpy).toHaveBeenCalledWith('checkpoints/test-model.safetensors') + }) + + it('should handle missing user_metadata.filename as error', async () => { + const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + const onSelectSpy = vi.fn() + const asset = createApiAsset({ id: 'asset-456' }) + + const detailAsset = createApiAsset({ + id: 'asset-456', + user_metadata: { filename: '' } // Invalid empty filename + }) + vi.mocked(assetService.getAssetDetails).mockResolvedValue(detailAsset) + + const { selectAssetWithCallback } = useAssetBrowser([asset]) + + await selectAssetWithCallback(asset.id, onSelectSpy) + + expect(assetService.getAssetDetails).toHaveBeenCalledWith('asset-456') + expect(onSelectSpy).not.toHaveBeenCalled() + expect(consoleErrorSpy).toHaveBeenCalledWith( + 'Invalid asset filename from user_metadata:', + null, + 'for asset:', + 'asset-456' + ) + + consoleErrorSpy.mockRestore() + }) + + it('should handle API errors gracefully', async () => { + const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + const onSelectSpy = vi.fn() + const asset = createApiAsset({ id: 'asset-789' }) + + const apiError = new Error('API Error') + vi.mocked(assetService.getAssetDetails).mockRejectedValue(apiError) + + const { selectAssetWithCallback } = useAssetBrowser([asset]) + + await selectAssetWithCallback(asset.id, onSelectSpy) + + expect(assetService.getAssetDetails).toHaveBeenCalledWith('asset-789') + expect(onSelectSpy).not.toHaveBeenCalled() + expect(consoleErrorSpy).toHaveBeenCalledWith( + expect.stringContaining('Failed to fetch asset details for asset-789'), + apiError + ) + + consoleErrorSpy.mockRestore() + }) + + it('should not fetch details when no callback provided', async () => { + const asset = createApiAsset({ id: 'asset-no-callback' }) + + const { selectAssetWithCallback } = useAssetBrowser([asset]) + + await selectAssetWithCallback(asset.id) + + expect(assetService.getAssetDetails).not.toHaveBeenCalled() }) }) diff --git a/tests-ui/platform/assets/composables/useAssetBrowserDialog.test.ts b/tests-ui/platform/assets/composables/useAssetBrowserDialog.test.ts index fefeeceac..06bbbd71a 100644 --- a/tests-ui/platform/assets/composables/useAssetBrowserDialog.test.ts +++ b/tests-ui/platform/assets/composables/useAssetBrowserDialog.test.ts @@ -10,7 +10,7 @@ vi.mock('@/stores/dialogStore') interface AssetBrowserProps { nodeType: string inputName: string - onAssetSelected?: ReturnType + onAssetSelected?: (filename: string) => void } function createAssetBrowserProps( @@ -85,5 +85,6 @@ describe('useAssetBrowserDialog', () => { key: 'global-asset-browser' }) }) + }) }) diff --git a/tests-ui/tests/platform/assets/utils/assetMetadataUtils.test.ts b/tests-ui/tests/platform/assets/utils/assetMetadataUtils.test.ts index 54551f595..a52abdbf7 100644 --- a/tests-ui/tests/platform/assets/utils/assetMetadataUtils.test.ts +++ b/tests-ui/tests/platform/assets/utils/assetMetadataUtils.test.ts @@ -3,7 +3,9 @@ import { describe, expect, it } from 'vitest' import type { AssetItem } from '@/platform/assets/schemas/assetSchema' import { getAssetBaseModel, - getAssetDescription + getAssetDescription, + getAssetFilename, + validateAssetFilename } from '@/platform/assets/utils/assetMetadataUtils' describe('assetMetadataUtils', () => { @@ -62,4 +64,50 @@ describe('assetMetadataUtils', () => { expect(getAssetBaseModel(mockAsset)).toBeNull() }) }) + + describe('getAssetFilename', () => { + it('should return trimmed filename when present', () => { + const asset = { + ...mockAsset, + user_metadata: { filename: ' checkpoints/model.safetensors ' } + } + expect(getAssetFilename(asset)).toBe('checkpoints/model.safetensors') + }) + + it('should return null when filename is empty string', () => { + const asset = { + ...mockAsset, + user_metadata: { filename: ' ' } + } + expect(getAssetFilename(asset)).toBeNull() + }) + + it('should return null when no metadata', () => { + expect(getAssetFilename(mockAsset)).toBeNull() + }) + }) + + describe('validateAssetFilename', () => { + it('should accept valid filenames', () => { + expect(validateAssetFilename('model.safetensors')).toBe(true) + expect(validateAssetFilename('checkpoints/model.safetensors')).toBe(true) + expect(validateAssetFilename('loras/style/anime.safetensors')).toBe(true) + }) + + it('should reject directory traversal attempts', () => { + expect(validateAssetFilename('../../../etc/passwd')).toBe(false) + expect(validateAssetFilename('models/../../../secret.txt')).toBe(false) + }) + + it('should reject dangerous characters', () => { + expect(validateAssetFilename('model