[SDK] Add refreshSession and migrate the session-pinned infra (#11381)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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/__tests__/session-test.ts

@@ -1,4 +1,3 @@
-import {type AtpAgent} from '@atproto/api'
 import {type SessionData} from '@atproto/lex-password-session'
 import {describe, expect, it, jest} from '@jest/globals'
 
@@ -18,21 +17,21 @@ jest.mock('../../../ageAssurance/state', () => ({
   unsafeGetAndComputeAgeAssurance: () => ({state: {}}),
 }))
 jest.mock('#/lib/notifications/notifications', () => ({
-  unregisterPushToken(_agents: AtpAgent[]) {
+  unregisterPushToken(_clients: unknown[]) {
     return Promise.resolve()
   },
 }))
 /*
  * The logout and account-removal reducer cases fire a push-token side effect
- * whose first step, `createTemporaryAgentsAndResume`, builds real `AtpAgent`s
- * and resumes them over the real network. Under jest that request outlives the
+ * whose first step, `createTemporaryClientsAndResume`, resumes real
+ * `PasswordSession`s over the real network. Under jest that request outlives the
  * suite: it rejects after teardown, and the resulting `logger.error` reaches
  * for `nanoid` in an environment that no longer has it, failing whichever suite
  * happens to be running at that moment. Stubbing the module keeps the side
  * effect synchronous and offline.
  */
 jest.mock('../util', () => ({
-  createTemporaryAgentsAndResume: () => Promise.resolve([]),
+  createTemporaryClientsAndResume: () => Promise.resolve([]),
 }))
 
 // Reuse a bundle within each test: session events are scoped by bundle identity.

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/reducer.ts

@@ -3,7 +3,7 @@ import {logger} from '#/lib/notifications/util'
 import {wrapSessionReducerForLogging} from './logging'
 import {createPublicSessionBundle} from './session-core'
 import {type AtpSessionEvent, type SessionAccount} from './types'
-import {createTemporaryAgentsAndResume} from './util'
+import {createTemporaryClientsAndResume} from './util'
 
 // Keep session internals outside the reducer's static view of a bundle.
 type OpaqueSessionBundle = {
@@ -155,8 +155,8 @@ let reducer = (state: State, action: Action): State => {
       // side effect
       const account = state.accounts.find(a => a.did === accountDid)
       if (account) {
-        createTemporaryAgentsAndResume([account])
-          .then(agents => unregisterPushToken(agents))
+        createTemporaryClientsAndResume([account])
+          .then(clients => unregisterPushToken(clients))
           .then(() =>
             logger.debug('Push token unregistered', {did: accountDid}),
           )
@@ -183,8 +183,8 @@ let reducer = (state: State, action: Action): State => {
       // side effect
       const account = state.accounts.find(a => a.did === accountDid)
       if (account && accountDid) {
-        createTemporaryAgentsAndResume([account])
-          .then(agents => unregisterPushToken(agents))
+        createTemporaryClientsAndResume([account])
+          .then(clients => unregisterPushToken(clients))
           .then(() =>
             logger.debug('Push token unregistered', {did: accountDid}),
           )
@@ -211,8 +211,8 @@ let reducer = (state: State, action: Action): State => {
       }
     }
     case 'logged-out-every-account': {
-      createTemporaryAgentsAndResume(state.accounts)
-        .then(agents => unregisterPushToken(agents))
+      createTemporaryClientsAndResume(state.accounts)
+        .then(clients => unregisterPushToken(clients))
         .then(() => logger.debug('Push token unregistered'))
         .catch(err => {
           logger.error('Failed to unregister push token', {

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>
 }

src/state/session/util.ts

@@ -1,7 +1,10 @@
-import AtpAgent from '@atproto/api'
+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 {sessionAccountToSession} from './session-data'
+import {networkAwareFetch} from './network'
+import {sessionAccountToSessionData} from './session-data'
 import {type SessionAccount} from './types'
 
 export {isSessionExpired, isSignupQueued} from './session-data'
@@ -12,30 +15,41 @@ export function readLastActiveAccount() {
 }
 
 /**
- * Creates and attempted to resumeSession for every stored session.
- * Intended to be used to send push token revokations just before logout.
+ * Resume a single-use session per stored account, for the push-token revocation
+ * sent just before logout.
+ *
+ * The sessions carry no lifecycle hooks - no `onUpdated`, no `onDeleted` - so a
+ * rotation one of them performs can neither persist over nor race the live
+ * session's tokens. That isolation is load-bearing: each exists only long enough
+ * to authenticate one `unregisterPush` call.
+ *
+ * PDS routing is left to the session rather than pinned from the stored
+ * `pdsUrl`, because `resume` refreshes (and fills in a missing didDoc from
+ * `getSession`) before the client issues anything, so the request already goes
+ * to the didDoc PDS.
+ *
+ * `resume` rejects only when a session is definitively dead; a transient network
+ * failure resolves with the stored tokens, which are the same ones the old agent
+ * path would have sent. Definitively dead sessions drop out of the settled list.
  */
-export async function createTemporaryAgentsAndResume(
+export async function createTemporaryClientsAndResume(
   accounts: SessionAccount[],
-) {
-  const agents = await Promise.allSettled(
+): Promise<TemporaryPushClient[]> {
+  const settled = await Promise.allSettled(
     accounts.map(async account => {
-      const agent: AtpAgent = new AtpAgent({service: account.service})
-      if (account.pdsUrl) {
-        agent.sessionManager.pdsUrl = new URL(account.pdsUrl)
-      }
-
-      const session = sessionAccountToSession(account)
-      const res = await agent.resumeSession(session)
-      if (!res.success) throw new Error('Failed to resume session')
-
-      agent.assertAuthenticated() // confirm auth success
-
-      return agent
+      const session = await PasswordSession.resume(
+        sessionAccountToSessionData(account),
+        {fetch: networkAwareFetch},
+      )
+      return {
+        client: createLexClient(session),
+        service: session.session.service,
+        handle: session.session.handle,
+      } satisfies TemporaryPushClient
     }),
   )
 
-  return agents
+  return settled
     .filter(x => x.status === 'fulfilled')
     .map(promise => promise.value)
 }