fix(mobile): keep one organization selection value per account

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<string, string>();
+    const gates: (() => void)[] = [];
+    /** Parks one storage call until the test releases it. */
+    async function nextGate(): Promise<void> {
+      const { promise, resolve } = Promise.withResolvers<undefined>();
+      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<void> {
+      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<string, string>;
+    park: () => void;
+    releaseAll: () => Promise<void>;
+  } {
+    const stored = new Map<string, string>();
+    const gates: (() => void)[] = [];
+    let parked = false;
+    async function waitWhileParked(): Promise<void> {
+      if (!parked) {
+        return;
+      }
+      const { promise, resolve } = Promise.withResolvers<undefined>();
+      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<void> {
+    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', () => {

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<void>>(Promise.resolve());
   const activeId = useRef<string | null>(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]);