[API Node] Show user state when logged in via API key (#3838)

Co-authored-by: github-actions <github-actions@github.com>
Co-authored-by: Chenlei Hu <hcl@comfy.org>

src/components/dialog/content/setting/UserPanel.vue

@@ -5,10 +5,10 @@
       <Divider class="mb-3" />
 
       <!-- Normal User Panel -->
-      <div v-if="user" class="flex flex-col gap-2">
+      <div v-if="isLoggedIn" class="flex flex-col gap-2">
         <UserAvatar
-          v-if="user.photoURL"
-          :photo-url="user.photoURL"
+          v-if="userPhotoUrl"
+          :photo-url="userPhotoUrl"
           shape="circle"
           size="large"
         />
@@ -18,7 +18,7 @@
             {{ $t('userSettings.name') }}
           </h3>
           <div class="text-muted">
-            {{ user.displayName || $t('userSettings.notSet') }}
+            {{ userDisplayName || $t('userSettings.notSet') }}
           </div>
         </div>
 
@@ -26,8 +26,8 @@
           <h3 class="font-medium">
             {{ $t('userSettings.email') }}
           </h3>
-          <a :href="'mailto:' + user.email" class="hover:underline">
-            {{ user.email }}
+          <a :href="'mailto:' + userEmail" class="hover:underline">
+            {{ userEmail }}
           </a>
         </div>
 
@@ -67,27 +67,6 @@
         />
       </div>
 
-      <!-- API Key Panel -->
-      <div v-else-if="hasApiKey" class="flex flex-col gap-4">
-        <div class="flex flex-col gap-0.5">
-          <h3 class="font-medium">
-            {{ $t('auth.apiKey.title') }}
-          </h3>
-          <div class="text-muted flex items-center gap-1">
-            <i class="pi pi-key" />
-            {{ $t('auth.apiKey.label') }}
-          </div>
-        </div>
-
-        <Button
-          class="mt-4 w-32"
-          severity="secondary"
-          :label="$t('auth.signOut.signOut')"
-          icon="pi pi-sign-out"
-          @click="handleApiKeySignOut"
-        />
-      </div>
-
       <!-- Login Section -->
       <div v-else class="flex flex-col gap-4">
         <p class="text-gray-600">
@@ -112,59 +91,22 @@ import Button from 'primevue/button'
 import Divider from 'primevue/divider'
 import ProgressSpinner from 'primevue/progressspinner'
 import TabPanel from 'primevue/tabpanel'
-import { computed } from 'vue'
 
 import UserAvatar from '@/components/common/UserAvatar.vue'
+import { useCurrentUser } from '@/composables/auth/useCurrentUser'
 import { useDialogService } from '@/services/dialogService'
-import { useApiKeyAuthStore } from '@/stores/apiKeyAuthStore'
-import { useCommandStore } from '@/stores/commandStore'
-import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
 
 const dialogService = useDialogService()
-const authStore = useFirebaseAuthStore()
-const commandStore = useCommandStore()
-const apiKeyStore = useApiKeyAuthStore()
-
-const user = computed(() => authStore.currentUser)
-const loading = computed(() => authStore.loading)
-const hasApiKey = computed(() => apiKeyStore.hasApiKey)
-
-const providerName = computed(() => {
-  const providerId = user.value?.providerData[0]?.providerId
-  if (providerId?.includes('google')) {
-    return 'Google'
-  }
-  if (providerId?.includes('github')) {
-    return 'GitHub'
-  }
-  return providerId
-})
-
-const providerIcon = computed(() => {
-  const providerId = user.value?.providerData[0]?.providerId
-  if (providerId?.includes('google')) {
-    return 'pi pi-google'
-  }
-  if (providerId?.includes('github')) {
-    return 'pi pi-github'
-  }
-  return 'pi pi-user'
-})
-
-const isEmailProvider = computed(() => {
-  const providerId = user.value?.providerData[0]?.providerId
-  return providerId === 'password'
-})
-
-const handleSignOut = async () => {
-  await commandStore.execute('Comfy.User.SignOut')
-}
-
-const handleApiKeySignOut = async () => {
-  await apiKeyStore.clearStoredApiKey()
-}
-
-const handleSignIn = async () => {
-  await commandStore.execute('Comfy.User.OpenSignInDialog')
-}
+const {
+  loading,
+  isLoggedIn,
+  isEmailProvider,
+  userDisplayName,
+  userEmail,
+  userPhotoUrl,
+  providerName,
+  providerIcon,
+  handleSignOut,
+  handleSignIn
+} = useCurrentUser()
 </script>

src/components/topbar/CurrentUserButton.vue

@@ -2,7 +2,7 @@
 <template>
   <div>
     <Button
-      v-if="isAuthenticated"
+      v-if="isLoggedIn"
       class="user-profile-button p-1"
       severity="secondary"
       text
@@ -30,15 +30,14 @@ import Popover from 'primevue/popover'
 import { computed, ref } from 'vue'
 
 import UserAvatar from '@/components/common/UserAvatar.vue'
-import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
+import { useCurrentUser } from '@/composables/auth/useCurrentUser'
 
 import CurrentUserPopover from './CurrentUserPopover.vue'
 
-const authStore = useFirebaseAuthStore()
+const { isLoggedIn, userPhotoUrl } = useCurrentUser()
 
 const popover = ref<InstanceType<typeof Popover> | null>(null)
-const isAuthenticated = computed(() => authStore.isAuthenticated)
 const photoURL = computed<string | undefined>(
-  () => authStore.currentUser?.photoURL ?? undefined
+  () => userPhotoUrl.value ?? undefined
 )
 </script>

src/components/topbar/CurrentUserPopover.vue

@@ -6,19 +6,19 @@
       <div class="flex flex-col items-center">
         <UserAvatar
           class="mb-3"
-          :photo-url="user?.photoURL"
+          :photo-url="userPhotoUrl"
           :pt:icon:class="{
-            '!text-2xl': !user?.photoURL
+            '!text-2xl': !userPhotoUrl
           }"
           size="large"
         />
 
         <!-- User Details -->
         <h3 class="text-lg font-semibold truncate my-0 mb-1">
-          {{ user?.displayName || $t('g.user') }}
+          {{ userDisplayName || $t('g.user') }}
         </h3>
-        <p v-if="user?.email" class="text-sm text-muted truncate my-0">
-          {{ user.email }}
+        <p v-if="userEmail" class="text-sm text-muted truncate my-0">
+          {{ userEmail }}
         </p>
       </div>
     </div>
@@ -64,20 +64,18 @@
 <script setup lang="ts">
 import Button from 'primevue/button'
 import Divider from 'primevue/divider'
-import { computed, onMounted } from 'vue'
+import { onMounted } from 'vue'
 
 import UserAvatar from '@/components/common/UserAvatar.vue'
 import UserCredit from '@/components/common/UserCredit.vue'
+import { useCurrentUser } from '@/composables/auth/useCurrentUser'
 import { useDialogService } from '@/services/dialogService'
 import { useFirebaseAuthService } from '@/services/firebaseAuthService'
-import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
 
-const authStore = useFirebaseAuthStore()
+const { userDisplayName, userEmail, userPhotoUrl } = useCurrentUser()
 const authService = useFirebaseAuthService()
 const dialogService = useDialogService()
 
-const user = computed(() => authStore.currentUser)
-
 const handleOpenUserSettings = () => {
   dialogService.showSettingsDialog('user')
 }

src/composables/auth/useCurrentUser.ts

@@ -0,0 +1,101 @@
+import { computed } from 'vue'
+
+import { useApiKeyAuthStore } from '@/stores/apiKeyAuthStore'
+import { useCommandStore } from '@/stores/commandStore'
+import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
+
+export const useCurrentUser = () => {
+  const authStore = useFirebaseAuthStore()
+  const commandStore = useCommandStore()
+  const apiKeyStore = useApiKeyAuthStore()
+
+  const firebaseUser = computed(() => authStore.currentUser)
+  const isApiKeyLogin = computed(() => apiKeyStore.isAuthenticated)
+  const isLoggedIn = computed(
+    () => !!isApiKeyLogin.value || firebaseUser.value !== null
+  )
+
+  const userDisplayName = computed(() => {
+    if (isApiKeyLogin.value) {
+      return apiKeyStore.currentUser?.name
+    }
+    return firebaseUser.value?.displayName
+  })
+
+  const userEmail = computed(() => {
+    if (isApiKeyLogin.value) {
+      return apiKeyStore.currentUser?.email
+    }
+    return firebaseUser.value?.email
+  })
+
+  const providerName = computed(() => {
+    if (isApiKeyLogin.value) {
+      return 'Comfy API Key'
+    }
+
+    const providerId = firebaseUser.value?.providerData[0]?.providerId
+    if (providerId?.includes('google')) {
+      return 'Google'
+    }
+    if (providerId?.includes('github')) {
+      return 'GitHub'
+    }
+    return providerId
+  })
+
+  const providerIcon = computed(() => {
+    if (isApiKeyLogin.value) {
+      return 'pi pi-key'
+    }
+
+    const providerId = firebaseUser.value?.providerData[0]?.providerId
+    if (providerId?.includes('google')) {
+      return 'pi pi-google'
+    }
+    if (providerId?.includes('github')) {
+      return 'pi pi-github'
+    }
+    return 'pi pi-user'
+  })
+
+  const isEmailProvider = computed(() => {
+    if (isApiKeyLogin.value) {
+      return false
+    }
+
+    const providerId = firebaseUser.value?.providerData[0]?.providerId
+    return providerId === 'password'
+  })
+
+  const userPhotoUrl = computed(() => {
+    if (isApiKeyLogin.value) return null
+    return firebaseUser.value?.photoURL
+  })
+
+  const handleSignOut = async () => {
+    if (isApiKeyLogin.value) {
+      await apiKeyStore.clearStoredApiKey()
+    } else {
+      await commandStore.execute('Comfy.User.SignOut')
+    }
+  }
+
+  const handleSignIn = async () => {
+    await commandStore.execute('Comfy.User.OpenSignInDialog')
+  }
+
+  return {
+    loading: authStore.loading,
+    isLoggedIn,
+    isApiKeyLogin,
+    isEmailProvider,
+    userDisplayName,
+    userEmail,
+    userPhotoUrl,
+    providerName,
+    providerIcon,
+    handleSignOut,
+    handleSignIn
+  }
+}

src/composables/setting/useSettingUI.ts

@@ -7,13 +7,14 @@ import {
 } from 'vue'
 import { useI18n } from 'vue-i18n'
 
-import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
 import { SettingTreeNode, useSettingStore } from '@/stores/settingStore'
 import type { SettingParams } from '@/types/settingTypes'
 import { isElectron } from '@/utils/envUtil'
 import { normalizeI18nKey } from '@/utils/formatUtil'
 import { buildTree } from '@/utils/treeUtil'
 
+import { useCurrentUser } from '../auth/useCurrentUser'
+
 interface SettingPanelItem {
   node: SettingTreeNode
   component: Component
@@ -29,7 +30,7 @@ export function useSettingUI(
     | 'credits'
 ) {
   const { t } = useI18n()
-  const firebaseAuthStore = useFirebaseAuthStore()
+  const { isLoggedIn } = useCurrentUser()
   const settingStore = useSettingStore()
   const activeCategory = ref<SettingTreeNode | null>(null)
 
@@ -165,7 +166,7 @@ export function useSettingUI(
       label: 'Account',
       children: [
         userPanel.node,
-        ...(firebaseAuthStore.isAuthenticated ? [creditsPanel.node] : [])
+        ...(isLoggedIn.value ? [creditsPanel.node] : [])
       ].map(translateCategory)
     },
     // Normal settings stored in the settingStore

src/locales/en/main.json

@@ -1308,7 +1308,8 @@
       "success": "Login successful",
       "failed": "Login failed",
       "insecureContextWarning": "This connection is insecure (HTTP) - your credentials may be intercepted by attackers if you proceed to login.",
-      "questionsContactPrefix": "Questions? Contact us at"
+      "questionsContactPrefix": "Questions? Contact us at",
+      "noAssociatedUser": "There is no Comfy user associated with the provided API key"
     },
     "signup": {
       "title": "Create an account",

src/locales/es/main.json

@@ -61,6 +61,7 @@
       "loginWithGithub": "Iniciar sesión con Github",
       "loginWithGoogle": "Iniciar sesión con Google",
       "newUser": "¿Eres nuevo aquí?",
+      "noAssociatedUser": "No hay ningún usuario de Comfy asociado con la clave API proporcionada",
       "orContinueWith": "O continuar con",
       "passwordLabel": "Contraseña",
       "passwordPlaceholder": "Ingresa tu contraseña",

src/locales/fr/main.json

@@ -61,6 +61,7 @@
       "loginWithGithub": "Se connecter avec Github",
       "loginWithGoogle": "Se connecter avec Google",
       "newUser": "Nouveau ici?",
+      "noAssociatedUser": "Aucun utilisateur Comfy n'est associé à la clé API fournie",
       "orContinueWith": "Ou continuer avec",
       "passwordLabel": "Mot de passe",
       "passwordPlaceholder": "Entrez votre mot de passe",

src/locales/ja/main.json

@@ -61,6 +61,7 @@
       "loginWithGithub": "Githubでログイン",
       "loginWithGoogle": "Googleでログイン",
       "newUser": "新規ユーザーですか？",
+      "noAssociatedUser": "指定されたAPIキーに関連付けられたComfyユーザーが存在しません",
       "orContinueWith": "または以下で続ける",
       "passwordLabel": "パスワード",
       "passwordPlaceholder": "パスワードを入力してください",

src/locales/ko/main.json

@@ -61,6 +61,7 @@
       "loginWithGithub": "Github로 로그인",
       "loginWithGoogle": "구글로 로그인",
       "newUser": "처음이신가요?",
+      "noAssociatedUser": "제공된 API 키와 연결된 Comfy 사용자가 없습니다",
       "orContinueWith": "또는 다음으로 계속",
       "passwordLabel": "비밀번호",
       "passwordPlaceholder": "비밀번호를 입력하세요",

src/locales/ru/main.json

@@ -61,6 +61,7 @@
       "loginWithGithub": "Войти через Github",
       "loginWithGoogle": "Войти через Google",
       "newUser": "Вы здесь впервые?",
+      "noAssociatedUser": "С предоставленным API-ключом не связан ни один пользователь Comfy",
       "orContinueWith": "Или продолжить с",
       "passwordLabel": "Пароль",
       "passwordPlaceholder": "Введите ваш пароль",

src/locales/zh/main.json

@@ -61,6 +61,7 @@
       "loginWithGithub": "使用Github登录",
       "loginWithGoogle": "使用Google登录",
       "newUser": "新来的？",
+      "noAssociatedUser": "所提供的 API 密钥未关联任何 Comfy 用户",
       "orContinueWith": "或者继续使用",
       "passwordLabel": "密码",
       "passwordPlaceholder": "输入您的密码",

src/stores/apiKeyAuthStore.ts

@@ -1,18 +1,51 @@
 import { useLocalStorage } from '@vueuse/core'
 import { defineStore } from 'pinia'
-import { computed } from 'vue'
+import { computed, ref, watch } from 'vue'
 
 import { useErrorHandling } from '@/composables/useErrorHandling'
 import { t } from '@/i18n'
+import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
 import { useToastStore } from '@/stores/toastStore'
+import { ApiKeyAuthHeader } from '@/types/authTypes'
+import { operations } from '@/types/comfyRegistryTypes'
+
+type ComfyApiUser =
+  operations['createCustomer']['responses']['201']['content']['application/json']
 
 const STORAGE_KEY = 'comfy_api_key'
 
 export const useApiKeyAuthStore = defineStore('apiKeyAuth', () => {
+  const firebaseAuthStore = useFirebaseAuthStore()
   const apiKey = useLocalStorage<string | null>(STORAGE_KEY, null)
   const toastStore = useToastStore()
   const { wrapWithErrorHandlingAsync, toastErrorHandler } = useErrorHandling()
 
+  const currentUser = ref<ComfyApiUser | null>(null)
+  const isAuthenticated = computed(() => !!currentUser.value)
+
+  const initializeUserFromApiKey = async () => {
+    const createCustomerResponse = await firebaseAuthStore.createCustomer()
+    if (!createCustomerResponse) {
+      apiKey.value = null
+      throw new Error(t('auth.login.noAssociatedUser'))
+    }
+    currentUser.value = createCustomerResponse
+  }
+
+  watch(
+    apiKey,
+    () => {
+      if (apiKey.value) {
+        // IF API key is set, initialize user
+        void initializeUserFromApiKey()
+      } else {
+        // IF API key is cleared, clear user
+        currentUser.value = null
+      }
+    },
+    { immediate: true }
+  )
+
   const reportError = (error: unknown) => {
     if (error instanceof Error && error.message === 'STORAGE_FAILED') {
       toastStore.add({
@@ -49,7 +82,11 @@ export const useApiKeyAuthStore = defineStore('apiKeyAuth', () => {
 
   const getApiKey = () => apiKey.value
 
-  const getAuthHeader = () => {
+  /**
+   * Retrieves the appropriate authentication header for API requests if an
+   * API key is available, otherwise returns null.
+   */
+  const getAuthHeader = (): ApiKeyAuthHeader | null => {
     const comfyOrgApiKey = getApiKey()
     if (comfyOrgApiKey) {
       return {
@@ -60,7 +97,11 @@ export const useApiKeyAuthStore = defineStore('apiKeyAuth', () => {
   }
 
   return {
-    hasApiKey: computed(() => !!apiKey.value),
+    // State
+    currentUser,
+    isAuthenticated,
+
+    // Actions
     storeApiKey,
     clearStoredApiKey,
     getAuthHeader,

src/stores/firebaseAuthStore.ts

@@ -107,6 +107,7 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
    *   - null if neither authentication method is available
    */
   const getAuthHeader = async (): Promise<AuthHeader | null> => {
+    // If available, set header with JWT used to identify the user to Firebase service
     const token = await getIdToken()
     if (token) {
       return {
@@ -114,8 +115,8 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
       }
     }
 
-    const apiKeyStore = useApiKeyAuthStore()
-    return apiKeyStore.getAuthHeader()
+    // If not authenticated with Firebase, try falling back to API key if available
+    return useApiKeyAuthStore().getAuthHeader()
   }
 
   const fetchBalance = async (): Promise<GetCustomerBalanceResponse | null> => {
@@ -356,6 +357,7 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
     login,
     register,
     logout,
+    createCustomer,
     getIdToken,
     loginWithGoogle,
     loginWithGithub,