Move session persistence into base API

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

@@ -6,7 +6,7 @@ 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'",
+      'Session state must be written through persisted.writeSession()',
     )
   })
 })

src/state/persisted/index.ts

@@ -15,6 +15,11 @@ import {
 import {type PersistedApi} from './types'
 import {normalizeData} from './util'
 
+export type {SessionCredentialMutation} from './session-merge'
+export {
+  runWithPersistedStorageLock as runWithCredentialLock,
+  runWithPersistedStorageLock,
+} from './storage-lock'
 export type {PersistedAccount, Schema} from '#/state/persisted/schema'
 export {defaults} from '#/state/persisted/schema'
 
@@ -59,7 +64,7 @@ export function write<K extends keyof Schema>(
 ): Promise<void> {
   if (key === 'session') {
     throw new Error(
-      "Session state must be written through '#/state/persisted/session'",
+      'Session state must be written through persisted.writeSession()',
     )
   }
   return enqueueWrite(async () => {
@@ -73,8 +78,7 @@ export function write<K extends keyof Schema>(
 }
 write satisfies PersistedApi['write']
 
-/** @internal Use `#/state/persisted/session` instead. */
+export function writeSession({
-export function writeSessionInternal({
   nextSession,
   credentialMutations,
 }: {
@@ -93,6 +97,8 @@ export function writeSessionInternal({
     return session
   })
 }
+writeSession satisfies PersistedApi['writeSession']
+
 export function onUpdate<K extends keyof Schema>(
   _key: K,
   _cb: (v: Schema[K]) => void,

src/state/persisted/index.web.ts

@@ -16,6 +16,11 @@ import {runWithPersistedStorageLock} from './storage-lock'
 import {type PersistedApi} from './types'
 import {normalizeData} from './util'
 
+export type {SessionCredentialMutation} from './session-merge'
+export {
+  runWithPersistedStorageLock as runWithCredentialLock,
+  runWithPersistedStorageLock,
+} from './storage-lock'
 export type {PersistedAccount, Schema} from '#/state/persisted/schema'
 export {defaults} from '#/state/persisted/schema'
 
@@ -73,7 +78,7 @@ export function write<K extends keyof Schema>(
 ): Promise<void> {
   if (key === 'session') {
     throw new Error(
-      "Session state must be written through '#/state/persisted/session'",
+      'Session state must be written through persisted.writeSession()',
     )
   }
   return runWithPersistedStorageLock({
@@ -107,9 +112,8 @@ 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 writeSessionInternal({
+export async function writeSession({
   nextSession,
   credentialMutations,
 }: {
@@ -129,6 +133,8 @@ export async function writeSessionInternal({
   broadcastUpdate({key: 'session'})
   return session
 }
+writeSession satisfies PersistedApi['writeSession']
+
 export function onUpdate<K extends keyof Schema>(
   key: K,
   cb: (v: Schema[K]) => void,

src/state/persisted/session.ts

@@ -1,35 +0,0 @@
-import * as persisted from './index'
-import {type Schema} from './schema'
-import {type SessionCredentialMutation} from './session-merge'
-
-export type {SessionCredentialMutation} from './session-merge'
-export {runWithPersistedStorageLock as runWithCredentialLock} from './storage-lock'
-
-export function read(): Schema['session'] {
-  return persisted.get('session')
-}
-
-/** On web, synchronously read the authoritative localStorage session. */
-export function readLatest(): Schema['session'] {
-  return persisted.readLatest('session')
-}
-
-/** Conditionally commit a session update inside {@link runWithCredentialLock}. */
-export function write({
-  nextSession,
-  credentialMutations,
-}: {
-  nextSession: Schema['session']
-  credentialMutations: SessionCredentialMutation[]
-}): Promise<Schema['session']> {
-  return persisted.writeSessionInternal({
-    nextSession,
-    credentialMutations,
-  })
-}
-
-export function onUpdate(
-  callback: (session: Schema['session']) => void,
-): () => void {
-  return persisted.onUpdate('session', callback)
-}

src/state/persisted/storage-lock.ts

@@ -1,3 +1,5 @@
+import {type PersistedApi} from './types'
+
 export function runWithPersistedStorageLock<T>({
   operation,
 }: {
@@ -13,3 +15,4 @@ export function runWithPersistedStorageLock<T>({
     )
   }
 }
+runWithPersistedStorageLock satisfies PersistedApi['runWithPersistedStorageLock']

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

@@ -1,3 +1,5 @@
+import {type PersistedApi} from './types'
+
 const PERSISTED_STORAGE_LOCK = 'bsky-persisted-storage'
 
 export function runWithPersistedStorageLock<T>({
@@ -20,6 +22,7 @@ export function runWithPersistedStorageLock<T>({
 
   return lockManager.request(PERSISTED_STORAGE_LOCK, operation)
 }
+runWithPersistedStorageLock satisfies PersistedApi['runWithPersistedStorageLock']
 
 function getLockManager(): LockManager | undefined {
   if (typeof navigator === 'undefined' || !('locks' in navigator)) {

src/state/persisted/types.ts

@@ -1,4 +1,5 @@
 import {type Schema} from './schema'
+import {type SessionCredentialMutation} from './session-merge'
 
 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>
+  /** Conditionally merges session credentials; generic session writes throw. */
+  writeSession(args: {
+    nextSession: Schema['session']
+    credentialMutations: SessionCredentialMutation[]
+  }): Promise<Schema['session']>
+  runWithPersistedStorageLock<T>(args: {
+    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,14 +6,16 @@ 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/session', () => {
+jest.mock('#/state/persisted', () => {
+  const actual = jest.requireActual<object>('#/state/persisted')
   const {
     defaults,
   }: typeof import('#/state/persisted/schema') = require('#/state/persisted/schema')
   return {
-    read: () => defaults.session,
+    ...actual,
+    get: () => defaults.session,
     readLatest: () => defaults.session,
-    write: ({nextSession}: {nextSession: typeof defaults.session}) =>
+    writeSession: ({nextSession}: {nextSession: typeof defaults.session}) =>
       Promise.resolve(nextSession),
     runWithCredentialLock: ({operation}: {operation: () => unknown}) =>
       Promise.resolve(operation()),

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

@@ -10,14 +10,16 @@ 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/session', () => {
+jest.mock('#/state/persisted', () => {
+  const actual = jest.requireActual<object>('#/state/persisted')
   const {
     defaults,
   }: typeof import('#/state/persisted/schema') = require('#/state/persisted/schema')
   return {
-    read: () => defaults.session,
+    ...actual,
+    get: () => defaults.session,
     readLatest: () => defaults.session,
-    write: ({nextSession}: {nextSession: typeof defaults.session}) =>
+    writeSession: ({nextSession}: {nextSession: typeof defaults.session}) =>
       Promise.resolve(nextSession),
     runWithCredentialLock: ({operation}: {operation: () => unknown}) =>
       Promise.resolve(operation()),

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

@@ -9,14 +9,16 @@ 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/session', () => {
+jest.mock('#/state/persisted', () => {
+  const actual = jest.requireActual<object>('#/state/persisted')
   const {
     defaults,
   }: typeof import('#/state/persisted/schema') = require('#/state/persisted/schema')
   return {
-    read: () => defaults.session,
+    ...actual,
+    get: () => defaults.session,
     readLatest: () => defaults.session,
-    write: ({nextSession}: {nextSession: typeof defaults.session}) =>
+    writeSession: ({nextSession}: {nextSession: typeof defaults.session}) =>
       Promise.resolve(nextSession),
     runWithCredentialLock: ({operation}: {operation: () => unknown}) =>
       Promise.resolve(operation()),

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

@@ -24,15 +24,15 @@ const mockPersisted: {session: Schema['session']; latest: Schema['session']} = {
  * exists to catch.
  */
 const mockPersistedListeners: ((value: Schema['session']) => void)[] = []
-jest.mock('#/state/persisted/session', () => ({
+jest.mock('#/state/persisted', () => ({
-  read: () => mockPersisted.session,
+  get: () => mockPersisted.session,
   readLatest: () => mockPersisted.latest,
-  write: ({
+  writeSession: ({
     nextSession,
     credentialMutations,
   }: {
     nextSession: Schema['session']
-    credentialMutations: import('#/state/persisted/session').SessionCredentialMutation[]
+    credentialMutations: import('#/state/persisted').SessionCredentialMutation[]
   }) => {
     const {
       applySessionUpdate,
@@ -48,7 +48,7 @@ jest.mock('#/state/persisted/session', () => ({
   },
   runWithCredentialLock: ({operation}: {operation: () => unknown}) =>
     Promise.resolve(operation()),
-  onUpdate: (callback: (value: Schema['session']) => void) => {
+  onUpdate: (_key: 'session', callback: (value: Schema['session']) => void) => {
     mockPersistedListeners.push(callback)
     return () => {}
   },

src/state/session/index.tsx

@@ -12,9 +12,8 @@ import {
 import {type Client} from '@atproto/lex'
 import {type SessionData} from '@atproto/lex-password-session'
 
-import {type Schema} from '#/state/persisted'
+import * as persistedSession from '#/state/persisted'
-import * as persistedSession from '#/state/persisted/session'
+import {type Schema, type SessionCredentialMutation} from '#/state/persisted'
-import {type SessionCredentialMutation} from '#/state/persisted/session'
 import {useCloseAllActiveElements} from '#/state/util'
 import {useGlobalDialogsControlContext} from '#/components/dialogs/Context'
 import {AnalyticsContext, useAnalyticsBase, utils} from '#/analytics'
@@ -90,7 +89,9 @@ class SessionStore {
 
   constructor() {
     // Careful: By the time this runs, persisted state must already be initialized.
-    const initialState = getInitialState(persistedSession.read().accounts)
+    const initialState = getInitialState(
+      persistedSession.get('session').accounts,
+    )
     addSessionDebugLog({type: 'reducer:init', state: redactState(initialState)})
     this.state = initialState
   }
@@ -127,7 +128,7 @@ class SessionStore {
         type: 'persisted:broadcast',
         data: redactPersistedSession(persistedData),
       })
-      persistence = persistedSession.write({
+      persistence = persistedSession.writeSession({
         nextSession: persistedData,
         credentialMutations,
       })
@@ -228,7 +229,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
               failedSet.add(dyingRefreshJwt)
 
               const persistedCandidate = persistedSession
-                .readLatest()
+                .readLatest('session')
                 .accounts.find(a => a.did === accountDid)
               const reducerCandidate = current.accounts.find(
                 a => a.did === accountDid,
@@ -479,7 +480,9 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const accountDids = [
         ...new Set([
           ...prevState.accounts.map(account => account.did),
-          ...persistedSession.readLatest().accounts.map(account => account.did),
+          ...persistedSession
+            .readLatest('session')
+            .accounts.map(account => account.did),
         ]),
       ]
       void persistedSession
@@ -525,7 +528,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       })
       const signal = cancelPendingTask()
       const latestStoredAccount = persistedSession
-        .readLatest()
+        .readLatest('session')
         .accounts.find(account => account.did === storedAccount.did)
       if (!latestStoredAccount?.refreshJwt) return
       const {bundle, account} = await createSessionBundleAndResume(
@@ -726,7 +729,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     [store, cancelPendingTask],
   )
   useEffect(() => {
-    return persistedSession.onUpdate(nextSession => {
+    return persistedSession.onUpdate('session', 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 persistedSession from '#/state/persisted/session'
+import * as persistedSession from '#/state/persisted'
 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} = persistedSession.read()
+  const {currentAccount, accounts} = persistedSession.get('session')
   return accounts.find(a => a.did === currentAccount?.did)
 }