From 861d737041b5de0f74f8de970d25b4ae33d47a75 Mon Sep 17 00:00:00 2001 From: Terry Jia Date: Sat, 16 May 2026 05:42:04 -0400 Subject: [PATCH] FE-702: rehydrate 3D viewer on subgraph re-entry via persistent ready hook (#12294) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary When a Preview3D / Load3D / SaveGLB node lives inside a subgraph, the 3D viewer correctly displays the model the first time you enter the subgraph but is blank after exiting and re-entering — even though `node.properties['Last Time Model File']` is still populated and the underlying file is on disk. Fix: introduce a persistent companion to `waitForLoad3d` in `useLoad3d.ts`: - `onLoad3dReady(callback)` — registers a callback that fires on *every* (re-)initialization of the `Load3d` instance for a given node, not just the first one. Cleared automatically when the node is removed from the graph (chained into `node.onRemoved` alongside the existing `pendingCallbacks` cleanup). - `waitForLoad3d` keeps its original one-shot semantics so callbacks that install per-node side effects (e.g. wrapping `node.onExecuted`, setting `sceneWidget.serializeValue`) do not chain on remount. - When `onLoad3dReady` is registered after a `Load3d` instance already exists, the callback fires synchronously as well, so the same code path covers both initial setup and subsequent rehydrations. Preview3D / Load3D / SaveGLB move the "reapply state from `node.properties` / `model_file` widget to the Load3d viewer" block from `waitForLoad3d` to `onLoad3dReady`. First mount and every subsequent remount now run identical rehydration code, with `node.properties['Last Time Model File']` (already workflow-JSON-serialised) as the single source of truth. ## Screenshots (if applicable) before https://github.com/user-attachments/assets/e4b0fe6f-c898-4210-b545-7ad6883ed722 after https://github.com/user-attachments/assets/a4a28490-071d-4694-87a8-5eaa501ac168 ┆Issue is synchronized with this [Notion page](https://www.notion.so/PR-12294-FE-702-rehydrate-3D-viewer-on-subgraph-re-entry-via-persistent-ready-hook-3616d73d3650811e93e7dedb32762711) by [Unito](https://www.unito.io) --- src/composables/useLoad3d.test.ts | 133 +++++++++++++++++++++++++++ src/composables/useLoad3d.ts | 61 ++++++++++-- src/extensions/core/load3d.test.ts | 34 ++++++- src/extensions/core/load3d.ts | 68 ++++++++------ src/extensions/core/saveMesh.test.ts | 61 +++++++++--- src/extensions/core/saveMesh.ts | 15 ++- 6 files changed, 313 insertions(+), 59 deletions(-) diff --git a/src/composables/useLoad3d.test.ts b/src/composables/useLoad3d.test.ts index c6ea3bb2e0..4f61241ee4 100644 --- a/src/composables/useLoad3d.test.ts +++ b/src/composables/useLoad3d.test.ts @@ -1559,4 +1559,137 @@ describe('useLoad3d', () => { expect(persistThumbnail).not.toHaveBeenCalled() }) }) + + describe('waitForLoad3d / onLoad3dReady', () => { + it('fires waitForLoad3d callback when load3d initializes, then drops it', async () => { + const composable = useLoad3d(mockNode) + const cb = vi.fn() + composable.waitForLoad3d(cb) + expect(cb).not.toHaveBeenCalled() + + await composable.initializeLoad3d(document.createElement('div')) + expect(cb).toHaveBeenCalledTimes(1) + + composable.cleanup() + await composable.initializeLoad3d(document.createElement('div')) + expect(cb).toHaveBeenCalledTimes(1) + }) + + it('fires onLoad3dReady callback on every (re-)initialization', async () => { + const composable = useLoad3d(mockNode) + const cb = vi.fn() + composable.onLoad3dReady(cb) + expect(cb).not.toHaveBeenCalled() + + await composable.initializeLoad3d(document.createElement('div')) + expect(cb).toHaveBeenCalledTimes(1) + + composable.cleanup() + await composable.initializeLoad3d(document.createElement('div')) + expect(cb).toHaveBeenCalledTimes(2) + + composable.cleanup() + await composable.initializeLoad3d(document.createElement('div')) + expect(cb).toHaveBeenCalledTimes(3) + }) + + it('fires onLoad3dReady synchronously when load3d already exists', async () => { + const composable = useLoad3d(mockNode) + await composable.initializeLoad3d(document.createElement('div')) + + const cb = vi.fn() + composable.onLoad3dReady(cb) + expect(cb).toHaveBeenCalledTimes(1) + }) + + it('clears persistent callbacks when the node is removed', async () => { + const composable = useLoad3d(mockNode) + const cb = vi.fn() + composable.onLoad3dReady(cb) + await composable.initializeLoad3d(document.createElement('div')) + expect(cb).toHaveBeenCalledTimes(1) + + mockNode.onRemoved?.() + + composable.cleanup() + await composable.initializeLoad3d(document.createElement('div')) + expect(cb).toHaveBeenCalledTimes(1) + }) + + it('isolates a throwing callback so subsequent callbacks and event wiring still run', async () => { + const consoleErrorSpy = vi + .spyOn(console, 'error') + .mockImplementation(() => {}) + + const composable = useLoad3d(mockNode) + const throwing = vi.fn(() => { + throw new Error('boom') + }) + const after = vi.fn() + composable.waitForLoad3d(throwing) + composable.onLoad3dReady(after) + + await composable.initializeLoad3d(document.createElement('div')) + + expect(throwing).toHaveBeenCalledTimes(1) + expect(after).toHaveBeenCalledTimes(1) + expect(mockLoad3d.addEventListener).toHaveBeenCalled() + expect(mockToastStore.addAlert).not.toHaveBeenCalled() + expect(consoleErrorSpy).toHaveBeenCalledWith( + 'Load3d ready callback failed:', + expect.any(Error) + ) + + consoleErrorSpy.mockRestore() + }) + + it('isolates a throwing callback in the synchronous already-mounted path', async () => { + const consoleErrorSpy = vi + .spyOn(console, 'error') + .mockImplementation(() => {}) + + const composable = useLoad3d(mockNode) + await composable.initializeLoad3d(document.createElement('div')) + + const throwing = vi.fn(() => { + throw new Error('boom') + }) + + expect(() => composable.waitForLoad3d(throwing)).not.toThrow() + expect(() => composable.onLoad3dReady(throwing)).not.toThrow() + expect(throwing).toHaveBeenCalledTimes(2) + + consoleErrorSpy.mockRestore() + }) + + it('cleans up callback maps when the node is removed before initializeLoad3d runs', async () => { + const leakedWait = vi.fn() + const leakedReady = vi.fn() + + const composable = useLoad3d(mockNode) + composable.waitForLoad3d(leakedWait) + composable.onLoad3dReady(leakedReady) + + mockNode.onRemoved?.() + + await composable.initializeLoad3d(document.createElement('div')) + + expect(leakedWait).not.toHaveBeenCalled() + expect(leakedReady).not.toHaveBeenCalled() + }) + + it('chains the onRemoved cleanup only once per node', () => { + const originalOnRemoved = vi.fn() + mockNode.onRemoved = originalOnRemoved + + const composable = useLoad3d(mockNode) + composable.waitForLoad3d(vi.fn()) + composable.onLoad3dReady(vi.fn()) + composable.onLoad3dReady(vi.fn()) + + mockNode.onRemoved?.() + + expect(originalOnRemoved).toHaveBeenCalledTimes(1) + }) + }) }) diff --git a/src/composables/useLoad3d.ts b/src/composables/useLoad3d.ts index bad8d77f38..405bc1e70b 100644 --- a/src/composables/useLoad3d.ts +++ b/src/composables/useLoad3d.ts @@ -39,6 +39,30 @@ import { useLoad3dService } from '@/services/load3dService' type Load3dReadyCallback = (load3d: Load3d) => void export const nodeToLoad3dMap = new Map() const pendingCallbacks = new Map() +const persistentReadyCallbacks = new Map() + +const nodesWithCleanup = new WeakSet() + +const ensureNodeCleanupChained = (node: LGraphNode): void => { + if (nodesWithCleanup.has(node)) return + nodesWithCleanup.add(node) + node.onRemoved = useChainCallback(node.onRemoved, () => { + useLoad3dService().removeLoad3d(node) + pendingCallbacks.delete(node) + persistentReadyCallbacks.delete(node) + }) +} + +const invokeReadyCallback = ( + callback: Load3dReadyCallback, + instance: Load3d +): void => { + try { + callback(instance) + } catch (error) { + console.error('Load3d ready callback failed:', error) + } +} export const useLoad3d = (nodeOrRef: MaybeRef) => { const nodeRef = toRef(nodeOrRef) @@ -177,10 +201,7 @@ export const useLoad3d = (nodeOrRef: MaybeRef) => { } ) - node.onRemoved = useChainCallback(node.onRemoved, () => { - useLoad3dService().removeLoad3d(node) - pendingCallbacks.delete(node) - }) + ensureNodeCleanupChained(node) nodeToLoad3dMap.set(node, load3d) @@ -188,13 +209,18 @@ export const useLoad3d = (nodeOrRef: MaybeRef) => { if (callbacks && load3d) { callbacks.forEach((callback) => { - if (load3d) { - callback(load3d) - } + if (load3d) invokeReadyCallback(callback, load3d) }) pendingCallbacks.delete(node) } + const persistent = persistentReadyCallbacks.get(node) + if (persistent && load3d) { + persistent.forEach((callback) => { + if (load3d) invokeReadyCallback(callback, load3d) + }) + } + handleEvents('add') } catch (error) { console.error('Error initializing Load3d:', error) @@ -351,8 +377,7 @@ export const useLoad3d = (nodeOrRef: MaybeRef) => { const existingInstance = nodeToLoad3dMap.get(node) if (existingInstance) { - callback(existingInstance) - + invokeReadyCallback(callback, existingInstance) return } @@ -361,6 +386,23 @@ export const useLoad3d = (nodeOrRef: MaybeRef) => { } pendingCallbacks.get(node)!.push(callback) + ensureNodeCleanupChained(node) + } + + const onLoad3dReady = (callback: Load3dReadyCallback) => { + const rawNode = toRaw(nodeRef.value) + if (!rawNode) return + + const node = rawNode as LGraphNode + + if (!persistentReadyCallbacks.has(node)) { + persistentReadyCallbacks.set(node, []) + } + persistentReadyCallbacks.get(node)!.push(callback) + ensureNodeCleanupChained(node) + + const existingInstance = nodeToLoad3dMap.get(node) + if (existingInstance) invokeReadyCallback(callback, existingInstance) } watch( @@ -979,6 +1021,7 @@ export const useLoad3d = (nodeOrRef: MaybeRef) => { // Methods initializeLoad3d, waitForLoad3d, + onLoad3dReady, handleMouseEnter, handleMouseLeave, handleStartRecording, diff --git a/src/extensions/core/load3d.test.ts b/src/extensions/core/load3d.test.ts index 563594b845..cd4570ff5a 100644 --- a/src/extensions/core/load3d.test.ts +++ b/src/extensions/core/load3d.test.ts @@ -7,12 +7,14 @@ import type { ComfyExtension } from '@/types/comfy' const { registerExtensionMock, waitForLoad3dMock, + onLoad3dReadyMock, configureMock, getLoad3dMock, toastAddAlertMock } = vi.hoisted(() => ({ registerExtensionMock: vi.fn(), waitForLoad3dMock: vi.fn(), + onLoad3dReadyMock: vi.fn(), configureMock: vi.fn(), getLoad3dMock: vi.fn(), toastAddAlertMock: vi.fn() @@ -30,7 +32,10 @@ vi.mock('@/services/load3dService', () => ({ })) vi.mock('@/composables/useLoad3d', () => ({ - useLoad3d: () => ({ waitForLoad3d: waitForLoad3dMock }), + useLoad3d: () => ({ + waitForLoad3d: waitForLoad3dMock, + onLoad3dReady: onLoad3dReadyMock + }), nodeToLoad3dMap: new Map() })) @@ -187,6 +192,9 @@ function setupBaseMocks() { waitForLoad3dMock.mockImplementation((cb: (load3d: FakeLoad3d) => void) => { cb(makeLoad3dMock()) }) + onLoad3dReadyMock.mockImplementation((cb: (load3d: FakeLoad3d) => void) => { + cb(makeLoad3dMock()) + }) } describe('load3d module registration', () => { @@ -271,6 +279,30 @@ describe('Comfy.Preview3D.nodeCreated', () => { }) }) + it('registers a persistent onLoad3dReady hook so subgraph re-entry rehydrates the model', async () => { + const onReadyCallbacks: Array<(load3d: FakeLoad3d) => void> = [] + onLoad3dReadyMock.mockImplementation((cb: (load3d: FakeLoad3d) => void) => { + onReadyCallbacks.push(cb) + }) + + const { preview3DExt } = await loadExtensionsFresh() + const node = makePreview3DNode({ + properties: { 'Last Time Model File': 'persisted/model.glb' } + }) + + await preview3DExt.nodeCreated(node) + expect(onReadyCallbacks).toHaveLength(1) + expect(configureMock).not.toHaveBeenCalled() + + // First mount. + onReadyCallbacks[0](makeLoad3dMock()) + expect(configureMock).toHaveBeenCalledTimes(1) + + // Subgraph exit + re-entry: same callback fires again with a fresh load3d. + onReadyCallbacks[0](makeLoad3dMock()) + expect(configureMock).toHaveBeenCalledTimes(2) + }) + it('persists Last Time Model File and normalizes backslashes after onExecuted', async () => { const { preview3DExt } = await loadExtensionsFresh() const node = makePreview3DNode() diff --git a/src/extensions/core/load3d.ts b/src/extensions/core/load3d.ts index 4626587523..79cdce8bc4 100644 --- a/src/extensions/core/load3d.ts +++ b/src/extensions/core/load3d.ts @@ -337,29 +337,34 @@ useExtensionService().registerExtension({ await nextTick() - useLoad3d(node).waitForLoad3d((load3d) => { + useLoad3d(node).onLoad3dReady((load3d) => { + const modelWidget = node.widgets?.find((w) => w.name === 'model_file') + const width = node.widgets?.find((w) => w.name === 'width') + const height = node.widgets?.find((w) => w.name === 'height') + if (!modelWidget || !width || !height) return + const cameraConfig = node.properties['Camera Config'] as | CameraConfig | undefined const cameraState = cameraConfig?.state const config = new Load3DConfiguration(load3d, node.properties) + config.configure({ + loadFolder: 'input', + modelWidget, + cameraState, + width, + height + }) + }) + useLoad3d(node).waitForLoad3d(() => { const modelWidget = node.widgets?.find((w) => w.name === 'model_file') const width = node.widgets?.find((w) => w.name === 'width') const height = node.widgets?.find((w) => w.name === 'height') const sceneWidget = node.widgets?.find((w) => w.name === 'image') if (modelWidget && width && height && sceneWidget) { - const settings = { - loadFolder: 'input', - modelWidget: modelWidget, - cameraState: cameraState, - width: width, - height: height - } - config.configure(settings) - sceneWidget.serializeValue = async () => { const currentLoad3d = nodeToLoad3dMap.get(node) if (!currentLoad3d) { @@ -477,32 +482,35 @@ useExtensionService().registerExtension({ const onExecuted = node.onExecuted + useLoad3d(node).onLoad3dReady((load3d) => { + const modelWidget = node.widgets?.find((w) => w.name === 'model_file') + if (!modelWidget) return + + const lastTimeModelFile = node.properties['Last Time Model File'] + if (!lastTimeModelFile) return + + modelWidget.value = lastTimeModelFile + + const cameraConfig = node.properties['Camera Config'] as + | CameraConfig + | undefined + const cameraState = cameraConfig?.state + + const config = new Load3DConfiguration(load3d, node.properties) + config.configure({ + loadFolder: 'output', + modelWidget, + cameraState, + silentOnNotFound: true + }) + }) + useLoad3d(node).waitForLoad3d((load3d) => { const config = new Load3DConfiguration(load3d, node.properties) const modelWidget = node.widgets?.find((w) => w.name === 'model_file') if (modelWidget) { - const lastTimeModelFile = node.properties['Last Time Model File'] - - if (lastTimeModelFile) { - modelWidget.value = lastTimeModelFile - - const cameraConfig = node.properties['Camera Config'] as - | CameraConfig - | undefined - const cameraState = cameraConfig?.state - - const settings = { - loadFolder: 'output', - modelWidget: modelWidget, - cameraState: cameraState, - silentOnNotFound: true - } - - config.configure(settings) - } - node.onExecuted = function (output: Load3dPreviewOutput) { onExecuted?.call(this, output) diff --git a/src/extensions/core/saveMesh.test.ts b/src/extensions/core/saveMesh.test.ts index eb29fcc560..911a78e370 100644 --- a/src/extensions/core/saveMesh.test.ts +++ b/src/extensions/core/saveMesh.test.ts @@ -3,12 +3,17 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' import type { LGraphNode } from '@/lib/litegraph/src/LGraphNode' import type { ComfyExtension } from '@/types/comfy' -const { registerExtensionMock, waitForLoad3dMock, configureForSaveMeshMock } = - vi.hoisted(() => ({ - registerExtensionMock: vi.fn(), - waitForLoad3dMock: vi.fn(), - configureForSaveMeshMock: vi.fn() - })) +const { + registerExtensionMock, + waitForLoad3dMock, + onLoad3dReadyMock, + configureForSaveMeshMock +} = vi.hoisted(() => ({ + registerExtensionMock: vi.fn(), + waitForLoad3dMock: vi.fn(), + onLoad3dReadyMock: vi.fn(), + configureForSaveMeshMock: vi.fn() +})) vi.mock('@/services/extensionService', () => ({ useExtensionService: () => ({ registerExtension: registerExtensionMock }) @@ -19,7 +24,10 @@ vi.mock('@/services/load3dService', () => ({ })) vi.mock('@/composables/useLoad3d', () => ({ - useLoad3d: () => ({ waitForLoad3d: waitForLoad3dMock }) + useLoad3d: () => ({ + waitForLoad3d: waitForLoad3dMock, + onLoad3dReady: onLoad3dReadyMock + }) })) vi.mock('@/extensions/core/load3d/Load3DConfiguration', () => ({ @@ -74,11 +82,15 @@ function makeNode( describe('saveMesh', () => { beforeEach(() => { vi.clearAllMocks() + const fakeLoad3d = () => ({ + whenLoadIdle: () => Promise.resolve(), + captureThumbnail: vi.fn() + }) waitForLoad3dMock.mockImplementation((cb: (load3d: unknown) => void) => { - cb({ - whenLoadIdle: () => Promise.resolve(), - captureThumbnail: vi.fn() - }) + cb(fakeLoad3d()) + }) + onLoad3dReadyMock.mockImplementation((cb: (load3d: unknown) => void) => { + cb(fakeLoad3d()) }) }) @@ -130,6 +142,33 @@ describe('saveMesh', () => { ) }) + it('registers a persistent onLoad3dReady hook so subgraph re-entry rehydrates the model', async () => { + const onReadyCallbacks: Array<(load3d: unknown) => void> = [] + onLoad3dReadyMock.mockImplementation((cb: (load3d: unknown) => void) => { + onReadyCallbacks.push(cb) + }) + + const ext = await loadSaveMeshExtensionFresh() + const node = makeNode({ + properties: { + 'Last Time Model File': 'sub/model.glb', + 'Last Time Model Folder': 'output' + } + }) + + await ext.nodeCreated(node) + expect(onReadyCallbacks).toHaveLength(1) + expect(configureForSaveMeshMock).not.toHaveBeenCalled() + + const fakeLoad3d = { whenLoadIdle: () => Promise.resolve() } + + onReadyCallbacks[0](fakeLoad3d) + expect(configureForSaveMeshMock).toHaveBeenCalledTimes(1) + + onReadyCallbacks[0]({ whenLoadIdle: () => Promise.resolve() }) + expect(configureForSaveMeshMock).toHaveBeenCalledTimes(2) + }) + it('defaults the load folder to output when only the file path is persisted', async () => { const ext = await loadSaveMeshExtensionFresh() const node = makeNode({ diff --git a/src/extensions/core/saveMesh.ts b/src/extensions/core/saveMesh.ts index a9998282bd..68147d72a3 100644 --- a/src/extensions/core/saveMesh.ts +++ b/src/extensions/core/saveMesh.ts @@ -81,7 +81,7 @@ useExtensionService().registerExtension({ await nextTick() - useLoad3d(node).waitForLoad3d((load3d) => { + useLoad3d(node).onLoad3dReady((load3d) => { if (!load3d) return const modelWidget = node.widgets?.find((w) => w.name === 'image') @@ -96,15 +96,14 @@ useExtensionService().registerExtension({ | 'output' | undefined) ?? 'output' - if (lastTimeModelFile) { - modelWidget.value = lastTimeModelFile + if (!lastTimeModelFile) return - const config = new Load3DConfiguration(load3d, node.properties) + modelWidget.value = lastTimeModelFile - config.configureForSaveMesh(lastTimeModelFolder, lastTimeModelFile, { - silentOnNotFound: true - }) - } + const config = new Load3DConfiguration(load3d, node.properties) + config.configureForSaveMesh(lastTimeModelFolder, lastTimeModelFile, { + silentOnNotFound: true + }) }) const onExecuted = node.onExecuted