From 1baa958923d760e1c8221250e377a4e96bd5f2a3 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 17 Feb 2024 19:03:22 -0800
Subject: [PATCH] (fix) ui - reload key spend info

---
 .../src/components/user_dashboard.tsx         | 20 ++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx
index bf95d09f73..f87303a125 100644
--- a/ui/litellm-dashboard/src/components/user_dashboard.tsx
+++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx
@@ -48,6 +48,11 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
   const token = searchParams.get("token");
   const [accessToken, setAccessToken] = useState<string | null>(null);
   const [userModels, setUserModels] = useState<string[]>([]);
+  window.addEventListener('beforeunload', function() {
+      // Clear session storage
+      sessionStorage.clear();
+  });
+
 
   function formatUserRole(userRole: string) {
     if (!userRole) {
@@ -70,6 +75,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
 
   // Moved useEffect inside the component and used a condition to run fetch only if the params are available
   useEffect(() => {
+
     if (token) {
       const decoded = jwtDecode(token) as { [key: string]: any };
       if (decoded) {
@@ -97,22 +103,22 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
       }
     }
     if (userID && accessToken && userRole && !data) {
-      const cachedData = localStorage.getItem("userData" + userID);
-      const cachedSpendData = localStorage.getItem("userSpendData" + userID);
-      const cachedUserModels = localStorage.getItem("userModels" + userID);
+      const cachedData = sessionStorage.getItem("userData" + userID);
+      const cachedSpendData = sessionStorage.getItem("userSpendData" + userID);
+      const cachedUserModels = sessionStorage.getItem("userModels" + userID);
       if (cachedData && cachedSpendData && cachedUserModels) {
         setData(JSON.parse(cachedData));
         setUserSpendData(JSON.parse(cachedSpendData));
         setUserModels(JSON.parse(cachedUserModels));
-
+  
       } else {
         const fetchData = async () => {
           try {
             const response = await userInfoCall(accessToken, userID, userRole);
             setUserSpendData(response["user_info"]);
             setData(response["keys"]); // Assuming this is the correct path to your data
-            localStorage.setItem("userData" + userID, JSON.stringify(response["keys"]));
-            localStorage.setItem(
+            sessionStorage.setItem("userData" + userID, JSON.stringify(response["keys"]));
+            sessionStorage.setItem(
               "userSpendData" + userID,
               JSON.stringify(response["user_info"])
             );
@@ -126,7 +132,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
 
             console.log("userModels:", userModels);
 
-            localStorage.setItem("userModels" + userID, JSON.stringify(available_model_names));
+            sessionStorage.setItem("userModels" + userID, JSON.stringify(available_model_names));