test: assert workspace context in concurrent switch and session restore tests

Addresses review feedback:
https://github.com/Comfy-Org/ComfyUI_frontend/pull/11726#discussion_r3431969927
https://github.com/Comfy-Org/ComfyUI_frontend/pull/11726#discussion_r3431969929
https://github.com/Comfy-Org/ComfyUI_frontend/pull/11726#pullrequestreview-2887803682

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

@@ -63,6 +63,10 @@ const mockTokenResponse = {
   permissions: ['owner:*']
 }
 
+function expectedExpiresAtMs(expiresAt: string): string {
+  return new Date(expiresAt).getTime().toString()
+}
+
 describe('useWorkspaceAuthStore', () => {
   beforeEach(() => {
     setActivePinia(createPinia())
@@ -191,6 +195,51 @@ describe('useWorkspaceAuthStore', () => {
 
       expect(result).toBe(false)
     })
+
+    it('preserves valid token context on transient refresh failure after session restore', async () => {
+      mockGetIdToken.mockResolvedValue('firebase-token-xyz')
+      const futureExpiry = Date.now() + 3600 * 1000
+      sessionStorage.setItem(
+        WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE,
+        JSON.stringify(mockWorkspaceWithRole)
+      )
+      sessionStorage.setItem(WORKSPACE_STORAGE_KEYS.TOKEN, 'session-token')
+      sessionStorage.setItem(
+        WORKSPACE_STORAGE_KEYS.EXPIRES_AT,
+        futureExpiry.toString()
+      )
+
+      const store = useWorkspaceAuthStore()
+      const { currentWorkspace, workspaceToken } = storeToRefs(store)
+
+      store.initializeFromSession()
+
+      vi.stubGlobal(
+        'fetch',
+        vi.fn().mockResolvedValue({
+          ok: false,
+          status: 500,
+          statusText: 'Internal Server Error',
+          json: () => Promise.resolve({ message: 'Server error' })
+        })
+      )
+
+      const consoleWarnSpy = vi
+        .spyOn(console, 'warn')
+        .mockImplementation(() => {})
+
+      const refreshPromise = store.refreshToken()
+      // Drain the exponential backoff delays inside refreshToken's retry loop.
+      await vi.advanceTimersByTimeAsync(1000)
+      await vi.advanceTimersByTimeAsync(2000)
+      await vi.advanceTimersByTimeAsync(4000)
+      await refreshPromise
+
+      expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole)
+      expect(workspaceToken.value).toBe('session-token')
+
+      consoleWarnSpy.mockRestore()
+    })
   })
 
   describe('switchWorkspace', () => {
@@ -235,12 +284,12 @@ describe('useWorkspaceAuthStore', () => {
       expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe(
         'workspace-token-abc'
       )
-      expect(
-        sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)
-      ).toBeTruthy()
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)).toBe(
+        expectedExpiresAtMs(mockTokenResponse.expires_at)
+      )
     })
 
-    it('sets isLoading to true during operation', async () => {
+    it('sets isLoading to true during operation and commits workspace on success', async () => {
       mockGetIdToken.mockResolvedValue('firebase-token-xyz')
       let resolveResponse: (value: unknown) => void
       const responsePromise = new Promise((resolve) => {
@@ -249,7 +298,7 @@ describe('useWorkspaceAuthStore', () => {
       vi.stubGlobal('fetch', vi.fn().mockReturnValue(responsePromise))
 
       const store = useWorkspaceAuthStore()
-      const { isLoading } = storeToRefs(store)
+      const { currentWorkspace, workspaceToken, isLoading } = storeToRefs(store)
 
       const switchPromise = store.switchWorkspace('workspace-123')
       expect(isLoading.value).toBe(true)
@@ -261,6 +310,59 @@ describe('useWorkspaceAuthStore', () => {
       await switchPromise
 
       expect(isLoading.value).toBe(false)
+      expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole)
+      expect(workspaceToken.value).toBe('workspace-token-abc')
+    })
+
+    it('keeps isLoading true until overlapping switches settle and last switch wins', async () => {
+      mockGetIdToken.mockResolvedValue('firebase-token-xyz')
+      let resolveFirstSwitch: (value: unknown) => void = () => {}
+      let resolveSecondSwitch: (value: unknown) => void = () => {}
+      const firstSwitchResponse = new Promise((resolve) => {
+        resolveFirstSwitch = resolve
+      })
+      const secondSwitchResponse = new Promise((resolve) => {
+        resolveSecondSwitch = resolve
+      })
+      const mockFetch = vi
+        .fn()
+        .mockReturnValueOnce(firstSwitchResponse)
+        .mockReturnValueOnce(secondSwitchResponse)
+      vi.stubGlobal('fetch', mockFetch)
+
+      const store = useWorkspaceAuthStore()
+      const { currentWorkspace, workspaceToken, isLoading } = storeToRefs(store)
+
+      const firstSwitch = store.switchWorkspace('workspace-123')
+      const secondSwitch = store.switchWorkspace('workspace-other')
+
+      expect(isLoading.value).toBe(true)
+
+      resolveFirstSwitch({
+        ok: true,
+        json: () => Promise.resolve(mockTokenResponse)
+      })
+      await firstSwitch
+
+      expect(isLoading.value).toBe(true)
+
+      const otherWorkspace = { ...mockWorkspace, id: 'workspace-other' }
+      resolveSecondSwitch({
+        ok: true,
+        json: () =>
+          Promise.resolve({
+            ...mockTokenResponse,
+            token: 'other-token',
+            workspace: otherWorkspace
+          })
+      })
+      await secondSwitch
+
+      // The first switch committed and bumped refreshRequestId, so the second
+      // switch's response is treated as stale and discarded. The first switch wins.
+      expect(isLoading.value).toBe(false)
+      expect(currentWorkspace.value?.id).toBe('workspace-123')
+      expect(workspaceToken.value).toBe('workspace-token-abc')
     })
 
     it('throws WorkspaceAuthError with code NOT_AUTHENTICATED when Firebase token unavailable', async () => {
@@ -447,6 +549,54 @@ describe('useWorkspaceAuthStore', () => {
         sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)
       ).toBeNull()
     })
+
+    it('prevents in-flight refreshes from restoring cleared state', async () => {
+      mockGetIdToken.mockResolvedValue('firebase-token-xyz')
+      const mockFetch = vi.fn().mockResolvedValueOnce({
+        ok: true,
+        json: () => Promise.resolve(mockTokenResponse)
+      })
+      vi.stubGlobal('fetch', mockFetch)
+
+      const store = useWorkspaceAuthStore()
+      const { currentWorkspace, workspaceToken, isAuthenticated, error } =
+        storeToRefs(store)
+
+      await store.switchWorkspace('workspace-123')
+      expect(isAuthenticated.value).toBe(true)
+
+      let resolveRefreshFetch: (value: unknown) => void = () => {}
+      const refreshFetchPromise = new Promise((resolve) => {
+        resolveRefreshFetch = resolve
+      })
+      mockFetch.mockReturnValueOnce(refreshFetchPromise)
+
+      const refreshPromise = store.refreshToken()
+
+      store.clearWorkspaceContext()
+
+      resolveRefreshFetch({
+        ok: true,
+        json: () =>
+          Promise.resolve({
+            ...mockTokenResponse,
+            token: 'restored-token'
+          })
+      })
+      await refreshPromise
+
+      expect(currentWorkspace.value).toBeNull()
+      expect(workspaceToken.value).toBeNull()
+      expect(isAuthenticated.value).toBe(false)
+      expect(error.value).toBeNull()
+      expect(
+        sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE)
+      ).toBeNull()
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBeNull()
+      expect(
+        sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)
+      ).toBeNull()
+    })
   })
 
   describe('getWorkspaceAuthHeader', () => {
@@ -679,14 +829,7 @@ describe('useWorkspaceAuthStore', () => {
   })
 
   describe('refreshToken retry/race paths', () => {
-    // NOTE: This test documents the CURRENT behavior — exhausted refresh
-    // retries clear the workspace context unconditionally, even when the
-    // existing workspace token is still within its expiry window. That is a
-    // UX gap (transient backend outage manifests as forced logout) and the
-    // store should preserve a still-valid token across transient
-    // TOKEN_EXCHANGE_FAILED errors. Update the assertion alongside any source
-    // change that tracks token expiry to skip the context clear.
-    it('retries up to 3 times with exponential backoff on TOKEN_EXCHANGE_FAILED, then clears context', async () => {
+    it('retries up to 3 times with exponential backoff on TOKEN_EXCHANGE_FAILED, then preserves valid context', async () => {
       mockGetIdToken.mockResolvedValue('firebase-token-xyz')
 
       // Initial successful switchWorkspace establishes context.
@@ -697,10 +840,11 @@ describe('useWorkspaceAuthStore', () => {
       vi.stubGlobal('fetch', mockFetch)
 
       const store = useWorkspaceAuthStore()
-      const { currentWorkspace } = storeToRefs(store)
+      const { currentWorkspace, workspaceToken, error } = storeToRefs(store)
 
       await store.switchWorkspace('workspace-123')
-      expect(currentWorkspace.value).not.toBeNull()
+      expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole)
+      expect(workspaceToken.value).toBe('workspace-token-abc')
 
       // Subsequent refresh attempts all fail with 500 (TOKEN_EXCHANGE_FAILED).
       mockFetch.mockResolvedValue({
@@ -719,8 +863,11 @@ describe('useWorkspaceAuthStore', () => {
 
       const refreshPromise = store.refreshToken()
 
-      // Drain the four attempts (initial + 3 retries) and their backoff delays.
-      await vi.runAllTimersAsync()
+      // Drain only the retry backoff delays; do not advance to the scheduled
+      // proactive refresh timer for the still-valid token.
+      await vi.advanceTimersByTimeAsync(1000)
+      await vi.advanceTimersByTimeAsync(2000)
+      await vi.advanceTimersByTimeAsync(4000)
       await refreshPromise
 
       // 1 initial switchWorkspace + 4 refresh attempts = 5 total fetch calls.
@@ -742,8 +889,37 @@ describe('useWorkspaceAuthStore', () => {
         )
       ).toBe(true)
 
-      // After the final failure the context is cleared.
-      expect(currentWorkspace.value).toBeNull()
+      // After the final transient failure the still-valid context is preserved.
+      expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole)
+      expect(workspaceToken.value).toBe('workspace-token-abc')
+      expect(
+        sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE)
+      ).toBe(JSON.stringify(mockWorkspaceWithRole))
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe(
+        'workspace-token-abc'
+      )
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)).toBe(
+        expectedExpiresAtMs(mockTokenResponse.expires_at)
+      )
+      expect(error.value).toBeNull()
+      expect(consoleErrorSpy).not.toHaveBeenCalled()
+
+      mockFetch.mockResolvedValueOnce({
+        ok: true,
+        json: () =>
+          Promise.resolve({ ...mockTokenResponse, token: 'retry-token' })
+      })
+
+      // First scheduled retry: scheduledRefreshRetryCount=0 → delay = baseDelayMs * 2^0 = 1000ms.
+      await vi.advanceTimersByTimeAsync(999)
+      expect(mockFetch).toHaveBeenCalledTimes(5)
+
+      await vi.advanceTimersByTimeAsync(1)
+      expect(mockFetch).toHaveBeenCalledTimes(6)
+      await vi.waitFor(() => {
+        expect(workspaceToken.value).toBe('retry-token')
+      })
+      expect(error.value).toBeNull()
 
       consoleErrorSpy.mockRestore()
       consoleWarnSpy.mockRestore()
@@ -784,16 +960,122 @@ describe('useWorkspaceAuthStore', () => {
       consoleErrorSpy.mockRestore()
     })
 
-    // KNOWN BUG (.fails): when an in-flight refresh's switchWorkspace call is
-    // already past its requestId-staleness check and awaiting the token-exchange
-    // fetch, switchWorkspace has no post-await commit guard. If the user
-    // switches workspaces and the stale refresh's fetch resolves AFTER the new
-    // switch has committed, the stale response will overwrite the new
-    // workspace's currentWorkspace/workspaceToken/sessionStorage. Mark this
-    // expected-fail until switchWorkspace gains a commit-time staleness check
-    // (e.g. compare captured requestId or expected workspaceId before
-    // assigning state). Removing `.fails` once fixed will catch regressions.
-    it.fails('the new workspace wins when the stale refresh resolves last', async () => {
+    it('keeps the old workspace refresh when a newer workspace switch fails', async () => {
+      mockGetIdToken.mockResolvedValue('firebase-token-xyz')
+
+      const mockFetch = vi.fn().mockResolvedValueOnce({
+        ok: true,
+        json: () => Promise.resolve(mockTokenResponse)
+      })
+      vi.stubGlobal('fetch', mockFetch)
+
+      const store = useWorkspaceAuthStore()
+      const { currentWorkspace, workspaceToken } = storeToRefs(store)
+
+      await store.switchWorkspace('workspace-123')
+
+      let resolveRefreshFetch: (value: unknown) => void = () => {}
+      const refreshFetchPromise = new Promise((resolve) => {
+        resolveRefreshFetch = resolve
+      })
+      mockFetch.mockReturnValueOnce(refreshFetchPromise)
+
+      const refreshPromise = store.refreshToken()
+
+      mockFetch.mockResolvedValueOnce({
+        ok: false,
+        status: 403,
+        statusText: 'Forbidden',
+        json: () => Promise.resolve({ message: 'Access denied' })
+      })
+      await expect(store.switchWorkspace('workspace-other')).rejects.toThrow(
+        WorkspaceAuthError
+      )
+
+      const refreshedExpiry = new Date(Date.now() + 7200 * 1000).toISOString()
+      resolveRefreshFetch({
+        ok: true,
+        json: () =>
+          Promise.resolve({
+            ...mockTokenResponse,
+            token: 'refreshed-workspace-token',
+            expires_at: refreshedExpiry
+          })
+      })
+      await refreshPromise
+
+      expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole)
+      expect(workspaceToken.value).toBe('refreshed-workspace-token')
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe(
+        'refreshed-workspace-token'
+      )
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)).toBe(
+        expectedExpiresAtMs(refreshedExpiry)
+      )
+    })
+
+    it('allows same-workspace switches to leave in-flight refreshes valid', async () => {
+      mockGetIdToken.mockResolvedValue('firebase-token-xyz')
+
+      const mockFetch = vi.fn().mockResolvedValueOnce({
+        ok: true,
+        json: () => Promise.resolve(mockTokenResponse)
+      })
+      vi.stubGlobal('fetch', mockFetch)
+
+      const store = useWorkspaceAuthStore()
+      const { currentWorkspace, workspaceToken } = storeToRefs(store)
+
+      await store.switchWorkspace('workspace-123')
+
+      let resolveRefreshFetch: (value: unknown) => void = () => {}
+      const refreshFetchPromise = new Promise((resolve) => {
+        resolveRefreshFetch = resolve
+      })
+      mockFetch.mockReturnValueOnce(refreshFetchPromise)
+
+      const refreshPromise = store.refreshToken()
+
+      const sameWorkspaceExpiry = new Date(
+        Date.now() + 7200 * 1000
+      ).toISOString()
+      mockFetch.mockResolvedValueOnce({
+        ok: true,
+        json: () =>
+          Promise.resolve({
+            ...mockTokenResponse,
+            token: 'same-workspace-token',
+            expires_at: sameWorkspaceExpiry
+          })
+      })
+      await store.switchWorkspace('workspace-123')
+
+      expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole)
+      expect(workspaceToken.value).toBe('same-workspace-token')
+
+      const refreshedExpiry = new Date(Date.now() + 9000 * 1000).toISOString()
+      resolveRefreshFetch({
+        ok: true,
+        json: () =>
+          Promise.resolve({
+            ...mockTokenResponse,
+            token: 'refreshed-workspace-token',
+            expires_at: refreshedExpiry
+          })
+      })
+      await refreshPromise
+
+      expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole)
+      expect(workspaceToken.value).toBe('refreshed-workspace-token')
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe(
+        'refreshed-workspace-token'
+      )
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)).toBe(
+        expectedExpiresAtMs(refreshedExpiry)
+      )
+    })
+
+    it('the new workspace wins when the stale refresh resolves last', async () => {
       mockGetIdToken.mockResolvedValue('firebase-token-xyz')
 
       const mockFetch = vi.fn().mockResolvedValueOnce({
@@ -817,6 +1099,149 @@ describe('useWorkspaceAuthStore', () => {
       const refreshPromise = store.refreshToken()
 
       // User switches workspace AND its fetch resolves first.
+      const newWorkspace = { ...mockWorkspace, id: 'workspace-other' }
+      const newExpiry = new Date(Date.now() + 7200 * 1000).toISOString()
+      mockFetch.mockResolvedValueOnce({
+        ok: true,
+        json: () =>
+          Promise.resolve({
+            ...mockTokenResponse,
+            token: 'new-workspace-token',
+            expires_at: newExpiry,
+            workspace: newWorkspace
+          })
+      })
+      await store.switchWorkspace('workspace-other')
+
+      // New workspace is committed at this point.
+      expect(currentWorkspace.value?.id).toBe('workspace-other')
+      expect(workspaceToken.value).toBe('new-workspace-token')
+      expect(
+        sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE)
+      ).toBe(JSON.stringify({ ...newWorkspace, role: 'owner' }))
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe(
+        'new-workspace-token'
+      )
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)).toBe(
+        expectedExpiresAtMs(newExpiry)
+      )
+
+      // Now resolve the stale refresh fetch — it carries an OLD-workspace
+      // token. It must not clobber the new workspace state or sessionStorage.
+      const staleExpiry = new Date(Date.now() + 1800 * 1000).toISOString()
+      resolveRefreshFetch({
+        ok: true,
+        json: () =>
+          Promise.resolve({
+            ...mockTokenResponse,
+            token: 'stale-token',
+            expires_at: staleExpiry
+          })
+      })
+      await refreshPromise
+
+      expect(currentWorkspace.value?.id).toBe('workspace-other')
+      expect(workspaceToken.value).toBe('new-workspace-token')
+      expect(
+        sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE)
+      ).toBe(JSON.stringify({ ...newWorkspace, role: 'owner' }))
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe(
+        'new-workspace-token'
+      )
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.EXPIRES_AT)).toBe(
+        expectedExpiresAtMs(newExpiry)
+      )
+    })
+
+    it('blocks a stale refresh that resolves after switch-away-and-back to same workspace', async () => {
+      mockGetIdToken.mockResolvedValue('firebase-token-xyz')
+
+      const mockFetch = vi.fn().mockResolvedValueOnce({
+        ok: true,
+        json: () => Promise.resolve(mockTokenResponse)
+      })
+      vi.stubGlobal('fetch', mockFetch)
+
+      const store = useWorkspaceAuthStore()
+      const { currentWorkspace, workspaceToken } = storeToRefs(store)
+
+      await store.switchWorkspace('workspace-123')
+
+      // Hang the refresh fetch so it resolves after both switches below.
+      let resolveRefreshFetch: (value: unknown) => void = () => {}
+      mockFetch.mockReturnValueOnce(
+        new Promise((resolve) => {
+          resolveRefreshFetch = resolve
+        })
+      )
+      const refreshPromise = store.refreshToken()
+
+      // Switch away to a different workspace...
+      const otherExpiry = new Date(Date.now() + 7200 * 1000).toISOString()
+      mockFetch.mockResolvedValueOnce({
+        ok: true,
+        json: () =>
+          Promise.resolve({
+            ...mockTokenResponse,
+            token: 'other-workspace-token',
+            expires_at: otherExpiry,
+            workspace: { ...mockTokenResponse.workspace, id: 'workspace-other' }
+          })
+      })
+      await store.switchWorkspace('workspace-other')
+      expect(workspaceToken.value).toBe('other-workspace-token')
+
+      // ...and back to the original workspace.
+      const backExpiry = new Date(Date.now() + 7200 * 1000).toISOString()
+      mockFetch.mockResolvedValueOnce({
+        ok: true,
+        json: () =>
+          Promise.resolve({
+            ...mockTokenResponse,
+            token: 'back-workspace-token',
+            expires_at: backExpiry
+          })
+      })
+      await store.switchWorkspace('workspace-123')
+      expect(workspaceToken.value).toBe('back-workspace-token')
+
+      // Stale refresh resolves with an old token — must not clobber state.
+      resolveRefreshFetch({
+        ok: true,
+        json: () =>
+          Promise.resolve({ ...mockTokenResponse, token: 'stale-token' })
+      })
+      await refreshPromise
+
+      expect(currentWorkspace.value?.id).toBe('workspace-123')
+      expect(workspaceToken.value).toBe('back-workspace-token')
+      expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe(
+        'back-workspace-token'
+      )
+    })
+
+    it('the new workspace keeps clean error state when a stale refresh fails last', async () => {
+      mockGetIdToken.mockResolvedValue('firebase-token-xyz')
+
+      const mockFetch = vi.fn().mockResolvedValueOnce({
+        ok: true,
+        json: () => Promise.resolve(mockTokenResponse)
+      })
+      vi.stubGlobal('fetch', mockFetch)
+
+      const store = useWorkspaceAuthStore()
+      const { currentWorkspace, workspaceToken, error } = storeToRefs(store)
+
+      await store.switchWorkspace('workspace-123')
+
+      let resolveRefreshFetch: (value: unknown) => void = () => {}
+      const refreshFetchPromise = new Promise((resolve) => {
+        resolveRefreshFetch = resolve
+      })
+      mockFetch.mockReturnValueOnce(refreshFetchPromise)
+
+      const refreshPromise = store.refreshToken()
+
       const newWorkspace = { ...mockWorkspace, id: 'workspace-other' }
       mockFetch.mockResolvedValueOnce({
         ok: true,
@@ -829,24 +1254,17 @@ describe('useWorkspaceAuthStore', () => {
       })
       await store.switchWorkspace('workspace-other')
 
-      // New workspace is committed at this point.
-      expect(currentWorkspace.value?.id).toBe('workspace-other')
-      expect(workspaceToken.value).toBe('new-workspace-token')
-
-      // Now resolve the stale refresh fetch — it carries an OLD-workspace
-      // token, and the source has no commit-time staleness check, so it
-      // clobbers the new workspace state.
       resolveRefreshFetch({
-        ok: true,
-        json: () =>
-          Promise.resolve({ ...mockTokenResponse, token: 'stale-token' })
+        ok: false,
+        status: 500,
+        statusText: 'Internal Server Error',
+        json: () => Promise.resolve({ message: 'Server error' })
       })
       await refreshPromise
 
-      // Once the source-side guard is added, both of these become true
-      // (the test stops failing) and `.fails` should be dropped.
       expect(currentWorkspace.value?.id).toBe('workspace-other')
       expect(workspaceToken.value).toBe('new-workspace-token')
+      expect(error.value).toBeNull()
     })
   })
 
@@ -861,27 +1279,42 @@ describe('useWorkspaceAuthStore', () => {
         })
       )
 
-      const setItemSpy = vi
-        .spyOn(sessionStorage, 'setItem')
-        .mockImplementation(() => {
+      const originalSessionStorage = globalThis.sessionStorage
+      // happy-dom Storage method spies can miss instance calls; replace the
+      // object so every setItem call deterministically throws.
+      const throwingSessionStorage = {
+        get length() {
+          return originalSessionStorage.length
+        },
+        key: originalSessionStorage.key.bind(originalSessionStorage),
+        getItem: originalSessionStorage.getItem.bind(originalSessionStorage),
+        setItem: vi.fn(() => {
           throw new Error('QuotaExceededError')
-        })
+        }),
+        removeItem: originalSessionStorage.removeItem.bind(
+          originalSessionStorage
+        ),
+        clear: originalSessionStorage.clear.bind(originalSessionStorage)
+      } satisfies Storage
+      vi.stubGlobal('sessionStorage', throwingSessionStorage)
       const consoleWarnSpy = vi
         .spyOn(console, 'warn')
         .mockImplementation(() => {})
 
-      const store = useWorkspaceAuthStore()
-      const { workspaceToken } = storeToRefs(store)
+      try {
+        const store = useWorkspaceAuthStore()
+        const { workspaceToken } = storeToRefs(store)
 
-      await store.switchWorkspace('workspace-123')
+        await store.switchWorkspace('workspace-123')
 
-      expect(workspaceToken.value).toBe('workspace-token-abc')
-      expect(consoleWarnSpy).toHaveBeenCalledWith(
-        'Failed to persist workspace context to sessionStorage'
-      )
-
-      setItemSpy.mockRestore()
-      consoleWarnSpy.mockRestore()
+        expect(workspaceToken.value).toBe('workspace-token-abc')
+        expect(consoleWarnSpy).toHaveBeenCalledWith(
+          'Failed to persist workspace context to sessionStorage'
+        )
+      } finally {
+        vi.stubGlobal('sessionStorage', originalSessionStorage)
+        consoleWarnSpy.mockRestore()
+      }
     })
   })
 

src/platform/workspace/stores/workspaceAuthStore.ts

@@ -37,6 +37,8 @@ export type WorkspaceTokenResponse = z.infer<
   typeof WorkspaceTokenResponseSchema
 >
 
+const MAX_SCHEDULED_REFRESH_RETRIES = 3
+
 export class WorkspaceAuthError extends Error {
   constructor(
     message: string,
@@ -59,6 +61,7 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
   // State
   const currentWorkspace = shallowRef<WorkspaceWithRole | null>(null)
   const workspaceToken = ref<string | null>(null)
+  const workspaceTokenExpiresAt = ref<number | null>(null)
   const isLoading = ref(false)
   const error = ref<Error | null>(null)
 
@@ -69,6 +72,8 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
 
   // Timer state
   let refreshTimerId: ReturnType<typeof setTimeout> | null = null
+  let inFlightSwitchCount = 0
+  let scheduledRefreshRetryCount = 0
   // The unified lifecycle keeps its own timer + request-id so it never shares
   // mutable state with the legacy switchWorkspace/refreshToken machinery.
   let unifiedRefreshTimerId: ReturnType<typeof setTimeout> | null = null
@@ -92,6 +97,7 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
 
   function scheduleTokenRefresh(expiresAt: number): void {
     stopRefreshTimer()
+    scheduledRefreshRetryCount = 0
     const now = Date.now()
     const refreshAt = expiresAt - TOKEN_REFRESH_BUFFER_MS
     const delay = Math.max(0, refreshAt - now)
@@ -101,6 +107,60 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
     }, delay)
   }
 
+  function scheduleClearAtExpiry(): void {
+    if (workspaceTokenExpiresAt.value === null) {
+      clearWorkspaceContext()
+      return
+    }
+
+    const timeUntilExpiry = workspaceTokenExpiresAt.value - Date.now()
+    if (timeUntilExpiry <= 0) {
+      clearWorkspaceContext()
+      return
+    }
+
+    stopRefreshTimer()
+    refreshTimerId = setTimeout(() => {
+      clearWorkspaceContext()
+    }, timeUntilExpiry)
+  }
+
+  function scheduleTokenRefreshRetry(delayMs: number): boolean {
+    if (workspaceTokenExpiresAt.value === null) {
+      clearWorkspaceContext()
+      return false
+    }
+
+    const timeUntilExpiry = workspaceTokenExpiresAt.value - Date.now()
+    if (timeUntilExpiry <= 0) {
+      clearWorkspaceContext()
+      return false
+    }
+
+    if (scheduledRefreshRetryCount >= MAX_SCHEDULED_REFRESH_RETRIES) {
+      scheduleClearAtExpiry()
+      return false
+    }
+
+    scheduledRefreshRetryCount += 1
+    stopRefreshTimer()
+    const timeUntilRefreshBuffer = Math.max(
+      0,
+      timeUntilExpiry - TOKEN_REFRESH_BUFFER_MS
+    )
+    refreshTimerId = setTimeout(
+      () => {
+        void refreshToken()
+      },
+      Math.min(delayMs, timeUntilRefreshBuffer)
+    )
+    return true
+  }
+
+  function isStaleWorkspaceRequest(capturedRequestId: number): boolean {
+    return capturedRequestId !== refreshRequestId
+  }
+
   function persistToSession(
     workspace: WorkspaceWithRole,
     token: string,
@@ -164,8 +224,9 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
         return false
       }
 
-      const parsedWorkspace = JSON.parse(workspaceJson)
-      const parseResult = WorkspaceWithRoleSchema.safeParse(parsedWorkspace)
+      const parseResult = WorkspaceWithRoleSchema.safeParse(
+        JSON.parse(workspaceJson)
+      )
 
       if (!parseResult.success) {
         clearSessionStorage()
@@ -174,6 +235,7 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
 
       currentWorkspace.value = parseResult.data
       workspaceToken.value = token
+      workspaceTokenExpiresAt.value = expiresAt
       error.value = null
 
       scheduleTokenRefresh(expiresAt)
@@ -275,29 +337,46 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
       return
     }
 
-    // Only increment request ID when switching to a different workspace
-    // This invalidates stale refresh operations for the old workspace
-    // but allows refresh operations for the same workspace to complete
-    if (currentWorkspace.value?.id !== workspaceId) {
-      refreshRequestId++
-    }
+    const capturedRequestId = refreshRequestId
 
+    inFlightSwitchCount += 1
     isLoading.value = true
-    error.value = null
 
     try {
       const { token, expiresAt, workspace } = await requestToken(workspaceId)
 
+      if (isStaleWorkspaceRequest(capturedRequestId)) {
+        console.warn(
+          'Aborting stale workspace switch: workspace context changed before commit'
+        )
+        return
+      }
+
+      if (currentWorkspace.value?.id !== workspaceId) {
+        refreshRequestId++
+      }
+      error.value = null
       currentWorkspace.value = workspace
       workspaceToken.value = token
+      workspaceTokenExpiresAt.value = expiresAt
+      scheduledRefreshRetryCount = 0
 
       persistToSession(workspace, token, expiresAt)
       scheduleTokenRefresh(expiresAt)
     } catch (err) {
+      if (isStaleWorkspaceRequest(capturedRequestId)) {
+        console.warn(
+          'Aborting stale workspace switch: workspace context changed before error commit',
+          err
+        )
+        return
+      }
+
       error.value = err instanceof Error ? err : new Error(String(err))
       throw error.value
     } finally {
-      isLoading.value = false
+      inFlightSwitchCount = Math.max(0, inFlightSwitchCount - 1)
+      isLoading.value = inFlightSwitchCount > 0
     }
   }
 
@@ -307,14 +386,15 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
     }
 
     const workspaceId = currentWorkspace.value.id
-    // Capture the current request ID to detect if workspace context changed during refresh
     const capturedRequestId = refreshRequestId
     const maxRetries = 3
     const baseDelayMs = 1000
+    // Clear any previous error optimistically; a stale-aborted refresh should
+    // not leave a stale error visible on the new workspace's context.
+    error.value = null
 
     for (let attempt = 0; attempt <= maxRetries; attempt++) {
-      // Check if workspace context changed since refresh started (user switched workspaces)
-      if (capturedRequestId !== refreshRequestId) {
+      if (isStaleWorkspaceRequest(capturedRequestId)) {
         console.warn(
           'Aborting stale token refresh: workspace context changed during refresh'
         )
@@ -335,8 +415,7 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
             err.code === 'NOT_AUTHENTICATED')
 
         if (isPermanentError) {
-          // Only clear context if this refresh is still for the current workspace
-          if (capturedRequestId === refreshRequestId) {
+          if (!isStaleWorkspaceRequest(capturedRequestId)) {
             console.error('Workspace access revoked or auth invalid:', err)
             clearWorkspaceContext()
           }
@@ -356,8 +435,21 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
           continue
         }
 
-        // Only clear context if this refresh is still for the current workspace
-        if (capturedRequestId === refreshRequestId) {
+        if (!isStaleWorkspaceRequest(capturedRequestId)) {
+          if (isTransientError && hasValidWorkspaceToken()) {
+            error.value = null
+            const retryScheduled = scheduleTokenRefreshRetry(
+              baseDelayMs * Math.pow(2, scheduledRefreshRetryCount)
+            )
+            console.warn(
+              retryScheduled
+                ? 'Failed to refresh workspace token after retries; preserving existing valid token and retrying later:'
+                : 'Failed to refresh workspace token after retries; preserving existing valid token until expiry:',
+              err
+            )
+            return
+          }
+
           console.error('Failed to refresh workspace token after retries:', err)
           clearWorkspaceContext()
         }
@@ -499,16 +591,23 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
     return workspaceToken.value ?? undefined
   }
 
+  function hasValidWorkspaceToken(): boolean {
+    return (
+      workspaceToken.value !== null &&
+      workspaceTokenExpiresAt.value !== null &&
+      workspaceTokenExpiresAt.value > Date.now()
+    )
+  }
+
   function clearWorkspaceContext(): void {
-    // Increment request ID to invalidate any in-flight stale refresh operations
     refreshRequestId++
     stopRefreshTimer()
     currentWorkspace.value = null
     workspaceToken.value = null
+    workspaceTokenExpiresAt.value = null
+    scheduledRefreshRetryCount = 0
     error.value = null
     clearSessionStorage()
-    // Drop the unified Cloud JWT and stop its timer on logout. Safe under any
-    // flag state: the slot is null when unified auth is off.
     clearUnifiedContext()
   }