redact session debug logs, fix stale bundle read and render-phase ref write

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

src/state/session/index.tsx

@@ -3,6 +3,7 @@ import {
   useCallback,
   useContext,
   useEffect,
+  useInsertionEffect,
   useMemo,
   useRef,
   useState,
@@ -31,7 +32,14 @@ import {
   sessionDataToSessionAccount,
 } from './session-core'
 export {isSignupQueued} from './session-data'
-import {addSessionDebugLog} from './logging'
+import {
+  addSessionDebugLog,
+  getBundleId,
+  redactAccount,
+  redactPersistedSession,
+  redactSessionData,
+  redactState,
+} from './logging'
 export type {SessionAccount} from '#/state/session/types'
 
 import {clearPersistedQueryStorage} from '#/lib/persisted-query-storage'
@@ -76,7 +84,7 @@ class SessionStore {
   constructor() {
     // Careful: By the time this runs, `persisted` needs to already be filled.
     const initialState = getInitialState(persisted.get('session').accounts)
-    addSessionDebugLog({type: 'reducer:init', state: initialState})
+    addSessionDebugLog({type: 'reducer:init', state: redactState(initialState)})
     this.state = initialState
   }
 
@@ -103,7 +111,10 @@ class SessionStore {
           a => a.did === nextState.currentBundleState.did,
         ),
       }
-      addSessionDebugLog({type: 'persisted:broadcast', data: persistedData})
+      addSessionDebugLog({
+        type: 'persisted:broadcast',
+        data: redactPersistedSession(persistedData),
+      })
       void persisted.write('session', persistedData)
     }
     this.listeners.forEach(listener => listener())
@@ -122,7 +133,9 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   const failedExpiryTokensRef = useRef<Map<string, Set<string>>>(new Map())
   /*
    * Rescued bundles need this callback for their own events. A ref avoids a
-   * self-reference in the callback's dependency list.
+   * self-reference in the callback's dependency list. It is filled by the
+   * 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<
     | ((
@@ -226,7 +239,15 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     },
     [store],
   )
-  onSessionChangeRef.current = onSessionChange
+  /*
+   * Writing the ref during render is forbidden under React Compiler. An
+   * insertion effect is the earliest commit-time slot, and the only reader
+   * (`onSessionChange`'s expiry-rescue path) runs from armed session hooks,
+   * which cannot fire before the first commit.
+   */
+  useInsertionEffect(() => {
+    onSessionChangeRef.current = onSessionChange
+  }, [onSessionChange])
 
   const createAccount = useCallback<SessionApiContext['createAccount']>(
     async (params, metrics) => {
@@ -251,7 +272,11 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       ax.metric('account:create:success', metrics, {
         session: utils.accountToSessionMetadata(account),
       })
-      addSessionDebugLog({type: 'method:end', method: 'createAccount', account})
+      addSessionDebugLog({
+        type: 'method:end',
+        method: 'createAccount',
+        account: redactAccount(account),
+      })
     },
     [ax, store, onSessionChange, cancelPendingTask],
   )
@@ -280,7 +305,11 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         {logContext, withPassword: true},
         {session: utils.accountToSessionMetadata(account)},
       )
-      addSessionDebugLog({type: 'method:end', method: 'login', account})
+      addSessionDebugLog({
+        type: 'method:end',
+        method: 'login',
+        account: redactAccount(account),
+      })
     },
     [ax, store, onSessionChange, cancelPendingTask],
   )
@@ -356,7 +385,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       addSessionDebugLog({
         type: 'method:start',
         method: 'resumeSession',
-        account: storedAccount,
+        account: redactAccount(storedAccount),
       })
       const signal = cancelPendingTask()
       const {bundle, account} = await createSessionBundleAndResume(
@@ -385,7 +414,11 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         newBundle: bundle,
         newAccount: account,
       })
-      addSessionDebugLog({type: 'method:end', method: 'resumeSession', account})
+      addSessionDebugLog({
+        type: 'method:end',
+        method: 'resumeSession',
+        account: redactAccount(account),
+      })
       if (isSwitchingAccounts) {
         // reset onboarding flow on switch account
         onboardingDispatch({type: 'skip'})
@@ -397,7 +430,13 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   const partialRefreshSession = useCallback<
     SessionApiContext['partialRefreshSession']
   >(async () => {
-    const bundle = state.currentBundleState.bundle as unknown as SessionBundle
+    /*
+     * Read the live bundle rather than the one captured by this render: a
+     * dispatch that lands before the next render would otherwise leave this
+     * holding a disposed bundle, whose agent dispatches unauthenticated.
+     */
+    const bundle = store.getState().currentBundleState
+      .bundle as unknown as SessionBundle
     const signal = cancelPendingTask()
     /* getSession targets the PDS; only the persisted account fields are patched. */
     const {data} = await bundle.agent.com.atproto.server.getSession()
@@ -415,21 +454,25 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         emailAuthFactor: data.emailAuthFactor,
       },
     })
-  }, [store, state, cancelPendingTask])
+  }, [store, cancelPendingTask])
 
   const removeAccount = useCallback<SessionApiContext['removeAccount']>(
     account => {
       addSessionDebugLog({
         type: 'method:start',
         method: 'removeAccount',
-        account,
+        account: redactAccount(account),
       })
       cancelPendingTask()
       store.dispatch({
         type: 'removed-account',
         accountDid: account.did,
       })
-      addSessionDebugLog({type: 'method:end', method: 'removeAccount', account})
+      addSessionDebugLog({
+        type: 'method:end',
+        method: 'removeAccount',
+        account: redactAccount(account),
+      })
       clearAgeAssuranceServerDataForDid({did: account.did})
     },
     [store, cancelPendingTask],
@@ -437,7 +480,10 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   useEffect(() => {
     return persisted.onUpdate('session', nextSession => {
       const synced = nextSession
-      addSessionDebugLog({type: 'persisted:receive', data: synced})
+      addSessionDebugLog({
+        type: 'persisted:receive',
+        data: redactPersistedSession(synced),
+      })
       store.dispatch({
         type: 'synced-accounts',
         syncedAccounts: synced.accounts,
@@ -498,12 +544,13 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
               if (isCurrent) {
                 addSessionDebugLog({
                   type: 'bundle:patch',
-                  bundle: newBundle,
+                  bundleId: getBundleId(newBundle),
-                  prevSession:
+                  prevSession: redactSessionData(
                     prevBundle.session && !prevBundle.session.destroyed
                       ? prevBundle.session.session
                       : undefined,
-                  nextSession: newBundle.session.session,
+                  ),
+                  nextSession: redactSessionData(newBundle.session.session),
                 })
               }
               return isCurrent
@@ -577,8 +624,8 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       currentBundleRef.current = bundle
       addSessionDebugLog({
         type: 'bundle:switch',
-        prevBundle,
+        prevBundleId: getBundleId(prevBundle),
-        nextBundle: bundle,
+        nextBundleId: getBundleId(bundle),
       })
       // Replaced bundles must never consume another refresh token.
       disposeBundle(prevBundle)

src/state/session/logging.ts

@@ -1,19 +1,140 @@
+import {type SessionData} from '@atproto/lex-password-session'
+
 import {type Schema} from '../persisted'
 import {type Action, type State} from './reducer'
 import {type AtpSessionEvent, type SessionAccount} from './types'
 
 type Reducer = (state: State, action: Action) => State
 
+/**
+ * An account reduced to the fields that are safe to ship off-device.
+ *
+ * Credentials become presence booleans and PII (email, service, pdsUrl) is
+ * dropped entirely. The log only ever needs to answer "which account was live,
+ * and did it still hold credentials?". The stubs at the bottom of this file are
+ * expected to be revived against Sentry or Bitdrift, so no payload type may be
+ * capable of carrying a live JWT in the first place.
+ */
+export type RedactedAccount = {
+  did: string
+  handle: string
+  active: boolean | undefined
+  status: string | undefined
+  signupQueued: boolean | undefined
+  hasAccessJwt: boolean
+  hasRefreshJwt: boolean
+}
+
+/** The account list plus the current did, in redacted form. */
+export type RedactedSessionSnapshot = {
+  accounts: RedactedAccount[]
+  currentDid: string | undefined
+}
+
+/** Live session data reduced to identity plus credential presence. */
+export type RedactedSessionData = {
+  did: string
+  handle: string
+  hasAccessJwt: boolean
+  hasRefreshJwt: boolean
+}
+
+function redact(account: SessionAccount): RedactedAccount {
+  return {
+    did: account.did,
+    handle: account.handle,
+    active: account.active,
+    status: account.status,
+    signupQueued: account.signupQueued,
+    hasAccessJwt: !!account.accessJwt,
+    hasRefreshJwt: !!account.refreshJwt,
+  }
+}
+
+export function redactAccount(
+  account: SessionAccount | undefined,
+): RedactedAccount | undefined {
+  return account ? redact(account) : undefined
+}
+
+export function redactState(state: State): RedactedSessionSnapshot {
+  return {
+    accounts: state.accounts.map(redact),
+    currentDid: state.currentBundleState.did,
+  }
+}
+
+export function redactPersistedSession(
+  data: Schema['session'],
+): RedactedSessionSnapshot {
+  return {
+    accounts: data.accounts.map(redact),
+    currentDid: data.currentAccount?.did,
+  }
+}
+
+export function redactSessionData(
+  data: SessionData | undefined,
+): RedactedSessionData | undefined {
+  if (!data) return undefined
+  return {
+    did: data.did,
+    handle: data.handle,
+    hasAccessJwt: !!data.accessJwt,
+    hasRefreshJwt: !!data.refreshJwt,
+  }
+}
+
+/*
+ * Bundles are logged for identity only - which one was live, and which one
+ * replaced it. Logging the object itself would reach its session, and through
+ * it the tokens, so each is mapped to an opaque per-run id instead.
+ */
+const bundleIds = new WeakMap<object, string>()
+const runId = Math.random().toString(36).slice(2)
+let nextBundleId = 1
+
+export function getBundleId(bundle: object): string {
+  let id = bundleIds.get(bundle)
+  if (id === undefined) {
+    id = runId + '::' + nextBundleId++
+    bundleIds.set(bundle, id)
+  }
+  return id
+}
+
+/** An action reduced to its discriminant plus the did it targets, if any. */
+type RedactedAction = {
+  type: Action['type']
+  accountDid?: string
+}
+
+function redactAction(action: Action): RedactedAction {
+  switch (action.type) {
+    case 'received-session-event':
+    case 'removed-account':
+    case 'partial-refresh-session':
+      return {type: action.type, accountDid: action.accountDid}
+    case 'switched-to-account':
+    case 'replaced-current-bundle':
+      return {type: action.type, accountDid: action.newAccount.did}
+    case 'synced-accounts':
+      return {type: action.type, accountDid: action.syncedCurrentDid}
+    default:
+      return {type: action.type}
+  }
+}
+
 type Log =
   | {
       type: 'reducer:init'
-      state: State
+      state: RedactedSessionSnapshot
     }
   | {
       type: 'reducer:call'
-      action: Action
+      action: RedactedAction
-      prevState: State
+      prevState: RedactedSessionSnapshot
-      nextState: State
+      nextState: RedactedSessionSnapshot
     }
   | {
       type: 'method:start'
@@ -23,7 +144,7 @@ type Log =
         | 'logout'
         | 'resumeSession'
         | 'removeAccount'
-      account?: SessionAccount
+      account?: RedactedAccount
     }
   | {
       type: 'method:end'
@@ -33,37 +154,41 @@ type Log =
         | 'logout'
         | 'resumeSession'
         | 'removeAccount'
-      account?: SessionAccount
+      account?: RedactedAccount
     }
   | {
       type: 'persisted:broadcast'
-      data: Schema['session']
+      data: RedactedSessionSnapshot
     }
   | {
       type: 'persisted:receive'
-      data: Schema['session']
+      data: RedactedSessionSnapshot
     }
   | {
       type: 'bundle:switch'
-      prevBundle: object
+      prevBundleId: string
-      nextBundle: object
+      nextBundleId: string
     }
   | {
       /*
-       * Dev-only bundle-swap log. The bundle is treated as an opaque object
+       * Dev-only bundle-swap log. Bundles are identified by their opaque ids;
-       * (the reducer never reads its internals); the session snapshots are
+       * the session snapshots record only identity and credential presence.
-       * plain objects captured for debugging.
        */
       type: 'bundle:patch'
-      bundle: object
+      bundleId: string
-      prevSession: object | undefined
+      prevSession: RedactedSessionData | undefined
-      nextSession: object | undefined
+      nextSession: RedactedSessionData | undefined
     }
 
 export function wrapSessionReducerForLogging(reducer: Reducer): Reducer {
   return function loggingWrapper(prevState: State, action: Action): State {
     const nextState = reducer(prevState, action)
-    addSessionDebugLog({type: 'reducer:call', prevState, action, nextState})
+    addSessionDebugLog({
+      type: 'reducer:call',
+      prevState: redactState(prevState),
+      action: redactAction(action),
+      nextState: redactState(nextState),
+    })
     return nextState
   }
 }