Fix session cookie creation race: skip initial token refresh, wrap extension auth hooks

src/platform/auth/session/useSessionCookie.ts

@@ -2,6 +2,15 @@ import { api } from '@/scripts/api'
 import { isCloud } from '@/platform/distribution/types'
 import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
 
+/**
+ * Tracks the in-flight createSession request to dedupe concurrent calls.
+ */
+let createSessionInFlight: Promise<void> | null = null
+/**
+ * Tracks the in-flight deleteSession request to dedupe concurrent calls.
+ */
+let deleteSessionInFlight: Promise<void> | null = null
+
 /**
  * Session cookie management for cloud authentication.
  * Creates and deletes session cookies on the ComfyUI server.
@@ -14,27 +23,40 @@ export const useSessionCookie = () => {
   const createSession = async (): Promise<void> => {
     if (!isCloud) return
 
-    const authStore = useFirebaseAuthStore()
-    const authHeader = await authStore.getAuthHeader()
-
-    if (!authHeader) {
-      throw new Error('No auth header available for session creation')
+    if (createSessionInFlight) {
+      await createSessionInFlight
+      return
     }
 
-    const response = await fetch(api.apiURL('/auth/session'), {
-      method: 'POST',
-      credentials: 'include',
-      headers: {
-        ...authHeader,
-        'Content-Type': 'application/json'
+    createSessionInFlight = (async () => {
+      const authStore = useFirebaseAuthStore()
+      const authHeader = await authStore.getAuthHeader()
+
+      if (!authHeader) {
+        throw new Error('No auth header available for session creation')
       }
-    })
 
-    if (!response.ok) {
-      const errorData = await response.json().catch(() => ({}))
-      throw new Error(
-        `Failed to create session: ${errorData.message || response.statusText}`
-      )
+      const response = await fetch(api.apiURL('/auth/session'), {
+        method: 'POST',
+        credentials: 'include',
+        headers: {
+          ...authHeader,
+          'Content-Type': 'application/json'
+        }
+      })
+
+      if (!response.ok) {
+        const errorData = await response.json().catch(() => ({}))
+        throw new Error(
+          `Failed to create session: ${errorData.message || response.statusText}`
+        )
+      }
+    })()
+
+    try {
+      await createSessionInFlight
+    } finally {
+      createSessionInFlight = null
     }
   }
 
@@ -45,16 +67,29 @@ export const useSessionCookie = () => {
   const deleteSession = async (): Promise<void> => {
     if (!isCloud) return
 
-    const response = await fetch(api.apiURL('/auth/session'), {
-      method: 'DELETE',
-      credentials: 'include'
-    })
+    if (deleteSessionInFlight) {
+      await deleteSessionInFlight
+      return
+    }
+
+    deleteSessionInFlight = (async () => {
+      const response = await fetch(api.apiURL('/auth/session'), {
+        method: 'DELETE',
+        credentials: 'include'
+      })
+
+      if (!response.ok) {
+        const errorData = await response.json().catch(() => ({}))
+        throw new Error(
+          `Failed to delete session: ${errorData.message || response.statusText}`
+        )
+      }
+    })()
 
-    if (!response.ok) {
-      const errorData = await response.json().catch(() => ({}))
-      throw new Error(
-        `Failed to delete session: ${errorData.message || response.statusText}`
-      )
+    try {
+      await deleteSessionInFlight
+    } finally {
+      deleteSessionInFlight = null
     }
   }
 

src/services/extensionService.ts

@@ -11,12 +11,14 @@ import { useMenuItemStore } from '@/stores/menuItemStore'
 import { useWidgetStore } from '@/stores/widgetStore'
 import { useBottomPanelStore } from '@/stores/workspace/bottomPanelStore'
 import type { ComfyExtension } from '@/types/comfy'
+import type { AuthUserInfo } from '@/types/authTypes'
 
 export const useExtensionService = () => {
   const extensionStore = useExtensionStore()
   const settingStore = useSettingStore()
   const keybindingStore = useKeybindingStore()
-  const { wrapWithErrorHandling } = useErrorHandling()
+  const { wrapWithErrorHandling, wrapWithErrorHandlingAsync } =
+    useErrorHandling()
 
   /**
    * Loads all extensions from the API into the window in parallel
@@ -77,22 +79,31 @@ export const useExtensionService = () => {
 
     if (extension.onAuthUserResolved) {
       const { onUserResolved } = useCurrentUser()
+      const handleUserResolved = wrapWithErrorHandlingAsync(
+        (user: AuthUserInfo) => extension.onAuthUserResolved?.(user, app)
+      )
       onUserResolved((user) => {
-        void extension.onAuthUserResolved?.(user, app)
+        void handleUserResolved(user)
       })
     }
 
     if (extension.onAuthTokenRefreshed) {
       const { onTokenRefreshed } = useCurrentUser()
+      const handleTokenRefreshed = wrapWithErrorHandlingAsync(() =>
+        extension.onAuthTokenRefreshed?.()
+      )
       onTokenRefreshed(() => {
-        void extension.onAuthTokenRefreshed?.()
+        void handleTokenRefreshed()
       })
     }
 
     if (extension.onAuthUserLogout) {
       const { onUserLogout } = useCurrentUser()
+      const handleUserLogout = wrapWithErrorHandlingAsync(() =>
+        extension.onAuthUserLogout?.()
+      )
       onUserLogout(() => {
-        void extension.onAuthUserLogout?.()
+        void handleUserLogout()
       })
     }
   }

src/stores/firebaseAuthStore.ts

@@ -64,6 +64,11 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
 
   // Token refresh trigger - increments when token is refreshed
   const tokenRefreshTrigger = ref(0)
+  /**
+   * The user ID for which the initial ID token has been observed.
+   * When a token changes for the same user, that is a refresh.
+   */
+  const lastTokenUserId = ref<string | null>(null)
 
   // Providers
   const googleProvider = new GoogleAuthProvider()
@@ -93,6 +98,9 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
   onAuthStateChanged(auth, (user) => {
     currentUser.value = user
     isInitialized.value = true
+    if (user === null) {
+      lastTokenUserId.value = null
+    }
 
     // Reset balance when auth state changes
     balance.value = null
@@ -102,6 +110,11 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => {
   // Listen for token refresh events
   onIdTokenChanged(auth, (user) => {
     if (user && isCloud) {
+      // Skip initial token change
+      if (lastTokenUserId.value !== user.uid) {
+        lastTokenUserId.value = user.uid
+        return
+      }
       tokenRefreshTrigger.value++
     }
   })

tests-ui/tests/store/firebaseAuthStore.tokenRefresh.test.ts

@@ -0,0 +1,110 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { createPinia, setActivePinia } from 'pinia'
+
+vi.mock('@/platform/distribution/types', () => ({
+  isCloud: true
+}))
+
+vi.mock('vuefire', () => ({
+  useFirebaseAuth: vi.fn()
+}))
+
+vi.mock('firebase/auth', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('firebase/auth')>()
+  return {
+    ...actual,
+    onAuthStateChanged: vi.fn(),
+    onIdTokenChanged: vi.fn(),
+    setPersistence: vi.fn().mockResolvedValue(undefined),
+    GoogleAuthProvider: class {
+      addScope = vi.fn()
+      setCustomParameters = vi.fn()
+    },
+    GithubAuthProvider: class {
+      addScope = vi.fn()
+      setCustomParameters = vi.fn()
+    }
+  }
+})
+
+import * as firebaseAuth from 'firebase/auth'
+import * as vuefire from 'vuefire'
+
+type MinimalUser = { uid: string }
+
+/** Create a minimal user-like object with a stable uid */
+function makeUser(uid: string): MinimalUser {
+  return { uid }
+}
+
+describe('firebaseAuthStore token refresh gating', () => {
+  let onAuthStateChangedCallback:
+    | ((user: MinimalUser | null) => void)
+    | undefined
+  let onIdTokenChangedCallback: ((user: MinimalUser | null) => void) | undefined
+  let store: any
+
+  beforeEach(async () => {
+    vi.resetModules()
+    vi.resetAllMocks()
+    setActivePinia(createPinia())
+
+    const authInstance = {}
+    vi.mocked(vuefire.useFirebaseAuth).mockReturnValue(authInstance as any)
+
+    vi.mocked(firebaseAuth.onAuthStateChanged).mockImplementation((...args) => {
+      const callback = args[1] as (user: MinimalUser | null) => void
+      onAuthStateChangedCallback = callback
+      return vi.fn()
+    })
+
+    vi.mocked(firebaseAuth.onIdTokenChanged).mockImplementation((...args) => {
+      const callback = args[1] as (user: MinimalUser | null) => void
+      onIdTokenChangedCallback = callback
+      return vi.fn()
+    })
+
+    const { useFirebaseAuthStore } = await import('@/stores/firebaseAuthStore')
+    store = useFirebaseAuthStore()
+  })
+
+  it('skips initial token for a user and increments on subsequent refresh', () => {
+    const user = makeUser('user-123')
+
+    onIdTokenChangedCallback?.(user)
+    expect(store.tokenRefreshTrigger).toBe(0)
+
+    onIdTokenChangedCallback?.(user)
+    expect(store.tokenRefreshTrigger).toBe(1)
+  })
+
+  it('does not increment when uid changes; increments on next refresh for new user', () => {
+    const userA = makeUser('user-a')
+    const userB = makeUser('user-b')
+
+    onIdTokenChangedCallback?.(userA)
+    expect(store.tokenRefreshTrigger).toBe(0)
+
+    onIdTokenChangedCallback?.(userB)
+    expect(store.tokenRefreshTrigger).toBe(0)
+
+    onIdTokenChangedCallback?.(userB)
+    expect(store.tokenRefreshTrigger).toBe(1)
+  })
+
+  it('resets gating after logout; first token after logout is skipped', () => {
+    const user = makeUser('user-x')
+
+    onIdTokenChangedCallback?.(user)
+    onIdTokenChangedCallback?.(user)
+    expect(store.tokenRefreshTrigger).toBe(1)
+
+    onAuthStateChangedCallback?.(null)
+
+    onIdTokenChangedCallback?.(user)
+    expect(store.tokenRefreshTrigger).toBe(1)
+
+    onIdTokenChangedCallback?.(user)
+    expect(store.tokenRefreshTrigger).toBe(2)
+  })
+})