Upgrade API, implement XRPC rework (#4857)

Co-authored-by: Matthieu Sieben <matthieu.sieben@gmail.com>
2024-08-12 14:00:15 -07:00 · 2024-08-12 14:00:15 -07:00 · 7df2327424
commit 7df2327424
parent ae883e2df7
19 changed files with 543 additions and 360 deletions
--- a/src/lib/api/api-polyfill.ts
+++ b/src/lib/api/api-polyfill.ts
@ -1,85 +0,0 @@
-import RNFS from 'react-native-fs'
-import {BskyAgent, jsonToLex, stringifyLex} from '@atproto/api'
-
-const GET_TIMEOUT = 15e3 // 15s
-const POST_TIMEOUT = 60e3 // 60s
-
-export function doPolyfill() {
-  BskyAgent.configure({fetch: fetchHandler})
-}
-
-interface FetchHandlerResponse {
-  status: number
-  headers: Record<string, string>
-  body: any
-}
-
-async function fetchHandler(
-  reqUri: string,
-  reqMethod: string,
-  reqHeaders: Record<string, string>,
-  reqBody: any,
-): Promise<FetchHandlerResponse> {
-  const reqMimeType = reqHeaders['Content-Type'] || reqHeaders['content-type']
-  if (reqMimeType && reqMimeType.startsWith('application/json')) {
-    reqBody = stringifyLex(reqBody)
-  } else if (
-    typeof reqBody === 'string' &&
-    (reqBody.startsWith('/') || reqBody.startsWith('file:'))
-  ) {
-    if (reqBody.endsWith('.jpeg') || reqBody.endsWith('.jpg')) {
-      // HACK
-      // React native has a bug that inflates the size of jpegs on upload
-      // we get around that by renaming the file ext to .bin
-      // see https://github.com/facebook/react-native/issues/27099
-      // -prf
-      const newPath = reqBody.replace(/\.jpe?g$/, '.bin')
-      await RNFS.moveFile(reqBody, newPath)
-      reqBody = newPath
-    }
-    // NOTE
-    // React native treats bodies with {uri: string} as file uploads to pull from cache
-    // -prf
-    reqBody = {uri: reqBody}
-  }
-
-  const controller = new AbortController()
-  const to = setTimeout(
-    () => controller.abort(),
-    reqMethod === 'post' ? POST_TIMEOUT : GET_TIMEOUT,
-  )
-
-  const res = await fetch(reqUri, {
-    method: reqMethod,
-    headers: reqHeaders,
-    body: reqBody,
-    signal: controller.signal,
-  })
-
-  const resStatus = res.status
-  const resHeaders: Record<string, string> = {}
-  res.headers.forEach((value: string, key: string) => {
-    resHeaders[key] = value
-  })
-  const resMimeType = resHeaders['Content-Type'] || resHeaders['content-type']
-  let resBody
-  if (resMimeType) {
-    if (resMimeType.startsWith('application/json')) {
-      resBody = jsonToLex(await res.json())
-    } else if (resMimeType.startsWith('text/')) {
-      resBody = await res.text()
-    } else if (resMimeType === 'application/vnd.ipld.car') {
-      resBody = await res.arrayBuffer()
-    } else {
-      throw new Error('Non-supported mime type')
-    }
-  }
-
-  clearTimeout(to)
-
-  return {
-    status: resStatus,
-    headers: resHeaders,
-    body: resBody,
-  }
-}
--- a/src/lib/api/api-polyfill.web.ts
+++ b/src/lib/api/api-polyfill.web.ts
@ -1,3 +0,0 @@
-export function doPolyfill() {
-  // no polyfill is needed on web
-}
--- a/src/lib/api/feed/custom.ts
+++ b/src/lib/api/feed/custom.ts
@ -1,7 +1,6 @@
 import {
  AppBskyFeedDefs,
  AppBskyFeedGetFeed as GetCustomFeed,
-  AtpAgent,
  BskyAgent,
 } from '@atproto/api'

@ -51,7 +50,7 @@ export class CustomFeedAPI implements FeedAPI {
    const agent = this.agent
    const isBlueskyOwned = isBlueskyOwnedFeed(this.params.feed)

-    const res = agent.session
+    const res = agent.did
      ? await this.agent.app.bsky.feed.getFeed(
          {
            ...this.params,
@ -106,34 +105,32 @@ async function loggedOutFetch({
  let contentLangs = getContentLanguages().join(',')

  // manually construct fetch call so we can add the `lang` cache-busting param
-  let res = await AtpAgent.fetch!(
+  let res = await fetch(
    `https://api.bsky.app/xrpc/app.bsky.feed.getFeed?feed=${feed}${
      cursor ? `&cursor=${cursor}` : ''
    }&limit=${limit}&lang=${contentLangs}`,
-    'GET',
-    {'Accept-Language': contentLangs},
-    undefined,
+    {method: 'GET', headers: {'Accept-Language': contentLangs}},
  )
-  if (res.body?.feed?.length) {
+  let data = res.ok ? await res.json() : null
+  if (data?.feed?.length) {
    return {
      success: true,
-      data: res.body,
+      data,
    }
  }

  // no data, try again with language headers removed
-  res = await AtpAgent.fetch!(
+  res = await fetch(
    `https://api.bsky.app/xrpc/app.bsky.feed.getFeed?feed=${feed}${
      cursor ? `&cursor=${cursor}` : ''
    }&limit=${limit}`,
-    'GET',
-    {'Accept-Language': ''},
-    undefined,
+    {method: 'GET', headers: {'Accept-Language': ''}},
  )
-  if (res.body?.feed?.length) {
+  data = res.ok ? await res.json() : null
+  if (data?.feed?.length) {
    return {
      success: true,
-      data: res.body,
+      data,
    }
  }

--- a/src/lib/api/index.ts
+++ b/src/lib/api/index.ts
@ -6,7 +6,6 @@ import {
  AppBskyFeedThreadgate,
  BskyAgent,
  ComAtprotoLabelDefs,
-  ComAtprotoRepoUploadBlob,
  RichText,
 } from '@atproto/api'
 import {AtUri} from '@atproto/api'
@ -15,10 +14,13 @@ import {logger} from '#/logger'
 import {ThreadgateSetting} from '#/state/queries/threadgate'
 import {isNetworkError} from 'lib/strings/errors'
 import {shortenLinks, stripInvalidMentions} from 'lib/strings/rich-text-manip'
-import {isNative, isWeb} from 'platform/detection'
+import {isNative} from 'platform/detection'
 import {ImageModel} from 'state/models/media/image'
 import {LinkMeta} from '../link-meta/link-meta'
 import {safeDeleteAsync} from '../media/manip'
+import {uploadBlob} from './upload-blob'
+
+export {uploadBlob}

 export interface ExternalEmbedDraft {
  uri: string
@ -28,25 +30,6 @@ export interface ExternalEmbedDraft {
  localThumb?: ImageModel
 }

-export async function uploadBlob(
-  agent: BskyAgent,
-  blob: string,
-  encoding: string,
-): Promise<ComAtprotoRepoUploadBlob.Response> {
-  if (isWeb) {
-    // `blob` should be a data uri
-    return agent.uploadBlob(convertDataURIToUint8Array(blob), {
-      encoding,
-    })
-  } else {
-    // `blob` should be a path to a file in the local FS
-    return agent.uploadBlob(
-      blob, // this will be special-cased by the fetch monkeypatch in /src/state/lib/api.ts
-      {encoding},
-    )
-  }
-}
-
 interface PostOpts {
  rawText: string
  replyTo?: string
@ -301,7 +284,7 @@ export async function createThreadgate(

  const postUrip = new AtUri(postUri)
  await agent.api.com.atproto.repo.putRecord({
-    repo: agent.session!.did,
+    repo: agent.accountDid,
    collection: 'app.bsky.feed.threadgate',
    rkey: postUrip.rkey,
    record: {
@ -312,15 +295,3 @@ export async function createThreadgate(
    },
  })
 }
-
-// helpers
-// =
-
-function convertDataURIToUint8Array(uri: string): Uint8Array {
-  var raw = window.atob(uri.substring(uri.indexOf(';base64,') + 8))
-  var binary = new Uint8Array(new ArrayBuffer(raw.length))
-  for (let i = 0; i < raw.length; i++) {
-    binary[i] = raw.charCodeAt(i)
-  }
-  return binary
-}
--- a/src/lib/api/upload-blob.ts
+++ b/src/lib/api/upload-blob.ts
@ -0,0 +1,82 @@
+import RNFS from 'react-native-fs'
+import {BskyAgent, ComAtprotoRepoUploadBlob} from '@atproto/api'
+
+/**
+ * @param encoding Allows overriding the blob's type
+ */
+export async function uploadBlob(
+  agent: BskyAgent,
+  input: string | Blob,
+  encoding?: string,
+): Promise<ComAtprotoRepoUploadBlob.Response> {
+  if (typeof input === 'string' && input.startsWith('file:')) {
+    const blob = await asBlob(input)
+    return agent.uploadBlob(blob, {encoding})
+  }
+
+  if (typeof input === 'string' && input.startsWith('/')) {
+    const blob = await asBlob(`file://${input}`)
+    return agent.uploadBlob(blob, {encoding})
+  }
+
+  if (typeof input === 'string' && input.startsWith('data:')) {
+    const blob = await fetch(input).then(r => r.blob())
+    return agent.uploadBlob(blob, {encoding})
+  }
+
+  if (input instanceof Blob) {
+    return agent.uploadBlob(input, {encoding})
+  }
+
+  throw new TypeError(`Invalid uploadBlob input: ${typeof input}`)
+}
+
+async function asBlob(uri: string): Promise<Blob> {
+  return withSafeFile(uri, async safeUri => {
+    // Note
+    // Android does not support `fetch()` on `file://` URIs. for this reason, we
+    // use XMLHttpRequest instead of simply calling:
+
+    // return fetch(safeUri.replace('file:///', 'file:/')).then(r => r.blob())
+
+    return await new Promise((resolve, reject) => {
+      const xhr = new XMLHttpRequest()
+      xhr.onload = () => resolve(xhr.response)
+      xhr.onerror = () => reject(new Error('Failed to load blob'))
+      xhr.responseType = 'blob'
+      xhr.open('GET', safeUri, true)
+      xhr.send(null)
+    })
+  })
+}
+
+// HACK
+// React native has a bug that inflates the size of jpegs on upload
+// we get around that by renaming the file ext to .bin
+// see https://github.com/facebook/react-native/issues/27099
+// -prf
+async function withSafeFile<T>(
+  uri: string,
+  fn: (path: string) => Promise<T>,
+): Promise<T> {
+  if (uri.endsWith('.jpeg') || uri.endsWith('.jpg')) {
+    // Since we don't "own" the file, we should avoid renaming or modifying it.
+    // Instead, let's copy it to a temporary file and use that (then remove the
+    // temporary file).
+    const newPath = uri.replace(/\.jpe?g$/, '.bin')
+    try {
+      await RNFS.copyFile(uri, newPath)
+    } catch {
+      // Failed to copy the file, just use the original
+      return await fn(uri)
+    }
+    try {
+      return await fn(newPath)
+    } finally {
+      // Remove the temporary file
+      await RNFS.unlink(newPath)
+    }
+  } else {
+    return fn(uri)
+  }
+}
--- a/src/lib/api/upload-blob.web.ts
+++ b/src/lib/api/upload-blob.web.ts
@ -0,0 +1,26 @@
+import {BskyAgent, ComAtprotoRepoUploadBlob} from '@atproto/api'
+
+/**
+ * @note It is recommended, on web, to use the `file` instance of the file
+ * selector input element, rather than a `data:` URL, to avoid
+ * loading the file into memory. `File` extends `Blob` "file" instances can
+ * be passed directly to this function.
+ */
+export async function uploadBlob(
+  agent: BskyAgent,
+  input: string | Blob,
+  encoding?: string,
+): Promise<ComAtprotoRepoUploadBlob.Response> {
+  if (typeof input === 'string' && input.startsWith('data:')) {
+    const blob = await fetch(input).then(r => r.blob())
+    return agent.uploadBlob(blob, {encoding})
+  }
+
+  if (input instanceof Blob) {
+    return agent.uploadBlob(input, {
+      encoding,
+    })
+  }
+
+  throw new TypeError(`Invalid uploadBlob input: ${typeof input}`)
+}
--- a/src/lib/media/manip.ts
+++ b/src/lib/media/manip.ts
@ -218,13 +218,7 @@ export async function safeDeleteAsync(path: string) {
  // Normalize is necessary for Android, otherwise it doesn't delete.
  const normalizedPath = normalizePath(path)
  try {
-    await Promise.allSettled([
-      deleteAsync(normalizedPath, {idempotent: true}),
-      // HACK: Try this one too. Might exist due to api-polyfill hack.
-      deleteAsync(normalizedPath.replace(/\.jpe?g$/, '.bin'), {
-        idempotent: true,
-      }),
-    ])
+    await deleteAsync(normalizedPath, {idempotent: true})
  } catch (e) {
    console.error('Failed to delete file', e)
  }
--- a/src/screens/SignupQueued.tsx
+++ b/src/screens/SignupQueued.tsx
@ -40,7 +40,7 @@ export function SignupQueued() {
      const res = await agent.com.atproto.temp.checkSignupQueue()
      if (res.data.activated) {
        // ready to go, exchange the access token for a usable one and kick off onboarding
-        await agent.refreshSession()
+        await agent.sessionManager.refreshSession()
        if (!isSignupQueued(agent.session?.accessJwt)) {
          onboardingDispatch({type: 'start'})
        }
--- a/src/state/queries/preferences/index.ts
+++ b/src/state/queries/preferences/index.ts
@ -37,14 +37,14 @@ export function usePreferencesQuery() {
    refetchOnWindowFocus: true,
    queryKey: preferencesQueryKey,
    queryFn: async () => {
-      if (agent.session?.did === undefined) {
+      if (!agent.did) {
        return DEFAULT_LOGGED_OUT_PREFERENCES
      } else {
        const res = await agent.getPreferences()

        // save to local storage to ensure there are labels on initial requests
        saveLabelers(
-          agent.session.did,
+          agent.did,
          res.moderationPrefs.labelers.map(l => l.did),
        )

--- a/src/state/session/tests/session-test.ts
+++ b/src/state/session/tests/session-test.ts
@ -27,7 +27,7 @@ describe('session', () => {
    `)

    const agent = new BskyAgent({service: 'https://alice.com'})
-    agent.session = {
+    agent.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -118,7 +118,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -166,7 +166,7 @@ describe('session', () => {
    `)

    const agent2 = new BskyAgent({service: 'https://bob.com'})
-    agent2.session = {
+    agent2.sessionManager.session = {
      active: true,
      did: 'bob-did',
      handle: 'bob.test',
@ -230,7 +230,7 @@ describe('session', () => {
    `)

    const agent3 = new BskyAgent({service: 'https://alice.com'})
-    agent3.session = {
+    agent3.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice-updated.test',
@ -294,7 +294,7 @@ describe('session', () => {
    `)

    const agent4 = new BskyAgent({service: 'https://jay.com'})
-    agent4.session = {
+    agent4.sessionManager.session = {
      active: true,
      did: 'jay-did',
      handle: 'jay.test',
@ -445,7 +445,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -502,7 +502,7 @@ describe('session', () => {
    `)

    const agent2 = new BskyAgent({service: 'https://alice.com'})
-    agent2.session = {
+    agent2.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -553,7 +553,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -598,7 +598,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -606,7 +606,7 @@ describe('session', () => {
      refreshJwt: 'alice-refresh-jwt-1',
    }
    const agent2 = new BskyAgent({service: 'https://bob.com'})
-    agent2.session = {
+    agent2.sessionManager.session = {
      active: true,
      did: 'bob-did',
      handle: 'bob.test',
@ -678,7 +678,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -695,7 +695,7 @@ describe('session', () => {
    expect(state.accounts.length).toBe(1)
    expect(state.currentAgentState.did).toBe('alice-did')

-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice-updated.test',
@ -748,7 +748,7 @@ describe('session', () => {
      }
    `)

-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice-updated.test',
@ -801,7 +801,7 @@ describe('session', () => {
      }
    `)

-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice-updated.test',
@ -859,7 +859,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -876,7 +876,7 @@ describe('session', () => {
    expect(state.accounts.length).toBe(1)
    expect(state.currentAgentState.did).toBe('alice-did')

-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice-updated.test',
@ -907,7 +907,7 @@ describe('session', () => {
    ])
    expect(lastState === state).toBe(true)

-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice-updated.test',
@ -931,7 +931,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -940,7 +940,7 @@ describe('session', () => {
    }

    const agent2 = new BskyAgent({service: 'https://bob.com'})
-    agent2.session = {
+    agent2.sessionManager.session = {
      active: true,
      did: 'bob-did',
      handle: 'bob.test',
@ -965,7 +965,7 @@ describe('session', () => {
    expect(state.accounts.length).toBe(2)
    expect(state.currentAgentState.did).toBe('bob-did')

-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice-updated.test',
@ -1032,7 +1032,7 @@ describe('session', () => {
      }
    `)

-    agent2.session = {
+    agent2.sessionManager.session = {
      active: true,
      did: 'bob-did',
      handle: 'bob-updated.test',
@ -1099,7 +1099,7 @@ describe('session', () => {

    // Ignore other events for inactive agent.
    const lastState = state
-    agent1.session = undefined
+    agent1.sessionManager.session = undefined
    state = run(state, [
      {
        type: 'received-agent-event',
@ -1126,7 +1126,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -1135,7 +1135,7 @@ describe('session', () => {
    }

    const agent2 = new BskyAgent({service: 'https://bob.com'})
-    agent2.session = {
+    agent2.sessionManager.session = {
      active: true,
      did: 'bob-did',
      handle: 'bob.test',
@ -1162,7 +1162,7 @@ describe('session', () => {
    expect(state.accounts.length).toBe(1)
    expect(state.currentAgentState.did).toBe('bob-did')

-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -1188,7 +1188,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -1206,7 +1206,7 @@ describe('session', () => {
    expect(state.accounts.length).toBe(1)
    expect(state.currentAgentState.did).toBe('alice-did')

-    agent1.session = undefined
+    agent1.sessionManager.session = undefined
    state = run(state, [
      {
        type: 'received-agent-event',
@ -1255,7 +1255,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -1273,7 +1273,7 @@ describe('session', () => {
    expect(state.accounts[0].accessJwt).toBe('alice-access-jwt-1')
    expect(state.currentAgentState.did).toBe('alice-did')

-    agent1.session = undefined
+    agent1.sessionManager.session = undefined
    state = run(state, [
      {
        type: 'received-agent-event',
@ -1320,7 +1320,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -1338,7 +1338,7 @@ describe('session', () => {
    expect(state.accounts[0].accessJwt).toBe('alice-access-jwt-1')
    expect(state.currentAgentState.did).toBe('alice-did')

-    agent1.session = undefined
+    agent1.sessionManager.session = undefined
    state = run(state, [
      {
        type: 'received-agent-event',
@ -1385,7 +1385,7 @@ describe('session', () => {
    let state = getInitialState([])

    const agent1 = new BskyAgent({service: 'https://alice.com'})
-    agent1.session = {
+    agent1.sessionManager.session = {
      active: true,
      did: 'alice-did',
      handle: 'alice.test',
@ -1393,7 +1393,7 @@ describe('session', () => {
      refreshJwt: 'alice-refresh-jwt-1',
    }
    const agent2 = new BskyAgent({service: 'https://bob.com'})
-    agent2.session = {
+    agent2.sessionManager.session = {
      active: true,
      did: 'bob-did',
      handle: 'bob.test',
@ -1416,7 +1416,7 @@ describe('session', () => {
    expect(state.currentAgentState.did).toBe('bob-did')

    const anotherTabAgent1 = new BskyAgent({service: 'https://jay.com'})
-    anotherTabAgent1.session = {
+    anotherTabAgent1.sessionManager.session = {
      active: true,
      did: 'jay-did',
      handle: 'jay.test',
@ -1424,7 +1424,7 @@ describe('session', () => {
      refreshJwt: 'jay-refresh-jwt-1',
    }
    const anotherTabAgent2 = new BskyAgent({service: 'https://alice.com'})
-    anotherTabAgent2.session = {
+    anotherTabAgent2.sessionManager.session = {
      active: true,
      did: 'bob-did',
      handle: 'bob.test',
@ -1492,7 +1492,7 @@ describe('session', () => {
    `)

    const anotherTabAgent3 = new BskyAgent({service: 'https://clarence.com'})
-    anotherTabAgent3.session = {
+    anotherTabAgent3.sessionManager.session = {
      active: true,
      did: 'clarence-did',
      handle: 'clarence.test',
--- a/src/state/session/agent.ts
+++ b/src/state/session/agent.ts
@ -1,4 +1,9 @@
-import {AtpSessionData, AtpSessionEvent, BskyAgent} from '@atproto/api'
+import {
+  AtpPersistSessionHandler,
+  AtpSessionData,
+  AtpSessionEvent,
+  BskyAgent,
+} from '@atproto/api'
 import {TID} from '@atproto/common-web'

 import {networkRetry} from '#/lib/async/retry'
@ -20,6 +25,8 @@ import {
 import {SessionAccount} from './types'
 import {isSessionExpired, isSignupQueued} from './util'

+type SetPersistSessionHandler = (cb: AtpPersistSessionHandler) => void
+
 export function createPublicAgent() {
  configureModerationForGuest() // Side effect but only relevant for tests
  return new BskyAgent({service: PUBLIC_BSKY_SERVICE})
@ -32,10 +39,11 @@ export async function createAgentAndResume(
    did: string,
    event: AtpSessionEvent,
  ) => void,
+  setPersistSessionHandler: SetPersistSessionHandler,
 ) {
  const agent = new BskyAgent({service: storedAccount.service})
  if (storedAccount.pdsUrl) {
-    agent.pdsUrl = agent.api.xrpc.uri = new URL(storedAccount.pdsUrl)
+    agent.sessionManager.pdsUrl = new URL(storedAccount.pdsUrl)
  }
  const gates = tryFetchGates(storedAccount.did, 'prefer-low-latency')
  const moderation = configureModerationForAccount(agent, storedAccount)
@ -43,9 +51,8 @@ export async function createAgentAndResume(
  if (isSessionExpired(storedAccount)) {
    await networkRetry(1, () => agent.resumeSession(prevSession))
  } else {
-    agent.session = prevSession
+    agent.sessionManager.session = prevSession
    if (!storedAccount.signupQueued) {
-      // Intentionally not awaited to unblock the UI:
      networkRetry(3, () => agent.resumeSession(prevSession)).catch(
        (e: any) => {
          logger.error(`networkRetry failed to resume session`, {
@ -60,7 +67,13 @@ export async function createAgentAndResume(
    }
  }

-  return prepareAgent(agent, gates, moderation, onSessionChange)
+  return prepareAgent(
+    agent,
+    gates,
+    moderation,
+    onSessionChange,
+    setPersistSessionHandler,
+  )
 }

 export async function createAgentAndLogin(
@ -80,6 +93,7 @@ export async function createAgentAndLogin(
    did: string,
    event: AtpSessionEvent,
  ) => void,
+  setPersistSessionHandler: SetPersistSessionHandler,
 ) {
  const agent = new BskyAgent({service})
  await agent.login({identifier, password, authFactorToken})
@ -87,7 +101,13 @@ export async function createAgentAndLogin(
  const account = agentToSessionAccountOrThrow(agent)
  const gates = tryFetchGates(account.did, 'prefer-fresh-gates')
  const moderation = configureModerationForAccount(agent, account)
-  return prepareAgent(agent, moderation, gates, onSessionChange)
+  return prepareAgent(
+    agent,
+    moderation,
+    gates,
+    onSessionChange,
+    setPersistSessionHandler,
+  )
 }

 export async function createAgentAndCreateAccount(
@ -115,6 +135,7 @@ export async function createAgentAndCreateAccount(
    did: string,
    event: AtpSessionEvent,
  ) => void,
+  setPersistSessionHandler: SetPersistSessionHandler,
 ) {
  const agent = new BskyAgent({service})
  await agent.createAccount({
@ -174,7 +195,13 @@ export async function createAgentAndCreateAccount(
    logger.error(e, {context: `session: failed snoozeEmailConfirmationPrompt`})
  }

-  return prepareAgent(agent, gates, moderation, onSessionChange)
+  return prepareAgent(
+    agent,
+    gates,
+    moderation,
+    onSessionChange,
+    setPersistSessionHandler,
+  )
 }

 async function prepareAgent(
@ -187,13 +214,14 @@ async function prepareAgent(
    did: string,
    event: AtpSessionEvent,
  ) => void,
+  setPersistSessionHandler: (cb: AtpPersistSessionHandler) => void,
 ) {
  // There's nothing else left to do, so block on them here.
  await Promise.all([gates, moderation])

  // Now the agent is ready.
  const account = agentToSessionAccountOrThrow(agent)
-  agent.setPersistSessionHandler(event => {
+  setPersistSessionHandler(event => {
    onSessionChange(agent, account.did, event)
    if (event !== 'create' && event !== 'update') {
      addSessionErrorLog(account.did, event)
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@ -1,5 +1,9 @@
 import React from 'react'
-import {AtpSessionEvent, BskyAgent} from '@atproto/api'
+import {
+  AtpPersistSessionHandler,
+  AtpSessionEvent,
+  BskyAgent,
+} from '@atproto/api'

 import {track} from '#/lib/analytics/analytics'
 import {logEvent} from '#/lib/statsig/statsig'
@ -47,6 +51,15 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
    return initialState
  })

+  const persistSessionHandler = React.useRef<
+    AtpPersistSessionHandler | undefined
+  >(undefined)
+  const setPersistSessionHandler = (
+    newHandler: AtpPersistSessionHandler | undefined,
+  ) => {
+    persistSessionHandler.current = newHandler
+  }
+
  const onAgentSessionChange = React.useCallback(
    (agent: BskyAgent, accountDid: string, sessionEvent: AtpSessionEvent) => {
      const refreshedAccount = agentToSessionAccount(agent) // Mutable, so snapshot it right away.
@ -73,6 +86,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
      const {agent, account} = await createAgentAndCreateAccount(
        params,
        onAgentSessionChange,
+        setPersistSessionHandler,
      )

      if (signal.aborted) {
@ -97,6 +111,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
      const {agent, account} = await createAgentAndLogin(
        params,
        onAgentSessionChange,
+        setPersistSessionHandler,
      )

      if (signal.aborted) {
@ -138,6 +153,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
      const {agent, account} = await createAgentAndResume(
        storedAccount,
        onAgentSessionChange,
+        setPersistSessionHandler,
      )

      if (signal.aborted) {
@ -202,7 +218,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
        } else {
          const agent = state.currentAgentState.agent as BskyAgent
          const prevSession = agent.session
-          agent.session = sessionAccountToSession(syncedAccount)
+          agent.sessionManager.session = sessionAccountToSession(syncedAccount)
          addSessionDebugLog({
            type: 'agent:patch',
            agent,
@ -249,8 +265,8 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
      addSessionDebugLog({type: 'agent:switch', prevAgent, nextAgent: agent})
      // We never reuse agents so let's fully neutralize the previous one.
      // This ensures it won't try to consume any refresh tokens.
-      prevAgent.session = undefined
-      prevAgent.setPersistSessionHandler(undefined)
+      prevAgent.sessionManager.session = undefined
+      setPersistSessionHandler(undefined)
    }
  }, [agent])

--- a/src/state/session/logging.ts
+++ b/src/state/session/logging.ts
@ -56,7 +56,7 @@ type Log =
      type: 'agent:patch'
      agent: object
      prevSession: AtpSessionData | undefined
-      nextSession: AtpSessionData
+      nextSession: AtpSessionData | undefined
    }

 export function wrapSessionReducerForLogging(reducer: Reducer): Reducer {