diff --git a/src/App.vue b/src/App.vue index 7c11c4c7beb..b5e17500fd1 100644 --- a/src/App.vue +++ b/src/App.vue @@ -1,11 +1,13 @@ diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue index 1122fe85f08..08e7573a9b1 100644 --- a/src/components/graph/GraphCanvas.vue +++ b/src/components/graph/GraphCanvas.vue @@ -502,19 +502,9 @@ onMounted(async () => { await workflowPersistence.loadTemplateFromUrlIfPresent() // Accept workspace invite from URL if present (e.g., ?invite=TOKEN) - // Uses watch because feature flags load asynchronously - flag may be false initially - // then become true once remoteConfig or websocket features are loaded - if (inviteUrlLoader) { - const stopWatching = watch( - () => flags.teamWorkspacesEnabled, - async (enabled) => { - if (enabled) { - stopWatching() - await inviteUrlLoader.loadInviteFromUrl() - } - }, - { immediate: true } - ) + // WorkspaceAuthGate ensures flag state is resolved before GraphCanvas mounts + if (inviteUrlLoader && flags.teamWorkspacesEnabled) { + await inviteUrlLoader.loadInviteFromUrl() } // Initialize release store to fetch releases from comfy-api (fire-and-forget) diff --git a/src/composables/useFeatureFlags.ts b/src/composables/useFeatureFlags.ts index ca54bb9c6f6..e0e69b4c9cc 100644 --- a/src/composables/useFeatureFlags.ts +++ b/src/composables/useFeatureFlags.ts @@ -1,7 +1,10 @@ import { computed, reactive, readonly } from 'vue' import { isCloud } from '@/platform/distribution/types' -import { remoteConfig } from '@/platform/remoteConfig/remoteConfig' +import { + isAuthenticatedConfigLoaded, + remoteConfig +} from '@/platform/remoteConfig/remoteConfig' import { api } from '@/scripts/api' /** @@ -95,9 +98,20 @@ export function useFeatureFlags() { ) ) }, + /** + * Whether team workspaces feature is enabled. + * IMPORTANT: Returns false until authenticated remote config is loaded. + * This ensures we never use workspace tokens when the feature is disabled, + * and prevents race conditions during initialization. + */ get teamWorkspacesEnabled() { if (!isCloud) return false + // Only return true if authenticated config has been loaded. + // This prevents race conditions where code checks this flag before + // WorkspaceAuthGate has refreshed the config with auth. + if (!isAuthenticatedConfigLoaded.value) return false + return ( remoteConfig.value.team_workspaces_enabled ?? api.getServerFeature(ServerFeatureFlag.TEAM_WORKSPACES_ENABLED, false) diff --git a/src/extensions/core/cloudRemoteConfig.ts b/src/extensions/core/cloudRemoteConfig.ts index 6d775389e9a..d3c1dc7ac58 100644 --- a/src/extensions/core/cloudRemoteConfig.ts +++ b/src/extensions/core/cloudRemoteConfig.ts @@ -16,13 +16,15 @@ useExtensionService().registerExtension({ const { isLoggedIn } = useCurrentUser() const { isActiveSubscription } = useSubscription() + // Refresh config when subscription status changes + // Initial auth-aware refresh happens in WorkspaceAuthGate before app renders watchDebounced( [isLoggedIn, isActiveSubscription], () => { if (!isLoggedIn.value) return void refreshRemoteConfig() }, - { debounce: 256, immediate: true } + { debounce: 256 } ) // Poll for config updates every 10 minutes (with auth) diff --git a/src/platform/remoteConfig/refreshRemoteConfig.ts b/src/platform/remoteConfig/refreshRemoteConfig.ts index 8c87bfcdc72..86b49909025 100644 --- a/src/platform/remoteConfig/refreshRemoteConfig.ts +++ b/src/platform/remoteConfig/refreshRemoteConfig.ts @@ -1,6 +1,6 @@ import { api } from '@/scripts/api' -import { remoteConfig } from './remoteConfig' +import { remoteConfig, remoteConfigState } from './remoteConfig' interface RefreshRemoteConfigOptions { /** @@ -12,7 +12,12 @@ interface RefreshRemoteConfigOptions { /** * Loads remote configuration from the backend /features endpoint - * and updates the reactive remoteConfig ref + * and updates the reactive remoteConfig ref. + * + * Sets remoteConfigState to: + * - 'anonymous' when loaded without auth + * - 'authenticated' when loaded with auth + * - 'error' when load fails */ export async function refreshRemoteConfig( options: RefreshRemoteConfigOptions = {} @@ -28,6 +33,7 @@ export async function refreshRemoteConfig( const config = await response.json() window.__CONFIG__ = config remoteConfig.value = config + remoteConfigState.value = useAuth ? 'authenticated' : 'anonymous' return } @@ -35,10 +41,12 @@ export async function refreshRemoteConfig( if (response.status === 401 || response.status === 403) { window.__CONFIG__ = {} remoteConfig.value = {} + remoteConfigState.value = 'error' } } catch (error) { console.error('Failed to fetch remote config:', error) window.__CONFIG__ = {} remoteConfig.value = {} + remoteConfigState.value = 'error' } } diff --git a/src/platform/remoteConfig/remoteConfig.ts b/src/platform/remoteConfig/remoteConfig.ts index 40f36522fb8..3b6ecb96c02 100644 --- a/src/platform/remoteConfig/remoteConfig.ts +++ b/src/platform/remoteConfig/remoteConfig.ts @@ -10,10 +10,32 @@ * This module is tree-shaken in OSS builds. */ -import { ref } from 'vue' +import { computed, ref } from 'vue' import type { RemoteConfig } from './types' +/** + * Load state for remote configuration. + * - 'unloaded': No config loaded yet + * - 'anonymous': Config loaded without auth (bootstrap) + * - 'authenticated': Config loaded with auth (user-specific flags available) + * - 'error': Failed to load config + */ +type RemoteConfigState = 'unloaded' | 'anonymous' | 'authenticated' | 'error' + +/** + * Current load state of remote configuration + */ +export const remoteConfigState = ref('unloaded') + +/** + * Whether the authenticated config has been loaded. + * Use this to gate access to user-specific feature flags like teamWorkspacesEnabled. + */ +export const isAuthenticatedConfigLoaded = computed( + () => remoteConfigState.value === 'authenticated' +) + /** * Reactive remote configuration * Updated whenever config is loaded from the server diff --git a/src/scripts/api.ts b/src/scripts/api.ts index 41310b0bdec..9a8c13f7dac 100644 --- a/src/scripts/api.ts +++ b/src/scripts/api.ts @@ -521,10 +521,11 @@ export class ComfyApi extends EventTarget { } // Get auth token and set cloud params if available + // Uses workspace token (if enabled) or Firebase token if (isCloud) { try { const authStore = await this.getAuthStore() - const authToken = await authStore?.getIdToken() + const authToken = await authStore?.getAuthToken() if (authToken) { params.set('token', authToken) } diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 106603a93e4..796965b27d5 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -1349,8 +1349,9 @@ export class ComfyApp { const executionStore = useExecutionStore() executionStore.lastNodeErrors = null - let comfyOrgAuthToken = await useFirebaseAuthStore().getIdToken() - let comfyOrgApiKey = useApiKeyAuthStore().getApiKey() + // Get auth token for backend nodes - uses workspace token if enabled, otherwise Firebase token + const comfyOrgAuthToken = await useFirebaseAuthStore().getAuthToken() + const comfyOrgApiKey = useApiKeyAuthStore().getApiKey() try { while (this.queueItems.length) { diff --git a/src/stores/firebaseAuthStore.ts b/src/stores/firebaseAuthStore.ts index 6073a43040f..41bae52a285 100644 --- a/src/stores/firebaseAuthStore.ts +++ b/src/stores/firebaseAuthStore.ts @@ -212,6 +212,31 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => { return token ? { Authorization: `Bearer ${token}` } : null } + /** + * Returns the raw auth token (not wrapped in a header object). + * Priority: workspace token > Firebase token. + * Use this for WebSocket connections and backend node auth. + */ + const getAuthToken = async (): Promise => { + if (flags.teamWorkspacesEnabled) { + const workspaceToken = sessionStorage.getItem( + WORKSPACE_STORAGE_KEYS.TOKEN + ) + const expiresAt = sessionStorage.getItem( + WORKSPACE_STORAGE_KEYS.EXPIRES_AT + ) + + if (workspaceToken && expiresAt) { + const expiryTime = parseInt(expiresAt, 10) + if (Date.now() < expiryTime) { + return workspaceToken + } + } + } + + return await getIdToken() + } + const fetchBalance = async (): Promise => { isFetchingBalance.value = true try { @@ -513,6 +538,7 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => { updatePassword: _updatePassword, deleteAccount: _deleteAccount, getAuthHeader, - getFirebaseAuthHeader + getFirebaseAuthHeader, + getAuthToken } }) diff --git a/src/views/GraphView.vue b/src/views/GraphView.vue index e543a13af70..ca03f22a60b 100644 --- a/src/views/GraphView.vue +++ b/src/views/GraphView.vue @@ -52,7 +52,6 @@ import { useProgressFavicon } from '@/composables/useProgressFavicon' import { SERVER_CONFIG_ITEMS } from '@/constants/serverConfig' import { i18n, loadLocale } from '@/i18n' import ModelImportProgressDialog from '@/platform/assets/components/ModelImportProgressDialog.vue' -import { useFeatureFlags } from '@/composables/useFeatureFlags' import { isCloud } from '@/platform/distribution/types' import { useSettingStore } from '@/platform/settings/settingStore' import { useTelemetry } from '@/platform/telemetry' @@ -246,27 +245,6 @@ const onReconnected = () => { } } -// Initialize workspace store when feature flag and auth become available -// Uses watch because remoteConfig loads asynchronously after component mount -if (isCloud) { - const { flags } = useFeatureFlags() - - watch( - () => [flags.teamWorkspacesEnabled, firebaseAuthStore.isAuthenticated], - async ([enabled, isAuthenticated]) => { - if (!enabled || !isAuthenticated) return - - const { useTeamWorkspaceStore } = - await import('@/platform/workspace/stores/teamWorkspaceStore') - const workspaceStore = useTeamWorkspaceStore() - if (workspaceStore.initState === 'uninitialized') { - await workspaceStore.initialize() - } - }, - { immediate: true } - ) -} - useEventListener(api, 'status', onStatus) useEventListener(api, 'execution_success', onExecutionSuccess) useEventListener(api, 'reconnecting', onReconnecting)