Add support for multiple changes in a single ChangeTracker state (#1022)

* wip

* Add tests

* Update package

* remove logs

* nit

* nit

---------

Co-authored-by: huchenlei <huchenlei@proton.me>

browser_tests/ComfyPage.ts

@@ -771,46 +771,43 @@ export class ComfyPage {
     await this.nextFrame()
   }
 
-  async ctrlC() {
+  async ctrlSend(keyToPress: string) {
     await this.page.keyboard.down('Control')
-    await this.page.keyboard.press('KeyC')
+    await this.page.keyboard.press(keyToPress)
     await this.page.keyboard.up('Control')
     await this.nextFrame()
   }
 
+  async ctrlA() {
+    await this.ctrlSend('KeyA')
+  }
+
+  async ctrlB() {
+    await this.ctrlSend('KeyB')
+  }
+
+  async ctrlC() {
+    await this.ctrlSend('KeyC')
+  }
+
   async ctrlV() {
-    await this.page.keyboard.down('Control')
-    await this.page.keyboard.press('KeyV')
-    await this.page.keyboard.up('Control')
-    await this.nextFrame()
+    await this.ctrlSend('KeyV')
   }
 
   async ctrlZ() {
-    await this.page.keyboard.down('Control')
-    await this.page.keyboard.press('KeyZ')
-    await this.page.keyboard.up('Control')
-    await this.nextFrame()
+    await this.ctrlSend('KeyZ')
   }
 
   async ctrlY() {
-    await this.page.keyboard.down('Control')
-    await this.page.keyboard.press('KeyY')
-    await this.page.keyboard.up('Control')
-    await this.nextFrame()
+    await this.ctrlSend('KeyY')
   }
 
   async ctrlArrowUp() {
-    await this.page.keyboard.down('Control')
-    await this.page.keyboard.press('ArrowUp')
-    await this.page.keyboard.up('Control')
-    await this.nextFrame()
+    await this.ctrlSend('ArrowUp')
   }
 
   async ctrlArrowDown() {
-    await this.page.keyboard.down('Control')
-    await this.page.keyboard.press('ArrowDown')
-    await this.page.keyboard.up('Control')
-    await this.nextFrame()
+    await this.ctrlSend('ArrowDown')
   }
 
   async closeMenu() {
@@ -941,6 +938,9 @@ export class ComfyPage {
     if (!id) return null
     return this.getNodeRefById(id)
   }
+  async moveMouseToEmptyArea() {
+    await this.page.mouse.move(10, 10)
+  }
 }
 
 export class NodeSlotReference {
@@ -1046,6 +1046,20 @@ export class NodeReference {
       y: pos[1]
     }
   }
+  async getBounding(): Promise<Position & Size> {
+    const [x, y, width, height]: [number, number, number, number] =
+      await this.comfyPage.page.evaluate((id) => {
+        const node = window['app'].graph.getNodeById(id)
+        if (!node) throw new Error('Node not found')
+        return node.getBounding()
+      }, this.id)
+    return {
+      x,
+      y,
+      width,
+      height
+    }
+  }
   async getSize(): Promise<Size> {
     const size = await this.getProperty<[number, number]>('size')
     return {
@@ -1053,6 +1067,18 @@ export class NodeReference {
       height: size[1]
     }
   }
+  async getFlags(): Promise<{ collapsed?: boolean; pinned?: boolean }> {
+    return await this.getProperty('flags')
+  }
+  async isPinned() {
+    return !!(await this.getFlags()).pinned
+  }
+  async isCollapsed() {
+    return !!(await this.getFlags()).collapsed
+  }
+  async isBypassed() {
+    return (await this.getProperty<number | null | undefined>('mode')) === 4
+  }
   async getProperty<T>(prop: string): Promise<T> {
     return await this.comfyPage.page.evaluate(
       ([id, prop]) => {
@@ -1072,22 +1098,37 @@ export class NodeReference {
   async getWidget(index: number) {
     return new NodeWidgetReference(index, this)
   }
-  async click(position: 'title', options?: Parameters<Page['click']>[1]) {
+  async click(
+    position: 'title' | 'collapse',
+    options?: Parameters<Page['click']>[1] & { moveMouseToEmptyArea?: boolean }
+  ) {
     const nodePos = await this.getPosition()
     const nodeSize = await this.getSize()
     let clickPos: Position
     switch (position) {
       case 'title':
-        clickPos = { x: nodePos.x + nodeSize.width / 2, y: nodePos.y + 15 }
+        clickPos = { x: nodePos.x + nodeSize.width / 2, y: nodePos.y - 15 }
+        break
+      case 'collapse':
+        clickPos = { x: nodePos.x + 5, y: nodePos.y - 10 }
         break
       default:
         throw new Error(`Invalid click position ${position}`)
     }
+
+    const moveMouseToEmptyArea = options?.moveMouseToEmptyArea
+    if (options) {
+      delete options.moveMouseToEmptyArea
+    }
+
     await this.comfyPage.canvas.click({
       ...options,
       position: clickPos
     })
     await this.comfyPage.nextFrame()
+    if (moveMouseToEmptyArea) {
+      await this.comfyPage.moveMouseToEmptyArea()
+    }
   }
   async connectWidget(
     originSlotIndex: number,
@@ -1151,3 +1192,35 @@ export const comfyPageFixture = base.extend<{ comfyPage: ComfyPage }>({
     await use(comfyPage)
   }
 })
+
+const makeMatcher = function <T>(
+  getValue: (node: NodeReference) => Promise<T> | T,
+  type: string
+) {
+  return async function (
+    node: NodeReference,
+    options?: { timeout?: number; intervals?: number[] }
+  ) {
+    const value = await getValue(node)
+    let assertion = expect(
+      value,
+      'Node is ' + (this.isNot ? '' : 'not ') + type
+    )
+    if (this.isNot) {
+      assertion = assertion.not
+    }
+    await expect(async () => {
+      assertion.toBeTruthy()
+    }).toPass({ timeout: 250, ...options })
+    return {
+      pass: !this.isNot,
+      message: () => 'Node is ' + (this.isNot ? 'not ' : '') + type
+    }
+  }
+}
+
+export const comfyExpect = expect.extend({
+  toBePinned: makeMatcher((n) => n.isPinned(), 'pinned'),
+  toBeBypassed: makeMatcher((n) => n.isBypassed(), 'bypassed'),
+  toBeCollapsed: makeMatcher((n) => n.isCollapsed(), 'collapsed')
+})

browser_tests/changeTracker.spec.ts

@@ -0,0 +1,100 @@
+import {
+  ComfyPage,
+  comfyPageFixture as test,
+  comfyExpect as expect
+} from './ComfyPage'
+
+async function beforeChange(comfyPage: ComfyPage) {
+  await comfyPage.page.evaluate(() => {
+    window['app'].canvas.emitBeforeChange()
+  })
+}
+async function afterChange(comfyPage: ComfyPage) {
+  await comfyPage.page.evaluate(() => {
+    window['app'].canvas.emitAfterChange()
+  })
+}
+
+test.describe('Change Tracker', () => {
+  test('Can group multiple change actions into a single transaction', async ({
+    comfyPage
+  }) => {
+    const node = (await comfyPage.getFirstNodeRef())!
+    expect(node).toBeTruthy()
+    await expect(node).not.toBeCollapsed()
+    await expect(node).not.toBeBypassed()
+
+    // Make changes outside set
+    // Bypass + collapse node
+    await node.click('collapse')
+    await comfyPage.ctrlB()
+    await expect(node).toBeCollapsed()
+    await expect(node).toBeBypassed()
+
+    // Undo, undo, ensure both changes undone
+    await comfyPage.ctrlZ()
+    await expect(node).not.toBeBypassed()
+    await expect(node).toBeCollapsed()
+    await comfyPage.ctrlZ()
+    await expect(node).not.toBeBypassed()
+    await expect(node).not.toBeCollapsed()
+
+    // Run again, but within a change transaction
+    beforeChange(comfyPage)
+
+    await node.click('collapse')
+    await comfyPage.ctrlB()
+    await expect(node).toBeCollapsed()
+    await expect(node).toBeBypassed()
+
+    // End transaction
+    afterChange(comfyPage)
+
+    // Ensure undo reverts both changes
+    await comfyPage.ctrlZ()
+    await expect(node).not.toBeBypassed()
+    await expect(node).not.toBeCollapsed()
+  })
+
+  test('Can group multiple transaction calls into a single one', async ({
+    comfyPage
+  }) => {
+    const node = (await comfyPage.getFirstNodeRef())!
+    const bypassAndPin = async () => {
+      await beforeChange(comfyPage)
+      await comfyPage.ctrlB()
+      await expect(node).toBeBypassed()
+      await comfyPage.page.keyboard.press('KeyP')
+      await comfyPage.nextFrame()
+      await expect(node).toBePinned()
+      await afterChange(comfyPage)
+    }
+
+    const collapse = async () => {
+      await beforeChange(comfyPage)
+      await node.click('collapse', { moveMouseToEmptyArea: true })
+      await expect(node).toBeCollapsed()
+      await afterChange(comfyPage)
+    }
+
+    const multipleChanges = async () => {
+      await beforeChange(comfyPage)
+      // Call other actions that uses begin/endChange
+      await collapse()
+      await bypassAndPin()
+      await afterChange(comfyPage)
+    }
+
+    await multipleChanges()
+
+    await comfyPage.ctrlZ()
+    await expect(node).not.toBeBypassed()
+    await expect(node).not.toBePinned()
+    await expect(node).not.toBeCollapsed()
+
+    await comfyPage.ctrlY()
+    await expect(node).toBeBypassed()
+    await expect(node).toBePinned()
+    await expect(node).toBeCollapsed()
+  })
+})

browser_tests/copyPaste.spec.ts

@@ -86,4 +86,23 @@ test.describe('Copy Paste', () => {
     await comfyPage.page.keyboard.up('Alt')
     await expect(comfyPage.canvas).toHaveScreenshot('drag-copy-copied-node.png')
   })
+
+  test('Can undo paste multiple nodes as single action', async ({
+    comfyPage
+  }) => {
+    const initialCount = await comfyPage.getGraphNodesCount()
+    expect(initialCount).toBeGreaterThan(1)
+    await comfyPage.canvas.click()
+    await comfyPage.ctrlA()
+    await comfyPage.page.mouse.move(10, 10)
+    await comfyPage.ctrlC()
+    await comfyPage.ctrlV()
+
+    const pasteCount = await comfyPage.getGraphNodesCount()
+    expect(pasteCount).toBe(initialCount * 2)
+
+    await comfyPage.ctrlZ()
+    const undoCount = await comfyPage.getGraphNodesCount()
+    expect(undoCount).toBe(initialCount)
+  })
 })

src/scripts/changeTracker.ts

@@ -12,6 +12,7 @@ export class ChangeTracker {
   activeState = null
   isOurLoad = false
   workflow: ComfyWorkflow | null
+  changeCount = 0
 
   ds: { scale: number; offset: [number, number] }
   nodeOutputs: any
@@ -46,7 +47,7 @@ export class ChangeTracker {
   }
 
   checkState() {
-    if (!this.app.graph) return
+    if (!this.app.graph || this.changeCount) return
 
     const currentState = this.app.graph.serialize()
     if (!this.activeState) {
@@ -100,6 +101,16 @@ export class ChangeTracker {
     }
   }
 
+  beforeChange() {
+    this.changeCount++
+  }
+
+  afterChange() {
+    if (!--this.changeCount) {
+      this.checkState()
+    }
+  }
+
   static init(app: ComfyApp) {
     const changeTracker = () =>
       app.workflowManager.activeWorkflow?.changeTracker ?? globalTracker
@@ -210,6 +221,15 @@ export class ChangeTracker {
       return v
     }
 
+    // Handle multiple commands as a single transaction
+    document.addEventListener('litegraph:canvas', (e: CustomEvent) => {
+      if (e.detail.subType === 'before-change') {
+        changeTracker().beforeChange()
+      } else if (e.detail.subType === 'after-change') {
+        changeTracker().afterChange()
+      }
+    })
+
     // Store node outputs
     api.addEventListener('executed', ({ detail }) => {
       const prompt =