From 968b8ccdbdd8cf80452ade2a2541825dc2b6e9bf Mon Sep 17 00:00:00 2001
From: Danny Avila <110412045+danny-avila@users.noreply.github.com>
Date: Sun, 10 Dec 2023 17:13:42 -0500
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=92=20fix:=20Robust=20Cache=20Reset=20?=
 =?UTF-8?q?on=20User=20Logout=20(#1324)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* refactor(Logout): rely on hooks for mutation behavior

* fix: logging out now correctly resets cache, disallowing any cache mixing between the next logged in user on the same browser

* chore: remove additional localStorage values on logout
---
 client/src/components/Nav/Logout.tsx          | 10 +--
 client/src/data-provider/mutations.ts         | 32 +++++++++-
 client/src/data-provider/queries.ts           | 11 ++++
 client/src/hooks/AuthContext.tsx              | 63 +++++++++----------
 client/src/hooks/Conversations/usePresets.ts  | 31 +++++----
 packages/data-provider/src/data-service.ts    |  4 +-
 packages/data-provider/src/index.ts           |  2 +-
 packages/data-provider/src/keys.ts            |  1 +
 .../data-provider/src/react-query-service.ts  | 39 +++---------
 .../src/types/{presets.ts => mutations.ts}    |  6 ++
 10 files changed, 109 insertions(+), 90 deletions(-)
 rename packages/data-provider/src/types/{presets.ts => mutations.ts} (73%)

diff --git a/client/src/components/Nav/Logout.tsx b/client/src/components/Nav/Logout.tsx
index 4a7eca5e44..a16f4f3266 100644
--- a/client/src/components/Nav/Logout.tsx
+++ b/client/src/components/Nav/Logout.tsx
@@ -7,18 +7,10 @@ const Logout = forwardRef(() => {
   const { logout } = useAuthContext();
   const localize = useLocalize();
 
-  const handleLogout = () => {
-    localStorage.removeItem('lastConversationSetup');
-    localStorage.removeItem('lastSelectedTools');
-    localStorage.removeItem('lastAssistant');
-    localStorage.removeItem('autoScroll');
-    logout();
-  };
-
   return (
     <button
       className="flex w-full cursor-pointer items-center gap-3 px-3 py-3 text-sm text-white transition-colors duration-200 hover:bg-gray-700"
-      onClick={handleLogout}
+      onClick={() => logout()}
     >
       <LogOutIcon />
       {localize('com_nav_log_out')}
diff --git a/client/src/data-provider/mutations.ts b/client/src/data-provider/mutations.ts
index 26c5185500..c38ca46c9e 100644
--- a/client/src/data-provider/mutations.ts
+++ b/client/src/data-provider/mutations.ts
@@ -1,4 +1,4 @@
-import { useMutation } from '@tanstack/react-query';
+import { useMutation, useQueryClient } from '@tanstack/react-query';
 import type { UseMutationResult } from '@tanstack/react-query';
 import type {
   FileUploadResponse,
@@ -10,9 +10,13 @@ import type {
   UpdatePresetOptions,
   DeletePresetOptions,
   PresetDeleteResponse,
+  LogoutOptions,
   TPreset,
 } from 'librechat-data-provider';
+
 import { dataService, MutationKeys } from 'librechat-data-provider';
+import { useSetRecoilState } from 'recoil';
+import store from '~/store';
 
 export const useUploadImageMutation = (
   options?: UploadMutationOptions,
@@ -69,3 +73,29 @@ export const useDeletePresetMutation = (
     ...(options || {}),
   });
 };
+
+/* login/logout */
+export const useLogoutUserMutation = (
+  options?: LogoutOptions,
+): UseMutationResult<unknown, unknown, undefined, unknown> => {
+  const queryClient = useQueryClient();
+  const setDefaultPreset = useSetRecoilState(store.defaultPreset);
+  return useMutation([MutationKeys.logoutUser], {
+    mutationFn: () => dataService.logout(),
+
+    ...(options || {}),
+    onSuccess: (...args) => {
+      options?.onSuccess?.(...args);
+    },
+    onMutate: (...args) => {
+      setDefaultPreset(null);
+      queryClient.removeQueries();
+      localStorage.removeItem('lastConversationSetup');
+      localStorage.removeItem('lastSelectedModel');
+      localStorage.removeItem('lastSelectedTools');
+      localStorage.removeItem('filesToDelete');
+      localStorage.removeItem('lastAssistant');
+      options?.onMutate?.(...args);
+    },
+  });
+};
diff --git a/client/src/data-provider/queries.ts b/client/src/data-provider/queries.ts
index ec10c9245d..4144b64812 100644
--- a/client/src/data-provider/queries.ts
+++ b/client/src/data-provider/queries.ts
@@ -1,5 +1,16 @@
 import { UseQueryOptions, useQuery, QueryObserverResult } from '@tanstack/react-query';
 import { QueryKeys, dataService } from 'librechat-data-provider';
+import type { TPreset } from 'librechat-data-provider';
+export const useGetPresetsQuery = (
+  config?: UseQueryOptions<TPreset[]>,
+): QueryObserverResult<TPreset[], unknown> => {
+  return useQuery<TPreset[]>([QueryKeys.presets], () => dataService.getPresets(), {
+    refetchOnWindowFocus: false,
+    refetchOnReconnect: false,
+    refetchOnMount: false,
+    ...config,
+  });
+};
 
 export const useGetEndpointsConfigOverride = <TData = unknown | boolean>(
   config?: UseQueryOptions<unknown | boolean, unknown, TData>,
diff --git a/client/src/hooks/AuthContext.tsx b/client/src/hooks/AuthContext.tsx
index 356df477b1..e30f74c84e 100644
--- a/client/src/hooks/AuthContext.tsx
+++ b/client/src/hooks/AuthContext.tsx
@@ -12,13 +12,13 @@ import {
   TLoginResponse,
   setTokenHeader,
   useLoginUserMutation,
-  useLogoutUserMutation,
   useGetUserQuery,
   useRefreshTokenMutation,
   TLoginUser,
 } from 'librechat-data-provider';
-import { TAuthConfig, TUserContext, TAuthContext, TResError } from '~/common';
 import { useNavigate } from 'react-router-dom';
+import { TAuthConfig, TUserContext, TAuthContext, TResError } from '~/common';
+import { useLogoutUserMutation } from '~/data-provider';
 import useTimeout from './useTimeout';
 
 const AuthContext = createContext<TAuthContext | undefined>(undefined);
@@ -30,20 +30,11 @@ const AuthContextProvider = ({
   authConfig?: TAuthConfig;
   children: ReactNode;
 }) => {
+  const navigate = useNavigate();
   const [user, setUser] = useState<TUser | undefined>(undefined);
   const [token, setToken] = useState<string | undefined>(undefined);
   const [error, setError] = useState<string | undefined>(undefined);
   const [isAuthenticated, setIsAuthenticated] = useState<boolean>(false);
-
-  const navigate = useNavigate();
-
-  const loginUser = useLoginUserMutation();
-  const logoutUser = useLogoutUserMutation();
-  const userQuery = useGetUserQuery({ enabled: !!token });
-  const refreshToken = useRefreshTokenMutation();
-
-  const doSetError = useTimeout({ callback: (error) => setError(error as string | undefined) });
-
   const setUserContext = useCallback(
     (userContext: TUserContext) => {
       const { token, isAuthenticated, user, redirect } = userContext;
@@ -60,6 +51,32 @@ const AuthContextProvider = ({
     },
     [navigate],
   );
+  const doSetError = useTimeout({ callback: (error) => setError(error as string | undefined) });
+
+  const loginUser = useLoginUserMutation();
+  const logoutUser = useLogoutUserMutation({
+    onSuccess: () => {
+      setUserContext({
+        token: undefined,
+        isAuthenticated: false,
+        user: undefined,
+        redirect: '/login',
+      });
+    },
+    onError: (error) => {
+      doSetError((error as Error).message);
+      setUserContext({
+        token: undefined,
+        isAuthenticated: false,
+        user: undefined,
+        redirect: '/login',
+      });
+    },
+  });
+
+  const logout = useCallback(() => logoutUser.mutate(undefined), [logoutUser]);
+  const userQuery = useGetUserQuery({ enabled: !!token });
+  const refreshToken = useRefreshTokenMutation();
 
   const login = (data: TLoginUser) => {
     loginUser.mutate(data, {
@@ -75,28 +92,6 @@ const AuthContextProvider = ({
     });
   };
 
-  const logout = useCallback(() => {
-    logoutUser.mutate(undefined, {
-      onSuccess: () => {
-        setUserContext({
-          token: undefined,
-          isAuthenticated: false,
-          user: undefined,
-          redirect: '/login',
-        });
-      },
-      onError: (error) => {
-        doSetError((error as Error).message);
-        setUserContext({
-          token: undefined,
-          isAuthenticated: false,
-          user: undefined,
-          redirect: '/login',
-        });
-      },
-    });
-  }, [setUserContext, doSetError, logoutUser]);
-
   const silentRefresh = useCallback(() => {
     if (authConfig?.test) {
       console.log('Test mode. Skipping silent refresh.');
diff --git a/client/src/hooks/Conversations/usePresets.ts b/client/src/hooks/Conversations/usePresets.ts
index 7af73657b4..ac6105c815 100644
--- a/client/src/hooks/Conversations/usePresets.ts
+++ b/client/src/hooks/Conversations/usePresets.ts
@@ -1,16 +1,15 @@
-import {
-  QueryKeys,
-  modularEndpoints,
-  useGetPresetsQuery,
-  useCreatePresetMutation,
-} from 'librechat-data-provider';
+import { QueryKeys, modularEndpoints, useCreatePresetMutation } from 'librechat-data-provider';
 import filenamify from 'filenamify';
 import { useCallback, useEffect, useRef } from 'react';
 import { useRecoilState, useSetRecoilState } from 'recoil';
 import exportFromJSON from 'export-from-json';
 import { useQueryClient } from '@tanstack/react-query';
 import type { TPreset } from 'librechat-data-provider';
-import { useUpdatePresetMutation, useDeletePresetMutation } from '~/data-provider';
+import {
+  useUpdatePresetMutation,
+  useDeletePresetMutation,
+  useGetPresetsQuery,
+} from '~/data-provider';
 import { useChatContext, useToastContext } from '~/Providers';
 import useNavigateToConvo from '~/hooks/useNavigateToConvo';
 import useDefaultConvo from '~/hooks/useDefaultConvo';
@@ -22,22 +21,28 @@ import store from '~/store';
 
 export default function usePresets() {
   const localize = useLocalize();
-  const { user } = useAuthContext();
+  const hasLoaded = useRef(false);
   const queryClient = useQueryClient();
   const { showToast } = useToastContext();
-  const hasLoaded = useRef(false);
+  const { user, isAuthenticated } = useAuthContext();
 
   const [_defaultPreset, setDefaultPreset] = useRecoilState(store.defaultPreset);
   const setPresetModalVisible = useSetRecoilState(store.presetModalVisible);
   const { preset, conversation, newConversation, setPreset } = useChatContext();
-  const presetsQuery = useGetPresetsQuery({ enabled: !!user });
+  const presetsQuery = useGetPresetsQuery({ enabled: !!user && isAuthenticated });
 
   useEffect(() => {
-    if (_defaultPreset || !presetsQuery.data || hasLoaded.current) {
+    const { data: presets } = presetsQuery;
+    if (_defaultPreset || !presets || hasLoaded.current) {
       return;
     }
 
-    const defaultPreset = presetsQuery.data.find((p) => p.defaultPreset);
+    if (presets && presets.length > 0 && user && presets[0].user !== user?.id) {
+      presetsQuery.refetch();
+      return;
+    }
+
+    const defaultPreset = presets.find((p) => p.defaultPreset);
     if (!defaultPreset) {
       hasLoaded.current = true;
       return;
@@ -49,7 +54,7 @@ export default function usePresets() {
     hasLoaded.current = true;
     // dependencies are stable and only needed once
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [presetsQuery.data]);
+  }, [presetsQuery.data, user]);
 
   const setPresets = useCallback(
     (presets: TPreset[]) => {
diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts
index 105783d395..b02565652b 100644
--- a/packages/data-provider/src/data-service.ts
+++ b/packages/data-provider/src/data-service.ts
@@ -1,5 +1,5 @@
 import * as f from './types/files';
-import * as p from './types/presets';
+import * as m from './types/mutations';
 import * as a from './types/assistants';
 import * as t from './types';
 import * as s from './schemas';
@@ -82,7 +82,7 @@ export function updatePreset(payload: s.TPreset): Promise<s.TPreset> {
   return request.post(endpoints.presets(), payload);
 }
 
-export function deletePreset(arg: s.TPreset | undefined): Promise<p.PresetDeleteResponse> {
+export function deletePreset(arg: s.TPreset | undefined): Promise<m.PresetDeleteResponse> {
   return request.post(endpoints.deletePreset(), arg);
 }
 
diff --git a/packages/data-provider/src/index.ts b/packages/data-provider/src/index.ts
index 2cb3aa2e24..ab91ff3440 100644
--- a/packages/data-provider/src/index.ts
+++ b/packages/data-provider/src/index.ts
@@ -2,7 +2,7 @@
 export * from './types';
 export * from './types/assistants';
 export * from './types/files';
-export * from './types/presets';
+export * from './types/mutations';
 /*
  * react query
  * TODO: move to client, or move schemas/types to their own package
diff --git a/packages/data-provider/src/keys.ts b/packages/data-provider/src/keys.ts
index 72faf08c35..6da00f7fc4 100644
--- a/packages/data-provider/src/keys.ts
+++ b/packages/data-provider/src/keys.ts
@@ -23,4 +23,5 @@ export enum MutationKeys {
   fileDelete = 'fileDelete',
   updatePreset = 'updatePreset',
   deletePreset = 'deletePreset',
+  logoutUser = 'logoutUser',
 }
diff --git a/packages/data-provider/src/react-query-service.ts b/packages/data-provider/src/react-query-service.ts
index e73023d151..11d1fd8596 100644
--- a/packages/data-provider/src/react-query-service.ts
+++ b/packages/data-provider/src/react-query-service.ts
@@ -8,7 +8,7 @@ import {
 } from '@tanstack/react-query';
 import * as t from './types';
 import * as s from './schemas';
-import * as p from './types/presets';
+import * as m from './types/mutations';
 import * as dataService from './data-service';
 import request from './request';
 import { QueryKeys } from './keys';
@@ -306,19 +306,8 @@ export const useUpdatePresetMutation = (): UseMutationResult<
   });
 };
 
-export const useGetPresetsQuery = (
-  config?: UseQueryOptions<s.TPreset[]>,
-): QueryObserverResult<s.TPreset[], unknown> => {
-  return useQuery<s.TPreset[]>([QueryKeys.presets], () => dataService.getPresets(), {
-    refetchOnWindowFocus: false,
-    refetchOnReconnect: false,
-    refetchOnMount: false,
-    ...config,
-  });
-};
-
 export const useDeletePresetMutation = (): UseMutationResult<
-  p.PresetDeleteResponse,
+  m.PresetDeleteResponse,
   unknown,
   s.TPreset | undefined,
   unknown
@@ -370,14 +359,8 @@ export const useLoginUserMutation = (): UseMutationResult<
 > => {
   const queryClient = useQueryClient();
   return useMutation((payload: t.TLoginUser) => dataService.login(payload), {
-    onSuccess: () => {
-      queryClient.invalidateQueries([QueryKeys.user]);
-      queryClient.invalidateQueries([QueryKeys.presets]);
-      queryClient.invalidateQueries([QueryKeys.conversation]);
-      queryClient.invalidateQueries([QueryKeys.allConversations]);
-    },
     onMutate: () => {
-      queryClient.invalidateQueries([QueryKeys.models]);
+      queryClient.removeQueries();
     },
   });
 };
@@ -396,15 +379,6 @@ export const useRegisterUserMutation = (): UseMutationResult<
   });
 };
 
-export const useLogoutUserMutation = (): UseMutationResult<unknown> => {
-  const queryClient = useQueryClient();
-  return useMutation(() => dataService.logout(), {
-    onSuccess: () => {
-      queryClient.invalidateQueries([QueryKeys.user]);
-    },
-  });
-};
-
 export const useRefreshTokenMutation = (): UseMutationResult<
   t.TRefreshTokenResponse,
   unknown,
@@ -414,7 +388,12 @@ export const useRefreshTokenMutation = (): UseMutationResult<
   const queryClient = useQueryClient();
   return useMutation(() => request.refreshToken(), {
     onMutate: () => {
-      queryClient.invalidateQueries([QueryKeys.models]);
+      queryClient.removeQueries();
+      localStorage.removeItem('lastConversationSetup');
+      localStorage.removeItem('lastSelectedModel');
+      localStorage.removeItem('lastSelectedTools');
+      localStorage.removeItem('filesToDelete');
+      localStorage.removeItem('lastAssistant');
     },
   });
 };
diff --git a/packages/data-provider/src/types/presets.ts b/packages/data-provider/src/types/mutations.ts
similarity index 73%
rename from packages/data-provider/src/types/presets.ts
rename to packages/data-provider/src/types/mutations.ts
index 2beaf307a6..569613f696 100644
--- a/packages/data-provider/src/types/presets.ts
+++ b/packages/data-provider/src/types/mutations.ts
@@ -20,3 +20,9 @@ export type DeletePresetOptions = {
   onMutate?: (variables: TPreset | undefined) => void | Promise<unknown>;
   onError?: (error: unknown, variables: TPreset | undefined, context?: unknown) => void;
 };
+
+export type LogoutOptions = {
+  onSuccess?: (data: unknown, variables: undefined, context?: unknown) => void;
+  onMutate?: (variables: undefined) => void | Promise<unknown>;
+  onError?: (error: unknown, variables: undefined, context?: unknown) => void;
+};