From edd15b09052c945f8b65569e3760614779b32c0a Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Wed, 23 Apr 2025 16:51:27 -0700 Subject: [PATCH] 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 | 6 +- .../src/components/networking.tsx | 3 + .../src/components/user_dashboard.tsx | 87 ++++++++++++++++--- 3 files changed, 84 insertions(+), 12 deletions(-) diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index bcfdfc280b..1fea83d054 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -139,6 +139,8 @@ export default function CreateKeyPage() { const [accessToken, setAccessToken] = useState(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 } diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index e518314e64..8d2c9f535f 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/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"); } diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index c0f4c96d86..4279c124d0 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/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 = ({ null, null ); + setUserSpendData(response["user_info"]); console.log(`userSpendData: ${JSON.stringify(userSpendData)}`) @@ -238,8 +240,11 @@ const UserDashboard: React.FC = ({ "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 = ({ } }, [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 = ({ ) } - - 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) {