refactor: extract auth-routing from workspaceApi to auth domain

- Add getAuthHeaderOrThrow() and getFirebaseAuthHeaderOrThrow() to authStore
- Delegate workspaceApi's auth-or-throw helpers to authStore methods
- Remove unused i18n import from workspaceApi
- Update shared mock factory with new methods
- Thrown errors now use AuthStoreError (auth domain)

browser_tests/fixtures/ComfyPage.ts

@@ -121,16 +121,6 @@ type KeysOfType<T, Match> = {
   [K in keyof T]: T[K] extends Match ? K : never
 }[keyof T]
 
-class MediaLightbox {
-  readonly root: Locator
-  readonly closeButton: Locator
-
-  constructor(public readonly page: Page) {
-    this.root = page.getByRole('dialog')
-    this.closeButton = this.root.getByLabel('Close')
-  }
-}
-
 class ConfirmDialog {
   public readonly root: Locator
   public readonly delete: Locator
@@ -192,7 +182,6 @@ export class ComfyPage {
   public readonly templates: ComfyTemplates
   public readonly settingDialog: SettingDialog
   public readonly confirmDialog: ConfirmDialog
-  public readonly mediaLightbox: MediaLightbox
   public readonly vueNodes: VueNodeHelpers
   public readonly appMode: AppModeHelper
   public readonly subgraph: SubgraphHelper
@@ -241,7 +230,6 @@ export class ComfyPage {
     this.templates = new ComfyTemplates(page)
     this.settingDialog = new SettingDialog(page, this)
     this.confirmDialog = new ConfirmDialog(page)
-    this.mediaLightbox = new MediaLightbox(page)
     this.vueNodes = new VueNodeHelpers(page)
     this.appMode = new AppModeHelper(this)
     this.subgraph = new SubgraphHelper(this)

browser_tests/fixtures/components/ComfyNodeSearchBox.ts

@@ -1,14 +1,11 @@
 import type { Locator, Page } from '@playwright/test'
 
 export class ComfyNodeSearchFilterSelectionPanel {
-  readonly root: Locator
-
-  constructor(public readonly page: Page) {
-    this.root = page.getByRole('dialog')
-  }
+  constructor(public readonly page: Page) {}
 
   get header() {
-    return this.root
+    return this.page
+      .getByRole('dialog')
       .locator('div')
       .filter({ hasText: 'Add node filter condition' })
   }

browser_tests/fixtures/helpers/AppModeHelper.ts

@@ -83,14 +83,6 @@ export class AppModeHelper {
     return this.page.locator('[data-testid="linear-widgets"]')
   }
 
-  /** The PrimeVue Popover for the image picker (renders with role="dialog"). */
-  get imagePickerPopover(): Locator {
-    return this.page
-      .getByRole('dialog')
-      .filter({ has: this.page.getByRole('button', { name: 'All' }) })
-      .first()
-  }
-
   /**
    * Get the actions menu trigger for a widget in the app mode widget list.
    * @param widgetName Text shown in the widget label (e.g. "seed").

browser_tests/helpers/templates.ts

@@ -10,12 +10,10 @@ import { TestIds } from '../fixtures/selectors'
 
 export class ComfyTemplates {
   readonly content: Locator
-  readonly dialog: Locator
   readonly allTemplateCards: Locator
 
   constructor(readonly page: Page) {
     this.content = page.getByTestId(TestIds.templates.content)
-    this.dialog = page.getByRole('dialog')
     this.allTemplateCards = page.locator('[data-testid^="template-workflow-"]')
   }
 

browser_tests/tests/appModeDropdownClipping.spec.ts

@@ -145,7 +145,10 @@ test.describe('App mode dropdown clipping', { tag: '@ui' }, () => {
 
     // The unstyled PrimeVue Popover renders with role="dialog".
     // Locate the one containing the image grid (filter buttons like "All", "Inputs").
-    const popover = comfyPage.appMode.imagePickerPopover
+    const popover = comfyPage.page
+      .getByRole('dialog')
+      .filter({ has: comfyPage.page.getByRole('button', { name: 'All' }) })
+      .first()
     await expect(popover).toBeVisible({ timeout: 5000 })
 
     const isInViewport = await popover.evaluate((el) => {

browser_tests/tests/confirmDialogTextWrap.spec.ts

@@ -18,7 +18,7 @@ test.describe('Confirm dialog text wrapping', { tag: ['@mobile'] }, () => {
         .catch(() => {})
     }, longFilename)
 
-    const dialog = comfyPage.confirmDialog.root
+    const dialog = comfyPage.page.getByRole('dialog')
     await expect(dialog).toBeVisible()
 
     const confirmButton = dialog.getByRole('button', { name: 'Confirm' })

browser_tests/tests/resultGallery.spec.ts

@@ -41,7 +41,7 @@ test.describe('MediaLightbox', { tag: ['@slow'] }, () => {
     await assetCard.hover()
     await assetCard.getByLabel('Zoom in').click()
 
-    const gallery = comfyPage.mediaLightbox.root
+    const gallery = comfyPage.page.getByRole('dialog')
     await expect(gallery).toBeVisible()
 
     return { gallery }
@@ -58,13 +58,13 @@ test.describe('MediaLightbox', { tag: ['@slow'] }, () => {
     await runAndOpenGallery(comfyPage)
 
     await comfyPage.page.keyboard.press('Escape')
-    await expect(comfyPage.mediaLightbox.root).not.toBeVisible()
+    await expect(comfyPage.page.getByRole('dialog')).not.toBeVisible()
   })
 
   test('closes gallery when clicking close button', async ({ comfyPage }) => {
     const { gallery } = await runAndOpenGallery(comfyPage)
 
     await gallery.getByLabel('Close').click()
-    await expect(comfyPage.mediaLightbox.root).not.toBeVisible()
+    await expect(comfyPage.page.getByRole('dialog')).not.toBeVisible()
   })
 })

browser_tests/tests/templates.spec.ts

@@ -116,7 +116,7 @@ test.describe('Templates', { tag: ['@slow', '@workflow'] }, () => {
 
     await comfyPage.command.executeCommand('Comfy.BrowseTemplates')
 
-    const dialog = comfyPage.templates.dialog.filter({
+    const dialog = comfyPage.page.getByRole('dialog').filter({
       has: comfyPage.page.getByRole('heading', { name: 'Modèles', exact: true })
     })
     await expect(dialog).toBeVisible()
@@ -220,9 +220,8 @@ test.describe('Templates', { tag: ['@slow', '@workflow'] }, () => {
       await expect(comfyPage.templates.content).toBeVisible()
 
       // Wait for filter bar select components to render
-      const sortBySelect = comfyPage.templates.dialog.getByRole('combobox', {
-        name: /Sort/
-      })
+      const dialog = comfyPage.page.getByRole('dialog')
+      const sortBySelect = dialog.getByRole('combobox', { name: /Sort/ })
       await expect(sortBySelect).toBeVisible()
 
       // Screenshot the filter bar containing MultiSelect and SingleSelect

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,7 @@ 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
-}
-
-async function getFirebaseHeaderOrThrow() {
-  const authHeader = await useAuthStore().getFirebaseAuthHeader()
-  if (!authHeader) {
-    throw new WorkspaceApiError(
-      t('toastMessages.userNotAuthenticated'),
-      401,
-      'NOT_AUTHENTICATED'
-    )
-  }
-  return authHeader
+  return useAuthStore().getAuthHeaderOrThrow()
 }
 
 function handleAxiosError(err: unknown): never {
@@ -500,7 +479,7 @@ export const workspaceApi = {
    * Uses Firebase auth (user identity) since the user isn't yet a workspace member.
    */
   async acceptInvite(token: string): Promise<AcceptInviteResponse> {
-    const headers = await getFirebaseHeaderOrThrow()
+    const headers = await useAuthStore().getFirebaseAuthHeaderOrThrow()
     try {
       const response = await workspaceApiClient.post<AcceptInviteResponse>(
         api.apiURL(`/invites/${token}/accept`),

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__/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,39 @@ 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.toMatchObject({
+        name: 'AuthStoreError',
+        message: '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.toMatchObject({
+        name: 'AuthStoreError',
+        message: '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 {
@@ -538,7 +522,9 @@ export const useAuthStore = defineStore('auth', () => {
     sendPasswordReset,
     updatePassword: _updatePassword,
     getAuthHeader,
+    getAuthHeaderOrThrow,
     getFirebaseAuthHeader,
+    getFirebaseAuthHeaderOrThrow,
     getAuthToken
   }
 })