diff --git a/src/App.vue b/src/App.vue index 6b7c56be079..f36f7472067 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 b2b526eeaee..9d483503365 100644 --- a/src/components/graph/GraphCanvas.vue +++ b/src/components/graph/GraphCanvas.vue @@ -466,19 +466,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 822773c0d66..453f87e0e2f 100644 --- a/src/scripts/api.ts +++ b/src/scripts/api.ts @@ -515,10 +515,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 258214dd6e1..ebf5e2412b1 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 85735e05c36..3dd6110901f 100644 --- a/src/views/GraphView.vue +++ b/src/views/GraphView.vue @@ -24,7 +24,7 @@