Fix(ui): refresh call concurrency for multiple browser tabs (#19303)

* fix(ui): refresh auth token for multi browser tabs

* update refresh logic

* fix multiple tab issue

* fix tests

* added tests

* fix ingestion bot failure

* fix sonar cloud

* update test description

* remove unused code and reset test on after all

* bump playwright

* avoid running refresh tests as it's been flaky for postgres

* revert playwright version bump changes

* Put 500 status

---------

Co-authored-by: mohitdeuex <mohit.y@deuexsolutions.com>

openmetadata-service/src/main/java/org/openmetadata/service/security/AuthenticationCodeFlowHandler.java

@@ -642,6 +642,7 @@ public class AuthenticationCodeFlowHandler {
 
   @SneakyThrows
   public static void getErrorMessage(HttpServletResponse resp, Exception e) {
+    resp.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
     resp.setContentType("text/html; charset=UTF-8");
     LOG.error("[Auth Callback Servlet] Failed in Auth Login : {}", e.getMessage());
     resp.getOutputStream()

openmetadata-ui/src/main/resources/ui/playwright/constant/login.ts

@@ -11,6 +11,7 @@
  *  limitations under the License.
  */
 export const JWT_EXPIRY_TIME_MAP = {
+  '3 minutes': 180,
   '1 hour': 3600,
   '2 hours': 7200,
   '3 hours': 10800,

openmetadata-ui/src/main/resources/ui/playwright/e2e/Flow/IngestionBot.spec.ts

@@ -57,7 +57,7 @@ const test = base.extend<{
       // Set a new value for a key in localStorage
       localStorage.setItem(
         'om-session',
-        JSON.stringify({ state: { oidcIdToken: token } })
+        JSON.stringify({ oidcIdToken: token })
       );
     }, tokenData.config.JWTToken);
 

openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/Login.spec.ts

@@ -11,26 +11,53 @@
  *  limitations under the License.
  */
 import { expect, test } from '@playwright/test';
-import { LOGIN_ERROR_MESSAGE } from '../../constant/login';
+import { JWT_EXPIRY_TIME_MAP, LOGIN_ERROR_MESSAGE } from '../../constant/login';
 import { UserClass } from '../../support/user/UserClass';
 import { performAdminLogin } from '../../utils/admin';
+import { redirectToHomePage } from '../../utils/common';
+import { updateJWTTokenExpiryTime } from '../../utils/login';
+import { visitUserProfilePage } from '../../utils/user';
 
 const user = new UserClass();
 const CREDENTIALS = user.data;
 const invalidEmail = 'userTest@openmetadata.org';
 const invalidPassword = 'testUsers@123';
 
+test.describe.configure({
+  // 5 minutes max for refresh token tests
+  timeout: 5 * 60 * 1000,
+});
+
 test.describe('Login flow should work properly', () => {
   test.afterAll('Cleanup', async ({ browser }) => {
     const { apiContext, afterAction, page } = await performAdminLogin(browser);
     const response = await page.request.get(
       `/api/v1/users/name/${user.getUserName()}`
     );
+
+    // reset token expiry to 4 hours
+    await updateJWTTokenExpiryTime(apiContext, JWT_EXPIRY_TIME_MAP['4 hours']);
+
     user.responseData = await response.json();
     await user.delete(apiContext);
     await afterAction();
   });
 
+  test.beforeAll(
+    'Update token timer to be 3 minutes for new token created',
+    async ({ browser }) => {
+      const { apiContext, afterAction } = await performAdminLogin(browser);
+
+      // update expiry for 3 mins
+      await updateJWTTokenExpiryTime(
+        apiContext,
+        JWT_EXPIRY_TIME_MAP['3 minutes']
+      );
+
+      await afterAction();
+    }
+  );
+
   test('Signup and Login with signed up credentials', async ({ page }) => {
     await page.goto('/');
 
@@ -111,4 +138,36 @@ test.describe('Login flow should work properly', () => {
     await page.getByRole('button', { name: 'Submit' }).click();
     await page.locator('[data-testid="go-back-button"]').click();
   });
+
+  test.fixme('Refresh should work', async ({ browser }) => {
+    const browserContext = await browser.newContext();
+    const { apiContext, afterAction } = await performAdminLogin(browser);
+    const page1 = await browserContext.newPage(),
+      page2 = await browserContext.newPage();
+
+    const testUser = new UserClass();
+    await testUser.create(apiContext);
+
+    await afterAction();
+
+    await test.step('Login and wait for refresh call is made', async () => {
+      // User login
+
+      await testUser.login(page1);
+      await redirectToHomePage(page1);
+      await redirectToHomePage(page2);
+
+      const refreshCall = page1.waitForResponse('**/refresh', {
+        timeout: 3 * 60 * 1000,
+      });
+
+      await refreshCall;
+
+      await redirectToHomePage(page1);
+
+      await visitUserProfilePage(page1, testUser.responseData.name);
+      await redirectToHomePage(page2);
+      await visitUserProfilePage(page2, testUser.responseData.name);
+    });
+  });
 });

openmetadata-ui/src/main/resources/ui/playwright/utils/common.ts

@@ -41,8 +41,7 @@ export const NAME_MAX_LENGTH_VALIDATION_ERROR =
 export const getToken = async (page: Page) => {
   return page.evaluate(
     () =>
-      JSON.parse(localStorage.getItem('om-session') ?? '{}')?.state
+      JSON.parse(localStorage.getItem('om-session') ?? '{}')?.oidcIdToken ?? ''
-        ?.oidcIdToken ?? ''
   );
 };
 

openmetadata-ui/src/main/resources/ui/src/components/AppBar/Appbar.tsx

@@ -26,6 +26,7 @@ import {
   isTourRoute,
 } from '../../utils/AuthProvider.util';
 import { addToRecentSearched } from '../../utils/CommonUtils';
+import { getOidcToken } from '../../utils/LocalStorageUtils';
 import searchClassBase from '../../utils/SearchClassBase';
 import NavBar from '../NavBar/NavBar';
 import './app-bar.style.less';
@@ -37,7 +38,7 @@ const Appbar: React.FC = (): JSX.Element => {
   const { isTourOpen, updateTourPage, updateTourSearch, tourSearchValue } =
     useTourProvider();
 
-  const { isAuthenticated, searchCriteria, getOidcToken, trySilentSignIn } =
+  const { isAuthenticated, searchCriteria, trySilentSignIn } =
     useApplicationStore();
 
   const parsedQueryString = Qs.parse(

openmetadata-ui/src/main/resources/ui/src/components/Auth/AppAuthenticators/Auth0Authenticator.tsx

@@ -22,6 +22,7 @@ import { useTranslation } from 'react-i18next';
 import { AuthProvider } from '../../../generated/settings/settings';
 
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
+import { setOidcToken } from '../../../utils/LocalStorageUtils';
 import { AuthenticatorRef } from '../AuthProviders/AuthProvider.interface';
 
 interface Props {
@@ -31,8 +32,7 @@ interface Props {
 
 const Auth0Authenticator = forwardRef<AuthenticatorRef, Props>(
   ({ children, onLogoutSuccess }: Props, ref) => {
-    const { setIsAuthenticated, authConfig, setOidcToken } =
+    const { setIsAuthenticated, authConfig } = useApplicationStore();
-      useApplicationStore();
     const { t } = useTranslation();
     const { loginWithRedirect, getAccessTokenSilently, getIdTokenClaims } =
       useAuth0();

openmetadata-ui/src/main/resources/ui/src/components/Auth/AppAuthenticators/BasicAuthAuthenticator.tsx

@@ -26,6 +26,11 @@ import {
 } from '../../../rest/auth-API';
 
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
+import {
+  getRefreshToken,
+  setOidcToken,
+  setRefreshToken,
+} from '../../../utils/LocalStorageUtils';
 import Loader from '../../common/Loader/Loader';
 import { useBasicAuth } from '../AuthProviders/BasicAuthProvider';
 
@@ -40,9 +45,7 @@ const BasicAuthenticator = forwardRef(
     const {
       setIsAuthenticated,
       authConfig,
-      getRefreshToken,
+
-      setRefreshToken,
-      setOidcToken,
       isApplicationLoading,
     } = useApplicationStore();
 
@@ -54,7 +57,13 @@ const BasicAuthenticator = forwardRef(
           authConfig?.provider !== AuthProvider.Basic &&
           authConfig?.provider !== AuthProvider.LDAP
         ) {
-          Promise.reject(t('message.authProvider-is-not-basic'));
+          return Promise.reject(
+            new Error(t('message.authProvider-is-not-basic'))
+          );
+        }
+
+        if (!refreshToken) {
+          return Promise.reject(new Error(t('message.no-token-available')));
         }
 
         const response = await getAccessTokenOnExpiry({

openmetadata-ui/src/main/resources/ui/src/components/Auth/AppAuthenticators/GenericAuthenticator.tsx

@@ -20,15 +20,11 @@ import { useHistory } from 'react-router-dom';
 import { ROUTES } from '../../../constants/constants';
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
 import { logoutUser, renewToken } from '../../../rest/LoginAPI';
+import { setOidcToken } from '../../../utils/LocalStorageUtils';
 
 export const GenericAuthenticator = forwardRef(
   ({ children }: { children: ReactNode }, ref) => {
-    const {
+    const { setIsAuthenticated, setIsSigningUp } = useApplicationStore();
-      setIsAuthenticated,
-      setIsSigningUp,
-      removeOidcToken,
-      setOidcToken,
-    } = useApplicationStore();
     const history = useHistory();
 
     const handleLogin = () => {
@@ -42,7 +38,7 @@ export const GenericAuthenticator = forwardRef(
       await logoutUser();
 
       history.push(ROUTES.SIGNIN);
-      removeOidcToken();
+      setOidcToken('');
       setIsAuthenticated(false);
     };
 

openmetadata-ui/src/main/resources/ui/src/components/Auth/AppAuthenticators/OidcAuthenticator.tsx

@@ -27,6 +27,8 @@ import { ROUTES } from '../../../constants/constants';
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
 import useCustomLocation from '../../../hooks/useCustomLocation/useCustomLocation';
 import SignInPage from '../../../pages/LoginPage/SignInPage';
+import TokenService from '../../../utils/Auth/TokenService/TokenServiceUtil';
+import { setOidcToken } from '../../../utils/LocalStorageUtils';
 import { showErrorToast } from '../../../utils/ToastUtils';
 import Loader from '../../common/Loader/Loader';
 import {
@@ -71,7 +73,6 @@ const OidcAuthenticator = forwardRef<AuthenticatorRef, Props>(
       updateAxiosInterceptors,
       currentUser,
       newUser,
-      setOidcToken,
       isApplicationLoading,
     } = useApplicationStore();
     const history = useHistory();
@@ -105,6 +106,9 @@ const OidcAuthenticator = forwardRef<AuthenticatorRef, Props>(
       // On success update token in store and update axios interceptors
       setOidcToken(user.id_token);
       updateAxiosInterceptors();
+      // Clear the refresh token in progress flag
+      // Since refresh token request completes with a callback
+      TokenService.getInstance().clearRefreshInProgress();
     };
 
     const handleSilentSignInFailure = (error: unknown) => {

openmetadata-ui/src/main/resources/ui/src/components/Auth/AppAuthenticators/OktaAuthenticator.tsx

@@ -20,6 +20,7 @@ import React, {
 } from 'react';
 
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
+import { setOidcToken } from '../../../utils/LocalStorageUtils';
 import { AuthenticatorRef } from '../AuthProviders/AuthProvider.interface';
 
 interface Props {
@@ -30,7 +31,7 @@ interface Props {
 const OktaAuthenticator = forwardRef<AuthenticatorRef, Props>(
   ({ children, onLogoutSuccess }: Props, ref) => {
     const { oktaAuth } = useOktaAuth();
-    const { setIsAuthenticated, setOidcToken } = useApplicationStore();
+    const { setIsAuthenticated } = useApplicationStore();
 
     const login = async () => {
       oktaAuth.signInWithRedirect();

openmetadata-ui/src/main/resources/ui/src/components/Auth/AppAuthenticators/SamlAuthenticator.tsx

@@ -36,6 +36,12 @@ import { showErrorToast } from '../../../utils/ToastUtils';
 import { ROUTES } from '../../../constants/constants';
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
 import { AccessTokenResponse, refreshSAMLToken } from '../../../rest/auth-API';
+import {
+  getOidcToken,
+  getRefreshToken,
+  setOidcToken,
+  setRefreshToken,
+} from '../../../utils/LocalStorageUtils';
 import { AuthenticatorRef } from '../AuthProviders/AuthProvider.interface';
 
 interface Props {
@@ -45,14 +51,7 @@ interface Props {
 
 const SamlAuthenticator = forwardRef<AuthenticatorRef, Props>(
   ({ children, onLogoutSuccess }: Props, ref) => {
-    const {
+    const { setIsAuthenticated, authConfig } = useApplicationStore();
-      setIsAuthenticated,
-      authConfig,
-      getOidcToken,
-      getRefreshToken,
-      setRefreshToken,
-      setOidcToken,
-    } = useApplicationStore();
     const config = authConfig?.samlConfiguration as SamlSSOClientConfig;
 
     const handleSilentSignIn = async (): Promise<AccessTokenResponse> => {

openmetadata-ui/src/main/resources/ui/src/components/Auth/AppCallbacks/Auth0Callback/Auth0Callback.test.tsx

@@ -53,11 +53,14 @@ jest.mock('../../../../hooks/useApplicationStore', () => {
     useApplicationStore: jest.fn(() => ({
       authConfig: {},
       handleSuccessfulLogin: mockHandleSuccessfulLogin,
-      setOidcToken: jest.fn(),
     })),
   };
 });
 
+jest.mock('../../../../utils/LocalStorageUtils', () => ({
+  setOidcToken: jest.fn(),
+}));
+
 describe('Test Auth0Callback component', () => {
   afterEach(() => {
     jest.clearAllMocks();

openmetadata-ui/src/main/resources/ui/src/components/Auth/AppCallbacks/Auth0Callback/Auth0Callback.tsx

@@ -16,11 +16,12 @@ import { t } from 'i18next';
 import React, { VFC } from 'react';
 
 import { useApplicationStore } from '../../../../hooks/useApplicationStore';
+import { setOidcToken } from '../../../../utils/LocalStorageUtils';
 import { OidcUser } from '../../AuthProviders/AuthProvider.interface';
 
 const Auth0Callback: VFC = () => {
   const { isAuthenticated, user, getIdTokenClaims, error } = useAuth0();
-  const { handleSuccessfulLogin, setOidcToken } = useApplicationStore();
+  const { handleSuccessfulLogin } = useApplicationStore();
   if (isAuthenticated) {
     getIdTokenClaims()
       .then((token) => {

openmetadata-ui/src/main/resources/ui/src/components/Auth/AuthProviders/AuthProvider.tsx

@@ -25,7 +25,7 @@ import {
   InternalAxiosRequestConfig,
 } from 'axios';
 import { CookieStorage } from 'cookie-storage';
-import { debounce, isEmpty, isNil, isNumber } from 'lodash';
+import { isEmpty, isNil, isNumber } from 'lodash';
 import Qs from 'qs';
 import React, {
   ComponentType,
@@ -72,7 +72,12 @@ import {
   isProtectedRoute,
   prepareUserProfileFromClaims,
 } from '../../../utils/AuthProvider.util';
-import { getOidcToken } from '../../../utils/LocalStorageUtils';
+import {
+  getOidcToken,
+  getRefreshToken,
+  setOidcToken,
+  setRefreshToken,
+} from '../../../utils/LocalStorageUtils';
 import { getPathNameFromWindowLocation } from '../../../utils/RouterUtils';
 import { escapeESReservedCharacters } from '../../../utils/StringsUtils';
 import { showErrorToast, showInfoToast } from '../../../utils/ToastUtils';
@@ -110,7 +115,9 @@ const isEmailVerifyField = 'isEmailVerified';
 
 let requestInterceptor: number | null = null;
 let responseInterceptor: number | null = null;
-let failedLoggedInUserRequest: boolean | null;
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+let pendingRequests: any[] = [];
 
 export const AuthProvider = ({
   childComponentType,
@@ -130,14 +137,11 @@ export const AuthProvider = ({
     jwtPrincipalClaimsMapping,
     setJwtPrincipalClaims,
     setJwtPrincipalClaimsMapping,
-    removeRefreshToken,
-    removeOidcToken,
-    getRefreshToken,
     isApplicationLoading,
     setApplicationLoading,
   } = useApplicationStore();
   const { updateDomains, updateDomainLoading } = useDomainStore();
-  const tokenService = useRef<TokenService>();
+  const tokenService = useRef<TokenService>(TokenService.getInstance());
 
   const location = useCustomLocation();
   const history = useHistory();
@@ -176,7 +180,7 @@ export const AuthProvider = ({
     removeSession();
 
     // remove the refresh token on logout
-    removeRefreshToken();
+    setRefreshToken('');
 
     setApplicationLoading(false);
 
@@ -184,14 +188,6 @@ export const AuthProvider = ({
     history.push(ROUTES.SIGNIN);
   }, [timeoutId]);
 
-  useEffect(() => {
-    if (authenticatorRef.current?.renewIdToken) {
-      tokenService.current = new TokenService(
-        authenticatorRef.current?.renewIdToken
-      );
-    }
-  }, [authenticatorRef.current?.renewIdToken]);
-
   const fetchDomainList = useCallback(async () => {
     try {
       updateDomainLoading(true);
@@ -228,7 +224,7 @@ export const AuthProvider = ({
 
   const resetUserDetails = (forceLogout = false) => {
     setCurrentUser({} as User);
-    removeOidcToken();
+    setOidcToken('');
     setIsAuthenticated(false);
     setApplicationLoading(false);
     clearTimeout(timeoutId);
@@ -268,39 +264,6 @@ export const AuthProvider = ({
     }
   };
 
-  /**
-   * This method will try to signIn silently when token is about to expire
-   * if it's not succeed then it will proceed for logout
-   */
-  const trySilentSignIn = async (forceLogout?: boolean) => {
-    const pathName = getPathNameFromWindowLocation();
-    // Do not try silent sign in for SignIn or SignUp route
-    if (
-      [ROUTES.SIGNIN, ROUTES.SIGNUP, ROUTES.SILENT_CALLBACK].includes(pathName)
-    ) {
-      return;
-    }
-
-    if (!tokenService.current?.isTokenUpdateInProgress()) {
-      // For OIDC we won't be getting newToken immediately hence not updating token here
-      const newToken = await tokenService.current?.refreshToken();
-      // Start expiry timer on successful silent signIn
-      if (newToken) {
-        // Start expiry timer on successful silent signIn
-        // eslint-disable-next-line @typescript-eslint/no-use-before-define
-        startTokenExpiryTimer();
-
-        // Retry the failed request after successful silent signIn
-        if (failedLoggedInUserRequest) {
-          await getLoggedInUserDetails();
-          failedLoggedInUserRequest = null;
-        }
-      } else if (forceLogout) {
-        resetUserDetails(true);
-      }
-    }
-  };
-
   /**
    * It will set an timer for 5 mins before Token will expire
    * If time if less then 5 mins then it will try to SilentSignIn
@@ -326,13 +289,25 @@ export const AuthProvider = ({
       // If token is about to expire then start silentSignIn
       // else just set timer to try for silentSignIn before token expires
       clearTimeout(timeoutId);
-      const timerId = setTimeout(() => {
+
-        trySilentSignIn();
+      const timerId = setTimeout(
-      }, timeoutExpiry);
+        tokenService.current?.refreshToken,
+        timeoutExpiry
+      );
       setTimeoutId(Number(timerId));
     }
   };
 
+  useEffect(() => {
+    if (authenticatorRef.current?.renewIdToken) {
+      tokenService.current.updateRenewToken(
+        authenticatorRef.current?.renewIdToken
+      );
+      // After every refresh success, start timer again
+      tokenService.current.updateRefreshSuccessCallback(startTokenExpiryTimer);
+    }
+  }, [authenticatorRef.current?.renewIdToken]);
+
   /**
    * Performs cleanup around timers
    * Clean silentSignIn activities if going on
@@ -542,13 +517,48 @@ export const AuthProvider = ({
         if (error.response) {
           const { status } = error.response;
           if (status === ClientErrors.UNAUTHORIZED) {
-            // store the failed request for retry after successful silent signIn
+            if (error.config.url === '/users/refresh') {
-            if (error.config.url === '/users/loggedInUser') {
+              return Promise.reject(error as Error);
-              failedLoggedInUserRequest = true;
             }
             handleStoreProtectedRedirectPath();
-            // try silent signIn if token is about to expire
+
-            debounce(() => trySilentSignIn(true), 100);
+            // If 401 error and refresh is not in progress, trigger the refresh
+            if (!tokenService.current?.isTokenUpdateInProgress()) {
+              // Start the refresh process
+              return new Promise((resolve, reject) => {
+                // Add this request to the pending queue
+                pendingRequests.push({
+                  resolve,
+                  reject,
+                  config: error.config,
+                });
+
+                // Refresh the token and retry the requests in the queue
+                tokenService.current.refreshToken().then((token) => {
+                  if (token) {
+                    // Retry the pending requests
+                    initializeAxiosInterceptors();
+                    pendingRequests.forEach(({ resolve, reject, config }) => {
+                      axiosClient(config).then(resolve).catch(reject);
+                    });
+                  } else {
+                    resetUserDetails(true);
+                  }
+                });
+
+                // Clear the queue after retrying
+                pendingRequests = [];
+              });
+            } else {
+              // If refresh is in progress, queue the request
+              return new Promise((resolve, reject) => {
+                pendingRequests.push({
+                  resolve,
+                  reject,
+                  config: error.config,
+                });
+              });
+            }
           }
         }
 
@@ -721,7 +731,6 @@ export const AuthProvider = ({
       onLoginHandler,
       onLogoutHandler,
       handleSuccessfulLogin,
-      trySilentSignIn,
       handleFailedLogin,
       updateAxiosInterceptors: initializeAxiosInterceptors,
     });
@@ -734,7 +743,6 @@ export const AuthProvider = ({
       onLoginHandler,
       onLogoutHandler,
       handleSuccessfulLogin,
-      trySilentSignIn,
       handleFailedLogin,
       updateAxiosInterceptors: initializeAxiosInterceptors,
     });

openmetadata-ui/src/main/resources/ui/src/components/Auth/AuthProviders/BasicAuthProvider.tsx

@@ -39,7 +39,13 @@ import {
 import { resetWebAnalyticSession } from '../../../utils/WebAnalyticsUtils';
 
 import { toLower } from 'lodash';
-import { useApplicationStore } from '../../../hooks/useApplicationStore';
+import { extractDetailsFromToken } from '../../../utils/AuthProvider.util';
+import {
+  getOidcToken,
+  getRefreshToken,
+  setOidcToken,
+  setRefreshToken,
+} from '../../../utils/LocalStorageUtils';
 import { OidcUser } from './AuthProvider.interface';
 
 interface BasicAuthProps {
@@ -84,13 +90,7 @@ const BasicAuthProvider = ({
   onLoginFailure,
 }: BasicAuthProps) => {
   const { t } = useTranslation();
-  const {
+
-    setRefreshToken,
-    setOidcToken,
-    getOidcToken,
-    removeOidcToken,
-    getRefreshToken,
-  } = useApplicationStore();
   const [loginError, setLoginError] = useState<string | null>(null);
   const history = useHistory();
 
@@ -176,10 +176,11 @@ const BasicAuthProvider = ({
   const handleLogout = async () => {
     const token = getOidcToken();
     const refreshToken = getRefreshToken();
-    if (token) {
+    const isExpired = extractDetailsFromToken(token).isExpired;
+    if (token && !isExpired) {
       try {
         await logoutUser({ token, refreshToken });
-        removeOidcToken();
+        setOidcToken('');
         history.push(ROUTES.SIGNIN);
       } catch (error) {
         showErrorToast(error as AxiosError);

openmetadata-ui/src/main/resources/ui/src/components/Auth/AuthProviders/OktaAuthProvider.tsx

@@ -20,6 +20,7 @@ import React, {
   useMemo,
 } from 'react';
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
+import { setOidcToken } from '../../../utils/LocalStorageUtils';
 import { OidcUser } from './AuthProvider.interface';
 
 interface Props {
@@ -31,7 +32,7 @@ export const OktaAuthProvider: FunctionComponent<Props> = ({
   children,
   onLoginSuccess,
 }: Props) => {
-  const { authConfig, setOidcToken } = useApplicationStore();
+  const { authConfig } = useApplicationStore();
   const { clientId, issuer, redirectUri, scopes, pkce } =
     authConfig as OktaAuthOptions;
 

openmetadata-ui/src/main/resources/ui/src/hooks/useApplicationStore.ts

@@ -11,7 +11,6 @@
  *  limitations under the License.
  */
 import { create } from 'zustand';
-import { persist } from 'zustand/middleware';
 import { AuthenticationConfigurationWithScope } from '../components/Auth/AuthProviders/AuthProvider.interface';
 import { EntityUnion } from '../components/Explore/ExplorePage.interface';
 import { AuthenticationConfiguration } from '../generated/configuration/authenticationConfiguration';
@@ -28,164 +27,133 @@ import { getThemeConfig } from '../utils/ThemeUtils';
 
 export const OM_SESSION_KEY = 'om-session';
 
-export const useApplicationStore = create<ApplicationStore>()(
+export const useApplicationStore = create<ApplicationStore>()((set, get) => ({
-  persist(
+  isApplicationLoading: false,
-    (set, get) => ({
+  theme: getThemeConfig(),
-      isApplicationLoading: false,
+  applicationConfig: {
-      theme: getThemeConfig(),
+    customTheme: getThemeConfig(),
-      applicationConfig: {
+  } as UIThemePreference,
-        customTheme: getThemeConfig(),
+  currentUser: undefined,
-      } as UIThemePreference,
+  newUser: undefined,
-      currentUser: undefined,
+  isAuthenticated: Boolean(getOidcToken()),
-      newUser: undefined,
+  authConfig: undefined,
-      isAuthenticated: Boolean(getOidcToken()),
+  authorizerConfig: undefined,
-      authConfig: undefined,
+  isSigningUp: false,
-      authorizerConfig: undefined,
+  jwtPrincipalClaims: [],
-      isSigningUp: false,
+  jwtPrincipalClaimsMapping: [],
-      jwtPrincipalClaims: [],
+  userProfilePics: {},
-      jwtPrincipalClaimsMapping: [],
+  cachedEntityData: {},
-      userProfilePics: {},
+  selectedPersona: {} as EntityReference,
-      cachedEntityData: {},
+  searchCriteria: '',
-      selectedPersona: {} as EntityReference,
+  inlineAlertDetails: undefined,
-      oidcIdToken: '',
+  applications: [],
-      refreshTokenKey: '',
+  appPreferences: {},
-      searchCriteria: '',
-      inlineAlertDetails: undefined,
-      applications: [],
-      appPreferences: {},
 
-      setInlineAlertDetails: (inlineAlertDetails) => {
+  setInlineAlertDetails: (inlineAlertDetails) => {
-        set({ inlineAlertDetails });
+    set({ inlineAlertDetails });
-      },
+  },
 
-      setHelperFunctionsRef: (helperFunctions: HelperFunctions) => {
+  setHelperFunctionsRef: (helperFunctions: HelperFunctions) => {
-        set({ ...helperFunctions });
+    set({ ...helperFunctions });
-      },
+  },
 
-      setSelectedPersona: (persona: EntityReference) => {
+  setSelectedPersona: (persona: EntityReference) => {
-        set({ selectedPersona: persona });
+    set({ selectedPersona: persona });
-      },
+  },
 
-      setApplicationConfig: (config: UIThemePreference) => {
+  setApplicationConfig: (config: UIThemePreference) => {
-        set({ applicationConfig: config, theme: config.customTheme });
+    set({ applicationConfig: config, theme: config.customTheme });
-      },
+  },
-      setCurrentUser: (user) => {
+  setCurrentUser: (user) => {
-        set({ currentUser: user });
+    set({ currentUser: user });
-      },
+  },
-      setAuthConfig: (authConfig: AuthenticationConfigurationWithScope) => {
+  setAuthConfig: (authConfig: AuthenticationConfigurationWithScope) => {
-        set({ authConfig });
+    set({ authConfig });
-      },
+  },
-      setAuthorizerConfig: (authorizerConfig: AuthorizerConfiguration) => {
+  setAuthorizerConfig: (authorizerConfig: AuthorizerConfiguration) => {
-        set({ authorizerConfig });
+    set({ authorizerConfig });
-      },
+  },
-      setJwtPrincipalClaims: (
+  setJwtPrincipalClaims: (
-        claims: AuthenticationConfiguration['jwtPrincipalClaims']
+    claims: AuthenticationConfiguration['jwtPrincipalClaims']
-      ) => {
+  ) => {
-        set({ jwtPrincipalClaims: claims });
+    set({ jwtPrincipalClaims: claims });
-      },
+  },
-      setJwtPrincipalClaimsMapping: (
+  setJwtPrincipalClaimsMapping: (
-        claimMapping: AuthenticationConfiguration['jwtPrincipalClaimsMapping']
+    claimMapping: AuthenticationConfiguration['jwtPrincipalClaimsMapping']
-      ) => {
+  ) => {
-        set({ jwtPrincipalClaimsMapping: claimMapping });
+    set({ jwtPrincipalClaimsMapping: claimMapping });
-      },
+  },
-      setIsAuthenticated: (authenticated: boolean) => {
+  setIsAuthenticated: (authenticated: boolean) => {
-        set({ isAuthenticated: authenticated });
+    set({ isAuthenticated: authenticated });
-      },
+  },
-      setIsSigningUp: (signingUp: boolean) => {
+  setIsSigningUp: (signingUp: boolean) => {
-        set({ isSigningUp: signingUp });
+    set({ isSigningUp: signingUp });
-      },
+  },
 
-      setApplicationLoading: (loading: boolean) => {
+  setApplicationLoading: (loading: boolean) => {
-        set({ isApplicationLoading: loading });
+    set({ isApplicationLoading: loading });
-      },
+  },
 
-      onLoginHandler: () => {
+  onLoginHandler: () => {
-        // This is a placeholder function that will be replaced by the actual function
+    // This is a placeholder function that will be replaced by the actual function
-      },
+  },
-      onLogoutHandler: () => {
+  onLogoutHandler: () => {
-        // This is a placeholder function that will be replaced by the actual function
+    // This is a placeholder function that will be replaced by the actual function
-      },
+  },
 
-      handleSuccessfulLogin: () => {
+  handleSuccessfulLogin: () => {
-        // This is a placeholder function that will be replaced by the actual function
+    // This is a placeholder function that will be replaced by the actual function
-      },
+  },
-      handleFailedLogin: () => {
+  handleFailedLogin: () => {
-        // This is a placeholder function that will be replaced by the actual function
+    // This is a placeholder function that will be replaced by the actual function
-      },
+  },
-      updateAxiosInterceptors: () => {
+  updateAxiosInterceptors: () => {
-        // This is a placeholder function that will be replaced by the actual function
+    // This is a placeholder function that will be replaced by the actual function
-      },
+  },
-      trySilentSignIn: (forceLogout?: boolean) => {
+  trySilentSignIn: (forceLogout?: boolean) => {
-        if (forceLogout) {
+    if (forceLogout) {
-          // This is a placeholder function that will be replaced by the actual function
+      // This is a placeholder function that will be replaced by the actual function
-        }
-      },
-      updateCurrentUser: (user) => {
-        set({ currentUser: user });
-      },
-      updateUserProfilePics: ({ id, user }: { id: string; user: User }) => {
-        set({
-          userProfilePics: { ...get()?.userProfilePics, [id]: user },
-        });
-      },
-      updateCachedEntityData: ({
-        id,
-        entityDetails,
-      }: {
-        id: string;
-        entityDetails: EntityUnion;
-      }) => {
-        set({
-          cachedEntityData: {
-            ...get()?.cachedEntityData,
-            [id]: entityDetails,
-          },
-        });
-      },
-      updateNewUser: (user) => {
-        set({ newUser: user });
-      },
-      getRefreshToken: () => {
-        return get()?.refreshTokenKey;
-      },
-      setRefreshToken: (refreshToken) => {
-        set({ refreshTokenKey: refreshToken });
-      },
-      setAppPreferences: (
-        preferences: Partial<ApplicationStore['appPreferences']>
-      ) => {
-        set((state) => ({
-          appPreferences: {
-            ...state.appPreferences,
-            ...preferences,
-          },
-        }));
-      },
-      getOidcToken: () => {
-        return get()?.oidcIdToken;
-      },
-      setOidcToken: (oidcToken) => {
-        set({ oidcIdToken: oidcToken });
-      },
-      removeOidcToken: () => {
-        set({ oidcIdToken: '' });
-      },
-      removeRefreshToken: () => {
-        set({ refreshTokenKey: '' });
-      },
-      updateSearchCriteria: (criteria) => {
-        set({ searchCriteria: criteria });
-      },
-      setApplicationsName: (applications: string[]) => {
-        set({ applications: applications });
-      },
-    }),
-    {
-      name: OM_SESSION_KEY, // name of item in the storage (must be unique)
-      partialize: (state) => ({
-        oidcIdToken: state.oidcIdToken,
-        refreshTokenKey: state.refreshTokenKey,
-      }),
     }
-  )
+  },
-);
+  updateCurrentUser: (user) => {
+    set({ currentUser: user });
+  },
+  updateUserProfilePics: ({ id, user }: { id: string; user: User }) => {
+    set({
+      userProfilePics: { ...get()?.userProfilePics, [id]: user },
+    });
+  },
+  updateCachedEntityData: ({
+    id,
+    entityDetails,
+  }: {
+    id: string;
+    entityDetails: EntityUnion;
+  }) => {
+    set({
+      cachedEntityData: {
+        ...get()?.cachedEntityData,
+        [id]: entityDetails,
+      },
+    });
+  },
+  updateNewUser: (user) => {
+    set({ newUser: user });
+  },
+  setAppPreferences: (
+    preferences: Partial<ApplicationStore['appPreferences']>
+  ) => {
+    set((state) => ({
+      appPreferences: {
+        ...state.appPreferences,
+        ...preferences,
+      },
+    }));
+  },
+  updateSearchCriteria: (criteria) => {
+    set({ searchCriteria: criteria });
+  },
+  setApplicationsName: (applications: string[]) => {
+    set({ applications: applications });
+  },
+}));

openmetadata-ui/src/main/resources/ui/src/interface/store.interface.ts

@@ -37,7 +37,6 @@ export interface HelperFunctions {
   handleSuccessfulLogin: (user: OidcUser) => Promise<void>;
   handleFailedLogin: () => void;
   updateAxiosInterceptors: () => void;
-  trySilentSignIn: (forceLogout?: boolean) => Promise<void>;
 }
 
 export interface AppPreferences {
@@ -53,8 +52,6 @@ export interface ApplicationStore
   userProfilePics: Record<string, User>;
   cachedEntityData: Record<string, EntityUnion>;
   selectedPersona: EntityReference;
-  oidcIdToken: string;
-  refreshTokenKey: string;
   authConfig?: AuthenticationConfigurationWithScope;
   applicationConfig?: UIThemePreference;
   searchCriteria: ExploreSearchIndex | '';
@@ -81,13 +78,6 @@ export interface ApplicationStore
     id: string;
     entityDetails: EntityUnion;
   }) => void;
-
-  getRefreshToken: () => string;
-  setRefreshToken: (refreshToken: string) => void;
-  getOidcToken: () => string;
-  setOidcToken: (oidcToken: string) => void;
-  removeOidcToken: () => void;
-  removeRefreshToken: () => void;
   updateSearchCriteria: (criteria: ExploreSearchIndex | '') => void;
   trySilentSignIn: (forceLogout?: boolean) => void;
   setApplicationsName: (applications: string[]) => void;

openmetadata-ui/src/main/resources/ui/src/pages/SamlCallback/index.tsx

@@ -19,12 +19,12 @@ import Loader from '../../components/common/Loader/Loader';
 import { REFRESH_TOKEN_KEY } from '../../constants/constants';
 import { useApplicationStore } from '../../hooks/useApplicationStore';
 import useCustomLocation from '../../hooks/useCustomLocation/useCustomLocation';
+import { setOidcToken, setRefreshToken } from '../../utils/LocalStorageUtils';
 
 const cookieStorage = new CookieStorage();
 
 const SamlCallback = () => {
-  const { handleSuccessfulLogin, setOidcToken, setRefreshToken } =
+  const { handleSuccessfulLogin } = useApplicationStore();
-    useApplicationStore();
   const location = useCustomLocation();
   const { t } = useTranslation();
 

openmetadata-ui/src/main/resources/ui/src/pages/SwaggerPage/index.test.tsx

@@ -23,9 +23,12 @@ jest.mock('./RapiDocReact', () => {
   ));
 });
 
+jest.mock('../../utils/LocalStorageUtils', () => ({
+  getOidcToken: jest.fn().mockReturnValue('fakeToken'),
+}));
+
 jest.mock('../../hooks/useApplicationStore', () => ({
   useApplicationStore: jest.fn().mockImplementation(() => ({
-    getOidcToken: () => 'fakeToken',
     theme: {
       primaryColor: '#9c27b0',
     },

openmetadata-ui/src/main/resources/ui/src/pages/SwaggerPage/index.tsx

@@ -17,11 +17,12 @@ import {
   TEXT_BODY_COLOR,
 } from '../../constants/constants';
 import { useApplicationStore } from '../../hooks/useApplicationStore';
+import { getOidcToken } from '../../utils/LocalStorageUtils';
 import RapiDocReact from './RapiDocReact';
 import './swagger.less';
 
 const SwaggerPage = () => {
-  const { getOidcToken, theme } = useApplicationStore();
+  const { theme } = useApplicationStore();
   const idToken = getOidcToken();
 
   return (

openmetadata-ui/src/main/resources/ui/src/utils/Auth/TokenService/TokenServiceUtil.ts

@@ -16,21 +16,23 @@ import { AccessTokenResponse } from '../../../rest/auth-API';
 import { extractDetailsFromToken } from '../../AuthProvider.util';
 import { getOidcToken } from '../../LocalStorageUtils';
 
+const REFRESH_IN_PROGRESS_KEY = 'refreshInProgress'; // Key to track if refresh is in progress
+
 type RenewTokenCallback = () =>
   | Promise<string>
   | Promise<AccessTokenResponse>
   | Promise<void>;
 
-class TokenService {
+const REFRESHED_KEY = 'tokenRefreshed';
-  channel: BroadcastChannel;
-  renewToken: RenewTokenCallback;
-  tokeUpdateInProgress: boolean;
 
-  constructor(renewToken: RenewTokenCallback) {
+class TokenService {
-    this.channel = new BroadcastChannel('auth_channel');
+  renewToken: RenewTokenCallback | null = null;
-    this.renewToken = renewToken;
+  refreshSuccessCallback: (() => void) | null = null;
-    this.channel.onmessage = this.handleTokenUpdate.bind(this);
+  private static _instance: TokenService;
-    this.tokeUpdateInProgress = false;
+
+  constructor() {
+    this.clearRefreshInProgress();
+    this.refreshToken = this.refreshToken.bind(this);
   }
 
   // This method will update token across tabs on receiving message to the channel
@@ -41,18 +43,40 @@ class TokenService {
       data: { type, token },
     } = event;
     if (type === 'TOKEN_UPDATE' && token) {
-      if (typeof token !== 'string') {
+      // Token is updated in localStorage hence no need to pass it
-        useApplicationStore.getState().setOidcToken(token.accessToken);
+      this.refreshSuccessCallback && this.refreshSuccessCallback();
-        useApplicationStore.getState().setRefreshToken(token.refreshToken);
-        useApplicationStore.getState().updateAxiosInterceptors();
-      } else {
-        useApplicationStore.getState().setOidcToken(token);
-      }
     }
   }
 
+  // Singleton instance of TokenService
+  static getInstance() {
+    if (!TokenService._instance) {
+      TokenService._instance = new TokenService();
+    }
+
+    return TokenService._instance;
+  }
+
+  public updateRenewToken(renewToken: RenewTokenCallback) {
+    this.renewToken = renewToken;
+  }
+
+  public updateRefreshSuccessCallback(callback: () => void) {
+    window.addEventListener('storage', (event) => {
+      if (event.key === REFRESHED_KEY && event.newValue === 'true') {
+        callback(); // Notify the tab that the token was refreshed
+        // Clear once notified
+        localStorage.removeItem(REFRESHED_KEY);
+      }
+    });
+  }
+
   // Refresh the token if it is expired
   async refreshToken() {
+    if (this.isTokenUpdateInProgress()) {
+      return;
+    }
+
     const token = getOidcToken();
     const { isExpired, timeoutExpiry } = extractDetailsFromToken(token);
 
@@ -61,11 +85,12 @@ class TokenService {
       // Logic to refresh the token
       const newToken = await this.fetchNewToken();
       // To update all the tabs on updating channel token
-      this.channel.postMessage({ type: 'TOKEN_UPDATE', token: newToken });
+      // Notify all tabs that the token has been refreshed
+      localStorage.setItem(REFRESHED_KEY, 'true');
 
       return newToken;
     } else {
-      return token;
+      return null;
     }
   }
 
@@ -74,26 +99,39 @@ class TokenService {
     let response: string | AccessTokenResponse | null | void = null;
     if (typeof this.renewToken === 'function') {
       try {
-        this.tokeUpdateInProgress = true;
+        this.setRefreshInProgress();
         response = await this.renewToken();
       } catch (error) {
         // Silent Frame window timeout error since it doesn't affect refresh token process
         if ((error as AxiosError).message !== 'Frame window timed out') {
           // Perform logout for any error
           useApplicationStore.getState().onLogoutHandler();
+          this.clearRefreshInProgress();
         }
         // Do nothing
       } finally {
-        this.tokeUpdateInProgress = false;
+        // If response is not null then clear the refresh flag
+        // For Callback based refresh token, response will be void
+        response && this.clearRefreshInProgress();
       }
     }
 
     return response;
   }
 
-  // Tracker for any ongoing token update
+  // Set refresh in progress (used by the tab that initiates the refresh)
+  setRefreshInProgress() {
+    localStorage.setItem(REFRESH_IN_PROGRESS_KEY, 'true');
+  }
+
+  // Clear the refresh flag (used after refresh is complete)
+  clearRefreshInProgress() {
+    localStorage.removeItem(REFRESH_IN_PROGRESS_KEY);
+  }
+
+  // Check if a refresh is already in progress (used by other tabs)
   isTokenUpdateInProgress() {
-    return this.tokeUpdateInProgress;
+    return localStorage.getItem(REFRESH_IN_PROGRESS_KEY) === 'true';
   }
 }
 

openmetadata-ui/src/main/resources/ui/src/utils/AuthProvider.util.ts

@@ -34,8 +34,8 @@ import {
   ClientType,
 } from '../generated/configuration/authenticationConfiguration';
 import { AuthProvider } from '../generated/settings/settings';
-import { useApplicationStore } from '../hooks/useApplicationStore';
 import { isDev } from './EnvironmentUtils';
+import { setOidcToken } from './LocalStorageUtils';
 
 const cookieStorage = new CookieStorage();
 
@@ -423,7 +423,6 @@ export const prepareUserProfileFromClaims = ({
 export const parseMSALResponse = (response: AuthenticationResult): OidcUser => {
   // Call your API with the access token and return the data you need to save in state
   const { idToken, scopes, account } = response;
-  const { setOidcToken } = useApplicationStore.getState();
 
   const user = {
     id_token: idToken,

openmetadata-ui/src/main/resources/ui/src/utils/LocalStorageUtils.ts

@@ -14,7 +14,27 @@ import { OM_SESSION_KEY } from '../hooks/useApplicationStore';
 
 export const getOidcToken = (): string => {
   return (
-    JSON.parse(localStorage.getItem(OM_SESSION_KEY) ?? '{}')?.state
+    JSON.parse(localStorage.getItem(OM_SESSION_KEY) ?? '{}')?.oidcIdToken ?? ''
-      ?.oidcIdToken ?? ''
   );
 };
+
+export const setOidcToken = (token: string) => {
+  const session = JSON.parse(localStorage.getItem(OM_SESSION_KEY) ?? '{}');
+
+  session.oidcIdToken = token;
+  localStorage.setItem(OM_SESSION_KEY, JSON.stringify(session));
+};
+
+export const getRefreshToken = (): string => {
+  return (
+    JSON.parse(localStorage.getItem(OM_SESSION_KEY) ?? '{}')?.refreshTokenKey ??
+    ''
+  );
+};
+
+export const setRefreshToken = (token: string) => {
+  const session = JSON.parse(localStorage.getItem(OM_SESSION_KEY) ?? '{}');
+
+  session.refreshTokenKey = token;
+  localStorage.setItem(OM_SESSION_KEY, JSON.stringify(session));
+};