fix(user_dashboard.tsx): add token expiry logic to user dashboard (#10250)

* fix(user_dashboard.tsx): add token expiry logic to user dashboard

if token expired redirect to `/sso/key/generate` for login

* fix(user_dashboard.tsx): check key health on login - if invalid -> redirect to login

handles invalid / expired key scenario

* fix(user_dashboard.tsx): fix linting error

* fix(page.tsx): fix invitation link flow

ui/litellm-dashboard/src/app/page.tsx

@@ -139,6 +139,8 @@ export default function CreateKeyPage() {
 
   const [accessToken, setAccessToken] = useState<string | null>(null);
 
+  const redirectToLogin = authLoading === false && token === null && invitation_id === null;
+
   useEffect(() => {
     const token = getCookie("token");
     setToken(token);
@@ -146,7 +148,7 @@ export default function CreateKeyPage() {
   }, []);
 
   useEffect(() => {
-    if (authLoading === false && token === null) {
+    if (redirectToLogin) {
       window.location.href = (proxyBaseUrl || "") + "/sso/key/generate"
     }
   }, [token, authLoading])
@@ -221,7 +223,7 @@ export default function CreateKeyPage() {
     }
   }, [accessToken, userID, userRole]);
 
-  if (authLoading || (authLoading == false && token === null)) {
+  if (authLoading) {
     return <LoadingScreen />
   }
 

ui/litellm-dashboard/src/components/networking.tsx

@@ -2493,6 +2493,9 @@ export const keyInfoCall = async (accessToken: String, keys: String[]) => {
 
     if (!response.ok) {
       const errorData = await response.text();
+      if (errorData.includes("Invalid proxy server token passed")) {
+        throw new Error("Invalid proxy server token passed");
+      }
       handleError(errorData);
       throw new Error("Network response was not ok");
     }

ui/litellm-dashboard/src/components/user_dashboard.tsx

@@ -7,7 +7,8 @@ import {
   getProxyUISettings,
   Organization,
   organizationListCall,
-  DEFAULT_ORGANIZATION
+  DEFAULT_ORGANIZATION,
+  keyInfoCall
 } from "./networking";
 import { fetchTeams } from "./common_components/fetch_teams";
 import { Grid, Col, Card, Text, Title } from "@tremor/react";
@@ -192,6 +193,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
               null,
               null
             );
+            
 
             setUserSpendData(response["user_info"]);
             console.log(`userSpendData: ${JSON.stringify(userSpendData)}`)
@@ -238,8 +240,11 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
               "userModels" + userID,
               JSON.stringify(available_model_names)
             );
-          } catch (error) {
+          } catch (error: any) {
             console.error("There was an error fetching the data", error);
+            if (error.message.includes("Invalid proxy server token passed")) {
+              gotoLogin();
+            }
             // Optionally, update your UI to reflect the error state here as well
           }
         };
@@ -249,6 +254,24 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
     }
   }, [userID, token, accessToken, keys, userRole]);
 
+
+  useEffect(() => {
+    // check key health - if it's invalid, redirect to login
+    if (accessToken) {
+      const fetchKeyInfo = async () => {
+        try {
+          const keyInfo = await keyInfoCall(accessToken, [accessToken]);
+          console.log("keyInfo: ", keyInfo);
+        } catch (error: any) {
+          if (error.message.includes("Invalid proxy server token passed")) {
+            gotoLogin();
+          }
+        }
+      }
+      fetchKeyInfo();
+    }
+  }, [accessToken]);
+
   useEffect(() => {
     console.log(`currentOrg: ${JSON.stringify(currentOrg)}, accessToken: ${accessToken}, userID: ${userID}, userRole: ${userRole}`)
     if (accessToken) {
@@ -295,24 +318,68 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
     )
   }
 
-
-  if (token == null) {
-    // user is not logged in as yet 
-    console.log("All cookies before redirect:", document.cookie);
-    
+  function gotoLogin() {
     // Clear token cookies using the utility function
     clearTokenCookies();
     
     const url = proxyBaseUrl
       ? `${proxyBaseUrl}/sso/key/generate`
       : `/sso/key/generate`;
-    
+
     console.log("Full URL:", url);
-    window.location.href = url;
+    window.location.href = url; 
 
     return null;
-  } else if (accessToken == null) {
+  }
+
+  if (token == null) {
+    // user is not logged in as yet 
+    console.log("All cookies before redirect:", document.cookie);
+    
+    // Clear token cookies using the utility function
+    gotoLogin();
     return null;
+  } else {
+    // Check if token is expired
+    try {
+      const decoded = jwtDecode(token) as { [key: string]: any };
+      console.log("Decoded token:", decoded);
+      const expTime = decoded.exp;
+      const currentTime = Math.floor(Date.now() / 1000);
+      
+      if (expTime && currentTime >= expTime) {
+        console.log("Token expired, redirecting to login");
+        
+        // Clear token cookies
+        clearTokenCookies();
+        
+        const url = proxyBaseUrl
+          ? `${proxyBaseUrl}/sso/key/generate`
+          : `/sso/key/generate`;
+        
+        console.log("Full URL for expired token:", url);
+        window.location.href = url;
+        
+        return null;
+      }
+    } catch (error) {
+      console.error("Error decoding token:", error);
+      // If there's an error decoding the token, consider it invalid
+      clearTokenCookies();
+      
+      const url = proxyBaseUrl
+        ? `${proxyBaseUrl}/sso/key/generate`
+        : `/sso/key/generate`;
+      
+      console.log("Full URL after token decode error:", url);
+      window.location.href = url;
+      
+      return null;
+    }
+    
+    if (accessToken == null) {
+      return null;
+    }
   }
 
   if (userID == null) {