fix(proxy): avoid N+1 key count queries in user list

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -1777,20 +1777,26 @@ async def get_user_key_counts(
     if not user_ids or len(user_ids) == 0:
         return {}
 
-    result = {}
+    result = {user_id: 0 for user_id in user_ids}
 
-    # Get count for each user_id individually
-    for user_id in user_ids:
-        count = await prisma_client.db.litellm_verificationtoken.count(
-            where={
-                "user_id": user_id,
-                "OR": [
-                    {"team_id": None},
-                    {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
-                ],
-            }
-        )
-        result[user_id] = count
+    grouped_counts = await prisma_client.db.litellm_verificationtoken.group_by(
+        by=["user_id"],
+        where={
+            "user_id": {"in": user_ids},
+            "OR": [
+                {"team_id": None},
+                {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
+            ],
+        },
+        count={"user_id": True},
+    )
+
+    for row in grouped_counts:
+        user_id = row.get("user_id")
+        if user_id is None:
+            continue
+        count_block = row.get("_count") or {}
+        result[user_id] = count_block.get("user_id", 0)
 
     return result
 

tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py

@@ -657,6 +657,56 @@ async def test_get_users_includes_timestamps(mocker):
     assert user_response.key_count == 0
 
 
+@pytest.mark.asyncio
+async def test_get_user_key_counts_returns_empty_for_no_user_ids(mocker):
+    mock_prisma_client = mocker.MagicMock()
+
+    assert await get_user_key_counts(mock_prisma_client, None) == {}
+    assert await get_user_key_counts(mock_prisma_client, []) == {}
+    assert mock_prisma_client.db.litellm_verificationtoken.group_by.call_count == 0
+
+
+@pytest.mark.asyncio
+async def test_get_user_key_counts_uses_group_by(mocker):
+    from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
+
+    mock_prisma_client = mocker.MagicMock()
+    captured_kwargs = {}
+
+    async def mock_group_by(**kwargs):
+        captured_kwargs.update(kwargs)
+        return [
+            {"user_id": "user-a", "_count": {"user_id": 2}},
+            {"user_id": "user-c", "_count": {"user_id": 1}},
+            {"user_id": "user-d"},
+            {"user_id": None, "_count": {"user_id": 99}},
+        ]
+
+    async def mock_count(*args, **kwargs):
+        raise AssertionError("count() should not be called for user key counts")
+
+    mock_prisma_client.db.litellm_verificationtoken.group_by = mock_group_by
+    mock_prisma_client.db.litellm_verificationtoken.count = mock_count
+
+    result = await get_user_key_counts(
+        mock_prisma_client,
+        ["user-a", "user-b", "user-c", "user-d"],
+    )
+
+    assert result == {"user-a": 2, "user-b": 0, "user-c": 1, "user-d": 0}
+    assert captured_kwargs == {
+        "by": ["user_id"],
+        "where": {
+            "user_id": {"in": ["user-a", "user-b", "user-c", "user-d"]},
+            "OR": [
+                {"team_id": None},
+                {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
+            ],
+        },
+        "count": {"user_id": True},
+    }
+
+
 def test_validate_sort_params():
     """
     Test that validate_sort_params returns None if sort_by is None