From 828e9eab78f946997b0ecf4e96602fd9dde3fea2 Mon Sep 17 00:00:00 2001
From: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com>
Date: Wed, 19 Mar 2025 10:20:55 -0300
Subject: [PATCH] fix: Improve token refresh logic and API interceptor handling
 (#7158)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 🐛 (api.tsx): fix logic to correctly handle refreshing access token and retrying requests on authentication errors
🐛 (use-post-refresh-access.ts): fix retry logic to always retry refreshing access token regardless of auto login setting

* ✨ (use-get-autologin.ts): add support for autoLogin state in useGetAutoLogin hook to handle auto login scenarios more effectively
📝 (use-post-refresh-access.ts): import and use autoLogin state from authStore in useRefreshAccessToken hook to determine retry count based on auto login status
---
 src/frontend/src/controllers/API/api.tsx      | 36 ++++++++++---------
 .../API/queries/auth/use-get-autologin.ts     | 10 ++++--
 .../queries/auth/use-post-refresh-access.ts   |  4 ++-
 3 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/src/frontend/src/controllers/API/api.tsx b/src/frontend/src/controllers/API/api.tsx
index faa35f05b..596424345 100644
--- a/src/frontend/src/controllers/API/api.tsx
+++ b/src/frontend/src/controllers/API/api.tsx
@@ -65,25 +65,27 @@ function ApiInterceptor() {
         const isAuthenticationError =
           error?.response?.status === 403 || error?.response?.status === 401;
 
-        if (isAuthenticationError && !IS_AUTO_LOGIN) {
-          if (autoLogin !== undefined && !autoLogin) {
-            if (
-              error?.config?.url?.includes("github") ||
-              error?.config?.url?.includes("public")
-            ) {
-              return Promise.reject(error);
-            }
-            const stillRefresh = checkErrorCount();
-            if (!stillRefresh) {
-              return Promise.reject(error);
-            }
+        const shouldRetryRefresh =
+          (isAuthenticationError && !IS_AUTO_LOGIN) ||
+          (isAuthenticationError && !autoLogin && autoLogin !== undefined);
 
-            await tryToRenewAccessToken(error);
+        if (shouldRetryRefresh) {
+          if (
+            error?.config?.url?.includes("github") ||
+            error?.config?.url?.includes("public")
+          ) {
+            return Promise.reject(error);
+          }
+          const stillRefresh = checkErrorCount();
+          if (!stillRefresh) {
+            return Promise.reject(error);
+          }
 
-            const accessToken = cookies.get(LANGFLOW_ACCESS_TOKEN);
-            if (!accessToken && error?.config?.url?.includes("login")) {
-              return Promise.reject(error);
-            }
+          await tryToRenewAccessToken(error);
+
+          const accessToken = cookies.get(LANGFLOW_ACCESS_TOKEN);
+          if (!accessToken && error?.config?.url?.includes("login")) {
+            return Promise.reject(error);
           }
         }
 
diff --git a/src/frontend/src/controllers/API/queries/auth/use-get-autologin.ts b/src/frontend/src/controllers/API/queries/auth/use-get-autologin.ts
index d6ced0565..4d2e799e6 100644
--- a/src/frontend/src/controllers/API/queries/auth/use-get-autologin.ts
+++ b/src/frontend/src/controllers/API/queries/auth/use-get-autologin.ts
@@ -31,6 +31,7 @@ export const useGetAutoLogin: useQueryFunctionType<undefined, undefined> = (
   const isLoginPage = location.pathname.includes("login");
   const navigate = useCustomNavigate();
   const { mutateAsync: mutationLogout } = useLogout();
+  const autoLogin = useAuthStore((state) => state.autoLogin);
 
   const retryCountRef = useRef(0);
   const retryTimerRef = useRef<NodeJS.Timeout | null>(null);
@@ -67,8 +68,13 @@ export const useGetAutoLogin: useQueryFunctionType<undefined, undefined> = (
   };
 
   const handleAutoLoginError = async () => {
-    const manualLoginNotAuthenticated = !isAuthenticated && !IS_AUTO_LOGIN;
-    const autoLoginNotAuthenticated = !isAuthenticated && IS_AUTO_LOGIN;
+    const manualLoginNotAuthenticated =
+      (!isAuthenticated && !IS_AUTO_LOGIN) ||
+      (!isAuthenticated && autoLogin !== undefined && !autoLogin);
+
+    const autoLoginNotAuthenticated =
+      (!isAuthenticated && IS_AUTO_LOGIN) ||
+      (!isAuthenticated && autoLogin !== undefined && autoLogin);
 
     if (manualLoginNotAuthenticated) {
       await mutationLogout();
diff --git a/src/frontend/src/controllers/API/queries/auth/use-post-refresh-access.ts b/src/frontend/src/controllers/API/queries/auth/use-post-refresh-access.ts
index 03c767cc0..5b3568aed 100644
--- a/src/frontend/src/controllers/API/queries/auth/use-post-refresh-access.ts
+++ b/src/frontend/src/controllers/API/queries/auth/use-post-refresh-access.ts
@@ -1,4 +1,5 @@
 import { IS_AUTO_LOGIN, LANGFLOW_REFRESH_TOKEN } from "@/constants/constants";
+import useAuthStore from "@/stores/authStore";
 import { useMutationFunctionType } from "@/types/api";
 import { Cookies } from "react-cookie";
 import { api } from "../../api";
@@ -17,6 +18,7 @@ export const useRefreshAccessToken: useMutationFunctionType<
 > = (options?) => {
   const { mutate } = UseRequestProcessor();
   const cookies = new Cookies();
+  const autoLogin = useAuthStore((state) => state.autoLogin);
 
   async function refreshAccess(): Promise<IRefreshAccessToken> {
     const res = await api.post<IRefreshAccessToken>(`${getURL("REFRESH")}`);
@@ -27,7 +29,7 @@ export const useRefreshAccessToken: useMutationFunctionType<
 
   const mutation = mutate(["useRefreshAccessToken"], refreshAccess, {
     ...options,
-    retry: IS_AUTO_LOGIN ? 0 : 3,
+    retry: IS_AUTO_LOGIN || autoLogin ? 0 : 2,
   });
 
   return mutation;