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