Include team alias in CLI JWT token

This commit is contained in:
Tyler Kalbach 2026-05-19 13:55:16 -04:00
parent a72414a061
commit cb19add560
4 changed files with 40 additions and 2 deletions

View file

@ -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),

View file

@ -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)

View file

@ -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,
):

View file

@ -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()