subclass agent to add setPersistSessionHandler (#4928)

Co-authored-by: Dan Abramov <dan.abramov@gmail.com>

2 files changed, 60 insertions, 82 deletions

diff --git a/src/state/session/agent.ts b/src/state/session/agent.ts
index 73be34bb2..ea6af677c 100644
--- a/src/state/session/agent.ts
+++ b/src/state/session/agent.ts
@@ -1,9 +1,4 @@
-import {
-  AtpPersistSessionHandler,
-  AtpSessionData,
-  AtpSessionEvent,
-  BskyAgent,
-} from '@atproto/api'
+import {AtpSessionData, AtpSessionEvent, BskyAgent} from '@atproto/api'
 import {TID} from '@atproto/common-web'
 
 import {networkRetry} from '#/lib/async/retry'
@@ -25,11 +20,9 @@ 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})
+  return new BskyAppAgent({service: PUBLIC_BSKY_SERVICE})
 }
 
 export async function createAgentAndResume(
@@ -39,9 +32,8 @@ export async function createAgentAndResume(
     did: string,
     event: AtpSessionEvent,
   ) => void,
-  setPersistSessionHandler: SetPersistSessionHandler,
 ) {
-  const agent = new BskyAgent({service: storedAccount.service})
+  const agent = new BskyAppAgent({service: storedAccount.service})
   if (storedAccount.pdsUrl) {
     agent.sessionManager.pdsUrl = new URL(storedAccount.pdsUrl)
   }
@@ -67,13 +59,7 @@ export async function createAgentAndResume(
     }
   }
 
-  return prepareAgent(
-    agent,
-    gates,
-    moderation,
-    onSessionChange,
-    setPersistSessionHandler,
-  )
+  return agent.prepare(gates, moderation, onSessionChange)
 }
 
 export async function createAgentAndLogin(
@@ -93,21 +79,14 @@ export async function createAgentAndLogin(
     did: string,
     event: AtpSessionEvent,
   ) => void,
-  setPersistSessionHandler: SetPersistSessionHandler,
 ) {
-  const agent = new BskyAgent({service})
+  const agent = new BskyAppAgent({service})
   await agent.login({identifier, password, authFactorToken})
 
   const account = agentToSessionAccountOrThrow(agent)
   const gates = tryFetchGates(account.did, 'prefer-fresh-gates')
   const moderation = configureModerationForAccount(agent, account)
-  return prepareAgent(
-    agent,
-    moderation,
-    gates,
-    onSessionChange,
-    setPersistSessionHandler,
-  )
+  return agent.prepare(gates, moderation, onSessionChange)
 }
 
 export async function createAgentAndCreateAccount(
@@ -135,9 +114,8 @@ export async function createAgentAndCreateAccount(
     did: string,
     event: AtpSessionEvent,
   ) => void,
-  setPersistSessionHandler: SetPersistSessionHandler,
 ) {
-  const agent = new BskyAgent({service})
+  const agent = new BskyAppAgent({service})
   await agent.createAccount({
     email,
     password,
@@ -195,39 +173,7 @@ export async function createAgentAndCreateAccount(
     logger.error(e, {context: `session: failed snoozeEmailConfirmationPrompt`})
   }
 
-  return prepareAgent(
-    agent,
-    gates,
-    moderation,
-    onSessionChange,
-    setPersistSessionHandler,
-  )
-}
-
-async function prepareAgent(
-  agent: BskyAgent,
-  // Not awaited in the calling code so we can delay blocking on them.
-  gates: Promise<void>,
-  moderation: Promise<void>,
-  onSessionChange: (
-    agent: BskyAgent,
-    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)
-  setPersistSessionHandler(event => {
-    onSessionChange(agent, account.did, event)
-    if (event !== 'create' && event !== 'update') {
-      addSessionErrorLog(account.did, event)
-    }
-  })
-  return {agent, account}
+  return agent.prepare(gates, moderation, onSessionChange)
 }
 
 export function agentToSessionAccountOrThrow(agent: BskyAgent): SessionAccount {
@@ -279,3 +225,51 @@ export function sessionAccountToSession(
     status: account.status,
   }
 }
+
+// Not exported. Use factories above to create it.
+class BskyAppAgent extends BskyAgent {
+  persistSessionHandler: ((event: AtpSessionEvent) => void) | undefined =
+    undefined
+
+  constructor({service}: {service: string}) {
+    super({
+      service,
+      persistSession: (event: AtpSessionEvent) => {
+        if (this.persistSessionHandler) {
+          this.persistSessionHandler(event)
+        }
+      },
+    })
+  }
+
+  async prepare(
+    // Not awaited in the calling code so we can delay blocking on them.
+    gates: Promise<void>,
+    moderation: Promise<void>,
+    onSessionChange: (
+      agent: BskyAgent,
+      did: string,
+      event: AtpSessionEvent,
+    ) => 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(this)
+    this.persistSessionHandler = event => {
+      onSessionChange(this, account.did, event)
+      if (event !== 'create' && event !== 'update') {
+        addSessionErrorLog(account.did, event)
+      }
+    }
+    return {account, agent: this}
+  }
+
+  dispose() {
+    this.sessionManager.session = undefined
+    this.persistSessionHandler = undefined
+  }
+}
+
+export type {BskyAppAgent}
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index 4f01f7165..ba12f4eae 100644
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@@ -1,9 +1,5 @@
 import React from 'react'
-import {
-  AtpPersistSessionHandler,
-  AtpSessionEvent,
-  BskyAgent,
-} from '@atproto/api'
+import {AtpSessionEvent, BskyAgent} from '@atproto/api'
 
 import {track} from '#/lib/analytics/analytics'
 import {logEvent} from '#/lib/statsig/statsig'
@@ -15,6 +11,7 @@ import {IS_DEV} from '#/env'
 import {emitSessionDropped} from '../events'
 import {
   agentToSessionAccount,
+  BskyAppAgent,
   createAgentAndCreateAccount,
   createAgentAndLogin,
   createAgentAndResume,
@@ -51,15 +48,6 @@ 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.
@@ -86,7 +74,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const {agent, account} = await createAgentAndCreateAccount(
         params,
         onAgentSessionChange,
-        setPersistSessionHandler,
       )
 
       if (signal.aborted) {
@@ -111,7 +98,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const {agent, account} = await createAgentAndLogin(
         params,
         onAgentSessionChange,
-        setPersistSessionHandler,
       )
 
       if (signal.aborted) {
@@ -153,7 +139,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const {agent, account} = await createAgentAndResume(
         storedAccount,
         onAgentSessionChange,
-        setPersistSessionHandler,
       )
 
       if (signal.aborted) {
@@ -255,7 +240,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   // @ts-ignore
   if (IS_DEV && isWeb) window.agent = state.currentAgentState.agent
 
-  const agent = state.currentAgentState.agent as BskyAgent
+  const agent = state.currentAgentState.agent as BskyAppAgent
   const currentAgentRef = React.useRef(agent)
   React.useEffect(() => {
     if (currentAgentRef.current !== agent) {
@@ -265,8 +250,7 @@ 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.sessionManager.session = undefined
-      setPersistSessionHandler(undefined)
+      prevAgent.dispose()
     }
   }, [agent])