From cb19add560faab86d0f26835b9dfa36b96fed47c Mon Sep 17 00:00:00 2001 From: Tyler Kalbach Date: Tue, 19 May 2026 13:55:16 -0400 Subject: [PATCH] Include team alias in CLI JWT token --- litellm/proxy/auth/auth_checks.py | 6 +++++- litellm/proxy/management_endpoints/ui_sso.py | 13 ++++++++++++- .../test_litellm/proxy/auth/test_auth_checks.py | 17 +++++++++++++++++ .../proxy/management_endpoints/test_ui_sso.py | 6 ++++++ 4 files changed, 40 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 13381c7a6c9..f0902fbf178 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2232,7 +2232,9 @@ class ExperimentalUIJWTToken: @staticmethod def get_cli_jwt_auth_token( - user_info: LiteLLM_UserTable, team_id: Optional[str] = None + user_info: LiteLLM_UserTable, + team_id: Optional[str] = None, + team_alias: Optional[str] = None, ) -> str: """ Generate a JWT token for CLI authentication with configurable expiration. @@ -2243,6 +2245,7 @@ class ExperimentalUIJWTToken: Args: user_info: User information from the database team_id: Team ID for the user (optional, uses user's team if available) + team_alias: Team alias for the selected team, if available Returns: Encrypted JWT token string @@ -2276,6 +2279,7 @@ class ExperimentalUIJWTToken: expires=expires, user_id=user_info.user_id, team_id=_team_id, + team_alias=team_alias, models=user_info.models, max_parallel_requests=None, user_role=LitellmUserRoles(user_info.user_role), diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 6e2e2bedac1..f55e2912838 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1889,6 +1889,17 @@ async def cli_poll_key( # If no team_id provided and user has 0 or 1 team, use first team (or None) team_id = user_teams[0] if len(user_teams) > 0 else None + team_alias = None + if team_id and isinstance(user_team_details, list): + team_alias = next( + ( + team.get("team_alias") + for team in user_team_details + if team.get("team_id") == team_id + ), + None, + ) + # Create user object for JWT generation user_info = LiteLLM_UserTable( user_id=user_id, @@ -1900,7 +1911,7 @@ async def cli_poll_key( # Generate CLI JWT on-demand (expiration configurable via LITELLM_CLI_JWT_EXPIRATION_HOURS) # Pass selected team_id to ensure JWT has correct team jwt_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info=user_info, team_id=team_id + user_info=user_info, team_id=team_id, team_alias=team_alias ) # Delete cache entry (single-use) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26f04a4abcb..d3b9c1bad2d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -127,6 +127,23 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v assert expires <= now + timedelta(minutes=10, seconds=2) +def test_get_cli_jwt_auth_token_includes_team_alias(valid_sso_user_defined_values): + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( + valid_sso_user_defined_values, + team_id="team-123", + team_alias="test-team", + ) + + decrypted_token = decrypt_value_helper( + token, key="ui_hash_key", exception_type="debug" + ) + assert decrypted_token is not None + token_data = json.loads(decrypted_token) + + assert token_data["team_id"] == "team-123" + assert token_data["team_alias"] == "test-team" + + def test_get_experimental_ui_login_jwt_auth_token_uses_10_min_expiry( valid_sso_user_defined_values, ): diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 83317157847..1dc75a43a12 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2488,6 +2488,11 @@ class TestCLIKeyRegenerationFlow: "user_id": "test-user-789", "user_role": "internal_user", "teams": ["team-a", "team-b", "team-c"], + "team_details": [ + {"team_id": "team-a", "team_alias": "Team A"}, + {"team_id": "team-b", "team_alias": "Team B"}, + {"team_id": "team-c", "team_alias": "Team C"}, + ], "models": ["gpt-4"], "user_email": "test@example.com", } @@ -2542,6 +2547,7 @@ class TestCLIKeyRegenerationFlow: mock_get_jwt.assert_called_once() jwt_call_args = mock_get_jwt.call_args assert jwt_call_args.kwargs["team_id"] == selected_team + assert jwt_call_args.kwargs["team_alias"] == "Team B" # Verify session was deleted after JWT generation mock_cache.delete_cache.assert_called_once()