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

src/services/extensionService.ts

@@ -11,12 +11,17 @@ 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,
+    toastErrorHandler
+  } = useErrorHandling()
 
   /**
    * Loads all extensions from the API into the window in parallel
@@ -77,22 +82,55 @@ export const useExtensionService = () => {
 
     if (extension.onAuthUserResolved) {
       const { onUserResolved } = useCurrentUser()
+      const handleUserResolved = wrapWithErrorHandlingAsync(
+        (user: AuthUserInfo) => extension.onAuthUserResolved?.(user, app),
+        (error) => {
+          console.error('[Extension Auth Hook Error]', {
+            extension: extension.name,
+            hook: 'onAuthUserResolved',
+            error
+          })
+          toastErrorHandler(error)
+        }
+      )
       onUserResolved((user) => {
-        void extension.onAuthUserResolved?.(user, app)
+        void handleUserResolved(user)
       })
     }
 
     if (extension.onAuthTokenRefreshed) {
       const { onTokenRefreshed } = useCurrentUser()
+      const handleTokenRefreshed = wrapWithErrorHandlingAsync(
+        () => extension.onAuthTokenRefreshed?.(),
+        (error) => {
+          console.error('[Extension Auth Hook Error]', {
+            extension: extension.name,
+            hook: 'onAuthTokenRefreshed',
+            error
+          })
+          toastErrorHandler(error)
+        }
+      )
       onTokenRefreshed(() => {
-        void extension.onAuthTokenRefreshed?.()
+        void handleTokenRefreshed()
       })
     }
 
     if (extension.onAuthUserLogout) {
       const { onUserLogout } = useCurrentUser()
+      const handleUserLogout = wrapWithErrorHandlingAsync(
+        () => extension.onAuthUserLogout?.(),
+        (error) => {
+          console.error('[Extension Auth Hook Error]', {
+            extension: extension.name,
+            hook: 'onAuthUserLogout',
+            error
+          })
+          toastErrorHandler(error)
+        }
+      )
       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.test.ts

@@ -83,6 +83,7 @@ vi.mock('@/services/dialogService')
 describe('useFirebaseAuthStore', () => {
   let store: ReturnType<typeof useFirebaseAuthStore>
   let authStateCallback: (user: any) => void
+  let idTokenCallback: (user: any) => void
 
   const mockAuth = {
     /* mock Auth object */
@@ -143,6 +144,55 @@ describe('useFirebaseAuthStore', () => {
     mockUser.getIdToken.mockResolvedValue('mock-id-token')
   })
 
+  describe('token refresh events', () => {
+    beforeEach(async () => {
+      vi.resetModules()
+      vi.doMock('@/platform/distribution/types', () => ({
+        isCloud: true,
+        isDesktop: true
+      }))
+
+      vi.mocked(firebaseAuth.onIdTokenChanged).mockImplementation(
+        (_auth, callback) => {
+          idTokenCallback = callback as (user: any) => void
+          return vi.fn()
+        }
+      )
+
+      vi.mocked(vuefire.useFirebaseAuth).mockReturnValue(mockAuth as any)
+
+      setActivePinia(createPinia())
+      const storeModule = await import('@/stores/firebaseAuthStore')
+      store = storeModule.useFirebaseAuthStore()
+    })
+
+    it("should not increment tokenRefreshTrigger on the user's first ID token event", () => {
+      idTokenCallback?.(mockUser)
+      expect(store.tokenRefreshTrigger).toBe(0)
+    })
+
+    it('should increment tokenRefreshTrigger on subsequent ID token events for the same user', () => {
+      idTokenCallback?.(mockUser)
+      idTokenCallback?.(mockUser)
+      expect(store.tokenRefreshTrigger).toBe(1)
+    })
+
+    it('should not increment when ID token event is for a different user UID', () => {
+      const otherUser = { uid: 'other-user-id' }
+      idTokenCallback?.(mockUser)
+      idTokenCallback?.(otherUser)
+      expect(store.tokenRefreshTrigger).toBe(0)
+    })
+
+    it('should increment after switching to a new UID and receiving a second event for that UID', () => {
+      const otherUser = { uid: 'other-user-id' }
+      idTokenCallback?.(mockUser)
+      idTokenCallback?.(otherUser)
+      idTokenCallback?.(otherUser)
+      expect(store.tokenRefreshTrigger).toBe(1)
+    })
+  })
+
   it('should initialize with the current user', () => {
     expect(store.currentUser).toEqual(mockUser)
     expect(store.isAuthenticated).toBe(true)