From 3a1bd1829a06518dde21be80614be69fd129d52a Mon Sep 17 00:00:00 2001 From: Christian Byrne Date: Sun, 22 Jun 2025 17:30:24 -0700 Subject: [PATCH] [feat] Add auto-refresh on task completion for RemoteWidget nodes (#4191) Co-authored-by: filtered <176114999+webfiltered@users.noreply.github.com> --- src/composables/widgets/useRemoteWidget.ts | 42 ++++ .../widgets/useRemoteWidget.test.ts | 184 +++++++++++++++++- 2 files changed, 225 insertions(+), 1 deletion(-) diff --git a/src/composables/widgets/useRemoteWidget.ts b/src/composables/widgets/useRemoteWidget.ts index e06aa39b3e..8cfcfda1fb 100644 --- a/src/composables/widgets/useRemoteWidget.ts +++ b/src/composables/widgets/useRemoteWidget.ts @@ -2,7 +2,9 @@ import { LGraphNode } from '@comfyorg/litegraph' import { IWidget } from '@comfyorg/litegraph' import axios from 'axios' +import { useChainCallback } from '@/composables/functional/useChainCallback' import type { RemoteWidgetConfig } from '@/schemas/nodeDefSchema' +import { api } from '@/scripts/api' const MAX_RETRIES = 5 const TIMEOUT = 4096 @@ -220,6 +222,46 @@ export function useRemoteWidget< node.addWidget('button', 'refresh', 'refresh', widget.refresh) } + /** + * Add auto-refresh toggle widget and execution success listener + */ + function addAutoRefreshToggle() { + let autoRefreshEnabled = false + + // Handler for execution success + const handleExecutionSuccess = () => { + if (autoRefreshEnabled && widget.refresh) { + widget.refresh() + } + } + + // Add toggle widget + const autoRefreshWidget = node.addWidget( + 'toggle', + 'Auto-refresh after generation', + false, + (value: boolean) => { + autoRefreshEnabled = value + }, + { + serialize: false + } + ) + + // Register event listener + api.addEventListener('execution_success', handleExecutionSuccess) + + // Cleanup on node removal + node.onRemoved = useChainCallback(node.onRemoved, function () { + api.removeEventListener('execution_success', handleExecutionSuccess) + }) + + return autoRefreshWidget + } + + // Always add auto-refresh toggle for remote widgets + addAutoRefreshToggle() + return { getCachedValue, getValue, diff --git a/tests-ui/tests/composables/widgets/useRemoteWidget.test.ts b/tests-ui/tests/composables/widgets/useRemoteWidget.test.ts index 22d19b779e..0bf7294537 100644 --- a/tests-ui/tests/composables/widgets/useRemoteWidget.test.ts +++ b/tests-ui/tests/composables/widgets/useRemoteWidget.test.ts @@ -26,6 +26,22 @@ vi.mock('@/stores/settingStore', () => ({ }) })) +vi.mock('@/scripts/api', () => ({ + api: { + addEventListener: vi.fn(), + removeEventListener: vi.fn() + } +})) + +vi.mock('@/composables/functional/useChainCallback', () => ({ + useChainCallback: vi.fn((original, ...callbacks) => { + return function (this: any, ...args: any[]) { + original?.apply(this, args) + callbacks.forEach((cb: any) => cb.apply(this, args)) + } + }) +})) + const FIRST_BACKOFF = 1000 // backoff is 1s on first retry const DEFAULT_VALUE = 'Loading...' @@ -40,7 +56,9 @@ function createMockConfig(overrides = {}): RemoteWidgetConfig { const createMockOptions = (inputOverrides = {}) => ({ remoteConfig: createMockConfig(inputOverrides), defaultValue: DEFAULT_VALUE, - node: {} as any, + node: { + addWidget: vi.fn() + } as any, widget: {} as any }) @@ -499,4 +517,168 @@ describe('useRemoteWidget', () => { expect(data2).toEqual(DEFAULT_VALUE) }) }) + + describe('auto-refresh on task completion', () => { + it('should add auto-refresh toggle widget', () => { + const mockNode = { + addWidget: vi.fn(), + widgets: [] + } + const mockWidget = { + refresh: vi.fn() + } + + useRemoteWidget({ + remoteConfig: createMockConfig(), + defaultValue: DEFAULT_VALUE, + node: mockNode as any, + widget: mockWidget as any + }) + + // Should add auto-refresh toggle widget + expect(mockNode.addWidget).toHaveBeenCalledWith( + 'toggle', + 'Auto-refresh after generation', + false, + expect.any(Function), + { + serialize: false + } + ) + }) + + it('should register event listener when enabled', async () => { + const { api } = await import('@/scripts/api') + + const mockNode = { + addWidget: vi.fn(), + widgets: [] + } + const mockWidget = { + refresh: vi.fn() + } + + useRemoteWidget({ + remoteConfig: createMockConfig(), + defaultValue: DEFAULT_VALUE, + node: mockNode as any, + widget: mockWidget as any + }) + + // Event listener should be registered immediately + expect(api.addEventListener).toHaveBeenCalledWith( + 'execution_success', + expect.any(Function) + ) + }) + + it('should refresh widget when workflow completes successfully', async () => { + const { api } = await import('@/scripts/api') + let executionSuccessHandler: (() => void) | undefined + + // Capture the event handler + vi.mocked(api.addEventListener).mockImplementation((event, handler) => { + if (event === 'execution_success') { + executionSuccessHandler = handler as () => void + } + }) + + const mockNode = { + addWidget: vi.fn(), + widgets: [] + } + const mockWidget = {} as any + + useRemoteWidget({ + remoteConfig: createMockConfig(), + defaultValue: DEFAULT_VALUE, + node: mockNode as any, + widget: mockWidget + }) + + // Spy on the refresh function that was added by useRemoteWidget + const refreshSpy = vi.spyOn(mockWidget, 'refresh') + + // Get the toggle callback and enable auto-refresh + const toggleCallback = mockNode.addWidget.mock.calls.find( + (call) => call[0] === 'toggle' + )?.[3] + toggleCallback?.(true) + + // Simulate workflow completion + executionSuccessHandler?.() + + expect(refreshSpy).toHaveBeenCalled() + }) + + it('should not refresh when toggle is disabled', async () => { + const { api } = await import('@/scripts/api') + let executionSuccessHandler: (() => void) | undefined + + // Capture the event handler + vi.mocked(api.addEventListener).mockImplementation((event, handler) => { + if (event === 'execution_success') { + executionSuccessHandler = handler as () => void + } + }) + + const mockNode = { + addWidget: vi.fn(), + widgets: [] + } + const mockWidget = {} as any + + useRemoteWidget({ + remoteConfig: createMockConfig(), + defaultValue: DEFAULT_VALUE, + node: mockNode as any, + widget: mockWidget + }) + + // Spy on the refresh function that was added by useRemoteWidget + const refreshSpy = vi.spyOn(mockWidget, 'refresh') + + // Toggle is disabled by default + // Simulate workflow completion + executionSuccessHandler?.() + + expect(refreshSpy).not.toHaveBeenCalled() + }) + + it('should cleanup event listener on node removal', async () => { + const { api } = await import('@/scripts/api') + let executionSuccessHandler: (() => void) | undefined + + // Capture the event handler + vi.mocked(api.addEventListener).mockImplementation((event, handler) => { + if (event === 'execution_success') { + executionSuccessHandler = handler as () => void + } + }) + + const mockNode = { + addWidget: vi.fn(), + widgets: [], + onRemoved: undefined as any + } + const mockWidget = { + refresh: vi.fn() + } + + useRemoteWidget({ + remoteConfig: createMockConfig(), + defaultValue: DEFAULT_VALUE, + node: mockNode as any, + widget: mockWidget as any + }) + + // Simulate node removal + mockNode.onRemoved?.() + + expect(api.removeEventListener).toHaveBeenCalledWith( + 'execution_success', + executionSuccessHandler + ) + }) + }) })