Move request for notifications permissions to `HomeReadyScreen` (#3977)

* cleanup the current logic

* add statsig logs

* implement requests for permissions where needed

* oops

* let `addPushTokenListener` handle the token registration

* place new log event type with the other `notifications` type

* place registration next to handler

* more organization

* only call `gate()` if permission is not yet granted

* be more specific to prevent gate pollution

* nit

* make `token` non-optional in `registerToken`

* remove `prevDid`, move `registerPushToken` into `useEffect`

* keep it outside actually

* nit

7 files changed, 105 insertions, 67 deletions

diff --git a/src/lib/notifications/notifications.ts b/src/lib/notifications/notifications.ts
index 38c18bf3f..66cedeaa6 100644
--- a/src/lib/notifications/notifications.ts
+++ b/src/lib/notifications/notifications.ts
@@ -1,27 +1,22 @@
+import React from 'react'
 import * as Notifications from 'expo-notifications'
 import {BskyAgent} from '@atproto/api'
 
 import {logger} from '#/logger'
-import {SessionAccount} from '#/state/session'
-import {devicePlatform} from 'platform/detection'
+import {SessionAccount, useAgent, useSession} from '#/state/session'
+import {logEvent, useGate} from 'lib/statsig/statsig'
+import {devicePlatform, isNative} from 'platform/detection'
 
 const SERVICE_DID = (serviceUrl?: string) =>
   serviceUrl?.includes('staging')
     ? 'did:web:api.staging.bsky.dev'
     : 'did:web:api.bsky.app'
 
-export async function requestPermissionsAndRegisterToken(
+async function registerPushToken(
   getAgent: () => BskyAgent,
   account: SessionAccount,
+  token: Notifications.DevicePushToken,
 ) {
-  // request notifications permission once the user has logged in
-  const perms = await Notifications.getPermissionsAsync()
-  if (!perms.granted) {
-    await Notifications.requestPermissionsAsync()
-  }
-
-  // register the push token with the server
-  const token = await Notifications.getDevicePushTokenAsync()
   try {
     await getAgent().api.app.bsky.notification.registerPush({
       serviceDid: SERVICE_DID(account.service),
@@ -42,38 +37,63 @@ export async function requestPermissionsAndRegisterToken(
   }
 }
 
-export function registerTokenChangeHandler(
-  getAgent: () => BskyAgent,
-  account: SessionAccount,
-): () => void {
-  // listens for new changes to the push token
-  // In rare situations, a push token may be changed by the push notification service while the app is running. When a token is rolled, the old one becomes invalid and sending notifications to it will fail. A push token listener will let you handle this situation gracefully by registering the new token with your backend right away.
-  const sub = Notifications.addPushTokenListener(async newToken => {
-    logger.debug(
-      'Notifications: Push token changed',
-      {tokenType: newToken.data, token: newToken.type},
-      logger.DebugContext.notifications,
-    )
-    try {
-      await getAgent().api.app.bsky.notification.registerPush({
-        serviceDid: SERVICE_DID(account.service),
-        platform: devicePlatform,
-        token: newToken.data,
-        appId: 'xyz.blueskyweb.app',
-      })
-      logger.debug(
-        'Notifications: Sent push token (event)',
-        {
-          tokenType: newToken.type,
-          token: newToken.data,
-        },
-        logger.DebugContext.notifications,
-      )
-    } catch (error) {
-      logger.error('Notifications: Failed to set push token', {message: error})
+export function useNotificationsRegistration() {
+  const [currentPermissions] = Notifications.usePermissions()
+  const {getAgent} = useAgent()
+  const {currentAccount} = useSession()
+
+  React.useEffect(() => {
+    if (!currentAccount || !currentPermissions?.granted) {
+      return
     }
-  })
-  return () => {
-    sub.remove()
-  }
+
+    // Whenever we all `getDevicePushTokenAsync()`, a change event will be fired below
+    Notifications.getDevicePushTokenAsync()
+
+    // According to the Expo docs, there is a chance that the token will change while the app is open in some rare
+    // cases. This will fire `registerPushToken` whenever that happens.
+    const subscription = Notifications.addPushTokenListener(async newToken => {
+      registerPushToken(getAgent, currentAccount, newToken)
+    })
+
+    return () => {
+      subscription.remove()
+    }
+  }, [currentAccount, currentPermissions?.granted, getAgent])
+}
+
+export function useRequestNotificationsPermission() {
+  const gate = useGate()
+  const [currentPermissions] = Notifications.usePermissions()
+
+  return React.useCallback(
+    async (context: 'StartOnboarding' | 'AfterOnboarding') => {
+      if (
+        !isNative ||
+        currentPermissions?.status === 'granted' ||
+        (currentPermissions?.status === 'denied' &&
+          !currentPermissions?.canAskAgain)
+      ) {
+        return
+      }
+      if (
+        context === 'StartOnboarding' &&
+        gate('request_notifications_permission_after_onboarding')
+      ) {
+        return
+      }
+      if (
+        context === 'AfterOnboarding' &&
+        !gate('request_notifications_permission_after_onboarding')
+      ) {
+        return
+      }
+
+      const res = await Notifications.requestPermissionsAsync()
+      logEvent('notifications:request', {
+        status: res.status,
+      })
+    },
+    [currentPermissions?.canAskAgain, currentPermissions?.status, gate],
+  )
 }
diff --git a/src/lib/statsig/events.ts b/src/lib/statsig/events.ts
index 3355377fe..d73d21a1a 100644
--- a/src/lib/statsig/events.ts
+++ b/src/lib/statsig/events.ts
@@ -16,6 +16,9 @@ export type LogEvents = {
     logContext: 'SwitchAccount' | 'Settings' | 'Deactivated'
   }
   'notifications:openApp': {}
+  'notifications:request': {
+    status: 'granted' | 'denied' | 'undetermined'
+  }
   'state:background': {
     secondsActive: number
   }
diff --git a/src/lib/statsig/gates.ts b/src/lib/statsig/gates.ts
index a2dbb4950..315706ad0 100644
--- a/src/lib/statsig/gates.ts
+++ b/src/lib/statsig/gates.ts
@@ -5,6 +5,7 @@ export type Gate =
   | 'disable_poll_on_discover_v2'
   | 'dms'
   | 'reduced_onboarding_and_home_algo'
+  | 'request_notifications_permission_after_onboarding'
   | 'show_follow_back_label_v2'
   | 'start_session_with_following_v2'
   | 'test_gate_1'
diff --git a/src/screens/Onboarding/StepInterests/index.tsx b/src/screens/Onboarding/StepInterests/index.tsx
index d6678f4b0..2711f6779 100644
--- a/src/screens/Onboarding/StepInterests/index.tsx
+++ b/src/screens/Onboarding/StepInterests/index.tsx
@@ -5,11 +5,12 @@ import {useLingui} from '@lingui/react'
 import {useQuery} from '@tanstack/react-query'
 
 import {useAnalytics} from '#/lib/analytics/analytics'
-import {logEvent} from '#/lib/statsig/statsig'
+import {logEvent, useGate} from '#/lib/statsig/statsig'
 import {capitalize} from '#/lib/strings/capitalize'
 import {logger} from '#/logger'
 import {useAgent} from '#/state/session'
 import {useOnboardingDispatch} from '#/state/shell'
+import {useRequestNotificationsPermission} from 'lib/notifications/notifications'
 import {
   DescriptionText,
   OnboardingControls,
@@ -33,6 +34,9 @@ export function StepInterests() {
   const t = useTheme()
   const {gtMobile} = useBreakpoints()
   const {track} = useAnalytics()
+  const gate = useGate()
+  const requestNotificationsPermission = useRequestNotificationsPermission()
+
   const {state, dispatch, interestsDisplayNames} = React.useContext(Context)
   const [saving, setSaving] = React.useState(false)
   const [interests, setInterests] = React.useState<string[]>(
@@ -129,6 +133,12 @@ export function StepInterests() {
     track('OnboardingV2:StepInterests:Start')
   }, [track])
 
+  React.useEffect(() => {
+    if (!gate('reduced_onboarding_and_home_algo')) {
+      requestNotificationsPermission('StartOnboarding')
+    }
+  }, [gate, requestNotificationsPermission])
+
   const title = isError ? (
     <Trans>Oh no! Something went wrong.</Trans>
   ) : (
diff --git a/src/screens/Onboarding/StepProfile/index.tsx b/src/screens/Onboarding/StepProfile/index.tsx
index bf47bbc95..d480a32af 100644
--- a/src/screens/Onboarding/StepProfile/index.tsx
+++ b/src/screens/Onboarding/StepProfile/index.tsx
@@ -10,11 +10,12 @@ import {msg, Trans} from '@lingui/macro'
 import {useLingui} from '@lingui/react'
 
 import {useAnalytics} from '#/lib/analytics/analytics'
-import {logEvent} from '#/lib/statsig/statsig'
+import {logEvent, useGate} from '#/lib/statsig/statsig'
 import {usePhotoLibraryPermission} from 'lib/hooks/usePermissions'
 import {compressIfNeeded} from 'lib/media/manip'
 import {openCropper} from 'lib/media/picker'
 import {getDataUriSize} from 'lib/media/util'
+import {useRequestNotificationsPermission} from 'lib/notifications/notifications'
 import {isNative, isWeb} from 'platform/detection'
 import {
   DescriptionText,
@@ -69,6 +70,9 @@ export function StepProfile() {
   const {gtMobile} = useBreakpoints()
   const {track} = useAnalytics()
   const {requestPhotoAccessIfNeeded} = usePhotoLibraryPermission()
+  const gate = useGate()
+  const requestNotificationsPermission = useRequestNotificationsPermission()
+
   const creatorControl = Dialog.useDialogControl()
   const [error, setError] = React.useState('')
 
@@ -86,6 +90,14 @@ export function StepProfile() {
     track('OnboardingV2:StepProfile:Start')
   }, [track])
 
+  React.useEffect(() => {
+    // We have an experiment running for redueced onboarding, where this screen shows up as the first in onboarding.
+    // We only want to request permissions when that gate is actually active to prevent pollution
+    if (gate('reduced_onboarding_and_home_algo')) {
+      requestNotificationsPermission('StartOnboarding')
+    }
+  }, [gate, requestNotificationsPermission])
+
   const openPicker = React.useCallback(
     async (opts?: ImagePickerOptions) => {
       const response = await launchImageLibraryAsync({
diff --git a/src/view/screens/Home.tsx b/src/view/screens/Home.tsx
index bd17e5fe4..d2d31ce6a 100644
--- a/src/view/screens/Home.tsx
+++ b/src/view/screens/Home.tsx
@@ -20,6 +20,7 @@ import {
 } from '#/state/shell'
 import {useSelectedFeed, useSetSelectedFeed} from '#/state/shell/selected-feed'
 import {useOTAUpdates} from 'lib/hooks/useOTAUpdates'
+import {useRequestNotificationsPermission} from 'lib/notifications/notifications'
 import {HomeTabNavigatorParams, NativeStackScreenProps} from 'lib/routes/types'
 import {FeedPage} from 'view/com/feeds/FeedPage'
 import {Pager, PagerRef, RenderTabBarFnProps} from 'view/com/pager/Pager'
@@ -58,7 +59,9 @@ function HomeScreenReady({
   preferences: UsePreferencesQueryResponse
   pinnedFeedInfos: SavedFeedSourceInfo[]
 }) {
-  useOTAUpdates()
+  const gate = useGate()
+  const requestNotificationsPermission = useRequestNotificationsPermission()
+
   const allFeeds = React.useMemo(
     () => pinnedFeedInfos.map(f => f.feedDescriptor),
     [pinnedFeedInfos],
@@ -70,6 +73,11 @@ function HomeScreenReady({
   const selectedFeed = allFeeds[selectedIndex]
 
   useSetTitle(pinnedFeedInfos[selectedIndex]?.displayName)
+  useOTAUpdates()
+
+  React.useEffect(() => {
+    requestNotificationsPermission('AfterOnboarding')
+  }, [requestNotificationsPermission])
 
   const pagerRef = React.useRef<PagerRef>(null)
   const lastPagerReportedIndexRef = React.useRef(selectedIndex)
@@ -109,7 +117,6 @@ function HomeScreenReady({
     }),
   )
 
-  const gate = useGate()
   const mode = useMinimalShellMode()
   const {isMobile} = useWebMediaQueries()
   useFocusEffect(
diff --git a/src/view/shell/index.tsx b/src/view/shell/index.tsx
index 425c1b3f8..7d080e57b 100644
--- a/src/view/shell/index.tsx
+++ b/src/view/shell/index.tsx
@@ -13,7 +13,7 @@ import * as NavigationBar from 'expo-navigation-bar'
 import {StatusBar} from 'expo-status-bar'
 import {useNavigationState} from '@react-navigation/native'
 
-import {useAgent, useSession} from '#/state/session'
+import {useSession} from '#/state/session'
 import {
   useIsDrawerOpen,
   useIsDrawerSwipeDisabled,
@@ -22,7 +22,7 @@ import {
 import {useCloseAnyActiveElement} from '#/state/util'
 import {useNotificationsHandler} from 'lib/hooks/useNotificationHandler'
 import {usePalette} from 'lib/hooks/usePalette'
-import * as notifications from 'lib/notifications/notifications'
+import {useNotificationsRegistration} from 'lib/notifications/notifications'
 import {isStateAtTabRoot} from 'lib/routes/helpers'
 import {useTheme} from 'lib/ThemeContext'
 import {isAndroid} from 'platform/detection'
@@ -57,13 +57,11 @@ function ShellInner() {
     [setIsDrawerOpen],
   )
   const canGoBack = useNavigationState(state => !isStateAtTabRoot(state))
-  const {hasSession, currentAccount} = useSession()
-  const {getAgent} = useAgent()
+  const {hasSession} = useSession()
   const closeAnyActiveElement = useCloseAnyActiveElement()
   const {importantForAccessibility} = useDialogStateContext()
-  // start undefined
-  const currentAccountDid = React.useRef<string | undefined>(undefined)
 
+  useNotificationsRegistration()
   useNotificationsHandler()
 
   React.useEffect(() => {
@@ -78,19 +76,6 @@ function ShellInner() {
     }
   }, [closeAnyActiveElement])
 
-  React.useEffect(() => {
-    // only runs when did changes
-    if (currentAccount && currentAccountDid.current !== currentAccount.did) {
-      currentAccountDid.current = currentAccount.did
-      notifications.requestPermissionsAndRegisterToken(getAgent, currentAccount)
-      const unsub = notifications.registerTokenChangeHandler(
-        getAgent,
-        currentAccount,
-      )
-      return unsub
-    }
-  }, [currentAccount, getAgent])
-
   return (
     <>
       <Animated.View