1 files changed, 56 insertions, 62 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}