diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d092fd8800..2e2720c104 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -43,7 +43,6 @@ from litellm.proxy.management_endpoints.sso_helper_utils import ( ) from litellm.proxy.management_endpoints.team_endpoints import team_member_add from litellm.proxy.management_endpoints.types import CustomOpenID -from litellm.proxy.management_helpers.ui_session_handler import UISessionHandler from litellm.secret_managers.main import str_to_bool if TYPE_CHECKING: @@ -691,10 +690,9 @@ async def auth_callback(request: Request): # noqa: PLR0915 ) if user_id is not None and isinstance(user_id, str): litellm_dashboard_ui += "?userID=" + user_id - - return UISessionHandler.generate_authenticated_redirect_response( - redirect_url=litellm_dashboard_ui, jwt_token=jwt_token - ) + redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) + redirect_response.set_cookie(key="token", value=jwt_token, secure=True) + return redirect_response async def insert_sso_user( diff --git a/litellm/proxy/management_helpers/ui_session_handler.py b/litellm/proxy/management_helpers/ui_session_handler.py deleted file mode 100644 index 6afe158c33..0000000000 --- a/litellm/proxy/management_helpers/ui_session_handler.py +++ /dev/null @@ -1,24 +0,0 @@ -import time - -from fastapi.responses import RedirectResponse - - -class UISessionHandler: - @staticmethod - def generate_authenticated_redirect_response( - redirect_url: str, jwt_token: str - ) -> RedirectResponse: - redirect_response = RedirectResponse(url=redirect_url, status_code=303) - redirect_response.set_cookie( - key=UISessionHandler._generate_token_name(), - value=jwt_token, - secure=True, - samesite="strict", - ) - return redirect_response - - @staticmethod - def _generate_token_name() -> str: - current_timestamp = int(time.time()) - cookie_name = f"token_{current_timestamp}" - return cookie_name diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bcca95b310..19d0ebe9ea 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7372,8 +7372,6 @@ async def login(request: Request): # noqa: PLR0915 import multipart except ImportError: subprocess.run(["pip", "install", "python-multipart"]) - from litellm.proxy.management_helpers.ui_session_handler import UISessionHandler - global master_key if master_key is None: raise ProxyException( @@ -7491,9 +7489,9 @@ async def login(request: Request): # noqa: PLR0915 algorithm="HS256", ) litellm_dashboard_ui += "?userID=" + user_id - return UISessionHandler.generate_authenticated_redirect_response( - redirect_url=litellm_dashboard_ui, jwt_token=jwt_token - ) + redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) + redirect_response.set_cookie(key="token", value=jwt_token) + return redirect_response elif _user_row is not None: """ When sharing invite links @@ -7559,9 +7557,11 @@ async def login(request: Request): # noqa: PLR0915 algorithm="HS256", ) litellm_dashboard_ui += "?userID=" + user_id - return UISessionHandler.generate_authenticated_redirect_response( - redirect_url=litellm_dashboard_ui, jwt_token=jwt_token + redirect_response = RedirectResponse( + url=litellm_dashboard_ui, status_code=303 ) + redirect_response.set_cookie(key="token", value=jwt_token) + return redirect_response else: raise ProxyException( message=f"Invalid credentials used to access UI.\nNot valid credentials for {username}", diff --git a/ui/litellm-dashboard/src/app/onboarding/page.tsx b/ui/litellm-dashboard/src/app/onboarding/page.tsx index 0c194ba0cb..e46e46fcb5 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 { getAuthToken, setAuthToken } from "@/utils/cookieUtils"; +import { getCookie } from "@/utils/cookieUtils"; export default function Onboarding() { const [form] = Form.useForm(); const searchParams = useSearchParams()!; - const token = getAuthToken(); + const token = getCookie('token'); const inviteID = searchParams.get("invitation_id"); const [accessToken, setAccessToken] = useState(null); const [defaultUserEmail, setDefaultUserEmail] = useState(""); @@ -88,7 +88,7 @@ export default function Onboarding() { litellm_dashboard_ui += "?userID=" + user_id; // set cookie "token" to jwtToken - setAuthToken(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/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index fa07c08615..2612cab594 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -30,7 +30,12 @@ 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"; -import { getAuthToken } from "@/utils/cookieUtils"; +function getCookie(name: string) { + const cookieValue = document.cookie + .split("; ") + .find((row) => row.startsWith(name + "=")); + return cookieValue ? cookieValue.split("=")[1] : null; +} function formatUserRole(userRole: string) { if (!userRole) { @@ -112,7 +117,7 @@ export default function CreateKeyPage() { const [accessToken, setAccessToken] = useState(null); useEffect(() => { - const token = getAuthToken(); + const token = getCookie("token"); setToken(token); }, []); diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 9f4230e6b2..22b47d525f 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -21,7 +21,6 @@ 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) { @@ -46,6 +45,14 @@ 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; @@ -87,7 +94,7 @@ const UserDashboard: React.FC = ({ // Assuming useSearchParams() hook exists and works in your setup const searchParams = useSearchParams()!; - const token = getAuthToken(); + const token = getCookie('token'); 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 e181606e93..a09cf4e97f 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.ts @@ -13,82 +13,32 @@ export function clearTokenCookies() { const paths = ['/', '/ui']; const sameSiteValues = ['Lax', 'Strict', 'None']; - // Get all cookies - const allCookies = document.cookie.split("; "); - const tokenPattern = /^token_\d+$/; - - // Find all token cookies - const tokenCookieNames = allCookies - .map(cookie => cookie.split("=")[0]) - .filter(name => name === "token" || tokenPattern.test(name)); - - // Clear each token cookie with various combinations - tokenCookieNames.forEach(cookieName => { - paths.forEach(path => { - // Basic clearing - document.cookie = `${cookieName}=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=${path};`; - - // With domain - document.cookie = `${cookieName}=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=${path}; domain=${domain};`; - - // Try different SameSite values - sameSiteValues.forEach(sameSite => { - const secureFlag = sameSite === 'None' ? ' Secure;' : ''; - document.cookie = `${cookieName}=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=${path}; SameSite=${sameSite};${secureFlag}`; - document.cookie = `${cookieName}=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=${path}; domain=${domain}; SameSite=${sameSite};${secureFlag}`; - }); + paths.forEach(path => { + // Basic clearing + document.cookie = `token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=${path};`; + + // With domain + document.cookie = `token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=${path}; domain=${domain};`; + + // Try different SameSite values + sameSiteValues.forEach(sameSite => { + const secureFlag = sameSite === 'None' ? ' Secure;' : ''; + document.cookie = `token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=${path}; SameSite=${sameSite};${secureFlag}`; + document.cookie = `token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=${path}; domain=${domain}; SameSite=${sameSite};${secureFlag}`; }); }); console.log("After clearing cookies:", document.cookie); } -export function setAuthToken(token: string) { - // Generate a token name with current timestamp - const currentTimestamp = Math.floor(Date.now() / 1000); - const tokenName = `token_${currentTimestamp}`; - - // Set the cookie with the timestamp-based name - document.cookie = `${tokenName}=${token}; path=/; domain=${window.location.hostname};`; -} - -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]; - - // Explicitly skip cookies named just "token" - if (name === "token") { - return null; - } - - // Only match cookies with the token_{timestamp} format - const match = name.match(tokenPattern); - if (match) { - return { - name, - timestamp: parseInt(match[1], 10), - value: parts.slice(1).join("=") - }; - } - return null; - }) - .filter((cookie): cookie is { name: string; timestamp: number; value: string } => 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; -} +/** + * 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