Yeet

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

@@ -0,0 +1,250 @@
+import {describe, expect, it} from '@jest/globals'
+
+import {type PersistedAccount, type Schema} from '../schema'
+import {
+  applySessionUpdate,
+  getCredentialState,
+  type SessionCredentialMutation,
+} from '../session'
+
+const DID = 'did:plc:example123'
+
+function jwt({jti, issuedAt}: {jti: string; issuedAt: number}) {
+  const encode = (value: object) =>
+    Buffer.from(JSON.stringify(value)).toString('base64url')
+  return `${encode({alg: 'none'})}.${encode({jti, iat: issuedAt})}.signature`
+}
+
+function account({
+  refreshJwt,
+  accessJwt = `access-${refreshJwt}`,
+}: {
+  refreshJwt?: string
+  accessJwt?: string
+}): PersistedAccount {
+  return {
+    service: 'https://bsky.social/',
+    did: DID,
+    handle: 'alice.test',
+    refreshJwt,
+    accessJwt,
+  }
+}
+
+function session({
+  account: sessionAccount,
+  credentialStates,
+}: {
+  account?: PersistedAccount
+  credentialStates?: Schema['session']['credentialStates']
+}): Schema['session'] {
+  return {
+    accounts: sessionAccount ? [sessionAccount] : [],
+    currentAccount: sessionAccount,
+    credentialStates,
+  }
+}
+
+function update({
+  storedSession,
+  nextAccount,
+  mutation,
+}: {
+  storedSession: Schema['session']
+  nextAccount?: PersistedAccount
+  mutation?: SessionCredentialMutation
+}) {
+  return applySessionUpdate({
+    storedSession,
+    nextSession: session({
+      account: nextAccount,
+      credentialStates: storedSession.credentialStates,
+    }),
+    credentialMutations: mutation ? [mutation] : [],
+  })
+}
+
+describe('versioned persisted sessions', () => {
+  const refreshA = jwt({jti: 'A', issuedAt: 1})
+  const refreshB = jwt({jti: 'B', issuedAt: 2})
+  const refreshBAlias = jwt({jti: 'B', issuedAt: 3})
+  const refreshC = jwt({jti: 'C', issuedAt: 4})
+
+  it('advances the credential version when refresh moves to a new jti', () => {
+    const storedSession = session({account: account({refreshJwt: refreshA})})
+    const next = account({refreshJwt: refreshB})
+
+    const result = update({
+      storedSession,
+      nextAccount: next,
+      mutation: {
+        type: 'refresh',
+        accountDid: DID,
+        baseRefreshJwt: refreshA,
+        resultRefreshJwt: refreshB,
+      },
+    })
+
+    expect(result.accounts[0].refreshJwt).toBe(refreshB)
+    expect(getCredentialState({session: result, accountDid: DID})).toEqual({
+      credentialVersion: 1,
+      refreshJti: 'B',
+      status: 'active',
+    })
+  })
+
+  it('keeps one version for convergent aliases with the same jti', () => {
+    const generationB = update({
+      storedSession: session({account: account({refreshJwt: refreshA})}),
+      nextAccount: account({refreshJwt: refreshB}),
+      mutation: {
+        type: 'refresh',
+        accountDid: DID,
+        baseRefreshJwt: refreshA,
+        resultRefreshJwt: refreshB,
+      },
+    })
+
+    const result = update({
+      storedSession: generationB,
+      nextAccount: account({refreshJwt: refreshBAlias}),
+      mutation: {
+        type: 'refresh',
+        accountDid: DID,
+        baseRefreshJwt: refreshA,
+        resultRefreshJwt: refreshBAlias,
+      },
+    })
+
+    expect(result.accounts[0].refreshJwt).toBe(refreshB)
+    expect(
+      getCredentialState({session: result, accountDid: DID}).credentialVersion,
+    ).toBe(1)
+  })
+
+  it('does not let a stale refresh overwrite a newer generation', () => {
+    const storedSession = session({
+      account: account({refreshJwt: refreshB}),
+      credentialStates: {
+        [DID]: {credentialVersion: 8, refreshJti: 'B', status: 'active'},
+      },
+    })
+
+    const result = update({
+      storedSession,
+      nextAccount: account({refreshJwt: refreshC}),
+      mutation: {
+        type: 'refresh',
+        accountDid: DID,
+        baseRefreshJwt: refreshA,
+        resultRefreshJwt: refreshC,
+      },
+    })
+
+    expect(result.accounts[0].refreshJwt).toBe(refreshB)
+    expect(
+      getCredentialState({session: result, accountDid: DID}).credentialVersion,
+    ).toBe(8)
+  })
+
+  it('patches metadata without replacing authoritative credentials', () => {
+    const storedSession = session({
+      account: account({refreshJwt: refreshB}),
+      credentialStates: {
+        [DID]: {credentialVersion: 8, refreshJti: 'B', status: 'active'},
+      },
+    })
+    const staleAccount = {
+      ...account({refreshJwt: refreshA}),
+      handle: 'alice-renamed.test',
+    }
+
+    const result = update({storedSession, nextAccount: staleAccount})
+
+    expect(result.accounts[0].handle).toBe('alice-renamed.test')
+    expect(result.accounts[0].refreshJwt).toBe(refreshB)
+    expect(result.accounts[0].accessJwt).toBe(`access-${refreshB}`)
+  })
+
+  it('does not let a stale expiry clear a newer generation', () => {
+    const storedSession = session({
+      account: account({refreshJwt: refreshB}),
+      credentialStates: {
+        [DID]: {credentialVersion: 8, refreshJti: 'B', status: 'active'},
+      },
+    })
+
+    const result = update({
+      storedSession,
+      nextAccount: account({refreshJwt: undefined, accessJwt: undefined}),
+      mutation: {
+        type: 'expire',
+        accountDid: DID,
+        baseRefreshJwt: refreshA,
+      },
+    })
+
+    expect(result.accounts[0].refreshJwt).toBe(refreshB)
+    expect(getCredentialState({session: result, accountDid: DID}).status).toBe(
+      'active',
+    )
+  })
+
+  it('preserves logout and removal tombstones until an explicit login', () => {
+    const active = session({
+      account: account({refreshJwt: refreshB}),
+      credentialStates: {
+        [DID]: {credentialVersion: 8, refreshJti: 'B', status: 'active'},
+      },
+    })
+    const loggedOut = update({
+      storedSession: active,
+      nextAccount: account({refreshJwt: undefined, accessJwt: undefined}),
+      mutation: {type: 'logout', accountDid: DID},
+    })
+
+    const staleRefresh = update({
+      storedSession: loggedOut,
+      nextAccount: account({refreshJwt: refreshC}),
+      mutation: {
+        type: 'refresh',
+        accountDid: DID,
+        baseRefreshJwt: refreshB,
+        resultRefreshJwt: refreshC,
+      },
+    })
+    expect(staleRefresh.accounts[0].refreshJwt).toBeUndefined()
+    expect(
+      getCredentialState({session: staleRefresh, accountDid: DID}),
+    ).toEqual({credentialVersion: 9, status: 'logged-out'})
+
+    const removed = update({
+      storedSession: staleRefresh,
+      mutation: {type: 'remove', accountDid: DID},
+    })
+    const staleSnapshot = update({
+      storedSession: removed,
+      nextAccount: account({refreshJwt: refreshB}),
+    })
+    expect(staleSnapshot.accounts).toEqual([])
+    expect(
+      getCredentialState({session: staleSnapshot, accountDid: DID}),
+    ).toEqual({credentialVersion: 10, status: 'removed'})
+
+    const loggedIn = update({
+      storedSession: staleSnapshot,
+      nextAccount: account({refreshJwt: refreshC}),
+      mutation: {
+        type: 'login',
+        accountDid: DID,
+        resultRefreshJwt: refreshC,
+      },
+    })
+    expect(loggedIn.accounts[0].refreshJwt).toBe(refreshC)
+    expect(getCredentialState({session: loggedIn, accountDid: DID})).toEqual({
+      credentialVersion: 11,
+      refreshJti: 'C',
+      status: 'active',
+    })
+  })
+})

src/state/persisted/index.ts

@@ -8,15 +8,20 @@ import {
   tryStringify,
 } from '#/state/persisted/schema'
 import {device} from '#/storage'
+import {applySessionUpdate} from './session'
+import {runWithSessionCredentialLock} from './session-lock'
 import {type PersistedApi} from './types'
 import {normalizeData} from './util'
 
+export type {SessionCredentialMutation} from './session'
+export {runWithSessionCredentialLock} from './session-lock'
 export type {PersistedAccount, Schema} from '#/state/persisted/schema'
 export {defaults} from '#/state/persisted/schema'
 
 const BSKY_STORAGE = 'BSKY_STORAGE'
 
 let _state: Schema = defaults
+let pendingWrite: Promise<unknown> = Promise.resolve()
 
 export async function init() {
   const stored = await readFromStorage()
@@ -43,18 +48,43 @@ export function readLatest<K extends keyof Schema>(key: K): Schema[K] {
 }
 readLatest satisfies PersistedApi['readLatest']
 
-export async function write<K extends keyof Schema>(
+export function write<K extends keyof Schema>(
   key: K,
   value: Schema[K],
 ): Promise<void> {
-  _state = normalizeData({
+  return enqueueWrite(async () => {
-    ..._state,
+    const next = normalizeData({
-    [key]: value,
+      ..._state,
+      [key]: value,
+    })
+    await persistWithRetry(next)
+    _state = next
   })
-  await writeToStorage(_state)
 }
 write satisfies PersistedApi['write']
 
+export function updateSession({
+  nextSession,
+  credentialMutations,
+}: {
+  nextSession: Schema['session']
+  credentialMutations: import('./session').SessionCredentialMutation[]
+}): Promise<Schema['session']> {
+  return enqueueWrite(async () => {
+    const session = applySessionUpdate({
+      storedSession: _state.session,
+      nextSession,
+      credentialMutations,
+    })
+    const next = normalizeData({..._state, session})
+    await persistWithRetry(next)
+    _state = next
+    return session
+  })
+}
+updateSession satisfies PersistedApi['updateSession']
+runWithSessionCredentialLock satisfies PersistedApi['runWithSessionCredentialLock']
+
 export function onUpdate<K extends keyof Schema>(
   _key: K,
   _cb: (v: Schema[K]) => void,
@@ -73,16 +103,36 @@ export async function clearStorage() {
 }
 clearStorage satisfies PersistedApi['clearStorage']
 
+function enqueueWrite<T>(operation: () => Promise<T>): Promise<T> {
+  const result = pendingWrite.then(operation, operation)
+  pendingWrite = result.then(
+    () => undefined,
+    () => undefined,
+  )
+  return result
+}
+
+async function persistWithRetry(value: Schema) {
+  try {
+    await writeToStorage(value)
+  } catch {
+    /* Retry the latest complete snapshot while the process still owns it. */
+    await writeToStorage(value)
+  }
+}
+
 async function writeToStorage(value: Schema) {
   const rawData = tryStringify(value)
-  if (rawData) {
+  if (!rawData) {
-    try {
+    throw new Error('Failed to serialize persisted state')
-      await AsyncStorage.setItem(BSKY_STORAGE, rawData)
+  }
-    } catch (e) {
+  try {
-      logger.error(`persisted state: failed writing root state to storage`, {
+    await AsyncStorage.setItem(BSKY_STORAGE, rawData)
-        message: e,
+  } catch (e) {
-      })
+    logger.error(`persisted state: failed writing root state to storage`, {
-    }
+      message: e,
+    })
+    throw e
   }
 }
 

src/state/persisted/index.web.ts

@@ -8,9 +8,17 @@ import {
   tryParse,
   tryStringify,
 } from '#/state/persisted/schema'
+import {
+  applySessionUpdate,
+  getCredentialState,
+  type SessionCredentialMutation,
+} from './session'
+import {runWithSessionCredentialLock} from './session-lock'
 import {type PersistedApi} from './types'
 import {normalizeData} from './util'
 
+export type {SessionCredentialMutation} from './session'
+export {runWithSessionCredentialLock} from './session-lock'
 export type {PersistedAccount, Schema} from '#/state/persisted/schema'
 export {defaults} from '#/state/persisted/schema'
 
@@ -62,38 +70,73 @@ export function readLatest<K extends keyof Schema>(key: K): Schema[K] {
 }
 readLatest satisfies PersistedApi['readLatest']
 
-// eslint-disable-next-line @typescript-eslint/require-await
+export function write<K extends keyof Schema>(
-export async function write<K extends keyof Schema>(
   key: K,
   value: Schema[K],
 ): Promise<void> {
-  const next = readFromStorage()
+  return runWithSessionCredentialLock({
-  if (next) {
+    accountDids: [],
-    // The storage could have been updated by a different tab before this tab is notified.
+    operation: () => {
-    // Make sure this write is applied on top of the latest data in the storage as long as it's valid.
+      const next = readFromStorage()
-    _state = next
+      if (next) {
-    // Don't fire the update listeners yet to avoid a loop.
+        // The storage could have been updated by a different tab before this tab is notified.
-    // If there was a change, we'll receive the broadcast event soon enough which will do that.
+        // Make sure this write is applied on top of the latest data in the storage as long as it's valid.
-  }
+        _state = next
-  try {
+        // Don't fire the update listeners yet to avoid a loop.
-    if (JSON.stringify({v: _state[key]}) === JSON.stringify({v: value})) {
+        // If there was a change, we'll receive the broadcast event soon enough which will do that.
-      // Fast path for updates that are guaranteed to be noops.
+      }
-      // This is good mostly because it avoids useless broadcasts to other tabs.
+      try {
-      return
+        if (JSON.stringify({v: _state[key]}) === JSON.stringify({v: value})) {
-    }
+          // Fast path for updates that are guaranteed to be noops.
-  } catch (e) {
+          // This is good mostly because it avoids useless broadcasts to other tabs.
-    // Ignore and go through the normal path.
+          return
-  }
+        }
-  _state = normalizeData({
+      } catch (e) {
-    ..._state,
+        // Ignore and go through the normal path.
-    [key]: value,
+      }
+      const updated = normalizeData({
+        ..._state,
+        [key]: value,
+      })
+      writeToStorage(updated)
+      _state = updated
+      broadcastUpdate({key})
+    },
   })
-  writeToStorage(_state)
-  broadcast.postMessage({event: {type: UPDATE_EVENT, key}})
-  broadcast.postMessage({event: UPDATE_EVENT}) // Backcompat while upgrading
 }
 write satisfies PersistedApi['write']
 
+// eslint-disable-next-line @typescript-eslint/require-await
+export async function updateSession({
+  nextSession,
+  credentialMutations,
+}: {
+  nextSession: Schema['session']
+  credentialMutations: SessionCredentialMutation[]
+}): Promise<Schema['session']> {
+  const stored = readFromStorage() ?? _state
+  const session = applySessionUpdate({
+    storedSession: stored.session,
+    nextSession,
+    credentialMutations,
+  })
+  const updated = normalizeData({...stored, session})
+  writeToStorage(updated)
+  _state = updated
+
+  const invalidations = credentialMutations.map(mutation => ({
+    accountDid: mutation.accountDid,
+    credentialVersion: getCredentialState({
+      session,
+      accountDid: mutation.accountDid,
+    }).credentialVersion,
+  }))
+  broadcastUpdate({key: 'session', invalidations})
+  return session
+}
+updateSession satisfies PersistedApi['updateSession']
+runWithSessionCredentialLock satisfies PersistedApi['runWithSessionCredentialLock']
+
 export function onUpdate<K extends keyof Schema>(
   key: K,
   cb: (v: Schema[K]) => void,
@@ -108,13 +151,17 @@ export function onUpdate<K extends keyof Schema>(
 }
 onUpdate satisfies PersistedApi['onUpdate']
 
-// eslint-disable-next-line @typescript-eslint/require-await
+export function clearStorage(): Promise<void> {
-export async function clearStorage() {
+  return runWithSessionCredentialLock({
-  try {
+    accountDids: [],
-    localStorage.removeItem(BSKY_STORAGE)
+    operation: () => {
-  } catch (e: any) {
+      try {
-    // Expected on the web in private mode.
+        localStorage.removeItem(BSKY_STORAGE)
-  }
+      } catch (e: any) {
+        // Expected on the web in private mode.
+      }
+    },
+  })
 }
 clearStorage satisfies PersistedApi['clearStorage']
 
@@ -156,14 +203,40 @@ async function onBroadcastMessage({data}: MessageEvent) {
   }
 }
 
+function broadcastUpdate({
+  key,
+  invalidations = [],
+}: {
+  key: keyof Schema
+  invalidations?: {accountDid: string; credentialVersion: number}[]
+}) {
+  if (invalidations.length > 0) {
+    for (const invalidation of invalidations) {
+      broadcast.postMessage({
+        event: {type: UPDATE_EVENT, key, ...invalidation},
+      })
+    }
+  } else {
+    broadcast.postMessage({event: {type: UPDATE_EVENT, key}})
+  }
+  broadcast.postMessage({event: UPDATE_EVENT}) // Backcompat while upgrading
+}
+
 function writeToStorage(value: Schema) {
   const rawData = tryStringify(value)
-  if (rawData) {
+  if (!rawData) {
-    try {
+    throw new Error('Failed to serialize persisted state')
-      localStorage.setItem(BSKY_STORAGE, rawData)
+  }
-    } catch (e) {
+  try {
-      // Expected on the web in private mode.
+    localStorage.setItem(BSKY_STORAGE, rawData)
+    if (localStorage.getItem(BSKY_STORAGE) !== rawData) {
+      throw new Error('Failed to verify persisted state')
     }
+  } catch (error) {
+    logger.error('persisted state: failed writing root state to storage', {
+      message: error,
+    })
+    throw error
   }
 }
 

src/state/persisted/schema.ts

@@ -57,12 +57,21 @@ const currentAccountSchema = accountSchema.extend({
 })
 export type PersistedCurrentAccount = z.infer<typeof currentAccountSchema>
 
+const credentialStateSchema = z.object({
+  credentialVersion: z.number().int().nonnegative(),
+  refreshJti: z.string().optional(),
+  status: z.enum(['active', 'logged-out', 'removed']),
+})
+export type PersistedCredentialState = z.infer<typeof credentialStateSchema>
+
 const schema = z.object({
   colorMode: z.enum(['system', 'light', 'dark']),
   darkTheme: z.enum(['dim', 'dark']).optional(),
   session: z.object({
     accounts: z.array(accountSchema),
     currentAccount: currentAccountSchema.optional(),
+    /** Per-account credential ordering, including logout and removal tombstones. */
+    credentialStates: z.record(z.string(), credentialStateSchema).optional(),
   }),
   reminders: z.object({
     lastEmailConfirm: z.string().optional(),

src/state/persisted/session-lock.ts

@@ -0,0 +1,18 @@
+export function runWithSessionCredentialLock<T>({
+  accountDids,
+  operation,
+}: {
+  accountDids: string[]
+  operation: () => T | Promise<T>
+}): Promise<T> {
+  void accountDids
+  try {
+    return Promise.resolve(operation())
+  } catch (error) {
+    return Promise.reject(
+      error instanceof Error
+        ? error
+        : new Error('Session credential operation failed', {cause: error}),
+    )
+  }
+}

src/state/persisted/session-lock.web.ts

@@ -0,0 +1,37 @@
+const PERSISTED_STORAGE_LOCK = 'bsky-persisted-storage'
+const SESSION_LOCK_PREFIX = 'bsky-session:'
+
+export function runWithSessionCredentialLock<T>({
+  accountDids,
+  operation,
+}: {
+  accountDids: string[]
+  operation: () => T | Promise<T>
+}): Promise<T> {
+  const lockManager = navigator.locks
+  if (!lockManager) {
+    try {
+      return Promise.resolve(operation())
+    } catch (error) {
+      return Promise.reject(
+        error instanceof Error
+          ? error
+          : new Error('Session credential operation failed', {cause: error}),
+      )
+    }
+  }
+
+  const accountLockNames = [...new Set(accountDids)]
+    .sort()
+    .map(did => `${SESSION_LOCK_PREFIX}${did}`)
+  /* All values share one localStorage blob, so every write also takes its lock. */
+  const lockNames = [PERSISTED_STORAGE_LOCK, ...accountLockNames]
+
+  const run = (index: number): Promise<T> => {
+    const lockName = lockNames[index]
+    if (!lockName) return Promise.resolve(operation())
+    return lockManager.request(lockName, () => run(index + 1))
+  }
+
+  return run(0)
+}

src/state/persisted/session.ts

@@ -0,0 +1,294 @@
+import {jwtDecode} from 'jwt-decode'
+
+import {
+  type PersistedAccount,
+  type PersistedCredentialState,
+  type Schema,
+} from './schema'
+
+export type SessionCredentialMutation =
+  | {
+      type: 'refresh'
+      accountDid: string
+      baseRefreshJwt: string | undefined
+      resultRefreshJwt: string | undefined
+    }
+  | {
+      type: 'expire'
+      accountDid: string
+      baseRefreshJwt: string | undefined
+    }
+  | {
+      type: 'login'
+      accountDid: string
+      resultRefreshJwt: string | undefined
+    }
+  | {
+      type: 'logout'
+      accountDid: string
+    }
+  | {
+      type: 'remove'
+      accountDid: string
+    }
+
+const CREDENTIAL_FIELDS = new Set<keyof PersistedAccount>([
+  'accessJwt',
+  'refreshJwt',
+])
+
+/** Read the server-side refresh-token generation, tolerating legacy bad data. */
+export function getRefreshJti({
+  refreshJwt,
+}: {
+  refreshJwt: string | undefined
+}): string | undefined {
+  if (!refreshJwt) return undefined
+  try {
+    const decoded = jwtDecode(refreshJwt)
+    if (typeof decoded.jti === 'string') return decoded.jti
+  } catch {}
+  /* A malformed legacy token still needs stable local identity. */
+  return refreshJwt
+}
+
+export function getCredentialState({
+  session,
+  accountDid,
+}: {
+  session: Schema['session']
+  accountDid: string
+}): PersistedCredentialState {
+  const stored = session.credentialStates?.[accountDid]
+  if (stored) return stored
+  const account = session.accounts.find(
+    candidate => candidate.did === accountDid,
+  )
+  return {
+    credentialVersion: 0,
+    refreshJti: getRefreshJti({refreshJwt: account?.refreshJwt}),
+    status: account?.refreshJwt ? 'active' : 'logged-out',
+  }
+}
+
+/**
+ * Apply a session snapshot to the latest persisted session without allowing its
+ * credential fields to overwrite newer credential generations.
+ */
+export function applySessionUpdate({
+  storedSession,
+  nextSession,
+  credentialMutations,
+}: {
+  storedSession: Schema['session']
+  nextSession: Schema['session']
+  credentialMutations: SessionCredentialMutation[]
+}): Schema['session'] {
+  const incomingByDid = new Map(
+    nextSession.accounts.map(account => [account.did, account]),
+  )
+  const storedByDid = new Map(
+    storedSession.accounts.map(account => [account.did, account]),
+  )
+
+  /* Incoming metadata wins, but stored credentials remain authoritative. */
+  const accounts = nextSession.accounts.map(incoming => {
+    const stored = storedByDid.get(incoming.did)
+    return stored
+      ? mergeAccountMetadata({storedAccount: stored, incomingAccount: incoming})
+      : incoming
+  })
+  for (const stored of storedSession.accounts) {
+    if (!incomingByDid.has(stored.did)) accounts.push(stored)
+  }
+
+  const credentialStates: Record<string, PersistedCredentialState> = {
+    ...storedSession.credentialStates,
+  }
+  for (const account of storedSession.accounts) {
+    credentialStates[account.did] = getCredentialState({
+      session: storedSession,
+      accountDid: account.did,
+    })
+  }
+
+  const result: Schema['session'] = {
+    accounts,
+    currentAccount: nextSession.currentAccount,
+    credentialStates,
+  }
+
+  for (const mutation of credentialMutations) {
+    applyCredentialMutation({session: result, nextSession, mutation})
+  }
+
+  result.accounts = result.accounts
+    .filter(
+      account =>
+        getCredentialState({session: result, accountDid: account.did})
+          .status !== 'removed',
+    )
+    .map(account => {
+      const credentialState = getCredentialState({
+        session: result,
+        accountDid: account.did,
+      })
+      return credentialState.status === 'logged-out'
+        ? {...account, accessJwt: undefined, refreshJwt: undefined}
+        : account
+    })
+
+  const currentDid = result.currentAccount?.did
+  const currentAccount = currentDid
+    ? result.accounts.find(account => account.did === currentDid)
+    : undefined
+  const currentCredentialState = currentDid
+    ? getCredentialState({session: result, accountDid: currentDid})
+    : undefined
+  result.currentAccount =
+    currentAccount && currentCredentialState?.status === 'active'
+      ? currentAccount
+      : undefined
+
+  return result
+}
+
+function applyCredentialMutation({
+  session,
+  nextSession,
+  mutation,
+}: {
+  session: Schema['session']
+  nextSession: Schema['session']
+  mutation: SessionCredentialMutation
+}) {
+  const previousState = getCredentialState({
+    session,
+    accountDid: mutation.accountDid,
+  })
+  const nextAccount = nextSession.accounts.find(
+    account => account.did === mutation.accountDid,
+  )
+
+  switch (mutation.type) {
+    case 'refresh': {
+      const baseRefreshJti = getRefreshJti({
+        refreshJwt: mutation.baseRefreshJwt,
+      })
+      const resultRefreshJti = getRefreshJti({
+        refreshJwt: mutation.resultRefreshJwt,
+      })
+      if (
+        previousState.status !== 'active' ||
+        previousState.refreshJti !== baseRefreshJti
+      ) {
+        /* The stored generation already advanced or became a tombstone. */
+        return
+      }
+      if (!nextAccount || !resultRefreshJti) return
+      replaceAccount({session, account: nextAccount})
+      session.credentialStates![mutation.accountDid] = {
+        credentialVersion:
+          resultRefreshJti === previousState.refreshJti
+            ? previousState.credentialVersion
+            : previousState.credentialVersion + 1,
+        refreshJti: resultRefreshJti,
+        status: 'active',
+      }
+      return
+    }
+    case 'expire': {
+      const baseRefreshJti = getRefreshJti({
+        refreshJwt: mutation.baseRefreshJwt,
+      })
+      if (
+        previousState.status !== 'active' ||
+        previousState.refreshJti !== baseRefreshJti
+      ) {
+        return
+      }
+      clearAccountCredentials({session, accountDid: mutation.accountDid})
+      session.credentialStates![mutation.accountDid] = {
+        credentialVersion: previousState.credentialVersion + 1,
+        status: 'logged-out',
+      }
+      return
+    }
+    case 'login': {
+      const resultRefreshJti = getRefreshJti({
+        refreshJwt: mutation.resultRefreshJwt,
+      })
+      if (!nextAccount || !resultRefreshJti) return
+      replaceAccount({session, account: nextAccount})
+      session.credentialStates![mutation.accountDid] = {
+        credentialVersion: previousState.credentialVersion + 1,
+        refreshJti: resultRefreshJti,
+        status: 'active',
+      }
+      return
+    }
+    case 'logout': {
+      clearAccountCredentials({session, accountDid: mutation.accountDid})
+      session.credentialStates![mutation.accountDid] = {
+        credentialVersion: previousState.credentialVersion + 1,
+        status: 'logged-out',
+      }
+      return
+    }
+    case 'remove': {
+      session.accounts = session.accounts.filter(
+        account => account.did !== mutation.accountDid,
+      )
+      session.credentialStates![mutation.accountDid] = {
+        credentialVersion: previousState.credentialVersion + 1,
+        status: 'removed',
+      }
+      return
+    }
+  }
+}
+
+function mergeAccountMetadata({
+  storedAccount,
+  incomingAccount,
+}: {
+  storedAccount: PersistedAccount
+  incomingAccount: PersistedAccount
+}): PersistedAccount {
+  const merged = {...storedAccount}
+  for (const key of Object.keys(
+    incomingAccount,
+  ) as (keyof PersistedAccount)[]) {
+    if (!CREDENTIAL_FIELDS.has(key)) {
+      Object.assign(merged, {[key]: incomingAccount[key]})
+    }
+  }
+  return merged
+}
+
+function replaceAccount({
+  session,
+  account,
+}: {
+  session: Schema['session']
+  account: PersistedAccount
+}) {
+  session.accounts = [
+    account,
+    ...session.accounts.filter(candidate => candidate.did !== account.did),
+  ]
+}
+
+function clearAccountCredentials({
+  session,
+  accountDid,
+}: {
+  session: Schema['session']
+  accountDid: string
+}) {
+  session.accounts = session.accounts.map(account =>
+    account.did === accountDid
+      ? {...account, accessJwt: undefined, refreshJwt: undefined}
+      : account,
+  )
+}

src/state/persisted/types.ts

@@ -1,4 +1,5 @@
 import {type Schema} from './schema'
+import {type SessionCredentialMutation} from './session'
 
 export type PersistedApi = {
   init(): Promise<void>
@@ -13,6 +14,14 @@ export type PersistedApi = {
    */
   readLatest<K extends keyof Schema>(key: K): Schema[K]
   write<K extends keyof Schema>(key: K, value: Schema[K]): Promise<void>
+  updateSession(params: {
+    nextSession: Schema['session']
+    credentialMutations: SessionCredentialMutation[]
+  }): Promise<Schema['session']>
+  runWithSessionCredentialLock<T>(params: {
+    accountDids: string[]
+    operation: () => T | Promise<T>
+  }): Promise<T>
   onUpdate<K extends keyof Schema>(
     key: K,
     cb: (v: Schema[K]) => void,

src/state/session/__tests__/provider-abort-test.tsx

@@ -14,6 +14,10 @@ jest.mock('#/state/persisted', () => {
     defaults,
     get: (key: keyof typeof defaults) => defaults[key],
     write: () => Promise.resolve(),
+    updateSession: ({nextSession}: {nextSession: typeof defaults.session}) =>
+      Promise.resolve(nextSession),
+    runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) =>
+      Promise.resolve(operation()),
     readLatest: (key: keyof typeof defaults) => defaults[key],
     onUpdate: () => () => {},
   }

src/state/session/__tests__/provider-clients-test.tsx

@@ -18,6 +18,10 @@ jest.mock('#/state/persisted', () => {
     defaults,
     get: (key: keyof typeof defaults) => defaults[key],
     write: () => Promise.resolve(),
+    updateSession: ({nextSession}: {nextSession: typeof defaults.session}) =>
+      Promise.resolve(nextSession),
+    runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) =>
+      Promise.resolve(operation()),
     readLatest: (key: keyof typeof defaults) => defaults[key],
     onUpdate: () => () => {},
   }

src/state/session/__tests__/provider-refresh-session-test.tsx

@@ -17,6 +17,10 @@ jest.mock('#/state/persisted', () => {
     defaults,
     get: (key: keyof typeof defaults) => defaults[key],
     write: () => Promise.resolve(),
+    updateSession: ({nextSession}: {nextSession: typeof defaults.session}) =>
+      Promise.resolve(nextSession),
+    runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) =>
+      Promise.resolve(operation()),
     readLatest: (key: keyof typeof defaults) => defaults[key],
     onUpdate: () => () => {},
   }

src/state/session/__tests__/provider-session-events-test.tsx

@@ -39,6 +39,27 @@ jest.mock('#/state/persisted', () => {
         ? mockPersisted.latest
         : defaults[key as keyof typeof defaults],
     write: () => Promise.resolve(),
+    updateSession: ({
+      nextSession,
+      credentialMutations,
+    }: {
+      nextSession: Schema['session']
+      credentialMutations: import('#/state/persisted/session').SessionCredentialMutation[]
+    }) => {
+      const {
+        applySessionUpdate,
+      }: typeof import('#/state/persisted/session') = require('#/state/persisted/session')
+      const committed = applySessionUpdate({
+        storedSession: mockPersisted.latest,
+        nextSession,
+        credentialMutations,
+      })
+      mockPersisted.session = committed
+      mockPersisted.latest = committed
+      return Promise.resolve(committed)
+    },
+    runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) =>
+      Promise.resolve(operation()),
     onUpdate: (_key: string, cb: (value: Schema['session']) => void) => {
       mockPersistedListeners.push(cb)
       return () => {}
@@ -308,7 +329,7 @@ describe('expiry rescue', () => {
     mockPersisted.latest = {accounts: [fresher], currentAccount: fresher}
 
     act(() => {
-      onSessionChange(
+      void onSessionChange(
         bundle as unknown as SessionBundle,
         DID,
         'expired',
@@ -337,7 +358,7 @@ describe('expiry rescue', () => {
     mockPersisted.latest = {accounts: [account], currentAccount: account}
 
     act(() => {
-      onSessionChange(
+      void onSessionChange(
         bundle as unknown as SessionBundle,
         DID,
         'expired',
@@ -365,7 +386,7 @@ describe('expiry rescue', () => {
 
     /* generation 1 dies and is rescued onto generation 2 */
     act(() => {
-      onSessionChange(
+      void onSessionChange(
         bundle as unknown as SessionBundle,
         DID,
         'expired',
@@ -382,7 +403,7 @@ describe('expiry rescue', () => {
      */
     mockPersisted.latest = {accounts: [account], currentAccount: account}
     act(() => {
-      onSessionChange(
+      void onSessionChange(
         rescued as unknown as SessionBundle,
         DID,
         'expired',
@@ -416,7 +437,7 @@ describe('refresh persistence', () => {
     )
 
     act(() => {
-      onSessionChange(
+      void onSessionChange(
         bundle as unknown as SessionBundle,
         DID,
         'update',
@@ -437,7 +458,7 @@ describe('refresh persistence', () => {
     )
 
     act(() => {
-      onSessionChange(
+      void onSessionChange(
         bundle as unknown as SessionBundle,
         DID,
         'update',
@@ -447,10 +468,81 @@ describe('refresh persistence', () => {
 
     expect(currentAccount()?.pdsUrl).toBe(`${DIDDOC_PDS_HOST}/`)
   })
+
+  it('adopts the committed generation when a stale refresh succeeds', async () => {
+    const account = makeAccount()
+    const bundle = makeBundle(account)
+    const {onSessionChange, currentAccount} = await renderLoggedIn(
+      account,
+      bundle,
+    )
+    const committed = makeAccount({
+      accessJwt: 'access-jwt-2',
+      refreshJwt: 'refresh-jwt-2',
+    })
+    mockPersisted.latest = mockPersisted.session = {
+      accounts: [committed],
+      currentAccount: committed,
+      credentialStates: {
+        [DID]: {
+          credentialVersion: 8,
+          refreshJti: 'refresh-jwt-2',
+          status: 'active',
+        },
+      },
+    }
+
+    await act(async () => {
+      await onSessionChange(
+        bundle as unknown as SessionBundle,
+        DID,
+        'update',
+        refreshedData('refresh-jwt-3'),
+      )
+    })
+
+    expect(mockRebuilds.at(-1)?.account.refreshJwt).toBe('refresh-jwt-2')
+    expect(currentAccount()?.refreshJwt).toBe('refresh-jwt-2')
+  })
+
+  it('preserves a logout tombstone against a delayed successful refresh', async () => {
+    const account = makeAccount()
+    const bundle = makeBundle(account)
+    const {onSessionChange, currentAccount, hasSession} = await renderLoggedIn(
+      account,
+      bundle,
+    )
+    const loggedOut = makeAccount({
+      accessJwt: undefined,
+      refreshJwt: undefined,
+    })
+    mockPersisted.latest = mockPersisted.session = {
+      accounts: [loggedOut],
+      currentAccount: undefined,
+      credentialStates: {
+        [DID]: {credentialVersion: 8, status: 'logged-out'},
+      },
+    }
+
+    await act(async () => {
+      await onSessionChange(
+        bundle as unknown as SessionBundle,
+        DID,
+        'update',
+        refreshedData('refresh-jwt-2'),
+      )
+    })
+
+    expect(hasSession()).toBe(false)
+    expect(currentAccount()).toBeUndefined()
+    expect(mockPersisted.latest.accounts[0].refreshJwt).toBeUndefined()
+  })
 })
 
 /** Deliver a cross-tab `persisted` update to the provider's newest listener. */
 function emitSynced(session: Schema['session']) {
+  mockPersisted.session = session
+  mockPersisted.latest = session
   mockPersistedListeners[mockPersistedListeners.length - 1](session)
 }
 

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

@@ -97,10 +97,12 @@ import {
   disposeBundle,
   finishPreparation,
   makeSessionHooks,
+  type OnSessionChange,
   registerBundleKillSwitch,
   sessionAccountToSessionData,
   type SessionBundle,
   sessionDataToSessionAccount,
+  takeSessionChangeError,
 } from '../session-core'
 import {
   asFetch,
@@ -367,7 +369,7 @@ describe('createSessionBundleFromStoredAccount', () => {
   it('builds three clients over one session', async () => {
     const result = createSessionBundleFromStoredAccount(
       makeAccount(),
-      jest.fn(),
+      jest.fn<OnSessionChange>(),
     )!
 
     /* every client reads its identity straight through the shared session */
@@ -384,7 +386,7 @@ describe('createSessionBundleFromStoredAccount', () => {
   })
 
   it('disposes a bundle rejected by the activation guard', async () => {
-    const onSessionChange = jest.fn()
+    const onSessionChange = jest.fn<OnSessionChange>()
     let rejectedBundle: SessionBundle | undefined
     const result = createSessionBundleFromStoredAccount(
       makeAccount(),
@@ -607,7 +609,7 @@ describe('session-hook payload threading (pre-commit ordering)', () => {
 describe('disposeBundle kill-switch', () => {
   it('the injected fetch throws after disposeBundle', () => {
     const hooks = makeSessionHooks({
-      onSessionChange: jest.fn(),
+      onSessionChange: jest.fn<OnSessionChange>(),
       getBundle: () => ({}) as SessionBundle,
       getDid: () => DID,
     })
@@ -850,7 +852,7 @@ describe('factory account snapshot after preparation', () => {
     await withFreshFactory(asFetch(fetchMock), async core => {
       const {account, bundle} = await core.createSessionBundleAndResume(
         makeAccount({accessJwt: 'valid-access-jwt'}),
-        jest.fn(),
+        jest.fn<OnSessionChange>(),
       )
       /* the moderation prep step ran the refresh */
       expect(mockConfigureModerationForAccount).toHaveBeenCalledTimes(1)
@@ -873,7 +875,7 @@ describe('factory account snapshot after preparation', () => {
     await withFreshFactory(asFetch(fetchMock), async core => {
       const {account} = await core.createSessionBundleAndResume(
         makeAccount({accessJwt: 'valid-access-jwt'}),
-        jest.fn(),
+        jest.fn<OnSessionChange>(),
       )
       expect(account.accessJwt).toBe('valid-access-jwt')
       expect(account.refreshJwt).toBe('refresh-jwt')
@@ -921,7 +923,7 @@ describe('a session destroyed or rejected during preparation', () => {
       await expect(
         core.createSessionBundleAndResume(
           makeAccount({accessJwt: 'valid-access-jwt'}),
-          jest.fn(),
+          jest.fn<OnSessionChange>(),
         ),
       ).rejects.toThrow('Session was revoked while it was being prepared')
 
@@ -946,7 +948,7 @@ describe('a session destroyed or rejected during preparation', () => {
       await expect(
         core.createSessionBundleAndResume(
           makeAccount({accessJwt: 'valid-access-jwt'}),
-          jest.fn(),
+          jest.fn<OnSessionChange>(),
         ),
       ).rejects.toThrow('prefetch blew up')
 
@@ -959,7 +961,7 @@ describe('a session destroyed or rejected during preparation', () => {
 
   it('finishPreparation disposes and rethrows without running the snapshot', async () => {
     const hooks = makeSessionHooks({
-      onSessionChange: jest.fn(),
+      onSessionChange: jest.fn<OnSessionChange>(),
       getBundle: () => bundle,
       getDid: () => DID,
     })
@@ -1024,6 +1026,9 @@ describe('a throwing onSessionChange does not brick the session', () => {
     expect(mockLoggerError.mock.calls[0][1]).toEqual({
       message: "session: onSessionChange threw for a 'update' event",
     })
+    expect(takeSessionChangeError({bundle})).toEqual(
+      new Error('reducer side effect exploded'),
+    )
 
     /* the session still committed the rotation, and can still refresh again */
     expect(session.session.accessJwt).toBe('access-jwt-2')

src/state/session/index.tsx

@@ -13,6 +13,7 @@ import {type Client} from '@atproto/lex'
 import {type SessionData} from '@atproto/lex-password-session'
 
 import * as persisted from '#/state/persisted'
+import {type Schema, type SessionCredentialMutation} from '#/state/persisted'
 import {useCloseAllActiveElements} from '#/state/util'
 import {useGlobalDialogsControlContext} from '#/components/dialogs/Context'
 import {AnalyticsContext, useAnalyticsBase, utils} from '#/analytics'
@@ -29,9 +30,11 @@ import {
   createSessionBundleAndResume,
   createSessionBundleFromStoredAccount,
   disposeBundle,
+  type OnSessionChange,
   type PublicSessionBundle,
   type SessionBundle,
   sessionDataToSessionAccount,
+  takeSessionChangeError,
 } from './session-core'
 export {isSignupQueued} from './session-data'
 import {
@@ -102,10 +105,15 @@ class SessionStore {
     }
   }
 
-  dispatch = (action: Action) => {
+  dispatch = (
+    action: Action,
+    credentialMutations: SessionCredentialMutation[] = [],
+  ): Promise<Schema['session'] | undefined> => {
     const nextState = reducer(this.state, action)
     this.state = nextState
-    // Persist synchronously without waiting for the React render cycle.
+    let persistence: Promise<Schema['session'] | undefined> =
+      Promise.resolve(undefined)
+    // Schedule persistence immediately without waiting for the React render cycle.
     if (nextState.needsPersist) {
       nextState.needsPersist = false
       const persistedData = {
@@ -118,9 +126,13 @@ class SessionStore {
         type: 'persisted:broadcast',
         data: redactPersistedSession(persistedData),
       })
-      void persisted.write('session', persistedData)
+      persistence = persisted.updateSession({
+        nextSession: persistedData,
+        credentialMutations,
+      })
     }
     this.listeners.forEach(listener => listener())
+    return persistence
   }
 }
 
@@ -140,15 +152,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
    * insertion effect below, which commits well before any session hook can
    * fire: hooks are armed only after an asynchronous session factory resolves.
    */
-  const onSessionChangeRef = useRef<
+  const onSessionChangeRef = useRef<OnSessionChange | null>(null)
-    | ((
-        bundle: SessionBundle,
-        accountDid: string,
-        sessionEvent: AtpSessionEvent,
-        sessionData?: SessionData,
-      ) => void)
-    | null
-  >(null)
 
   const onSessionChange = useCallback(
     (
@@ -157,106 +161,169 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       sessionEvent: AtpSessionEvent,
       sessionData?: SessionData,
     ) => {
-      /*
+      const baseRefreshJwt =
-       * Only the live bundle may reset the expiry-rescue bookkeeping: a stale
+        sessionEvent === 'update' && !bundle.session.destroyed
-       * bundle's late update would otherwise clear the failed-generation set
+          ? bundle.session.session.refreshJwt
-       * that bounds the rescue loop. (Its dispatch below is separately dropped
+          : sessionData?.refreshJwt
-       * by the reducer's identity guard.)
-       */
-      if (
-        sessionEvent === 'update' &&
-        sessionData &&
-        (store.getState().currentBundleState.bundle as unknown as
-          SessionBundle | PublicSessionBundle) === bundle
-      ) {
-        failedExpiryTokensRef.current.get(accountDid)?.clear()
-      }
 
-      /*
+      return persisted.runWithSessionCredentialLock({
-       * PasswordSession invokes its hooks before updating its live getter. Use
+        accountDids: [accountDid],
-       * the delivered payload so a refresh persists the newly rotated tokens.
+        operation: async () => {
-       *
+          /*
-       * A refresh payload carries no didDoc unless the server sends one, so the
+           * Only the live bundle may reset the expiry-rescue bookkeeping: a stale
-       * stored account's `pdsUrl` is threaded in as the fallback. Without it the
+           * bundle's late update would otherwise clear the failed-generation set
-       * refresh would persist `pdsUrl: undefined` and the next cold start would
+           * that bounds the rescue loop. (Its dispatch below is separately dropped
-       * route pre-refresh requests to the entryway instead of the PDS.
+           * by the reducer's identity guard.)
-       */
+           */
-      const refreshedAccount =
+          if (
-        sessionEvent === 'update' && sessionData
+            sessionEvent === 'update' &&
-          ? sessionDataToSessionAccount(
+            sessionData &&
-              sessionData,
+            (store.getState().currentBundleState.bundle as unknown as
-              sessionData.service,
+              SessionBundle | PublicSessionBundle) === bundle
-              store.getState().accounts.find(a => a.did === accountDid)?.pdsUrl,
+          ) {
-            )
+            failedExpiryTokensRef.current.get(accountDid)?.clear()
-          : undefined
-
-      /*
-       * A stale tab may expire a token after another tab has already rotated it.
-       * Prefer a newer persisted or reducer generation over logging every tab
-       * out. Failed generations are recorded and bounded to guarantee that a
-       * repeatedly expiring session eventually falls through to logout.
-       */
-      if (sessionEvent === 'expired') {
-        const current = store.getState()
-        const currentBundle = current.currentBundleState.bundle as unknown as
-          SessionBundle | PublicSessionBundle
-        const dyingRefreshJwt = sessionData?.refreshJwt
-        // Stale bundle events are handled by the reducer's identity guard.
-        if (
-          currentBundle === bundle &&
-          current.currentBundleState.did === accountDid &&
-          dyingRefreshJwt
-        ) {
-          let failedSet = failedExpiryTokensRef.current.get(accountDid)
-          if (!failedSet) {
-            failedSet = new Set()
-            failedExpiryTokensRef.current.set(accountDid, failedSet)
           }
-          failedSet.add(dyingRefreshJwt)
 
-          const persistedCandidate = persisted
+          /*
-            .readLatest('session')
+           * PasswordSession invokes its hooks before updating its live getter. Use
-            .accounts.find(a => a.did === accountDid)
+           * the delivered payload so a refresh persists the newly rotated tokens.
-          const reducerCandidate = current.accounts.find(
+           *
-            a => a.did === accountDid,
+           * A refresh payload carries no didDoc unless the server sends one, so the
-          )
+           * stored account's `pdsUrl` is threaded in as the fallback. Without it the
-          const candidate = pickExpiryRescueCandidate({
+           * refresh would persist `pdsUrl: undefined` and the next cold start would
-            dyingRefreshJwt,
+           * route pre-refresh requests to the entryway instead of the PDS.
-            candidates: [persistedCandidate, reducerCandidate],
+           */
-            failedRefreshJwts: failedSet,
+          const refreshedAccount =
-          })
+            sessionEvent === 'update' && sessionData
+              ? sessionDataToSessionAccount(
+                  sessionData,
+                  sessionData.service,
+                  store.getState().accounts.find(a => a.did === accountDid)
+                    ?.pdsUrl,
+                )
+              : undefined
 
-          if (candidate) {
+          /*
-            const rebuilt = createSessionBundleFromStoredAccount(
+           * A stale tab may expire a token after another tab has already rotated it.
-              candidate,
+           * Prefer a newer persisted or reducer generation over logging every tab
-              onSessionChangeRef.current!,
+           * out. Failed generations are recorded and bounded to guarantee that a
-            )
+           * repeatedly expiring session eventually falls through to logout.
-            if (rebuilt) {
+           */
-              store.dispatch({
+          if (sessionEvent === 'expired') {
-                type: 'replaced-current-bundle',
+            const current = store.getState()
-                newBundle: rebuilt.bundle,
+            const currentBundle = current.currentBundleState
-                newAccount: rebuilt.account,
+              .bundle as unknown as SessionBundle | PublicSessionBundle
+            const dyingRefreshJwt = sessionData?.refreshJwt
+            // Stale bundle events are handled by the reducer's identity guard.
+            if (
+              currentBundle === bundle &&
+              current.currentBundleState.did === accountDid &&
+              dyingRefreshJwt
+            ) {
+              let failedSet = failedExpiryTokensRef.current.get(accountDid)
+              if (!failedSet) {
+                failedSet = new Set()
+                failedExpiryTokensRef.current.set(accountDid, failedSet)
+              }
+              failedSet.add(dyingRefreshJwt)
+
+              const persistedCandidate = persisted
+                .readLatest('session')
+                .accounts.find(a => a.did === accountDid)
+              const reducerCandidate = current.accounts.find(
+                a => a.did === accountDid,
+              )
+              const candidate = pickExpiryRescueCandidate({
+                dyingRefreshJwt,
+                candidates: [persistedCandidate, reducerCandidate],
+                failedRefreshJwts: failedSet,
               })
-              return
+
+              if (candidate) {
+                const rebuilt = createSessionBundleFromStoredAccount(
+                  candidate,
+                  onSessionChangeRef.current!,
+                )
+                if (rebuilt) {
+                  await store.dispatch({
+                    type: 'replaced-current-bundle',
+                    newBundle: rebuilt.bundle,
+                    newAccount: rebuilt.account,
+                  })
+                  return
+                }
+              }
             }
           }
-        }
-      }
 
-      // Only the current bundle may report that its session was dropped.
+          // Only the current bundle may report that its session was dropped.
-      if (
+          if (
-        sessionEvent === 'expired' &&
+            sessionEvent === 'expired' &&
-        store.getState().currentBundleState.bundle === bundle
+            store.getState().currentBundleState.bundle === bundle
-      ) {
+          ) {
-        emitSessionDropped()
+            emitSessionDropped()
-      }
+          }
-      // Bundle identity prevents stale sessions from changing the active account.
+
-      store.dispatch({
+          const credentialMutations: SessionCredentialMutation[] = []
-        type: 'received-session-event',
+          if (sessionEvent === 'update' && refreshedAccount) {
-        bundle,
+            credentialMutations.push({
-        refreshedAccount,
+              type: 'refresh',
-        accountDid,
+              accountDid,
-        sessionEvent,
+              baseRefreshJwt,
+              resultRefreshJwt: refreshedAccount.refreshJwt,
+            })
+          } else if (sessionEvent === 'expired') {
+            credentialMutations.push({
+              type: 'expire',
+              accountDid,
+              baseRefreshJwt,
+            })
+          }
+
+          // Bundle identity prevents stale sessions from changing the active account.
+          const committedSession = await store.dispatch(
+            {
+              type: 'received-session-event',
+              bundle,
+              refreshedAccount,
+              accountDid,
+              sessionEvent,
+            },
+            credentialMutations,
+          )
+
+          if (sessionEvent === 'update' && committedSession) {
+            const committedAccount = committedSession.accounts.find(
+              account => account.did === accountDid,
+            )
+            if (
+              committedAccount?.refreshJwt &&
+              committedAccount.refreshJwt !== refreshedAccount?.refreshJwt &&
+              store.getState().currentBundleState.bundle === bundle
+            ) {
+              const rebuilt = createSessionBundleFromStoredAccount(
+                committedAccount,
+                onSessionChangeRef.current!,
+              )
+              if (rebuilt) {
+                await store.dispatch({
+                  type: 'replaced-current-bundle',
+                  newBundle: rebuilt.bundle,
+                  newAccount: rebuilt.account,
+                })
+              }
+            } else if (
+              !committedAccount?.refreshJwt &&
+              store.getState().currentBundleState.bundle === bundle
+            ) {
+              await store.dispatch({
+                type: 'synced-accounts',
+                syncedAccounts: committedSession.accounts,
+                syncedCurrentDid: committedSession.currentAccount?.did,
+              })
+            }
+          }
+        },
       })
     },
     [store],
@@ -286,10 +353,23 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         disposeBundle(bundle)
         return
       }
-      store.dispatch({
+      await persisted.runWithSessionCredentialLock({
-        type: 'switched-to-account',
+        accountDids: [account.did],
-        newBundle: bundle,
+        operation: () =>
-        newAccount: account,
+          store.dispatch(
+            {
+              type: 'switched-to-account',
+              newBundle: bundle,
+              newAccount: account,
+            },
+            [
+              {
+                type: 'login',
+                accountDid: account.did,
+                resultRefreshJwt: account.refreshJwt,
+              },
+            ],
+          ),
       })
       ax.metric('account:create:success', metrics, {
         session: utils.accountToSessionMetadata(account),
@@ -317,10 +397,23 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         disposeBundle(bundle)
         return
       }
-      store.dispatch({
+      await persisted.runWithSessionCredentialLock({
-        type: 'switched-to-account',
+        accountDids: [account.did],
-        newBundle: bundle,
+        operation: () =>
-        newAccount: account,
+          store.dispatch(
+            {
+              type: 'switched-to-account',
+              newBundle: bundle,
+              newAccount: account,
+            },
+            [
+              {
+                type: 'login',
+                accountDid: account.did,
+                resultRefreshJwt: account.refreshJwt,
+              },
+            ],
+          ),
       })
       ax.metric(
         'account:loggedIn',
@@ -343,9 +436,18 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       addSessionDebugLog({type: 'method:start', method: 'logout'})
       cancelPendingTask()
       const prevState = store.getState()
-      store.dispatch({
+      const accountDid = prevState.currentBundleState.did
-        type: 'logged-out-current-account',
+      if (accountDid) {
-      })
+        void persisted
+          .runWithSessionCredentialLock({
+            accountDids: [accountDid],
+            operation: () =>
+              store.dispatch({type: 'logged-out-current-account', accountDid}, [
+                {type: 'logout', accountDid},
+              ]),
+          })
+          .catch(() => {})
+      }
       ax.metric(
         'account:loggedOut',
         {logContext, scope: 'current'},
@@ -377,9 +479,27 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       addSessionDebugLog({type: 'method:start', method: 'logout'})
       cancelPendingTask()
       const prevState = store.getState()
-      store.dispatch({
+      const accountDids = [
-        type: 'logged-out-every-account',
+        ...new Set([
-      })
+          ...prevState.accounts.map(account => account.did),
+          ...persisted
+            .readLatest('session')
+            .accounts.map(account => account.did),
+        ]),
+      ]
+      void persisted
+        .runWithSessionCredentialLock({
+          accountDids,
+          operation: () =>
+            store.dispatch(
+              {type: 'logged-out-every-account'},
+              accountDids.map(accountDid => ({
+                type: 'logout' as const,
+                accountDid,
+              })),
+            ),
+        })
+        .catch(() => {})
       ax.metric(
         'account:loggedOut',
         {logContext, scope: 'every'},
@@ -410,8 +530,12 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         account: redactAccount(storedAccount),
       })
       const signal = cancelPendingTask()
+      const latestStoredAccount = persisted
+        .readLatest('session')
+        .accounts.find(account => account.did === storedAccount.did)
+      if (!latestStoredAccount?.refreshJwt) return
       const {bundle, account} = await createSessionBundleAndResume(
-        storedAccount,
+        latestStoredAccount,
         onSessionChange,
       )
 
@@ -431,11 +555,48 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         disposeBundle(bundle)
         return
       }
-      store.dispatch({
+      const committedSession = await persisted.runWithSessionCredentialLock({
-        type: 'switched-to-account',
+        accountDids: [account.did],
-        newBundle: bundle,
+        operation: () =>
-        newAccount: account,
+          store.dispatch(
+            {
+              type: 'switched-to-account',
+              newBundle: bundle,
+              newAccount: account,
+            },
+            [
+              {
+                type: 'refresh',
+                accountDid: account.did,
+                baseRefreshJwt: latestStoredAccount.refreshJwt,
+                resultRefreshJwt: account.refreshJwt,
+              },
+            ],
+          ),
       })
+      const committedAccount = committedSession?.accounts.find(
+        candidate => candidate.did === account.did,
+      )
+      if (!committedAccount?.refreshJwt) {
+        await store.dispatch({
+          type: 'synced-accounts',
+          syncedAccounts: committedSession?.accounts ?? [],
+          syncedCurrentDid: committedSession?.currentAccount?.did,
+        })
+        return
+      }
+      if (committedAccount.refreshJwt !== account.refreshJwt) {
+        const rebuilt = createSessionBundleFromStoredAccount(
+          committedAccount,
+          onSessionChange,
+        )
+        if (!rebuilt) return
+        await store.dispatch({
+          type: 'replaced-current-bundle',
+          newBundle: rebuilt.bundle,
+          newAccount: rebuilt.account,
+        })
+      }
       addSessionDebugLog({
         type: 'method:end',
         method: 'resumeSession',
@@ -464,18 +625,22 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     /* getSession targets the PDS; only the persisted account fields are patched. */
     const data = await bundle.pdsClient.call(com.atproto.server.getSession, {})
     if (signal.aborted) return
-    store.dispatch({
+    await persisted.runWithSessionCredentialLock({
-      type: 'partial-refresh-session',
+      accountDids: [data.did],
-      /*
+      operation: () =>
-       * Read the did off the response rather than the session: the bundle may
+        store.dispatch({
-       * have been disposed while the request was in flight, and the live
+          type: 'partial-refresh-session',
-       * getters throw in that state.
+          /*
-       */
+           * Read the did off the response rather than the session: the bundle may
-      accountDid: data.did,
+           * have been disposed while the request was in flight, and the live
-      patch: {
+           * getters throw in that state.
-        emailConfirmed: data.emailConfirmed,
+           */
-        emailAuthFactor: data.emailAuthFactor,
+          accountDid: data.did,
-      },
+          patch: {
+            emailConfirmed: data.emailConfirmed,
+            emailAuthFactor: data.emailAuthFactor,
+          },
+        }),
     })
   }, [store, cancelPendingTask])
 
@@ -511,6 +676,12 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     if (!bundle.session) return undefined // logged out: nothing to refresh
     const before = bundle.session.session
     const after = await bundle.session.refresh()
+    const sessionChangeError = takeSessionChangeError({bundle})
+    if (sessionChangeError) {
+      throw sessionChangeError instanceof Error
+        ? sessionChangeError
+        : new Error('Session persistence failed', {cause: sessionChangeError})
+    }
     if (after === before) {
       throw new Error('Failed to refresh session')
     }
@@ -541,10 +712,19 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         account: redactAccount(account),
       })
       cancelPendingTask()
-      store.dispatch({
+      void persisted
-        type: 'removed-account',
+        .runWithSessionCredentialLock({
-        accountDid: account.did,
+          accountDids: [account.did],
-      })
+          operation: () =>
+            store.dispatch(
+              {
+                type: 'removed-account',
+                accountDid: account.did,
+              },
+              [{type: 'remove', accountDid: account.did}],
+            ),
+        })
+        .catch(() => {})
       addSessionDebugLog({
         type: 'method:end',
         method: 'removeAccount',
@@ -561,7 +741,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         type: 'persisted:receive',
         data: redactPersistedSession(synced),
       })
-      store.dispatch({
+      void store.dispatch({
         type: 'synced-accounts',
         syncedAccounts: synced.accounts,
         syncedCurrentDid: synced.currentAccount?.did,
@@ -636,7 +816,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
             return
           }
           const {bundle: newBundle, account: newAccount} = rebuilt
-          store.dispatch({
+          void store.dispatch({
             type: 'replaced-current-bundle',
             newBundle,
             newAccount,

src/state/session/reducer.ts

@@ -46,6 +46,7 @@ export type Action =
     }
   | {
       type: 'logged-out-current-account'
+      accountDid?: string
     }
   | {
       type: 'logged-out-every-account'
@@ -179,7 +180,7 @@ let reducer = (state: State, action: Action): State => {
     }
     case 'logged-out-current-account': {
       const {currentBundleState} = state
-      const accountDid = currentBundleState.did
+      const accountDid = action.accountDid ?? currentBundleState.did
       // side effect
       const account = state.accounts.find(a => a.did === accountDid)
       if (account && accountDid) {
@@ -206,7 +207,10 @@ let reducer = (state: State, action: Action): State => {
               }
             : a,
         ),
-        currentBundleState: createPublicBundleState(),
+        currentBundleState:
+          currentBundleState.did === accountDid
+            ? createPublicBundleState()
+            : currentBundleState,
         needsPersist: true,
       }
     }

src/state/session/session-core.ts

@@ -68,6 +68,7 @@ export type SessionBundle = {
  * state private and tied to bundle identity.
  */
 const bundleKillSwitches = new WeakMap<SessionBundle, () => void>()
+const bundleSessionChangeErrors = new WeakMap<SessionBundle, unknown>()
 
 /**
  * Register the lifecycle closure used by {@link disposeBundle}.
@@ -83,6 +84,17 @@ export function registerBundleKillSwitch(
   bundleKillSwitches.set(bundle, kill)
 }
 
+/** Return and clear the latest persistence/session-change hook failure. */
+export function takeSessionChangeError({
+  bundle,
+}: {
+  bundle: SessionBundle
+}): unknown {
+  const error = bundleSessionChangeErrors.get(bundle)
+  bundleSessionChangeErrors.delete(bundle)
+  return error
+}
+
 /**
  * Build the three clients over a session.
  *
@@ -126,7 +138,7 @@ export type OnSessionChange = (
   did: string,
   event: AtpSessionEvent,
   sessionData?: SessionData,
-) => void
+) => void | Promise<void>
 
 /**
  * Hooks stay inert during initial session preparation. `kill()` disarms them
@@ -146,7 +158,10 @@ export function makeSessionHooks({
 }) {
   let armed = false
   let killed = false
-  const dispatch = (event: AtpSessionEvent, sessionData?: SessionData) => {
+  const dispatch = async (
+    event: AtpSessionEvent,
+    sessionData?: SessionData,
+  ) => {
     if (!armed) {
       return
     }
@@ -159,12 +174,15 @@ export function makeSessionHooks({
      * effects and event emitters, so treat it as capable of throwing.
      */
     try {
+      const bundle = getBundle()
+      bundleSessionChangeErrors.delete(bundle)
       const did = getDid()
-      onSessionChange(getBundle(), did, event, sessionData)
+      await onSessionChange(bundle, did, event, sessionData)
       if (event !== 'update') {
         addSessionErrorLog(did, event)
       }
     } catch (e) {
+      bundleSessionChangeErrors.set(getBundle(), e)
       logger.error(e instanceof Error ? e : String(e), {
         message: `session: onSessionChange threw for a '${event}' event`,
       })
@@ -178,13 +196,13 @@ export function makeSessionHooks({
       return networkAwareFetch(input, init)
     },
     onUpdated(data) {
-      dispatch('update', data)
+      return dispatch('update', data)
     },
     onDeleted(data) {
-      dispatch('expired', data)
+      return dispatch('expired', data)
     },
     onUpdateFailure() {
-      dispatch('network-error')
+      return dispatch('network-error')
     },
   }
   return Object.assign(hooks, {