add refreshSession to the session api and migrate its callers

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