From 2e677ef6853705a339d0c1a66b127cdb9a663e71 Mon Sep 17 00:00:00 2001
From: --list <18093452+simula-r@users.noreply.github.com>
Date: Tue, 20 Jan 2026 12:10:10 -0800
Subject: [PATCH] feat(workspace): add core flow integration

- Add teamWorkspacesEnabled feature flag with isCloud guard
- Update firebaseAuthStore.getAuthHeader() to prioritize workspace token
- Add getFirebaseAuthHeader() for user-scoped endpoints
- Initialize workspace store in router guard before app loads
- Load invite from URL in GraphCanvas after app initialization
- Update CurrentUserButton to show workspace popover when enabled

BREAKING: When teamWorkspacesEnabled flag is true:
- API calls use workspace-scoped tokens
- User menu shows workspace context
- Settings panel shows workspace mode layout
---
 src/components/graph/GraphCanvas.vue        | 10 +++++
 src/components/topbar/CurrentUserButton.vue | 48 ++++++++++++++++++---
 src/composables/useFeatureFlags.ts          |  3 ++
 src/router.ts                               | 22 ++++++++++
 src/stores/firebaseAuthStore.ts             | 29 ++++++++-----
 src/views/GraphView.vue                     |  2 +
 6 files changed, 98 insertions(+), 16 deletions(-)

diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue
index 48d82e44b..60b0e6d26 100644
--- a/src/components/graph/GraphCanvas.vue
+++ b/src/components/graph/GraphCanvas.vue
@@ -126,11 +126,13 @@ import { useNodeBadge } from '@/composables/node/useNodeBadge'
 import { useCanvasDrop } from '@/composables/useCanvasDrop'
 import { useContextMenuTranslation } from '@/composables/useContextMenuTranslation'
 import { useCopy } from '@/composables/useCopy'
+import { useFeatureFlags } from '@/composables/useFeatureFlags'
 import { useGlobalLitegraph } from '@/composables/useGlobalLitegraph'
 import { usePaste } from '@/composables/usePaste'
 import { useVueFeatureFlags } from '@/composables/useVueFeatureFlags'
 import { mergeCustomNodesI18n, t } from '@/i18n'
 import { LiteGraph } from '@/lib/litegraph/src/litegraph'
+import { isCloud } from '@/platform/distribution/types'
 import { useLitegraphSettings } from '@/platform/settings/composables/useLitegraphSettings'
 import { CORE_SETTINGS } from '@/platform/settings/constants/coreSettings'
 import { useSettingStore } from '@/platform/settings/settingStore'
@@ -394,6 +396,7 @@ const loadCustomNodesI18n = async () => {
 
 const comfyAppReady = ref(false)
 const workflowPersistence = useWorkflowPersistence()
+const { flags } = useFeatureFlags()
 useCanvasDrop(canvasRef)
 useLitegraphSettings()
 useNodeBadge()
@@ -459,6 +462,13 @@ onMounted(async () => {
   // Load template from URL if present
   await workflowPersistence.loadTemplateFromUrlIfPresent()
 
+  // Accept workspace invite from URL if present (e.g., ?invite=TOKEN)
+  if (isCloud && flags.teamWorkspacesEnabled) {
+    const { useInviteUrlLoader } =
+      await import('@/platform/workspace/composables/useInviteUrlLoader')
+    await useInviteUrlLoader().loadInviteFromUrl()
+  }
+
   // Initialize release store to fetch releases from comfy-api (fire-and-forget)
   const { useReleaseStore } =
     await import('@/platform/updates/common/releaseStore')
diff --git a/src/components/topbar/CurrentUserButton.vue b/src/components/topbar/CurrentUserButton.vue
index 151dfd405..849a9ba25 100644
--- a/src/components/topbar/CurrentUserButton.vue
+++ b/src/components/topbar/CurrentUserButton.vue
@@ -1,4 +1,4 @@
-<!-- A button that shows current authenticated user's avatar -->
+<!-- A button that shows workspace icon (Cloud) or user avatar -->
 <template>
   <div>
     <Button
@@ -16,7 +16,16 @@
           )
         "
       >
-        <UserAvatar :photo-url="photoURL" :class="compact && 'size-full'" />
+        <WorkspaceProfilePic
+          v-if="showWorkspaceIcon"
+          :workspace-name="workspaceName"
+          :class="compact && 'size-full'"
+        />
+        <UserAvatar
+          v-else
+          :photo-url="photoURL"
+          :class="compact && 'size-full'"
+        />
 
         <i v-if="showArrow" class="icon-[lucide--chevron-down] size-3 px-1" />
       </div>
@@ -27,38 +36,65 @@
       :show-arrow="false"
       :pt="{
         root: {
-          class: 'rounded-lg'
+          class: 'rounded-lg w-80'
         }
       }"
     >
-      <CurrentUserPopover @close="closePopover" />
+      <!-- Workspace mode: workspace-aware popover -->
+      <CurrentUserPopoverWorkspace
+        v-if="teamWorkspacesEnabled"
+        @close="closePopover"
+      />
+      <!-- Legacy mode: original popover -->
+      <CurrentUserPopover v-else @close="closePopover" />
     </Popover>
   </div>
 </template>
 
 <script setup lang="ts">
+import { storeToRefs } from 'pinia'
 import Popover from 'primevue/popover'
-import { computed, ref } from 'vue'
+import { computed, defineAsyncComponent, ref } from 'vue'
 
 import UserAvatar from '@/components/common/UserAvatar.vue'
+import WorkspaceProfilePic from '@/components/common/WorkspaceProfilePic.vue'
 import Button from '@/components/ui/button/Button.vue'
 import { useCurrentUser } from '@/composables/auth/useCurrentUser'
+import { useFeatureFlags } from '@/composables/useFeatureFlags'
+import { isCloud } from '@/platform/distribution/types'
+import { useTeamWorkspaceStore } from '@/platform/workspace/stores/teamWorkspaceStore'
 import { cn } from '@/utils/tailwindUtil'
 
 import CurrentUserPopover from './CurrentUserPopover.vue'
 
+const CurrentUserPopoverWorkspace = defineAsyncComponent(
+  () => import('./CurrentUserPopoverWorkspace.vue')
+)
+
 const { showArrow = true, compact = false } = defineProps<{
   showArrow?: boolean
   compact?: boolean
 }>()
 
+const { flags } = useFeatureFlags()
+const teamWorkspacesEnabled = computed(() => flags.teamWorkspacesEnabled)
+
 const { isLoggedIn, userPhotoUrl } = useCurrentUser()
 
-const popover = ref<InstanceType<typeof Popover> | null>(null)
 const photoURL = computed<string | undefined>(
   () => userPhotoUrl.value ?? undefined
 )
 
+const showWorkspaceIcon = computed(() => isCloud && teamWorkspacesEnabled.value)
+
+const workspaceName = computed(() => {
+  if (!showWorkspaceIcon.value) return ''
+  const { workspaceName } = storeToRefs(useTeamWorkspaceStore())
+  return workspaceName.value
+})
+
+const popover = ref<InstanceType<typeof Popover> | null>(null)
+
 const closePopover = () => {
   popover.value?.hide()
 }
diff --git a/src/composables/useFeatureFlags.ts b/src/composables/useFeatureFlags.ts
index 136b7ccd1..ca54bb9c6 100644
--- a/src/composables/useFeatureFlags.ts
+++ b/src/composables/useFeatureFlags.ts
@@ -1,5 +1,6 @@
 import { computed, reactive, readonly } from 'vue'
 
+import { isCloud } from '@/platform/distribution/types'
 import { remoteConfig } from '@/platform/remoteConfig/remoteConfig'
 import { api } from '@/scripts/api'
 
@@ -95,6 +96,8 @@ export function useFeatureFlags() {
       )
     },
     get teamWorkspacesEnabled() {
+      if (!isCloud) return false
+
       return (
         remoteConfig.value.team_workspaces_enabled ??
         api.getServerFeature(ServerFeatureFlag.TEAM_WORKSPACES_ENABLED, false)
diff --git a/src/router.ts b/src/router.ts
index 102ec8f2a..a8ffad40b 100644
--- a/src/router.ts
+++ b/src/router.ts
@@ -86,6 +86,10 @@ installPreservedQueryTracker(router, [
   {
     namespace: PRESERVED_QUERY_NAMESPACES.TEMPLATE,
     keys: ['template', 'source', 'mode']
+  },
+  {
+    namespace: PRESERVED_QUERY_NAMESPACES.INVITE,
+    keys: ['invite']
   }
 ])
 
@@ -178,6 +182,24 @@ if (isCloud) {
       })
     }
 
+    // Initialize workspace context for logged-in users navigating to root
+    // This must happen before the app loads to ensure workspace context is ready
+
+    if (to.path === '/' && flags.teamWorkspacesEnabled) {
+      const { useTeamWorkspaceStore } =
+        await import('@/platform/workspace/stores/teamWorkspaceStore')
+      const workspaceStore = useTeamWorkspaceStore()
+
+      if (workspaceStore.initState === 'uninitialized') {
+        try {
+          await workspaceStore.initialize()
+        } catch (error) {
+          console.error('Workspace initialization failed:', error)
+          // Continue anyway - workspace features will be degraded
+        }
+      }
+    }
+
     // User is logged in - check if they need onboarding (when enabled)
     // For root path, check actual user status to handle waitlisted users
     if (!isElectron() && isLoggedIn && to.path === '/') {
diff --git a/src/stores/firebaseAuthStore.ts b/src/stores/firebaseAuthStore.ts
index baa1840a8..943318240 100644
--- a/src/stores/firebaseAuthStore.ts
+++ b/src/stores/firebaseAuthStore.ts
@@ -25,12 +25,12 @@ import { getComfyApiBaseUrl } from '@/config/comfyApi'
 import { t } from '@/i18n'
 import { WORKSPACE_STORAGE_KEYS } from '@/platform/auth/workspace/workspaceConstants'
 import { isCloud } from '@/platform/distribution/types'
-import { remoteConfig } from '@/platform/remoteConfig/remoteConfig'
 import { useTelemetry } from '@/platform/telemetry'
 import { useDialogService } from '@/services/dialogService'
 import { useApiKeyAuthStore } from '@/stores/apiKeyAuthStore'
 import type { AuthHeader } from '@/types/authTypes'
 import type { operations } from '@/types/comfyRegistryTypes'
+import { useFeatureFlags } from '@/composables/useFeatureFlags'
 
 type CreditPurchaseResponse =
   operations['InitiateCreditPurchase']['responses']['201']['content']['application/json']
@@ -58,6 +58,8 @@ export class FirebaseAuthStoreError extends Error {
 }
 
 export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
+  const { flags } = useFeatureFlags()
+
   // State
   const loading = ref(false)
   const currentUser = ref<User | null>(null)
@@ -173,7 +175,7 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
    *   - null if no authentication method is available
    */
   const getAuthHeader = async (): Promise<AuthHeader | null> => {
-    if (remoteConfig.value.team_workspaces_enabled) {
+    if (flags.teamWorkspacesEnabled) {
       const workspaceToken = sessionStorage.getItem(
         WORKSPACE_STORAGE_KEYS.TOKEN
       )
@@ -201,10 +203,19 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
     return useApiKeyAuthStore().getAuthHeader()
   }
 
+  /**
+   * Returns Firebase auth header for user-scoped endpoints (e.g., /customers/*).
+   * Use this for endpoints that need user identity, not workspace context.
+   */
+  const getFirebaseAuthHeader = async (): Promise<AuthHeader | null> => {
+    const token = await getIdToken()
+    return token ? { Authorization: `Bearer ${token}` } : null
+  }
+
   const fetchBalance = async (): Promise<GetCustomerBalanceResponse | null> => {
     isFetchingBalance.value = true
     try {
-      const authHeader = await getAuthHeader()
+      const authHeader = await getFirebaseAuthHeader()
       if (!authHeader) {
         throw new FirebaseAuthStoreError(
           t('toastMessages.userNotAuthenticated')
@@ -242,7 +253,7 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
   }
 
   const createCustomer = async (): Promise<CreateCustomerResponse> => {
-    const authHeader = await getAuthHeader()
+    const authHeader = await getFirebaseAuthHeader()
     if (!authHeader) {
       throw new FirebaseAuthStoreError(t('toastMessages.userNotAuthenticated'))
     }
@@ -404,7 +415,7 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
   const addCredits = async (
     requestBodyContent: CreditPurchasePayload
   ): Promise<CreditPurchaseResponse> => {
-    const authHeader = await getAuthHeader()
+    const authHeader = await getFirebaseAuthHeader()
     if (!authHeader) {
       throw new FirebaseAuthStoreError(t('toastMessages.userNotAuthenticated'))
     }
@@ -444,21 +455,19 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
   const accessBillingPortal = async (
     targetTier?: BillingPortalTargetTier
   ): Promise<AccessBillingPortalResponse> => {
-    const authHeader = await getAuthHeader()
+    const authHeader = await getFirebaseAuthHeader()
     if (!authHeader) {
       throw new FirebaseAuthStoreError(t('toastMessages.userNotAuthenticated'))
     }
 
-    const requestBody = targetTier ? { target_tier: targetTier } : undefined
-
     const response = await fetch(buildApiUrl('/customers/billing'), {
       method: 'POST',
       headers: {
         ...authHeader,
         'Content-Type': 'application/json'
       },
-      ...(requestBody && {
-        body: JSON.stringify(requestBody)
+      ...(targetTier && {
+        body: JSON.stringify({ target_tier: targetTier })
       })
     })
 
diff --git a/src/views/GraphView.vue b/src/views/GraphView.vue
index 6dc309c1a..9df7d60bd 100644
--- a/src/views/GraphView.vue
+++ b/src/views/GraphView.vue
@@ -17,6 +17,7 @@
 
   <GlobalToast />
   <RerouteMigrationToast />
+  <WorkspaceCreatedToast />
   <ModelImportProgressDialog />
   <ManagerProgressToast />
   <UnloadWindowConfirmDialog v-if="!isElectron()" />
@@ -45,6 +46,7 @@ import UnloadWindowConfirmDialog from '@/components/dialog/UnloadWindowConfirmDi
 import GraphCanvas from '@/components/graph/GraphCanvas.vue'
 import GlobalToast from '@/components/toast/GlobalToast.vue'
 import RerouteMigrationToast from '@/components/toast/RerouteMigrationToast.vue'
+import WorkspaceCreatedToast from '@/components/toast/WorkspaceCreatedToast.vue'
 import { useBrowserTabTitle } from '@/composables/useBrowserTabTitle'
 import { useCoreCommands } from '@/composables/useCoreCommands'
 import { useErrorHandling } from '@/composables/useErrorHandling'