Merge branch 'main' into synap5e/fix/context-menu-download-multi-output

Amp-Thread-ID: https://ampcode.com/threads/T-019d6f4f-c2c0-734c-8230-f6896b1e4350
Co-authored-by: Amp <amp@ampcode.com>

browser_tests/tests/sidebar/assets.cloud.spec.ts

@@ -0,0 +1,173 @@
+import type {
+  CreateAssetExportData,
+  CreateAssetExportResponse
+} from '@comfyorg/ingest-types'
+import { expect } from '@playwright/test'
+import type { Page } from '@playwright/test'
+
+import { ComfyPage, comfyPageFixture } from '@e2e/fixtures/ComfyPage'
+import { createMockJob } from '@e2e/fixtures/helpers/AssetsHelper'
+import type { RemoteConfig } from '@/platform/remoteConfig/types'
+import type { operations } from '@/types/comfyRegistryTypes'
+
+const TEST_API_KEY = 'playwright-cloud-api-key'
+const TEST_USERNAME = 'playwright-cloud-assets'
+const PLAYWRIGHT_TEST_URL =
+  process.env.PLAYWRIGHT_TEST_URL || 'http://localhost:8188'
+const ASSET_EXPORT_ENDPOINT = '/api/assets/export'
+
+type CreateCustomerResponse =
+  operations['createCustomer']['responses']['201']['content']['application/json']
+
+type CloudAssetsFixtures = {
+  comfyPage: ComfyPage
+}
+
+const test = comfyPageFixture.extend<CloudAssetsFixtures>({
+  comfyPage: async ({ page, request }, use) => {
+    const comfyPage = new ComfyPage(page, request)
+
+    await use(comfyPage)
+
+    await comfyPage.assetApi.clearMocks()
+  }
+})
+
+type SeedCloudSessionArgs = {
+  apiKey: string
+  userId: string
+  username: string
+}
+
+const featuresResponse: RemoteConfig = {
+  comfy_api_base_url: `${PLAYWRIGHT_TEST_URL}/api`
+}
+
+const customerResponse: CreateCustomerResponse = {
+  id: 'customer-1'
+}
+
+const assetExportResponse: CreateAssetExportResponse = {
+  task_id: 'asset-export-task-1',
+  status: 'created'
+}
+
+const SAMPLE_JOBS = [
+  createMockJob({
+    id: 'job-gamma',
+    create_time: 3000,
+    execution_start_time: 3000,
+    execution_end_time: 3020,
+    preview_output: {
+      filename: 'abstract_art.png',
+      subfolder: '',
+      type: 'output',
+      nodeId: '3',
+      mediaType: 'images'
+    },
+    outputs_count: 2
+  })
+]
+
+async function seedCloudSession(userId: string, page: Page) {
+  await page.addInitScript(
+    ({ apiKey, userId: seededUserId, username }: SeedCloudSessionArgs) => {
+      localStorage.clear()
+      sessionStorage.clear()
+      localStorage.setItem('comfy_api_key', apiKey)
+      localStorage.setItem('Comfy.userId', seededUserId)
+      localStorage.setItem('Comfy.userName', username)
+    },
+    {
+      apiKey: TEST_API_KEY,
+      userId,
+      username: TEST_USERNAME
+    } satisfies SeedCloudSessionArgs
+  )
+}
+
+test.describe(
+  'Assets sidebar - cloud multi-output download',
+  {
+    tag: '@cloud'
+  },
+  () => {
+    test('context menu Download creates a ZIP export for a multi-output job', async ({
+      page,
+      comfyPage
+    }) => {
+      const userId = await comfyPage.setupUser(TEST_USERNAME)
+
+      await seedCloudSession(userId, page)
+
+      await comfyPage.setupSettings({
+        'Comfy.UseNewMenu': 'Top',
+        'Comfy.Graph.CanvasInfo': false,
+        'Comfy.Graph.CanvasMenu': false,
+        'Comfy.Canvas.SelectionToolbox': false,
+        'Comfy.EnableTooltips': false,
+        'Comfy.userId': userId,
+        'Comfy.TutorialCompleted': true,
+        'Comfy.VersionCompatibility.DisableWarnings': true,
+        'Comfy.RightSidePanel.ShowErrorsTab': false
+      })
+
+      await page.route('**/api/features', async (route) => {
+        await route.fulfill({
+          status: 200,
+          contentType: 'application/json',
+          body: JSON.stringify(featuresResponse)
+        })
+      })
+
+      await page.route('**/customers', async (route) => {
+        await route.fulfill({
+          status: 201,
+          contentType: 'application/json',
+          body: JSON.stringify(customerResponse)
+        })
+      })
+
+      await page.route(`**${ASSET_EXPORT_ENDPOINT}`, async (route) => {
+        await route.fulfill({
+          status: 202,
+          contentType: 'application/json',
+          body: JSON.stringify(assetExportResponse)
+        })
+      })
+
+      await comfyPage.assets.mockOutputHistory(SAMPLE_JOBS)
+      await comfyPage.assets.mockInputFiles([])
+      await comfyPage.setup({ clearStorage: false })
+
+      const tab = comfyPage.menu.assetsTab
+      await tab.open()
+      await tab.waitForAssets(1)
+
+      const exportRequestPromise = page.waitForRequest((request) => {
+        const url = new URL(request.url())
+        return (
+          request.method() === 'POST' &&
+          url.pathname.endsWith(ASSET_EXPORT_ENDPOINT)
+        )
+      })
+
+      await tab.assetCards.first().dispatchEvent('contextmenu', {
+        bubbles: true,
+        cancelable: true,
+        button: 2
+      })
+
+      await expect(page.locator('.p-contextmenu')).toBeVisible()
+      await tab.contextMenuItem('Download').click()
+
+      const exportRequest = await exportRequestPromise
+      const payload =
+        exportRequest.postDataJSON() as CreateAssetExportData['body']
+
+      expect(payload.job_ids).toEqual(['job-gamma'])
+      expect(payload.naming_strategy).toBe('preserve')
+      expect(payload).not.toHaveProperty('job_asset_name_filters')
+    })
+  }
+)

src/platform/assets/components/MediaAssetContextMenu.test.ts

@@ -1,7 +1,8 @@
-/* eslint-disable vue/one-component-per-file */
-import { render } from '@testing-library/vue'
+import { mount } from '@vue/test-utils'
+import type { MenuItem } from 'primevue/menuitem'
 import { afterEach, describe, expect, it, vi } from 'vitest'
-import { defineComponent, nextTick, onMounted, ref } from 'vue'
+import { defineComponent, nextTick } from 'vue'
+import type { ComponentPublicInstance } from 'vue'
 
 import MediaAssetContextMenu from '@/platform/assets/components/MediaAssetContextMenu.vue'
 import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
@@ -31,6 +32,7 @@ vi.mock('@/utils/loaderNodeUtil', () => ({
 const mediaAssetActions = {
   addWorkflow: vi.fn(),
   downloadAsset: vi.fn(),
+  downloadMultipleAssets: vi.fn(),
   openWorkflow: vi.fn(),
   exportWorkflow: vi.fn(),
   copyJobId: vi.fn(),
@@ -47,6 +49,10 @@ const contextMenuStub = defineComponent({
     pt: {
       type: Object,
       default: undefined
+    },
+    model: {
+      type: Array,
+      default: () => []
     }
   },
   emits: ['hide'],
@@ -84,73 +90,81 @@ const buttonStub = {
   template: '<div class="button-stub"><slot /></div>'
 }
 
-interface MediaAssetContextMenuExposed {
+type MediaAssetContextMenuExposed = ComponentPublicInstance & {
   show: (event: MouseEvent) => void
 }
 
-let capturedRef: MediaAssetContextMenuExposed | null = null
-
-function mountComponent() {
-  const onHide = vi.fn()
-  const { container, unmount } = render(
-    defineComponent({
-      components: { MediaAssetContextMenu },
-      setup() {
-        const menuRef = ref<MediaAssetContextMenuExposed | null>(null)
-        onMounted(() => {
-          capturedRef = menuRef.value
-        })
-        return { menuRef, asset, onHide }
-      },
-      template:
-        '<MediaAssetContextMenu ref="menuRef" :asset="asset" asset-type="output" file-kind="image" @hide="onHide" />'
-    }),
-    {
-      global: {
-        stubs: {
-          ContextMenu: contextMenuStub,
-          Button: buttonStub
-        }
+const mountComponent = () =>
+  mount(MediaAssetContextMenu, {
+    attachTo: document.body,
+    props: {
+      asset,
+      assetType: 'output',
+      fileKind: 'image'
+    },
+    global: {
+      stubs: {
+        ContextMenu: contextMenuStub,
+        Button: buttonStub
       }
     }
-  )
-  return { container, unmount, onHide }
-}
+  })
 
-async function showMenu(container: Element): Promise<HTMLElement> {
+async function showMenu(
+  wrapper: ReturnType<typeof mountComponent>
+): Promise<HTMLElement> {
+  const exposed = wrapper.vm as MediaAssetContextMenuExposed
   const event = new MouseEvent('contextmenu', { bubbles: true })
-  capturedRef!.show(event)
+  exposed.show(event)
   await nextTick()
-  // eslint-disable-next-line testing-library/no-container
-  return container.querySelector('.context-menu-stub') as HTMLElement
+
+  return wrapper.get('.context-menu-stub').element as HTMLElement
 }
 
 afterEach(() => {
   vi.clearAllMocks()
-  capturedRef = null
   document.body.innerHTML = ''
 })
 
 describe('MediaAssetContextMenu', () => {
   it('dismisses outside pointerdown using the rendered root id', async () => {
-    const { container, unmount, onHide } = mountComponent()
+    const wrapper = mountComponent()
     const outside = document.createElement('div')
     document.body.append(outside)
 
-    const menu = await showMenu(container)
+    const menu = await showMenu(wrapper)
     const menuId = menu.id
 
     expect(menuId).not.toBe('')
-    // eslint-disable-next-line testing-library/no-node-access
     expect(document.getElementById(menuId)).toBe(menu)
 
     outside.dispatchEvent(new Event('pointerdown', { bubbles: true }))
     await nextTick()
 
-    // eslint-disable-next-line testing-library/no-container, testing-library/no-node-access
-    expect(container.querySelector('.context-menu-stub')).toBeNull()
-    expect(onHide).toHaveBeenCalledOnce()
+    expect(wrapper.find('.context-menu-stub').exists()).toBe(false)
+    expect(wrapper.emitted('hide')).toEqual([[]])
 
-    unmount()
+    wrapper.unmount()
+  })
+
+  it('routes Download through downloadMultipleAssets so multi-output jobs zip', async () => {
+    const wrapper = mountComponent()
+    await showMenu(wrapper)
+
+    const menu = wrapper.findComponent(contextMenuStub)
+    const items = menu.props('model') as MenuItem[]
+    const download = items.find(
+      (item) => item.label === 'mediaAsset.actions.download'
+    )
+
+    expect(download?.command).toBeTypeOf('function')
+    download?.command?.({ originalEvent: new Event('click'), item: download })
+
+    expect(mediaAssetActions.downloadMultipleAssets).toHaveBeenCalledWith([
+      asset
+    ])
+    expect(mediaAssetActions.downloadAsset).not.toHaveBeenCalled()
+
+    wrapper.unmount()
   })
 })

src/platform/assets/components/MediaAssetContextMenu.vue

@@ -217,7 +217,7 @@ const contextMenuItems = computed<MenuItem[]>(() => {
   items.push({
     label: t('mediaAsset.actions.download'),
     icon: 'icon-[lucide--download]',
-    command: () => actions.downloadAsset(asset)
+    command: () => actions.downloadMultipleAssets([asset])
   })
 
   // Separator before workflow actions (only if there are workflow actions)

src/stores/bootstrapStore.test.ts

@@ -59,6 +59,13 @@ vi.mock('@/stores/authStore', () => ({
   }))
 }))
 
+const mockIsApiKeyAuthenticated = ref(false)
+vi.mock('@/stores/apiKeyAuthStore', () => ({
+  useApiKeyAuthStore: vi.fn(() => ({
+    isAuthenticated: mockIsApiKeyAuthenticated
+  }))
+}))
+
 const mockDistributionTypes = vi.hoisted(() => ({
   isCloud: false
 }))
@@ -69,6 +76,7 @@ describe('bootstrapStore', () => {
     mockIsSettingsReady.value = false
     mockIsAuthInitialized.value = false
     mockIsAuthAuthenticated.value = false
+    mockIsApiKeyAuthenticated.value = false
     mockNeedsLogin.value = false
     mockDistributionTypes.isCloud = false
     setActivePinia(createTestingPinia({ stubActions: false }))
@@ -122,5 +130,28 @@ describe('bootstrapStore', () => {
         expect(settingStore.isReady).toBe(true)
       })
     })
+
+    it('allows API key auth to satisfy cloud bootstrap', async () => {
+      const store = useBootstrapStore()
+      const settingStore = useSettingStore()
+      const bootstrapPromise = store.startStoreBootstrap()
+
+      expect(store.isI18nReady).toBe(false)
+      expect(settingStore.isReady).toBe(false)
+
+      mockIsAuthInitialized.value = true
+      await nextTick()
+
+      expect(store.isI18nReady).toBe(false)
+      expect(settingStore.isReady).toBe(false)
+
+      mockIsApiKeyAuthenticated.value = true
+      await bootstrapPromise
+
+      await vi.waitFor(() => {
+        expect(store.isI18nReady).toBe(true)
+        expect(settingStore.isReady).toBe(true)
+      })
+    })
   })
 })

src/stores/bootstrapStore.ts

@@ -5,6 +5,7 @@ import { isCloud } from '@/platform/distribution/types'
 import { useSettingStore } from '@/platform/settings/settingStore'
 import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
 import { api } from '@/scripts/api'
+import { useApiKeyAuthStore } from '@/stores/apiKeyAuthStore'
 import { useAuthStore } from '@/stores/authStore'
 import { useUserStore } from '@/stores/userStore'
 
@@ -38,8 +39,13 @@ export const useBootstrapStore = defineStore('bootstrap', () => {
   async function startStoreBootstrap() {
     if (isCloud) {
       const { isInitialized, isAuthenticated } = storeToRefs(useAuthStore())
+      const { isAuthenticated: isApiKeyAuthenticated } =
+        storeToRefs(useApiKeyAuthStore())
+
       await until(isInitialized).toBe(true)
-      await until(isAuthenticated).toBe(true)
+      await until(
+        () => isAuthenticated.value || isApiKeyAuthenticated.value
+      ).toBe(true)
     }
 
     const userStore = useUserStore()