1 files changed, 20 insertions, 4 deletions
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index 09fcf8664..4f01f7165 100644
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@@ -1,5 +1,9 @@
 import React from 'react'
-import {AtpSessionEvent, BskyAgent} from '@atproto/api'
+import {
+  AtpPersistSessionHandler,
+  AtpSessionEvent,
+  BskyAgent,
+} from '@atproto/api'
 
 import {track} from '#/lib/analytics/analytics'
 import {logEvent} from '#/lib/statsig/statsig'
@@ -47,6 +51,15 @@ 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.
@@ -73,6 +86,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const {agent, account} = await createAgentAndCreateAccount(
         params,
         onAgentSessionChange,
+        setPersistSessionHandler,
       )
 
       if (signal.aborted) {
@@ -97,6 +111,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const {agent, account} = await createAgentAndLogin(
         params,
         onAgentSessionChange,
+        setPersistSessionHandler,
       )
 
       if (signal.aborted) {
@@ -138,6 +153,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       const {agent, account} = await createAgentAndResume(
         storedAccount,
         onAgentSessionChange,
+        setPersistSessionHandler,
       )
 
       if (signal.aborted) {
@@ -202,7 +218,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         } else {
           const agent = state.currentAgentState.agent as BskyAgent
           const prevSession = agent.session
-          agent.session = sessionAccountToSession(syncedAccount)
+          agent.sessionManager.session = sessionAccountToSession(syncedAccount)
           addSessionDebugLog({
             type: 'agent:patch',
             agent,
@@ -249,8 +265,8 @@ 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.session = undefined
-      prevAgent.setPersistSessionHandler(undefined)
+      prevAgent.sessionManager.session = undefined
+      setPersistSessionHandler(undefined)
     }
   }, [agent])