From fdcdee8775d170c18c212f6ce1bdaa58fcbcc25e Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 4 Jul 2024 15:40:49 -0700
Subject: [PATCH] fix invite id flow

---
 .../src/app/onboarding/page.tsx               |  9 +++--
 ui/litellm-dashboard/src/app/page.tsx         | 34 +++++++++++++++----
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/onboarding/page.tsx b/ui/litellm-dashboard/src/app/onboarding/page.tsx
index a31afb0f9..373944a98 100644
--- a/ui/litellm-dashboard/src/app/onboarding/page.tsx
+++ b/ui/litellm-dashboard/src/app/onboarding/page.tsx
@@ -33,7 +33,7 @@ export default function Onboarding() {
   const [form] = Form.useForm();
   const searchParams = useSearchParams();
   const token = getCookie('token');
-  const inviteID = searchParams.get("id");
+  const inviteID = searchParams.get("invitation_id");
   const [accessToken, setAccessToken] = useState<string | null>(null);
   const [defaultUserEmail, setDefaultUserEmail] = useState<string>("");
   const [userEmail, setUserEmail] = useState<string>("");
@@ -92,7 +92,10 @@ export default function Onboarding() {
     ).then((data) => {
       let litellm_dashboard_ui = "/ui/";
       const user_id = data.data?.user_id || data.user_id;
-      litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwtToken;
+      litellm_dashboard_ui += "?userID=" + user_id;
+
+      // set cookie "token" to jwtToken
+      document.cookie = "token=" + jwtToken;
       console.log("redirecting to:", litellm_dashboard_ui);
 
       window.location.href = litellm_dashboard_ui;
@@ -101,7 +104,7 @@ export default function Onboarding() {
     // redirect to login page
   };
   return (
-    <div className="mx-auto max-w-md mt-10">
+    <div className="mx-auto w-full max-w-md mt-10">
       <Card>
         <Title className="text-sm mb-5 text-center">🚅 LiteLLM</Title>
         <Title className="text-xl">Sign up</Title>
diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 41f0e8529..1aaf8cbc3 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -77,6 +77,7 @@ const CreateKeyPage = () => {
   const searchParams = useSearchParams();
   const [modelData, setModelData] = useState<any>({ data: [] });
   const userID = searchParams.get("userID");
+  const invitation_id = searchParams.get("invitation_id");
   const token = getCookie('token');
 
   const [page, setPage] = useState("api-keys");
@@ -128,7 +129,23 @@ const CreateKeyPage = () => {
 
   return (
     <Suspense fallback={<div>Loading...</div>}>
-      <div className="flex flex-col min-h-screen">
+      {
+        invitation_id ? (
+          <UserDashboard
+              userID={userID}
+              userRole={userRole}
+              teams={teams}
+              keys={keys}
+              setUserRole={setUserRole}
+              userEmail={userEmail}
+              setUserEmail={setUserEmail}
+              setTeams={setTeams}
+              setKeys={setKeys}
+              setProxySettings={setProxySettings}
+              proxySettings={proxySettings}
+            />
+        ) : (
+        <div className="flex flex-col min-h-screen">
         <Navbar
           userID={userID}
           userRole={userRole}
@@ -140,11 +157,11 @@ const CreateKeyPage = () => {
         />
         <div className="flex flex-1 overflow-auto">
           <div className="mt-8">
-            <Sidebar
-              setPage={setPage}
-              userRole={userRole}
-              defaultSelectedKey={null}
-            />
+                <Sidebar
+                setPage={setPage}
+                userRole={userRole}
+                defaultSelectedKey={null}
+              />            
           </div>
 
           {page == "api-keys" ? (
@@ -250,7 +267,10 @@ const CreateKeyPage = () => {
             />
           )}
         </div>
-      </div>
+      </div>         
+        )
+      }
+      
     </Suspense>
   );
 };