diff --git a/src/services/extensionService.ts b/src/services/extensionService.ts index 4f3f523d3c..75b159e11e 100644 --- a/src/services/extensionService.ts +++ b/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() }) } } diff --git a/src/stores/firebaseAuthStore.ts b/src/stores/firebaseAuthStore.ts index 08e4419258..dc379fce95 100644 --- a/src/stores/firebaseAuthStore.ts +++ b/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(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++ } }) diff --git a/tests-ui/tests/store/firebaseAuthStore.test.ts b/tests-ui/tests/store/firebaseAuthStore.test.ts index 3065e5bf0b..fa67700bce 100644 --- a/tests-ui/tests/store/firebaseAuthStore.test.ts +++ b/tests-ui/tests/store/firebaseAuthStore.test.ts @@ -83,6 +83,7 @@ vi.mock('@/services/dialogService') describe('useFirebaseAuthStore', () => { let store: ReturnType 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)