Feat/workspaces 5 auth gate check (#8350)

## Summary

- Fix auth related race conditions with a new WorkspaceAuthGate in
App.vue
- De dup initialization calls 
- Add state machine to track state of refreshRemoteConfig
- Fix websocket not using new workspace jwt
- Misc improvments

## Changes

- **What**: Mainly WorkspaceAuthGate.vue
- **Breaking**: <!-- Any breaking changes (if none, remove this line)
-->
- **Dependencies**: <!-- New dependencies (if none, remove this line)
-->

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-8350-Feat-workspaces-5-auth-gate-check-2f66d73d365081b1a49afcd418fab3e7)
by [Unito](https://www.unito.io)

src/App.vue

@@ -1,11 +1,13 @@
 <template>
-  <router-view />
-  <ProgressSpinner
-    v-if="isLoading"
-    class="absolute inset-0 flex h-[unset] items-center justify-center"
-  />
-  <GlobalDialog />
-  <BlockUI full-screen :blocked="isLoading" />
+  <WorkspaceAuthGate>
+    <router-view />
+    <ProgressSpinner
+      v-if="isLoading"
+      class="absolute inset-0 flex h-[unset] items-center justify-center"
+    />
+    <GlobalDialog />
+    <BlockUI full-screen :blocked="isLoading" />
+  </WorkspaceAuthGate>
 </template>
 
 <script setup lang="ts">
@@ -14,6 +16,7 @@ import BlockUI from 'primevue/blockui'
 import ProgressSpinner from 'primevue/progressspinner'
 import { computed, onMounted } from 'vue'
 
+import WorkspaceAuthGate from '@/components/auth/WorkspaceAuthGate.vue'
 import GlobalDialog from '@/components/dialog/GlobalDialog.vue'
 import config from '@/config'
 import { useWorkspaceStore } from '@/stores/workspaceStore'

src/components/auth/WorkspaceAuthGate.test.ts

@@ -0,0 +1,206 @@
+import { flushPromises, mount } from '@vue/test-utils'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { ref } from 'vue'
+
+import WorkspaceAuthGate from './WorkspaceAuthGate.vue'
+
+const mockIsInitialized = ref(false)
+const mockCurrentUser = ref<object | null>(null)
+
+vi.mock('@/stores/firebaseAuthStore', () => ({
+  useFirebaseAuthStore: () => ({
+    isInitialized: mockIsInitialized,
+    currentUser: mockCurrentUser
+  })
+}))
+
+const mockRefreshRemoteConfig = vi.fn()
+vi.mock('@/platform/remoteConfig/refreshRemoteConfig', () => ({
+  refreshRemoteConfig: (options: unknown) => mockRefreshRemoteConfig(options)
+}))
+
+const mockTeamWorkspacesEnabled = vi.hoisted(() => ({ value: false }))
+vi.mock('@/composables/useFeatureFlags', () => ({
+  useFeatureFlags: () => ({
+    flags: {
+      get teamWorkspacesEnabled() {
+        return mockTeamWorkspacesEnabled.value
+      }
+    }
+  })
+}))
+
+const mockWorkspaceStoreInitialize = vi.fn()
+const mockWorkspaceStoreInitState = vi.hoisted(() => ({
+  value: 'uninitialized' as string
+}))
+vi.mock('@/platform/workspace/stores/teamWorkspaceStore', () => ({
+  useTeamWorkspaceStore: () => ({
+    get initState() {
+      return mockWorkspaceStoreInitState.value
+    },
+    initialize: mockWorkspaceStoreInitialize
+  })
+}))
+
+const mockIsCloud = vi.hoisted(() => ({ value: true }))
+vi.mock('@/platform/distribution/types', () => ({
+  get isCloud() {
+    return mockIsCloud.value
+  }
+}))
+
+vi.mock('primevue/progressspinner', () => ({
+  default: { template: '<div class="progress-spinner" />' }
+}))
+
+describe('WorkspaceAuthGate', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockIsCloud.value = true
+    mockIsInitialized.value = false
+    mockCurrentUser.value = null
+    mockTeamWorkspacesEnabled.value = false
+    mockWorkspaceStoreInitState.value = 'uninitialized'
+    mockRefreshRemoteConfig.mockResolvedValue(undefined)
+    mockWorkspaceStoreInitialize.mockResolvedValue(undefined)
+  })
+
+  const mountComponent = () =>
+    mount(WorkspaceAuthGate, {
+      slots: {
+        default: '<div data-testid="slot-content">App Content</div>'
+      }
+    })
+
+  describe('non-cloud builds', () => {
+    it('renders slot immediately when isCloud is false', async () => {
+      mockIsCloud.value = false
+
+      const wrapper = mountComponent()
+      await flushPromises()
+
+      expect(wrapper.find('[data-testid="slot-content"]').exists()).toBe(true)
+      expect(wrapper.find('.progress-spinner').exists()).toBe(false)
+      expect(mockRefreshRemoteConfig).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('cloud builds - unauthenticated user', () => {
+    it('shows spinner while waiting for Firebase auth', () => {
+      mockIsInitialized.value = false
+
+      const wrapper = mountComponent()
+
+      expect(wrapper.find('.progress-spinner').exists()).toBe(true)
+      expect(wrapper.find('[data-testid="slot-content"]').exists()).toBe(false)
+    })
+
+    it('renders slot when Firebase initializes with no user', async () => {
+      mockIsInitialized.value = false
+
+      const wrapper = mountComponent()
+      expect(wrapper.find('.progress-spinner').exists()).toBe(true)
+
+      mockIsInitialized.value = true
+      mockCurrentUser.value = null
+      await flushPromises()
+
+      expect(wrapper.find('[data-testid="slot-content"]').exists()).toBe(true)
+      expect(mockRefreshRemoteConfig).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('cloud builds - authenticated user', () => {
+    beforeEach(() => {
+      mockIsInitialized.value = true
+      mockCurrentUser.value = { uid: 'user-123' }
+    })
+
+    it('refreshes remote config with auth after Firebase init', async () => {
+      mountComponent()
+      await flushPromises()
+
+      expect(mockRefreshRemoteConfig).toHaveBeenCalledWith({ useAuth: true })
+    })
+
+    it('renders slot when teamWorkspacesEnabled is false', async () => {
+      mockTeamWorkspacesEnabled.value = false
+
+      const wrapper = mountComponent()
+      await flushPromises()
+
+      expect(wrapper.find('[data-testid="slot-content"]').exists()).toBe(true)
+      expect(mockWorkspaceStoreInitialize).not.toHaveBeenCalled()
+    })
+
+    it('initializes workspace store when teamWorkspacesEnabled is true', async () => {
+      mockTeamWorkspacesEnabled.value = true
+
+      const wrapper = mountComponent()
+      await flushPromises()
+
+      expect(mockWorkspaceStoreInitialize).toHaveBeenCalled()
+      expect(wrapper.find('[data-testid="slot-content"]').exists()).toBe(true)
+    })
+
+    it('skips workspace init when store is already initialized', async () => {
+      mockTeamWorkspacesEnabled.value = true
+      mockWorkspaceStoreInitState.value = 'ready'
+
+      const wrapper = mountComponent()
+      await flushPromises()
+
+      expect(mockWorkspaceStoreInitialize).not.toHaveBeenCalled()
+      expect(wrapper.find('[data-testid="slot-content"]').exists()).toBe(true)
+    })
+  })
+
+  describe('error handling - graceful degradation', () => {
+    beforeEach(() => {
+      mockIsInitialized.value = true
+      mockCurrentUser.value = { uid: 'user-123' }
+    })
+
+    it('renders slot when remote config refresh fails', async () => {
+      mockRefreshRemoteConfig.mockRejectedValue(new Error('Network error'))
+
+      const wrapper = mountComponent()
+      await flushPromises()
+
+      expect(wrapper.find('[data-testid="slot-content"]').exists()).toBe(true)
+    })
+
+    it('renders slot when remote config refresh times out', async () => {
+      vi.useFakeTimers()
+      // Never-resolving promise simulates a hanging request
+      mockRefreshRemoteConfig.mockReturnValue(new Promise(() => {}))
+
+      const wrapper = mountComponent()
+      await flushPromises()
+
+      // Still showing spinner before timeout
+      expect(wrapper.find('.progress-spinner').exists()).toBe(true)
+
+      // Advance past the 10 second timeout
+      await vi.advanceTimersByTimeAsync(10_001)
+      await flushPromises()
+
+      // Should render slot after timeout (graceful degradation)
+      expect(wrapper.find('[data-testid="slot-content"]').exists()).toBe(true)
+      vi.useRealTimers()
+    })
+
+    it('renders slot when workspace store initialization fails', async () => {
+      mockTeamWorkspacesEnabled.value = true
+      mockWorkspaceStoreInitialize.mockRejectedValue(
+        new Error('Workspace init failed')
+      )
+
+      const wrapper = mountComponent()
+      await flushPromises()
+
+      expect(wrapper.find('[data-testid="slot-content"]').exists()).toBe(true)
+    })
+  })
+})

src/components/auth/WorkspaceAuthGate.vue

@@ -0,0 +1,126 @@
+<template>
+  <slot v-if="isReady" />
+  <div
+    v-else
+    class="fixed inset-0 z-[1100] flex items-center justify-center bg-[var(--p-mask-background)]"
+  >
+    <ProgressSpinner />
+  </div>
+</template>
+
+<script setup lang="ts">
+/**
+ * WorkspaceAuthGate - Conditional auth checkpoint for workspace mode.
+ *
+ * This gate ensures proper initialization order for workspace-scoped auth:
+ * 1. Wait for Firebase auth to resolve
+ * 2. Check if teamWorkspacesEnabled feature flag is on
+ * 3. If YES: Initialize workspace token and store before rendering
+ * 4. If NO: Render immediately using existing Firebase auth
+ *
+ * This prevents race conditions where API calls use Firebase tokens
+ * instead of workspace tokens when the workspace feature is enabled.
+ */
+import { promiseTimeout, until } from '@vueuse/core'
+import { storeToRefs } from 'pinia'
+import ProgressSpinner from 'primevue/progressspinner'
+import { ref } from 'vue'
+
+import { useFeatureFlags } from '@/composables/useFeatureFlags'
+import { isCloud } from '@/platform/distribution/types'
+import { refreshRemoteConfig } from '@/platform/remoteConfig/refreshRemoteConfig'
+import { useTeamWorkspaceStore } from '@/platform/workspace/stores/teamWorkspaceStore'
+import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
+
+const FIREBASE_INIT_TIMEOUT_MS = 16_000
+const CONFIG_REFRESH_TIMEOUT_MS = 10_000
+
+const isReady = ref(!isCloud)
+
+async function initialize(): Promise<void> {
+  if (!isCloud) return
+
+  const authStore = useFirebaseAuthStore()
+  const { isInitialized, currentUser } = storeToRefs(authStore)
+
+  try {
+    // Step 1: Wait for Firebase auth to resolve
+    // This is shared with router guard - both wait for the same thing,
+    // but this gate blocks rendering while router guard blocks navigation
+    if (!isInitialized.value) {
+      await until(isInitialized).toBe(true, {
+        timeout: FIREBASE_INIT_TIMEOUT_MS
+      })
+    }
+
+    // Step 2: If not authenticated, nothing more to do
+    // Unauthenticated users don't have workspace context
+    if (!currentUser.value) {
+      isReady.value = true
+      return
+    }
+
+    // Step 3: Refresh feature flags with auth context
+    // This ensures teamWorkspacesEnabled reflects the authenticated user's state
+    // Timeout prevents hanging if server is slow/unresponsive
+    try {
+      await Promise.race([
+        refreshRemoteConfig({ useAuth: true }),
+        promiseTimeout(CONFIG_REFRESH_TIMEOUT_MS).then(() => {
+          throw new Error('Config refresh timeout')
+        })
+      ])
+    } catch (error) {
+      console.warn(
+        '[WorkspaceAuthGate] Failed to refresh remote config:',
+        error
+      )
+      // Continue - feature flags will use defaults (teamWorkspacesEnabled=false)
+      // App will render with Firebase auth fallback
+    }
+
+    // Step 4: THE CHECKPOINT - Are we in workspace mode?
+    const { flags } = useFeatureFlags()
+    if (!flags.teamWorkspacesEnabled) {
+      // Not in workspace mode - use existing Firebase auth flow
+      // No additional initialization needed
+      isReady.value = true
+      return
+    }
+
+    // Step 5: WORKSPACE MODE - Full initialization
+    await initializeWorkspaceMode()
+  } catch (error) {
+    console.error('[WorkspaceAuthGate] Initialization failed:', error)
+  } finally {
+    // Always render (graceful degradation)
+    // If workspace init failed, API calls fall back to Firebase token
+    isReady.value = true
+  }
+}
+
+async function initializeWorkspaceMode(): Promise<void> {
+  // Initialize the full workspace store which handles:
+  // - Restoring workspace token from session (fast path for refresh)
+  // - Fetching workspace list
+  // - Switching to last used workspace if needed
+  // - Setting active workspace
+  try {
+    const workspaceStore = useTeamWorkspaceStore()
+    if (workspaceStore.initState === 'uninitialized') {
+      await workspaceStore.initialize()
+    }
+  } catch (error) {
+    // Log but don't block - workspace UI features may not work but app will render
+    // API calls will fall back to Firebase token
+    console.warn(
+      '[WorkspaceAuthGate] Failed to initialize workspace store:',
+      error
+    )
+  }
+}
+
+// Start initialization immediately during component setup
+// (not in onMounted, so initialization starts before DOM is ready)
+void initialize()
+</script>

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)

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)

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)

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'
   }
 }

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<RemoteConfigState>('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

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)
         }

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) {

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<string | undefined> => {
+    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<GetCustomerBalanceResponse | null> => {
     isFetchingBalance.value = true
     try {
@@ -513,6 +538,7 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
     updatePassword: _updatePassword,
     deleteAccount: _deleteAccount,
     getAuthHeader,
-    getFirebaseAuthHeader
+    getFirebaseAuthHeader,
+    getAuthToken
   }
 })

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)