diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 205aeca328d..7e04b1afe90 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -45,6 +45,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router +from litellm.utils import get_utc_datetime from .auth_checks_organization import organization_role_based_access_check @@ -957,7 +958,7 @@ async def get_team_object( class ExperimentalUIJWTToken: @staticmethod def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: - from datetime import UTC, datetime, timedelta + from datetime import timedelta from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, @@ -967,7 +968,7 @@ class ExperimentalUIJWTToken: raise Exception("User role is required for experimental UI login") # Calculate expiration time (10 minutes from now) - expiration_time = datetime.now(UTC) + timedelta(minutes=10) + expiration_time = get_utc_datetime() + timedelta(minutes=10) # 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/tests/litellm/proxy/auth/test_auth_checks.py b/tests/litellm/proxy/auth/test_auth_checks.py index 7e8d99d4eed..24e8506e971 100644 --- a/tests/litellm/proxy/auth/test_auth_checks.py +++ b/tests/litellm/proxy/auth/test_auth_checks.py @@ -8,7 +8,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from datetime import UTC, datetime, timedelta +from datetime import datetime, timedelta import pytest @@ -20,6 +20,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.utils import get_utc_datetime @pytest.fixture(autouse=True) @@ -68,8 +69,8 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v # Verify expiration time is set and valid assert "expires" in token_data expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) - assert expires > datetime.now(UTC) - assert expires <= datetime.now(UTC) + timedelta(minutes=10) + assert expires > get_utc_datetime() + assert expires <= get_utc_datetime() + timedelta(minutes=10) def test_get_experimental_ui_login_jwt_auth_token_invalid(