Fix stale cross-tab session writes

src/state/persisted/index.ts

@@ -8,7 +8,7 @@ import {
   tryStringify,
 } from '#/state/persisted/schema'
 import {device} from '#/storage'
-import {type PersistedApi} from './types'
+import {type PersistedApi, type PersistedWriteValue} from './types'
 import {normalizeData} from './util'
 
 export type {PersistedAccount, Schema} from '#/state/persisted/schema'
@@ -45,16 +45,22 @@ readLatest satisfies PersistedApi['readLatest']
 
 export async function write<K extends keyof Schema>(
   key: K,
-  value: Schema[K],
+  value: PersistedWriteValue<Schema[K]>,
 ): Promise<void> {
   _state = normalizeData({
     ..._state,
-    [key]: value,
+    [key]: resolveWriteValue(value, _state[key]),
   })
   await writeToStorage(_state)
 }
 write satisfies PersistedApi['write']
 
+function resolveWriteValue<T>(value: PersistedWriteValue<T>, latest: T): T {
+  return typeof value === 'function'
+    ? (value as (current: T) => T)(latest)
+    : value
+}
+
 export function onUpdate<K extends keyof Schema>(
   _key: K,
   _cb: (v: Schema[K]) => void,

src/state/persisted/index.web.ts

@@ -8,7 +8,7 @@ import {
   tryParse,
   tryStringify,
 } from '#/state/persisted/schema'
-import {type PersistedApi} from './types'
+import {type PersistedApi, type PersistedWriteValue} from './types'
 import {normalizeData} from './util'
 
 export type {PersistedAccount, Schema} from '#/state/persisted/schema'
@@ -65,7 +65,7 @@ readLatest satisfies PersistedApi['readLatest']
 // eslint-disable-next-line @typescript-eslint/require-await
 export async function write<K extends keyof Schema>(
   key: K,
-  value: Schema[K],
+  value: PersistedWriteValue<Schema[K]>,
 ): Promise<void> {
   const next = readFromStorage()
   if (next) {
@@ -75,8 +75,9 @@ export async function write<K extends keyof Schema>(
     // Don't fire the update listeners yet to avoid a loop.
     // If there was a change, we'll receive the broadcast event soon enough which will do that.
   }
+  const nextValue = resolveWriteValue(value, _state[key])
   try {
-    if (JSON.stringify({v: _state[key]}) === JSON.stringify({v: value})) {
+    if (JSON.stringify({v: _state[key]}) === JSON.stringify({v: nextValue})) {
       // Fast path for updates that are guaranteed to be noops.
       // This is good mostly because it avoids useless broadcasts to other tabs.
       return
@@ -86,7 +87,7 @@ export async function write<K extends keyof Schema>(
   }
   _state = normalizeData({
     ..._state,
-    [key]: value,
+    [key]: nextValue,
   })
   writeToStorage(_state)
   broadcast.postMessage({event: {type: UPDATE_EVENT, key}})
@@ -94,6 +95,12 @@ export async function write<K extends keyof Schema>(
 }
 write satisfies PersistedApi['write']
 
+function resolveWriteValue<T>(value: PersistedWriteValue<T>, latest: T): T {
+  return typeof value === 'function'
+    ? (value as (current: T) => T)(latest)
+    : value
+}
+
 export function onUpdate<K extends keyof Schema>(
   key: K,
   cb: (v: Schema[K]) => void,

src/state/persisted/types.ts

@@ -1,5 +1,11 @@
 import {type Schema} from './schema'
 
+/**
+ * A value to persist, or a function that derives one from the freshest
+ * persisted value while the write operation is in progress.
+ */
+export type PersistedWriteValue<T> = T | ((latest: T) => T)
+
 export type PersistedApi = {
   init(): Promise<void>
   get<K extends keyof Schema>(key: K): Schema[K]
@@ -12,7 +18,10 @@ export type PersistedApi = {
    * it is identical to {@link get} (single-instance, no other writer).
    */
   readLatest<K extends keyof Schema>(key: K): Schema[K]
-  write<K extends keyof Schema>(key: K, value: Schema[K]): Promise<void>
+  write<K extends keyof Schema>(
+    key: K,
+    value: PersistedWriteValue<Schema[K]>,
+  ): Promise<void>
   onUpdate<K extends keyof Schema>(
     key: K,
     cb: (v: Schema[K]) => void,

src/state/session/__tests__/session-test.ts

@@ -1,7 +1,13 @@
 import {type SessionData} from '@atproto/lex-password-session'
 import {describe, expect, it, jest} from '@jest/globals'
 
-import {type Action, getInitialState, reducer, type State} from '../reducer'
+import {
+  type Action,
+  getInitialState,
+  rebasePersistedSession,
+  reducer,
+  type State,
+} from '../reducer'
 import {sessionDataToSessionAccount} from '../session-core'
 import {type SessionAccount} from '../types'
 
@@ -1777,8 +1783,233 @@ describe('session', () => {
     expect(state.currentBundleState.did).toBe('alice-did')
     expect(state.needsPersist).toBe(true)
   })
+
+  describe('rebasePersistedSession', () => {
+    it('keeps a newer token generation when a stale metadata update is persisted', () => {
+      const stale = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-1',
+        refreshJwt: 'alice-refresh-jwt-1',
+        emailConfirmed: false,
+        emailAuthFactor: false,
+      })
+      const fresh = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-2',
+        refreshJwt: 'alice-refresh-jwt-2',
+        emailConfirmed: false,
+        emailAuthFactor: false,
+      })
+      const bob = makeAccount('https://bob.com', {
+        active: true,
+        did: 'bob-did',
+        handle: 'bob.test',
+        accessJwt: 'bob-access-jwt-1',
+        refreshJwt: 'bob-refresh-jwt-1',
+      })
+      const desired = snapshot([{...stale, emailConfirmed: true}], 'alice-did')
+      const latest = snapshot([fresh, bob], 'alice-did')
+
+      const rebased = rebasePersistedSession(latest, desired, {
+        type: 'partial-refresh-session',
+        accountDid: 'alice-did',
+        patch: {emailConfirmed: true, emailAuthFactor: false},
+      })
+
+      expect(rebased.accounts).toEqual([{...fresh, emailConfirmed: true}, bob])
+      expect(rebased.currentAccount?.did).toBe('alice-did')
+    })
+
+    it('accepts a fresh token pair delivered by the active session', () => {
+      const stale = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-1',
+        refreshJwt: 'alice-refresh-jwt-1',
+      })
+      const fresh = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-2',
+        refreshJwt: 'alice-refresh-jwt-2',
+      })
+
+      const rebased = rebasePersistedSession(
+        snapshot([stale], 'alice-did'),
+        snapshot([fresh], 'alice-did'),
+        {
+          type: 'received-session-event',
+          bundle: makeBundle('https://alice.com'),
+          accountDid: 'alice-did',
+          refreshedAccount: fresh,
+          sessionEvent: 'update',
+        },
+      )
+
+      expect(rebased.accounts).toEqual([fresh])
+    })
+
+    it('preserves a concurrent newer generation instead of clearing it on expiry', () => {
+      const dying = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-1',
+        refreshJwt: 'alice-refresh-jwt-1',
+      })
+      const fresh = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-2',
+        refreshJwt: 'alice-refresh-jwt-2',
+      })
+      const expired = {...dying, accessJwt: undefined, refreshJwt: undefined}
+
+      const rebased = rebasePersistedSession(
+        snapshot([fresh], 'alice-did'),
+        snapshot([expired]),
+        {
+          type: 'received-session-event',
+          bundle: makeBundle('https://alice.com'),
+          accountDid: 'alice-did',
+          refreshedAccount: undefined,
+          sessionEvent: 'expired',
+          expiredRefreshJwt: dying.refreshJwt,
+        },
+      )
+
+      expect(rebased.accounts).toEqual([fresh])
+      expect(rebased.currentAccount?.did).toBe('alice-did')
+    })
+
+    it('keeps explicit logout authoritative over a newer stored token', () => {
+      const fresh = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-2',
+        refreshJwt: 'alice-refresh-jwt-2',
+      })
+      const loggedOut = {...fresh, accessJwt: undefined, refreshJwt: undefined}
+
+      const rebased = rebasePersistedSession(
+        snapshot([fresh], 'alice-did'),
+        snapshot([loggedOut]),
+        {type: 'logged-out-current-account'},
+      )
+
+      expect(rebased.accounts).toEqual([loggedOut])
+      expect(rebased.currentAccount).toBeUndefined()
+    })
+
+    it('clears only the explicitly logged-out account', () => {
+      const alice = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-2',
+        refreshJwt: 'alice-refresh-jwt-2',
+      })
+      const bobStale = makeAccount('https://bob.com', {
+        active: true,
+        did: 'bob-did',
+        handle: 'bob.test',
+        accessJwt: 'bob-access-jwt-1',
+        refreshJwt: 'bob-refresh-jwt-1',
+      })
+      const bobFresh = {
+        ...bobStale,
+        accessJwt: 'bob-access-jwt-2',
+        refreshJwt: 'bob-refresh-jwt-2',
+      }
+      const loggedOutAlice = {
+        ...alice,
+        accessJwt: undefined,
+        refreshJwt: undefined,
+      }
+
+      const rebased = rebasePersistedSession(
+        snapshot([alice, bobFresh], 'alice-did'),
+        snapshot([loggedOutAlice, bobStale]),
+        {type: 'logged-out-current-account', accountDid: 'alice-did'},
+      )
+
+      expect(rebased.accounts).toEqual([loggedOutAlice, bobFresh])
+    })
+
+    it('does not restore an account another tab removed', () => {
+      const stale = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-1',
+        refreshJwt: 'alice-refresh-jwt-1',
+      })
+
+      const rebased = rebasePersistedSession(
+        snapshot([]),
+        snapshot([{...stale, emailConfirmed: true}], 'alice-did'),
+        {
+          type: 'partial-refresh-session',
+          accountDid: 'alice-did',
+          patch: {emailConfirmed: true, emailAuthFactor: false},
+        },
+      )
+
+      expect(rebased.accounts).toEqual([])
+      expect(rebased.currentAccount).toBeUndefined()
+    })
+
+    it('clears credentials for accounts added before an all-account logout', () => {
+      const alice = makeAccount('https://alice.com', {
+        active: true,
+        did: 'alice-did',
+        handle: 'alice.test',
+        accessJwt: 'alice-access-jwt-2',
+        refreshJwt: 'alice-refresh-jwt-2',
+      })
+      const bob = makeAccount('https://bob.com', {
+        active: true,
+        did: 'bob-did',
+        handle: 'bob.test',
+        accessJwt: 'bob-access-jwt-2',
+        refreshJwt: 'bob-refresh-jwt-2',
+      })
+      const loggedOutAlice = {
+        ...alice,
+        accessJwt: undefined,
+        refreshJwt: undefined,
+      }
+
+      const rebased = rebasePersistedSession(
+        snapshot([alice, bob], 'alice-did'),
+        snapshot([loggedOutAlice]),
+        {type: 'logged-out-every-account'},
+      )
+
+      expect(rebased.accounts).toEqual([
+        loggedOutAlice,
+        {...bob, accessJwt: undefined, refreshJwt: undefined},
+      ])
+      expect(rebased.currentAccount).toBeUndefined()
+    })
+  })
 })
 
+function snapshot(accounts: SessionAccount[], currentDid?: string) {
+  return {
+    accounts,
+    currentAccount: accounts.find(account => account.did === currentDid),
+  }
+}
+
 function run(initialState: State, actions: Action[]): State {
   let state = initialState
   for (let action of actions) {

src/state/session/index.tsx

@@ -22,7 +22,13 @@ import {emitSessionDropped} from '../events'
 import {getPublicAppviewClient} from './clients'
 import {createSessionBundleAndCreateAccount} from './create-account'
 import {pickExpiryRescueCandidate} from './expiry-rescue'
-import {type Action, getInitialState, reducer, type State} from './reducer'
+import {
+  type Action,
+  getInitialState,
+  rebasePersistedSession,
+  reducer,
+  type State,
+} from './reducer'
 import {
   type AtpSessionEvent,
   createSessionBundleAndLogin,
@@ -108,17 +114,29 @@ class SessionStore {
     // Persist synchronously without waiting for the React render cycle.
     if (nextState.needsPersist) {
       nextState.needsPersist = false
-      const persistedData = {
+      const desiredPersistedData = {
         accounts: nextState.accounts,
         currentAccount: nextState.accounts.find(
           a => a.did === nextState.currentBundleState.did,
         ),
       }
-      addSessionDebugLog({
+      /*
-        type: 'persisted:broadcast',
+       * A background tab can have missed a token-rotation broadcast. Derive the
-        data: redactPersistedSession(persistedData),
+       * outgoing value inside `write`, after its final web storage re-read, so an
+       * unrelated metadata update cannot restore an older token generation.
+       */
+      void persisted.write('session', latest => {
+        const persistedData = rebasePersistedSession(
+          latest,
+          desiredPersistedData,
+          action,
+        )
+        addSessionDebugLog({
+          type: 'persisted:broadcast',
+          data: redactPersistedSession(persistedData),
+        })
+        return persistedData
       })
-      void persisted.write('session', persistedData)
     }
     this.listeners.forEach(listener => listener())
   }
@@ -257,6 +275,8 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         refreshedAccount,
         accountDid,
         sessionEvent,
+        expiredRefreshJwt:
+          sessionEvent === 'expired' ? sessionData?.refreshJwt : undefined,
       })
     },
     [store],
@@ -290,6 +310,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         type: 'switched-to-account',
         newBundle: bundle,
         newAccount: account,
+        tokenUpdate: 'login',
       })
       ax.metric('account:create:success', metrics, {
         session: utils.accountToSessionMetadata(account),
@@ -321,6 +342,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         type: 'switched-to-account',
         newBundle: bundle,
         newAccount: account,
+        tokenUpdate: 'login',
       })
       ax.metric(
         'account:loggedIn',
@@ -345,6 +367,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const prevState = store.getState()
       store.dispatch({
         type: 'logged-out-current-account',
+        accountDid: prevState.currentBundleState.did,
       })
       ax.metric(
         'account:loggedOut',
@@ -435,6 +458,10 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         type: 'switched-to-account',
         newBundle: bundle,
         newAccount: account,
+        tokenUpdate:
+          account.refreshJwt !== storedAccount.refreshJwt
+            ? 'refresh'
+            : undefined,
       })
       addSessionDebugLog({
         type: 'method:end',

src/state/session/reducer.ts

@@ -1,5 +1,6 @@
 import {unregisterPushToken} from '#/lib/notifications/notifications'
 import {logger} from '#/lib/notifications/util'
+import {type Schema} from '#/state/persisted/schema'
 import {wrapSessionReducerForLogging} from './logging'
 import {createPublicSessionBundle} from './session-core'
 import {type AtpSessionEvent, type SessionAccount} from './types'
@@ -28,11 +29,15 @@ export type Action =
       accountDid: string
       refreshedAccount: SessionAccount | undefined
       sessionEvent: AtpSessionEvent
+      /** The token that failed, used to preserve a newer concurrent generation. */
+      expiredRefreshJwt?: string
     }
   | {
       type: 'switched-to-account'
       newBundle: OpaqueSessionBundle
       newAccount: SessionAccount
+      /** Whether this action obtained a fresh credential generation. */
+      tokenUpdate?: 'login' | 'refresh'
     }
   | {
       // Replace an immutable session from synced or rescued tokens without rebroadcasting.
@@ -46,6 +51,8 @@ export type Action =
     }
   | {
       type: 'logged-out-current-account'
+      /** Omitted only by legacy reducer tests; production always supplies it. */
+      accountDid?: string
     }
   | {
       type: 'logged-out-every-account'
@@ -61,6 +68,164 @@ export type Action =
       patch: Pick<SessionAccount, 'emailConfirmed' | 'emailAuthFactor'>
     }
 
+/**
+ * Rebase a pending session write over the newest persisted snapshot.
+ *
+ * Browser tabs hold independent in-memory session state. A tab which missed a
+ * token-rotation broadcast may still need to persist unrelated metadata. The
+ * persisted layer can protect other root keys, but it cannot tell whether the
+ * complete `session` value it receives contains an older credential generation.
+ * Keep the newest stored credential pair unless this action is known to have
+ * obtained a fresh pair or intentionally cleared it.
+ */
+export function rebasePersistedSession(
+  latest: Schema['session'],
+  desired: Schema['session'],
+  action: Action,
+): Schema['session'] {
+  const latestByDid = new Map(
+    latest.accounts.map(account => [account.did, account]),
+  )
+  const desiredDids = new Set(desired.accounts.map(account => account.did))
+  const removedDid =
+    action.type === 'removed-account' ? action.accountDid : undefined
+  const preservedNewerTokenDids = new Set<string>()
+
+  const accounts = desired.accounts.flatMap(desiredAccount => {
+    const latestAccount = latestByDid.get(desiredAccount.did)
+    if (!latestAccount) {
+      /* Only an explicit fresh login may introduce an account absent from storage. */
+      return shouldAddMissingAccount(action, desiredAccount.did)
+        ? [desiredAccount]
+        : []
+    }
+
+    const tokenMode = getTokenWriteMode(
+      action,
+      desiredAccount.did,
+      latestAccount,
+    )
+    if (tokenMode === 'keep-latest') {
+      if (latestAccount.refreshJwt !== desiredAccount.refreshJwt) {
+        preservedNewerTokenDids.add(desiredAccount.did)
+      }
+      return [
+        {
+          ...desiredAccount,
+          accessJwt: latestAccount.accessJwt,
+          refreshJwt: latestAccount.refreshJwt,
+        },
+      ]
+    }
+    return [desiredAccount]
+  })
+
+  /* A stale tab must not discard an account another tab added meanwhile. */
+  for (const latestAccount of latest.accounts) {
+    if (
+      !desiredDids.has(latestAccount.did) &&
+      latestAccount.did !== removedDid
+    ) {
+      accounts.push(
+        action.type === 'logged-out-every-account'
+          ? {...latestAccount, accessJwt: undefined, refreshJwt: undefined}
+          : latestAccount,
+      )
+    }
+  }
+
+  const desiredCurrentDid = desired.currentAccount?.did
+  const latestCurrentDid = latest.currentAccount?.did
+  const currentDid = selectCurrentDid({
+    action,
+    desiredCurrentDid,
+    latestCurrentDid,
+    preservedNewerTokenDids,
+  })
+
+  return {
+    accounts,
+    currentAccount: accounts.find(account => account.did === currentDid),
+  }
+}
+
+function shouldAddMissingAccount(action: Action, did: string): boolean {
+  return (
+    action.type === 'switched-to-account' &&
+    action.newAccount.did === did &&
+    action.tokenUpdate === 'login'
+  )
+}
+
+function getTokenWriteMode(
+  action: Action,
+  did: string,
+  latestAccount: SessionAccount,
+): 'replace' | 'clear' | 'keep-latest' {
+  switch (action.type) {
+    case 'logged-out-current-account':
+      return action.accountDid === undefined || action.accountDid === did
+        ? 'clear'
+        : 'keep-latest'
+    case 'logged-out-every-account':
+      return 'clear'
+    case 'received-session-event':
+      if (action.sessionEvent === 'update' && action.refreshedAccount) {
+        /* A remote logout wins over an in-flight local refresh. */
+        return latestAccount.refreshJwt ? 'replace' : 'keep-latest'
+      }
+      if (action.sessionEvent === 'expired') {
+        return latestAccount.refreshJwt &&
+          latestAccount.refreshJwt !== action.expiredRefreshJwt
+          ? 'keep-latest'
+          : 'clear'
+      }
+      return 'keep-latest'
+    case 'switched-to-account':
+      if (action.newAccount.did !== did) return 'keep-latest'
+      if (action.tokenUpdate === 'login') return 'replace'
+      if (action.tokenUpdate === 'refresh' && latestAccount.refreshJwt) {
+        return 'replace'
+      }
+      return 'keep-latest'
+    default:
+      return 'keep-latest'
+  }
+}
+
+function selectCurrentDid({
+  action,
+  desiredCurrentDid,
+  latestCurrentDid,
+  preservedNewerTokenDids,
+}: {
+  action: Action
+  desiredCurrentDid: string | undefined
+  latestCurrentDid: string | undefined
+  preservedNewerTokenDids: Set<string>
+}): string | undefined {
+  switch (action.type) {
+    case 'switched-to-account':
+    case 'removed-account':
+    case 'logged-out-current-account':
+    case 'logged-out-every-account':
+      return desiredCurrentDid
+    case 'received-session-event':
+      if (
+        action.sessionEvent === 'expired' &&
+        action.accountDid !== undefined &&
+        preservedNewerTokenDids.has(action.accountDid)
+      ) {
+        return latestCurrentDid
+      }
+      return action.sessionEvent === 'expired'
+        ? desiredCurrentDid
+        : (latestCurrentDid ?? desiredCurrentDid)
+    default:
+      return latestCurrentDid ?? desiredCurrentDid
+  }
+}
+
 function createPublicBundleState(): BundleState {
   return {
     bundle: createPublicSessionBundle(),