Cloud Auth Backport

.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.

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(

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

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 = {}
-    }
-    if (!options.headers) {
-      options.headers = {}
+  /**
+   * 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.cache) {
-      options.cache = 'no-cache'
+  }
+
+  /**
+   * 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')
+      }
     }
+  }
+
+  async fetchApi(route: string, options?: RequestInit) {
+    const headers: HeadersInit = options?.headers ?? {}
+
+    if (isCloud) {
+      await this.waitForAuthInitialization()
 
-    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
+      // 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)
+        }
+      }
     }
-    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()
   }
 
   /**