From c7345e4bfd3f34baeb3bb0d3d51319bc868ca8f5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 5 Jan 2026 12:24:34 -0800 Subject: [PATCH] Use timedelta to calculate key expiry --- .../key_management_endpoints.py | 4 +- .../test_key_management_endpoints.py | 111 +++++++++++++++--- 2 files changed, 98 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4d72bc8625..8d45493bd9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2097,7 +2097,9 @@ async def generate_key_helper_fn( # noqa: PLR0915 if duration is None: # allow tokens that never expire expires = None else: - expires = get_budget_reset_time(budget_duration=duration) + # Add duration to current time for exact expiration (not standardized reset time) + duration_seconds = duration_in_seconds(duration) + expires = datetime.now(timezone.utc) + timedelta(seconds=duration_seconds) if key_budget_duration is None: # one-time budget key_reset_at = None diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 3e69b4e0ca..33fa7fc7bd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -279,7 +279,9 @@ async def test_key_token_handling(monkeypatch): @pytest.mark.asyncio async def test_budget_reset_and_expires_at_first_of_month(monkeypatch): """ - Test that when budget_duration, duration, and key_budget_duration are "1mo", budget_reset_at and expires are set to first of next month + Test that when budget_duration, duration, and key_budget_duration are "1mo": + - budget_reset_at is set to first of next month (standardized reset time) + - expires is set to approximately 1 month from creation time (exact duration) """ mock_prisma_client = AsyncMock() mock_insert_data = AsyncMock( @@ -299,7 +301,7 @@ async def test_budget_reset_and_expires_at_first_of_month(monkeypatch): return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None) ) - from datetime import datetime, timezone + from datetime import datetime, timedelta, timezone import pytest @@ -324,7 +326,7 @@ async def test_budget_reset_and_expires_at_first_of_month(monkeypatch): # Get the current date now = datetime.now(timezone.utc) - # Calculate expected reset date (first of next month) + # Calculate expected reset date (first of next month) for budget_reset_at if now.month == 12: expected_month = 1 expected_year = now.year + 1 @@ -332,19 +334,96 @@ async def test_budget_reset_and_expires_at_first_of_month(monkeypatch): expected_month = now.month + 1 expected_year = now.year - # Verify budget_reset_at, expires is set to first of next month - for key in ["budget_reset_at", "expires"]: - response_date = response.get(key) - assert response_date is not None, f"{key} not found in response" - assert ( - response_date.year == expected_year - ), f"Expected year {expected_year}, got {response_date.year} for {key}" - assert ( - response_date.month == expected_month - ), f"Expected month {expected_month}, got {response_date.month} for {key}" - assert ( - response_date.day == 1 - ), f"Expected day 1, got {response_date.day} for {key}" + # Verify budget_reset_at is set to first of next month (standardized reset time) + budget_reset_at = response.get("budget_reset_at") + assert budget_reset_at is not None, "budget_reset_at not found in response" + assert ( + budget_reset_at.year == expected_year + ), f"Expected year {expected_year}, got {budget_reset_at.year} for budget_reset_at" + assert ( + budget_reset_at.month == expected_month + ), f"Expected month {expected_month}, got {budget_reset_at.month} for budget_reset_at" + assert ( + budget_reset_at.day == 1 + ), f"Expected day 1, got {budget_reset_at.day} for budget_reset_at" + + # Verify expires is set to approximately 1 month from creation time (exact duration, not standardized) + expires = response.get("expires") + assert expires is not None, "expires not found in response" + # expires should be approximately 1 month from now (same day next month, same time) + # Allow for some variance due to test execution time + expected_expires_min = now + timedelta(days=28) + expected_expires_max = now + timedelta(days=32) + assert ( + expected_expires_min <= expires <= expected_expires_max + ), f"Expected expires to be approximately 1 month from now, got {expires}" + + +@pytest.mark.asyncio +async def test_key_expiration_exact_duration_hours(monkeypatch): + """ + Test that key expiration uses exact duration addition, not standardized reset times. + Specifically tests the bug where "12h" duration would expire at midnight instead of 12 hours from creation. + """ + mock_prisma_client = AsyncMock() + mock_insert_data = AsyncMock( + return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None) + ) + mock_prisma_client.insert_data = mock_insert_data + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[] + ) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None) + ) + + from datetime import datetime, timedelta, timezone + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_helper_fn, + ) + + # Use monkeypatch to set the prisma_client + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Test key generation with duration="12h" + # This should expire exactly 12 hours from creation, not at the next midnight/noon boundary + response = await generate_key_helper_fn( + request_type="user", + duration="12h", + user_id="test_user", + ) + + expires = response.get("expires") + assert expires is not None, "expires not found in response" + + # Calculate expected expiration (approximately 12 hours from now) + # Allow for small variance due to test execution time + now = datetime.now(timezone.utc) + expected_expires_min = now + timedelta(hours=11, minutes=59) + expected_expires_max = now + timedelta(hours=12, minutes=1) + + assert ( + expected_expires_min <= expires <= expected_expires_max + ), f"Expected expires to be approximately 12 hours from now ({now}), got {expires}. Duration should be exact, not aligned to time boundaries." + + # Verify it's NOT aligned to hour boundaries (e.g., not exactly at :00 minutes) + # If created at 2:30 PM, it should expire at 2:30 AM, not midnight + expires_minute = expires.minute + expires_second = expires.second + # If the expiration is exactly at :00:00, it might be aligned (though could be coincidence) + # More importantly, verify the duration is correct + time_diff = expires - now + hours_diff = time_diff.total_seconds() / 3600 + assert ( + 11.9 <= hours_diff <= 12.1 + ), f"Expected expiration to be approximately 12 hours from creation, got {hours_diff} hours" @pytest.mark.asyncio