add refreshSession to the session api and migrate its callers

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

oxlint-suppressions.json

@@ -848,9 +848,6 @@
     },
     "typescript/no-misused-promises": {
       "count": 1
-    },
-    "typescript/no-unsafe-member-access": {
-      "count": 1
     }
   },
   "src/screens/E2E/SharedPreferencesTesterScreen.tsx": {

src/components/dialogs/EmailDialog/data/useConfirmEmail.ts

@@ -1,13 +1,15 @@
 import {useMutation} from '@tanstack/react-query'
 
-import {useAgent, useSession} from '#/state/session'
+import {usePdsClient, useSession, useSessionApi} from '#/state/session'
+import {com} from '#/lexicons'
 
 export function useConfirmEmail({
   onSuccess,
   onError,
 }: {onSuccess?: () => void; onError?: () => void} = {}) {
-  const agent = useAgent()
+  const pdsClient = usePdsClient()
   const {currentAccount} = useSession()
+  const {refreshSession} = useSessionApi()
 
   return useMutation({
     mutationFn: async ({token}: {token: string}) => {
@@ -15,12 +17,12 @@ export function useConfirmEmail({
         throw new Error('No email found for the current account')
       }
 
-      await agent.com.atproto.server.confirmEmail({
+      await pdsClient.call(com.atproto.server.confirmEmail, {
         email: currentAccount.email.trim(),
         token: token.trim(),
       })
       // will update session state at root of app
-      await agent.resumeSession(agent.session!)
+      await refreshSession()
     },
     onSuccess,
     onError,

src/components/dialogs/EmailDialog/data/useManageEmail2FA.ts

@@ -1,10 +1,12 @@
 import {useMutation} from '@tanstack/react-query'
 
-import {useAgent, useSession} from '#/state/session'
+import {usePdsClient, useSession, useSessionApi} from '#/state/session'
+import {com} from '#/lexicons'
 
 export function useManageEmail2FA() {
-  const agent = useAgent()
+  const pdsClient = usePdsClient()
   const {currentAccount} = useSession()
+  const {refreshSession} = useSessionApi()
 
   return useMutation({
     mutationFn: async ({
@@ -17,13 +19,13 @@ export function useManageEmail2FA() {
         throw new Error('No email found for the current account')
       }
 
-      await agent.com.atproto.server.updateEmail({
+      await pdsClient.call(com.atproto.server.updateEmail, {
         email: currentAccount.email,
         emailAuthFactor: enabled,
         token,
       })
       // will update session state at root of app
-      await agent.resumeSession(agent.session!)
+      await refreshSession()
     },
   })
 }

src/components/dialogs/EmailDialog/data/useUpdateEmail.ts

@@ -1,19 +1,26 @@
+import {type Client} from '@atproto/lex'
 import {useMutation} from '@tanstack/react-query'
 
-import {useAgent} from '#/state/session'
+import {usePdsClient, useSessionApi} from '#/state/session'
 import {useRequestEmailUpdate} from '#/components/dialogs/EmailDialog/data/useRequestEmailUpdate'
+import {com} from '#/lexicons'
 
 async function updateEmailAndRefreshSession(
-  agent: ReturnType<typeof useAgent>,
+  pdsClient: Client,
+  refreshSession: () => Promise<unknown>,
   email: string,
   token?: string,
 ) {
-  await agent.com.atproto.server.updateEmail({email: email.trim(), token})
-  await agent.resumeSession(agent.session!)
+  await pdsClient.call(com.atproto.server.updateEmail, {
+    email: email.trim(),
+    token,
+  })
+  await refreshSession()
 }
 
 export function useUpdateEmail() {
-  const agent = useAgent()
+  const pdsClient = usePdsClient()
+  const {refreshSession} = useSessionApi()
   const {mutateAsync: requestEmailUpdate} = useRequestEmailUpdate()
 
   return useMutation<
@@ -23,7 +30,12 @@ export function useUpdateEmail() {
   >({
     mutationFn: async ({email, token}: {email: string; token?: string}) => {
       if (token) {
-        await updateEmailAndRefreshSession(agent, email, token)
+        await updateEmailAndRefreshSession(
+          pdsClient,
+          refreshSession,
+          email,
+          token,
+        )
         return {
           status: 'success',
         }
@@ -34,7 +46,12 @@ export function useUpdateEmail() {
             status: 'tokenRequired',
           }
         } else {
-          await updateEmailAndRefreshSession(agent, email, token)
+          await updateEmailAndRefreshSession(
+            pdsClient,
+            refreshSession,
+            email,
+            token,
+          )
           return {
             status: 'success',
           }

src/screens/Deactivated.tsx

@@ -7,10 +7,11 @@ import {Trans} from '@lingui/react/macro'
 import {useQueryClient} from '@tanstack/react-query'
 
 import {useAccountSwitcher} from '#/lib/hooks/useAccountSwitcher'
+import {isErrorMaybeAppPasswordPermissions} from '#/lib/strings/errors'
 import {logger} from '#/logger'
 import {
   type SessionAccount,
-  useAgent,
+  usePdsClient,
   useSession,
   useSessionApi,
 } from '#/state/session'
@@ -25,6 +26,7 @@ import * as Layout from '#/components/Layout'
 import {Loader} from '#/components/Loader'
 import {Text} from '#/components/Typography'
 import {IS_WEB} from '#/env'
+import {com} from '#/lexicons'
 
 const COL_WIDTH = 400
 
@@ -36,8 +38,8 @@ export function Deactivated() {
   const {onPressSwitchAccount, pendingDid} = useAccountSwitcher()
   const {setShowLoggedOut} = useLoggedOutViewControls()
   const hasOtherAccounts = accounts.length > 1
-  const {logoutCurrentAccount} = useSessionApi()
-  const agent = useAgent()
+  const {logoutCurrentAccount, refreshSession} = useSessionApi()
+  const pdsClient = usePdsClient()
   const [pending, setPending] = useState(false)
   const [error, setError] = useState<string | undefined>()
   const queryClient = useQueryClient()
@@ -70,21 +72,24 @@ export function Deactivated() {
   const handleActivate = useCallback(async () => {
     try {
       setPending(true)
-      await agent.com.atproto.server.activateAccount()
+      await pdsClient.call(com.atproto.server.activateAccount)
       await queryClient.resetQueries()
-      await agent.resumeSession(agent.session!)
+      await refreshSession()
     } catch (e: any) {
-      switch (e.message) {
-        case 'Bad token scope':
-          setError(
-            _(
-              msg`You're signed in with an App Password. Please sign in with your main password to continue deactivating your account.`,
-            ),
-          )
-          break
-        default:
-          setError(_(msg`Something went wrong, please try again`))
-          break
+      /*
+       * `activateAccount` declares no lexicon errors, so the app-password case
+       * arrives as an undeclared code plus a message. The shared helper matches
+       * both that and the plain-string form the old exact `e.message` switch
+       * relied on.
+       */
+      if (isErrorMaybeAppPasswordPermissions(e)) {
+        setError(
+          _(
+            msg`You're signed in with an App Password. Please sign in with your main password to continue deactivating your account.`,
+          ),
+        )
+      } else {
+        setError(_(msg`Something went wrong, please try again`))
       }
 
       logger.error(e, {
@@ -93,7 +98,7 @@ export function Deactivated() {
     } finally {
       setPending(false)
     }
-  }, [_, agent, setPending, setError, queryClient])
+  }, [_, pdsClient, refreshSession, setPending, setError, queryClient])
 
   return (
     <View style={[a.util_screen_outer, a.flex_1]}>

src/screens/Settings/components/ChangeHandleDialog.tsx

@@ -27,7 +27,7 @@ import {useFetchDid, useUpdateHandleMutation} from '#/state/queries/handle'
 import {RQKEY as RQKEY_PROFILE} from '#/state/queries/profile'
 import {useServiceQuery} from '#/state/queries/service'
 import {useCurrentAccountProfile} from '#/state/queries/useCurrentAccountProfile'
-import {useAgent, useSession} from '#/state/session'
+import {useAgent, useSession, useSessionApi} from '#/state/session'
 import {ErrorScreen} from '#/view/com/util/error/ErrorScreen'
 import {atoms as a, native, useBreakpoints, useTheme} from '#/alf'
 import {Admonition} from '#/components/Admonition'
@@ -152,7 +152,7 @@ function ProvidedHandlePage({
 }) {
   const {_} = useLingui()
   const [subdomain, setSubdomain] = useState('')
-  const agent = useAgent()
+  const {refreshSession} = useSessionApi()
   const control = Dialog.useDialogContext()
   const {currentAccount} = useSession()
   const queryClient = useQueryClient()
@@ -173,7 +173,7 @@ function ProvidedHandlePage({
           queryKey: RQKEY_PROFILE(currentAccount.did),
         })
       }
-      agent.resumeSession(agent.session!).then(() => control.close())
+      refreshSession().then(() => control.close())
     },
   })
 
@@ -311,7 +311,7 @@ function OwnHandlePage({goToServiceHandle}: {goToServiceHandle: () => void}) {
   const {currentAccount} = useSession()
   const [dnsPanel, setDNSPanel] = useState(true)
   const [domain, setDomain] = useState('')
-  const agent = useAgent()
+  const {refreshSession} = useSessionApi()
   const control = Dialog.useDialogContext()
   const fetchDid = useFetchDid()
   const queryClient = useQueryClient()
@@ -328,7 +328,7 @@ function OwnHandlePage({goToServiceHandle}: {goToServiceHandle: () => void}) {
           queryKey: RQKEY_PROFILE(currentAccount.did),
         })
       }
-      agent.resumeSession(agent.session!).then(() => control.close())
+      refreshSession().then(() => control.close())
     },
   })
 

src/screens/Settings/components/DisableEmail2FADialog.tsx

@@ -5,7 +5,8 @@ import {useLingui} from '@lingui/react'
 import {Trans} from '@lingui/react/macro'
 
 import {cleanError} from '#/lib/strings/errors'
-import {useAgent, useSession} from '#/state/session'
+import {matchXrpcError} from '#/lib/xrpc-error'
+import {usePdsClient, useSession, useSessionApi} from '#/state/session'
 import {ErrorMessage} from '#/view/com/util/error/ErrorMessage'
 import {atoms as a, useBreakpoints, useTheme} from '#/alf'
 import {Button, ButtonIcon, ButtonText} from '#/components/Button'
@@ -16,6 +17,7 @@ import {Loader} from '#/components/Loader'
 import * as Toast from '#/components/Toast'
 import {P, Text} from '#/components/Typography'
 import {IS_NATIVE} from '#/env'
+import {com} from '#/lexicons'
 
 enum Stages {
   Email,
@@ -31,7 +33,8 @@ export function DisableEmail2FADialog({
   const t = useTheme()
   const {gtMobile} = useBreakpoints()
   const {currentAccount} = useSession()
-  const agent = useAgent()
+  const pdsClient = usePdsClient()
+  const {refreshSession} = useSessionApi()
 
   const [stage, setStage] = useState<Stages>(Stages.Email)
   const [confirmationCode, setConfirmationCode] = useState<string>('')
@@ -42,7 +45,7 @@ export function DisableEmail2FADialog({
     setError('')
     setIsProcessing(true)
     try {
-      await agent.com.atproto.server.requestEmailUpdate()
+      await pdsClient.call(com.atproto.server.requestEmailUpdate)
       setStage(Stages.ConfirmCode)
     } catch (e) {
       setError(cleanError(String(e)))
@@ -56,21 +59,26 @@ export function DisableEmail2FADialog({
     setIsProcessing(true)
     try {
       if (currentAccount?.email) {
-        await agent.com.atproto.server.updateEmail({
+        await pdsClient.call(com.atproto.server.updateEmail, {
           email: currentAccount.email,
           token: confirmationCode.trim(),
           emailAuthFactor: false,
         })
-        await agent.resumeSession(agent.session!)
+        await refreshSession()
         Toast.show(_(msg({message: 'Email 2FA disabled', context: 'toast'})))
       }
       control.close()
     } catch (e) {
-      const errMsg = String(e)
-      if (errMsg.includes('Token is invalid')) {
+      /*
+       * The old check matched the PDS message "Token is invalid"; the lexicon
+       * declares that case as `InvalidToken`, so match the code instead.
+       */
+      if (
+        matchXrpcError(e, com.atproto.server.updateEmail) === 'InvalidToken'
+      ) {
         setError(_(msg`Invalid 2FA confirmation code.`))
       } else {
-        setError(cleanError(errMsg))
+        setError(cleanError(e))
       }
     } finally {
       setIsProcessing(false)

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

@@ -0,0 +1,204 @@
+import {PasswordSession} from '@atproto/lex-password-session'
+import {beforeEach, describe, expect, it, jest} from '@jest/globals'
+import {act, render} from '@testing-library/react-native'
+
+import {type SessionAccount} from '../types'
+
+/*
+ * The provider pulls the whole app shell in through `#/state/util` and the
+ * account factories. These mocks cut the tree back to the session lifecycle
+ * itself, mirroring provider-clients-test.tsx.
+ */
+jest.mock('#/state/persisted', () => {
+  const {
+    defaults,
+  }: typeof import('#/state/persisted/schema') = require('#/state/persisted/schema')
+  return {
+    defaults,
+    get: (key: keyof typeof defaults) => defaults[key],
+    write: () => Promise.resolve(),
+    readLatest: (key: keyof typeof defaults) => defaults[key],
+    onUpdate: () => () => {},
+  }
+})
+jest.mock('#/state/util', () => ({useCloseAllActiveElements: () => () => {}}))
+jest.mock('#/components/dialogs/Context', () => ({
+  useGlobalDialogsControlContext: () => ({signinDialogControl: {open() {}}}),
+}))
+jest.mock('#/analytics', () => ({
+  AnalyticsContext: ({children}: {children: React.ReactNode}) => children,
+  useAnalyticsBase: () => ({metric() {}, logger: {debug() {}, error() {}}}),
+  utils: {accountToSessionMetadata: () => ({}), useMeta: () => undefined},
+}))
+jest.mock('#/state/shell/onboarding', () => ({
+  useOnboardingDispatch: () => () => {},
+}))
+jest.mock('#/ageAssurance/data', () => ({
+  clearAgeAssuranceServerDataForAll: () => {},
+  clearAgeAssuranceServerDataForDid: () => {},
+}))
+jest.mock('#/lib/persisted-query-storage', () => ({
+  clearPersistedQueryStorage: () => Promise.resolve(),
+}))
+jest.mock('#/lib/notifications/notifications', () => ({
+  unregisterPushToken: () => Promise.resolve(),
+}))
+jest.mock('jwt-decode', () => ({
+  jwtDecode: () => ({scope: 'com.atproto.access'}),
+}))
+jest.mock('#/state/events', () => ({
+  emitSessionDropped: () => {},
+  emitNetworkConfirmed: () => {},
+  emitNetworkLost: () => {},
+}))
+
+const mockLogin = jest.fn<(...args: unknown[]) => Promise<unknown>>()
+jest.mock('../session-core', () => ({
+  ...jest.requireActual<object>('../session-core'),
+  createSessionBundleAndLogin: (...args: unknown[]) => mockLogin(...args),
+}))
+jest.mock('../create-account', () => ({
+  createSessionBundleAndCreateAccount: () => new Promise(() => {}),
+}))
+
+import {Provider, useSession, useSessionApi} from '#/state/session'
+import {type SessionApiContext} from '#/state/session/types'
+import {BskyAppAgent, PasswordSessionManager} from '../bridge-agent'
+import {type SessionBundle} from '../session-core'
+import {sessionAccountToSessionData} from '../session-data'
+import {
+  asFetch,
+  DID,
+  HANDLE,
+  json,
+  makeAccount,
+  makeMockFetch,
+  type MockFetch,
+} from './mock-fetch'
+
+/**
+ * Build a bundle whose session is a real `PasswordSession` over the stubbed
+ * network, since `refreshSession` drives the session's own refresh machinery.
+ */
+function makeBundle(
+  account: SessionAccount,
+  fetchMock: MockFetch,
+): SessionBundle {
+  const session = new PasswordSession(sessionAccountToSessionData(account), {
+    fetch: asFetch(fetchMock),
+  })
+  const manager = new PasswordSessionManager(session, {
+    service: account.service,
+  })
+  manager.setFetch(asFetch(fetchMock))
+  return {
+    session,
+    agent: new BskyAppAgent(manager),
+    service: new URL(account.service),
+  }
+}
+
+type Harness = {
+  api: SessionApiContext
+  currentAccount: () => SessionAccount | undefined
+}
+
+function renderProvider(): Harness {
+  let api!: SessionApiContext
+  let currentAccount: SessionAccount | undefined
+  function Probe() {
+    api = useSessionApi()
+    currentAccount = useSession().currentAccount
+    return null
+  }
+  render(
+    <Provider>
+      <Probe />
+    </Provider>,
+  )
+  return {api, currentAccount: () => currentAccount}
+}
+
+/** Render the provider and log `account` in through the stubbed login factory. */
+async function renderLoggedIn(
+  account: SessionAccount,
+  fetchMock: MockFetch,
+): Promise<Harness> {
+  const bundle = makeBundle(account, fetchMock)
+  const harness = renderProvider()
+  mockLogin.mockResolvedValueOnce({bundle, account})
+  await act(async () => {
+    await harness.api.login({} as never, 'LoginForm')
+  })
+  return harness
+}
+
+beforeEach(() => {
+  mockLogin.mockReset()
+})
+
+describe('refreshSession', () => {
+  it('resolves with the rotated account snapshot', async () => {
+    const fetchMock = makeMockFetch()
+    const {api} = await renderLoggedIn(makeAccount(), fetchMock)
+
+    let refreshed: SessionAccount | undefined
+    await act(async () => {
+      refreshed = await api.refreshSession()
+    })
+
+    /* the mock's refresh response rotates both tokens */
+    expect(refreshed?.accessJwt).toBe('access-jwt-2')
+    expect(refreshed?.refreshJwt).toBe('refresh-jwt-2')
+    expect(refreshed?.did).toBe(DID)
+    expect(refreshed?.handle).toBe(HANDLE)
+  })
+
+  it('exposes the fresh tokens before the store has caught up', async () => {
+    const fetchMock = makeMockFetch()
+    const {api, currentAccount} = await renderLoggedIn(makeAccount(), fetchMock)
+
+    /*
+     * The point of the return value: `SignupQueued` branches on the fresh
+     * accessJwt synchronously, without waiting for `onUpdated` -> dispatch ->
+     * re-render.
+     */
+    let refreshed: SessionAccount | undefined
+    const before = currentAccount()?.accessJwt
+    await act(async () => {
+      refreshed = await api.refreshSession()
+    })
+    expect(before).toBe('access-jwt')
+    expect(refreshed?.accessJwt).toBe('access-jwt-2')
+  })
+
+  it('resolves with undefined when logged out', async () => {
+    const {api} = renderProvider()
+
+    let refreshed: SessionAccount | undefined = makeAccount()
+    await act(async () => {
+      refreshed = await api.refreshSession()
+    })
+
+    expect(refreshed).toBeUndefined()
+  })
+
+  it('rejects when the refresh rotated nothing', async () => {
+    /*
+     * A transient failure: `PasswordSession.refresh()` reports through
+     * `onUpdateFailure` and resolves with the SAME data object. Callers read
+     * resolution as "tokens rotated", so this must reject.
+     */
+    const fetchMock = makeMockFetch({
+      'com.atproto.server.refreshSession': () =>
+        json({error: 'InternalServerError'}, 500),
+    })
+    const {api} = await renderLoggedIn(makeAccount(), fetchMock)
+
+    await expect(
+      act(async () => {
+        await api.refreshSession()
+      }),
+    ).rejects.toThrow('Failed to refresh session')
+  })
+})

src/state/session/index.tsx

@@ -82,6 +82,7 @@ const ApiContext = createContext<SessionApiContext>({
   resumeSession: async () => {},
   removeAccount: () => {},
   partialRefreshSession: async () => {},
+  refreshSession: () => Promise.resolve(undefined),
 })
 ApiContext.displayName = 'SessionApiContext'
 
@@ -473,6 +474,50 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     })
   }, [store, cancelPendingTask])
 
+  /**
+   * Rotate the session's tokens and hand back the resulting account snapshot.
+   *
+   * Rejects when the rotation was a no-op, restoring the contract the
+   * `agent.resumeSession(agent.session!)` call sites were written against (the
+   * bridge agent's `refreshSession` override does the same, for the same
+   * reason). `PasswordSession.refresh()` resolves with the
+   * unchanged `SessionData` on a transient failure - a 500 or a network error
+   * reported through `onUpdateFailure` - and reserves rejection for a
+   * definitively dead session. Callers here all read resolution as "tokens
+   * rotated": the verification dialogs close, `Deactivated` clears its error
+   * state, and `SignupQueued` re-checks the token scope, so a resolved no-op
+   * would report success or loop silently. Identity, not a field comparison, is
+   * the signal: `PasswordSession` allocates a new object per successful
+   * rotation and returns the existing one untouched otherwise. Capturing the
+   * data immediately before the call also handles concurrent refreshes, since a
+   * rotation another caller's queued refresh performed still differs from what
+   * we captured.
+   *
+   * Like {@link partialRefreshSession}, the bundle comes from
+   * `store.getState()` rather than the render's `state`: a dispatch landing
+   * before the next render would otherwise leave this holding a disposed
+   * bundle, and reading live also keeps the callback's identity stable across
+   * unrelated state updates.
+   */
+  const refreshSession = useCallback<
+    SessionApiContext['refreshSession']
+  >(async () => {
+    const bundle = store.getState().currentBundleState.bundle as unknown as
+      | SessionBundle
+      | PublicSessionBundle
+    if (!bundle.session) return undefined // logged out: nothing to refresh
+    const before = bundle.session.session
+    const after = await bundle.session.refresh()
+    if (after === before) {
+      throw new Error('Failed to refresh session')
+    }
+    /*
+     * The session's `onUpdated` hook dispatches the new tokens into the store,
+     * but that lands a render away; this snapshot exposes them immediately.
+     */
+    return sessionDataToSessionAccount(after, after.service)
+  }, [store])
+
   const removeAccount = useCallback<SessionApiContext['removeAccount']>(
     account => {
       addSessionDebugLog({
@@ -607,6 +652,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       resumeSession,
       removeAccount,
       partialRefreshSession,
+      refreshSession,
     }),
     [
       createAccount,
@@ -616,6 +662,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       resumeSession,
       removeAccount,
       partialRefreshSession,
+      refreshSession,
     ],
   )
 

src/state/session/types.ts

@@ -52,4 +52,17 @@ export type SessionApiContext = {
    * so it produces no session-change side effects.
    */
   partialRefreshSession: () => Promise<void>
+  /**
+   * Rotates the session's tokens and resolves with the resulting account
+   * snapshot, or `undefined` when logged out.
+   *
+   * Rejects when nothing was rotated, so a resolved promise means "tokens
+   * rotated". Every caller relies on that: the verification dialogs close on
+   * resolution, and `SignupQueued` re-checks the token scope.
+   *
+   * The snapshot is returned rather than read off `currentAccount`, because the
+   * session's `onUpdated` hook -> `store.dispatch` path is a render cycle away
+   * and `SignupQueued` branches synchronously on the fresh `accessJwt`.
+   */
+  refreshSession: () => Promise<SessionAccount | undefined>
 }