From b6eff682d38750a2da3d3b5a9b66d1bfcbf4b529 Mon Sep 17 00:00:00 2001
From: Mathijs de Bruin <mathijs@mathijsfietst.nl>
Date: Mon, 25 Nov 2024 18:26:46 +0000
Subject: [PATCH] WIP ugly WIP

---
 frontend/src/pages/AuthCallback.tsx           |  2 +
 libs/react-client/src/api/hooks/auth/index.ts |  1 +
 libs/react-client/src/auth/token.ts           |  2 +
 libs/react-client/src/auth/user.ts            | 66 +++++++++----------
 libs/react-client/src/useConfig.ts            |  3 +
 5 files changed, 41 insertions(+), 33 deletions(-)

diff --git a/frontend/src/pages/AuthCallback.tsx b/frontend/src/pages/AuthCallback.tsx
index 753651686c..1b56bad19d 100644
--- a/frontend/src/pages/AuthCallback.tsx
+++ b/frontend/src/pages/AuthCallback.tsx
@@ -11,6 +11,8 @@ export default function AuthCallback() {
   const navigate = useNavigate();
 
   useEffect(() => {
+    console.log('AuthCallbackEffect');
+
     if (!cookieAuth) {
       // Legacy auth token from request query parameter
       const token = query.get('access_token');
diff --git a/libs/react-client/src/api/hooks/auth/index.ts b/libs/react-client/src/api/hooks/auth/index.ts
index aaa1537d4d..e18763226e 100644
--- a/libs/react-client/src/api/hooks/auth/index.ts
+++ b/libs/react-client/src/api/hooks/auth/index.ts
@@ -7,6 +7,7 @@ import { useUser } from 'src/auth/user';
 import { accessTokenState } from 'src/state';
 
 export const useAuth = (): IUseAuth => {
+  console.log('useAuth');
   const { authConfig, isLoading, cookieAuth } = useAuthConfig();
   const { logout } = useSessionManagement();
   const { user, setUserFromAPI } = useUser();
diff --git a/libs/react-client/src/auth/token.ts b/libs/react-client/src/auth/token.ts
index f1b404ad98..2ca77a4ae9 100644
--- a/libs/react-client/src/auth/token.ts
+++ b/libs/react-client/src/auth/token.ts
@@ -41,6 +41,8 @@ export function ensureTokenPrefix(token: string): string {
 }
 
 export const useTokenManagement = () => {
+  console.log('useTokenManagement');
+
   const [, setUser] = useRecoilState(userState);
   const [, setAccessToken] = useRecoilState(accessTokenState);
 
diff --git a/libs/react-client/src/auth/user.ts b/libs/react-client/src/auth/user.ts
index 05b56ac6e6..c39ec6ade7 100644
--- a/libs/react-client/src/auth/user.ts
+++ b/libs/react-client/src/auth/user.ts
@@ -1,54 +1,54 @@
-import { useContext, useEffect } from 'react';
+import { useEffect } from 'react';
 import { useRecoilState } from 'recoil';
-import { ChainlitContext } from 'src/context';
 import { userState } from 'src/state';
 
+import { IUser, useApi } from '..';
 import { useAuthConfig } from './config';
 import { getToken, useTokenManagement } from './token';
 
 export const useUser = () => {
-  const apiClient = useContext(ChainlitContext);
+  console.log('useUser');
+
   const [user, setUser] = useRecoilState(userState);
   const { cookieAuth } = useAuthConfig();
   const { handleSetAccessToken } = useTokenManagement();
 
-  // Legacy token auth; initialize the token from local storage
-  const setUserFromLocalStore = () => {
-    {
-      const storedAccessToken = getToken();
-      if (storedAccessToken) handleSetAccessToken(storedAccessToken);
-    }
-  };
-
-  // Cookie-based auth; use API to set user.
-  const setUserFromAPI = async () => {
-    // Get user from cookie, return true when successful
-    try {
-      const apiUser = await apiClient.getUser();
-      if (apiUser) {
-        setUser(apiUser);
+  const { data: userData, mutate: mutateUserData } = useApi<IUser>('/user', {
+    revalidateOnMount: false
+  });
+
+  // Attempt to get user when cookieAuth are available.
+  useEffect(() => {
+    if (!user) {
+      if (cookieAuth) {
+        console.log('cookieAuth', user, cookieAuth);
+        mutateUserData();
+        return;
       }
-    } catch (_) {
-      return;
-    }
-  };
 
-  const initUser = () => {
-    // Already logged in
-    if (user) return;
+      // Not using cookie auth, callback to header tokens
+      console.log('tokenAuth', user, cookieAuth);
+      const token = getToken();
+      if (token) handleSetAccessToken(token);
+    }
+  }, [user, cookieAuth]);
 
-    // Legacy fallback
-    if (!cookieAuth) return setUserFromLocalStore();
+  // When user data is available, set the user object.
+  useEffect(() => {
+    console.log('userData effect');
 
-    // Request user from API
-    setUserFromAPI();
-  };
+    if (userData) {
+      console.log('setUser', userData);
+      setUser(userData);
+    }
+  }, [userData]);
 
-  // Attempt to initialize user on app start.
-  useEffect(initUser, [cookieAuth]);
+  useEffect(() => {
+    console.log('useUser effect');
+  });
 
   return {
     user,
-    setUserFromAPI
+    setUserFromAPI: mutateUserData
   };
 };
diff --git a/libs/react-client/src/useConfig.ts b/libs/react-client/src/useConfig.ts
index a802199136..11efb0cc2a 100644
--- a/libs/react-client/src/useConfig.ts
+++ b/libs/react-client/src/useConfig.ts
@@ -6,6 +6,8 @@ import { configState } from './state';
 import { IChainlitConfig } from './types';
 
 const useConfig = (accessToken?: string) => {
+  console.log('useConfig');
+
   const [config, setConfig] = useRecoilState(configState);
   const { isAuthenticated } = useAuth();
   const language = navigator.language || 'en-US';
@@ -19,6 +21,7 @@ const useConfig = (accessToken?: string) => {
 
   useEffect(() => {
     if (!data) return;
+    console.log('setConfig');
     setConfig(data);
   }, [data, setConfig]);