update generate_authenticated_redirect_response

This commit is contained in:
Ishaan Jaff 2025-03-03 17:20:58 -08:00
parent 4f2c6640d5
commit 082bc0c959
2 changed files with 29 additions and 3 deletions

View file

@ -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(

View file

@ -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