Add virtual scroll to Queue Tab for improved performance (#2108)

browser_tests/fixtures/ComfyPage.ts

@@ -14,11 +14,13 @@ import { ComfyNodeSearchBox } from './components/ComfyNodeSearchBox'
 import { SettingDialog } from './components/SettingDialog'
 import {
   NodeLibrarySidebarTab,
+  QueueSidebarTab,
   WorkflowsSidebarTab
 } from './components/SidebarTab'
 import { Topbar } from './components/Topbar'
 import type { Position, Size } from './types'
 import { NodeReference } from './utils/litegraphUtils'
+import TaskHistory from './utils/taskHistory'
 
 dotenv.config()
 
@@ -45,6 +47,10 @@ class ComfyMenu {
     return new WorkflowsSidebarTab(this.page)
   }
 
+  get queueTab() {
+    return new QueueSidebarTab(this.page)
+  }
+
   get topbar() {
     return new Topbar(this.page)
   }
@@ -233,6 +239,10 @@ export class ComfyPage {
     }
   }
 
+  setupHistory(): TaskHistory {
+    return new TaskHistory(this)
+  }
+
   async setup({ clearStorage = true }: { clearStorage?: boolean } = {}) {
     await this.goto()
     if (clearStorage) {

browser_tests/fixtures/components/SidebarTab.ts

@@ -22,6 +22,12 @@ class SidebarTab {
     }
     await this.tabButton.click()
   }
+  async close() {
+    if (!this.tabButton.isVisible()) {
+      return
+    }
+    await this.tabButton.click()
+  }
 }
 
 export class NodeLibrarySidebarTab extends SidebarTab {
@@ -147,3 +153,93 @@ export class WorkflowsSidebarTab extends SidebarTab {
     await this.page.waitForTimeout(300)
   }
 }
+
+export class QueueSidebarTab extends SidebarTab {
+  constructor(public readonly page: Page) {
+    super(page, 'queue')
+  }
+
+  get root() {
+    return this.page.locator('.sidebar-content-container', { hasText: 'Queue' })
+  }
+
+  get tasks() {
+    return this.root.locator('[data-virtual-grid-item]')
+  }
+
+  get visibleTasks() {
+    return this.tasks.locator('visible=true')
+  }
+
+  get clearButton() {
+    return this.root.locator('.clear-all-button')
+  }
+
+  get collapseTasksButton() {
+    return this.getToggleExpandButton(false)
+  }
+
+  get expandTasksButton() {
+    return this.getToggleExpandButton(true)
+  }
+
+  get noResultsPlaceholder() {
+    return this.root.locator('.no-results-placeholder')
+  }
+
+  private getToggleExpandButton(isExpanded: boolean) {
+    const iconSelector = isExpanded ? '.pi-image' : '.pi-images'
+    return this.root.locator(`.toggle-expanded-button ${iconSelector}`)
+  }
+
+  async open() {
+    await super.open()
+    return this.root.waitFor({ state: 'visible' })
+  }
+
+  async close() {
+    await super.close()
+    await this.root.waitFor({ state: 'hidden' })
+  }
+
+  async expandTasks() {
+    await this.expandTasksButton.click()
+    await this.collapseTasksButton.waitFor({ state: 'visible' })
+  }
+
+  async collapseTasks() {
+    await this.collapseTasksButton.click()
+    await this.expandTasksButton.waitFor({ state: 'visible' })
+  }
+
+  async waitForTasks() {
+    return Promise.all([
+      this.tasks.first().waitFor({ state: 'visible' }),
+      this.tasks.last().waitFor({ state: 'visible' })
+    ])
+  }
+
+  async scrollTasks(direction: 'up' | 'down') {
+    const scrollToEl =
+      direction === 'up' ? this.tasks.last() : this.tasks.first()
+    await scrollToEl.scrollIntoViewIfNeeded()
+    await this.waitForTasks()
+  }
+
+  async clearTasks() {
+    await this.clearButton.click()
+    const confirmButton = this.page.getByLabel('Delete')
+    await confirmButton.click()
+    await this.noResultsPlaceholder.waitFor({ state: 'visible' })
+  }
+
+  /** Set the width of the tab (out of 100). Must call before opening the tab */
+  async setTabWidth(width: number) {
+    if (width < 0 || width > 100) {
+      throw new Error('Width must be between 0 and 100')
+    }
+    return this.page.evaluate((width) => {
+      localStorage.setItem('queue', JSON.stringify([width, 100 - width]))
+    }, width)
+  }
+}

browser_tests/fixtures/utils/taskHistory.ts

@@ -0,0 +1,161 @@
+import fs from 'fs'
+import _ from 'lodash'
+import path from 'path'
+import type { Request, Route } from 'playwright'
+import { v4 as uuidv4 } from 'uuid'
+
+import type {
+  HistoryTaskItem,
+  TaskItem,
+  TaskOutput
+} from '../../../src/types/apiTypes'
+import type { ComfyPage } from '../ComfyPage'
+
+/** keyof TaskOutput[string] */
+type OutputFileType = 'images' | 'audio' | 'animated'
+
+const DEFAULT_IMAGE = 'example.webp'
+
+const getFilenameParam = (request: Request) => {
+  const url = new URL(request.url())
+  return url.searchParams.get('filename') || DEFAULT_IMAGE
+}
+
+const getContentType = (filename: string, fileType: OutputFileType) => {
+  const subtype = path.extname(filename).slice(1)
+  switch (fileType) {
+    case 'images':
+      return `image/${subtype}`
+    case 'audio':
+      return `audio/${subtype}`
+    case 'animated':
+      return `video/${subtype}`
+  }
+}
+
+const setQueueIndex = (task: TaskItem) => {
+  task.prompt[0] = TaskHistory.queueIndex++
+}
+
+const setPromptId = (task: TaskItem) => {
+  task.prompt[1] = uuidv4()
+}
+
+export default class TaskHistory {
+  static queueIndex = 0
+  static readonly defaultTask: Readonly<HistoryTaskItem> = {
+    prompt: [0, 'prompt-id', {}, { client_id: uuidv4() }, []],
+    outputs: {},
+    status: {
+      status_str: 'success',
+      completed: true,
+      messages: []
+    },
+    taskType: 'History'
+  }
+  private tasks: HistoryTaskItem[] = []
+  private outputContentTypes: Map<string, string> = new Map()
+
+  constructor(readonly comfyPage: ComfyPage) {}
+
+  private loadAsset: (filename: string) => Buffer = _.memoize(
+    (filename: string) => {
+      const filePath = this.comfyPage.assetPath(filename)
+      return fs.readFileSync(filePath)
+    }
+  )
+
+  private async handleGetHistory(route: Route) {
+    return route.fulfill({
+      status: 200,
+      contentType: 'application/json',
+      body: JSON.stringify(this.tasks)
+    })
+  }
+
+  private async handleGetView(route: Route) {
+    const fileName = getFilenameParam(route.request())
+    if (!this.outputContentTypes.has(fileName)) route.continue()
+
+    const asset = this.loadAsset(fileName)
+    return route.fulfill({
+      status: 200,
+      contentType: this.outputContentTypes.get(fileName),
+      body: asset,
+      headers: {
+        'Cache-Control': 'public, max-age=31536000',
+        'Content-Length': asset.byteLength.toString()
+      }
+    })
+  }
+
+  async setupRoutes() {
+    return this.comfyPage.page.route(
+      /.*\/api\/(view|history)(\?.*)?$/,
+      async (route) => {
+        const request = route.request()
+        const method = request.method()
+
+        const isViewReq = request.url().includes('view') && method === 'GET'
+        if (isViewReq) return this.handleGetView(route)
+
+        const isHistoryPath = request.url().includes('history')
+        const isGetHistoryReq = isHistoryPath && method === 'GET'
+        if (isGetHistoryReq) return this.handleGetHistory(route)
+
+        const isClearReq =
+          method === 'POST' &&
+          isHistoryPath &&
+          request.postDataJSON()?.clear === true
+        if (isClearReq) return this.clearTasks()
+
+        return route.continue()
+      }
+    )
+  }
+
+  private createOutputs(
+    filenames: string[],
+    filetype: OutputFileType
+  ): TaskOutput {
+    return filenames.reduce((outputs, filename, i) => {
+      const nodeId = `${i + 1}`
+      outputs[nodeId] = {
+        [filetype]: [{ filename, subfolder: '', type: 'output' }]
+      }
+      const contentType = getContentType(filename, filetype)
+      this.outputContentTypes.set(filename, contentType)
+      return outputs
+    }, {})
+  }
+
+  private addTask(task: HistoryTaskItem) {
+    setPromptId(task)
+    setQueueIndex(task)
+    this.tasks.push(task)
+  }
+
+  clearTasks() {
+    this.tasks = []
+  }
+
+  withTask(
+    outputFilenames: string[],
+    outputFiletype: OutputFileType = 'images',
+    overrides: Partial<HistoryTaskItem> = {}
+  ): this {
+    this.addTask({
+      ...TaskHistory.defaultTask,
+      outputs: this.createOutputs(outputFilenames, outputFiletype),
+      ...overrides
+    })
+    return this
+  }
+
+  /** Repeats the last task in the task history a specified number of times. */
+  repeat(n: number): this {
+    for (let i = 0; i < n; i++)
+      this.addTask(structuredClone(this.tasks.at(-1)) as HistoryTaskItem)
+    return this
+  }
+}

browser_tests/menu.spec.ts

@@ -704,3 +704,143 @@ test.describe('Menu', () => {
     })
   })
 })
+
+test.describe('Queue sidebar', () => {
+  test.beforeEach(async ({ comfyPage }) => {
+    await comfyPage.setSetting('Comfy.UseNewMenu', 'Top')
+  })
+
+  test('can display tasks', async ({ comfyPage }) => {
+    await comfyPage.setupHistory().withTask(['example.webp']).setupRoutes()
+    await comfyPage.menu.queueTab.open()
+    await comfyPage.menu.queueTab.waitForTasks()
+    expect(await comfyPage.menu.queueTab.visibleTasks.count()).toBe(1)
+  })
+
+  test('can display tasks after closing then opening', async ({
+    comfyPage
+  }) => {
+    await comfyPage.setupHistory().withTask(['example.webp']).setupRoutes()
+    await comfyPage.menu.queueTab.open()
+    await comfyPage.menu.queueTab.close()
+    await comfyPage.menu.queueTab.open()
+    await comfyPage.menu.queueTab.waitForTasks()
+    expect(await comfyPage.menu.queueTab.visibleTasks.count()).toBe(1)
+  })
+
+  test.describe('Virtual scroll', () => {
+    const layouts = [
+      { description: 'Five columns layout', width: 95, rows: 3, cols: 5 },
+      { description: 'Three columns layout', width: 55, rows: 3, cols: 3 },
+      { description: 'Two columns layout', width: 40, rows: 3, cols: 2 }
+    ]
+
+    test.beforeEach(async ({ comfyPage }) => {
+      await comfyPage
+        .setupHistory()
+        .withTask(['example.webp'])
+        .repeat(50)
+        .setupRoutes()
+    })
+
+    layouts.forEach(({ description, width, rows, cols }) => {
+      const preRenderedRows = 1
+      const preRenderedTasks = preRenderedRows * cols * 2
+      const visibleTasks = rows * cols
+      const expectRenderLimit = visibleTasks + preRenderedTasks
+
+      test.describe(description, () => {
+        test.beforeEach(async ({ comfyPage }) => {
+          await comfyPage.menu.queueTab.setTabWidth(width)
+          await comfyPage.menu.queueTab.open()
+          await comfyPage.menu.queueTab.waitForTasks()
+        })
+
+        test('should not render items outside of view', async ({
+          comfyPage
+        }) => {
+          const renderedCount =
+            await comfyPage.menu.queueTab.visibleTasks.count()
+          expect(renderedCount).toBeLessThanOrEqual(expectRenderLimit)
+        })
+
+        test('should teardown items after scrolling away', async ({
+          comfyPage
+        }) => {
+          await comfyPage.menu.queueTab.scrollTasks('down')
+          const renderedCount =
+            await comfyPage.menu.queueTab.visibleTasks.count()
+          expect(renderedCount).toBeLessThanOrEqual(expectRenderLimit)
+        })
+
+        test('should re-render items after scrolling away then back', async ({
+          comfyPage
+        }) => {
+          await comfyPage.menu.queueTab.scrollTasks('down')
+          await comfyPage.menu.queueTab.scrollTasks('up')
+          const renderedCount =
+            await comfyPage.menu.queueTab.visibleTasks.count()
+          expect(renderedCount).toBeLessThanOrEqual(expectRenderLimit)
+        })
+      })
+    })
+  })
+
+  test.describe('Expand tasks', () => {
+    test.beforeEach(async ({ comfyPage }) => {
+      // 2-item batch and 3-item batch -> 3 additional items when expanded
+      await comfyPage
+        .setupHistory()
+        .withTask(['example.webp', 'example.webp', 'example.webp'])
+        .withTask(['example.webp', 'example.webp'])
+        .setupRoutes()
+      await comfyPage.menu.queueTab.open()
+      await comfyPage.menu.queueTab.waitForTasks()
+    })
+
+    test('can expand tasks with multiple outputs', async ({ comfyPage }) => {
+      const initialCount = await comfyPage.menu.queueTab.visibleTasks.count()
+      await comfyPage.menu.queueTab.expandTasks()
+      expect(await comfyPage.menu.queueTab.visibleTasks.count()).toBe(
+        initialCount + 3
+      )
+    })
+
+    test('can collapse flat tasks', async ({ comfyPage }) => {
+      const initialCount = await comfyPage.menu.queueTab.visibleTasks.count()
+      await comfyPage.menu.queueTab.expandTasks()
+      await comfyPage.menu.queueTab.collapseTasks()
+      expect(await comfyPage.menu.queueTab.visibleTasks.count()).toBe(
+        initialCount
+      )
+    })
+  })
+
+  test.describe('Clear tasks', () => {
+    test.beforeEach(async ({ comfyPage }) => {
+      await comfyPage
+        .setupHistory()
+        .withTask(['example.webp'])
+        .repeat(6)
+        .setupRoutes()
+      await comfyPage.menu.queueTab.open()
+    })
+
+    test('can clear all tasks', async ({ comfyPage }) => {
+      await comfyPage.menu.queueTab.clearTasks()
+      expect(await comfyPage.menu.queueTab.visibleTasks.count()).toBe(0)
+      expect(
+        await comfyPage.menu.queueTab.noResultsPlaceholder.isVisible()
+      ).toBe(true)
+    })
+
+    test('can load new tasks after clearing all', async ({ comfyPage }) => {
+      await comfyPage.menu.queueTab.clearTasks()
+      await comfyPage.menu.queueTab.close()
+      await comfyPage.setupHistory().withTask(['example.webp']).setupRoutes()
+      await comfyPage.menu.queueTab.open()
+      await comfyPage.menu.queueTab.waitForTasks()
+      expect(await comfyPage.menu.queueTab.visibleTasks.count()).toBe(1)
+    })
+  })
+})