Fix auth refresh in certain edge cases (#10218)

This PR should fix a bug when session doesn't refresh when a computer comes back from sleep mode

app/ide-desktop/lib/dashboard/src/components/StatelessSpinner.tsx

@@ -21,10 +21,14 @@ export default function StatelessSpinner(props: StatelessSpinnerProps) {
   const [state, setState] = React.useState(spinner.SpinnerState.initial)
 
   React.useEffect(() => {
-    window.setTimeout(() => {
+    const timeout = window.setTimeout(() => {
       setState(rawState)
     })
-  }, [/* should never change */ rawState])
+
+    return () => {
+      window.clearTimeout(timeout)
+    }
+  }, [rawState])
 
   return <Spinner state={state} {...(size != null ? { size } : {})} />
 }

app/ide-desktop/lib/dashboard/src/providers/AuthProvider.tsx

@@ -171,7 +171,7 @@ export default function AuthProvider(props: AuthProviderProps) {
   const { children, projectManagerUrl, projectManagerRootDirectory } = props
   const logger = loggerProvider.useLogger()
   const { cognito } = authService ?? {}
-  const { session, deinitializeSession, onSessionError } = sessionProvider.useSession()
+  const { session, onSessionError } = sessionProvider.useSession()
   const { setBackendWithoutSavingType } = backendProvider.useStrictSetBackend()
   const { localStorage } = localStorageProvider.useLocalStorage()
   const { getText } = textProvider.useText()
@@ -628,7 +628,6 @@ export default function AuthProvider(props: AuthProviderProps) {
       gtagEvent('cloud_sign_out')
       cognito.saveAccessToken(null)
       localStorage.clearUserSpecificEntries()
-      deinitializeSession()
       setInitialized(false)
       sentry.setUser(null)
       setUserSession(null)

app/ide-desktop/lib/dashboard/src/providers/SessionProvider.tsx

@@ -4,9 +4,6 @@ import * as React from 'react'
 
 import * as reactQuery from '@tanstack/react-query'
 
-import * as asyncEffectHooks from '#/hooks/asyncEffectHooks'
-import * as refreshHooks from '#/hooks/refreshHooks'
-
 import * as errorModule from '#/utilities/error'
 
 import type * as cognito from '#/authentication/cognito'
@@ -19,8 +16,6 @@ import * as listen from '#/authentication/listen'
 /** State contained in a {@link SessionContext}. */
 interface SessionContextType {
   readonly session: cognito.UserSession | null
-  /** Set `initialized` to false. Must be called when logging out. */
-  readonly deinitializeSession: () => void
   readonly onSessionError: (callback: (error: Error) => void) => () => void
 }
 
@@ -51,14 +46,14 @@ export interface SessionProviderProps {
 }
 
 const FIVE_MINUTES_MS = 300_000
+// const SIX_HOURS_MS = 21_600_000
 const SIX_HOURS_MS = 21_600_000
 
 /** A React provider for the session of the authenticated user. */
 export default function SessionProvider(props: SessionProviderProps) {
   const { mainPageUrl, children, userSession, registerAuthEventListener, refreshUserSession } =
     props
-  const [refresh, doRefresh] = refreshHooks.useRefresh()
-  const [initialized, setInitialized] = React.useState(false)
+
   const errorCallbacks = React.useRef(new Set<(error: Error) => void>())
 
   /** Returns a function to unregister the listener. */
@@ -69,47 +64,40 @@ export default function SessionProvider(props: SessionProviderProps) {
     }
   }, [])
 
-  // Register an async effect that will fetch the user's session whenever the `refresh` state is
-  // set. This is useful when a user has just logged in (as their cached credentials are
-  // out of date, so this will update them).
-  const session = asyncEffectHooks.useAsyncEffect(
-    null,
-    async () => {
-      if (userSession == null) {
-        setInitialized(true)
-        return null
-      } else {
-        try {
-          const innerSession = await userSession()
-          setInitialized(true)
-          return innerSession
-        } catch (error) {
-          if (error instanceof Error) {
-            for (const listener of errorCallbacks.current) {
-              listener(error)
-            }
-          }
-          throw error
-        }
-      }
-    },
-    [refresh]
-  )
+  const queryClient = reactQuery.useQueryClient()
 
-  const timeUntilRefresh = session
+  const session = reactQuery.useSuspenseQuery({
+    queryKey: ['userSession', userSession],
+    queryFn: userSession
+      ? () =>
+          userSession().catch(error => {
+            if (error instanceof Error) {
+              for (const listener of errorCallbacks.current) {
+                listener(error)
+              }
+            }
+            throw error
+          })
+      : reactQuery.skipToken,
+    refetchOnWindowFocus: true,
+    refetchIntervalInBackground: true,
+  })
+
+  const timeUntilRefresh = session.data
     ? // If the session has not expired, we should refresh it when it is 5 minutes from expiring.
-      new Date(session.expireAt).getTime() - Date.now() - FIVE_MINUTES_MS
+      new Date(session.data.expireAt).getTime() - Date.now() - FIVE_MINUTES_MS
     : Infinity
 
+  const refreshUserSessionMutation = reactQuery.useMutation({
+    mutationKey: ['refreshUserSession', session.data],
+    mutationFn: () => refreshUserSession?.().then(() => null) ?? Promise.resolve(),
+    meta: { invalidates: [['userSession']], awaitInvalidates: true },
+  })
+
   reactQuery.useQuery({
-    queryKey: ['userSession'],
+    queryKey: ['refreshUserSession'],
     queryFn: refreshUserSession
-      ? () =>
-          refreshUserSession()
-            .then(() => {
-              doRefresh()
-            })
-            .then(() => null)
+      ? () => refreshUserSessionMutation.mutateAsync()
       : reactQuery.skipToken,
     refetchOnWindowFocus: true,
     refetchIntervalInBackground: true,
@@ -119,7 +107,6 @@ export default function SessionProvider(props: SessionProviderProps) {
   // Register an effect that will listen for authentication events. When the event occurs, we
   // will refresh or clear the user's session, forcing a re-render of the page with the new
   // session.
-  //
   // For example, if a user clicks the "sign out" button, this will clear the user's session, which
   // means the login screen (which is a child of this provider) should render.
   React.useEffect(
@@ -128,7 +115,7 @@ export default function SessionProvider(props: SessionProviderProps) {
         switch (event) {
           case listen.AuthEvent.signIn:
           case listen.AuthEvent.signOut: {
-            doRefresh()
+            void queryClient.invalidateQueries({ queryKey: ['userSession'] })
             break
           }
           case listen.AuthEvent.customOAuthState:
@@ -139,7 +126,7 @@ export default function SessionProvider(props: SessionProviderProps) {
             // will not work.
             // See https://github.com/aws-amplify/amplify-js/issues/3391#issuecomment-756473970
             history.replaceState({}, '', mainPageUrl)
-            doRefresh()
+            void queryClient.invalidateQueries({ queryKey: ['userSession'] })
             break
           }
           default: {
@@ -147,16 +134,12 @@ export default function SessionProvider(props: SessionProviderProps) {
           }
         }
       }),
-    [doRefresh, registerAuthEventListener, mainPageUrl]
+    [registerAuthEventListener, mainPageUrl, queryClient]
   )
 
-  const deinitializeSession = () => {
-    setInitialized(false)
-  }
-
   return (
-    <SessionContext.Provider value={{ session, deinitializeSession, onSessionError }}>
-      {initialized && children}
+    <SessionContext.Provider value={{ session: session.data, onSessionError }}>
+      {children}
     </SessionContext.Provider>
   )
 }