Reduce diff

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

@@ -0,0 +1,12 @@
+import {describe, expect, it} from '@jest/globals'
+
+import * as persisted from '../index'
+import {defaults} from '../schema'
+
+describe('generic persisted API', () => {
+  it('rejects session writes', () => {
+    expect(() => persisted.write('session', defaults.session)).toThrow(
+      "Session state must be written through '#/state/persisted/session'",
+    )
+  })
+})

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

@@ -5,7 +5,7 @@ import {
   applySessionUpdate,
   getCredentialState,
   type SessionCredentialMutation,
-} from '../session'
+} from '../session-merge'
 
 const DID = 'did:plc:example123'
 

src/state/persisted/index.ts

@@ -8,13 +8,13 @@ import {
   tryStringify,
 } from '#/state/persisted/schema'
 import {device} from '#/storage'
-import {applySessionUpdate} from './session'
-import {runWithSessionCredentialLock} from './session-lock'
+import {
+  applySessionUpdate,
+  type SessionCredentialMutation,
+} from './session-merge'
 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'
 
@@ -52,6 +52,11 @@ export function write<K extends keyof Schema>(
   key: K,
   value: Schema[K],
 ): Promise<void> {
+  if (key === 'session') {
+    throw new Error(
+      "Session state must be written through '#/state/persisted/session'",
+    )
+  }
   return enqueueWrite(async () => {
     const next = normalizeData({
       ..._state,
@@ -63,12 +68,13 @@ export function write<K extends keyof Schema>(
 }
 write satisfies PersistedApi['write']
 
-export function updateSession({
+/** @internal Use `#/state/persisted/session` instead. */
+export function writeSessionInternal({
   nextSession,
   credentialMutations,
 }: {
   nextSession: Schema['session']
-  credentialMutations: import('./session').SessionCredentialMutation[]
+  credentialMutations: SessionCredentialMutation[]
 }): Promise<Schema['session']> {
   return enqueueWrite(async () => {
     const session = applySessionUpdate({
@@ -82,9 +88,6 @@ export function updateSession({
     return session
   })
 }
-updateSession satisfies PersistedApi['updateSession']
-runWithSessionCredentialLock satisfies PersistedApi['runWithSessionCredentialLock']
-
 export function onUpdate<K extends keyof Schema>(
   _key: K,
   _cb: (v: Schema[K]) => void,

src/state/persisted/index.web.ts

@@ -8,17 +8,15 @@ import {
   tryParse,
   tryStringify,
 } from '#/state/persisted/schema'
+import {runWithSessionCredentialLock} from './session-lock'
 import {
   applySessionUpdate,
   getCredentialState,
   type SessionCredentialMutation,
-} from './session'
-import {runWithSessionCredentialLock} from './session-lock'
+} from './session-merge'
 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'
 
@@ -74,6 +72,11 @@ export function write<K extends keyof Schema>(
   key: K,
   value: Schema[K],
 ): Promise<void> {
+  if (key === 'session') {
+    throw new Error(
+      "Session state must be written through '#/state/persisted/session'",
+    )
+  }
   return runWithSessionCredentialLock({
     accountDids: [],
     operation: () => {
@@ -106,8 +109,9 @@ export function write<K extends keyof Schema>(
 }
 write satisfies PersistedApi['write']
 
+/** @internal Use `#/state/persisted/session` instead. */
 // eslint-disable-next-line @typescript-eslint/require-await
-export async function updateSession({
+export async function writeSessionInternal({
   nextSession,
   credentialMutations,
 }: {
@@ -134,9 +138,6 @@ export async function updateSession({
   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,

src/state/persisted/session-merge.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/session.ts

@@ -1,294 +1,45 @@
-import {jwtDecode} from 'jwt-decode'
+import * as persisted from './index'
+import {type Schema} from './schema'
+import {runWithSessionCredentialLock} from './session-lock'
+import {type SessionCredentialMutation} from './session-merge'
 
-import {
-  type PersistedAccount,
-  type PersistedCredentialState,
-  type Schema,
-} from './schema'
+export type {SessionCredentialMutation} from './session-merge'
 
-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 read(): Schema['session'] {
+  return persisted.get('session')
 }
 
-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',
-  }
+/** On web, synchronously read the authoritative localStorage session. */
+export function readLatest(): Schema['session'] {
+  return persisted.readLatest('session')
 }
 
-/**
- * Apply a session snapshot to the latest persisted session without allowing its
- * credential fields to overwrite newer credential generations.
- */
-export function applySessionUpdate({
-  storedSession,
+/** Conditionally commit a session update inside {@link runWithCredentialLock}. */
+export function write({
   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
+}): Promise<Schema['session']> {
+  return persisted.writeSessionInternal({
+    nextSession,
+    credentialMutations,
   })
-  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,
+export function runWithCredentialLock<T>({
+  accountDids,
+  operation,
 }: {
-  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
-    }
-  }
+  accountDids: string[]
+  operation: () => T | Promise<T>
+}): Promise<T> {
+  return runWithSessionCredentialLock({accountDids, operation})
 }
 
-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,
-  )
+export function onUpdate(
+  callback: (session: Schema['session']) => void,
+): () => void {
+  return persisted.onUpdate('session', callback)
 }

src/state/persisted/types.ts

@@ -1,5 +1,4 @@
 import {type Schema} from './schema'
-import {type SessionCredentialMutation} from './session'
 
 export type PersistedApi = {
   init(): Promise<void>
@@ -14,14 +13,6 @@ 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

@@ -6,19 +6,17 @@ import {act, render} from '@testing-library/react-native'
  * account factories. These mocks cut the tree back to the session lifecycle
  * itself, which is all these tests drive.
  */
-jest.mock('#/state/persisted', () => {
+jest.mock('#/state/persisted/session', () => {
   const {
     defaults,
   }: typeof import('#/state/persisted/schema') = require('#/state/persisted/schema')
   return {
-    defaults,
-    get: (key: keyof typeof defaults) => defaults[key],
-    write: () => Promise.resolve(),
-    updateSession: ({nextSession}: {nextSession: typeof defaults.session}) =>
+    read: () => defaults.session,
+    readLatest: () => defaults.session,
+    write: ({nextSession}: {nextSession: typeof defaults.session}) =>
       Promise.resolve(nextSession),
-    runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) =>
+    runWithCredentialLock: ({operation}: {operation: () => unknown}) =>
       Promise.resolve(operation()),
-    readLatest: (key: keyof typeof defaults) => defaults[key],
     onUpdate: () => () => {},
   }
 })

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

@@ -10,19 +10,17 @@ import {type SessionAccount} from '../types'
  * account factories. These mocks cut the tree back to the session lifecycle
  * itself, mirroring provider-abort-test.tsx.
  */
-jest.mock('#/state/persisted', () => {
+jest.mock('#/state/persisted/session', () => {
   const {
     defaults,
   }: typeof import('#/state/persisted/schema') = require('#/state/persisted/schema')
   return {
-    defaults,
-    get: (key: keyof typeof defaults) => defaults[key],
-    write: () => Promise.resolve(),
-    updateSession: ({nextSession}: {nextSession: typeof defaults.session}) =>
+    read: () => defaults.session,
+    readLatest: () => defaults.session,
+    write: ({nextSession}: {nextSession: typeof defaults.session}) =>
       Promise.resolve(nextSession),
-    runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) =>
+    runWithCredentialLock: ({operation}: {operation: () => unknown}) =>
       Promise.resolve(operation()),
-    readLatest: (key: keyof typeof defaults) => defaults[key],
     onUpdate: () => () => {},
   }
 })

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

@@ -9,19 +9,17 @@ import {type SessionAccount} from '../types'
  * account factories. These mocks cut the tree back to the session lifecycle
  * itself, mirroring provider-clients-test.tsx.
  */
-jest.mock('#/state/persisted', () => {
+jest.mock('#/state/persisted/session', () => {
   const {
     defaults,
   }: typeof import('#/state/persisted/schema') = require('#/state/persisted/schema')
   return {
-    defaults,
-    get: (key: keyof typeof defaults) => defaults[key],
-    write: () => Promise.resolve(),
-    updateSession: ({nextSession}: {nextSession: typeof defaults.session}) =>
+    read: () => defaults.session,
+    readLatest: () => defaults.session,
+    write: ({nextSession}: {nextSession: typeof defaults.session}) =>
       Promise.resolve(nextSession),
-    runWithSessionCredentialLock: ({operation}: {operation: () => unknown}) =>
+    runWithCredentialLock: ({operation}: {operation: () => unknown}) =>
       Promise.resolve(operation()),
-    readLatest: (key: keyof typeof defaults) => defaults[key],
     onUpdate: () => () => {},
   }
 })

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

@@ -24,48 +24,35 @@ const mockPersisted: {session: Schema['session']; latest: Schema['session']} = {
  * exists to catch.
  */
 const mockPersistedListeners: ((value: Schema['session']) => void)[] = []
-jest.mock('#/state/persisted', () => {
-  const {
-    defaults,
-  }: typeof import('#/state/persisted/schema') = require('#/state/persisted/schema')
-  return {
-    defaults,
-    get: (key: string) =>
-      key === 'session'
-        ? mockPersisted.session
-        : defaults[key as keyof typeof defaults],
-    readLatest: (key: string) =>
-      key === 'session'
-        ? mockPersisted.latest
-        : defaults[key as keyof typeof defaults],
-    write: () => Promise.resolve(),
-    updateSession: ({
+jest.mock('#/state/persisted/session', () => ({
+  read: () => mockPersisted.session,
+  readLatest: () => mockPersisted.latest,
+  write: ({
+    nextSession,
+    credentialMutations,
+  }: {
+    nextSession: Schema['session']
+    credentialMutations: import('#/state/persisted/session').SessionCredentialMutation[]
+  }) => {
+    const {
+      applySessionUpdate,
+    }: typeof import('#/state/persisted/session-merge') = require('#/state/persisted/session-merge')
+    const committed = applySessionUpdate({
+      storedSession: mockPersisted.latest,
       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 () => {}
-    },
-  }
-})
+    })
+    mockPersisted.session = committed
+    mockPersisted.latest = committed
+    return Promise.resolve(committed)
+  },
+  runWithCredentialLock: ({operation}: {operation: () => unknown}) =>
+    Promise.resolve(operation()),
+  onUpdate: (callback: (value: Schema['session']) => void) => {
+    mockPersistedListeners.push(callback)
+    return () => {}
+  },
+}))
 jest.mock('#/state/util', () => ({useCloseAllActiveElements: () => () => {}}))
 jest.mock('#/components/dialogs/Context', () => ({
   useGlobalDialogsControlContext: () => ({signinDialogControl: {open() {}}}),

src/state/session/index.tsx

@@ -12,8 +12,9 @@ import {
 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 {type Schema} from '#/state/persisted'
+import * as persistedSession from '#/state/persisted/session'
+import {type SessionCredentialMutation} from '#/state/persisted/session'
 import {useCloseAllActiveElements} from '#/state/util'
 import {useGlobalDialogsControlContext} from '#/components/dialogs/Context'
 import {AnalyticsContext, useAnalyticsBase, utils} from '#/analytics'
@@ -88,8 +89,8 @@ class SessionStore {
   private listeners = new Set<() => void>()
 
   constructor() {
-    // Careful: By the time this runs, `persisted` needs to already be filled.
-    const initialState = getInitialState(persisted.get('session').accounts)
+    // Careful: By the time this runs, persisted state must already be initialized.
+    const initialState = getInitialState(persistedSession.read().accounts)
     addSessionDebugLog({type: 'reducer:init', state: redactState(initialState)})
     this.state = initialState
   }
@@ -126,7 +127,7 @@ class SessionStore {
         type: 'persisted:broadcast',
         data: redactPersistedSession(persistedData),
       })
-      persistence = persisted.updateSession({
+      persistence = persistedSession.write({
         nextSession: persistedData,
         credentialMutations,
       })
@@ -166,7 +167,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
           ? bundle.session.session.refreshJwt
           : sessionData?.refreshJwt
 
-      return persisted.runWithSessionCredentialLock({
+      return persistedSession.runWithCredentialLock({
         accountDids: [accountDid],
         operation: async () => {
           /*
@@ -227,8 +228,8 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
               }
               failedSet.add(dyingRefreshJwt)
 
-              const persistedCandidate = persisted
-                .readLatest('session')
+              const persistedCandidate = persistedSession
+                .readLatest()
                 .accounts.find(a => a.did === accountDid)
               const reducerCandidate = current.accounts.find(
                 a => a.did === accountDid,
@@ -353,7 +354,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         disposeBundle(bundle)
         return
       }
-      await persisted.runWithSessionCredentialLock({
+      await persistedSession.runWithCredentialLock({
         accountDids: [account.did],
         operation: () =>
           store.dispatch(
@@ -397,7 +398,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         disposeBundle(bundle)
         return
       }
-      await persisted.runWithSessionCredentialLock({
+      await persistedSession.runWithCredentialLock({
         accountDids: [account.did],
         operation: () =>
           store.dispatch(
@@ -438,8 +439,8 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const prevState = store.getState()
       const accountDid = prevState.currentBundleState.did
       if (accountDid) {
-        void persisted
-          .runWithSessionCredentialLock({
+        void persistedSession
+          .runWithCredentialLock({
             accountDids: [accountDid],
             operation: () =>
               store.dispatch({type: 'logged-out-current-account', accountDid}, [
@@ -482,13 +483,11 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const accountDids = [
         ...new Set([
           ...prevState.accounts.map(account => account.did),
-          ...persisted
-            .readLatest('session')
-            .accounts.map(account => account.did),
+          ...persistedSession.readLatest().accounts.map(account => account.did),
         ]),
       ]
-      void persisted
-        .runWithSessionCredentialLock({
+      void persistedSession
+        .runWithCredentialLock({
           accountDids,
           operation: () =>
             store.dispatch(
@@ -530,8 +529,8 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         account: redactAccount(storedAccount),
       })
       const signal = cancelPendingTask()
-      const latestStoredAccount = persisted
-        .readLatest('session')
+      const latestStoredAccount = persistedSession
+        .readLatest()
         .accounts.find(account => account.did === storedAccount.did)
       if (!latestStoredAccount?.refreshJwt) return
       const {bundle, account} = await createSessionBundleAndResume(
@@ -555,7 +554,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         disposeBundle(bundle)
         return
       }
-      const committedSession = await persisted.runWithSessionCredentialLock({
+      const committedSession = await persistedSession.runWithCredentialLock({
         accountDids: [account.did],
         operation: () =>
           store.dispatch(
@@ -625,7 +624,7 @@ 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
-    await persisted.runWithSessionCredentialLock({
+    await persistedSession.runWithCredentialLock({
       accountDids: [data.did],
       operation: () =>
         store.dispatch({
@@ -712,8 +711,8 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         account: redactAccount(account),
       })
       cancelPendingTask()
-      void persisted
-        .runWithSessionCredentialLock({
+      void persistedSession
+        .runWithCredentialLock({
           accountDids: [account.did],
           operation: () =>
             store.dispatch(
@@ -735,7 +734,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     [store, cancelPendingTask],
   )
   useEffect(() => {
-    return persisted.onUpdate('session', nextSession => {
+    return persistedSession.onUpdate(nextSession => {
       const synced = nextSession
       addSessionDebugLog({
         type: 'persisted:receive',

src/state/session/util.ts

@@ -2,7 +2,7 @@ import {PasswordSession} from '@atproto/lex-password-session'
 
 import {createLexClient} from '#/lib/lexClient'
 import {type TemporaryPushClient} from '#/lib/notifications/notifications'
-import * as persisted from '#/state/persisted'
+import * as persistedSession from '#/state/persisted/session'
 import {networkAwareFetch} from './network'
 import {sessionAccountToSessionData} from './session-data'
 import {type SessionAccount} from './types'
@@ -10,7 +10,7 @@ import {type SessionAccount} from './types'
 export {isSessionExpired, isSignupQueued} from './session-data'
 
 export function readLastActiveAccount() {
-  const {currentAccount, accounts} = persisted.get('session')
+  const {currentAccount, accounts} = persistedSession.read()
   return accounts.find(a => a.did === currentAccount?.did)
 }