From 748212e000f963976bca5d63e0961d75a7e8b296 Mon Sep 17 00:00:00 2001
From: dan <dan.abramov@gmail.com>
Date: Wed, 6 Dec 2023 18:32:14 +0000
Subject: [PATCH] Remove getSession waterfall (#2112)

---
 package.json                |  1 +
 src/state/session/index.tsx | 80 ++++++++++++++++++++++++++-----------
 yarn.lock                   |  5 +++
 3 files changed, 62 insertions(+), 24 deletions(-)

diff --git a/package.json b/package.json
index fdca3e89..4d526dac 100644
--- a/package.json
+++ b/package.json
@@ -110,6 +110,7 @@
     "fast-text-encoding": "^1.0.6",
     "history": "^5.3.0",
     "js-sha256": "^0.9.0",
+    "jwt-decode": "^4.0.0",
     "lande": "^1.0.10",
     "lodash.chunk": "^4.2.0",
     "lodash.debounce": "^4.0.8",
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index 0d52d252..49c1326d 100644
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@@ -1,6 +1,7 @@
 import React from 'react'
 import {BskyAgent, AtpPersistSessionHandler} from '@atproto/api'
 import {useQueryClient} from '@tanstack/react-query'
+import {jwtDecode} from 'jwt-decode'
 
 import {networkRetry} from '#/lib/async/retry'
 import {logger} from '#/logger'
@@ -327,34 +328,65 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         ),
       })
 
-      await networkRetry(3, () =>
-        agent.resumeSession({
-          accessJwt: account.accessJwt || '',
-          refreshJwt: account.refreshJwt || '',
-          did: account.did,
-          handle: account.handle,
-        }),
-      )
-
-      if (!agent.session) {
-        throw new Error(`session: initSession failed to establish a session`)
+      let canReusePrevSession = false
+      try {
+        if (account.accessJwt) {
+          const decoded = jwtDecode(account.accessJwt)
+          if (decoded.exp) {
+            const didExpire = Date.now() >= decoded.exp * 1000
+            if (!didExpire) {
+              canReusePrevSession = true
+            }
+          }
+        }
+      } catch (e) {
+        console.error('Could not decode jwt.')
       }
 
-      // ensure changes in handle/email etc are captured on reload
-      const freshAccount: SessionAccount = {
-        service: agent.service.toString(),
-        did: agent.session.did,
-        handle: agent.session.handle,
-        email: agent.session.email,
-        emailConfirmed: agent.session.emailConfirmed || false,
-        refreshJwt: agent.session.refreshJwt,
-        accessJwt: agent.session.accessJwt,
+      const prevSession = {
+        accessJwt: account.accessJwt || '',
+        refreshJwt: account.refreshJwt || '',
+        did: account.did,
+        handle: account.handle,
       }
 
-      __globalAgent = agent
-      queryClient.clear()
-      upsertAccount(freshAccount)
-      emitSessionLoaded(freshAccount, agent)
+      if (canReusePrevSession) {
+        agent.session = prevSession
+        __globalAgent = agent
+        queryClient.clear()
+        upsertAccount(account)
+        emitSessionLoaded(account, agent)
+        // Intentionally not awaited to unblock the UI:
+        resumeSessionWithFreshAccount().then(async freshAccount => {
+          if (JSON.stringify(account) !== JSON.stringify(freshAccount)) {
+            upsertAccount(freshAccount)
+            emitSessionLoaded(freshAccount, agent)
+          }
+        })
+      } else {
+        const freshAccount = await resumeSessionWithFreshAccount()
+        __globalAgent = agent
+        queryClient.clear()
+        upsertAccount(freshAccount)
+        emitSessionLoaded(freshAccount, agent)
+      }
+
+      async function resumeSessionWithFreshAccount(): Promise<SessionAccount> {
+        await networkRetry(3, () => agent.resumeSession(prevSession))
+        if (!agent.session) {
+          throw new Error(`session: initSession failed to establish a session`)
+        }
+        // ensure changes in handle/email etc are captured on reload
+        return {
+          service: agent.service.toString(),
+          did: agent.session.did,
+          handle: agent.session.handle,
+          email: agent.session.email,
+          emailConfirmed: agent.session.emailConfirmed || false,
+          refreshJwt: agent.session.refreshJwt,
+          accessJwt: agent.session.accessJwt,
+        }
+      }
     },
     [upsertAccount, queryClient],
   )
diff --git a/yarn.lock b/yarn.lock
index af78ffc7..b9606320 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -14279,6 +14279,11 @@ jws@^3.2.2:
     jwa "^1.4.1"
     safe-buffer "^5.0.1"
 
+jwt-decode@^4.0.0:
+  version "4.0.0"
+  resolved "https://registry.yarnpkg.com/jwt-decode/-/jwt-decode-4.0.0.tgz#2270352425fd413785b2faf11f6e755c5151bd4b"
+  integrity sha512-+KJGIyHgkGuIq3IEBNftfhW/LfWhXUIY6OmyVWjliu5KH1y0fw7VQ8YndE2O4qZdMSd9SqbnC8GOcZEy0Om7sA==
+
 key-encoder@^2.0.3:
   version "2.0.3"
   resolved "https://registry.yarnpkg.com/key-encoder/-/key-encoder-2.0.3.tgz#77073bb48ff1fe2173bb2088b83b91152c8fa4ba"