feat: fix spinner after logout (#5572)

webui/react/src/components/Router.tsx

@@ -18,10 +18,7 @@ const Router: React.FC<Props> = (props: Props) => {
     Loaded: (auth) => auth.isAuthenticated,
     NotLoaded: () => false,
   });
-  const authChecked = Loadable.match(loadableAuth.auth, {
-    Loaded: (auth) => auth.checked,
-    NotLoaded: () => false,
-  });
+  const authChecked = loadableAuth.authChecked;
   const [canceler] = useState(new AbortController());
   const { actions: uiActions } = useUI();
   const location = useLocation();

webui/react/src/hooks/useModal/JupyterLab/useModalJupyterLab.test.tsx

@@ -49,11 +49,12 @@ jest.mock('components/MonacoEditor', () => ({
 }));
 
 const ModalTrigger: React.FC = () => {
-  const { setAuth } = useAuth();
+  const { setAuth, setAuthCheck } = useAuth();
   const { contextHolder, modalOpen } = useModalJupyterLab();
 
   useEffect(() => {
     setAuth({ isAuthenticated: true });
+    setAuthCheck();
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
 

webui/react/src/hooks/useSettings.test.tsx

@@ -106,10 +106,11 @@ const extraConfig: hook.SettingsConfig<ExtraSettings> = {
 };
 
 const Container: React.FC<{ children: JSX.Element }> = ({ children }) => {
-  const { setAuth } = useAuth();
+  const { setAuth, setAuthCheck } = useAuth();
 
   useEffect(() => {
     setAuth({ isAuthenticated: true });
+    setAuthCheck();
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
 

webui/react/src/hooks/useSettingsProvider.tsx

@@ -38,10 +38,7 @@ export const SettingsProvider: React.FC<React.PropsWithChildren> = ({ children }
     Loaded: (auth) => auth.user,
     NotLoaded: () => undefined,
   });
-  const checked = Loadable.match(loadableAuth.auth, {
-    Loaded: (auth) => auth.checked,
-    NotLoaded: () => false,
-  });
+  const checked = loadableAuth.authChecked;
   const [canceler] = useState(new AbortController());
   const [isLoading, setIsLoading] = useState(true);
   const querySettings = useRef('');

webui/react/src/pages/ExperimentDetails/CodeViewer/CodeViewer.test.tsx

@@ -110,10 +110,11 @@ const experimentIdMock = 123;
 const user = userEvent.setup();
 
 const Container: React.FC<Props> = (props) => {
-  const { setAuth } = useAuth();
+  const { setAuth, setAuthCheck } = useAuth();
 
   useEffect(() => {
     setAuth({ isAuthenticated: true });
+    setAuthCheck();
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
 

webui/react/src/pages/Settings/UserManagement.test.tsx

@@ -69,13 +69,14 @@ const Container: React.FC = () => {
   const { updateCurrentUser } = useCurrentUsers();
   const [canceler] = useState(new AbortController());
   const fetchUsers = useFetchUsers(canceler);
-  const { setAuth } = useAuth();
+  const { setAuth, setAuthCheck } = useAuth();
 
   const loadUsers = useCallback(async () => {
     await fetchUsers();
     setAuth({ isAuthenticated: true });
+    setAuthCheck();
     updateCurrentUser(currentUser);
-  }, [fetchUsers, updateCurrentUser, setAuth]);
+  }, [fetchUsers, setAuthCheck, updateCurrentUser, setAuth]);
 
   useEffect(() => {
     loadUsers();

webui/react/src/pages/SignIn.tsx

@@ -42,10 +42,7 @@ const SignIn: React.FC = () => {
   const { actions: uiActions } = useUI();
   const location = useLocation();
   const loadableAuth = useAuth();
-  const authChecked = Loadable.match(loadableAuth.auth, {
-    Loaded: (auth) => auth.checked,
-    NotLoaded: () => false,
-  });
+  const authChecked = loadableAuth.authChecked;
   const isAuthenticated = Loadable.match(loadableAuth.auth, {
     Loaded: (auth) => auth.isAuthenticated,
     NotLoaded: () => false,

webui/react/src/stores/auth.tsx

@@ -5,15 +5,18 @@ import { Auth } from 'types';
 import { getCookie, setCookie } from 'utils/browser';
 import { Loadable, Loaded, NotLoaded } from 'utils/loadable';
 
-export type CurrentUser = Auth & { checked: boolean };
+export type CurrentUser = Auth;
 
 type AuthContext = {
   auth: Loadable<CurrentUser>;
+  authChecked: boolean;
   updateAuth: (fn: (users: Loadable<CurrentUser>) => Loadable<CurrentUser>) => void;
+  updateAuthChecked: (fn: (aChecked: boolean) => boolean) => void;
 };
 
 type UseAuthReturn = {
   auth: Loadable<CurrentUser>;
+  authChecked: boolean;
   resetAuth: () => void;
   setAuth: (auth: Auth) => void;
   setAuthCheck: () => void;
@@ -25,12 +28,15 @@ const AuthContext = createContext<AuthContext | null>(null);
 
 export const AuthProvider: React.FC<PropsWithChildren> = ({ children }) => {
   const [auth, setAuth] = useState<Loadable<CurrentUser>>(NotLoaded);
+  const [authChecked, setAuthChecked] = useState(false);
 
   return (
     <AuthContext.Provider
       value={{
         auth,
+        authChecked,
         updateAuth: setAuth,
+        updateAuthChecked: setAuthChecked,
       }}>
       {children}
     </AuthContext.Provider>
@@ -55,14 +61,15 @@ export const useAuth = (): UseAuthReturn => {
   if (context === null) {
     throw new Error('Attempted to use useAuth outside of Auth Context');
   }
-  const { auth, updateAuth } = context;
+  const { auth, authChecked, updateAuth, updateAuthChecked } = context;
 
   const resetAuth = useCallback(() => {
     clearAuthCookie();
     globalStorage.removeAuthToken();
 
     updateAuth(() => NotLoaded);
-  }, [updateAuth]);
+    updateAuthChecked(() => false);
+  }, [updateAuth, updateAuthChecked]);
 
   const setAuth = useCallback(
     (auth: Auth) => {
@@ -82,19 +89,14 @@ export const useAuth = (): UseAuthReturn => {
   );
 
   const setAuthCheck = useCallback(() => {
-    updateAuth((prevState) => {
-      if (!Loadable.isLoaded(prevState)) return prevState;
-
-      const auth = prevState.data;
-
-      if (auth.checked) return prevState;
-
-      return Loaded({ ...auth, checked: true });
+    updateAuthChecked(() => {
+      return true;
     });
-  }, [updateAuth]);
+  }, [updateAuthChecked]);
 
   return {
     auth,
+    authChecked,
     resetAuth,
     setAuth,
     setAuthCheck,