From 298495979ddfcd96a9b0eff3e889563bced05eb4 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 3 Mar 2025 17:21:14 -0800
Subject: [PATCH] use common getAuthToken

---
 .../src/app/onboarding/page.tsx               |  4 +-
 ui/litellm-dashboard/src/app/page.tsx         |  9 +---
 .../src/components/user_dashboard.tsx         | 11 +----
 ui/litellm-dashboard/src/utils/cookieUtils.ts | 44 ++++++++++++++-----
 4 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/onboarding/page.tsx b/ui/litellm-dashboard/src/app/onboarding/page.tsx
index e46e46fcb54..6b06468a397 100644
--- a/ui/litellm-dashboard/src/app/onboarding/page.tsx
+++ b/ui/litellm-dashboard/src/app/onboarding/page.tsx
@@ -20,12 +20,12 @@ import {
 } from "@/components/networking";
 import { jwtDecode } from "jwt-decode";
 import { Form, Button as Button2, message } from "antd";
-import { getCookie } from "@/utils/cookieUtils";
+import { getAuthToken } from "@/utils/cookieUtils";
 
 export default function Onboarding() {
   const [form] = Form.useForm();
   const searchParams = useSearchParams()!;
-  const token = getCookie('token');
+  const token = getAuthToken();
   const inviteID = searchParams.get("invitation_id");
   const [accessToken, setAccessToken] = useState<string | null>(null);
   const [defaultUserEmail, setDefaultUserEmail] = useState<string>("");
diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 2612cab5946..fa07c08615c 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -30,12 +30,7 @@ import { Organization } from "@/components/networking";
 import GuardrailsPanel from "@/components/guardrails";
 import { fetchUserModels } from "@/components/create_key_button";
 import { fetchTeams } from "@/components/common_components/fetch_teams";
-function getCookie(name: string) {
-  const cookieValue = document.cookie
-    .split("; ")
-    .find((row) => row.startsWith(name + "="));
-  return cookieValue ? cookieValue.split("=")[1] : null;
-}
+import { getAuthToken } from "@/utils/cookieUtils";
 
 function formatUserRole(userRole: string) {
   if (!userRole) {
@@ -117,7 +112,7 @@ export default function CreateKeyPage() {
   const [accessToken, setAccessToken] = useState<string | null>(null);
 
   useEffect(() => {
-    const token = getCookie("token");
+    const token = getAuthToken();
     setToken(token);
   }, []);
 
diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx
index 22b47d525fc..9f4230e6b28 100644
--- a/ui/litellm-dashboard/src/components/user_dashboard.tsx
+++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx
@@ -21,6 +21,7 @@ import { useSearchParams, useRouter } from "next/navigation";
 import { Team } from "./key_team_helpers/key_list";
 import { jwtDecode } from "jwt-decode";
 import { Typography } from "antd";
+import { getAuthToken } from "@/utils/cookieUtils";
 import { clearTokenCookies } from "@/utils/cookieUtils";
 const isLocal = process.env.NODE_ENV === "development";
 if (isLocal != true) {
@@ -45,14 +46,6 @@ export type UserInfo = {
   spend: number;
 }
 
-function getCookie(name: string) {
-  console.log("COOKIES", document.cookie)
-  const cookieValue = document.cookie
-      .split('; ')
-      .find(row => row.startsWith(name + '='));
-  return cookieValue ? cookieValue.split('=')[1] : null;
-}
-
 interface UserDashboardProps {
   userID: string | null;
   userRole: string | null;
@@ -94,7 +87,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
   // Assuming useSearchParams() hook exists and works in your setup
   const searchParams = useSearchParams()!;
 
-  const token = getCookie('token');
+  const token = getAuthToken();
 
   const invitation_id = searchParams.get("invitation_id");
 
diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.ts b/ui/litellm-dashboard/src/utils/cookieUtils.ts
index a09cf4e97f6..857deda4bb8 100644
--- a/ui/litellm-dashboard/src/utils/cookieUtils.ts
+++ b/ui/litellm-dashboard/src/utils/cookieUtils.ts
@@ -31,14 +31,36 @@ export function clearTokenCookies() {
   console.log("After clearing cookies:", document.cookie);
 }
 
-/**
- * Gets a cookie value by name
- * @param name The name of the cookie to retrieve
- * @returns The cookie value or null if not found
- */
-export function getCookie(name: string) {
-  const cookieValue = document.cookie
-    .split('; ')
-    .find(row => row.startsWith(name + '='));
-  return cookieValue ? cookieValue.split('=')[1] : null;
-} 
\ No newline at end of file
+export function getAuthToken() {
+    // Check if we're in a browser environment
+    if (typeof window === 'undefined' || typeof document === 'undefined') {
+      return null;
+    }
+    
+    const tokenPattern = /^token_(\d+)$/;
+    const allCookies = document.cookie.split("; ");
+    
+    const tokenCookies = allCookies
+      .map(cookie => {
+        const parts = cookie.split("=");
+        const name = parts[0];
+        const match = name.match(tokenPattern);
+        if (match) {
+          return {
+            name,
+            timestamp: parseInt(match[1], 10),
+            value: parts.slice(1).join("=")
+          };
+        }
+        return null;
+      })
+      .filter(cookie => cookie !== null);
+    
+    if (tokenCookies.length > 0) {
+      // Sort by timestamp (newest first)
+      tokenCookies.sort((a, b) => b.timestamp - a.timestamp);
+      return tokenCookies[0].value;
+    }
+    
+    return null;
+}
\ No newline at end of file