support create automatic aks cluster (Azure#846)

src/commands/aksCreateCluster/aksCreateCluster.ts

@@ -1,18 +1,23 @@
 import { IActionContext } from "@microsoft/vscode-azext-utils";
-import { getAksClusterSubscriptionNode } from "../utils/clusters";
-import { failed } from "../utils/errorable";
 import * as vscode from "vscode";
 import * as k8s from "vscode-kubernetes-tools-api";
-import { getExtension } from "../utils/host";
-import { CreateClusterDataProvider, CreateClusterPanel } from "../../panels/CreateClusterPanel";
 import { getReadySessionProvider } from "../../auth/azureAuth";
+import { CreateClusterDataProvider, CreateClusterPanel } from "../../panels/CreateClusterPanel";
+import { getAksClusterSubscriptionNode } from "../utils/clusters";
+import { Errorable, failed } from "../utils/errorable";
+import { getExtension } from "../utils/host";
+import { getSubscription } from "../utils/subscriptions";
+import { SubscriptionTreeNode } from "../../tree/subscriptionTreeItem";
 
 /**
  * A multi-step input using window.createQuickPick() and window.createInputBox().
  *
  * This first part uses the helper class `MultiStepInput` that wraps the API for the multi-step case.
  */
 export default async function aksCreateCluster(_context: IActionContext, target: unknown): Promise<void> {
+    let subscriptionNode: Errorable<SubscriptionTreeNode>;
+    let subscriptionId: string | undefined;
+    let subscriptionName: string | undefined;
     const cloudExplorer = await k8s.extension.cloudExplorer.v1;
 
     const sessionProvider = await getReadySessionProvider();
@@ -21,10 +26,33 @@ export default async function aksCreateCluster(_context: IActionContext, target:
         return;
     }
 
-    const subscriptionNode = getAksClusterSubscriptionNode(target, cloudExplorer);
-    if (failed(subscriptionNode)) {
-        vscode.window.showErrorMessage(subscriptionNode.error);
-        return;
+    switch (typeof target) {
+        case "string": {
+            const subscriptionResult = await getSubscription(sessionProvider.result, target);
+            if (failed(subscriptionResult)) {
+                vscode.window.showErrorMessage(subscriptionResult.error);
+                return;
+            }
+            subscriptionId = subscriptionResult.result?.id;
+            subscriptionName = subscriptionResult.result?.displayName;
+            break;
+        }
+
+        default: {
+            subscriptionNode = getAksClusterSubscriptionNode(target, cloudExplorer);
+            if (failed(subscriptionNode)) {
+                vscode.window.showErrorMessage(subscriptionNode.error);
+                return;
+            }
+            if (!subscriptionNode.result || !subscriptionNode.result.subscriptionId || !subscriptionNode.result.name) {
+                vscode.window.showErrorMessage("Subscription not found.");
+                return;
+            }
+            subscriptionId = subscriptionNode.result?.subscriptionId;
+            subscriptionName = subscriptionNode.result?.name;
+            break;
+        }
+
     }
 
     const extension = getExtension();
@@ -35,11 +63,13 @@ export default async function aksCreateCluster(_context: IActionContext, target:
 
     const panel = new CreateClusterPanel(extension.result.extensionUri);
 
-    const dataProvider = new CreateClusterDataProvider(
-        sessionProvider.result,
-        subscriptionNode.result.subscriptionId,
-        subscriptionNode.result.name,
-        () => vscode.commands.executeCommand("aks.refreshSubscription", target),
+    if (!subscriptionId || !subscriptionName) {
+        vscode.window.showErrorMessage("Subscription ID or Name is undefined.");
+        return;
+    }
+
+    const dataProvider = new CreateClusterDataProvider(sessionProvider.result, subscriptionId, subscriptionName, () =>
+        vscode.commands.executeCommand("aks.refreshSubscription", target),
     );
 
     panel.show(dataProvider);

src/commands/aksCreateCluster/aksCreateClusterFromCopilot.ts

@@ -0,0 +1,6 @@
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import * as vscode from "vscode";
+
+export async function aksCreateClusterFromCopilot(_context: IActionContext, subscriptionId: string): Promise<void> {
+    vscode.commands.executeCommand("aks.createCluster", subscriptionId);
+}

src/commands/utils/featureRegistrations.ts

@@ -0,0 +1,85 @@
+import { FeatureClient } from "@azure/arm-features";
+import { longRunning } from "./host";
+
+const MAX_RETRIES = 5;
+const RETRY_DELAY_MS = 60000; // 1 minute
+
+export enum FeatureRegistrationState {
+    Registered = "Registered",
+    Registering = "Registering",
+    Failed = "Failed",
+}
+
+export type MultipleFeatureRegistration = {
+    resourceProviderNamespace: string;
+    featureName: string;
+};
+
+export type FeatureRegistrationResult = {
+    resourceProviderNamespace: string;
+    featureName: string;
+    registrationStatus: string;
+};
+
+export async function createFeatureRegistrationsWithRetry(
+    featureClient: FeatureClient,
+    resourceProviderNamespace: string,
+    featureName: string,
+): Promise<void> {
+    let retries = 0;
+    //register the feature
+    try {
+        await featureClient.features.register(resourceProviderNamespace, featureName);
+    } catch (error) {
+        throw new Error(
+            `Failed to initiate registration for feature ${featureName} in ${resourceProviderNamespace}: ${error}`,
+        );
+    }
+    while (retries < MAX_RETRIES) {
+        // get the registration status
+        const featureRegistrationResult = await featureClient.features.get(resourceProviderNamespace, featureName);
+        const result = featureRegistrationResult?.properties?.state ?? FeatureRegistrationState.Failed;
+        switch (result) {
+            case FeatureRegistrationState.Registered:
+                console.log(`Feature ${featureName} registered successfully for ${resourceProviderNamespace}.`);
+                return;
+
+            case FeatureRegistrationState.Registering:
+                retries++;
+                console.log(`Feature ${featureName} is still registering. Retry ${retries}/${MAX_RETRIES}.`);
+                await delay(RETRY_DELAY_MS);
+                break;
+
+            default:
+                throw new Error(
+                    `Failed to register the preview feature ${featureName} for ${resourceProviderNamespace}. Current state: ${result}, please try again.`,
+                );
+        }
+    }
+}
+
+async function delay(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+export async function createMultipleFeatureRegistrations(
+    featureClient: FeatureClient,
+    featureRegistrations: MultipleFeatureRegistration[],
+): Promise<void> {
+    await longRunning(`Registering the preview features.`, async () => {
+        try {
+            const featureRegistrationResults = featureRegistrations.map(async (featureRegistration) => {
+                return createFeatureRegistrationsWithRetry(
+                    featureClient,
+                    featureRegistration.resourceProviderNamespace,
+                    featureRegistration.featureName,
+                );
+            });
+            await Promise.all(featureRegistrationResults);
+            console.log("All features registered successfully.");
+        } catch (error) {
+            console.error("Error registering features:", error);
+            throw error;
+        }
+    });
+}

src/commands/utils/subscriptions.ts

@@ -1,8 +1,8 @@
-import { Subscription } from "@azure/arm-resources-subscriptions";
+import { Subscription, SubscriptionsGetResponse } from "@azure/arm-resources-subscriptions";
+import { ReadyAzureSessionProvider } from "../../auth/types";
 import { getSubscriptionClient, listAll } from "./arm";
-import { Errorable, map as errmap } from "./errorable";
 import { getFilteredSubscriptions } from "./config";
-import { ReadyAzureSessionProvider } from "../../auth/types";
+import { Errorable, map as errmap } from "./errorable";
 
 export enum SelectionType {
     Filtered,
@@ -41,3 +41,19 @@ function sortAndFilter(subscriptions: DefinedSubscription[], selectionType: Sele
 function isDefinedSubscription(sub: Subscription): sub is DefinedSubscription {
     return sub.subscriptionId !== undefined && sub.displayName !== undefined;
 }
+
+function isDefinedSubscriptionGetResponse(sub: SubscriptionsGetResponse): sub is DefinedSubscription {
+    return sub.subscriptionId !== undefined && sub.displayName !== undefined;
+}
+
+export async function getSubscription(
+    sessionProvider: ReadyAzureSessionProvider,
+    subscriptionId: string,
+): Promise<Errorable<DefinedSubscription>> {
+    const client = getSubscriptionClient(sessionProvider);
+    const subResult: SubscriptionsGetResponse = await client.subscriptions.get(subscriptionId);
+    if (!isDefinedSubscriptionGetResponse(subResult)) {
+        return { succeeded: false, error: "Subscription is not found" };
+    }
+    return { succeeded: true, result: subResult };
+}

src/panels/CreateClusterPanel.ts

@@ -1,17 +1,22 @@
 import { ContainerServiceClient, KubernetesVersion } from "@azure/arm-containerservice";
+import { FeatureClient } from "@azure/arm-features";
 import { ResourceGroup as ARMResourceGroup, ResourceManagementClient } from "@azure/arm-resources";
 import { RestError } from "@azure/storage-blob";
 import { Uri, window } from "vscode";
 import { getEnvironment } from "../auth/azureAuth";
 import { ReadyAzureSessionProvider } from "../auth/types";
-import { getAksClient, getResourceManagementClient } from "../commands/utils/arm";
+import { getAksClient, getFeatureClient, getResourceManagementClient } from "../commands/utils/arm";
 import { getPortalResourceUrl } from "../commands/utils/env";
 import { failed, getErrorMessage } from "../commands/utils/errorable";
+import {
+    createMultipleFeatureRegistrations,
+    MultipleFeatureRegistration,
+} from "../commands/utils/featureRegistrations";
 import { getResourceGroups } from "../commands/utils/resourceGroups";
 import { MessageHandler, MessageSink } from "../webview-contract/messaging";
 import {
     InitialState,
-    Preset,
+    PresetType,
     ProgressEventType,
     ToVsCodeMsgDef,
     ToWebViewMsgDef,
@@ -34,6 +39,7 @@ export class CreateClusterPanel extends BasePanel<"createCluster"> {
 export class CreateClusterDataProvider implements PanelDataProvider<"createCluster"> {
     private readonly resourceManagementClient: ResourceManagementClient;
     private readonly containerServiceClient: ContainerServiceClient;
+    private readonly featureClient: FeatureClient;
 
     public constructor(
         readonly sessionProvider: ReadyAzureSessionProvider,
@@ -43,6 +49,7 @@ export class CreateClusterDataProvider implements PanelDataProvider<"createClust
     ) {
         this.resourceManagementClient = getResourceManagementClient(sessionProvider, this.subscriptionId);
         this.containerServiceClient = getAksClient(sessionProvider, this.subscriptionId);
+        this.featureClient = getFeatureClient(sessionProvider, this.subscriptionId);
     }
 
     getTitle(): string {
@@ -129,7 +136,7 @@ export class CreateClusterDataProvider implements PanelDataProvider<"createClust
         groupName: string,
         location: string,
         name: string,
-        preset: Preset,
+        preset: PresetType,
         webview: MessageSink<ToWebViewMsgDef>,
     ) {
         if (isNewResourceGroup) {
@@ -150,6 +157,7 @@ export class CreateClusterDataProvider implements PanelDataProvider<"createClust
             webview,
             this.containerServiceClient,
             this.resourceManagementClient,
+            this.featureClient,
         );
 
         this.refreshTree();
@@ -196,10 +204,11 @@ async function createCluster(
     groupName: string,
     location: string,
     name: string,
-    preset: Preset,
+    preset: PresetType,
     webview: MessageSink<ToWebViewMsgDef>,
     containerServiceClient: ContainerServiceClient,
     resourceManagementClient: ResourceManagementClient,
+    featureClient: FeatureClient,
 ) {
     const operationDescription = `Creating cluster ${name}`;
     webview.postProgressUpdate({
@@ -252,12 +261,27 @@ async function createCluster(
     // Create a unique deployment name.
     const deploymentName = `${name}-${Math.random().toString(36).substring(5)}`;
     const deploymentSpec = new ClusterDeploymentBuilder()
-        .buildCommonParameters(clusterSpec)
+        .buildCommonParameters(clusterSpec, preset)
         .buildTemplate(preset)
         .getDeployment();
 
     const environment = getEnvironment();
 
+    // feature registration
+    try {
+        await doFeatureRegistration(preset, featureClient);
+    } catch (error) {
+        window.showErrorMessage(`Error Registering preview features for AKS cluster ${name}: ${error}`);
+        webview.postProgressUpdate({
+            event: ProgressEventType.Failed,
+            operationDescription: "Error Registering preview features for AKS cluster",
+            errorMessage: getErrorMessage(error),
+            deploymentPortalUrl: null,
+            createdCluster: null,
+        });
+        return;
+    }
+
     try {
         const poller = await resourceManagementClient.deployments.beginCreateOrUpdate(
             groupName,
@@ -323,6 +347,41 @@ async function createCluster(
     }
 }
 
+async function doFeatureRegistration(preset: PresetType, featureClient: FeatureClient) {
+    if (preset !== PresetType.Automatic) {
+        return;
+    }
+    //Doc link - https://learn.microsoft.com/en-us/azure/aks/learn/quick-kubernetes-automatic-deploy?pivots=azure-cli#register-the-feature-flags
+    const features: MultipleFeatureRegistration[] = [
+        {
+            resourceProviderNamespace: "Microsoft.ContainerService",
+            featureName: "EnableAPIServerVnetIntegrationPreview",
+        },
+        {
+            resourceProviderNamespace: "Microsoft.ContainerService",
+            featureName: "NRGLockdownPreview",
+        },
+        {
+            resourceProviderNamespace: "Microsoft.ContainerService",
+            featureName: "SafeguardsPreview",
+        },
+        {
+            resourceProviderNamespace: "Microsoft.ContainerService",
+            featureName: "NodeAutoProvisioningPreview",
+        },
+        {
+            resourceProviderNamespace: "Microsoft.ContainerService",
+            featureName: "DisableSSHPreview",
+        },
+        {
+            resourceProviderNamespace: "Microsoft.ContainerService",
+            featureName: "AutomaticSKUPreview",
+        },
+    ];
+
+    await createMultipleFeatureRegistrations(featureClient, features);
+}
+
 function getInvalidTemplateErrorMessage(ex: InvalidTemplateDeploymentRestError): string {
     const innerDetails = ex.details.error?.details || [];
     if (innerDetails.length > 0) {