fix: preserve key alias and team_id after key regeneration, deletion

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,

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,
         ]:

tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py

@@ -11,6 +11,7 @@ sys.path.insert(
 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 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
     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