Persist PDS URL for session resumption (#3620)

* fix: store PDS URL for session resumption

* fix: handle a few more cases

* fix: blocking resumption should also use pds url

* refactor: do it in the construct itself

* fix: revert ce96223

2 files changed, 21 insertions, 4 deletions

diff --git a/src/state/persisted/schema.ts b/src/state/persisted/schema.ts
index 714719c2b..f090365a3 100644
--- a/src/state/persisted/schema.ts
+++ b/src/state/persisted/schema.ts
@@ -15,6 +15,7 @@ const accountSchema = z.object({
   refreshJwt: z.string().optional(), // optional because it can expire
   accessJwt: z.string().optional(), // optional because it can expire
   deactivated: z.boolean().optional(),
+  pdsUrl: z.string().optional(),
 })
 export type PersistedAccount = z.infer<typeof accountSchema>
 
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index ad5af130a..adf6ccf6d 100644
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@@ -117,6 +117,7 @@ const ApiContext = React.createContext<ApiContext>({
 })
 
 function createPersistSessionHandler(
+  agent: BskyAgent,
   account: SessionAccount,
   persistSessionCallback: (props: {
     expired: boolean
@@ -144,6 +145,7 @@ function createPersistSessionHandler(
       email: session?.email || account.email,
       emailConfirmed: session?.emailConfirmed || account.emailConfirmed,
       deactivated: isSessionDeactivated(session?.accessJwt),
+      pdsUrl: agent.pdsUrl?.toString(),
 
       /*
        * Tokens are undefined if the session expires, or if creation fails for
@@ -276,12 +278,14 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         refreshJwt: agent.session.refreshJwt,
         accessJwt: agent.session.accessJwt,
         deactivated,
+        pdsUrl: agent.pdsUrl?.toString(),
       }
 
       await configureModeration(agent, account)
 
       agent.setPersistSessionHandler(
         createPersistSessionHandler(
+          agent,
           account,
           ({expired, refreshedAccount}) => {
             upsertAccount(refreshedAccount, expired)
@@ -327,12 +331,14 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         refreshJwt: agent.session.refreshJwt,
         accessJwt: agent.session.accessJwt,
         deactivated: isSessionDeactivated(agent.session.accessJwt),
+        pdsUrl: agent.pdsUrl?.toString(),
       }
 
       await configureModeration(agent, account)
 
       agent.setPersistSessionHandler(
         createPersistSessionHandler(
+          agent,
           account,
           ({expired, refreshedAccount}) => {
             upsertAccount(refreshedAccount, expired)
@@ -379,16 +385,24 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       logger.debug(`session: initSession`, {}, logger.DebugContext.session)
       const fetchingGates = tryFetchGates(account.did, 'prefer-low-latency')
 
-      const agent = new BskyAgent({
-        service: account.service,
-        persistSession: createPersistSessionHandler(
+      const agent = new BskyAgent({service: account.service})
+
+      // restore the correct PDS URL if available
+      if (account.pdsUrl) {
+        agent.pdsUrl = agent.api.xrpc.uri = new URL(account.pdsUrl)
+      }
+
+      agent.setPersistSessionHandler(
+        createPersistSessionHandler(
+          agent,
           account,
           ({expired, refreshedAccount}) => {
             upsertAccount(refreshedAccount, expired)
           },
           {networkErrorCallback: clearCurrentAccount},
         ),
-      })
+      )
+
       // @ts-ignore
       if (IS_DEV && isWeb) window.agent = agent
       await configureModeration(agent, account)
@@ -421,6 +435,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         logger.debug(`session: attempting to reuse previous session`)
 
         agent.session = prevSession
+
         __globalAgent = agent
         await fetchingGates
         upsertAccount(account)
@@ -498,6 +513,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
           refreshJwt: agent.session.refreshJwt,
           accessJwt: agent.session.accessJwt,
           deactivated: isSessionDeactivated(agent.session.accessJwt),
+          pdsUrl: agent.pdsUrl?.toString(),
         }
       }
     },