From 30e5308cb3223981c715e723cdd3a6d2ee3aaa4b Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 3 Mar 2025 17:50:45 -0800
Subject: [PATCH] fix invitation link sign in logic

---
 litellm/proxy/proxy_server.py                      | 10 +++++-----
 ui/litellm-dashboard/src/app/onboarding/page.tsx   | 13 ++++++-------
 ui/litellm-dashboard/src/components/networking.tsx |  4 ++--
 3 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index bcca95b310..d60b7b80a8 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -7588,6 +7588,8 @@ async def onboarding(invite_link: str):
     - Pass in user_email if set
     """
     global prisma_client, master_key, general_settings
+    from litellm.proxy.management_helpers.ui_session_handler import UISessionHandler
+
     if master_key is None:
         raise ProxyException(
             message="Master Key not set for Proxy. Please set Master Key to use Admin UI. Set `LITELLM_MASTER_KEY` in .env or set general_settings:master_key in config.yaml.  https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.",
@@ -7688,11 +7690,9 @@ async def onboarding(invite_link: str):
     )
 
     litellm_dashboard_ui += "?token={}&user_email={}".format(jwt_token, user_email)
-    return {
-        "login_url": litellm_dashboard_ui,
-        "token": jwt_token,
-        "user_email": user_email,
-    }
+    return UISessionHandler.generate_authenticated_redirect_response(
+        redirect_url=litellm_dashboard_ui, jwt_token=jwt_token
+    )
 
 
 @app.post("/onboarding/claim_token", include_in_schema=False)
diff --git a/ui/litellm-dashboard/src/app/onboarding/page.tsx b/ui/litellm-dashboard/src/app/onboarding/page.tsx
index 6b06468a39..d7e9a72dba 100644
--- a/ui/litellm-dashboard/src/app/onboarding/page.tsx
+++ b/ui/litellm-dashboard/src/app/onboarding/page.tsx
@@ -12,6 +12,7 @@ import {
   Col,
 } from "@tremor/react";
 import { RiAlarmWarningLine, RiCheckboxCircleLine } from "@remixicon/react";
+import { clearTokenCookies, getAuthToken } from "@/utils/cookieUtils";
 import {
   invitationClaimCall,
   userUpdateUserCall,
@@ -20,7 +21,6 @@ import {
 } from "@/components/networking";
 import { jwtDecode } from "jwt-decode";
 import { Form, Button as Button2, message } from "antd";
-import { getAuthToken } from "@/utils/cookieUtils";
 
 export default function Onboarding() {
   const [form] = Form.useForm();
@@ -38,15 +38,18 @@ export default function Onboarding() {
     if (!inviteID) {
       return;
     }
+    clearTokenCookies();
     getOnboardingCredentials(inviteID).then((data) => {
       const login_url = data.login_url;
       console.log("login_url:", login_url);
       setLoginUrl(login_url);
 
-      const token = data.token;
+      const token = getAuthToken();
+      if (!token) {
+        return;
+      }
       const decoded = jwtDecode(token) as { [key: string]: any };
       setJwtToken(token);
-
       console.log("decoded:", decoded);
       setAccessToken(decoded.key);
 
@@ -86,11 +89,7 @@ export default function Onboarding() {
       let litellm_dashboard_ui = "/ui/";
       const user_id = data.data?.user_id || data.user_id;
       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;
     });
 
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index 4673e064f5..8c05e6d160 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/ui/litellm-dashboard/src/components/networking.tsx
@@ -2,7 +2,7 @@
  * Helper file for calls being made to proxy
  */
 import { message } from "antd";
-
+import { clearTokenCookies } from "@/utils/cookieUtils";
 const isLocal = process.env.NODE_ENV === "development";
 export const proxyBaseUrl = isLocal ? "http://localhost:4000" : null;
 if (isLocal != true) {
@@ -50,7 +50,7 @@ const handleError = async (errorData: string) => {
       message.info("UI Session Expired. Logging out.");
       lastErrorTime = currentTime;
       await sleep(3000); // 5 second sleep
-      document.cookie = "token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=/;";
+      clearTokenCookies();
       window.location.href = baseUrl;
     }
     lastErrorTime = currentTime;