From 22a7174e3c88dc5658277eb64328d00431bc95b2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 02:21:14 -0700 Subject: [PATCH] fix(ui): stamp exp on the UI session cookie so bounded-lifetime readers accept it --- litellm/proxy/auth/login_utils.py | 26 +++++++++ litellm/proxy/management_endpoints/ui_sso.py | 8 +-- litellm/proxy/proxy_server.py | 30 ++-------- .../proxy/auth/test_login_utils.py | 55 +++++++++++++++++++ 4 files changed, 90 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 11f12e597b9..f35d94c986e 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -7,12 +7,15 @@ login endpoints (e.g., /login and /v2/login). import os import secrets +from datetime import datetime, timedelta, timezone from typing import Literal, Optional, cast +import jwt from fastapi import HTTPException import litellm from litellm.constants import LITELLM_PROXY_ADMIN_NAME, LITELLM_UI_SESSION_DURATION +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( LiteLLM_UserTable, LitellmUserRoles, @@ -313,6 +316,29 @@ async def authenticate_user( ) +def _ui_session_exp_timestamp() -> int: + """The ``exp`` claim (unix seconds) for a UI session cookie, ``LITELLM_UI_SESSION_DURATION`` + from now. The virtual key sealed inside the cookie already expires after this same + duration; stamping the JWT itself gives the cookie the bounded lifetime the dashboard's + client-side expiry check and the server-side session-cookie readers both assume, instead + of a token that stays signature-valid until the master key rotates.""" + ttl_seconds = duration_in_seconds(LITELLM_UI_SESSION_DURATION) + return int((datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds)).timestamp()) + + +def encode_ui_session_jwt(returned_ui_token_object: ReturnedUITokenObject, master_key: str) -> str: + """Encode a UI session cookie JWT with a bounded ``exp``. + + The single choke point every UI login path (SSO and username/password /login, /v2, + /v3) uses to mint the ``token`` cookie, so the cookie's lifetime is set in exactly one + place and cannot drift between paths. Without the ``exp`` the cookie is valid until the + master key rotates, and the session-cookie readers that require a bounded lifetime + (the MCP interactive sign-in) reject it. + """ + claims = {**cast(dict, returned_ui_token_object), "exp": _ui_session_exp_timestamp()} + return jwt.encode(claims, master_key, algorithm="HS256") + + def create_ui_token_object( login_result: LoginResult, general_settings: dict, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 0475566192e..fe6682e4221 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -3216,11 +3216,9 @@ class SSOAuthenticationHandler: server_root_path=get_server_root_path(), ) - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - master_key or "", - algorithm="HS256", - ) + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + jwt_token = encode_ui_session_jwt(returned_ui_token_object, master_key or "") # Control-plane cross-origin: store JWT behind a single-use opaque # code (60s TTL) so the token never appears in browser history / logs. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0eb42b266e..0dbb3cf94b2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13339,7 +13339,7 @@ async def fallback_login(request: Request): @router.post("/login", include_in_schema=False) # hidden since this is a helper for UI sso login async def login(request: Request): global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt from litellm.proxy.utils import get_custom_url form = await request.form() @@ -13362,13 +13362,7 @@ async def login(request: Request): ) # Generate JWT token - import jwt - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) # Build redirect URL litellm_dashboard_ui = get_custom_url(str(request.base_url)) @@ -13387,7 +13381,7 @@ async def login(request: Request): @router.post("/v2/login", include_in_schema=False) # hidden helper for UI logins via API async def login_v2(request: Request): global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt from litellm.proxy.utils import get_custom_url try: @@ -13408,13 +13402,7 @@ async def login_v2(request: Request): premium_user=premium_user, ) - import jwt - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): @@ -13458,7 +13446,7 @@ async def login_v2(request: Request): ) # control-plane login — always returns token in body for cross-origin use async def login_v3(request: Request): global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt from litellm.proxy.utils import get_custom_url try: @@ -13487,13 +13475,7 @@ async def login_v3(request: Request): premium_user=premium_user, ) - import jwt - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 288e2533b72..e301e878bf4 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -559,3 +559,58 @@ async def test_authenticate_user_database_login_with_non_ascii_password(): assert isinstance(result, LoginResult) assert result.user_id == "test-user-123" assert result.user_email == user_email + + +class TestEncodeUiSessionJwt: + """The UI session cookie must carry a bounded exp so it does not stay + signature-valid until the master key rotates, and so the session-cookie readers + that require a bounded lifetime (the MCP interactive sign-in) accept it.""" + + def _decode(self, token: str) -> dict: + import jwt + + return jwt.decode(token, "sk-master-for-tests", algorithms=["HS256"]) + + def test_encoded_cookie_carries_bounded_exp(self): + import time + + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + token_object = {"user_id": "u1", "key": "sk-abc", "login_method": "username_password"} + with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "24h"): + token = encode_ui_session_jwt(token_object, "sk-master-for-tests") + claims = self._decode(token) + assert claims["user_id"] == "u1" + assert claims["login_method"] == "username_password" + remaining = claims["exp"] - int(time.time()) + assert 23 * 3600 < remaining <= 24 * 3600 + + def test_duration_is_honored_from_env(self): + import time + + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "1h"): + token = encode_ui_session_jwt({"user_id": "u1"}, "sk-master-for-tests") + remaining = self._decode(token)["exp"] - int(time.time()) + assert 0 < remaining <= 3600 + + def test_cookie_is_accepted_by_the_exp_requiring_session_reader(self): + """The regression this change exists for: before it, the UI cookie carried no + exp and _user_id_from_session_cookie (require=["exp"]) rejected every real login, + so the MCP interactive sign-in could never capture identity. A cookie minted by + this helper must now be accepted.""" + from unittest.mock import MagicMock + + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + _user_id_from_session_cookie, + ) + from litellm.proxy.auth.login_utils import encode_ui_session_jwt + + token_object = {"user_id": "cornell-user", "key": "sk-abc", "login_method": "sso"} + with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "24h"): + token = encode_ui_session_jwt(token_object, "sk-master-for-tests") + request = MagicMock() + request.cookies = {"token": token} + with patch("litellm.proxy.proxy_server.master_key", "sk-master-for-tests"): + assert _user_id_from_session_cookie(request) == "cornell-user"