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..9a74c52828 --- /dev/null +++ b/litellm/proxy/management_helpers/ui_session_handler.py @@ -0,0 +1,23 @@ +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, + ) + return redirect_response + + @staticmethod + def _generate_token_name() -> str: + current_timestamp = int(time.time()) + cookie_name = f"token_{current_timestamp}" + return cookie_name