From 4c89aed4d9f64f4a62d30b49a6ea441188909d77 Mon Sep 17 00:00:00 2001
From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com>
Date: Mon, 3 Feb 2025 20:04:17 -0800
Subject: [PATCH] Add check for expired refresh or session token that exists

---
 client/src/lib/hooks/useConnection.ts | 34 ++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts
index 58ea0a8..e08e0ae 100644
--- a/client/src/lib/hooks/useConnection.ts
+++ b/client/src/lib/hooks/useConnection.ts
@@ -131,10 +131,13 @@ export function useConnection({
       return tokens.access_token;
     } catch (error) {
       console.error("Token refresh failed:", error);
-      // Clear tokens and redirect to home
+      // If refresh token is expired/invalid (401) or any other error,
+      // clear tokens and redirect to home to trigger re-authentication
       sessionStorage.removeItem(SESSION_KEYS.ACCESS_TOKEN);
       sessionStorage.removeItem(SESSION_KEYS.REFRESH_TOKEN);
-      window.location.href = "/";
+      sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl);
+      const redirectUrl = await startOAuthFlow(sseUrl);
+      window.location.href = redirectUrl;
       throw error;
     }
   };
@@ -178,12 +181,27 @@ export function useConnection({
           fetch: async (url, init) => {
             const response = await fetch(url, { ...init, headers });
             
-            if (response.status === 401 && sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) {
-              // Try to refresh the token
-              const newAccessToken = await handleTokenRefresh();
-              headers["Authorization"] = `Bearer ${newAccessToken}`;
-              // Retry the request with new token
-              return fetch(url, { ...init, headers });
+            if (response.status === 401) {
+              // First try to refresh if we have a refresh token
+              if (sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) {
+                try {
+                  const newAccessToken = await handleTokenRefresh();
+                  headers["Authorization"] = `Bearer ${newAccessToken}`;
+                  // Retry the request with new token
+                  return fetch(url, { ...init, headers });
+                } catch (error) {
+                  console.error("Token refresh failed:", error);
+                }
+              }
+              
+              // If we have an access token but refresh failed or wasn't available,
+              // we need to re-authenticate since the token is invalid
+              if (sessionStorage.getItem(SESSION_KEYS.ACCESS_TOKEN)) {
+                sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl);
+                const redirectUrl = await startOAuthFlow(sseUrl);
+                window.location.href = redirectUrl;
+                return new Response(); // This won't actually be used due to redirect
+              }
             }
             
             return response;