diff --git a/src/platform/auth/session/useSessionCookie.ts b/src/platform/auth/session/useSessionCookie.ts index 49f6fec46..38712b115 100644 --- a/src/platform/auth/session/useSessionCookie.ts +++ b/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 | null = null +/** + * Tracks the in-flight deleteSession request to dedupe concurrent calls. + */ +let deleteSessionInFlight: Promise | 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 => { 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 (!response.ok) { - const errorData = await response.json().catch(() => ({})) - throw new Error( - `Failed to create session: ${errorData.message || response.statusText}` - ) + if (!authHeader) { + throw new Error('No auth header available for session creation') + } + + 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 => { if (!isCloud) return - const response = await fetch(api.apiURL('/auth/session'), { - method: 'DELETE', - credentials: 'include' - }) + if (deleteSessionInFlight) { + await deleteSessionInFlight + return + } - if (!response.ok) { - const errorData = await response.json().catch(() => ({})) - throw new Error( - `Failed to delete session: ${errorData.message || response.statusText}` - ) + 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}` + ) + } + })() + + try { + await deleteSessionInFlight + } finally { + deleteSessionInFlight = null } } diff --git a/src/services/extensionService.ts b/src/services/extensionService.ts index 4f3f523d3..33099588e 100644 --- a/src/services/extensionService.ts +++ b/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() }) } } diff --git a/src/stores/firebaseAuthStore.ts b/src/stores/firebaseAuthStore.ts index 08e441925..dc379fce9 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.tokenRefresh.test.ts b/tests-ui/tests/store/firebaseAuthStore.tokenRefresh.test.ts new file mode 100644 index 000000000..5527e609c --- /dev/null +++ b/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() + 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) + }) +})