Upgrade API, implement XRPC rework (#4857)

Co-authored-by: Matthieu Sieben <matthieu.sieben@gmail.com>

5 files changed, 95 insertions, 51 deletions

diff --git a/src/state/queries/preferences/index.ts b/src/state/queries/preferences/index.ts
index 6991f8647..ab866d5e2 100644
--- 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),
         )
 
diff --git a/src/state/session/__tests__/session-test.ts b/src/state/session/__tests__/session-test.ts
index 486604169..731b66b0e 100644
--- 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',
diff --git a/src/state/session/agent.ts b/src/state/session/agent.ts
index 4456ab0bf..73be34bb2 100644
--- 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)
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index 09fcf8664..4f01f7165 100644
--- 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])
 
diff --git a/src/state/session/logging.ts b/src/state/session/logging.ts
index b57f1fa0b..7e1df500b 100644
--- 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 {