diff --git a/src/platform/assets/composables/useOutputStacks.test.ts b/src/platform/assets/composables/useOutputStacks.test.ts new file mode 100644 index 000000000..d770b3c75 --- /dev/null +++ b/src/platform/assets/composables/useOutputStacks.test.ts @@ -0,0 +1,200 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { ref } from 'vue' + +import type { AssetItem } from '@/platform/assets/schemas/assetSchema' +import { useOutputStacks } from '@/platform/assets/composables/useOutputStacks' + +const mocks = vi.hoisted(() => ({ + resolveOutputAssetItems: vi.fn() +})) + +vi.mock('@/platform/assets/utils/outputAssetUtil', () => ({ + resolveOutputAssetItems: mocks.resolveOutputAssetItems +})) + +type Deferred = { + promise: Promise + resolve: (value: T) => void + reject: (reason?: unknown) => void +} + +function createDeferred(): Deferred { + let resolve!: (value: T) => void + let reject!: (reason?: unknown) => void + const promise = new Promise((resolveFn, rejectFn) => { + resolve = resolveFn + reject = rejectFn + }) + return { promise, resolve, reject } +} + +function createAsset(overrides: Partial = {}): AssetItem { + return { + id: 'asset-1', + name: 'parent.png', + tags: [], + created_at: '2025-01-01T00:00:00.000Z', + user_metadata: { + promptId: 'prompt-1', + nodeId: 'node-1', + subfolder: 'outputs' + }, + ...overrides + } +} + +describe('useOutputStacks', () => { + beforeEach(() => { + vi.resetAllMocks() + }) + + it('expands stacks and exposes children as selectable assets', async () => { + const parent = createAsset({ id: 'parent', name: 'parent.png' }) + const childA = createAsset({ + id: 'child-a', + name: 'child-a.png', + user_metadata: undefined + }) + const childB = createAsset({ + id: 'child-b', + name: 'child-b.png', + user_metadata: undefined + }) + + vi.mocked(mocks.resolveOutputAssetItems).mockResolvedValue([childA, childB]) + + const { assetItems, isStackExpanded, selectableAssets, toggleStack } = + useOutputStacks({ assets: ref([parent]) }) + + await toggleStack(parent) + + expect(mocks.resolveOutputAssetItems).toHaveBeenCalledWith( + expect.objectContaining({ promptId: 'prompt-1' }), + { + createdAt: parent.created_at, + excludeOutputKey: parent.name + } + ) + expect(isStackExpanded(parent)).toBe(true) + expect(assetItems.value.map((item) => item.asset.id)).toEqual([ + parent.id, + childA.id, + childB.id + ]) + expect(assetItems.value[1]).toMatchObject({ + asset: childA, + isChild: true + }) + expect(assetItems.value[2]).toMatchObject({ + asset: childB, + isChild: true + }) + expect(selectableAssets.value).toEqual([parent, childA, childB]) + }) + + it('collapses an expanded stack when toggled again', async () => { + const parent = createAsset({ id: 'parent', name: 'parent.png' }) + const child = createAsset({ + id: 'child', + name: 'child.png', + user_metadata: undefined + }) + + vi.mocked(mocks.resolveOutputAssetItems).mockResolvedValue([child]) + + const { assetItems, isStackExpanded, toggleStack } = useOutputStacks({ + assets: ref([parent]) + }) + + await toggleStack(parent) + await toggleStack(parent) + + expect(isStackExpanded(parent)).toBe(false) + expect(assetItems.value.map((item) => item.asset.id)).toEqual([parent.id]) + }) + + it('ignores assets without stack metadata', async () => { + const asset = createAsset({ + id: 'no-meta', + name: 'no-meta.png', + user_metadata: undefined + }) + + const { assetItems, isStackExpanded, toggleStack } = useOutputStacks({ + assets: ref([asset]) + }) + + await toggleStack(asset) + + expect(mocks.resolveOutputAssetItems).not.toHaveBeenCalled() + expect(isStackExpanded(asset)).toBe(false) + expect(assetItems.value).toHaveLength(1) + expect(assetItems.value[0].asset).toMatchObject(asset) + }) + + it('does not expand when no children are resolved', async () => { + const parent = createAsset({ id: 'parent', name: 'parent.png' }) + + vi.mocked(mocks.resolveOutputAssetItems).mockResolvedValue([]) + + const { assetItems, isStackExpanded, toggleStack } = useOutputStacks({ + assets: ref([parent]) + }) + + await toggleStack(parent) + + expect(isStackExpanded(parent)).toBe(false) + expect(assetItems.value.map((item) => item.asset.id)).toEqual([parent.id]) + }) + + it('does not expand when resolving children throws', async () => { + const parent = createAsset({ id: 'parent', name: 'parent.png' }) + const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + + vi.mocked(mocks.resolveOutputAssetItems).mockRejectedValue( + new Error('resolve failed') + ) + + const { assetItems, isStackExpanded, toggleStack } = useOutputStacks({ + assets: ref([parent]) + }) + + await toggleStack(parent) + + expect(isStackExpanded(parent)).toBe(false) + expect(assetItems.value.map((item) => item.asset.id)).toEqual([parent.id]) + + errorSpy.mockRestore() + }) + + it('guards against duplicate loads while a stack is resolving', async () => { + const parent = createAsset({ id: 'parent', name: 'parent.png' }) + const child = createAsset({ + id: 'child', + name: 'child.png', + user_metadata: undefined + }) + const deferred = createDeferred() + + vi.mocked(mocks.resolveOutputAssetItems).mockReturnValue(deferred.promise) + + const { assetItems, toggleStack } = useOutputStacks({ + assets: ref([parent]) + }) + + const firstToggle = toggleStack(parent) + const secondToggle = toggleStack(parent) + + expect(mocks.resolveOutputAssetItems).toHaveBeenCalledTimes(1) + + deferred.resolve([child]) + + await firstToggle + await secondToggle + + expect(assetItems.value.map((item) => item.asset.id)).toEqual([ + parent.id, + child.id + ]) + }) +}) diff --git a/src/platform/assets/composables/useOutputStacks.ts b/src/platform/assets/composables/useOutputStacks.ts new file mode 100644 index 000000000..5c53c23e4 --- /dev/null +++ b/src/platform/assets/composables/useOutputStacks.ts @@ -0,0 +1,127 @@ +import { computed, ref } from 'vue' +import type { Ref } from 'vue' + +import { getOutputAssetMetadata } from '@/platform/assets/schemas/assetMetadataSchema' +import type { AssetItem } from '@/platform/assets/schemas/assetSchema' +import { resolveOutputAssetItems } from '@/platform/assets/utils/outputAssetUtil' + +type OutputStackListItem = { + key: string + asset: AssetItem + isChild?: boolean +} + +type UseOutputStacksOptions = { + assets: Ref +} + +export function useOutputStacks({ assets }: UseOutputStacksOptions) { + const expandedStackPromptIds = ref>(new Set()) + const stackChildrenByPromptId = ref>({}) + const loadingStackPromptIds = ref>(new Set()) + + const assetItems = computed(() => { + const items: OutputStackListItem[] = [] + + for (const asset of assets.value) { + const promptId = getStackPromptId(asset) + items.push({ + key: `asset-${asset.id}`, + asset + }) + + if (!promptId || !expandedStackPromptIds.value.has(promptId)) { + continue + } + + const children = stackChildrenByPromptId.value[promptId] ?? [] + for (const child of children) { + items.push({ + key: `asset-${child.id}`, + asset: child, + isChild: true + }) + } + } + + return items + }) + + const selectableAssets = computed(() => + assetItems.value.map((item) => item.asset) + ) + + function getStackPromptId(asset: AssetItem): string | null { + const metadata = getOutputAssetMetadata(asset.user_metadata) + return metadata?.promptId ?? null + } + + function isStackExpanded(asset: AssetItem): boolean { + const promptId = getStackPromptId(asset) + if (!promptId) return false + return expandedStackPromptIds.value.has(promptId) + } + + async function toggleStack(asset: AssetItem) { + const promptId = getStackPromptId(asset) + if (!promptId) return + + if (expandedStackPromptIds.value.has(promptId)) { + const next = new Set(expandedStackPromptIds.value) + next.delete(promptId) + expandedStackPromptIds.value = next + return + } + + if (!stackChildrenByPromptId.value[promptId]?.length) { + if (loadingStackPromptIds.value.has(promptId)) { + return + } + const nextLoading = new Set(loadingStackPromptIds.value) + nextLoading.add(promptId) + loadingStackPromptIds.value = nextLoading + + const children = await resolveStackChildren(asset) + + const afterLoading = new Set(loadingStackPromptIds.value) + afterLoading.delete(promptId) + loadingStackPromptIds.value = afterLoading + + if (!children.length) { + return + } + + stackChildrenByPromptId.value = { + ...stackChildrenByPromptId.value, + [promptId]: children + } + } + + const nextExpanded = new Set(expandedStackPromptIds.value) + nextExpanded.add(promptId) + expandedStackPromptIds.value = nextExpanded + } + + async function resolveStackChildren(asset: AssetItem): Promise { + const metadata = getOutputAssetMetadata(asset.user_metadata) + if (!metadata) { + return [] + } + try { + return await resolveOutputAssetItems(metadata, { + createdAt: asset.created_at, + excludeOutputKey: asset.name + }) + } catch (error) { + console.error('Failed to resolve stack children:', error) + return [] + } + } + + return { + assetItems, + selectableAssets, + isStackExpanded, + toggleStack + } +}