diff --git a/src/composables/auth/useCurrentUser.ts b/src/composables/auth/useCurrentUser.ts index 37b9e4866..cc5634cae 100644 --- a/src/composables/auth/useCurrentUser.ts +++ b/src/composables/auth/useCurrentUser.ts @@ -1,5 +1,5 @@ import { whenever } from '@vueuse/core' -import { computed } from 'vue' +import { computed, watch } from 'vue' import { useFirebaseAuthActions } from '@/composables/auth/useFirebaseAuthActions' import { t } from '@/i18n' @@ -37,6 +37,15 @@ export const useCurrentUser = () => { const onUserResolved = (callback: (user: AuthUserInfo) => void) => whenever(resolvedUserInfo, callback, { immediate: true }) + const onTokenRefreshed = (callback: () => void) => + whenever(() => authStore.tokenRefreshTrigger, callback) + + const onUserLogout = (callback: () => void) => { + watch(resolvedUserInfo, (user) => { + if (!user) callback() + }) + } + const userDisplayName = computed(() => { if (isApiKeyLogin.value) { return apiKeyStore.currentUser?.name @@ -133,6 +142,8 @@ export const useCurrentUser = () => { handleSignOut, handleSignIn, handleDeleteAccount, - onUserResolved + onUserResolved, + onTokenRefreshed, + onUserLogout } } diff --git a/src/extensions/core/cloudSessionCookie.ts b/src/extensions/core/cloudSessionCookie.ts new file mode 100644 index 000000000..fc9ffcb80 --- /dev/null +++ b/src/extensions/core/cloudSessionCookie.ts @@ -0,0 +1,25 @@ +import { useSessionCookie } from '@/platform/auth/session/useSessionCookie' +import { useExtensionService } from '@/services/extensionService' + +/** + * Cloud-only extension that manages session cookies for authentication. + * Creates session cookie on login, refreshes it when token refreshes, and deletes on logout. + */ +useExtensionService().registerExtension({ + name: 'Comfy.Cloud.SessionCookie', + + onAuthUserResolved: async () => { + const { createSession } = useSessionCookie() + await createSession() + }, + + onAuthTokenRefreshed: async () => { + const { createSession } = useSessionCookie() + await createSession() + }, + + onAuthUserLogout: async () => { + const { deleteSession } = useSessionCookie() + await deleteSession() + } +}) diff --git a/src/extensions/core/index.ts b/src/extensions/core/index.ts index 40ec459b5..a11a61abb 100644 --- a/src/extensions/core/index.ts +++ b/src/extensions/core/index.ts @@ -28,6 +28,7 @@ import './widgetInputs' if (isCloud) { await import('./cloudRemoteConfig') await import('./cloudBadges') + await import('./cloudSessionCookie') if (window.__CONFIG__?.subscription_required) { await import('./cloudSubscription') diff --git a/src/platform/auth/session/useSessionCookie.ts b/src/platform/auth/session/useSessionCookie.ts new file mode 100644 index 000000000..49f6fec46 --- /dev/null +++ b/src/platform/auth/session/useSessionCookie.ts @@ -0,0 +1,65 @@ +import { api } from '@/scripts/api' +import { isCloud } from '@/platform/distribution/types' +import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore' + +/** + * Session cookie management for cloud authentication. + * Creates and deletes session cookies on the ComfyUI server. + */ +export const useSessionCookie = () => { + /** + * Creates or refreshes the session cookie. + * Called after login and on token refresh. + */ + 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') + } + + 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}` + ) + } + } + + /** + * Deletes the session cookie. + * Called on logout. + */ + const deleteSession = async (): Promise => { + if (!isCloud) return + + 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}` + ) + } + } + + return { + createSession, + deleteSession + } +} diff --git a/src/services/extensionService.ts b/src/services/extensionService.ts index 26b67fba7..1e51991d6 100644 --- a/src/services/extensionService.ts +++ b/src/services/extensionService.ts @@ -80,6 +80,20 @@ export const useExtensionService = () => { void extension.onAuthUserResolved?.(user, app) }) } + + if (extension.onAuthTokenRefreshed) { + const { onTokenRefreshed } = useCurrentUser() + onTokenRefreshed(() => { + void extension.onAuthTokenRefreshed?.() + }) + } + + if (extension.onAuthUserLogout) { + const { onUserLogout } = useCurrentUser() + onUserLogout(() => { + void extension.onAuthUserLogout?.() + }) + } } /** diff --git a/src/stores/firebaseAuthStore.ts b/src/stores/firebaseAuthStore.ts index 208ffe523..08e441925 100644 --- a/src/stores/firebaseAuthStore.ts +++ b/src/stores/firebaseAuthStore.ts @@ -8,6 +8,7 @@ import { deleteUser, getAdditionalUserInfo, onAuthStateChanged, + onIdTokenChanged, sendPasswordResetEmail, setPersistence, signInWithEmailAndPassword, @@ -61,6 +62,9 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => { const balance = ref(null) const lastBalanceUpdateTime = ref(null) + // Token refresh trigger - increments when token is refreshed + const tokenRefreshTrigger = ref(0) + // Providers const googleProvider = new GoogleAuthProvider() googleProvider.addScope('email') @@ -95,6 +99,13 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => { lastBalanceUpdateTime.value = null }) + // Listen for token refresh events + onIdTokenChanged(auth, (user) => { + if (user && isCloud) { + tokenRefreshTrigger.value++ + } + }) + const getIdToken = async (): Promise => { if (!currentUser.value) return try { @@ -421,6 +432,7 @@ export const useFirebaseAuthStore = defineStore('firebaseAuth', () => { balance, lastBalanceUpdateTime, isFetchingBalance, + tokenRefreshTrigger, // Getters isAuthenticated, diff --git a/src/types/comfy.ts b/src/types/comfy.ts index 513724ab5..e8652c94f 100644 --- a/src/types/comfy.ts +++ b/src/types/comfy.ts @@ -219,5 +219,17 @@ export interface ComfyExtension { */ onAuthUserResolved?(user: AuthUserInfo, app: ComfyApp): Promise | void + /** + * Fired whenever the auth token is refreshed. + * This is an experimental API and may be changed or removed in the future. + */ + onAuthTokenRefreshed?(): Promise | void + + /** + * Fired when user logs out. + * This is an experimental API and may be changed or removed in the future. + */ + onAuthUserLogout?(): Promise | void + [key: string]: any } diff --git a/tests-ui/tests/store/firebaseAuthStore.test.ts b/tests-ui/tests/store/firebaseAuthStore.test.ts index d37c3857a..3065e5bf0 100644 --- a/tests-ui/tests/store/firebaseAuthStore.test.ts +++ b/tests-ui/tests/store/firebaseAuthStore.test.ts @@ -56,6 +56,7 @@ vi.mock('firebase/auth', async (importOriginal) => { createUserWithEmailAndPassword: vi.fn(), signOut: vi.fn(), onAuthStateChanged: vi.fn(), + onIdTokenChanged: vi.fn(), signInWithPopup: vi.fn(), GoogleAuthProvider: class { addScope = vi.fn()