From 8b7b580ed4efae0318d6a05e6b49b08285639342 Mon Sep 17 00:00:00 2001 From: Arjan Singh <1598641+arjansingh@users.noreply.github.com> Date: Wed, 22 Oct 2025 17:06:37 -0700 Subject: [PATCH] Cloud Auth Backport (#6195) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 Co-authored-by: GitHub Action --- .env_example | 4 + .../auth/useFirebaseAuthActions.ts | 12 ++ src/router.ts | 42 ++++- src/scripts/api.ts | 142 ++++++++++++--- tests-ui/tests/api.fetchApi.test.ts | 171 ++++++++++++++++++ vite.config.mts | 91 ++++++++-- 6 files changed, 421 insertions(+), 41 deletions(-) create mode 100644 tests-ui/tests/api.fetchApi.test.ts diff --git a/.env_example b/.env_example index 520521fe03..81c31f2258 100644 --- 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. diff --git a/src/composables/auth/useFirebaseAuthActions.ts b/src/composables/auth/useFirebaseAuthActions.ts index 19db8b3a23..f03c20da12 100644 --- 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( diff --git a/src/router.ts b/src/router.ts index cc51f7715e..e468b02f37 100644 --- 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 diff --git a/src/scripts/api.ts b/src/scripts/api.ts index e372b8fa92..271ce17eee 100644 --- 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 /** Keys (names) of API events that pass a {@link CustomEvent} `detail` object. */ type ComplexApiEvents = keyof NeverNever +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( @@ -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() /** @@ -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 { + 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 => { + 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( @@ -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() } /** diff --git a/tests-ui/tests/api.fetchApi.test.ts b/tests-ui/tests/api.fetchApi.test.ts new file mode 100644 index 0000000000..bc5b3a0cdc --- /dev/null +++ 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) + ) + }) + }) +}) diff --git a/vite.config.mts b/vite.config.mts index dbd4f20c50..57571dd755 100644 --- 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 } } : {}),