WIP this is insane.

Chainlit · Nov 25, 2024 · 83a4474 · 83a4474
1 parent b6eff68
commit 83a4474
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 34 deletions.
diff --git a/libs/react-client/src/api/hooks/api.ts b/libs/react-client/src/api/hooks/api.ts
@@ -8,9 +8,10 @@ import useSWR, { SWRConfiguration } from 'swr';
 const fetcher = async (
   client: ChainlitAPI,
   endpoint: string,
-  token?: string
+  token?: string,
+  expect401?: boolean
 ) => {
-  const res = await client.get(endpoint, token);
+  const res = await client.get(endpoint, token, expect401);
   return res?.json();
 };
 
@@ -40,7 +41,11 @@ const fetcher = async (
  */
 function useApi<T>(
   path?: string | null,
-  { token, ...swrConfig }: SWRConfiguration & { token?: string } = {}
+  {
+    token,
+    expect401,
+    ...swrConfig
+  }: SWRConfiguration & { token?: string; expect401?: boolean } = {}
 ) {
   const client = useContext(ChainlitContext);
   let accessToken = useRecoilValue(accessTokenState);
@@ -50,7 +55,7 @@ function useApi<T>(
   const memoizedFetcher = useMemo(
     () =>
       ([url, token]: [url: string, token: string]) =>
-        fetcher(client, url, token),
+        fetcher(client, url, token, expect401),
     [client]
   );
 

diff --git a/libs/react-client/src/api/hooks/auth/index.ts b/libs/react-client/src/api/hooks/auth/index.ts
@@ -12,7 +12,6 @@ export const useAuth = (): IUseAuth => {
   const { logout } = useSessionManagement();
   const { user, setUserFromAPI } = useUser();
   const [accessToken] = useRecoilState(accessTokenState);
-
   const { handleSetAccessToken } = useTokenManagement();
 
   const isReady = !!(!isLoading && authConfig);

diff --git a/libs/react-client/src/api/index.tsx b/libs/react-client/src/api/index.tsx
@@ -1,5 +1,5 @@
 import { ensureTokenPrefix, removeToken } from 'src/auth/token';
-import { IThread, IUser } from 'src/types';
+import { IThread } from 'src/types';
 
 import { IFeedback } from 'src/types/feedback';
 
@@ -81,7 +81,8 @@ export class APIBase {
     path: string,
     token?: string,
     data?: Payload,
-    signal?: AbortSignal
+    signal?: AbortSignal,
+    expect401?: boolean
   ): Promise<Response> {
     try {
       const headers: { Authorization?: string; 'Content-Type'?: string } = {};
@@ -105,7 +106,7 @@ export class APIBase {
 
       if (!res.ok) {
         const body = await res.json();
-        if (res.status === 401 && this.on401) {
+        if (res.status === 401 && !expect401 && this.on401) {
           removeToken();
           this.on401();
         }
@@ -122,8 +123,8 @@ export class APIBase {
     }
   }
 
-  async get(endpoint: string, token?: string) {
-    return await this.fetch('GET', endpoint, token);
+  async get(endpoint: string, token?: string, expect401?: boolean) {
+    return await this.fetch('GET', endpoint, token, { expect401 });
   }
 
   async post(
@@ -159,11 +160,6 @@ 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();

diff --git a/libs/react-client/src/auth/user.ts b/libs/react-client/src/auth/user.ts
@@ -6,45 +6,57 @@ import { IUser, useApi } from '..';
 import { useAuthConfig } from './config';
 import { getToken, useTokenManagement } from './token';
 
+// Crappy logic due to lack of something like an `<AuthProvider/>` component.
+let isInitialized = false;
+
 export const useUser = () => {
   console.log('useUser');
 
   const [user, setUser] = useRecoilState(userState);
   const { cookieAuth } = useAuthConfig();
   const { handleSetAccessToken } = useTokenManagement();
 
-  const { data: userData, mutate: mutateUserData } = useApi<IUser>('/user', {
-    revalidateOnMount: false
-  });
+  const { data: userData, mutate: mutateUserData } = useApi<IUser>(
+    cookieAuth ? '/user' : null,
+    {
+      expect401: true
+      // onErrorRetry: (error, key, config, revalidate, { retryCount }) => {
+      //   console.log('onError');
+      //   // 401 is a normal response until we're logged in.
+      //   if (error.status === 401) {
+      //     // Make sure we remove user though
+      //     setUser(null);
+      //     return;
+      //   }
+
+      //   // On other errors, retry after 5 seconds.
+      //   setTimeout(() => revalidate({ retryCount }), config.errorRetryInterval)
+      // }
+    }
+  );
 
-  // Attempt to get user when cookieAuth are available.
+  // Not using cookie auth, callback to header tokens
   useEffect(() => {
-    if (!user) {
-      if (cookieAuth) {
-        console.log('cookieAuth', user, cookieAuth);
-        mutateUserData();
-        return;
-      }
-
-      // Not using cookie auth, callback to header tokens
-      console.log('tokenAuth', user, cookieAuth);
+    if (!isInitialized && !user && !cookieAuth) {
       const token = getToken();
       if (token) handleSetAccessToken(token);
     }
-  }, [user, cookieAuth]);
+  }, [cookieAuth]);
 
-  // When user data is available, set the user object.
+  // When user data is available from the API, set the user object.
   useEffect(() => {
-    console.log('userData effect');
-
-    if (userData) {
+    if (!isInitialized && userData) {
       console.log('setUser', userData);
       setUser(userData);
     }
   }, [userData]);
 
+  // Make sure effects are only executed once.
   useEffect(() => {
-    console.log('useUser effect');
+    if (!isInitialized) {
+      isInitialized = true;
+      console.log('isInitialized');
+    }
   });
 
   return {