feat: add link store lifecycle rehydration contract

- add passive LinkStore for links, floating links, and reroutes

- rehydrate LinkStore on graph clear and configure lifecycle paths

- add unit and browser tests for ChangeTracker/link topology undo-redo parity

Amp-Thread-ID: https://ampcode.com/threads/T-019c9840-2f49-701c-bdf2-ae7f2f50acef
Co-authored-by: Amp <amp@ampcode.com>

browser_tests/tests/changeTracker.spec.ts

@@ -3,6 +3,7 @@ import {
   comfyExpect as expect,
   comfyPageFixture as test
 } from '../fixtures/ComfyPage'
+import type { WorkspaceStore } from '../types/globals'
 
 async function beforeChange(comfyPage: ComfyPage) {
   await comfyPage.page.evaluate(() => {
@@ -63,6 +64,66 @@ test.describe('Change Tracker', { tag: '@workflow' }, () => {
       expect(await comfyPage.workflow.getUndoQueueSize()).toBe(0)
       expect(await comfyPage.workflow.getRedoQueueSize()).toBe(2)
     })
+
+    test('undo/redo restores link topology with reroutes and floating links', async ({
+      comfyPage
+    }) => {
+      await comfyPage.workflow.loadWorkflow('links/batch_move_links')
+
+      const readTopology = () =>
+        comfyPage.page.evaluate(() => {
+          const graph = window.app!.rootGraph
+          return {
+            links: graph.links.size,
+            floatingLinks: graph.floatingLinks.size,
+            reroutes: graph.reroutes.size,
+            serialised: graph.serialize()
+          }
+        })
+
+      const baseline = await readTopology()
+
+      await beforeChange(comfyPage)
+      await comfyPage.page.evaluate(() => {
+        const graph = window.app!.rootGraph
+        const firstLink = graph.links.values().next().value
+        if (!firstLink) throw new Error('Expected at least one link')
+
+        const reroute = graph.createReroute(
+          [firstLink.id * 5, firstLink.id * 3],
+          firstLink
+        )
+        graph.addFloatingLink(firstLink.toFloating('output', reroute.id))
+      })
+      await afterChange(comfyPage)
+
+      const mutated = await readTopology()
+      expect(mutated.floatingLinks).toBeGreaterThan(baseline.floatingLinks)
+      expect(mutated.reroutes).toBeGreaterThan(baseline.reroutes)
+      await expect.poll(() => comfyPage.workflow.getUndoQueueSize()).toBe(1)
+
+      await comfyPage.page.evaluate(async () => {
+        await (
+          window.app!.extensionManager as WorkspaceStore
+        ).workflow.activeWorkflow?.changeTracker.undo()
+      })
+      const afterUndo = await readTopology()
+      expect(afterUndo.links).toBe(baseline.links)
+      expect(afterUndo.floatingLinks).toBe(baseline.floatingLinks)
+      expect(afterUndo.reroutes).toBe(baseline.reroutes)
+      expect(afterUndo.serialised).toEqual(baseline.serialised)
+
+      await comfyPage.page.evaluate(async () => {
+        await (
+          window.app!.extensionManager as WorkspaceStore
+        ).workflow.activeWorkflow?.changeTracker.redo()
+      })
+      const afterRedo = await readTopology()
+      expect(afterRedo.links).toBe(mutated.links)
+      expect(afterRedo.floatingLinks).toBe(mutated.floatingLinks)
+      expect(afterRedo.reroutes).toBe(mutated.reroutes)
+      expect(afterRedo.serialised).toEqual(mutated.serialised)
+    })
   })
 
   test('Can group multiple change actions into a single transaction', async ({

src/lib/litegraph/src/LGraph.test.ts

@@ -39,6 +39,44 @@ class DummyNode extends LGraphNode {
   }
 }
 
+function createNumberNode(title: string): LGraphNode {
+  const node = new LGraphNode(title)
+  node.addOutput('out', 'number')
+  node.addInput('in', 'number')
+  return node
+}
+
+function buildLinkTopology(graph: LGraph): {
+  floatingLinkId: number
+  linkedNodeId: NodeId
+  rerouteId: number
+} {
+  const source = createNumberNode('source')
+  const floatingTarget = createNumberNode('floating-target')
+  const linkedTarget = createNumberNode('linked-target')
+  graph.add(source)
+  graph.add(floatingTarget)
+  graph.add(linkedTarget)
+
+  source.connect(0, floatingTarget, 0)
+  source.connect(0, linkedTarget, 0)
+
+  const linkToDisconnect = graph.getLink(floatingTarget.inputs[0].link)
+  if (!linkToDisconnect) throw new Error('Expected link to disconnect')
+
+  const reroute = graph.createReroute([120, 80], linkToDisconnect)
+  graph.addFloatingLink(linkToDisconnect.toFloating('output', reroute.id))
+
+  const floatingLinkId = [...graph.floatingLinks.keys()][0]
+  if (floatingLinkId == null) throw new Error('Expected floating link')
+
+  return {
+    floatingLinkId,
+    linkedNodeId: linkedTarget.id,
+    rerouteId: reroute.id
+  }
+}
+
 describe('LGraph', () => {
   it('should serialize deterministic node order', async () => {
     LiteGraph.registerNodeType('dummy', DummyNode)
@@ -184,6 +222,76 @@ describe('Floating Links / Reroutes', () => {
   })
 })
 
+describe('LinkStore Lifecycle Rehydration', () => {
+  it('tracks links, floating links, and reroutes after configure', () => {
+    const graph = new LGraph()
+    const { floatingLinkId, linkedNodeId, rerouteId } = buildLinkTopology(graph)
+    const serialised = graph.asSerialisable()
+
+    const restored = new LGraph(serialised)
+    const linkedInput = restored.getNodeById(linkedNodeId)?.inputs[0]
+    const linkedInputLink = restored.getLink(linkedInput?.link)
+
+    expect(restored.linkStore.links.size).toBe(restored.links.size)
+    expect(restored.linkStore.floatingLinks.size).toBe(
+      restored.floatingLinks.size
+    )
+    expect(restored.linkStore.reroutes.size).toBe(restored.reroutes.size)
+    expect(restored.linkStore.getFloatingLink(floatingLinkId)).toBeDefined()
+    expect(restored.linkStore.getReroute(rerouteId)).toBeDefined()
+    expect(restored.linkStore.getLink(linkedInput?.link)).toBe(linkedInputLink)
+  })
+
+  it('clears and rehydrates the store on graph.clear()', () => {
+    const graph = new LGraph()
+    buildLinkTopology(graph)
+
+    expect(graph.linkStore.links.size).toBeGreaterThan(0)
+    expect(graph.linkStore.floatingLinks.size).toBeGreaterThan(0)
+    expect(graph.linkStore.reroutes.size).toBeGreaterThan(0)
+
+    graph.clear()
+
+    expect(graph.linkStore.links.size).toBe(0)
+    expect(graph.linkStore.floatingLinks.size).toBe(0)
+    expect(graph.linkStore.reroutes.size).toBe(0)
+  })
+
+  it('preserves root/subgraph store isolation after round-trip', () => {
+    setActivePinia(createTestingPinia({ stubActions: false }))
+
+    const root = new LGraph()
+    const rootTopology = buildLinkTopology(root)
+
+    const subgraph = root.createSubgraph(createTestSubgraphData())
+    const subgraphSource = createNumberNode('subgraph-source')
+    const subgraphTarget = createNumberNode('subgraph-target')
+    subgraph.add(subgraphSource)
+    subgraph.add(subgraphTarget)
+    subgraphSource.connect(0, subgraphTarget, 0)
+
+    root.add(createTestSubgraphNode(subgraph, { pos: [500, 200] }))
+
+    const serialised = root.asSerialisable()
+    const restoredRoot = new LGraph(serialised)
+    const restoredSubgraph = [...restoredRoot.subgraphs.values()][0]
+
+    if (!restoredSubgraph) throw new Error('Expected restored subgraph')
+
+    const subgraphLinkId = [...restoredSubgraph.links.keys()][0]
+
+    expect(
+      restoredRoot.linkStore.getFloatingLink(rootTopology.floatingLinkId)
+    ).toBeDefined()
+    expect(
+      restoredRoot.linkStore.getReroute(rootTopology.rerouteId)
+    ).toBeDefined()
+    expect(restoredRoot.linkStore.getLink(subgraphLinkId)).toBeUndefined()
+    expect(restoredSubgraph.linkStore.getLink(subgraphLinkId)).toBeDefined()
+    expect(restoredSubgraph.linkStore.floatingLinks.size).toBe(0)
+  })
+})
+
 describe('Graph Clearing and Callbacks', () => {
   test('clear() calls both node.onRemoved() and graph.onNodeRemoved()', ({
     expect

src/lib/litegraph/src/LGraph.ts

@@ -17,6 +17,7 @@ import type { DragAndScaleState } from './DragAndScale'
 import { LGraphCanvas } from './LGraphCanvas'
 import { LGraphGroup } from './LGraphGroup'
 import { LGraphNode } from './LGraphNode'
+import { LinkStore } from './LinkStore'
 import type { NodeId } from './LGraphNode'
 import { LLink } from './LLink'
 import type { LinkId, SerialisedLLinkArray } from './LLink'
@@ -270,6 +271,12 @@ export class LGraph
   private _lastFloatingLinkId: number = 0
 
   private readonly floatingLinksInternal: Map<LinkId, LLink> = new Map()
+  private readonly linkStoreInternal = new LinkStore()
+
+  get linkStore(): LinkStore {
+    return this.linkStoreInternal
+  }
+
   get floatingLinks(): ReadonlyMap<LinkId, LLink> {
     return this.floatingLinksInternal
   }
@@ -393,6 +400,7 @@ export class LGraph
     this._links.clear()
     this.reroutes.clear()
     this.floatingLinksInternal.clear()
+    this.rehydrateLinkStore()
 
     this._lastFloatingLinkId = 0
 
@@ -429,6 +437,14 @@ export class LGraph
     this.canvasAction((c) => c.clear())
   }
 
+  private rehydrateLinkStore(): void {
+    this.linkStoreInternal.rehydrate({
+      links: this._links,
+      floatingLinks: this.floatingLinksInternal,
+      reroutes: this.reroutesInternal
+    })
+  }
+
   get subgraphs(): Map<UUID, Subgraph> {
     return this.rootGraph._subgraphs
   }
@@ -2569,6 +2585,7 @@ export class LGraph
       this.setDirtyCanvas(true, true)
       return error
     } finally {
+      this.rehydrateLinkStore()
       this.events.dispatch('configured')
     }
   }

src/lib/litegraph/src/LinkStore.ts

@@ -0,0 +1,56 @@
+import type { LinkId, LLink } from './LLink'
+import type { Reroute, RerouteId } from './Reroute'
+
+interface LinkStoreTopology {
+  links: ReadonlyMap<LinkId, LLink>
+  floatingLinks: ReadonlyMap<LinkId, LLink>
+  reroutes: ReadonlyMap<RerouteId, Reroute>
+}
+
+/**
+ * Passive graph-scoped topology store.
+ *
+ * Slice 1 contract: this store owns no mutation logic and is rehydrated from
+ * graph lifecycle boundaries (`clear` and `configure`).
+ */
+export class LinkStore {
+  private _links: ReadonlyMap<LinkId, LLink> = new Map()
+  private _floatingLinks: ReadonlyMap<LinkId, LLink> = new Map()
+  private _reroutes: ReadonlyMap<RerouteId, Reroute> = new Map()
+
+  get links(): ReadonlyMap<LinkId, LLink> {
+    return this._links
+  }
+
+  get floatingLinks(): ReadonlyMap<LinkId, LLink> {
+    return this._floatingLinks
+  }
+
+  get reroutes(): ReadonlyMap<RerouteId, Reroute> {
+    return this._reroutes
+  }
+
+  getLink(id: LinkId | null | undefined): LLink | undefined {
+    return id == null ? undefined : this._links.get(id)
+  }
+
+  getFloatingLink(id: LinkId | null | undefined): LLink | undefined {
+    return id == null ? undefined : this._floatingLinks.get(id)
+  }
+
+  getReroute(id: RerouteId | null | undefined): Reroute | undefined {
+    return id == null ? undefined : this._reroutes.get(id)
+  }
+
+  clear(): void {
+    this._links = new Map()
+    this._floatingLinks = new Map()
+    this._reroutes = new Map()
+  }
+
+  rehydrate(topology: LinkStoreTopology): void {
+    this._links = topology.links
+    this._floatingLinks = topology.floatingLinks
+    this._reroutes = topology.reroutes
+  }
+}

src/scripts/changeTracker.test.ts

@@ -0,0 +1,66 @@
+import { describe, expect, it } from 'vitest'
+
+import { LGraph, LGraphNode } from '@/lib/litegraph/src/litegraph'
+import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/workflowSchema'
+
+import { ChangeTracker } from './changeTracker'
+
+function createTopologyGraph() {
+  const graph = new LGraph()
+
+  const source = new LGraphNode('source')
+  source.addOutput('out', 'number')
+
+  const floatingTarget = new LGraphNode('floating-target')
+  floatingTarget.addInput('in', 'number')
+
+  const linkedTarget = new LGraphNode('linked-target')
+  linkedTarget.addInput('in', 'number')
+
+  graph.add(source)
+  graph.add(floatingTarget)
+  graph.add(linkedTarget)
+
+  source.connect(0, floatingTarget, 0)
+  source.connect(0, linkedTarget, 0)
+
+  const link = graph.getLink(floatingTarget.inputs[0].link)
+  if (!link) throw new Error('Expected link')
+
+  graph.createReroute([100, 100], link)
+  floatingTarget.disconnectInput(0, true)
+
+  return graph
+}
+
+describe('ChangeTracker.graphEqual', () => {
+  it('returns false when links differ', () => {
+    const graph = createTopologyGraph()
+    const stateA = graph.asSerialisable() as unknown as ComfyWorkflowJSON
+    const stateB = structuredClone(stateA)
+
+    stateB.links = []
+
+    expect(ChangeTracker.graphEqual(stateA, stateB)).toBe(false)
+  })
+
+  it('returns false when floatingLinks differ', () => {
+    const graph = createTopologyGraph()
+    const stateA = graph.asSerialisable() as unknown as ComfyWorkflowJSON
+    const stateB = structuredClone(stateA)
+
+    stateB.floatingLinks = []
+
+    expect(ChangeTracker.graphEqual(stateA, stateB)).toBe(false)
+  })
+
+  it('returns false when reroutes differ', () => {
+    const graph = createTopologyGraph()
+    const stateA = graph.asSerialisable() as unknown as ComfyWorkflowJSON
+    const stateB = structuredClone(stateA)
+
+    stateB.reroutes = []
+
+    expect(ChangeTracker.graphEqual(stateA, stateB)).toBe(false)
+  })
+})