From 15055cb8c46f2b8bb09e3b7a28bbf1f8ec58fbe0 Mon Sep 17 00:00:00 2001
From: Mary <148872143+mary-ext@users.noreply.github.com>
Date: Thu, 25 Apr 2024 03:57:16 +0700
Subject: 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
---
 src/state/session/index.tsx | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

(limited to 'src/state/session/index.tsx')

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(),
         }
       }
     },
-- 
cgit 1.4.1

