From e952c666f3a1d4e5faba85c16ca87c8856bbfa2f Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 25 Nov 2024 22:41:45 -0800
Subject: [PATCH] (UI fix) UI does not reload when you login / open a new tab
 (#6909)

* store current page on url

* update menu history
---
 ui/litellm-dashboard/src/app/page.tsx         | 30 ++++++++++++-------
 .../src/components/leftnav.tsx                | 16 +++++++---
 2 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 9448be82a..94af12574 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -82,16 +82,26 @@ const CreateKeyPage = () => {
   const invitation_id = searchParams.get("invitation_id");
   const token = getCookie('token');
 
+  // Get page from URL, default to 'api-keys' if not present
   const [page, setPage] = useState(() => {
-    if (typeof window !== 'undefined') {
-      return localStorage.getItem('selectedPage') || "api-keys";
-    }
-    return "api-keys";
+    return searchParams.get('page') || 'api-keys';
   });
 
-  useEffect(() => {
-    localStorage.setItem('selectedPage', page);
-  }, [page]);
+  // Custom setPage function that updates URL
+  const updatePage = (newPage: string) => {
+    // Update URL without full page reload
+    const newSearchParams = new URLSearchParams(searchParams);
+    newSearchParams.set('page', newPage);
+    
+    // Use Next.js router to update URL
+    window.history.pushState(
+      null, 
+      '', 
+      `?${newSearchParams.toString()}`
+    );
+    
+    setPage(newPage);
+  };
 
   const [accessToken, setAccessToken] = useState<string | null>(null);
 
@@ -172,8 +182,8 @@ const CreateKeyPage = () => {
         />
         <div className="flex flex-1 overflow-auto">
           <div className="mt-8">
-                <Sidebar
-                setPage={setPage}
+              <Sidebar
+                setPage={updatePage}
                 userRole={userRole}
                 defaultSelectedKey={page}
               />            
@@ -289,7 +299,7 @@ const CreateKeyPage = () => {
             />
           )}
         </div>
-      </div>         
+      </div>
         )
       }
       
diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx
index 4304fe836..ba7519324 100644
--- a/ui/litellm-dashboard/src/components/leftnav.tsx
+++ b/ui/litellm-dashboard/src/components/leftnav.tsx
@@ -7,7 +7,7 @@ const { Sider } = Layout;
 
 // Define the props type
 interface SidebarProps {
-  setPage: React.Dispatch<React.SetStateAction<string>>;
+  setPage: (page: string) => void;
   userRole: string;
   defaultSelectedKey: string;
 }
@@ -67,9 +67,17 @@ const Sidebar: React.FC<SidebarProps> = ({
           style={{ height: "100%", borderRight: 0 }}
         >
           {filteredMenuItems.map(item => (
-            <Menu.Item key={item.key} onClick={() => setPage(item.page)}>
-              <Text>{item.label}</Text>
-            </Menu.Item>
+            <Menu.Item 
+            key={item.key} 
+            onClick={() => {
+              const newSearchParams = new URLSearchParams(window.location.search);
+              newSearchParams.set('page', item.page);
+              window.history.pushState(null, '', `?${newSearchParams.toString()}`);
+              setPage(item.page);
+            }}
+          >
+            <Text>{item.label}</Text>
+          </Menu.Item>
           ))}
         </Menu>
       </Sider>