From 8696c33d7bd973d5dae98a64c0cc5b4fc97fa23a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 25 Feb 2026 17:55:26 -0800 Subject: [PATCH] [Fix] /key/aliases: Add pagination and search to prevent OOMs The /key/aliases endpoint previously fetched all key aliases from the database without limit, causing OOM crashes with large key sets. Added page, size, and search query parameters with database-level filtering to enable paginated and searchable key alias retrieval. Updated the response to include pagination metadata (total_count, current_page, total_pages, size) matching the /v2/model/info pattern. Co-Authored-By: Claude Haiku 4.5 --- .../key_management_endpoints.py | 59 ++++++++++++------- .../test_key_generate_prisma.py | 34 +++++++---- 2 files changed, 61 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b489369071f..9f5b8222292 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4109,13 +4109,23 @@ async def list_keys( dependencies=[Depends(user_api_key_auth)], ) @management_endpoint_wrapper -async def key_aliases() -> Dict[str, List[str]]: +async def key_aliases( + page: int = Query(1, ge=1, description="Page number"), + size: int = Query(50, ge=1, le=100, description="Page size"), + search: Optional[str] = Query( + None, description="Search key aliases (case-insensitive partial match)" + ), +) -> Dict[str, Any]: """ - Lists all key aliases + Lists key aliases with pagination and optional search. Returns: { - "aliases": List[str] + "aliases": List[str], + "total_count": int, + "current_page": int, + "total_pages": int, + "size": int, } """ try: @@ -4127,36 +4137,43 @@ async def key_aliases() -> Dict[str, List[str]]: verbose_proxy_logger.error("Database not connected") raise Exception("Database not connected") - where: Dict[str, Any] = {} + conditions: List[Dict[str, Any]] = [{"key_alias": {"not": None}}] try: - where.update(_get_condition_to_filter_out_ui_session_tokens()) + conditions.append(_get_condition_to_filter_out_ui_session_tokens()) except NameError: # Helper may not exist in some builds; ignore if missing pass + if search: + conditions.append( + {"key_alias": {"contains": search, "mode": "insensitive"}} + ) + where: Dict[str, Any] = {"AND": conditions} + total_count = await prisma_client.db.litellm_verificationtoken.count( + where=where, + ) rows = await prisma_client.db.litellm_verificationtoken.find_many( where=where, order=[{"key_alias": "asc"}], + skip=(page - 1) * size, + take=size, ) - seen = set() - aliases: List[str] = [] - for row in rows: - alias = getattr(row, "key_alias", None) - if alias is None and isinstance(row, dict): - alias = row.get("key_alias") + aliases: List[str] = [row.key_alias for row in rows if row.key_alias] # type: ignore[misc] - if not alias: - continue + total_pages = -(-total_count // size) if total_count > 0 else 0 + verbose_proxy_logger.debug( + f"key_aliases: page={page}, size={size}, search={search!r}, " + f"total_count={total_count}, total_pages={total_pages}" + ) - alias_str = str(alias).strip() - if alias_str and alias_str not in seen: - seen.add(alias_str) - aliases.append(alias_str) - - verbose_proxy_logger.debug(f"Returning {len(aliases)} key aliases") - - return {"aliases": aliases} + return { + "aliases": aliases, + "total_count": total_count, + "current_page": page, + "total_pages": total_pages, + "size": size, + } except Exception as e: verbose_proxy_logger.exception(f"Error in key_aliases: {e}") diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index c3f68762810..ed528f21e0d 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3668,9 +3668,10 @@ async def test_list_keys(prisma_client): async def test_key_aliases(prisma_client): """ Test the key_aliases function: - - Returns a list + - Returns a paginated response - Includes alias from a newly created key - - Aliases are unique and sorted + - Aliases are sorted + - Pagination and search params work correctly """ import asyncio import uuid @@ -3682,10 +3683,16 @@ async def test_key_aliases(prisma_client): setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") await litellm.proxy.proxy_server.prisma_client.connect() - # Basic call - response = await key_aliases() + # Basic call - check pagination response shape + response = await key_aliases(page=1, size=50) assert "aliases" in response assert isinstance(response["aliases"], list) + assert "total_count" in response + assert "current_page" in response + assert "total_pages" in response + assert "size" in response + assert response["current_page"] == 1 + assert response["size"] == 50 # Create a new user (and key) with a unique alias unique_id = str(uuid.uuid4()) @@ -3704,17 +3711,22 @@ async def test_key_aliases(prisma_client): # Allow async DB writes to settle await asyncio.sleep(2) - # Call again and validate - response_after = await key_aliases() + # Call again and validate alias is present + response_after = await key_aliases(page=1, size=50) aliases = response_after["aliases"] - - # Contains the new alias assert test_alias in aliases - - # Unique & sorted (endpoint dedupes and orders ascending) - assert len(aliases) == len(set(aliases)) assert aliases == sorted(aliases) + # Search by partial alias + partial = test_alias[:10] + search_response = await key_aliases(page=1, size=50, search=partial) + assert test_alias in search_response["aliases"] + + # Search with no match + no_match_response = await key_aliases(page=1, size=50, search="__no_match_xyz__") + assert len(no_match_response["aliases"]) == 0 + assert no_match_response["total_count"] == 0 + @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") @pytest.mark.asyncio