From ca0aef1a4d8899a1e9d3b2b744329b2b57e68d43 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 5 Mar 2026 16:54:09 -0800 Subject: [PATCH] [Fix] Key Expiry: remove check_db_only, add service account guard, fix uncheck UI bug MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _validate_regenerate_key_duration_against_team: remove check_db_only=True from get_team_object call so cache is used on hits and DB only on misses, consistent with the guideline of avoiding unnecessary direct DB queries. - Add service account guard (user_id is None → skip validation), mirroring the enterprise add_team_member_key_duration guard so service account keys are not subject to team member duration limits during regeneration. - regenerate_key_modal.tsx: add else branch to the Never Expires onChange handler so unchecking resets the form field to undefined instead of leaving it as null (which the backend would interpret as never-expires). - Tests: set mock_key.user_id = "user-123" explicitly in all regeneration tests where validation should fire; add test_service_account_key_skips_validation. Co-Authored-By: Claude Sonnet 4.6 --- .../key_management_endpoints.py | 6 +++- .../test_key_management_endpoints.py | 30 +++++++++++++++++++ .../organisms/regenerate_key_modal.tsx | 2 ++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 45542cba7a3..8b402fcff77 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3408,6 +3408,11 @@ async def _validate_regenerate_key_duration_against_team( if data is None: return + # Service account keys are exempt from team member duration limits, + # mirroring the enterprise add_team_member_key_duration guard. + if getattr(key_in_db, "user_id", None) is None: + return + # Determine the user's requested duration in seconds. # Distinguish "duration not sent" (leave unchanged) from "duration: null" (never expires). if data.duration is None: @@ -3427,7 +3432,6 @@ async def _validate_regenerate_key_duration_against_team( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=None, - check_db_only=True, ) except Exception as e: verbose_proxy_logger.warning( 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 3fb4b1149d6..a1605734d8e 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 @@ -6602,6 +6602,7 @@ class TestValidateRegenerateKeyDurationAgainstTeam: mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.team_id = "team-123" + mock_key.user_id = "user-123" mock_prisma = AsyncMock() mock_cache = MagicMock() # Should not raise @@ -6622,6 +6623,7 @@ class TestValidateRegenerateKeyDurationAgainstTeam: mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.team_id = "team-123" + mock_key.user_id = "user-123" # No duration kwarg at all → not in model_fields_set data = RegenerateKeyRequest() mock_prisma = AsyncMock() @@ -6644,6 +6646,7 @@ class TestValidateRegenerateKeyDurationAgainstTeam: mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.team_id = "team-123" + mock_key.user_id = "user-123" # Explicitly set duration=None → "duration" IS in model_fields_set data = RegenerateKeyRequest(duration=None) assert "duration" in data.model_fields_set @@ -6678,6 +6681,7 @@ class TestValidateRegenerateKeyDurationAgainstTeam: mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.team_id = "team-123" + mock_key.user_id = "user-123" data = RegenerateKeyRequest(duration="999d") mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj) @@ -6708,6 +6712,7 @@ class TestValidateRegenerateKeyDurationAgainstTeam: mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.team_id = None + mock_key.user_id = "user-123" data = RegenerateKeyRequest(duration="10d") mock_prisma = AsyncMock() mock_cache = MagicMock() @@ -6719,6 +6724,28 @@ class TestValidateRegenerateKeyDurationAgainstTeam: user_api_key_cache=mock_cache, ) + @pytest.mark.asyncio + async def test_service_account_key_skips_validation(self): + """Service account keys (user_id=None) are exempt from team duration limits.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_regenerate_key_duration_against_team, + ) + from litellm.proxy._types import RegenerateKeyRequest + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.team_id = "team-123" + mock_key.user_id = None # service account key has no user_id + data = RegenerateKeyRequest(duration="999d") + mock_prisma = AsyncMock() + mock_cache = MagicMock() + # Should not raise even though "999d" would exceed a finite team max + await _validate_regenerate_key_duration_against_team( + data=data, + key_in_db=mock_key, + prisma_client=mock_prisma, + user_api_key_cache=mock_cache, + ) + @pytest.mark.asyncio async def test_duration_within_team_max_passes(self): from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -6728,6 +6755,7 @@ class TestValidateRegenerateKeyDurationAgainstTeam: mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.team_id = "team-123" + mock_key.user_id = "user-123" data = RegenerateKeyRequest(duration="3d") mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj) @@ -6758,6 +6786,7 @@ class TestValidateRegenerateKeyDurationAgainstTeam: mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.team_id = "team-123" + mock_key.user_id = "user-123" data = RegenerateKeyRequest(duration="10d") mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj) @@ -6790,6 +6819,7 @@ class TestValidateRegenerateKeyDurationAgainstTeam: mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.team_id = "team-123" + mock_key.user_id = "user-123" data = RegenerateKeyRequest(duration=None) # null = never expires = infinite mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj) diff --git a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx b/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx index 6142bcffb78..ad01d8e2603 100644 --- a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx +++ b/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx @@ -240,6 +240,8 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat if (checked) { form.setFieldValue("duration", null); setRegenerateFormData((prev: any) => ({ ...prev, duration: null })); + } else { + form.setFieldValue("duration", undefined); } }} >