mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Include team alias in CLI JWT token
This commit is contained in:
parent
a72414a061
commit
cb19add560
4 changed files with 40 additions and 2 deletions
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue