From 59a52e55e783574bd201360c5ab5bc9b95cf55e2 Mon Sep 17 00:00:00 2001 From: Benjamin Lu Date: Tue, 28 Apr 2026 06:51:10 -0700 Subject: [PATCH 1/2] fix: guard workspace auth refresh races --- .../workspace/stores/useWorkspaceAuth.test.ts | 102 ++++++++++-------- .../workspace/stores/workspaceAuthStore.ts | 35 ++++++ 2 files changed, 95 insertions(+), 42 deletions(-) diff --git a/src/platform/workspace/stores/useWorkspaceAuth.test.ts b/src/platform/workspace/stores/useWorkspaceAuth.test.ts index 74728daa089..c310b178e4e 100644 --- a/src/platform/workspace/stores/useWorkspaceAuth.test.ts +++ b/src/platform/workspace/stores/useWorkspaceAuth.test.ts @@ -671,14 +671,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. @@ -689,10 +682,11 @@ describe('useWorkspaceAuthStore', () => { vi.stubGlobal('fetch', mockFetch) const store = useWorkspaceAuthStore() - const { currentWorkspace } = storeToRefs(store) + const { currentWorkspace, workspaceToken } = 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({ @@ -711,8 +705,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. @@ -734,8 +731,16 @@ 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(consoleErrorSpy).not.toHaveBeenCalled() consoleErrorSpy.mockRestore() consoleWarnSpy.mockRestore() @@ -776,16 +781,7 @@ 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('the new workspace wins when the stale refresh resolves last', async () => { mockGetIdToken.mockResolvedValue('firebase-token-xyz') const mockFetch = vi.fn().mockResolvedValueOnce({ @@ -824,10 +820,15 @@ describe('useWorkspaceAuthStore', () => { // 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' + ) // 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. + // token. It must not clobber the new workspace state or sessionStorage. resolveRefreshFetch({ ok: true, json: () => @@ -835,10 +836,14 @@ describe('useWorkspaceAuthStore', () => { }) 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( + sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE) + ).toBe(JSON.stringify({ ...newWorkspace, role: 'owner' })) + expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe( + 'new-workspace-token' + ) }) }) @@ -853,27 +858,40 @@ describe('useWorkspaceAuthStore', () => { }) ) - const setItemSpy = vi - .spyOn(sessionStorage, 'setItem') - .mockImplementation(() => { + const originalSessionStorage = globalThis.sessionStorage + 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) - - await store.switchWorkspace('workspace-123') + try { + const store = useWorkspaceAuthStore() + const { workspaceToken } = storeToRefs(store) - expect(workspaceToken.value).toBe('workspace-token-abc') - expect(consoleWarnSpy).toHaveBeenCalledWith( - 'Failed to persist workspace context to sessionStorage' - ) + await store.switchWorkspace('workspace-123') - 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() + } }) }) diff --git a/src/platform/workspace/stores/workspaceAuthStore.ts b/src/platform/workspace/stores/workspaceAuthStore.ts index 4a0ddb09b4f..b1d85e3038a 100644 --- a/src/platform/workspace/stores/workspaceAuthStore.ts +++ b/src/platform/workspace/stores/workspaceAuthStore.ts @@ -49,6 +49,7 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => { // State const currentWorkspace = shallowRef(null) const workspaceToken = ref(null) + const workspaceTokenExpiresAt = ref(null) const isLoading = ref(false) const error = ref(null) @@ -155,6 +156,7 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => { currentWorkspace.value = parseResult.data workspaceToken.value = token + workspaceTokenExpiresAt.value = expiresAt error.value = null scheduleTokenRefresh(expiresAt) @@ -176,6 +178,7 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => { if (currentWorkspace.value?.id !== workspaceId) { refreshRequestId++ } + const capturedRequestId = refreshRequestId isLoading.value = true error.value = null @@ -257,12 +260,27 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => { role: data.role } + if (capturedRequestId !== refreshRequestId) { + console.warn( + 'Aborting stale workspace switch: workspace context changed before commit' + ) + return + } + currentWorkspace.value = workspaceWithRole workspaceToken.value = data.token + workspaceTokenExpiresAt.value = expiresAt persistToSession(workspaceWithRole, data.token, expiresAt) scheduleTokenRefresh(expiresAt) } catch (err) { + if (capturedRequestId !== refreshRequestId) { + console.warn( + 'Aborting stale workspace switch: workspace context changed before error commit' + ) + return + } + error.value = err instanceof Error ? err : new Error(String(err)) throw error.value } finally { @@ -327,6 +345,14 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => { // Only clear context if this refresh is still for the current workspace if (capturedRequestId === refreshRequestId) { + if (isTransientError && hasValidWorkspaceToken()) { + console.warn( + 'Failed to refresh workspace token after retries; preserving existing valid token:', + err + ) + return + } + console.error('Failed to refresh workspace token after retries:', err) clearWorkspaceContext() } @@ -347,12 +373,21 @@ 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 error.value = null clearSessionStorage() } From c3bcc28fbca8ea22271b29758866facd8f8db38e Mon Sep 17 00:00:00 2001 From: Benjamin Lu Date: Tue, 28 Apr 2026 08:17:31 -0700 Subject: [PATCH 2/2] fix: retry preserved workspace token refreshes --- .../workspace/stores/useWorkspaceAuth.test.ts | 18 +++++++++++++++- .../workspace/stores/workspaceAuthStore.ts | 21 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/platform/workspace/stores/useWorkspaceAuth.test.ts b/src/platform/workspace/stores/useWorkspaceAuth.test.ts index c310b178e4e..ae0f5446160 100644 --- a/src/platform/workspace/stores/useWorkspaceAuth.test.ts +++ b/src/platform/workspace/stores/useWorkspaceAuth.test.ts @@ -682,7 +682,7 @@ describe('useWorkspaceAuthStore', () => { vi.stubGlobal('fetch', mockFetch) const store = useWorkspaceAuthStore() - const { currentWorkspace, workspaceToken } = storeToRefs(store) + const { currentWorkspace, workspaceToken, error } = storeToRefs(store) await store.switchWorkspace('workspace-123') expect(currentWorkspace.value).toEqual(mockWorkspaceWithRole) @@ -740,8 +740,24 @@ describe('useWorkspaceAuthStore', () => { expect(sessionStorage.getItem(WORKSPACE_STORAGE_KEYS.TOKEN)).toBe( 'workspace-token-abc' ) + expect(error.value).toBeNull() expect(consoleErrorSpy).not.toHaveBeenCalled() + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => + Promise.resolve({ ...mockTokenResponse, token: 'retry-token' }) + }) + + await vi.advanceTimersByTimeAsync(7999) + expect(mockFetch).toHaveBeenCalledTimes(5) + + await vi.advanceTimersByTimeAsync(1) + expect(mockFetch).toHaveBeenCalledTimes(6) + await vi.waitFor(() => { + expect(workspaceToken.value).toBe('retry-token') + }) + consoleErrorSpy.mockRestore() consoleWarnSpy.mockRestore() }) diff --git a/src/platform/workspace/stores/workspaceAuthStore.ts b/src/platform/workspace/stores/workspaceAuthStore.ts index b1d85e3038a..5f80fe4dcc5 100644 --- a/src/platform/workspace/stores/workspaceAuthStore.ts +++ b/src/platform/workspace/stores/workspaceAuthStore.ts @@ -83,6 +83,25 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => { }, delay) } + function scheduleTokenRefreshRetry(delayMs: number): void { + if (workspaceTokenExpiresAt.value === null) { + return + } + + const timeUntilExpiry = workspaceTokenExpiresAt.value - Date.now() + if (timeUntilExpiry <= 0) { + return + } + + stopRefreshTimer() + refreshTimerId = setTimeout( + () => { + void refreshToken() + }, + Math.min(delayMs, timeUntilExpiry) + ) + } + function persistToSession( workspace: WorkspaceWithRole, token: string, @@ -346,6 +365,8 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => { // Only clear context if this refresh is still for the current workspace if (capturedRequestId === refreshRequestId) { if (isTransientError && hasValidWorkspaceToken()) { + error.value = null + scheduleTokenRefreshRetry(baseDelayMs * Math.pow(2, maxRetries)) console.warn( 'Failed to refresh workspace token after retries; preserving existing valid token:', err