diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b489369071f..5b56133f1ce 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,55 @@ async def key_aliases() -> Dict[str, List[str]]: verbose_proxy_logger.error("Database not connected") raise Exception("Database not connected") - where: Dict[str, Any] = {} - try: - where.update(_get_condition_to_filter_out_ui_session_tokens()) - except NameError: - # Helper may not exist in some builds; ignore if missing - pass + # Build a parameterized WHERE clause to avoid loading full rows into + # memory. Raw SQL is used because the Prisma client wrapper does not + # support column-level SELECT projection on find_many. + # + # $1 is always UI_SESSION_TOKEN_TEAM_ID (filters out UI session tokens). + query_params: List[Any] = [UI_SESSION_TOKEN_TEAM_ID] + where_parts = [ + "key_alias IS NOT NULL", + "key_alias != ''", + "(team_id IS NULL OR team_id != $1)", + ] + if search: + query_params.append(f"%{search}%") + where_parts.append(f"key_alias ILIKE ${len(query_params)}") - rows = await prisma_client.db.litellm_verificationtoken.find_many( - where=where, - order=[{"key_alias": "asc"}], + where_sql = " AND ".join(where_parts) + + count_sql = ( + f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' + ) + count_rows = await prisma_client.db.query_raw(count_sql, *query_params) + total_count = int(count_rows[0]["count"]) if count_rows else 0 + + aliases_params = query_params + [size, (page - 1) * size] + limit_idx = len(aliases_params) - 1 + offset_idx = len(aliases_params) + aliases_sql = ( + f"SELECT key_alias" + f' FROM "LiteLLM_VerificationToken"' + f" WHERE {where_sql}" + f" ORDER BY key_alias ASC" + f" LIMIT ${limit_idx} OFFSET ${offset_idx}" + ) + alias_rows = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + aliases: List[str] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] + + 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}" ) - 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") - - if not alias: - continue - - 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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index fb71adc1085..7565e901ecd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -45,6 +45,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( check_team_key_model_specific_limits, delete_verification_tokens, generate_key_helper_fn, + key_aliases, list_keys, prepare_key_update_data, reset_key_spend_fn, @@ -6210,3 +6211,97 @@ async def test_generate_key_helper_fn_agent_id(): assert key_data.get("agent_id") == "test-agent-456", ( f"Expected agent_id='test-agent-456' in key_data, got: {key_data.get('agent_id')}" ) + + +@pytest.mark.asyncio +async def test_key_aliases_response_shape(): + """Test that key_aliases returns the correct paginated response shape.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 2}], + [{"key_alias": "alias-alpha"}, {"key_alias": "alias-beta"}], + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await key_aliases(page=1, size=50, search=None) + + assert result["aliases"] == ["alias-alpha", "alias-beta"] + assert result["total_count"] == 2 + assert result["current_page"] == 1 + assert result["total_pages"] == 1 + assert result["size"] == 50 + + # Both SQL calls must filter out null/empty aliases + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + aliases_sql = mock_prisma_client.db.query_raw.call_args_list[1].args[0] + assert "key_alias IS NOT NULL" in count_sql + assert "key_alias IS NOT NULL" in aliases_sql + + +@pytest.mark.asyncio +async def test_key_aliases_pagination_skip_take(): + """Test that LIMIT and OFFSET are correctly derived from page and size.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 120}], + [], + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await key_aliases(page=3, size=25, search=None) + + assert result["current_page"] == 3 + assert result["size"] == 25 + assert result["total_count"] == 120 + assert result["total_pages"] == 5 # ceil(120 / 25) + + # aliases query params: [UI_SESSION_TOKEN_TEAM_ID, size=25, offset=50] + aliases_call_args = mock_prisma_client.db.query_raw.call_args_list[1].args + assert aliases_call_args[-2] == 25 # LIMIT = size + assert aliases_call_args[-1] == 50 # OFFSET = (3 - 1) * 25 + + +@pytest.mark.asyncio +async def test_key_aliases_search_filter(): + """Test that the search param adds a case-insensitive ILIKE condition.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 0}], + [], + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await key_aliases(page=1, size=50, search="my-key") + + count_call = mock_prisma_client.db.query_raw.call_args_list[0] + count_sql = count_call.args[0] + count_params = count_call.args[1:] + + assert "ILIKE" in count_sql + assert "%my-key%" in count_params + + +@pytest.mark.asyncio +async def test_key_aliases_no_search_omits_ilike_filter(): + """Test that without a search term no ILIKE condition is added.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 0}], + [], + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await key_aliases(page=1, size=50, search=None) + + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + assert "ILIKE" not in count_sql + +