Add support for Feature Flags (#4439)

browser_tests/tests/featureFlags.spec.ts

@@ -0,0 +1,382 @@
+import { expect } from '@playwright/test'
+
+import { comfyPageFixture as test } from '../fixtures/ComfyPage'
+
+test.describe('Feature Flags', () => {
+  test('Client and server exchange feature flags on connection', async ({
+    comfyPage
+  }) => {
+    // Navigate to a new page to capture the initial WebSocket connection
+    const newPage = await comfyPage.page.context().newPage()
+
+    // Set up monitoring before navigation
+    await newPage.addInitScript(() => {
+      // This runs before any page scripts
+      window.__capturedMessages = {
+        clientFeatureFlags: null,
+        serverFeatureFlags: null
+      }
+
+      // Capture outgoing client messages
+      const originalSend = WebSocket.prototype.send
+      WebSocket.prototype.send = function (data) {
+        try {
+          const parsed = JSON.parse(data)
+          if (parsed.type === 'feature_flags') {
+            window.__capturedMessages.clientFeatureFlags = parsed
+          }
+        } catch (e) {
+          // Not JSON, ignore
+        }
+        return originalSend.call(this, data)
+      }
+
+      // Monitor for server feature flags
+      const checkInterval = setInterval(() => {
+        if (
+          window['app']?.api?.serverFeatureFlags &&
+          Object.keys(window['app'].api.serverFeatureFlags).length > 0
+        ) {
+          window.__capturedMessages.serverFeatureFlags =
+            window['app'].api.serverFeatureFlags
+          clearInterval(checkInterval)
+        }
+      }, 100)
+
+      // Clear after 10 seconds
+      setTimeout(() => clearInterval(checkInterval), 10000)
+    })
+
+    // Navigate to the app
+    await newPage.goto(comfyPage.url)
+
+    // Wait for both client and server feature flags
+    await newPage.waitForFunction(
+      () =>
+        window.__capturedMessages.clientFeatureFlags !== null &&
+        window.__capturedMessages.serverFeatureFlags !== null,
+      { timeout: 10000 }
+    )
+
+    // Get the captured messages
+    const messages = await newPage.evaluate(() => window.__capturedMessages)
+
+    // Verify client sent feature flags
+    expect(messages.clientFeatureFlags).toBeTruthy()
+    expect(messages.clientFeatureFlags).toHaveProperty('type', 'feature_flags')
+    expect(messages.clientFeatureFlags).toHaveProperty('data')
+    expect(messages.clientFeatureFlags.data).toHaveProperty(
+      'supports_preview_metadata'
+    )
+    expect(
+      typeof messages.clientFeatureFlags.data.supports_preview_metadata
+    ).toBe('boolean')
+
+    // Verify server sent feature flags back
+    expect(messages.serverFeatureFlags).toBeTruthy()
+    expect(messages.serverFeatureFlags).toHaveProperty(
+      'supports_preview_metadata'
+    )
+    expect(typeof messages.serverFeatureFlags.supports_preview_metadata).toBe(
+      'boolean'
+    )
+    expect(messages.serverFeatureFlags).toHaveProperty('max_upload_size')
+    expect(typeof messages.serverFeatureFlags.max_upload_size).toBe('number')
+    expect(Object.keys(messages.serverFeatureFlags).length).toBeGreaterThan(0)
+
+    await newPage.close()
+  })
+
+  test('Server feature flags are received and accessible', async ({
+    comfyPage
+  }) => {
+    // Wait for connection to establish
+    await comfyPage.page.waitForTimeout(1000)
+
+    // Get the actual server feature flags from the backend
+    const serverFlags = await comfyPage.page.evaluate(() => {
+      return window['app'].api.serverFeatureFlags
+    })
+
+    // Verify we received real feature flags from the backend
+    expect(serverFlags).toBeTruthy()
+    expect(Object.keys(serverFlags).length).toBeGreaterThan(0)
+
+    // The backend should send feature flags
+    expect(serverFlags).toHaveProperty('supports_preview_metadata')
+    expect(typeof serverFlags.supports_preview_metadata).toBe('boolean')
+    expect(serverFlags).toHaveProperty('max_upload_size')
+    expect(typeof serverFlags.max_upload_size).toBe('number')
+  })
+
+  test('serverSupportsFeature method works with real backend flags', async ({
+    comfyPage
+  }) => {
+    // Wait for connection
+    await comfyPage.page.waitForTimeout(1000)
+
+    // Test serverSupportsFeature with real backend flags
+    const supportsPreviewMetadata = await comfyPage.page.evaluate(() => {
+      return window['app'].api.serverSupportsFeature(
+        'supports_preview_metadata'
+      )
+    })
+    // The method should return a boolean based on the backend's value
+    expect(typeof supportsPreviewMetadata).toBe('boolean')
+
+    // Test non-existent feature - should always return false
+    const supportsNonExistent = await comfyPage.page.evaluate(() => {
+      return window['app'].api.serverSupportsFeature('non_existent_feature_xyz')
+    })
+    expect(supportsNonExistent).toBe(false)
+
+    // Test that the method only returns true for boolean true values
+    const testResults = await comfyPage.page.evaluate(() => {
+      // Temporarily modify serverFeatureFlags to test behavior
+      const original = window['app'].api.serverFeatureFlags
+      window['app'].api.serverFeatureFlags = {
+        bool_true: true,
+        bool_false: false,
+        string_value: 'yes',
+        number_value: 1,
+        null_value: null
+      }
+
+      const results = {
+        bool_true: window['app'].api.serverSupportsFeature('bool_true'),
+        bool_false: window['app'].api.serverSupportsFeature('bool_false'),
+        string_value: window['app'].api.serverSupportsFeature('string_value'),
+        number_value: window['app'].api.serverSupportsFeature('number_value'),
+        null_value: window['app'].api.serverSupportsFeature('null_value')
+      }
+
+      // Restore original
+      window['app'].api.serverFeatureFlags = original
+      return results
+    })
+
+    // serverSupportsFeature should only return true for boolean true values
+    expect(testResults.bool_true).toBe(true)
+    expect(testResults.bool_false).toBe(false)
+    expect(testResults.string_value).toBe(false)
+    expect(testResults.number_value).toBe(false)
+    expect(testResults.null_value).toBe(false)
+  })
+
+  test('getServerFeature method works with real backend data', async ({
+    comfyPage
+  }) => {
+    // Wait for connection
+    await comfyPage.page.waitForTimeout(1000)
+
+    // Test getServerFeature method
+    const previewMetadataValue = await comfyPage.page.evaluate(() => {
+      return window['app'].api.getServerFeature('supports_preview_metadata')
+    })
+    expect(typeof previewMetadataValue).toBe('boolean')
+
+    // Test getting max_upload_size
+    const maxUploadSize = await comfyPage.page.evaluate(() => {
+      return window['app'].api.getServerFeature('max_upload_size')
+    })
+    expect(typeof maxUploadSize).toBe('number')
+    expect(maxUploadSize).toBeGreaterThan(0)
+
+    // Test getServerFeature with default value for non-existent feature
+    const defaultValue = await comfyPage.page.evaluate(() => {
+      return window['app'].api.getServerFeature(
+        'non_existent_feature_xyz',
+        'default'
+      )
+    })
+    expect(defaultValue).toBe('default')
+  })
+
+  test('getServerFeatures returns all backend feature flags', async ({
+    comfyPage
+  }) => {
+    // Wait for connection
+    await comfyPage.page.waitForTimeout(1000)
+
+    // Test getServerFeatures returns all flags
+    const allFeatures = await comfyPage.page.evaluate(() => {
+      return window['app'].api.getServerFeatures()
+    })
+
+    expect(allFeatures).toBeTruthy()
+    expect(allFeatures).toHaveProperty('supports_preview_metadata')
+    expect(typeof allFeatures.supports_preview_metadata).toBe('boolean')
+    expect(allFeatures).toHaveProperty('max_upload_size')
+    expect(Object.keys(allFeatures).length).toBeGreaterThan(0)
+  })
+
+  test('Client feature flags are immutable', async ({ comfyPage }) => {
+    // Test that getClientFeatureFlags returns a copy
+    const immutabilityTest = await comfyPage.page.evaluate(() => {
+      const flags1 = window['app'].api.getClientFeatureFlags()
+      const flags2 = window['app'].api.getClientFeatureFlags()
+
+      // Modify the first object
+      flags1.test_modification = true
+
+      // Get flags again to check if original was modified
+      const flags3 = window['app'].api.getClientFeatureFlags()
+
+      return {
+        areEqual: flags1 === flags2,
+        hasModification: flags3.test_modification !== undefined,
+        hasSupportsPreview: flags1.supports_preview_metadata !== undefined,
+        supportsPreviewValue: flags1.supports_preview_metadata
+      }
+    })
+
+    // Verify they are different objects (not the same reference)
+    expect(immutabilityTest.areEqual).toBe(false)
+
+    // Verify modification didn't affect the original
+    expect(immutabilityTest.hasModification).toBe(false)
+
+    // Verify the flags contain expected properties
+    expect(immutabilityTest.hasSupportsPreview).toBe(true)
+    expect(typeof immutabilityTest.supportsPreviewValue).toBe('boolean') // From clientFeatureFlags.json
+  })
+
+  test('Server features are immutable when accessed via getServerFeatures', async ({
+    comfyPage
+  }) => {
+    // Wait for connection to establish
+    await comfyPage.page.waitForTimeout(1000)
+
+    const immutabilityTest = await comfyPage.page.evaluate(() => {
+      // Get a copy of server features
+      const features1 = window['app'].api.getServerFeatures()
+
+      // Try to modify it
+      features1.supports_preview_metadata = false
+      features1.new_feature = 'added'
+
+      // Get another copy
+      const features2 = window['app'].api.getServerFeatures()
+
+      return {
+        modifiedValue: features1.supports_preview_metadata,
+        originalValue: features2.supports_preview_metadata,
+        hasNewFeature: features2.new_feature !== undefined,
+        hasSupportsPreview: features2.supports_preview_metadata !== undefined
+      }
+    })
+
+    // The modification should only affect the copy
+    expect(immutabilityTest.modifiedValue).toBe(false)
+    expect(typeof immutabilityTest.originalValue).toBe('boolean') // Backend sends boolean for supports_preview_metadata
+    expect(immutabilityTest.hasNewFeature).toBe(false)
+    expect(immutabilityTest.hasSupportsPreview).toBe(true)
+  })
+
+  test('Feature flags are negotiated early in connection lifecycle', async ({
+    comfyPage
+  }) => {
+    // This test verifies that feature flags are available early in the app lifecycle
+    // which is important for protocol negotiation
+
+    // Create a new page to ensure clean state
+    const newPage = await comfyPage.page.context().newPage()
+
+    // Set up monitoring before navigation
+    await newPage.addInitScript(() => {
+      // Track when various app components are ready
+      ;(window as any).__appReadiness = {
+        featureFlagsReceived: false,
+        apiInitialized: false,
+        appInitialized: false
+      }
+
+      // Monitor when feature flags arrive by checking periodically
+      const checkFeatureFlags = setInterval(() => {
+        if (
+          window['app']?.api?.serverFeatureFlags?.supports_preview_metadata !==
+          undefined
+        ) {
+          ;(window as any).__appReadiness.featureFlagsReceived = true
+          clearInterval(checkFeatureFlags)
+        }
+      }, 10)
+
+      // Monitor API initialization
+      const checkApi = setInterval(() => {
+        if (window['app']?.api) {
+          ;(window as any).__appReadiness.apiInitialized = true
+          clearInterval(checkApi)
+        }
+      }, 10)
+
+      // Monitor app initialization
+      const checkApp = setInterval(() => {
+        if (window['app']?.graph) {
+          ;(window as any).__appReadiness.appInitialized = true
+          clearInterval(checkApp)
+        }
+      }, 10)
+
+      // Clean up after 10 seconds
+      setTimeout(() => {
+        clearInterval(checkFeatureFlags)
+        clearInterval(checkApi)
+        clearInterval(checkApp)
+      }, 10000)
+    })
+
+    // Navigate to the app
+    await newPage.goto(comfyPage.url)
+
+    // Wait for feature flags to be received
+    await newPage.waitForFunction(
+      () =>
+        window['app']?.api?.serverFeatureFlags?.supports_preview_metadata !==
+        undefined,
+      {
+        timeout: 10000
+      }
+    )
+
+    // Get readiness state
+    const readiness = await newPage.evaluate(() => {
+      return {
+        ...(window as any).__appReadiness,
+        currentFlags: window['app'].api.serverFeatureFlags
+      }
+    })
+
+    // Verify feature flags are available
+    expect(readiness.currentFlags).toHaveProperty('supports_preview_metadata')
+    expect(typeof readiness.currentFlags.supports_preview_metadata).toBe(
+      'boolean'
+    )
+    expect(readiness.currentFlags).toHaveProperty('max_upload_size')
+
+    // Verify feature flags were received (we detected them via polling)
+    expect(readiness.featureFlagsReceived).toBe(true)
+
+    // Verify API was initialized (feature flags require API)
+    expect(readiness.apiInitialized).toBe(true)
+
+    await newPage.close()
+  })
+
+  test('Backend /features endpoint returns feature flags', async ({
+    comfyPage
+  }) => {
+    // Test the HTTP endpoint directly
+    const response = await comfyPage.page.request.get(
+      `${comfyPage.url}/api/features`
+    )
+    expect(response.ok()).toBe(true)
+
+    const features = await response.json()
+    expect(features).toBeTruthy()
+    expect(features).toHaveProperty('supports_preview_metadata')
+    expect(typeof features.supports_preview_metadata).toBe('boolean')
+    expect(features).toHaveProperty('max_upload_size')
+    expect(Object.keys(features).length).toBeGreaterThan(0)
+  })
+})

src/config/clientFeatureFlags.json

@@ -0,0 +1,3 @@
+{
+    "supports_preview_metadata": false
+}

src/schemas/apiSchema.ts

@@ -113,6 +113,8 @@ const zLogRawResponse = z.object({
   entries: z.array(zLogEntry)
 })
 
+const zFeatureFlagsWsMessage = z.record(z.string(), z.any())
+
 export type StatusWsMessageStatus = z.infer<typeof zStatusWsMessageStatus>
 export type StatusWsMessage = z.infer<typeof zStatusWsMessage>
 export type ProgressWsMessage = z.infer<typeof zProgressWsMessage>
@@ -132,6 +134,7 @@ export type ProgressTextWsMessage = z.infer<typeof zProgressTextWsMessage>
 export type DisplayComponentWsMessage = z.infer<
   typeof zDisplayComponentWsMessage
 >
+export type FeatureFlagsWsMessage = z.infer<typeof zFeatureFlagsWsMessage>
 // End of ws messages
 
 const zPromptInputItem = z.object({

src/scripts/api.ts

@@ -1,5 +1,6 @@
 import axios from 'axios'
 
+import defaultClientFeatureFlags from '@/config/clientFeatureFlags.json'
 import type {
   DisplayComponentWsMessage,
   EmbeddingsResponse,
@@ -11,6 +12,7 @@ import type {
   ExecutionStartWsMessage,
   ExecutionSuccessWsMessage,
   ExtensionsResponse,
+  FeatureFlagsWsMessage,
   HistoryTaskItem,
   LogsRawResponse,
   LogsWsMessage,
@@ -105,6 +107,7 @@ interface BackendApiCalls {
   b_preview: Blob
   progress_text: ProgressTextWsMessage
   display_component: DisplayComponentWsMessage
+  feature_flags: FeatureFlagsWsMessage
 }
 
 /** Dictionary of all api calls */
@@ -234,6 +237,19 @@ export class ComfyApi extends EventTarget {
 
   reportedUnknownMessageTypes = new Set<string>()
 
+  /**
+   * Get feature flags supported by this frontend client.
+   * Returns a copy to prevent external modification.
+   */
+  getClientFeatureFlags(): Record<string, unknown> {
+    return { ...defaultClientFeatureFlags }
+  }
+
+  /**
+   * Feature flags received from the backend server.
+   */
+  serverFeatureFlags: Record<string, unknown> = {}
+
   /**
    * The auth token for the comfy org account if the user is logged in.
    * This is only used for {@link queuePrompt} now. It is not directly
@@ -375,6 +391,15 @@ export class ComfyApi extends EventTarget {
 
     this.socket.addEventListener('open', () => {
       opened = true
+
+      // Send feature flags as the first message
+      this.socket!.send(
+        JSON.stringify({
+          type: 'feature_flags',
+          data: this.getClientFeatureFlags()
+        })
+      )
+
       if (isReconnect) {
         this.dispatchCustomEvent('reconnected')
       }
@@ -468,6 +493,14 @@ export class ComfyApi extends EventTarget {
             case 'b_preview':
               this.dispatchCustomEvent(msg.type, msg.data)
               break
+            case 'feature_flags':
+              // Store server feature flags
+              this.serverFeatureFlags = msg.data
+              console.log(
+                'Server feature flags received:',
+                this.serverFeatureFlags
+              )
+              break
             default:
               if (this.#registered.has(msg.type)) {
                 // Fallback for custom types - calls super direct.
@@ -962,6 +995,33 @@ export class ComfyApi extends EventTarget {
   async getCustomNodesI18n(): Promise<Record<string, any>> {
     return (await axios.get(this.apiURL('/i18n'))).data
   }
+
+  /**
+   * Checks if the server supports a specific feature.
+   * @param featureName The name of the feature to check
+   * @returns true if the feature is supported, false otherwise
+   */
+  serverSupportsFeature(featureName: string): boolean {
+    return this.serverFeatureFlags[featureName] === true
+  }
+
+  /**
+   * Gets a server feature flag value.
+   * @param featureName The name of the feature to get
+   * @param defaultValue The default value if the feature is not found
+   * @returns The feature value or default
+   */
+  getServerFeature<T = unknown>(featureName: string, defaultValue?: T): T {
+    return (this.serverFeatureFlags[featureName] ?? defaultValue) as T
+  }
+
+  /**
+   * Gets all server feature flags.
+   * @returns Copy of all server feature flags
+   */
+  getServerFeatures(): Record<string, unknown> {
+    return { ...this.serverFeatureFlags }
+  }
 }
 
 export const api = new ComfyApi()

tests-ui/tests/api.featureFlags.test.ts

@@ -0,0 +1,208 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { api } from '@/scripts/api'
+
+describe('API Feature Flags', () => {
+  let mockWebSocket: any
+  const wsEventHandlers: { [key: string]: (event: any) => void } = {}
+
+  beforeEach(() => {
+    // Use fake timers
+    vi.useFakeTimers()
+
+    // Mock WebSocket
+    mockWebSocket = {
+      readyState: 1, // WebSocket.OPEN
+      send: vi.fn(),
+      close: vi.fn(),
+      addEventListener: vi.fn(
+        (event: string, handler: (event: any) => void) => {
+          wsEventHandlers[event] = handler
+        }
+      ),
+      removeEventListener: vi.fn()
+    }
+
+    // Mock WebSocket constructor
+    global.WebSocket = vi.fn().mockImplementation(() => mockWebSocket) as any
+
+    // Reset API state
+    api.serverFeatureFlags = {}
+
+    // Mock getClientFeatureFlags to return test feature flags
+    vi.spyOn(api, 'getClientFeatureFlags').mockReturnValue({
+      supports_preview_metadata: true,
+      api_version: '1.0.0',
+      capabilities: ['bulk_operations', 'async_nodes']
+    })
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+    vi.restoreAllMocks()
+  })
+
+  describe('Feature flags negotiation', () => {
+    it('should send client feature flags as first message on connection', async () => {
+      // Initialize API connection
+      const initPromise = api.init()
+
+      // Simulate connection open
+      wsEventHandlers['open'](new Event('open'))
+
+      // Check that feature flags were sent as first message
+      expect(mockWebSocket.send).toHaveBeenCalledTimes(1)
+      const sentMessage = JSON.parse(mockWebSocket.send.mock.calls[0][0])
+      expect(sentMessage).toEqual({
+        type: 'feature_flags',
+        data: {
+          supports_preview_metadata: true,
+          api_version: '1.0.0',
+          capabilities: ['bulk_operations', 'async_nodes']
+        }
+      })
+
+      // Simulate server response with status message
+      wsEventHandlers['message']({
+        data: JSON.stringify({
+          type: 'status',
+          data: {
+            status: { exec_info: { queue_remaining: 0 } },
+            sid: 'test-sid'
+          }
+        })
+      })
+
+      // Simulate server feature flags response
+      wsEventHandlers['message']({
+        data: JSON.stringify({
+          type: 'feature_flags',
+          data: {
+            supports_preview_metadata: true,
+            async_execution: true,
+            supported_formats: ['webp', 'jpeg', 'png'],
+            api_version: '1.0.0',
+            max_upload_size: 104857600,
+            capabilities: ['isolated_nodes', 'dynamic_models']
+          }
+        })
+      })
+
+      await initPromise
+
+      // Check that server features were stored
+      expect(api.serverFeatureFlags).toEqual({
+        supports_preview_metadata: true,
+        async_execution: true,
+        supported_formats: ['webp', 'jpeg', 'png'],
+        api_version: '1.0.0',
+        max_upload_size: 104857600,
+        capabilities: ['isolated_nodes', 'dynamic_models']
+      })
+    })
+
+    it('should handle server without feature flags support', async () => {
+      // Initialize API connection
+      const initPromise = api.init()
+
+      // Simulate connection open
+      wsEventHandlers['open'](new Event('open'))
+
+      // Clear the send mock to reset
+      mockWebSocket.send.mockClear()
+
+      // Simulate server response with status but no feature flags
+      wsEventHandlers['message']({
+        data: JSON.stringify({
+          type: 'status',
+          data: {
+            status: { exec_info: { queue_remaining: 0 } },
+            sid: 'test-sid'
+          }
+        })
+      })
+
+      // Simulate some other message (not feature flags)
+      wsEventHandlers['message']({
+        data: JSON.stringify({
+          type: 'execution_start',
+          data: {}
+        })
+      })
+
+      await initPromise
+
+      // Server features should remain empty
+      expect(api.serverFeatureFlags).toEqual({})
+    })
+  })
+
+  describe('Feature checking methods', () => {
+    beforeEach(() => {
+      // Set up some test features
+      api.serverFeatureFlags = {
+        supports_preview_metadata: true,
+        async_execution: false,
+        capabilities: ['isolated_nodes', 'dynamic_models']
+      }
+    })
+
+    it('should check if server supports a boolean feature', () => {
+      expect(api.serverSupportsFeature('supports_preview_metadata')).toBe(true)
+      expect(api.serverSupportsFeature('async_execution')).toBe(false)
+      expect(api.serverSupportsFeature('non_existent_feature')).toBe(false)
+    })
+
+    it('should get server feature value', () => {
+      expect(api.getServerFeature('supports_preview_metadata')).toBe(true)
+      expect(api.getServerFeature('capabilities')).toEqual([
+        'isolated_nodes',
+        'dynamic_models'
+      ])
+      expect(api.getServerFeature('non_existent_feature')).toBeUndefined()
+    })
+  })
+
+  describe('Client feature flags configuration', () => {
+    it('should use mocked client feature flags', () => {
+      // Verify mocked flags are returned
+      const clientFlags = api.getClientFeatureFlags()
+      expect(clientFlags).toEqual({
+        supports_preview_metadata: true,
+        api_version: '1.0.0',
+        capabilities: ['bulk_operations', 'async_nodes']
+      })
+    })
+
+    it('should return a copy of client feature flags', () => {
+      // Temporarily restore the real implementation for this test
+      vi.mocked(api.getClientFeatureFlags).mockRestore()
+
+      // Verify that modifications to returned object don't affect original
+      const clientFlags1 = api.getClientFeatureFlags()
+      const clientFlags2 = api.getClientFeatureFlags()
+
+      // Should be different objects
+      expect(clientFlags1).not.toBe(clientFlags2)
+
+      // But with same content
+      expect(clientFlags1).toEqual(clientFlags2)
+
+      // Modifying one should not affect the other
+      clientFlags1.test_flag = true
+      expect(api.getClientFeatureFlags()).not.toHaveProperty('test_flag')
+    })
+  })
+
+  describe('Integration with preview messages', () => {
+    it('should affect preview message handling based on feature support', () => {
+      // Test with metadata support
+      api.serverFeatureFlags = { supports_preview_metadata: true }
+      expect(api.serverSupportsFeature('supports_preview_metadata')).toBe(true)
+
+      // Test without metadata support
+      api.serverFeatureFlags = {}
+      expect(api.serverSupportsFeature('supports_preview_metadata')).toBe(false)
+    })
+  })
+})