Phoenix runtime - Reset workspace (#1136)

* Phoenix runtime - Reset workspace

* Format and Lint issues

* Typo issue

* Reset warning text change and create new context on allcation of new container

* Closing only notebook related

* resolved comments from previous PR

* On Schema Analyser allocate call

Co-authored-by: Srinath Narayanan <srnara@microsoft.com>

src/Explorer/Notebook/NotebookContainerClient.ts

@@ -1,14 +1,21 @@
 /**
  * Notebook container related stuff
  */
+import { PhoenixClient } from "Phoenix/PhoenixClient";
 import * as Constants from "../../Common/Constants";
-import { ConnectionStatusType } from "../../Common/Constants";
+import { ConnectionStatusType, HttpHeaders, HttpStatusCodes } from "../../Common/Constants";
 import { getErrorMessage } from "../../Common/ErrorHandlingUtils";
 import * as Logger from "../../Common/Logger";
 import * as DataModels from "../../Contracts/DataModels";
-import { ContainerConnectionInfo } from "../../Contracts/DataModels";
+import {
+  ContainerConnectionInfo,
+  IPhoenixConnectionInfoResult,
+  IProvisionData,
+  IResponse,
+} from "../../Contracts/DataModels";
 import { userContext } from "../../UserContext";
 import { createOrUpdate, destroy } from "../../Utils/arm/generatedClients/cosmosNotebooks/notebookWorkspaces";
+import { getAuthorizationHeader } from "../../Utils/AuthorizationUtils";
 import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { NotebookUtil } from "./NotebookUtil";
 import { useNotebook } from "./useNotebook";
@@ -103,12 +110,9 @@ export class NotebookContainerClient {
 
   private checkStatus(): boolean {
     if (NotebookUtil.isPhoenixEnabled()) {
-      if (
-        useNotebook.getState().containerStatus?.status &&
-        useNotebook.getState().containerStatus?.status === Constants.ContainerStatusType.InActive
-      ) {
+      if (useNotebook.getState().containerStatus?.status === Constants.ContainerStatusType.Disconnected) {
         const connectionStatus: ContainerConnectionInfo = {
-          status: ConnectionStatusType.ReConnect,
+          status: ConnectionStatusType.Reconnect,
         };
         useNotebook.getState().resetContainerConnection(connectionStatus);
         useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
@@ -117,17 +121,21 @@ export class NotebookContainerClient {
     }
     return true;
   }
-  public async resetWorkspace(): Promise<void> {
+
+  public async resetWorkspace(): Promise<IResponse<IPhoenixConnectionInfoResult>> {
     this.isResettingWorkspace = true;
+    let response: IResponse<IPhoenixConnectionInfoResult>;
     try {
-      await this._resetWorkspace();
+      response = await this._resetWorkspace();
     } catch (error) {
       Promise.reject(error);
+      return response;
     }
     this.isResettingWorkspace = false;
+    return response;
   }
 
-  private async _resetWorkspace(): Promise<void> {
+  private async _resetWorkspace(): Promise<IResponse<IPhoenixConnectionInfoResult>> {
     const notebookServerInfo = useNotebook.getState().notebookServerInfo;
     if (!notebookServerInfo || !notebookServerInfo.notebookServerEndpoint) {
       const error = "No server endpoint detected";
@@ -137,13 +145,40 @@ export class NotebookContainerClient {
 
     const { notebookServerEndpoint, authToken } = this.getNotebookServerConfig();
     try {
-      await fetch(`${notebookServerEndpoint}/api/shutdown`, {
-        method: "POST",
-        headers: { Authorization: authToken },
-      });
+      let data: IPhoenixConnectionInfoResult;
+      let response: Response;
+      if (NotebookUtil.isPhoenixEnabled()) {
+        const provisionData: IProvisionData = {
+          aadToken: userContext.authorizationToken,
+          subscriptionId: userContext.subscriptionId,
+          resourceGroup: userContext.resourceGroup,
+          dbAccountName: userContext.databaseAccount.name,
+          cosmosEndpoint: userContext.databaseAccount.properties.documentEndpoint,
+        };
+        response = await fetch(`${PhoenixClient.getPhoenixEndpoint()}/api/controlplane/toolscontainer/reset`, {
+          method: "POST",
+          headers: this.getHeaders(),
+          body: JSON.stringify(provisionData),
+        });
+        if (response.status === HttpStatusCodes.OK) {
+          data = await response.json();
+        }
+      } else {
+        response = await fetch(`${notebookServerEndpoint}/api/shutdown`, {
+          method: "POST",
+          headers: { Authorization: authToken },
+        });
+      }
+      return {
+        status: response.status,
+        data,
+      };
     } catch (error) {
       Logger.logError(getErrorMessage(error), "NotebookContainerClient/resetWorkspace");
-      await this.recreateNotebookWorkspaceAsync();
+      if (!NotebookUtil.isPhoenixEnabled()) {
+        await this.recreateNotebookWorkspaceAsync();
+      }
+      throw error;
     }
   }
 
@@ -175,4 +210,12 @@ export class NotebookContainerClient {
       return Promise.reject(error);
     }
   }
+
+  private getHeaders(): HeadersInit {
+    const authorizationHeader = getAuthorizationHeader();
+    return {
+      [authorizationHeader.header]: authorizationHeader.token,
+      [HttpHeaders.contentType]: "application/json",
+    };
+  }
 }

src/Explorer/Notebook/useNotebook.ts

@@ -8,6 +8,7 @@ import * as Logger from "../../Common/Logger";
 import { configContext } from "../../ConfigContext";
 import * as DataModels from "../../Contracts/DataModels";
 import { ContainerConnectionInfo, ContainerInfo } from "../../Contracts/DataModels";
+import { useTabs } from "../../hooks/useTabs";
 import { IPinnedRepo } from "../../Juno/JunoClient";
 import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
@@ -88,7 +89,7 @@ export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
   isRefreshed: false,
   containerStatus: {
     status: undefined,
-    durationLeftMin: undefined,
+    durationLeftInMinutes: undefined,
     notebookServerInfo: undefined,
   },
   setIsNotebookEnabled: (isNotebookEnabled: boolean) => set({ isNotebookEnabled }),
@@ -279,16 +280,13 @@ export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
   setConnectionInfo: (connectionInfo: ContainerConnectionInfo) => set({ connectionInfo }),
   setIsAllocating: (isAllocating: boolean) => set({ isAllocating }),
   resetContainerConnection: (connectionStatus: ContainerConnectionInfo): void => {
+    useTabs.getState().closeAllNotebookTabs(true);
     useNotebook.getState().setConnectionInfo(connectionStatus);
-    useNotebook.getState().setNotebookServerInfo({
-      notebookServerEndpoint: undefined,
-      authToken: undefined,
-      forwardingId: undefined,
-    });
+    useNotebook.getState().setNotebookServerInfo(undefined);
     useNotebook.getState().setIsAllocating(false);
     useNotebook.getState().setContainerStatus({
       status: undefined,
-      durationLeftMin: undefined,
+      durationLeftInMinutes: undefined,
       notebookServerInfo: undefined,
     });
   },