[feat] Add auto-refresh on task completion for RemoteWidget nodes (#4191)

Co-authored-by: filtered <[email protected]>

Author: christian-byrne

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,

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
+      )
+    })
+  })
 })