From dfcf1ef1644d804945206da99daaf873803cf52e Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Tue, 11 Aug 2026 16:14:49 +0200 Subject: [PATCH] fix: validate regeneration before deletion history --- .../key_management_endpoints.py | 15 +++-- .../test_key_management_endpoints.py | 60 +++++++++++++++++++ 2 files changed, 69 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 384a5413eb3..2a68a7771e0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4817,12 +4817,6 @@ async def _execute_virtual_key_regeneration( existing_key_row=key_in_db, user_api_key_dict=user_api_key_dict, ) - _check_key_type_transition_against_existing_routes( - data=data, - existing_key_row=key_in_db, - user_api_key_dict=user_api_key_dict, - ) - # Apply the same membership rule used on /key/update: when the caller # asks to point the regenerated key at a different organization_id, # require they are a member of (or proxy admin over) the target org. @@ -5122,6 +5116,15 @@ async def regenerate_key_fn( detail={"error": "You are not authorized to regenerate this key"}, ) + # Validate the requested route transition before recording the old + # key as deleted. A rejected regeneration must not leave historical + # deletion state for a key that remains active. + _check_key_type_transition_against_existing_routes( + data=data, + existing_key_row=_key_in_db, + user_api_key_dict=user_api_key_dict, + ) + if data is not None and (data.access_group_ids or data.object_permission is not None): regenerate_team_table: LiteLLM_TeamTableCachedObj | None = None if _key_in_db.team_id is not 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 bf14f82d114..ef34fed014e 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 @@ -11393,6 +11393,66 @@ async def test_execute_virtual_key_regeneration_rejects_preset_over_custom_route assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 +@pytest.mark.asyncio +async def test_regenerate_rejects_route_transition_before_recording_deletion(): + """A rejected regeneration must not create a deleted-key history row.""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + regenerate_key_fn, + ) + + existing_key = _make_regenerate_existing_key() + existing_key.allowed_routes = ["/custom/admin-defined-route"] + mock_repo = MagicMock() + mock_repo.table.find_unique = AsyncMock(return_value=existing_key) + persist_deleted_mock = AsyncMock() + execute_mock = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: module-global proxy setting for test + patch("litellm.proxy.proxy_server.master_key", None), # test-quality-ok: module-global proxy setting for test + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), # test-quality-ok: module-global prisma client for test + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: module-global proxy logger for test + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: module-global user key cache for test + patch("litellm.proxy.proxy_server.hash_token", lambda token: "hashed-old"), # test-quality-ok: module-global token hasher for test + patch( # test-quality-ok: repository dependency isolation for test + "litellm.proxy.management_endpoints.key_management_endpoints.VerificationTokenRepository", + return_value=mock_repo, + ), + patch( # test-quality-ok: team permission checks isolation for test + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: token modification authorization isolation for test + "litellm.proxy.management_endpoints.key_management_endpoints.can_modify_verification_token", + new_callable=AsyncMock, + return_value=True, + ), + patch( # test-quality-ok: verify deletion persistence is not called before validation + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + persist_deleted_mock, + ), + patch( # test-quality-ok: verify regeneration execution is not called before validation + "litellm.proxy.management_endpoints.key_management_endpoints._execute_virtual_key_regeneration", + execute_mock, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await regenerate_key_fn( + key="sk-old", + data=RegenerateKeyRequest(key_type=LiteLLMKeyType.LLM_API), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user-1", + ), + ) + + assert exc_info.value.status_code == 403 + assert "allowed_routes" in str(exc_info.value.detail) + persist_deleted_mock.assert_not_awaited() + execute_mock.assert_not_awaited() + + @pytest.mark.asyncio async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monkeypatch): """Regenerate must reject durations exceeding upperbound_key_generate_params.duration."""