From b41311bb21865ec47653487cb8c63bd994a5496c Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 8 Mar 2025 16:13:18 -0800
Subject: [PATCH] (UI) - Fix show correct count of internal user keys on Users
 Page (#9082)

* get_user_key_counts

* fix get_user_key_counts

* fix get_user_key_counts

* test_get_users_filters_dashboard_keys

* remove unused func
---
 .../internal_user_endpoints.py                | 46 ++++++++--
 .../test_key_management.py                    | 86 +++++++++++++++++++
 2 files changed, 125 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py
index 36c12ada3fd..37b6f0bbf54 100644
--- a/litellm/proxy/management_endpoints/internal_user_endpoints.py
+++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py
@@ -739,6 +739,43 @@ async def user_update(
         )
 
 
+async def get_user_key_counts(
+    prisma_client,
+    user_ids: Optional[List[str]] = None,
+):
+    """
+    Helper function to get the count of keys for each user using Prisma's count method.
+
+    Args:
+        prisma_client: The Prisma client instance
+        user_ids: List of user IDs to get key counts for
+
+    Returns:
+        Dictionary mapping user_id to key count
+    """
+    from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
+
+    if not user_ids or len(user_ids) == 0:
+        return {}
+
+    result = {}
+
+    # 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
+
+    return result
+
+
 @router.get(
     "/user/get_users",
     tags=["Internal User management"],
@@ -830,14 +867,9 @@ async def get_users(
 
     # Get key count for each user
     if users is not None:
-        user_keys = await prisma_client.db.litellm_verificationtoken.group_by(
-            by=["user_id"],
-            count={"user_id": True},
-            where={"user_id": {"in": [user.user_id for user in users]}},
+        user_key_counts = await get_user_key_counts(
+            prisma_client, [user.user_id for user in users]
         )
-        user_key_counts = {
-            item["user_id"]: item["_count"]["user_id"] for item in user_keys
-        }
     else:
         user_key_counts = {}
 
diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py
index 46fdc0b2cf1..c1548096f28 100644
--- a/tests/proxy_admin_ui_tests/test_key_management.py
+++ b/tests/proxy_admin_ui_tests/test_key_management.py
@@ -378,6 +378,92 @@ async def test_get_users(prisma_client):
         await prisma_client.db.litellm_usertable.delete(where={"user_id": user.user_id})
 
 
+@pytest.mark.asyncio
+async def test_get_users_filters_dashboard_keys(prisma_client):
+    """
+    Tests that /users/list endpoint doesn't return keys with team_id='litellm-dashboard'
+
+    The dashboard keys should be filtered out from the response
+    """
+    litellm.set_verbose = True
+    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    await litellm.proxy.proxy_server.prisma_client.connect()
+
+    # Create a test user
+    new_user_id = f"test_user_with_keys-{uuid.uuid4()}"
+    test_user = NewUserRequest(
+        user_id=new_user_id,
+        user_role=LitellmUserRoles.INTERNAL_USER.value,
+        auto_create_key=False,
+    )
+
+    await new_user(
+        test_user,
+        UserAPIKeyAuth(
+            user_role=LitellmUserRoles.PROXY_ADMIN,
+            api_key="sk-1234",
+            user_id="admin",
+        ),
+    )
+
+    # Create two keys for the user - one with team_id="litellm-dashboard" and one without
+    regular_key = await generate_key_helper_fn(
+        user_id=test_user.user_id,
+        request_type="key",
+        team_id="litellm-dashboard",  # This key should be included in the response
+        models=[],
+        aliases={},
+        config={},
+        spend=0,
+        duration=None,
+    )
+
+    regular_key = await generate_key_helper_fn(
+        user_id=test_user.user_id,
+        request_type="key",
+        team_id="NEW_TEAM",  # This key should be included in the response
+        models=[],
+        aliases={},
+        config={},
+        spend=0,
+        duration=None,
+    )
+
+    regular_key = await generate_key_helper_fn(
+        user_id=test_user.user_id,
+        request_type="key",
+        team_id=None,  # This key should be included in the response
+        models=[],
+        aliases={},
+        config={},
+        spend=0,
+        duration=None,
+    )
+
+    # Test get_users for the specific user
+    result = await get_users(
+        user_ids=test_user.user_id,
+        role=None,
+        page=1,
+        page_size=20,
+    )
+
+    print("get users result", result)
+    assert "users" in result
+    assert len(result["users"]) == 1
+
+    # Verify the key count is correct (should be 1, not counting dashboard keys)
+    user = result["users"][0]
+    assert user.user_id == test_user.user_id
+    assert user.key_count == 2  # Only count the regular keys, not the UI dashboard key
+
+    # Clean up test user and keys
+    await prisma_client.db.litellm_usertable.delete(
+        where={"user_id": test_user.user_id}
+    )
+
+
 @pytest.mark.asyncio
 async def test_get_users_key_count(prisma_client):
     """