From b7bc36d32ab85eda5efe2f9db3d46509c9d4220b Mon Sep 17 00:00:00 2001
From: anovazzi1 <otavio2204@gmail.com>
Date: Tue, 23 Jul 2024 11:59:13 -0300
Subject: [PATCH] refactor: authentication logic and move isAdmin to zustand
 store (#2888)

* refactor: remove isAuthenticated from context and move it to zustand store

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Cristhian Zanforlin Lousa <72977554+Cristhianzl@users.noreply.github.com>
---
 src/frontend/src/components/authAdminGuard/index.tsx  | 4 ++--
 src/frontend/src/components/headerComponent/index.tsx | 3 ++-
 src/frontend/src/contexts/authContext.tsx             | 9 ++-------
 src/frontend/src/types/contexts/auth.ts               | 2 --
 4 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/src/frontend/src/components/authAdminGuard/index.tsx b/src/frontend/src/components/authAdminGuard/index.tsx
index 972b96eb46b0..a08e1f4a9167 100644
--- a/src/frontend/src/components/authAdminGuard/index.tsx
+++ b/src/frontend/src/components/authAdminGuard/index.tsx
@@ -4,10 +4,10 @@ import { Navigate } from "react-router-dom";
 import { AuthContext } from "../../contexts/authContext";
 
 export const ProtectedAdminRoute = ({ children }) => {
-  const { isAdmin, logout, userData } = useContext(AuthContext);
+  const { logout, userData } = useContext(AuthContext);
   const isAuthenticated = useAuthStore((state) => state.isAuthenticated);
   const autoLogin = useAuthStore((state) => state.autoLogin);
-
+  const isAdmin = useAuthStore((state) => state.isAdmin);
   if (!isAuthenticated) {
     logout();
   } else if ((userData && !isAdmin) || autoLogin) {
diff --git a/src/frontend/src/components/headerComponent/index.tsx b/src/frontend/src/components/headerComponent/index.tsx
index ed034a778454..ff83341a7c3a 100644
--- a/src/frontend/src/components/headerComponent/index.tsx
+++ b/src/frontend/src/components/headerComponent/index.tsx
@@ -35,7 +35,8 @@ export default function Header(): JSX.Element {
   const notificationCenter = useAlertStore((state) => state.notificationCenter);
   const location = useLocation();
 
-  const { logout, isAdmin, userData } = useContext(AuthContext);
+  const { logout, userData } = useContext(AuthContext);
+  const isAdmin = useAuthStore((state) => state.isAdmin);
   const autoLogin = useAuthStore((state) => state.autoLogin);
 
   const navigate = useNavigate();
diff --git a/src/frontend/src/contexts/authContext.tsx b/src/frontend/src/contexts/authContext.tsx
index 1ae2adeb98f8..b1d415aca1a6 100644
--- a/src/frontend/src/contexts/authContext.tsx
+++ b/src/frontend/src/contexts/authContext.tsx
@@ -21,8 +21,6 @@ import { Users } from "../types/api";
 import { AuthContextType } from "../types/contexts/auth";
 
 const initialValue: AuthContextType = {
-  isAdmin: false,
-  setIsAdmin: () => false,
   accessToken: null,
   login: () => {},
   logout: () => new Promise(() => {}),
@@ -43,7 +41,6 @@ export function AuthProvider({ children }): React.ReactElement {
   const [accessToken, setAccessToken] = useState<string | null>(
     cookies.get(LANGFLOW_ACCESS_TOKEN) ?? null,
   );
-  const [isAdmin, setIsAdmin] = useState<boolean>(false);
   const [userData, setUserData] = useState<Users | null>(null);
   const setLoading = useAlertStore((state) => state.setLoading);
   const [apiKey, setApiKey] = useState<string | null>(
@@ -78,7 +75,7 @@ export function AuthProvider({ children }): React.ReactElement {
       .then(async (user) => {
         setUserData(user);
         const isSuperUser = user!.is_superuser;
-        setIsAdmin(isSuperUser);
+        useAuthStore.getState().setIsAdmin(isSuperUser);
         getFoldersApi(true, true);
         const res = await getGlobalVariables();
         setGlobalVariables(res);
@@ -105,7 +102,7 @@ export function AuthProvider({ children }): React.ReactElement {
       await requestLogout();
       cookies.remove(LANGFLOW_API_TOKEN, { path: "/" });
       cookies.remove(LANGFLOW_AUTO_LOGIN_OPTION, { path: "/" });
-      setIsAdmin(false);
+      useAuthStore.getState().setIsAdmin(false);
       setUserData(null);
       setAccessToken(null);
       useAuthStore.getState().setIsAuthenticated(false);
@@ -128,8 +125,6 @@ export function AuthProvider({ children }): React.ReactElement {
     // !! to convert string to boolean
     <AuthContext.Provider
       value={{
-        isAdmin,
-        setIsAdmin,
         accessToken,
         login,
         logout,
diff --git a/src/frontend/src/types/contexts/auth.ts b/src/frontend/src/types/contexts/auth.ts
index cbfbe6895769..32dbef825484 100644
--- a/src/frontend/src/types/contexts/auth.ts
+++ b/src/frontend/src/types/contexts/auth.ts
@@ -1,8 +1,6 @@
 import { Users } from "../api";
 
 export type AuthContextType = {
-  isAdmin: boolean;
-  setIsAdmin: (isAdmin: boolean) => void;
   accessToken: string | null;
   login: (accessToken: string, autoLogin: string) => void;
   logout: () => Promise<void>;