fix: address workspace switcher review feedback

Amp-Thread-ID: https://ampcode.com/threads/T-01a016ea-c8f0-77da-842d-58fe769bdd11
Co-authored-by: Amp <amp@ampcode.com>

Author: dante01yoon

src/components/topbar/CurrentUserPopoverLegacy.test.ts

@@ -448,14 +448,35 @@ describe('CurrentUserPopoverLegacy', () => {
 
       const trigger = screen.getByTestId('workspace-switcher-trigger')
       expect(trigger).toHaveAttribute('aria-expanded', 'false')
+      expect(trigger).toHaveAttribute('aria-haspopup', 'menu')
+      expect(trigger).toHaveAttribute(
+        'aria-controls',
+        'workspace-switcher-panel'
+      )
       expect(screen.queryByTestId('workspace-switcher-panel')).toBeNull()
 
       await user.click(trigger)
 
-      expect(screen.getByTestId('workspace-switcher-panel')).toBeInTheDocument()
+      const panel = screen.getByTestId('workspace-switcher-panel')
+      expect(panel).toBeInTheDocument()
+      expect(panel).toHaveAttribute('id', 'workspace-switcher-panel')
+      expect(panel).toHaveAttribute('role', 'menu')
       expect(trigger).toHaveAttribute('aria-expanded', 'true')
     })
 
+    it('closes the switcher on Escape or a click elsewhere', async () => {
+      const { user } = renderComponent(readyWorkspaceState)
+      const trigger = screen.getByTestId('workspace-switcher-trigger')
+
+      await user.click(trigger)
+      await user.keyboard('{Escape}')
+      expect(screen.queryByTestId('workspace-switcher-panel')).toBeNull()
+
+      await user.click(trigger)
+      await user.click(screen.getByText('Test User'))
+      expect(screen.queryByTestId('workspace-switcher-panel')).toBeNull()
+    })
+
     it('keeps credits visible but hides top-up for workspace members', () => {
       mockCanAccessSubscriptionFeatures.value = false
       renderComponent({

src/components/topbar/CurrentUserPopoverLegacy.vue

@@ -32,12 +32,16 @@
 
     <div v-if="showWorkspaceSwitcher" class="relative">
       <Button
+        ref="workspaceSwitcherTrigger"
         v-tooltip="{ value: workspaceName, showDelay: 300 }"
         variant="muted-textonly"
         class="flex h-auto w-full items-center justify-between rounded-lg px-4 py-2 hover:bg-secondary-background-hover"
         :aria-expanded="isWorkspaceSwitcherOpen"
+        aria-haspopup="menu"
+        aria-controls="workspace-switcher-panel"
         data-testid="workspace-switcher-trigger"
         @click="isWorkspaceSwitcherOpen = !isWorkspaceSwitcherOpen"
+        @keydown.escape.stop="isWorkspaceSwitcherOpen = false"
       >
         <div class="flex w-0 flex-1 items-center gap-2">
           <WorkspaceProfilePic
@@ -53,6 +57,9 @@
 
       <div
         v-if="isWorkspaceSwitcherOpen"
+        id="workspace-switcher-panel"
+        ref="workspaceSwitcherPanel"
+        role="menu"
         class="absolute top-0 right-full z-10 mr-4 rounded-lg border border-border-default bg-base-background shadow-[1px_1px_8px_0_rgba(0,0,0,0.4)]"
         data-testid="workspace-switcher-panel"
       >
@@ -145,10 +152,11 @@
 </template>
 
 <script setup lang="ts">
+import { onClickOutside } from '@vueuse/core'
 import { storeToRefs } from 'pinia'
 import Divider from 'primevue/divider'
 import Skeleton from 'primevue/skeleton'
-import { computed, onMounted, ref } from 'vue'
+import { computed, onMounted, ref, useTemplateRef } from 'vue'
 import { useI18n } from 'vue-i18n'
 
 import { formatCreditsFromCents } from '@/base/credits/comfyCredits'
@@ -191,6 +199,17 @@ const { initState, workspaces, workspaceName } = storeToRefs(
   useTeamWorkspaceStore()
 )
 const isWorkspaceSwitcherOpen = ref(false)
+const workspaceSwitcherTrigger = useTemplateRef('workspaceSwitcherTrigger')
+const workspaceSwitcherPanel = useTemplateRef('workspaceSwitcherPanel')
+
+onClickOutside(
+  workspaceSwitcherPanel,
+  () => {
+    isWorkspaceSwitcherOpen.value = false
+  },
+  { ignore: [workspaceSwitcherTrigger] }
+)
+
 const showWorkspaceSwitcher = computed(
   () => initState.value === 'ready' && workspaces.value.length > 0
 )

src/platform/workspace/components/CurrentUserPopoverWorkspace.test.ts

@@ -174,15 +174,35 @@ describe('CurrentUserPopoverWorkspace', () => {
   it('toggles the workspace switcher panel from the selector row', async () => {
     const user = userEvent.setup()
     renderComponent()
+    const trigger = screen.getByTestId('workspace-switcher-trigger')
 
     expect(
       screen.queryByTestId('workspace-switcher-panel')
     ).not.toBeInTheDocument()
+    expect(trigger).toHaveAttribute('aria-expanded', 'false')
+    expect(trigger).toHaveAttribute('aria-haspopup', 'menu')
+    expect(trigger).toHaveAttribute('aria-controls', 'workspace-switcher-panel')
 
-    await user.click(screen.getByTestId('workspace-switcher-trigger'))
-    expect(screen.getByTestId('workspace-switcher-panel')).toBeInTheDocument()
+    await user.click(trigger)
+    const panel = screen.getByTestId('workspace-switcher-panel')
+    expect(panel).toHaveAttribute('id', 'workspace-switcher-panel')
+    expect(panel).toHaveAttribute('role', 'menu')
+    expect(trigger).toHaveAttribute('aria-expanded', 'true')
+
+    await user.click(trigger)
+    expect(
+      screen.queryByTestId('workspace-switcher-panel')
+    ).not.toBeInTheDocument()
+  })
+
+  it('closes the workspace switcher panel on Escape', async () => {
+    const user = userEvent.setup()
+    renderComponent()
+    const trigger = screen.getByTestId('workspace-switcher-trigger')
+
+    await user.click(trigger)
+    await user.keyboard('{Escape}')
 
-    await user.click(screen.getByTestId('workspace-switcher-trigger'))
     expect(
       screen.queryByTestId('workspace-switcher-panel')
     ).not.toBeInTheDocument()

src/platform/workspace/components/CurrentUserPopoverWorkspace.vue

@@ -26,12 +26,17 @@
 
     <!-- Workspace Selector -->
     <div v-if="!accountActionsOnly" class="relative">
-      <div
+      <button
         ref="workspaceSwitcherTrigger"
         v-tooltip="{ value: workspaceName, showDelay: 300 }"
-        class="flex cursor-pointer items-center justify-between rounded-lg px-4 py-2 hover:bg-secondary-background-hover"
+        type="button"
+        class="flex w-full cursor-pointer items-center justify-between rounded-lg px-4 py-2 hover:bg-secondary-background-hover"
+        :aria-expanded="isWorkspaceSwitcherOpen"
+        aria-haspopup="menu"
+        aria-controls="workspace-switcher-panel"
         data-testid="workspace-switcher-trigger"
         @click="toggleWorkspaceSwitcher"
+        @keydown.escape.stop="isWorkspaceSwitcherOpen = false"
       >
         <div class="flex w-0 flex-1 items-center gap-2">
           <WorkspaceProfilePic
@@ -43,11 +48,13 @@
           </span>
         </div>
         <i class="pi pi-chevron-down shrink-0 text-sm text-muted-foreground" />
-      </div>
+      </button>
 
       <div
         v-if="isWorkspaceSwitcherOpen"
+        id="workspace-switcher-panel"
         ref="workspaceSwitcherPanel"
+        role="menu"
         class="absolute top-0 right-full z-10 mr-4 rounded-lg border border-border-default bg-base-background shadow-[1px_1px_8px_0_rgba(0,0,0,0.4)]"
         data-testid="workspace-switcher-panel"
       >

src/platform/workspace/stores/teamWorkspaceStore.test.ts

@@ -370,6 +370,19 @@ describe('useTeamWorkspaceStore', () => {
       expect(store.activeWorkspaceId).toBe(mockMemberWorkspace.id)
     })
 
+    it('clears cached billing rails when identity state is reset', async () => {
+      const store = useTeamWorkspaceStore()
+      await store.initialize()
+      store.setWorkspaceBillingRail(mockPersonalWorkspace.id, 'legacy_stripe')
+      expect(store.activeWorkspaceBillingRail).toBe('legacy_stripe')
+
+      store.resetForIdentityChange()
+      await store.initialize()
+
+      expect(store.activeWorkspaceId).toBe(mockPersonalWorkspace.id)
+      expect(store.activeWorkspaceBillingRail).toBeNull()
+    })
+
     it('does not let a previous user initialization overwrite the next user', async () => {
       let resolveFirstList: (value: unknown) => void = () => {}
       mockWorkspaceApi.list
@@ -475,6 +488,8 @@ describe('useTeamWorkspaceStore', () => {
       )
       const store = useTeamWorkspaceStore()
       await store.initialize()
+      store.setWorkspaceBillingRail(mockPersonalWorkspace.id, 'legacy_stripe')
+      expect(store.activeWorkspaceBillingRail).toBe('legacy_stripe')
 
       await store.switchWorkspace(mockTeamWorkspace.id)
 
@@ -488,6 +503,10 @@ describe('useTeamWorkspaceStore', () => {
       ).not.toHaveBeenCalled()
       expect(mockReload).not.toHaveBeenCalled()
       expect(store.isSwitching).toBe(false)
+      expect(store.activeWorkspaceBillingRail).toBeNull()
+
+      store.setWorkspaceBillingRail(mockTeamWorkspace.id, 'metronome')
+      expect(store.activeWorkspaceBillingRail).toBe('metronome')
     })
 
     it('rejects an overlapping local switch', async () => {

src/platform/workspace/stores/teamWorkspaceStore.ts

@@ -905,6 +905,7 @@ export const useTeamWorkspaceStore = defineStore('teamWorkspace', () => {
     initState.value = 'uninitialized'
     workspaces.value = []
     mutableActiveWorkspaceId.value = null
+    billingRailByWorkspaceId.value = {}
     error.value = null
     isCreating.value = false
     isDeleting.value = false

src/stores/__tests__/authTokenPriority.test.ts

@@ -323,6 +323,22 @@ describe('auth token priority chain', () => {
       expect(mockUser.getIdToken).not.toHaveBeenCalled()
       expect(mockEnsureWorkspaceToken).not.toHaveBeenCalled()
     })
+
+    it('retries queue authentication after a previous initialization error', async () => {
+      mockDistributionTypes.isCloud = false
+      mockTeamWorkspaceInitState = 'error'
+      mockInitializeWorkspaces.mockImplementation(async () => {
+        mockActiveWorkspaceId = 'workspace-123'
+        mockTeamWorkspaceInitState = 'ready'
+      })
+      mockEnsureWorkspaceToken.mockResolvedValue('workspace-raw-token')
+
+      await expect(store.getWorkspaceAuthToken()).resolves.toBe(
+        'workspace-raw-token'
+      )
+      expect(mockInitializeWorkspaces).toHaveBeenCalledOnce()
+      expect(mockEnsureWorkspaceToken).toHaveBeenCalledWith('workspace-123')
+    })
   })
 
   describe('unified login mint wiring', () => {

src/stores/authStore.ts

@@ -337,7 +337,8 @@ export const useAuthStore = defineStore('auth', () => {
       currentUser.value &&
       !teamWorkspaceStore.activeWorkspaceId &&
       (teamWorkspaceStore.initState === 'uninitialized' ||
-        teamWorkspaceStore.initState === 'loading')
+        teamWorkspaceStore.initState === 'loading' ||
+        teamWorkspaceStore.initState === 'error')
     ) {
       try {
         await teamWorkspaceStore.initialize()