diff --git a/docs/my-website/docs/proxy/cli_sso.md b/docs/my-website/docs/proxy/cli_sso.md index cde6bf266d4..ad0f033f802 100644 --- a/docs/my-website/docs/proxy/cli_sso.md +++ b/docs/my-website/docs/proxy/cli_sso.md @@ -28,6 +28,37 @@ EXPERIMENTAL_UI_LOGIN="True" litellm --config config.yaml ::: +### Configuration + +#### JWT Token Expiration + +By default, CLI authentication tokens expire after **24 hours**. You can customize this expiration time by setting the `LITELLM_CLI_JWT_EXPIRATION_HOURS` environment variable when starting your LiteLLM Proxy: + +```bash +# Set CLI JWT tokens to expire after 48 hours +export LITELLM_CLI_JWT_EXPIRATION_HOURS=48 +export EXPERIMENTAL_UI_LOGIN="True" +litellm --config config.yaml +``` + +Or in a single command: + +```bash +LITELLM_CLI_JWT_EXPIRATION_HOURS=48 EXPERIMENTAL_UI_LOGIN="True" litellm --config config.yaml +``` + +**Examples:** +- `LITELLM_CLI_JWT_EXPIRATION_HOURS=12` - Tokens expire after 12 hours +- `LITELLM_CLI_JWT_EXPIRATION_HOURS=168` - Tokens expire after 7 days (168 hours) +- `LITELLM_CLI_JWT_EXPIRATION_HOURS=720` - Tokens expire after 30 days (720 hours) + +:::tip +You can check your current token's age and expiration status using: +```bash +litellm-proxy whoami +``` +::: + ### Steps 1. **Install the CLI** diff --git a/litellm/constants.py b/litellm/constants.py index 49ca3a509b1..0525bf3843e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1165,6 +1165,7 @@ LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token" CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session" CLI_JWT_TOKEN_NAME = "cli-jwt-token" +CLI_JWT_EXPIRATION_HOURS = int(os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", 24)) ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job" diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5e0a211906e..61d9044f925 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -21,6 +21,8 @@ from litellm.caching.caching import DualCache from litellm.caching.dual_cache import LimitedSizeOrderedDict from litellm.constants import ( + CLI_JWT_EXPIRATION_HOURS, + CLI_JWT_TOKEN_NAME, DEFAULT_IN_MEMORY_TTL, DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, DEFAULT_MAX_RECURSE_DEPTH, @@ -1602,7 +1604,10 @@ def get_cli_jwt_auth_token( user_info: LiteLLM_UserTable, team_id: Optional[str] = None ) -> str: """ - Generate a JWT token for CLI authentication with 24-hour expiration. + Generate a JWT token for CLI authentication with configurable expiration. + + The expiration time can be controlled via the LITELLM_CLI_JWT_EXPIRATION_HOURS + environment variable (defaults to 24 hours). Args: user_info: User information from the database @@ -1613,7 +1618,6 @@ def get_cli_jwt_auth_token( """ from datetime import timedelta - from litellm.constants import CLI_JWT_TOKEN_NAME from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) @@ -1621,8 +1625,8 @@ def get_cli_jwt_auth_token( if user_info.user_role is None: raise Exception("User role is required for CLI JWT login") - # Calculate expiration time (24 hours from now - matching old CLI key behavior) - expiration_time = get_utc_datetime() + timedelta(hours=24) + # Calculate expiration time (configurable via LITELLM_CLI_JWT_EXPIRATION_HOURS env var) + expiration_time = get_utc_datetime() + timedelta(hours=CLI_JWT_EXPIRATION_HOURS) # Format the expiration time as ISO 8601 string expires = expiration_time.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "+00:00" diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index bdc8d56d1c3..2345b0263e3 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -11,6 +11,8 @@ from rich.console import Console from rich.table import Table +from litellm.constants import CLI_JWT_EXPIRATION_HOURS + # Token storage utilities def get_token_file_path() -> str: @@ -592,8 +594,8 @@ def whoami(): age_hours = (time.time() - timestamp) / 3600 click.echo(f"Token age: {age_hours:.1f} hours") - if age_hours > 24: - click.echo("⚠️ Warning: Token is more than 24 hours old and may have expired.") + if age_hours > CLI_JWT_EXPIRATION_HOURS: + click.echo(f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.") # Export functions for use by other CLI commands diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 5adf54c1627..613390c0ec1 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1156,7 +1156,7 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None): max_budget=litellm.max_ui_session_budget, ) - # Generate CLI JWT on-demand (24hr expiration) + # 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 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 807559207e6..3df0dc881e8 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -15,10 +15,10 @@ import litellm from litellm.proxy._types import ( CallInfo, + Litellm_EntityType, LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable, - Litellm_EntityType, LitellmUserRoles, ProxyErrorTypes, ProxyException, @@ -131,6 +131,60 @@ def test_get_key_object_from_ui_hash_key_invalid(): assert key_object is None +def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values): + """Test generating CLI JWT token with default 24-hour expiration""" + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + + # Decrypt and verify token contents + 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["user_id"] == "test_user" + assert token_data["user_role"] == LitellmUserRoles.PROXY_ADMIN.value + assert token_data["models"] == ["gpt-3.5-turbo"] + assert token_data["max_budget"] == litellm.max_ui_session_budget + + # Verify expiration time is set to 24 hours (default) + assert "expires" in token_data + expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) + assert expires > get_utc_datetime() + assert expires <= get_utc_datetime() + timedelta(hours=24, minutes=1) + assert expires >= get_utc_datetime() + timedelta(hours=23, minutes=59) + + +def test_get_cli_jwt_auth_token_custom_expiration( + valid_sso_user_defined_values, monkeypatch +): + """Test generating CLI JWT token with custom expiration via environment variable""" + # Set custom expiration to 48 hours + monkeypatch.setenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", "48") + + # Reload the constants module to pick up the new env var + import importlib + + from litellm import constants + importlib.reload(constants) + + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + + # Decrypt and verify token contents + 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) + + # Verify expiration time is set to 48 hours + assert "expires" in token_data + expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) + assert expires > get_utc_datetime() + timedelta(hours=47, minutes=59) + assert expires <= get_utc_datetime() + timedelta(hours=48, minutes=1) + + + @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"""