Fix logout (#12399)

Author: MrFlashAccount

app/gui/src/dashboard/providers/SessionProvider.tsx

@@ -115,9 +115,9 @@ export default function SessionProvider(props: SessionProviderProps) {
   const queryClient = reactQuery.useQueryClient()
   const toastAndLog = useToastAndLog()
 
-  const sessionQuery = createSessionQuery(authService)
+  const sessionQueryOptions = createSessionQuery(authService)
 
-  const session = reactQuery.useSuspenseQuery(sessionQuery)
+  const session = reactQuery.useSuspenseQuery(sessionQueryOptions)
 
   const refreshUserSessionMutation = reactQuery.useMutation({
     mutationKey: ['refreshUserSession', { expireAt: session.data?.expireAt }],
@@ -126,12 +126,12 @@ export default function SessionProvider(props: SessionProviderProps) {
       if (data) {
         httpClient?.setSessionToken(data.accessToken)
       }
-      return queryClient.invalidateQueries({ queryKey: sessionQuery.queryKey })
+      return queryClient.invalidateQueries({ queryKey: sessionQueryOptions.queryKey })
     },
     onError: (error) => {
       // Something went wrong with the refresh token, so we need to sign the user out.
       toastAndLog('sessionExpiredError', error)
-      queryClient.setQueryData(sessionQuery.queryKey, null)
+      queryClient.setQueryData(sessionQueryOptions.queryKey, null)
       return logoutMutation.mutateAsync()
     },
   })
@@ -162,6 +162,7 @@ export default function SessionProvider(props: SessionProviderProps) {
       void queryClient.clearWithPersister()
     },
     onError: () => toast.error(getText('signOutError')),
+    meta: { invalidates: [sessionQueryOptions.queryKey], awaitInvalidates: true },
   })
 
   const signUp = useEventCallback(
@@ -209,7 +210,7 @@ export default function SessionProvider(props: SessionProviderProps) {
       }
 
       return queryClient
-        .invalidateQueries({ queryKey: sessionQuery.queryKey })
+        .invalidateQueries({ queryKey: sessionQueryOptions.queryKey })
         .then(() => ({ challenge, user }) as const)
     } else {
       throw new Error(result.val.message)
@@ -282,7 +283,7 @@ export default function SessionProvider(props: SessionProviderProps) {
         switch (event) {
           case listen.AuthEvent.signIn:
           case listen.AuthEvent.signOut: {
-            void queryClient.invalidateQueries({ queryKey: sessionQuery.queryKey })
+            void queryClient.invalidateQueries({ queryKey: sessionQueryOptions.queryKey })
             break
           }
           case listen.AuthEvent.customOAuthState:
@@ -293,15 +294,15 @@ export default function SessionProvider(props: SessionProviderProps) {
             // will not work.
             // See https://github.com/aws-amplify/amplify-js/issues/3391#issuecomment-756473970
             history.replaceState({}, '', mainPageUrl)
-            void queryClient.invalidateQueries({ queryKey: sessionQuery.queryKey })
+            void queryClient.invalidateQueries({ queryKey: sessionQueryOptions.queryKey })
             break
           }
           default: {
             throw new errorModule.UnreachableCaseError(event)
           }
         }
       }),
-    [registerAuthEventListener, mainPageUrl, queryClient, sessionQuery.queryKey],
+    [registerAuthEventListener, mainPageUrl, queryClient, sessionQueryOptions.queryKey],
   )
 
   const organizationId = useEventCallback(() => authService.organizationId())