Fix changeTracker modified state (#1481)

* Add jsondiffpatch

* Add logs

* Add graphDiff helper

* Fix changeTracker

* Add loglevel

* Add playwright test

* Fix jest test

* nit

* nit

* Fix test url

* nit

Author: huchenlei

browser_tests/changeTracker.spec.ts

@@ -16,6 +16,43 @@ async function afterChange(comfyPage: ComfyPage) {
 }
 
 test.describe('Change Tracker', () => {
+  test.describe('Undo/Redo', () => {
+    test.beforeEach(async ({ comfyPage }) => {
+      await comfyPage.setSetting('Comfy.UseNewMenu', 'Top')
+    })
+
+    // Flaky https://github.com/Comfy-Org/ComfyUI_frontend/pull/1481
+    // The collapse can be recognized as several changes.
+    test.skip('Can undo multiple operations', async ({ comfyPage }) => {
+      function isModified() {
+        return comfyPage.page.evaluate(async () => {
+          return window['app'].extensionManager.workflow.activeWorkflow
+            .isModified
+        })
+      }
+
+      await comfyPage.menu.topbar.saveWorkflow('undo-redo-test')
+      expect(await isModified()).toBe(false)
+
+      const node = (await comfyPage.getFirstNodeRef())!
+      await node.click('collapse')
+      await expect(node).toBeCollapsed()
+      expect(await isModified()).toBe(true)
+
+      await comfyPage.ctrlB()
+      await expect(node).toBeBypassed()
+      expect(await isModified()).toBe(true)
+
+      await comfyPage.ctrlZ()
+      await expect(node).not.toBeBypassed()
+      expect(await isModified()).toBe(true)
+
+      await comfyPage.ctrlZ()
+      await expect(node).not.toBeCollapsed()
+      expect(await isModified()).toBe(false)
+    })
+  })
+
   test('Can group multiple change actions into a single transaction', async ({
     comfyPage
   }) => {

browser_tests/fixtures/ComfyPage.ts

@@ -167,7 +167,7 @@ export class ComfyPage {
   }
 
   async setupUser(username: string) {
-    const res = await this.request.get(`${this.url}/users`)
+    const res = await this.request.get(`${this.url}/api/users`)
     if (res.status() !== 200)
       throw new Error(`Failed to retrieve users: ${await res.text()}`)
 
@@ -181,7 +181,7 @@ export class ComfyPage {
   }
 
   async createUser(username: string) {
-    const resp = await this.request.post(`${this.url}/users`, {
+    const resp = await this.request.post(`${this.url}/api/users`, {
       data: { username }
     })
 

package-lock.json

@@ -79,7 +79,9 @@
     "axios": "^1.7.4",
     "dotenv": "^16.4.5",
     "fuse.js": "^7.0.0",
+    "jsondiffpatch": "^0.6.0",
     "lodash": "^4.17.21",
+    "loglevel": "^1.9.2",
     "pinia": "^2.1.7",
     "primeicons": "^7.0.0",
     "primevue": "^4.0.5",

package.json

@@ -2363,7 +2363,7 @@ export class ComfyApp {
     }
     await this.#invokeExtensionsAsync('afterConfigureGraph', missingNodeTypes)
     // @ts-expect-error zod types issue. Will be fixed after we enable ts-strict
-    workflowService.afterLoadNewGraph(workflow, this.graph.serialize())
+    await workflowService.afterLoadNewGraph(workflow, this.graph.serialize())
     requestAnimationFrame(() => {
       this.graph.setDirtyCanvas(true, true)
     })

src/scripts/app.ts

@@ -7,6 +7,8 @@ import type { ComfyWorkflowJSON } from '@/types/comfyWorkflow'
 import type { ExecutedWsMessage } from '@/types/apiTypes'
 import { useExecutionStore } from '@/stores/executionStore'
 import _ from 'lodash'
+import * as jsondiffpatch from 'jsondiffpatch'
+import log from 'loglevel'
 
 function clone(obj: any) {
   try {
@@ -20,6 +22,10 @@ function clone(obj: any) {
   return JSON.parse(JSON.stringify(obj))
 }
 
+const logger = log.getLogger('ChangeTracker')
+// Change to debug for more verbose logging
+logger.setLevel('info')
+
 export class ChangeTracker {
   static MAX_HISTORY = 50
   /**
@@ -29,6 +35,10 @@ export class ChangeTracker {
   undoQueue: ComfyWorkflowJSON[] = []
   redoQueue: ComfyWorkflowJSON[] = []
   changeCount: number = 0
+  /**
+   * Whether the redo/undo restoring is in progress.
+   */
+  private restoringState: boolean = false
 
   ds?: { scale: number; offset: [number, number] }
   nodeOutputs?: Record<string, any>
@@ -55,8 +65,12 @@ export class ChangeTracker {
    * Save the current state as the initial state.
    */
   reset(state?: ComfyWorkflowJSON) {
+    // Do not reset the state if we are restoring.
+    if (this.restoringState) return
+
+    logger.debug('Reset State')
     this.activeState = state ?? this.activeState
-    this.initialState = this.activeState
+    this.initialState = clone(this.activeState)
   }
 
   store() {
@@ -85,6 +99,13 @@ export class ChangeTracker {
         this.initialState,
         this.activeState
       )
+      if (workflow.isModified) {
+        const diff = ChangeTracker.graphDiff(
+          this.initialState,
+          this.activeState
+        )
+        logger.debug('Graph diff:', diff)
+      }
     }
   }
 
@@ -101,6 +122,8 @@ export class ChangeTracker {
       if (this.undoQueue.length > ChangeTracker.MAX_HISTORY) {
         this.undoQueue.shift()
       }
+      logger.debug('Diff detected. Undo queue length:', this.undoQueue.length)
+
       this.activeState = clone(currentState)
       this.redoQueue.length = 0
       api.dispatchEvent(
@@ -114,21 +137,38 @@ export class ChangeTracker {
     const prevState = source.pop()
     if (prevState) {
       target.push(this.activeState!)
-      await this.app.loadGraphData(prevState, false, false, this.workflow, {
-        showMissingModelsDialog: false,
-        showMissingNodesDialog: false
-      })
-      this.activeState = prevState
-      this.updateModified()
+      this.restoringState = true
+      try {
+        await this.app.loadGraphData(prevState, false, false, this.workflow, {
+          showMissingModelsDialog: false,
+          showMissingNodesDialog: false
+        })
+        this.activeState = prevState
+        this.updateModified()
+      } finally {
+        this.restoringState = false
+      }
     }
   }
 
   async undo() {
     await this.updateState(this.undoQueue, this.redoQueue)
+    logger.debug(
+      'Undo. Undo queue length:',
+      this.undoQueue.length,
+      'Redo queue length:',
+      this.redoQueue.length
+    )
   }
 
   async redo() {
     await this.updateState(this.redoQueue, this.undoQueue)
+    logger.debug(
+      'Redo. Undo queue length:',
+      this.undoQueue.length,
+      'Redo queue length:',
+      this.redoQueue.length
+    )
   }
 
   async undoRedo(e: KeyboardEvent) {
@@ -198,6 +238,7 @@ export class ChangeTracker {
 
           // If our active element is some type of input then handle changes after they're done
           if (ChangeTracker.bindInput(app, bindInputEl)) return
+          logger.debug('checkState on keydown')
           changeTracker.checkState()
         })
       },
@@ -207,34 +248,40 @@ export class ChangeTracker {
     window.addEventListener('keyup', (e) => {
       if (keyIgnored) {
         keyIgnored = false
+        logger.debug('checkState on keyup')
         checkState()
       }
     })
 
     // Handle clicking DOM elements (e.g. widgets)
     window.addEventListener('mouseup', () => {
+      logger.debug('checkState on mouseup')
       checkState()
     })
 
     // Handle prompt queue event for dynamic widget changes
     api.addEventListener('promptQueued', () => {
+      logger.debug('checkState on promptQueued')
       checkState()
     })
 
     api.addEventListener('graphCleared', () => {
+      logger.debug('checkState on graphCleared')
       checkState()
     })
 
     // Handle litegraph clicks
     const processMouseUp = LGraphCanvas.prototype.processMouseUp
     LGraphCanvas.prototype.processMouseUp = function (e) {
       const v = processMouseUp.apply(this, [e])
+      logger.debug('checkState on processMouseUp')
       checkState()
       return v
     }
     const processMouseDown = LGraphCanvas.prototype.processMouseDown
     LGraphCanvas.prototype.processMouseDown = function (e) {
       const v = processMouseDown.apply(this, [e])
+      logger.debug('checkState on processMouseDown')
       checkState()
       return v
     }
@@ -251,13 +298,15 @@ export class ChangeTracker {
         callback(v)
         checkState()
       }
+      logger.debug('checkState on prompt')
       return prompt.apply(this, [title, value, extendedCallback, event])
     }
 
     // Handle litegraph context menu for COMBO widgets
     const close = LiteGraph.ContextMenu.prototype.close
     LiteGraph.ContextMenu.prototype.close = function (e: MouseEvent) {
       const v = close.apply(this, [e])
+      logger.debug('checkState on contextMenuClose')
       checkState()
       return v
     }
@@ -267,6 +316,7 @@ export class ChangeTracker {
     LiteGraph.LGraph.prototype.onNodeAdded = function (node: LGraphNode) {
       const v = onNodeAdded?.apply(this, [node])
       if (!app?.configuringGraph) {
+        logger.debug('checkState on onNodeAdded')
         checkState()
       }
       return v
@@ -357,4 +407,20 @@ export class ChangeTracker {
 
     return false
   }
+
+  static graphDiff(a: ComfyWorkflowJSON, b: ComfyWorkflowJSON) {
+    function sortGraphNodes(graph: ComfyWorkflowJSON) {
+      return {
+        links: graph.links,
+        groups: graph.groups,
+        nodes: graph.nodes.sort((a, b) => {
+          if (typeof a.id === 'number' && typeof b.id === 'number') {
+            return a.id - b.id
+          }
+          return 0
+        })
+      }
+    }
+    return jsondiffpatch.diff(sortGraphNodes(a), sortGraphNodes(b))
+  }
 }

src/scripts/changeTracker.ts

@@ -19,4 +19,10 @@ module.exports = async function () {
       useI18n: jest.fn()
     }
   })
+
+  jest.mock('jsondiffpatch', () => {
+    return {
+      diff: jest.fn()
+    }
+  })
 }

tests-ui/tests/fast/globalSetup.ts

@@ -74,4 +74,10 @@ module.exports = async function () {
       useI18n: jest.fn()
     }
   })
+
+  jest.mock('jsondiffpatch', () => {
+    return {
+      diff: jest.fn()
+    }
+  })
 }