Factor out user setting, cleanup auth API.

libs/react-client/src/api/hooks/auth/index.ts

@@ -1,41 +1,21 @@
-import { useContext, useEffect } from 'react';
 import { useRecoilState } from 'recoil';
 import { useAuthConfig } from 'src/auth/config';
 import { useSessionManagement } from 'src/auth/session';
-import { getToken, useTokenManagement } from 'src/auth/token';
+import { useTokenManagement } from 'src/auth/token';
 import { IUseAuth } from 'src/auth/types';
-import { ChainlitContext } from 'src/context';
-import { accessTokenState, userState } from 'src/state';
+import { useUser } from 'src/auth/user';
+import { accessTokenState } from 'src/state';
 
-// Define useAuth hook
 export const useAuth = (): IUseAuth => {
-  const apiClient = useContext(ChainlitContext);
-
   const { authConfig, isLoading } = useAuthConfig();
-  const { logout } = useSessionManagement(apiClient);
-
-  const [user] = useRecoilState(userState);
+  const { logout } = useSessionManagement();
+  const { user } = useUser();
   const [accessToken] = useRecoilState(accessTokenState);
 
   const { handleSetAccessToken } = useTokenManagement();
 
   const isReady = !!(!isLoading && authConfig);
 
-  useEffect(() => {
-    if (authConfig?.cookieAuth) {
-      // TODO: Do something to get user object.
-      // setUser(user as IUser)
-      return;
-    }
-
-    const storedAccessToken = getToken();
-    if (!!user && storedAccessToken) {
-      // Initialize the token from local storage
-      handleSetAccessToken(storedAccessToken);
-      return;
-    }
-  }, []);
-
   if (!authConfig?.requireLogin) {
     return {
       data: authConfig,

libs/react-client/src/api/index.tsx

@@ -1,5 +1,5 @@
 import { ensureTokenPrefix, removeToken } from 'src/auth/token';
-import { IThread } from 'src/types';
+import { IThread, IUser } from 'src/types';
 
 import { IFeedback } from 'src/types/feedback';
 
@@ -159,6 +159,11 @@ export class ChainlitAPI extends APIBase {
     return res.json();
   }
 
+  async getUser(accessToken?: string): Promise<IUser> {
+    const res = await this.get(`/user`, accessToken);
+    return res.json();
+  }
+
   async logout() {
     const res = await this.post(`/logout`, {});
     return res.json();

libs/react-client/src/auth/config.ts

@@ -17,9 +17,13 @@ export const useAuthConfig = () => {
     }
   }, [authConfigData, setAuthConfig]);
 
+  // Secure default: only set false if explicitly defined.
+  const cookieAuth: boolean = authConfig?.cookieAuth !== false;
+
   return {
     authConfig,
     isLoading,
-    setAuthConfig
+    setAuthConfig,
+    cookieAuth
   };
 };

libs/react-client/src/auth/session.ts

@@ -1,11 +1,13 @@
-import { useRecoilState } from 'recoil';
-import { useSetRecoilState } from 'recoil';
+import { useContext } from 'react';
+import { useRecoilState, useSetRecoilState } from 'recoil';
+import { ChainlitContext } from 'src/context';
 import { accessTokenState, threadHistoryState, userState } from 'src/state';
 
 import { useAuthConfig } from './config';
 import { removeToken } from './token';
 
-export const useSessionManagement = (apiClient: any) => {
+export const useSessionManagement = () => {
+  const apiClient = useContext(ChainlitContext);
   const [, setUser] = useRecoilState(userState);
   const [, setAccessToken] = useRecoilState(accessTokenState);
 

libs/react-client/src/auth/token.ts

@@ -1,7 +1,5 @@
 import jwt_decode from 'jwt-decode';
-import { useContext } from 'react';
 import { useRecoilState } from 'recoil';
-import { ChainlitContext } from 'src/context';
 import { accessTokenState, userState } from 'src/state';
 
 import { useSessionManagement } from './session';
@@ -46,7 +44,7 @@ export const useTokenManagement = () => {
   const [, setUser] = useRecoilState(userState);
   const [, setAccessToken] = useRecoilState(accessTokenState);
 
-  const { logout } = useSessionManagement(useContext(ChainlitContext));
+  const { logout } = useSessionManagement();
 
   const processToken = (token: string): void => {
     try {

libs/react-client/src/auth/user.ts

@@ -0,0 +1,54 @@
+import { useContext, useEffect } from 'react';
+import { useRecoilState } from 'recoil';
+import { ChainlitContext } from 'src/context';
+import { userState } from 'src/state';
+
+import { useAuthConfig } from './config';
+import { getToken, useTokenManagement } from './token';
+
+export const useUser = () => {
+  const apiClient = useContext(ChainlitContext);
+  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 succesful
+    try {
+      const apiUser = await apiClient.getUser();
+      if (apiUser) {
+        setUser(apiUser);
+      }
+    } catch (e) {
+      return;
+    }
+  };
+
+  const initUser = () => {
+    // Already logged in
+    if (user) return;
+
+    // Legacy fallback
+    if (!cookieAuth) return setUserFromLocalStore();
+
+    // Request user from API
+    setUserFromAPI();
+  };
+
+  // Attempt to initialise user on app start.
+  useEffect(initUser, [cookieAuth]);
+
+  return {
+    user,
+    setUserFromAPI
+  };
+};