Improve effects, memoize

src/ageAssurance/index.tsx

@@ -1,4 +1,4 @@
-import {createContext, useContext, useEffect, useMemo} from 'react'
+import {createContext, useCallback, useContext, useEffect, useMemo} from 'react'
 
 import {useGetAndRegisterPushToken} from '#/lib/notifications/notifications'
 import {Provider as RedirectOverlayProvider} from '#/ageAssurance/components/RedirectOverlay'
@@ -56,16 +56,20 @@ function InnerProvider({children}: {children: React.ReactNode}) {
   const state = useAgeAssuranceState()
   const getAndRegisterPushToken = useGetAndRegisterPushToken()
 
+  const handleAccessUpdate = useCallback(
+    (s: AgeAssuranceState) => {
+      getAndRegisterPushToken({
+        isAgeRestricted: s.access !== AgeAssuranceAccess.Full,
+      })
+    },
+    [getAndRegisterPushToken],
+  )
+  useOnAgeAssuranceAccessUpdate(handleAccessUpdate)
+
   useEffect(() => {
     logger.debug(`useAgeAssuranceState`, {state})
   }, [state])
 
-  useOnAgeAssuranceAccessUpdate(state => {
-    getAndRegisterPushToken({
-      isAgeRestricted: state.access !== AgeAssuranceAccess.Full,
-    })
-  })
-
   return (
     <AgeAssuranceStateContext.Provider
       value={useMemo(

src/ageAssurance/state.ts

@@ -1,4 +1,4 @@
-import {useMemo, useRef} from 'react'
+import {useEffect, useMemo, useState} from 'react'
 import {computeAgeAssuranceRegionAccess} from '@atproto/api'
 
 import {useSession} from '#/state/session'
@@ -81,11 +81,14 @@ export function useOnAgeAssuranceAccessUpdate(
   cb: (state: AgeAssuranceState) => void,
 ) {
   const state = useAgeAssuranceState()
-  const prevState = useRef<AgeAssuranceAccess | null>(null)
+  // start with null to ensure callback is called on first render
+  const [prevAccess, setPrevAccess] = useState<AgeAssuranceAccess | null>(null)
 
-  if (prevState.current !== state.access) {
-    prevState.current = state.access
-    cb(state)
-    logger.debug(`useOnAgeAssuranceAccessUpdate`, {state})
-  }
+  useEffect(() => {
+    if (prevAccess !== state.access) {
+      setPrevAccess(state.access)
+      cb(state)
+      logger.debug(`useOnAgeAssuranceAccessUpdate`, {state})
+    }
+  }, [cb, state, prevAccess])
 }