(UI) - Security Improvement, move to JWT Auth for Admin UI Sessions (#8995)

* (UI) - Improvements to session handling logic  (#8970)

* add cookieUtils

* use utils for clearing cookies

* on logout use clearTokenCookies

* ui use correct clearTokenCookies

* navbar show userEmail on UserID page

* add timestamp on token cookie

* update generate_authenticated_redirect_response

* use common getAuthToken

* fix clearTokenCookies

* fixes for get auth token

* fix invitation link sign in logic

* Revert "fix invitation link sign in logic"

This reverts commit 30e5308cb3.

* fix getAuthToken

* update setAuthToken

* fix ui session handling

* fix ui session handler

* bug fix stop generating LiteLLM Virtual keys for access

* working JWT insert into cookies

* use central place to build UI JWT token

* add _validate_ui_token

* fix ui session handler

* fix fetchWithCredentials

* check allowed routes for ui session tokens

* expose validate_session endpoint

* validate session endpoint

* call sso/session/validate

* getUISessionDetails

* ui move to getUISessionDetails

* /sso/session/validate

* fix cookie utils

* use getUISessionDetails

* use ui_session_id

* "/spend/logs/ui" in spend_tracking_routes

* working sign in JWT flow for proxy admin

* allow proxy admin to access ui routes

* use check_route_access

* update types

* update login method

* fixes to ui session handler

* working flow for admin and internal users

* fixes for invite links

* use JWTs for SSO sign in

* fix /invitation/new flow

* fix code quality checks

* fix _get_ui_session_token_from_cookies

* /organization/list

* ui sso sign in

* TestUISessionHandler

* TestUISessionHandler

litellm/proxy/management_helpers/ui_session_handler.py

@@ -0,0 +1,141 @@
+import time
+from datetime import datetime, timedelta, timezone
+from typing import Literal, Optional
+
+from fastapi.requests import Request
+from fastapi.responses import RedirectResponse
+
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import LiteLLM_JWTAuth, LitellmUserRoles
+
+
+class UISessionHandler:
+
+    @staticmethod
+    def _get_latest_ui_cookie_name(cookies: dict) -> Optional[str]:
+        """
+        Get the name of the most recent LiteLLM UI cookie
+        """
+        # Find all LiteLLM UI cookies (format: litellm_ui_token_{timestamp})
+        litellm_ui_cookies = [
+            k for k in cookies.keys() if k.startswith("litellm_ui_token_")
+        ]
+        if not litellm_ui_cookies:
+            return None
+
+        # Sort by timestamp (descending) to get the most recent one
+        try:
+            # Extract timestamps and sort numerically
+            sorted_cookies = sorted(
+                litellm_ui_cookies,
+                key=lambda x: int(x.split("_")[-1]),
+                reverse=True,
+            )
+            return sorted_cookies[0]
+        except (ValueError, IndexError):
+            # Fallback to simple string sort if timestamp extraction fails
+            litellm_ui_cookies.sort(reverse=True)
+            return litellm_ui_cookies[0]
+
+    @staticmethod
+    # Add this function to extract auth token from cookies
+    def _get_ui_session_token_from_cookies(request: Request) -> Optional[str]:
+        """
+        Extract authentication token from cookies if present
+        """
+        try:
+            cookies = request.cookies
+            verbose_proxy_logger.debug(f"AUTH COOKIES: {cookies}")
+
+            cookie_name = UISessionHandler._get_latest_ui_cookie_name(cookies)
+            if cookie_name:
+                return cookies[cookie_name]
+
+            return None
+        except Exception as e:
+            verbose_proxy_logger.error(
+                f"Error getting UI session token from cookies: {e}"
+            )
+            return None
+
+    @staticmethod
+    def build_authenticated_ui_jwt_token(
+        user_id: str,
+        user_role: Optional[LitellmUserRoles],
+        user_email: Optional[str],
+        premium_user: bool,
+        disabled_non_admin_personal_key_creation: bool,
+        login_method: Literal["username_password", "sso"],
+    ) -> str:
+        """
+        Build a JWT token for the authenticated UI session
+
+        This token is used to authenticate the user's session when they are redirected to the UI
+        """
+        import jwt
+
+        from litellm.proxy.proxy_server import general_settings, master_key
+
+        if master_key is None:
+            raise ValueError("Master key is not set")
+
+        expiration = datetime.now(timezone.utc) + timedelta(hours=24)
+        initial_payload = {
+            "user_id": user_id,
+            "user_email": user_email,
+            "user_role": user_role,  # this is the path without sso - we can assume only admins will use this
+            "login_method": login_method,
+            "premium_user": premium_user,
+            "auth_header_name": general_settings.get(
+                "litellm_key_header_name", "Authorization"
+            ),
+            "iss": "litellm-proxy",  # Issuer - identifies this as an internal token
+            "aud": "litellm-ui",  # Audience - identifies this as a UI token
+            "exp": expiration,
+            "disabled_non_admin_personal_key_creation": disabled_non_admin_personal_key_creation,
+        }
+
+        if (
+            user_role == LitellmUserRoles.PROXY_ADMIN
+            or user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
+        ):
+            initial_payload["scope"] = [
+                LiteLLM_JWTAuth().admin_jwt_scope,
+            ]
+
+        jwt_token = jwt.encode(
+            initial_payload,
+            master_key,
+            algorithm="HS256",
+        )
+        return jwt_token
+
+    @staticmethod
+    def is_ui_session_token(token_dict: dict) -> bool:
+        """
+        Returns True if the token is a UI session token
+        """
+        return (
+            token_dict.get("iss") == "litellm-proxy"
+            and token_dict.get("aud") == "litellm-ui"
+        )
+
+    @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,
+            httponly=True,
+            samesite="strict",
+        )
+        return redirect_response
+
+    @staticmethod
+    def _generate_token_name() -> str:
+        current_timestamp = int(time.time())
+        cookie_name = f"litellm_ui_token_{current_timestamp}"
+        return cookie_name