Age assurance tweaks (#8671)

* Refetch state while in pending status

* Add success state to redirect, just have user close, nicer experience

* Stop refetching after pending

2 files changed, 78 insertions, 7 deletions

diff --git a/src/components/ageAssurance/AgeAssuranceRedirectDialog.tsx b/src/components/ageAssurance/AgeAssuranceRedirectDialog.tsx
index ff2e0bfd0..b1c287e1b 100644
--- a/src/components/ageAssurance/AgeAssuranceRedirectDialog.tsx
+++ b/src/components/ageAssurance/AgeAssuranceRedirectDialog.tsx
@@ -14,6 +14,7 @@ import {AgeAssuranceBadge} from '#/components/ageAssurance/AgeAssuranceBadge'
 import {Button, ButtonText} from '#/components/Button'
 import * as Dialog from '#/components/Dialog'
 import {useGlobalDialogsControlContext} from '#/components/dialogs/Context'
+import {CheckThick_Stroke2_Corner0_Rounded as SuccessIcon} from '#/components/icons/Check'
 import {CircleInfo_Stroke2_Corner0_Rounded as ErrorIcon} from '#/components/icons/CircleInfo'
 import {Loader} from '#/components/Loader'
 import {Text} from '#/components/Typography'
@@ -66,7 +67,7 @@ export function AgeAssuranceRedirectDialog() {
   // Dialog.useAutoOpen(control.control, 3e3)
 
   return (
-    <Dialog.Outer control={control.control}>
+    <Dialog.Outer control={control.control} onClose={() => control.clear()}>
       <Dialog.Handle />
 
       <Dialog.ScrollableInner
@@ -86,6 +87,7 @@ export function Inner({}: {optimisticState?: AgeAssuranceRedirectDialogState}) {
   const unmounted = useRef(false)
   const control = useAgeAssuranceRedirectDialogControl()
   const [error, setError] = useState(false)
+  const [success, setSuccess] = useState(false)
   const {refetch: refreshAgeAssuranceState} = useAgeAssuranceAPIContext()
 
   useEffect(() => {
@@ -125,8 +127,7 @@ export function Inner({}: {optimisticState?: AgeAssuranceRedirectDialogState}) {
         // success! update state
         await refreshAgeAssuranceState()
 
-        control.clear()
-        control.control.close()
+        setSuccess(true)
 
         logger.metric('ageAssurance:redirectDialogSuccess', {})
       })
@@ -143,6 +144,55 @@ export function Inner({}: {optimisticState?: AgeAssuranceRedirectDialogState}) {
     }
   }, [agent, control, refreshAgeAssuranceState])
 
+  if (success) {
+    return (
+      <>
+        <View style={[a.align_start, a.w_full]}>
+          <AgeAssuranceBadge />
+
+          <View
+            style={[
+              a.flex_row,
+              a.justify_between,
+              a.align_center,
+              a.gap_sm,
+              a.pt_lg,
+              a.pb_md,
+            ]}>
+            <SuccessIcon size="sm" fill={t.palette.positive_600} />
+            <Text style={[a.text_xl, a.font_heavy]}>
+              <Trans>Success</Trans>
+            </Text>
+          </View>
+
+          <Text style={[a.text_md, a.leading_snug]}>
+            <Trans>
+              We've confirmed your age assurance status. You can now close this
+              dialog.
+            </Trans>
+          </Text>
+
+          {isNative && (
+            <View style={[a.w_full, a.pt_lg]}>
+              <Button
+                label={_(msg`Close`)}
+                size="large"
+                variant="solid"
+                color="secondary"
+                onPress={() => control.control.close()}>
+                <ButtonText>
+                  <Trans>Close</Trans>
+                </ButtonText>
+              </Button>
+            </View>
+          )}
+        </View>
+
+        <Dialog.Close />
+      </>
+    )
+  }
+
   return (
     <>
       <View style={[a.align_start, a.w_full]}>
@@ -175,8 +225,8 @@ export function Inner({}: {optimisticState?: AgeAssuranceRedirectDialogState}) {
             </Trans>
           ) : (
             <Trans>
-              We're confirming your status with our servers. This dialog should
-              close in a few seconds.
+              We're confirming your age assurance status with our servers. This
+              should only take a few seconds.
             </Trans>
           )}
         </Text>
@@ -187,7 +237,8 @@ export function Inner({}: {optimisticState?: AgeAssuranceRedirectDialogState}) {
               label={_(msg`Close`)}
               size="large"
               variant="solid"
-              color="secondary">
+              color="secondary"
+              onPress={() => control.control.close()}>
               <ButtonText>
                 <Trans>Close</Trans>
               </ButtonText>
diff --git a/src/state/ageAssurance/index.tsx b/src/state/ageAssurance/index.tsx
index eded74773..3451b1139 100644
--- a/src/state/ageAssurance/index.tsx
+++ b/src/state/ageAssurance/index.tsx
@@ -1,4 +1,4 @@
-import {createContext, useContext, useMemo} from 'react'
+import {createContext, useContext, useMemo, useState} from 'react'
 import {type AppBskyUnspeccedDefs} from '@atproto/api'
 import {useQuery} from '@tanstack/react-query'
 
@@ -46,6 +46,7 @@ export function Provider({children}: {children: React.ReactNode}) {
   const {geolocation} = useGeolocation()
   const isAgeAssuranceEnabled = useIsAgeAssuranceEnabled()
   const getAndRegisterPushToken = useGetAndRegisterPushToken()
+  const [refetchWhilePending, setRefetchWhilePending] = useState(false)
 
   const {data, isFetched, refetch} = useQuery({
     /**
@@ -56,6 +57,7 @@ export function Provider({children}: {children: React.ReactNode}) {
      * However, it only needs to run if AA is enabled.
      */
     enabled: isAgeAssuranceEnabled,
+    refetchOnWindowFocus: refetchWhilePending,
     queryKey: createAgeAssuranceQueryKey(agent.session?.did ?? 'never'),
     async queryFn() {
       if (!agent.session) return null
@@ -109,6 +111,24 @@ export function Provider({children}: {children: React.ReactNode}) {
     return ctx
   }, [isFetched, data, isAgeAssuranceEnabled])
 
+  if (
+    !!ageAssuranceContext.lastInitiatedAt &&
+    ageAssuranceContext.status === 'pending' &&
+    !refetchWhilePending
+  ) {
+    /*
+     * If we have a pending state, we want to refetch on window focus to ensure
+     * that we get the latest state when the user returns to the app.
+     */
+    setRefetchWhilePending(true)
+  } else if (
+    !!ageAssuranceContext.lastInitiatedAt &&
+    ageAssuranceContext.status !== 'pending' &&
+    refetchWhilePending
+  ) {
+    setRefetchWhilePending(false)
+  }
+
   const ageAssuranceAPIContext = useMemo<AgeAssuranceAPIContextType>(
     () => ({
       refetch,