subclass agent to add setPersistSessionHandler (#4928)

Co-authored-by: Dan Abramov <dan.abramov@gmail.com>
zio/stable
Hailey 2024-08-12 19:43:06 -07:00 committed by GitHub
parent 99d1a881f2
commit 3c04d9bd84
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 60 additions and 82 deletions

View File

@ -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}

View File

@ -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])