refactor: reconcile workspaceAuthStore with authStore

- Delegate workspace token reads from authStore to workspaceAuthStore
- getAuthHeader(): use workspaceAuthStore.getWorkspaceAuthHeader() instead of sessionStorage
- getAuthToken(): use workspaceAuthStore.getWorkspaceToken() instead of sessionStorage
- onAuthStateChanged logout: use workspaceAuthStore.clearWorkspaceContext()
- Add getWorkspaceToken() to workspaceAuthStore
- Remove WORKSPACE_STORAGE_KEYS import from authStore
- Add authTokenPriority.test.ts with 7 priority chain scenarios

src/platform/workspace/api/workspaceApi.ts

@@ -1,6 +1,5 @@
 import axios from 'axios'
 
-import { t } from '@/i18n'
 import { api } from '@/scripts/api'
 import { useAuthStore } from '@/stores/authStore'
 
@@ -288,27 +287,11 @@ const workspaceApiClient = axios.create({
 })
 
 async function getAuthHeaderOrThrow() {
-  const authHeader = await useAuthStore().getAuthHeader()
-  if (!authHeader) {
-    throw new WorkspaceApiError(
-      t('toastMessages.userNotAuthenticated'),
-      401,
-      'NOT_AUTHENTICATED'
-    )
-  }
-  return authHeader
+  return useAuthStore().getAuthHeaderOrThrow()
 }
 
 async function getFirebaseHeaderOrThrow() {
-  const authHeader = await useAuthStore().getFirebaseAuthHeader()
-  if (!authHeader) {
-    throw new WorkspaceApiError(
-      t('toastMessages.userNotAuthenticated'),
-      401,
-      'NOT_AUTHENTICATED'
-    )
-  }
-  return authHeader
+  return useAuthStore().getFirebaseAuthHeaderOrThrow()
 }
 
 function handleAxiosError(err: unknown): never {

src/platform/workspace/stores/workspaceAuthStore.ts

@@ -343,6 +343,10 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
     }
   }
 
+  function getWorkspaceToken(): string | undefined {
+    return workspaceToken.value ?? undefined
+  }
+
   function clearWorkspaceContext(): void {
     // Increment request ID to invalidate any in-flight stale refresh operations
     refreshRequestId++
@@ -370,6 +374,7 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
     switchWorkspace,
     refreshToken,
     getWorkspaceAuthHeader,
+    getWorkspaceToken,
     clearWorkspaceContext
   }
 })

src/stores/__tests__/authStoreMock.ts

@@ -28,7 +28,9 @@ export interface AuthStoreMockControls {
   logout: Mock
   getIdToken: Mock
   getAuthHeader: Mock
+  getAuthHeaderOrThrow: Mock
   getFirebaseAuthHeader: Mock
+  getFirebaseAuthHeaderOrThrow: Mock
   getAuthToken: Mock
   createCustomer: Mock
   fetchBalance: Mock
@@ -63,7 +65,13 @@ export function createAuthStoreMock(): {
     logout: vi.fn(),
     getIdToken: vi.fn().mockResolvedValue('mock-id-token'),
     getAuthHeader: vi.fn().mockResolvedValue(null),
+    getAuthHeaderOrThrow: vi.fn().mockResolvedValue({
+      Authorization: 'Bearer mock-id-token'
+    }),
     getFirebaseAuthHeader: vi.fn().mockResolvedValue(null),
+    getFirebaseAuthHeaderOrThrow: vi.fn().mockResolvedValue({
+      Authorization: 'Bearer mock-id-token'
+    }),
     getAuthToken: vi.fn().mockResolvedValue(undefined),
     createCustomer: vi.fn(),
     fetchBalance: vi.fn(),

src/stores/__tests__/authTokenPriority.test.ts

@@ -0,0 +1,211 @@
+import type { User } from 'firebase/auth'
+import * as firebaseAuth from 'firebase/auth'
+import { setActivePinia } from 'pinia'
+import type { Mock } from 'vitest'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import * as vuefire from 'vuefire'
+
+import { useAuthStore } from '@/stores/authStore'
+import { createTestingPinia } from '@pinia/testing'
+
+const { mockFeatureFlags } = vi.hoisted(() => ({
+  mockFeatureFlags: {
+    teamWorkspacesEnabled: false
+  }
+}))
+
+const { mockDistributionTypes } = vi.hoisted(() => ({
+  mockDistributionTypes: {
+    isCloud: true,
+    isDesktop: true
+  }
+}))
+
+const mockWorkspaceAuthHeader = vi.fn().mockReturnValue(null)
+const mockGetWorkspaceToken = vi.fn().mockReturnValue(undefined)
+const mockClearWorkspaceContext = vi.fn()
+
+vi.mock('@/platform/workspace/stores/workspaceAuthStore', () => ({
+  useWorkspaceAuthStore: () => ({
+    getWorkspaceAuthHeader: mockWorkspaceAuthHeader,
+    getWorkspaceToken: mockGetWorkspaceToken,
+    clearWorkspaceContext: mockClearWorkspaceContext
+  })
+}))
+
+vi.mock('@/composables/useFeatureFlags', () => ({
+  useFeatureFlags: () => ({
+    flags: mockFeatureFlags
+  })
+}))
+
+vi.mock('vuefire', () => ({
+  useFirebaseAuth: vi.fn()
+}))
+
+vi.mock('vue-i18n', () => ({
+  useI18n: () => ({ t: (key: string) => key }),
+  createI18n: () => ({ global: { t: (key: string) => key } })
+}))
+
+vi.mock('firebase/auth', async (importOriginal) => {
+  const actual = await importOriginal<typeof firebaseAuth>()
+  return {
+    ...actual,
+    signInWithEmailAndPassword: vi.fn(),
+    createUserWithEmailAndPassword: vi.fn(),
+    signOut: vi.fn(),
+    onAuthStateChanged: vi.fn(),
+    onIdTokenChanged: vi.fn(),
+    signInWithPopup: vi.fn(),
+    GoogleAuthProvider: class {
+      addScope = vi.fn()
+      setCustomParameters = vi.fn()
+    },
+    GithubAuthProvider: class {
+      addScope = vi.fn()
+      setCustomParameters = vi.fn()
+    },
+    getAdditionalUserInfo: vi.fn(),
+    setPersistence: vi.fn().mockResolvedValue(undefined)
+  }
+})
+
+vi.mock('@/platform/telemetry', () => ({
+  useTelemetry: () => ({ trackAuth: vi.fn() })
+}))
+
+vi.mock('@/stores/toastStore', () => ({
+  useToastStore: () => ({ add: vi.fn() })
+}))
+
+vi.mock('@/services/dialogService')
+vi.mock('@/platform/distribution/types', () => mockDistributionTypes)
+
+const mockApiKeyGetAuthHeader = vi.fn().mockReturnValue(null)
+vi.mock('@/stores/apiKeyAuthStore', () => ({
+  useApiKeyAuthStore: () => ({
+    getAuthHeader: mockApiKeyGetAuthHeader,
+    getApiKey: vi.fn(),
+    currentUser: null,
+    isAuthenticated: false,
+    storeApiKey: vi.fn(),
+    clearStoredApiKey: vi.fn()
+  })
+}))
+
+type MockUser = Omit<User, 'getIdToken'> & { getIdToken: Mock }
+
+describe('auth token priority chain', () => {
+  let store: ReturnType<typeof useAuthStore>
+  let authStateCallback: (user: User | null) => void
+
+  const mockAuth: Record<string, unknown> = {}
+
+  const mockUser: MockUser = {
+    uid: 'test-user-id',
+    email: 'test@example.com',
+    getIdToken: vi.fn().mockResolvedValue('firebase-token')
+  } as Partial<User> as MockUser
+
+  beforeEach(() => {
+    vi.resetAllMocks()
+
+    mockFeatureFlags.teamWorkspacesEnabled = false
+    mockWorkspaceAuthHeader.mockReturnValue(null)
+    mockGetWorkspaceToken.mockReturnValue(undefined)
+    mockApiKeyGetAuthHeader.mockReturnValue(null)
+    mockUser.getIdToken.mockResolvedValue('firebase-token')
+
+    vi.mocked(vuefire.useFirebaseAuth).mockReturnValue(
+      mockAuth as unknown as ReturnType<typeof vuefire.useFirebaseAuth>
+    )
+
+    vi.mocked(firebaseAuth.onAuthStateChanged).mockImplementation(
+      (_, callback) => {
+        authStateCallback = callback as (user: User | null) => void
+        ;(callback as (user: User | null) => void)(mockUser)
+        return vi.fn()
+      }
+    )
+
+    setActivePinia(createTestingPinia({ stubActions: false }))
+    store = useAuthStore()
+  })
+
+  describe('getAuthHeader priority', () => {
+    it('returns workspace auth header when workspace is active and feature enabled', async () => {
+      mockFeatureFlags.teamWorkspacesEnabled = true
+      mockWorkspaceAuthHeader.mockReturnValue({
+        Authorization: 'Bearer workspace-token'
+      })
+
+      const header = await store.getAuthHeader()
+
+      expect(header).toEqual({
+        Authorization: 'Bearer workspace-token'
+      })
+    })
+
+    it('returns Firebase token when workspace is not active but user is authenticated', async () => {
+      mockFeatureFlags.teamWorkspacesEnabled = true
+      mockWorkspaceAuthHeader.mockReturnValue(null)
+
+      const header = await store.getAuthHeader()
+
+      expect(header).toEqual({
+        Authorization: 'Bearer firebase-token'
+      })
+    })
+
+    it('returns API key when neither workspace nor Firebase are available', async () => {
+      authStateCallback(null)
+      mockApiKeyGetAuthHeader.mockReturnValue({ 'X-API-KEY': 'test-key' })
+
+      const header = await store.getAuthHeader()
+
+      expect(header).toEqual({ 'X-API-KEY': 'test-key' })
+    })
+
+    it('returns null when no auth method is available', async () => {
+      authStateCallback(null)
+
+      const header = await store.getAuthHeader()
+
+      expect(header).toBeNull()
+    })
+
+    it('skips workspace header when team_workspaces feature is disabled', async () => {
+      mockFeatureFlags.teamWorkspacesEnabled = false
+      mockWorkspaceAuthHeader.mockReturnValue({
+        Authorization: 'Bearer workspace-token'
+      })
+
+      const header = await store.getAuthHeader()
+
+      expect(header).toEqual({
+        Authorization: 'Bearer firebase-token'
+      })
+    })
+  })
+
+  describe('getAuthToken priority', () => {
+    it('returns workspace token when workspace is active and feature enabled', async () => {
+      mockFeatureFlags.teamWorkspacesEnabled = true
+      mockGetWorkspaceToken.mockReturnValue('workspace-raw-token')
+
+      const token = await store.getAuthToken()
+
+      expect(token).toBe('workspace-raw-token')
+    })
+
+    it('returns Firebase token when workspace token is not available', async () => {
+      mockFeatureFlags.teamWorkspacesEnabled = true
+      mockGetWorkspaceToken.mockReturnValue(undefined)
+
+      const token = await store.getAuthToken()
+
+      expect(token).toBe('firebase-token')
+    })
+  })
+})

src/stores/authStore.test.ts

@@ -730,6 +730,37 @@ describe('useAuthStore', () => {
     })
   })
 
+  describe('getAuthHeaderOrThrow', () => {
+    it('returns auth header when authenticated', async () => {
+      const header = await store.getAuthHeaderOrThrow()
+      expect(header).toEqual({ Authorization: 'Bearer mock-id-token' })
+    })
+
+    it('throws AuthStoreError when not authenticated', async () => {
+      authStateCallback(null)
+      mockApiKeyGetAuthHeader.mockReturnValue(null)
+
+      await expect(store.getAuthHeaderOrThrow()).rejects.toThrow(
+        'toastMessages.userNotAuthenticated'
+      )
+    })
+  })
+
+  describe('getFirebaseAuthHeaderOrThrow', () => {
+    it('returns Firebase auth header when authenticated', async () => {
+      const header = await store.getFirebaseAuthHeaderOrThrow()
+      expect(header).toEqual({ Authorization: 'Bearer mock-id-token' })
+    })
+
+    it('throws AuthStoreError when not authenticated', async () => {
+      authStateCallback(null)
+
+      await expect(store.getFirebaseAuthHeaderOrThrow()).rejects.toThrow(
+        'toastMessages.userNotAuthenticated'
+      )
+    })
+  })
+
   describe('createCustomer', () => {
     it('should succeed with API key auth when no Firebase user is present', async () => {
       authStateCallback(null)

src/stores/authStore.ts

@@ -22,10 +22,10 @@ import { useFirebaseAuth } from 'vuefire'
 
 import { getComfyApiBaseUrl } from '@/config/comfyApi'
 import { t } from '@/i18n'
-import { WORKSPACE_STORAGE_KEYS } from '@/platform/workspace/workspaceConstants'
 import { isCloud } from '@/platform/distribution/types'
 import { useTelemetry } from '@/platform/telemetry'
 import { useDialogService } from '@/services/dialogService'
+import { useWorkspaceAuthStore } from '@/platform/workspace/stores/workspaceAuthStore'
 import { useApiKeyAuthStore } from '@/stores/apiKeyAuthStore'
 import type { AuthHeader } from '@/types/authTypes'
 import type { operations } from '@/types/comfyRegistryTypes'
@@ -110,15 +110,7 @@ export const useAuthStore = defineStore('auth', () => {
     isInitialized.value = true
     if (user === null) {
       lastTokenUserId.value = null
-
-      // Clear workspace sessionStorage on logout to prevent stale tokens
-      try {
-        sessionStorage.removeItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE)
-        sessionStorage.removeItem(WORKSPACE_STORAGE_KEYS.TOKEN)
-        sessionStorage.removeItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)
-      } catch {
-        // Ignore sessionStorage errors (e.g., in private browsing mode)
-      }
+      useWorkspaceAuthStore().clearWorkspaceContext()
     }
 
     // Reset balance when auth state changes
@@ -175,21 +167,8 @@ export const useAuthStore = defineStore('auth', () => {
    */
   const getAuthHeader = async (): Promise<AuthHeader | null> => {
     if (flags.teamWorkspacesEnabled) {
-      const workspaceToken = sessionStorage.getItem(
-        WORKSPACE_STORAGE_KEYS.TOKEN
-      )
-      const expiresAt = sessionStorage.getItem(
-        WORKSPACE_STORAGE_KEYS.EXPIRES_AT
-      )
-
-      if (workspaceToken && expiresAt) {
-        const expiryTime = parseInt(expiresAt, 10)
-        if (Date.now() < expiryTime) {
-          return {
-            Authorization: `Bearer ${workspaceToken}`
-          }
-        }
-      }
+      const wsHeader = useWorkspaceAuthStore().getWorkspaceAuthHeader()
+      if (wsHeader) return wsHeader
     }
 
     const token = await getIdToken()
@@ -218,24 +197,29 @@ export const useAuthStore = defineStore('auth', () => {
    */
   const getAuthToken = async (): Promise<string | undefined> => {
     if (flags.teamWorkspacesEnabled) {
-      const workspaceToken = sessionStorage.getItem(
-        WORKSPACE_STORAGE_KEYS.TOKEN
-      )
-      const expiresAt = sessionStorage.getItem(
-        WORKSPACE_STORAGE_KEYS.EXPIRES_AT
-      )
-
-      if (workspaceToken && expiresAt) {
-        const expiryTime = parseInt(expiresAt, 10)
-        if (Date.now() < expiryTime) {
-          return workspaceToken
-        }
-      }
+      const wsToken = useWorkspaceAuthStore().getWorkspaceToken()
+      if (wsToken) return wsToken
     }
 
     return await getIdToken()
   }
 
+  const getAuthHeaderOrThrow = async (): Promise<AuthHeader> => {
+    const authHeader = await getAuthHeader()
+    if (!authHeader) {
+      throw new AuthStoreError(t('toastMessages.userNotAuthenticated'))
+    }
+    return authHeader
+  }
+
+  const getFirebaseAuthHeaderOrThrow = async (): Promise<AuthHeader> => {
+    const authHeader = await getFirebaseAuthHeader()
+    if (!authHeader) {
+      throw new AuthStoreError(t('toastMessages.userNotAuthenticated'))
+    }
+    return authHeader
+  }
+
   const fetchBalance = async (): Promise<GetCustomerBalanceResponse | null> => {
     isFetchingBalance.value = true
     try {
@@ -534,7 +518,9 @@ export const useAuthStore = defineStore('auth', () => {
     sendPasswordReset,
     updatePassword: _updatePassword,
     getAuthHeader,
+    getAuthHeaderOrThrow,
     getFirebaseAuthHeader,
+    getFirebaseAuthHeaderOrThrow,
     getAuthToken
   }
 })