diff --git a/apps/mobile/src/lib/organization-context.test.ts b/apps/mobile/src/lib/organization-context.test.ts index 5c6e84b7899..263fb9463a5 100644 --- a/apps/mobile/src/lib/organization-context.test.ts +++ b/apps/mobile/src/lib/organization-context.test.ts @@ -1,6 +1,7 @@ /* oxlint-disable @typescript-eslint/no-unsafe-call @typescript-eslint/no-unsafe-member-access */ +/* eslint-disable max-lines -- selection persistence and membership recovery share one provider harness */ import { type QueryClient } from '@tanstack/react-query'; -import { createElement } from 'react'; +import { type ComponentType, createElement, type ReactNode, useState } from 'react'; import { act } from '@/test/renderer'; import { renderWithProviders, waitFor } from '@/test/render-with-providers'; import { beforeEach, describe, expect, it, vi } from 'vitest'; @@ -65,7 +66,22 @@ type OrganizationContextValue = { setOrganizationId: (id: string | null) => void; }; -async function mountProvider(): Promise<{ +/** + * Hands the provider a new token while the tree stays mounted, the way an + * access-token rotation does: the same account and no auth-epoch bump. + */ +const tokenRotation: { set: (token: string) => void } = { set: () => undefined }; + +function RotatingTokenAuth({ children }: { children: ReactNode }): ReactNode { + const [token, setToken] = useState('t'); + tokenRotation.set = setToken; + hoisted.useAuth.mockReturnValue({ token }); + return createElement(OrganizationProvider, null, children); +} + +async function mountProvider( + wrapper: ComponentType<{ children: ReactNode }> = OrganizationProvider +): Promise<{ getCtx: () => OrganizationContextValue; queryClient: QueryClient; unmount: () => void; @@ -76,7 +92,7 @@ async function mountProvider(): Promise<{ return null; } - const ui = await renderWithProviders(createElement(Consumer), { wrapper: OrganizationProvider }); + const ui = await renderWithProviders(createElement(Consumer), { wrapper }); // Let the provider's restore read settle. The list may still be pending, so // this only guarantees the read decided whether a default resolution is owed. await act(async () => { @@ -147,6 +163,66 @@ describe('OrganizationProvider.setOrganizationId', () => { unmount(); }); + it('lands each selection after the previous one settled', async () => { + const stored = new Map(); + const gates: (() => void)[] = []; + /** Parks one storage call until the test releases it. */ + async function nextGate(): Promise { + const { promise, resolve } = Promise.withResolvers(); + gates.push(() => { + resolve(undefined); + }); + await promise; + } + // Every storage call parks until released, so the test decides when each + // write completes: the first selection completes last. + hoisted.setAccountMetadata.mockImplementation(async (key: string, value: string) => { + await nextGate(); + stored.set(key, value); + }); + hoisted.deleteAccountMetadata.mockImplementation(async (key: string) => { + await nextGate(); + stored.delete(key); + }); + + const { getCtx, unmount } = await mountProvider(); + + await act(() => { + getCtx().setOrganizationId('org-a'); + }); + await act(() => { + getCtx().setOrganizationId('org-b'); + }); + + // One selection at a time: the second one is written after the first pair. + expect(hoisted.setAccountMetadata).toHaveBeenCalledTimes(1); + expect(hoisted.setAccountMetadata).toHaveBeenCalledWith(ORGANIZATION_STORAGE_KEY, 'org-a'); + + async function releaseNext(): Promise { + const release = gates.shift(); + if (release) { + await act(async () => { + release(); + await Promise.resolve(); + }); + } + } + await releaseNext(); + await releaseNext(); + await releaseNext(); + await releaseNext(); + await releaseNext(); + + expect(stored.get(ORGANIZATION_STORAGE_KEY)).toBe('org-b'); + expect(stored.has(ORGANIZATION_PERSONAL_STORAGE_KEY)).toBe(false); + expect(hoisted.setAccountMetadata.mock.calls.at(-1)).toEqual([ + ORGANIZATION_STORAGE_KEY, + 'org-b', + ]); + + unmount(); + }); + it('clears the persisted org, unregisters tokens, and marks Personal on the way to personal', async () => { const { getCtx, unmount } = await mountProvider(); @@ -189,6 +265,125 @@ describe('OrganizationProvider.setOrganizationId', () => { unmount(); }); + + /** + * Storage double backed by one map whose calls park on demand, so a test can + * rotate the token while a save is in flight and then release it. + */ + function installStorageDouble(): { + stored: Map; + park: () => void; + releaseAll: () => Promise; + } { + const stored = new Map(); + const gates: (() => void)[] = []; + let parked = false; + async function waitWhileParked(): Promise { + if (!parked) { + return; + } + const { promise, resolve } = Promise.withResolvers(); + gates.push(() => { + resolve(undefined); + }); + await promise; + } + hoisted.getItemAsync.mockImplementation(async (key: string) => { + await Promise.resolve(); + return stored.get(key) ?? null; + }); + hoisted.setAccountMetadata.mockImplementation(async (key: string, value: string) => { + await waitWhileParked(); + stored.set(key, value); + }); + hoisted.deleteAccountMetadata.mockImplementation(async (key: string) => { + await waitWhileParked(); + stored.delete(key); + }); + return { + stored, + park: () => { + parked = true; + }, + releaseAll: async () => { + parked = false; + // Every parked call is released at once; callers await the observable + // state with `waitFor` rather than a fixed delay. + await act(async () => { + for (const release of gates.splice(0)) { + release(); + } + await Promise.resolve(); + }); + }, + }; + } + + /** Rotates the access token without a sign-out, as the refresh path does. */ + async function rotateToken(): Promise { + await act(async () => { + tokenRotation.set('t2'); + await Promise.resolve(); + }); + } + + it('keeps a fresh Personal choice when the token rotates while its save is in flight', async () => { + const storage = installStorageDouble(); + hoisted.list.mockResolvedValue([ + { organizationId: 'org-a', organizationName: 'A org', role: 'member' }, + ]); + const { getCtx, unmount } = await mountProvider(RotatingTokenAuth); + // The default organization is the stored selection before the switch. + await waitFor(() => storage.stored.get(ORGANIZATION_STORAGE_KEY) === 'org-a'); + expect(getCtx().organizationId).toBe('org-a'); + + // The explicit Personal choice parks on its first write, and a rotation + // lands while that save is still in flight. + storage.park(); + await act(() => { + getCtx().setOrganizationId(null); + }); + await rotateToken(); + await storage.releaseAll(); + + // The pair the user chose is the one stored, and the restore publishes it + // instead of falling back to the default organization. + await waitFor(() => storage.stored.get(ORGANIZATION_PERSONAL_STORAGE_KEY) === 'personal'); + expect(storage.stored.has(ORGANIZATION_STORAGE_KEY)).toBe(false); + await waitFor(() => getCtx().organizationId === null); + expect(getCtx().organizationId).toBeNull(); + + unmount(); + }); + + it('lands the last of two queued selections when the token rotates', async () => { + const storage = installStorageDouble(); + hoisted.list.mockResolvedValue([ + { organizationId: 'org-a', organizationName: 'A org', role: 'member' }, + { organizationId: 'org-b', organizationName: 'B org', role: 'member' }, + { organizationId: 'org-c', organizationName: 'C org', role: 'member' }, + ]); + const { getCtx, unmount } = await mountProvider(RotatingTokenAuth); + await waitFor(() => storage.stored.get(ORGANIZATION_STORAGE_KEY) === 'org-a'); + + storage.park(); + await act(() => { + getCtx().setOrganizationId('org-b'); + }); + await act(() => { + getCtx().setOrganizationId('org-c'); + }); + expect(getCtx().organizationId).toBe('org-c'); + await rotateToken(); + await storage.releaseAll(); + + await waitFor(() => storage.stored.get(ORGANIZATION_STORAGE_KEY) === 'org-c'); + expect(storage.stored.has(ORGANIZATION_PERSONAL_STORAGE_KEY)).toBe(false); + await waitFor(() => getCtx().organizationId === 'org-c'); + expect(getCtx().organizationId).toBe('org-c'); + + unmount(); + }); }); describe('OrganizationProvider lost-membership recovery', () => { diff --git a/apps/mobile/src/lib/organization-context.tsx b/apps/mobile/src/lib/organization-context.tsx index 3cd348b5032..b385f933218 100644 --- a/apps/mobile/src/lib/organization-context.tsx +++ b/apps/mobile/src/lib/organization-context.tsx @@ -21,6 +21,7 @@ import { unregisterActivityTokensAndTombstone } from '@/lib/auth/logout-cleanup' import { writePrivacySnapshotAndEnd } from '@/lib/glanceable/cleanup'; import { useOrganizationsList } from '@/lib/hooks/use-organizations-list'; import { type OrgListEntry } from '@/lib/hooks/use-organization-queries'; +import { chainSave } from '@/lib/hooks/save-chain'; import { announcingToast } from '@/lib/a11y/announcing-toast'; import { i18n } from '@/i18n'; import { ORGANIZATION_PERSONAL_STORAGE_KEY, ORGANIZATION_STORAGE_KEY } from '@/lib/storage-keys'; @@ -36,6 +37,14 @@ import { ORGANIZATION_PERSONAL_STORAGE_KEY, ORGANIZATION_STORAGE_KEY } from '@/l */ const PERSONAL_MARKER = 'personal'; +/** + * One queue slot for the whole selection: a save runs after the previous save + * settled, so the two keys above are written as one pair and in the order the + * selections were made, and the stored pair always names the selection the + * provider publishes. + */ +const SELECTION_QUEUE_KEY = 'organization-selection'; + type OrganizationContextValue = { /** null = personal, string = org UUID */ organizationId: string | null; @@ -73,6 +82,14 @@ export function OrganizationProvider({ children }: { readonly children: ReactNod error: null, }); const generation = useRef(0); + // Orders the stored pair's writers: a newer selection, a sign-out, or an + // account switch supersedes an in-flight save. `restore` deliberately does + // not bump it: a token rotation re-runs `restore` for the same account, and + // that must not cancel the save of a selection the user just made. + const saveGeneration = useRef(0); + // Settles when the last scheduled save has settled. A restore reads the pair + // a save that began before it left behind, instead of a half-written one. + const savesSettled = useRef>(Promise.resolve()); const activeId = useRef(null); // Set by `restore` when nothing explicit is stored: the organization list // owns the publish then. A ref (not state) so `setOrganizationId` clears it @@ -107,11 +124,15 @@ export function OrganizationProvider({ children }: { readonly children: ReactNod // A stored organization is an explicit choice; the marker beside it is an // explicit Personal choice. An absent value is 'not chosen yet': the // organization list decides the default and the effect below publishes it. - // A legacy absent value is treated the same way. + // A legacy absent value is treated the same way. A save that began before + // this restore owns the pair until it settles, so it is awaited first; a + // save that begins later supersedes this restore (it bumps `generation`), + // so the reads below are discarded rather than published half-written. // The two reads are kept apart: a stored organization stands alone, so a // marker-read failure must not discard it. The marker is only consulted // when no organization is stored, and a failure there stays a restore // error — the selection is then genuinely undetermined. + await savesSettled.current; const stored = await SecureStore.getItemAsync(ORGANIZATION_STORAGE_KEY); const marker = stored === null ? await SecureStore.getItemAsync(ORGANIZATION_PERSONAL_STORAGE_KEY) : null; @@ -142,22 +163,52 @@ export function OrganizationProvider({ children }: { readonly children: ReactNod }, [token]); const persist = useCallback(async (id: string | null) => { + // A selection supersedes any restore still deciding what to publish. generation.current += 1; - const operation = generation.current; + saveGeneration.current += 1; + const save = saveGeneration.current; const epoch = currentAuthEpoch(); setState(current => ({ ...current, isSaving: true })); let error: 'save' | null = null; - try { + // Registered before the first await, so a restore that starts now waits for + // this whole pair: it reads what the save leaves behind, never a key of it. + const writes = chainSave(SELECTION_QUEUE_KEY, async () => { + // Only the selection the provider still publishes is written: a newer + // selection, a sign-in, or a sign-out owns the stored state, and its + // own save (or the sign-out cleanup, whose deletes chain behind these + // keys) writes the pair this save would have written. A restore does + // not own the pair: it reads it after this save settles. + if (saveGeneration.current !== save || !isCurrentAuthEpoch(epoch)) { + return; + } + // Keep the pair ordered with saves, but stop between keys if a new + // selection or auth epoch supersedes this one. Its queued save, or + // sign-out cleanup, owns the remaining storage update. await (id ? setAccountMetadata(ORGANIZATION_STORAGE_KEY, id) : deleteAccountMetadata(ORGANIZATION_STORAGE_KEY)); + if (saveGeneration.current !== save || !isCurrentAuthEpoch(epoch)) { + return; + } await (id ? deleteAccountMetadata(ORGANIZATION_PERSONAL_STORAGE_KEY) : setAccountMetadata(ORGANIZATION_PERSONAL_STORAGE_KEY, PERSONAL_MARKER)); + }); + // A restore only needs the pair to have settled, so its own rejection is + // swallowed here; this save reports the failure in its state below. + savesSettled.current = (async () => { + try { + await writes; + } catch { + // Reported by this save's own `error` state. + } + })(); + try { + await writes; } catch { error = 'save'; } - if (generation.current === operation && isCurrentAuthEpoch(epoch)) { + if (saveGeneration.current === save && isCurrentAuthEpoch(epoch)) { setState(current => ({ ...current, error, isSaving: false })); } }, []); @@ -268,6 +319,9 @@ export function OrganizationProvider({ children }: { readonly children: ReactNod void restore(); } else { generation.current += 1; + // A signed-out account owns no stored selection: a save still in flight + // must not write the pair after the sign-out cleanup deleted it. + saveGeneration.current += 1; activeId.current = null; awaitingDefault.current = false; defaultFence.current = null; @@ -277,6 +331,8 @@ export function OrganizationProvider({ children }: { readonly children: ReactNod setState({ token, organizationId: null, isLoaded: true, isSaving: false, error: null }); } return () => { + // A token change cancels the previous restore's publication, not a save: + // a rotation is the same account, and its save must still land. generation.current += 1; }; }, [token, restore]);