diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 86dec9fcaf..d903e2665c 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -7,6 +7,7 @@ Has all /sso/* routes import asyncio import os +import time import uuid from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast @@ -44,6 +45,7 @@ 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,9 +693,10 @@ 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 - redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token, secure=True) - return redirect_response + + return UISessionHandler.generate_authenticated_redirect_response( + redirect_url=litellm_dashboard_ui, jwt_token=jwt_token + ) 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 new file mode 100644 index 0000000000..6afe158c33 --- /dev/null +++ b/litellm/proxy/management_helpers/ui_session_handler.py @@ -0,0 +1,24 @@ +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 19d0ebe9ea..bcca95b310 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7372,6 +7372,8 @@ 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( @@ -7489,9 +7491,9 @@ async def login(request: Request): # noqa: PLR0915 algorithm="HS256", ) litellm_dashboard_ui += "?userID=" + user_id - redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token) - return redirect_response + return UISessionHandler.generate_authenticated_redirect_response( + redirect_url=litellm_dashboard_ui, jwt_token=jwt_token + ) elif _user_row is not None: """ When sharing invite links @@ -7557,11 +7559,9 @@ async def login(request: Request): # noqa: PLR0915 algorithm="HS256", ) litellm_dashboard_ui += "?userID=" + user_id - redirect_response = RedirectResponse( - url=litellm_dashboard_ui, status_code=303 + return UISessionHandler.generate_authenticated_redirect_response( + redirect_url=litellm_dashboard_ui, jwt_token=jwt_token ) - 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 e46e46fcb5..0c194ba0cb 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, setAuthToken } 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(null); const [defaultUserEmail, setDefaultUserEmail] = useState(""); @@ -88,7 +88,7 @@ export default function Onboarding() { litellm_dashboard_ui += "?userID=" + user_id; // set cookie "token" to jwtToken - document.cookie = "token=" + jwtToken; + setAuthToken(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 2612cab594..fa07c08615 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(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 22b47d525f..9f4230e6b2 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 = ({ // 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 a09cf4e97f..e181606e93 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.ts @@ -13,32 +13,82 @@ export function clearTokenCookies() { const paths = ['/', '/ui']; const sameSiteValues = ['Lax', 'Strict', 'None']; - 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}`; + // 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}`; + }); }); }); 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 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; +}