diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 021b59e295f..630aa3f3e15 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7502,6 +7502,12 @@ async def login(request: Request): litellm_dashboard_ui += "/ui/" import jwt + if litellm_master_key_hash is None: + raise HTTPException( + status_code=500, + detail={"error": "No master key set, please set LITELLM_MASTER_KEY"}, + ) + jwt_token = jwt.encode( { "user_id": user_id, @@ -7511,11 +7517,13 @@ async def login(request: Request): "login_method": "username_password", "premium_user": premium_user, }, - "secret", + litellm_master_key_hash, algorithm="HS256", ) - litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token - return RedirectResponse(url=litellm_dashboard_ui, status_code=303) + 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 elif _user_row is not None: """ When sharing invite links @@ -7564,6 +7572,14 @@ async def login(request: Request): litellm_dashboard_ui += "/ui/" import jwt + if litellm_master_key_hash is None: + raise HTTPException( + status_code=500, + detail={ + "error": "No master key set, please set LITELLM_MASTER_KEY" + }, + ) + jwt_token = jwt.encode( { "user_id": user_id, @@ -7573,11 +7589,15 @@ async def login(request: Request): "login_method": "username_password", "premium_user": premium_user, }, - "secret", + litellm_master_key_hash, algorithm="HS256", ) - litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token - return RedirectResponse(url=litellm_dashboard_ui, status_code=303) + 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 else: raise ProxyException( message=f"Invalid credentials used to access UI. Passed in username: {username}, passed in password: {password}.\nNot valid credentials for {username}", @@ -7688,6 +7708,12 @@ async def onboarding(invite_link: str): litellm_dashboard_ui += "/ui/onboarding" import jwt + if litellm_master_key_hash is None: + raise HTTPException( + status_code=500, + detail={"error": "No master key set, please set LITELLM_MASTER_KEY"}, + ) + jwt_token = jwt.encode( { "user_id": user_obj.user_id, @@ -7697,7 +7723,7 @@ async def onboarding(invite_link: str): "login_method": "username_password", "premium_user": premium_user, }, - "secret", + litellm_master_key_hash, algorithm="HS256", ) @@ -8108,6 +8134,12 @@ async def auth_callback(request: Request): import jwt + if litellm_master_key_hash is None: + raise HTTPException( + status_code=500, + detail={"error": "No master key set, please set LITELLM_MASTER_KEY"}, + ) + jwt_token = jwt.encode( { "user_id": user_id, @@ -8117,11 +8149,13 @@ async def auth_callback(request: Request): "login_method": "sso", "premium_user": premium_user, }, - "secret", + litellm_master_key_hash, algorithm="HS256", ) - litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token - return RedirectResponse(url=litellm_dashboard_ui) + 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 #### INVITATION MANAGEMENT #### diff --git a/ui/litellm-dashboard/src/app/onboarding/page.tsx b/ui/litellm-dashboard/src/app/onboarding/page.tsx index 44c37610cec..a31afb0f9e5 100644 --- a/ui/litellm-dashboard/src/app/onboarding/page.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/page.tsx @@ -20,10 +20,19 @@ import { } from "@/components/networking"; import { jwtDecode } from "jwt-decode"; import { Form, Button as Button2, message } from "antd"; + +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; +} + export default function Onboarding() { const [form] = Form.useForm(); const searchParams = useSearchParams(); - const token = searchParams.get("token"); + const token = getCookie('token'); const inviteID = searchParams.get("id"); const [accessToken, setAccessToken] = useState(null); const [defaultUserEmail, setDefaultUserEmail] = useState(""); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index ed9c98df7c6..41f0e8529b3 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -19,6 +19,15 @@ import CacheDashboard from "@/components/cache_dashboard"; import { jwtDecode } from "jwt-decode"; import { Typography } from "antd"; +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; +} + + function formatUserRole(userRole: string) { if (!userRole) { return "Undefined Role"; @@ -68,7 +77,7 @@ const CreateKeyPage = () => { const searchParams = useSearchParams(); const [modelData, setModelData] = useState({ data: [] }); const userID = searchParams.get("userID"); - const token = searchParams.get("token"); + const token = getCookie('token'); const [page, setPage] = useState("api-keys"); const [accessToken, setAccessToken] = useState(null); diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 5aa0c92dd53..ec66bcea0fa 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -24,6 +24,14 @@ type UserSpendData = { max_budget?: number | null; }; +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; @@ -66,7 +74,8 @@ const UserDashboard: React.FC = ({ const viewSpend = searchParams.get("viewSpend"); const router = useRouter(); - const token = searchParams.get("token"); + const token = getCookie('token'); + const [accessToken, setAccessToken] = useState(null); const [teamSpend, setTeamSpend] = useState(null); const [userModels, setUserModels] = useState([]);