Cloud Auth Backport (#6195)

## Summary Backports Firebase authentication with cloud environments. Changes only work when developing for cloud environment locally. ## Changes - Router guards to force unauthenticated users to sign in. - Configure auth headers for REST and Websocket connections. - Code implemented in a way that enables build tree-shaking based on distribution - Updates to build process to build cloud distribution and simplify development workflow ## Review Focus 1. Idomatic Vue/codebase patterns. 2. Build logic (please double check that I integrated correctly with: https://github.com/Comfy-Org/ComfyUI_frontend/blob/rh-test/vite.config.mts) ## Screenshots (if applicable) https://github.com/user-attachments/assets/ee4ea3f7-afa6-4da0-ba43-d62ed8ba4e18 ┆Issue is synchronized with this [Notion page](https://www.notion.so/PR-6195-Feat-cloud-auth-backport-2946d73d365081f395f5f2a89fb7d800) by [Unito](https://www.unito.io) --------- Co-authored-by: Alexander Brown <drjkl@comfy.org> Co-authored-by: GitHub Action <action@github.com>
2026-03-02 11:40:00 +00:00 · 2025-10-22 17:06:37 -07:00
parent bfe53d7721
commit 8b7b580ed4
6 changed files with 421 additions and 41 deletions
--- a/.env_example
+++ b/.env_example
@@ -5,6 +5,10 @@ PLAYWRIGHT_TEST_URL=http://localhost:5173

 # Proxy target of the local development server
 # Note: localhost:8188 does not work.
+# Cloud auto-detection: Setting this to any *.comfy.org URL automatically enables
+# cloud mode (DISTRIBUTION=cloud) without needing to set DISTRIBUTION separately.
+# Examples: https://testcloud.comfy.org/, https://stagingcloud.comfy.org/,
+#           https://pr-123.testenvs.comfy.org/, https://cloud.comfy.org/
 DEV_SERVER_COMFYUI_URL=http://127.0.0.1:8188

 # Allow dev server access from remote IP addresses.
--- a/src/composables/auth/useFirebaseAuthActions.ts
+++ b/src/composables/auth/useFirebaseAuthActions.ts
@@ -1,10 +1,12 @@
 import { FirebaseError } from 'firebase/app'
 import { AuthErrorCodes } from 'firebase/auth'
 import { ref } from 'vue'
+import { useRouter } from 'vue-router'

 import { useErrorHandling } from '@/composables/useErrorHandling'
 import type { ErrorRecoveryStrategy } from '@/composables/useErrorHandling'
 import { t } from '@/i18n'
+import { isCloud } from '@/platform/distribution/types'
 import { useToastStore } from '@/platform/updates/common/toastStore'
 import { useDialogService } from '@/services/dialogService'
 import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
@@ -54,6 +56,16 @@ export const useFirebaseAuthActions = () => {
      detail: t('auth.signOut.successDetail'),
      life: 5000
    })
+
+    if (isCloud) {
+      try {
+        const router = useRouter()
+        await router.push({ name: 'cloud-login' })
+      } catch (error) {
+        // needed for local development until we bring in cloud login pages.
+        window.location.reload()
+      }
+    }
  }, reportError)

  const sendPasswordReset = wrapWithErrorHandlingAsync(
--- a/src/router.ts
+++ b/src/router.ts
@@ -1,14 +1,18 @@
+import { until } from '@vueuse/core'
+import { storeToRefs } from 'pinia'
 import {
  createRouter,
  createWebHashHistory,
  createWebHistory
 } from 'vue-router'

+import { isCloud } from '@/platform/distribution/types'
+import { useDialogService } from '@/services/dialogService'
+import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
+import { useUserStore } from '@/stores/userStore'
+import { isElectron } from '@/utils/envUtil'
 import LayoutDefault from '@/views/layouts/LayoutDefault.vue'

-import { useUserStore } from './stores/userStore'
-import { isElectron } from './utils/envUtil'
-
 const isFileProtocol = window.location.protocol === 'file:'
 const basePath = isElectron() ? '/' : window.location.pathname

@@ -56,4 +60,36 @@ const router = createRouter({
  }
 })

+if (isCloud) {
+  // Global authentication guard
+  router.beforeEach(async (_to, _from, next) => {
+    const authStore = useFirebaseAuthStore()
+
+    // Wait for Firebase auth to initialize
+    // Timeout after 16 seconds
+    if (!authStore.isInitialized) {
+      try {
+        const { isInitialized } = storeToRefs(authStore)
+        await until(isInitialized).toBe(true, { timeout: 16_000 })
+      } catch (error) {
+        console.error('Auth initialization failed:', error)
+        return next({ name: 'cloud-auth-timeout' })
+      }
+    }
+
+    // Pass authenticated users
+    const authHeader = await authStore.getAuthHeader()
+    if (authHeader) {
+      return next()
+    }
+
+    // Show sign-in for unauthenticated users
+    const dialogService = useDialogService()
+    const loginSuccess = await dialogService.showSignInDialog()
+
+    if (loginSuccess) return next()
+    return next(false)
+  })
+}
+
 export default router
--- a/src/scripts/api.ts
+++ b/src/scripts/api.ts
@@ -1,3 +1,4 @@
+import { promiseTimeout, until } from '@vueuse/core'
 import axios from 'axios'
 import { get } from 'es-toolkit/compat'

@@ -6,6 +7,7 @@ import type {
  ModelFile,
  ModelFolderInfo
 } from '@/platform/assets/schemas/assetSchema'
+import { isCloud } from '@/platform/distribution/types'
 import { useToastStore } from '@/platform/updates/common/toastStore'
 import { type WorkflowTemplates } from '@/platform/workflow/templates/types/template'
 import type {
@@ -42,6 +44,8 @@ import type {
  UserDataFullInfo
 } from '@/schemas/apiSchema'
 import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
+import type { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
+import type { AuthHeader } from '@/types/authTypes'
 import type { NodeExecutionId } from '@/types/nodeIdentification'

 interface QueuePromptRequestBody {
@@ -200,6 +204,16 @@ type SimpleApiEvents = keyof PickNevers<ApiEventTypes>
 /** Keys (names) of API events that pass a {@link CustomEvent} `detail` object. */
 type ComplexApiEvents = keyof NeverNever<ApiEventTypes>

+function addHeaderEntry(headers: HeadersInit, key: string, value: string) {
+  if (Array.isArray(headers)) {
+    headers.push([key, value])
+  } else if (headers instanceof Headers) {
+    headers.set(key, value)
+  } else {
+    headers[key] = value
+  }
+}
+
 /** EventTarget typing has no generic capability. */
 export interface ComfyApi extends EventTarget {
  addEventListener<TEvent extends keyof ApiEvents>(
@@ -263,6 +277,11 @@ export class ComfyApi extends EventTarget {
  user: string
  socket: WebSocket | null = null

+  /**
+   * Cache Firebase auth store composable function.
+   */
+  private authStoreComposable?: typeof useFirebaseAuthStore
+
  reportedUnknownMessageTypes = new Set<string>()

  /**
@@ -317,25 +336,77 @@ export class ComfyApi extends EventTarget {
    return this.api_base + route
  }

-  fetchApi(route: string, options?: RequestInit) {
-    if (!options) {
-      options = {}
+  /**
+   * Gets the Firebase auth store instance using cached composable function.
+   * Caches the composable function on first call, then reuses it.
+   * Returns null for non-cloud distributions.
+   * @returns The Firebase auth store instance, or null if not in cloud
+   */
+  private async getAuthStore() {
+    if (isCloud) {
+      if (!this.authStoreComposable) {
+        const module = await import('@/stores/firebaseAuthStore')
+        this.authStoreComposable = module.useFirebaseAuthStore
+      }
+
+      return this.authStoreComposable()
    }
-    if (!options.headers) {
-      options.headers = {}
+  }
+
+  /**
+   * Waits for Firebase auth to be initialized before proceeding.
+   * Includes 10-second timeout to prevent infinite hanging.
+   */
+  private async waitForAuthInitialization(): Promise<void> {
+    if (isCloud) {
+      const authStore = await this.getAuthStore()
+      if (!authStore) return
+
+      if (authStore.isInitialized) return
+
+      try {
+        await Promise.race([
+          until(authStore.isInitialized),
+          promiseTimeout(10000)
+        ])
+      } catch {
+        console.warn('Firebase auth initialization timeout after 10 seconds')
+      }
    }
-    if (!options.cache) {
-      options.cache = 'no-cache'
+  }
+
+  async fetchApi(route: string, options?: RequestInit) {
+    const headers: HeadersInit = options?.headers ?? {}
+
+    if (isCloud) {
+      await this.waitForAuthInitialization()
+
+      // Get Firebase JWT token if user is logged in
+      const getAuthHeaderIfAvailable = async (): Promise<AuthHeader | null> => {
+        try {
+          const authStore = await this.getAuthStore()
+          return authStore ? await authStore.getAuthHeader() : null
+        } catch (error) {
+          console.warn('Failed to get auth header:', error)
+          return null
+        }
+      }
+
+      const authHeader = await getAuthHeaderIfAvailable()
+
+      if (authHeader) {
+        for (const [key, value] of Object.entries(authHeader)) {
+          addHeaderEntry(headers, key, value)
+        }
+      }
    }

-    if (Array.isArray(options.headers)) {
-      options.headers.push(['Comfy-User', this.user])
-    } else if (options.headers instanceof Headers) {
-      options.headers.set('Comfy-User', this.user)
-    } else {
-      options.headers['Comfy-User'] = this.user
-    }
-    return fetch(this.apiURL(route), options)
+    addHeaderEntry(headers, 'Comfy-User', this.user)
+    return fetch(this.apiURL(route), {
+      cache: 'no-cache',
+      ...options,
+      headers
+    })
  }

  override addEventListener<TEvent extends keyof ApiEvents>(
@@ -402,19 +473,44 @@ export class ComfyApi extends EventTarget {
   * Creates and connects a WebSocket for realtime updates
   * @param {boolean} isReconnect If the socket is connection is a reconnect attempt
   */
-  #createSocket(isReconnect?: boolean) {
+  private async createSocket(isReconnect?: boolean) {
    if (this.socket) {
      return
    }

    let opened = false
    let existingSession = window.name
+
+    // Build WebSocket URL with query parameters
+    const params = new URLSearchParams()
+
    if (existingSession) {
-      existingSession = '?clientId=' + existingSession
+      params.set('clientId', existingSession)
    }
-    this.socket = new WebSocket(
-      `ws${window.location.protocol === 'https:' ? 's' : ''}://${this.api_host}${this.api_base}/ws${existingSession}`
-    )
+
+    // Get auth token and set cloud params if available
+    if (isCloud) {
+      try {
+        const authStore = await this.getAuthStore()
+        const authToken = await authStore?.getIdToken()
+        if (authToken) {
+          params.set('token', authToken)
+        }
+      } catch (error) {
+        // Continue without auth token if there's an error
+        console.warn(
+          'Could not get auth token for WebSocket connection:',
+          error
+        )
+      }
+    }
+
+    const protocol = window.location.protocol === 'https:' ? 'wss' : 'ws'
+    const baseUrl = `${protocol}://${this.api_host}${this.api_base}/ws`
+    const query = params.toString()
+    const wsUrl = query ? `${baseUrl}?${query}` : baseUrl
+
+    this.socket = new WebSocket(wsUrl)
    this.socket.binaryType = 'arraybuffer'

    this.socket.addEventListener('open', () => {
@@ -441,9 +537,9 @@ export class ComfyApi extends EventTarget {
    })

    this.socket.addEventListener('close', () => {
-      setTimeout(() => {
+      setTimeout(async () => {
        this.socket = null
-        this.#createSocket(true)
+        await this.createSocket(true)
      }, 300)
      if (opened) {
        this.dispatchCustomEvent('status', null)
@@ -579,7 +675,7 @@ export class ComfyApi extends EventTarget {
   * Initialises sockets and realtime updates
   */
  init() {
-    this.#createSocket()
+    this.createSocket()
  }

  /**
--- a/tests-ui/tests/api.fetchApi.test.ts
+++ b/tests-ui/tests/api.fetchApi.test.ts
@@ -0,0 +1,171 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { api } from '@/scripts/api'
+
+// Mock global fetch
+vi.stubGlobal('fetch', vi.fn())
+
+describe('api.fetchApi', () => {
+  beforeEach(() => {
+    vi.resetAllMocks()
+
+    // Reset api state
+    api.user = 'test-user'
+  })
+
+  describe('header handling', () => {
+    it('should add Comfy-User header with plain object headers', async () => {
+      const mockFetch = vi
+        .mocked(global.fetch)
+        .mockResolvedValue(new Response())
+
+      await api.fetchApi('/test', {
+        headers: {}
+      })
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        expect.stringContaining('/test'),
+        expect.objectContaining({
+          headers: {
+            'Comfy-User': 'test-user'
+          }
+        })
+      )
+    })
+
+    it('should add Comfy-User header with Headers instance', async () => {
+      const mockFetch = vi
+        .mocked(global.fetch)
+        .mockResolvedValue(new Response())
+      const headers = new Headers()
+
+      await api.fetchApi('/test', { headers })
+
+      expect(mockFetch).toHaveBeenCalled()
+      const callHeaders = mockFetch.mock.calls[0][1]?.headers
+      expect(callHeaders).toEqual(headers)
+    })
+
+    it('should add Comfy-User header with array headers', async () => {
+      const mockFetch = vi
+        .mocked(global.fetch)
+        .mockResolvedValue(new Response())
+      const headers: [string, string][] = []
+
+      await api.fetchApi('/test', { headers })
+
+      expect(mockFetch).toHaveBeenCalled()
+      const callHeaders = mockFetch.mock.calls[0][1]?.headers
+      expect(callHeaders).toContainEqual(['Comfy-User', 'test-user'])
+    })
+
+    it('should preserve existing headers when adding Comfy-User', async () => {
+      const mockFetch = vi
+        .mocked(global.fetch)
+        .mockResolvedValue(new Response())
+
+      await api.fetchApi('/test', {
+        headers: {
+          'Content-Type': 'application/json',
+          'X-Custom': 'value'
+        }
+      })
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        expect.stringContaining('/test'),
+        expect.objectContaining({
+          headers: {
+            'Content-Type': 'application/json',
+            'X-Custom': 'value',
+            'Comfy-User': 'test-user'
+          }
+        })
+      )
+    })
+
+    it('should not allow developer-specified headers to be overridden by options', async () => {
+      const mockFetch = vi
+        .mocked(global.fetch)
+        .mockResolvedValue(new Response())
+
+      await api.fetchApi('/test', {
+        headers: {
+          'Comfy-User': 'fennec-girl'
+        }
+      })
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        expect.stringContaining('/test'),
+        expect.objectContaining({
+          headers: {
+            'Comfy-User': 'test-user'
+          }
+        })
+      )
+    })
+  })
+
+  describe('default options', () => {
+    it('should set cache to no-cache by default', async () => {
+      const mockFetch = vi
+        .mocked(global.fetch)
+        .mockResolvedValue(new Response())
+
+      await api.fetchApi('/test')
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        expect.any(String),
+        expect.objectContaining({
+          cache: 'no-cache'
+        })
+      )
+    })
+
+    it('should include required headers even when no headers option is provided', async () => {
+      const mockFetch = vi
+        .mocked(global.fetch)
+        .mockResolvedValue(new Response())
+
+      await api.fetchApi('/test')
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        expect.any(String),
+        expect.objectContaining({
+          headers: expect.objectContaining({
+            'Comfy-User': 'test-user'
+          })
+        })
+      )
+    })
+
+    it('should not override existing cache option', async () => {
+      const mockFetch = vi
+        .mocked(global.fetch)
+        .mockResolvedValue(new Response())
+
+      await api.fetchApi('/test', { cache: 'force-cache' })
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        expect.any(String),
+        expect.objectContaining({
+          cache: 'force-cache'
+        })
+      )
+    })
+  })
+
+  describe('URL construction', () => {
+    it('should use apiURL for route construction', async () => {
+      const mockFetch = vi
+        .mocked(global.fetch)
+        .mockResolvedValue(new Response())
+
+      await api.fetchApi('/test/route')
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        expect.stringContaining('/api/test/route'),
+        expect.any(Object)
+      )
+    })
+  })
+})
--- a/vite.config.mts
+++ b/vite.config.mts
@@ -21,16 +21,61 @@ const ANALYZE_BUNDLE = process.env.ANALYZE_BUNDLE === 'true'
 // vite dev server will listen on all addresses, including LAN and public addresses
 const VITE_REMOTE_DEV = process.env.VITE_REMOTE_DEV === 'true'
 const DISABLE_TEMPLATES_PROXY = process.env.DISABLE_TEMPLATES_PROXY === 'true'
-const DISABLE_VUE_PLUGINS = process.env.DISABLE_VUE_PLUGINS === 'true'
 const GENERATE_SOURCEMAP = process.env.GENERATE_SOURCEMAP !== 'false'

-const DEV_SERVER_COMFYUI_URL =
-  process.env.DEV_SERVER_COMFYUI_URL || 'http://127.0.0.1:8188'
+// Auto-detect cloud mode from DEV_SERVER_COMFYUI_URL
+const DEV_SERVER_COMFYUI_ENV_URL = process.env.DEV_SERVER_COMFYUI_URL
+const IS_CLOUD_URL = DEV_SERVER_COMFYUI_ENV_URL?.includes('.comfy.org')

-const DISTRIBUTION = (process.env.DISTRIBUTION || 'localhost') as
-  | 'desktop'
-  | 'localhost'
-  | 'cloud'
+const DISTRIBUTION = (process.env.DISTRIBUTION ||
+  (IS_CLOUD_URL ? 'cloud' : 'localhost')) as 'desktop' | 'localhost' | 'cloud'
+
+// Disable Vue DevTools for production cloud distribution
+const DISABLE_VUE_PLUGINS =
+  process.env.DISABLE_VUE_PLUGINS === 'true' ||
+  (DISTRIBUTION === 'cloud' && !IS_DEV)
+
+const DEV_SEVER_FALLBACK_URL =
+  DISTRIBUTION === 'cloud'
+    ? 'https://stagingcloud.comfy.org'
+    : 'http://127.0.0.1:8188'
+
+const DEV_SERVER_COMFYUI_URL =
+  DEV_SERVER_COMFYUI_ENV_URL || DEV_SEVER_FALLBACK_URL
+
+// Optional: Add API key to .env as STAGING_API_KEY if needed for cloud authentication
+const addAuthHeaders = (proxy: any) => {
+  proxy.on('proxyReq', (proxyReq: any, _req: any, _res: any) => {
+    const apiKey = process.env.STAGING_API_KEY
+    if (apiKey) {
+      proxyReq.setHeader('X-API-KEY', apiKey)
+    }
+  })
+}
+
+// Cloud endpoint configuration overrides
+const cloudSecureConfig =
+  DISTRIBUTION === 'cloud'
+    ? {
+        secure: false
+      }
+    : {}
+
+const cloudProxyConfig =
+  DISTRIBUTION === 'cloud'
+    ? {
+        ...cloudSecureConfig,
+        changeOrigin: true
+      }
+    : {}
+
+const cloudProxyConfigWithAuth =
+  DISTRIBUTION === 'cloud'
+    ? {
+        ...cloudProxyConfig,
+        configure: addAuthHeaders
+      }
+    : {}

 const BUILD_FLAGS = {
  REQUIRE_SUBSCRIPTION: process.env.REQUIRE_SUBSCRIPTION === 'true'
@@ -59,46 +104,62 @@ export default defineConfig({
    },
    proxy: {
      '/internal': {
-        target: DEV_SERVER_COMFYUI_URL
+        target: DEV_SERVER_COMFYUI_URL,
+        ...cloudProxyConfigWithAuth
      },

      '/api': {
        target: DEV_SERVER_COMFYUI_URL,
-        // Return empty array for extensions API as these modules
-        // are not on vite's dev server.
+        ...cloudProxyConfigWithAuth,
        bypass: (req, res, _options) => {
+          // Return empty array for extensions API as these modules
+          // are not on vite's dev server.
          if (req.url === '/api/extensions') {
            res.end(JSON.stringify([]))
+            return false
          }
+
+          // Bypass multi-user auth check from staging (cloud only)
+          if (DISTRIBUTION === 'cloud' && req.url === '/api/users') {
+            res.setHeader('Content-Type', 'application/json')
+            res.end(JSON.stringify({})) // Return empty object to simulate single-user mode
+            return false
+          }
+
          return null
        }
      },

      '/ws': {
        target: DEV_SERVER_COMFYUI_URL,
-        ws: true
+        ws: true,
+        ...cloudProxyConfigWithAuth
      },

      '/workflow_templates': {
-        target: DEV_SERVER_COMFYUI_URL
+        target: DEV_SERVER_COMFYUI_URL,
+        ...cloudProxyConfigWithAuth
      },

      // Proxy extension assets (images/videos) under /extensions to the ComfyUI backend
      '/extensions': {
        target: DEV_SERVER_COMFYUI_URL,
-        changeOrigin: true
+        changeOrigin: true,
+        ...cloudSecureConfig
      },

      // Proxy docs markdown from backend
      '/docs': {
        target: DEV_SERVER_COMFYUI_URL,
-        changeOrigin: true
+        changeOrigin: true,
+        ...cloudSecureConfig
      },

      ...(!DISABLE_TEMPLATES_PROXY
        ? {
            '/templates': {
-              target: DEV_SERVER_COMFYUI_URL
+              target: DEV_SERVER_COMFYUI_URL,
+              ...cloudProxyConfig
            }
          }
        : {}),