From 92ed63092c6bdc225c47cca829ca8de7c5250b6e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 24 Jun 2026 22:58:45 +0530 Subject: [PATCH] fix(auth): set max_budget on CLI session token to enforce max_ui_session_budget CLI session tokens were missing max_budget, so _virtual_key_max_budget_check had no per-session ceiling to enforce. Operators relying on max_ui_session_budget could be bypassed for the full token lifetime. Mirrors the existing UI token path. --- litellm/proxy/auth/auth_checks.py | 4 +--- tests/test_litellm/proxy/auth/test_auth_checks.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 267d87965e4..b7f8a916b30 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2454,9 +2454,6 @@ class ExperimentalUIJWTToken: # Use first team if user has teams _team_id = user_info.teams[0] if len(user_info.teams) > 0 else None - # Unique per-login identity so concurrent sessions get isolated spend - # attribution; budget is enforced via the shared team/user counters - # (team_id/user_id), not a per-key max_budget. session_token = f"{CLI_SESSION_KEY_PREFIX}-{secrets.token_urlsafe(16)}" session_alias = f"{CLI_SESSION_KEY_PREFIX}-{user_info.user_id}" @@ -2465,6 +2462,7 @@ class ExperimentalUIJWTToken: key_name=session_alias, key_alias=session_alias, expires=expires, + max_budget=litellm.max_ui_session_budget, user_id=user_info.user_id, team_id=_team_id, team_alias=team_alias, diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 6d343af5b15..4721bf346c9 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -538,6 +538,18 @@ def test_get_cli_jwt_auth_token_unique_per_session(valid_sso_user_defined_values assert first["key_name"] == second["key_name"] == expected_alias +def test_get_cli_jwt_auth_token_enforces_max_ui_session_budget( + valid_sso_user_defined_values, +): + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + + decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + assert decrypted is not None + token_data = json.loads(decrypted) + + assert token_data.get("max_budget") == litellm.max_ui_session_budget + + @pytest.mark.asyncio async def test_default_internal_user_params_with_get_user_object(monkeypatch): """Test that default_internal_user_params is used when creating a new user via get_user_object"""