From fef5c5fd3ef3c6dc5376749fac226df59de55a87 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 7 Feb 2026 17:08:01 -0800 Subject: [PATCH 1/5] fix: preserve key alias and team_id after key regeneration, deletion --- .../common_daily_activity.py | 35 ++- .../key_management_endpoints.py | 16 +- .../test_common_daily_activity.py | 254 ++++++++++++++++++ 3 files changed, 301 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 99a732f9efb..79d80abf272 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -327,14 +327,43 @@ async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: Set[str], ) -> Dict[str, Dict[str, Any]]: - """Update api key metadata for a single record.""" + """Get api key metadata, falling back to deleted keys table for keys not found in active table. + + This ensures that key_alias and team_id are preserved in historical activity logs + even after a key is deleted or regenerated. + """ key_records = await prisma_client.db.litellm_verificationtoken.find_many( where={"token": {"in": list(api_keys)}} ) - return { - k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records + result = { + k.token: {"key_alias": k.key_alias, "team_id": k.team_id} + for k in key_records } + # For any keys not found in the active table, check the deleted keys table + missing_keys = api_keys - set(result.keys()) + if missing_keys: + try: + deleted_key_records = ( + await prisma_client.db.litellm_deletedverificationtoken.find_many( + where={"token": {"in": list(missing_keys)}}, + order={"deleted_at": "desc"}, + ) + ) + # Use the most recent deleted record for each token (ordered by deleted_at desc) + for k in deleted_key_records: + if k.token not in result: + result[k.token] = { + "key_alias": k.key_alias, + "team_id": k.team_id, + } + except Exception: + verbose_proxy_logger.debug( + "Failed to fetch deleted key metadata for missing keys" + ) + + return result + def _adjust_dates_for_timezone( start_date: str, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2eb6cf65281..264dff5aa54 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3313,6 +3313,20 @@ async def regenerate_key_fn( verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) + # Save the old key record to deleted table before regeneration + # This preserves key_alias and team_id metadata for historical spend records + try: + await _persist_deleted_verification_tokens( + keys=[_key_in_db], + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + except Exception: + verbose_proxy_logger.debug( + "Failed to persist old key record to deleted table during regeneration" + ) + new_token = get_new_token(data=data) new_token_hash = hash_token(new_token) @@ -3749,7 +3763,7 @@ async def list_keys( else: admin_team_ids = None - if not user_id and user_api_key_dict.user_role not in [ + if user_id is None and user_api_key_dict.user_role not in [ LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, ]: diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 93457631d2d..48869803b20 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -11,6 +11,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( _is_user_agent_tag, compute_tag_metadata_totals, + get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, ) @@ -208,3 +209,256 @@ def __init__(self, date, endpoint, api_key, model, spend, prompt_tokens, complet assert chat_endpoint.api_key_breakdown["key-1"].metrics.spend == 15.0 assert "key-2" in embeddings_endpoint.api_key_breakdown assert embeddings_endpoint.api_key_breakdown["key-2"].metrics.spend == 3.0 + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_returns_active_key_metadata(): + """Test that get_api_key_metadata should return metadata for active keys.""" + mock_prisma = MagicMock() + + # Mock active key record + mock_active_key = MagicMock() + mock_active_key.token = "active-key-hash-123" + mock_active_key.key_alias = "my-active-key" + mock_active_key.team_id = "team-abc" + + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[mock_active_key] + ) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={"active-key-hash-123"}, + ) + + assert "active-key-hash-123" in result + assert result["active-key-hash-123"]["key_alias"] == "my-active-key" + assert result["active-key-hash-123"]["team_id"] == "team-abc" + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_falls_back_to_deleted_keys(): + """Test that get_api_key_metadata should fall back to deleted keys table for missing keys.""" + mock_prisma = MagicMock() + + # No active keys found + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + # Deleted key record exists + mock_deleted_key = MagicMock() + mock_deleted_key.token = "deleted-key-hash-456" + mock_deleted_key.key_alias = "toto-test-2" + mock_deleted_key.team_id = "team-xyz" + + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( + return_value=[mock_deleted_key] + ) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={"deleted-key-hash-456"}, + ) + + assert "deleted-key-hash-456" in result + assert result["deleted-key-hash-456"]["key_alias"] == "toto-test-2" + assert result["deleted-key-hash-456"]["team_id"] == "team-xyz" + + # Verify deleted table was queried with the missing key + mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_called_once_with( + where={"token": {"in": ["deleted-key-hash-456"]}}, + order={"deleted_at": "desc"}, + ) + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_mixed_active_and_deleted_keys(): + """Test that get_api_key_metadata should return metadata for both active and deleted keys.""" + mock_prisma = MagicMock() + + # One active key found + mock_active_key = MagicMock() + mock_active_key.token = "active-key-hash" + mock_active_key.key_alias = "active-alias" + mock_active_key.team_id = "team-active" + + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[mock_active_key] + ) + + # One deleted key found + mock_deleted_key = MagicMock() + mock_deleted_key.token = "deleted-key-hash" + mock_deleted_key.key_alias = "deleted-alias" + mock_deleted_key.team_id = "team-deleted" + + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( + return_value=[mock_deleted_key] + ) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={"active-key-hash", "deleted-key-hash"}, + ) + + # Both keys should have metadata + assert len(result) == 2 + assert result["active-key-hash"]["key_alias"] == "active-alias" + assert result["active-key-hash"]["team_id"] == "team-active" + assert result["deleted-key-hash"]["key_alias"] == "deleted-alias" + assert result["deleted-key-hash"]["team_id"] == "team-deleted" + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_deleted_table_not_queried_when_all_keys_found(): + """Test that get_api_key_metadata should not query deleted table when all keys are active.""" + mock_prisma = MagicMock() + + mock_active_key = MagicMock() + mock_active_key.token = "key-hash-1" + mock_active_key.key_alias = "alias-1" + mock_active_key.team_id = "team-1" + + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[mock_active_key] + ) + mock_prisma.db.litellm_deletedverificationtoken = MagicMock() + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( + return_value=[] + ) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={"key-hash-1"}, + ) + + assert len(result) == 1 + assert result["key-hash-1"]["key_alias"] == "alias-1" + # Deleted table should NOT have been queried + mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_deleted_table_error_handled_gracefully(): + """Test that get_api_key_metadata should handle errors from deleted table gracefully.""" + mock_prisma = MagicMock() + + # No active keys found + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + # Deleted table raises an error (e.g., table doesn't exist in older schema) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( + side_effect=Exception("Table not found") + ) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={"missing-key-hash"}, + ) + + # Should return empty dict without raising + assert result == {} + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_regenerated_key_uses_most_recent_deleted_record(): + """Test that get_api_key_metadata should use the most recent deleted record for regenerated keys.""" + mock_prisma = MagicMock() + + # No active keys found (old hash no longer in active table after regeneration) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + # Multiple deleted records for same token (e.g., regenerated multiple times) + mock_deleted_1 = MagicMock() + mock_deleted_1.token = "old-key-hash" + mock_deleted_1.key_alias = "latest-alias" + mock_deleted_1.team_id = "latest-team" + + mock_deleted_2 = MagicMock() + mock_deleted_2.token = "old-key-hash" + mock_deleted_2.key_alias = "older-alias" + mock_deleted_2.team_id = "older-team" + + # Ordered by deleted_at desc, so first record is the most recent + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( + return_value=[mock_deleted_1, mock_deleted_2] + ) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={"old-key-hash"}, + ) + + # Should use the first (most recent) record + assert result["old-key-hash"]["key_alias"] == "latest-alias" + assert result["old-key-hash"]["team_id"] == "latest-team" + + +@pytest.mark.asyncio +async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): + """Test that the full aggregation pipeline should preserve metadata for deleted keys.""" + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + + class MockRecord: + def __init__(self, date, endpoint, api_key, model, spend, prompt_tokens, completion_tokens): + self.date = date + self.endpoint = endpoint + self.api_key = api_key + self.model = model + self.model_group = None + self.custom_llm_provider = "openai" + self.mcp_namespaced_tool_name = None + self.spend = spend + self.prompt_tokens = prompt_tokens + self.completion_tokens = completion_tokens + self.total_tokens = prompt_tokens + completion_tokens + self.cache_read_input_tokens = 0 + self.cache_creation_input_tokens = 0 + self.api_requests = 1 + self.successful_requests = 1 + self.failed_requests = 0 + + # Records reference a deleted key + mock_records = [ + MockRecord("2024-01-01", "/v1/chat/completions", "deleted-key-hash", "gpt-4", 10.0, 100, 50), + ] + + mock_table = MagicMock() + mock_table.find_many = AsyncMock(return_value=mock_records) + mock_prisma.db.litellm_dailyuserspend = mock_table + + # Active table returns nothing for this key + mock_prisma.db.litellm_verificationtoken = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + # Deleted table returns the metadata + mock_deleted_key = MagicMock() + mock_deleted_key.token = "deleted-key-hash" + mock_deleted_key.key_alias = "toto-test-2" + mock_deleted_key.team_id = "69cd4b77-b095-4489-8c46-4f2f31d840a2" + + mock_prisma.db.litellm_deletedverificationtoken = MagicMock() + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( + return_value=[mock_deleted_key] + ) + + result = await get_daily_activity_aggregated( + prisma_client=mock_prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + ) + + # Verify the deleted key's metadata is preserved + daily_data = result.results[0] + chat_endpoint = daily_data.breakdown.endpoints["/v1/chat/completions"] + assert "deleted-key-hash" in chat_endpoint.api_key_breakdown + key_data = chat_endpoint.api_key_breakdown["deleted-key-hash"] + assert key_data.metadata.key_alias == "toto-test-2" + assert key_data.metadata.team_id == "69cd4b77-b095-4489-8c46-4f2f31d840a2" + assert key_data.metrics.spend == 10.0 From fa9585371890f485c8c83f672fed8179b945b103 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 7 Feb 2026 17:17:23 -0800 Subject: [PATCH 2/5] resolved greptile issue related to regeneration persistence and exception handling --- .../common_daily_activity.py | 8 ++++--- .../key_management_endpoints.py | 23 ++++++++----------- 2 files changed, 15 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 79d80abf272..e5df2f82f69 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -357,9 +357,11 @@ async def get_api_key_metadata( "key_alias": k.key_alias, "team_id": k.team_id, } - except Exception: - verbose_proxy_logger.debug( - "Failed to fetch deleted key metadata for missing keys" + except Exception as e: + verbose_proxy_logger.warning( + "Failed to fetch deleted key metadata for %d missing keys: %s", + len(missing_keys), + e, ) return result diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 264dff5aa54..48e0c32880b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3313,19 +3313,16 @@ async def regenerate_key_fn( verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) - # Save the old key record to deleted table before regeneration - # This preserves key_alias and team_id metadata for historical spend records - try: - await _persist_deleted_verification_tokens( - keys=[_key_in_db], - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - except Exception: - verbose_proxy_logger.debug( - "Failed to persist old key record to deleted table during regeneration" - ) + # Save the old key record to deleted table before regeneration. + # This preserves key_alias and team_id metadata for historical spend records. + # If this fails, abort the regeneration to avoid permanently losing the + # old hash→metadata mapping. + await _persist_deleted_verification_tokens( + keys=[_key_in_db], + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) new_token = get_new_token(data=data) From ef25ecb68eaf1509b45e4a20e922498e9540df19 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 7 Feb 2026 18:21:33 -0800 Subject: [PATCH 3/5] fixed user id --- litellm/proxy/management_endpoints/key_management_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 48e0c32880b..099c4543362 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3760,7 +3760,7 @@ async def list_keys( else: admin_team_ids = None - if user_id is None and user_api_key_dict.user_role not in [ + if not user_id and user_api_key_dict.user_role not in [ LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, ]: From 6f49261651e14e1d39882f449cd5f46578551265 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 14 Feb 2026 16:36:01 -0800 Subject: [PATCH 4/5] fix failing test and lint --- .../proxy/management_endpoints/key_management_endpoints.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 099c4543362..9cc0c67de93 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3160,7 +3160,7 @@ def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: dependencies=[Depends(user_api_key_auth)], ) @management_endpoint_wrapper -async def regenerate_key_fn( +async def regenerate_key_fn( # noqa: PLR0915 key: Optional[str] = None, data: Optional[RegenerateKeyRequest] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -3313,6 +3313,10 @@ async def regenerate_key_fn( verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) + # Normalize litellm_changed_by: if it's a Header object or not a string, convert to None + if litellm_changed_by is not None and not isinstance(litellm_changed_by, str): + litellm_changed_by = None + # Save the old key record to deleted table before regeneration. # This preserves key_alias and team_id metadata for historical spend records. # If this fails, abort the regeneration to avoid permanently losing the From 88a631d32abc3af2e55d86735ef66f8bb841de1f Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 14 Feb 2026 16:48:46 -0800 Subject: [PATCH 5/5] fixed ruff --- .../key_management_endpoints.py | 128 ++++++++++-------- 1 file changed, 69 insertions(+), 59 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 9cc0c67de93..1c110f6982b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -68,6 +68,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key from litellm.proxy.utils import ( PrismaClient, + ProxyLogging, _hash_token_if_needed, handle_exception_on_proxy, is_valid_api_key, @@ -3149,6 +3150,63 @@ def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: return new_token +async def _execute_virtual_key_regeneration( + *, + prisma_client: PrismaClient, + key_in_db: LiteLLM_VerificationToken, + hashed_api_key: str, + key: str, + data: Optional[RegenerateKeyRequest], + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str], + user_api_key_cache: DualCache, + proxy_logging_obj: ProxyLogging, +) -> GenerateKeyResponse: + """Generate new token, update DB, invalidate cache, and return response.""" + from litellm.proxy.proxy_server import hash_token + + new_token = get_new_token(data=data) + new_token_hash = hash_token(new_token) + new_token_key_name = f"sk-...{new_token[-4:]}" + update_data = {"token": new_token_hash, "key_name": new_token_key_name} + + non_default_values = {} + if data is not None: + non_default_values = await prepare_key_update_data( + data=data, existing_key_row=key_in_db + ) + verbose_proxy_logger.debug("non_default_values: %s", non_default_values) + update_data.update(non_default_values) + update_data = prisma_client.jsonify_object(data=update_data) + + updated_token = await prisma_client.db.litellm_verificationtoken.update( + where={"token": hashed_api_key}, + data=update_data, # type: ignore + ) + updated_token_dict = dict(updated_token) if updated_token is not None else {} + updated_token_dict["key"] = new_token + updated_token_dict["token_id"] = updated_token_dict.pop("token") + + if hashed_api_key or key: + await _delete_cache_key_object( + hashed_token=hash_token(key), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + response = GenerateKeyResponse(**updated_token_dict) + asyncio.create_task( + KeyManagementEventHooks.async_key_rotated_hook( + data=data, + existing_key_row=key_in_db, + response=response, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + ) + return response + + @router.post( "/key/{key:path}/regenerate", tags=["key management"], @@ -3160,7 +3218,7 @@ def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: dependencies=[Depends(user_api_key_auth)], ) @management_endpoint_wrapper -async def regenerate_key_fn( # noqa: PLR0915 +async def regenerate_key_fn( key: Optional[str] = None, data: Optional[RegenerateKeyRequest] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -3328,65 +3386,17 @@ async def regenerate_key_fn( # noqa: PLR0915 litellm_changed_by=litellm_changed_by, ) - new_token = get_new_token(data=data) - - new_token_hash = hash_token(new_token) - new_token_key_name = f"sk-...{new_token[-4:]}" - - # Prepare the update data - update_data = { - "token": new_token_hash, - "key_name": new_token_key_name, - } - - non_default_values = {} - if data is not None: - # Update with any provided parameters from GenerateKeyRequest - non_default_values = await prepare_key_update_data( - data=data, existing_key_row=_key_in_db - ) - verbose_proxy_logger.debug("non_default_values: %s", non_default_values) - - update_data.update(non_default_values) - update_data = prisma_client.jsonify_object(data=update_data) - # Update the token in the database - updated_token = await prisma_client.db.litellm_verificationtoken.update( - where={"token": hashed_api_key}, - data=update_data, # type: ignore - ) - - updated_token_dict = {} - if updated_token is not None: - updated_token_dict = dict(updated_token) - - updated_token_dict["key"] = new_token - updated_token_dict["token_id"] = updated_token_dict.pop("token") - - ### 3. remove existing key entry from cache - ###################################################################### - - if hashed_api_key or key: - await _delete_cache_key_object( - hashed_token=hash_token(key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - response = GenerateKeyResponse( - **updated_token_dict, - ) - - asyncio.create_task( - KeyManagementEventHooks.async_key_rotated_hook( - data=data, - existing_key_row=_key_in_db, - response=response, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) + return await _execute_virtual_key_regeneration( + prisma_client=prisma_client, + key_in_db=_key_in_db, + hashed_api_key=hashed_api_key, + key=key, + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) - - return response except Exception as e: verbose_proxy_logger.exception("Error regenerating key: %s", e) raise handle_exception_on_proxy(e)