Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix session persistence handlers #4927

Closed
wants to merge 2 commits into from
Closed

Fix session persistence handlers #4927

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
894a99d
Readd deleted comment
gaearon Aug 13, 2024
afb1ba8
Maintain persist handler per agent
gaearon Aug 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 50 additions & 54 deletions src/state/session/agent.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,6 @@ 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})
Expand All @@ -39,9 +37,14 @@ export async function createAgentAndResume(
did: string,
event: AtpSessionEvent,
) => void,
setPersistSessionHandler: SetPersistSessionHandler,
) {
const agent = new BskyAgent({service: storedAccount.service})
let currentPersistHandler: AtpPersistSessionHandler | undefined
const agent = new BskyAgent({
service: storedAccount.service,
persistSession(event, session) {
currentPersistHandler?.(event, session)
},
})
if (storedAccount.pdsUrl) {
agent.sessionManager.pdsUrl = new URL(storedAccount.pdsUrl)
}
Expand All @@ -53,6 +56,7 @@ export async function createAgentAndResume(
} else {
agent.sessionManager.session = prevSession
if (!storedAccount.signupQueued) {
// Intentionally not awaited to unblock the UI:
networkRetry(3, () => agent.resumeSession(prevSession)).catch(
(e: any) => {
logger.error(`networkRetry failed to resume session`, {
Expand All @@ -67,13 +71,16 @@ export async function createAgentAndResume(
}
}

return prepareAgent(
agent,
gates,
moderation,
onSessionChange,
setPersistSessionHandler,
)
const account = agentToSessionAccountOrThrow(agent)
currentPersistHandler = event => {
onSessionChange(agent, account.did, event)
if (event !== 'create' && event !== 'update') {
addSessionErrorLog(account.did, event)
}
}

await Promise.all([gates, moderation])
return {account, agent}
}

export async function createAgentAndLogin(
Expand All @@ -93,21 +100,28 @@ export async function createAgentAndLogin(
did: string,
event: AtpSessionEvent,
) => void,
setPersistSessionHandler: SetPersistSessionHandler,
) {
const agent = new BskyAgent({service})
let currentPersistHandler: AtpPersistSessionHandler | undefined
const agent = new BskyAgent({
service,
persistSession(event, session) {
currentPersistHandler?.(event, session)
},
})
await agent.login({identifier, password, authFactorToken})

const account = agentToSessionAccountOrThrow(agent)
currentPersistHandler = event => {
onSessionChange(agent, account.did, event)
if (event !== 'create' && event !== 'update') {
addSessionErrorLog(account.did, event)
}
}

const gates = tryFetchGates(account.did, 'prefer-fresh-gates')
const moderation = configureModerationForAccount(agent, account)
return prepareAgent(
agent,
moderation,
gates,
onSessionChange,
setPersistSessionHandler,
)
await Promise.all([gates, moderation])
return {account, agent}
}

export async function createAgentAndCreateAccount(
Expand Down Expand Up @@ -135,9 +149,14 @@ export async function createAgentAndCreateAccount(
did: string,
event: AtpSessionEvent,
) => void,
setPersistSessionHandler: SetPersistSessionHandler,
) {
const agent = new BskyAgent({service})
let currentPersistHandler: AtpPersistSessionHandler | undefined
const agent = new BskyAgent({
service,
persistSession(event, session) {
currentPersistHandler?.(event, session)
},
})
await agent.createAccount({
email,
password,
Expand All @@ -146,7 +165,15 @@ export async function createAgentAndCreateAccount(
verificationPhone,
verificationCode,
})

const account = agentToSessionAccountOrThrow(agent)
currentPersistHandler = event => {
onSessionChange(agent, account.did, event)
if (event !== 'create' && event !== 'update') {
addSessionErrorLog(account.did, event)
}
}

const gates = tryFetchGates(account.did, 'prefer-fresh-gates')
const moderation = configureModerationForAccount(agent, account)

Expand Down Expand Up @@ -195,39 +222,8 @@ 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 {account, agent}
}

export function agentToSessionAccountOrThrow(agent: BskyAgent): SessionAccount {
Expand Down
19 changes: 1 addition & 18 deletions src/state/session/index.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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'
Expand Down Expand Up @@ -51,15 +47,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.
Expand All @@ -86,7 +73,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
const {agent, account} = await createAgentAndCreateAccount(
params,
onAgentSessionChange,
setPersistSessionHandler,
)

if (signal.aborted) {
Expand All @@ -111,7 +97,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
const {agent, account} = await createAgentAndLogin(
params,
onAgentSessionChange,
setPersistSessionHandler,
)

if (signal.aborted) {
Expand Down Expand Up @@ -153,7 +138,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
const {agent, account} = await createAgentAndResume(
storedAccount,
onAgentSessionChange,
setPersistSessionHandler,
)

if (signal.aborted) {
Expand Down Expand Up @@ -266,7 +250,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
// 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)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note I am no longer cleaning up the persistence handler when the current agent gets swapped. I think this should be fine. It didn't used to be cleaned up in the past either, I only added this at some intermediate stage for peace of mind. But we already have a condition here that prevents inactive agents from clearing persisted data. For inactive agents, refreshedAccount should be undefined because we already set prevAgent.sessionManager.session = undefined here. And even if somehow it had a session object, the reducer always matches events with state by the did so it won't update the wrong thing.

}
}, [agent])

Expand Down
Loading