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/Common/Constants.ts

@@ -344,12 +344,12 @@ export enum ConnectionStatusType {
   Connecting = "Connecting",
   Connected = "Connected",
   Failed = "Connection Failed",
-  ReConnect = "Reconnect",
+  Reconnect = "Reconnect",
 }
 
 export enum ContainerStatusType {
   Active = "Active",
-  InActive = "InActive",
+  Disconnected = "Disconnected",
 }
 
 export const EmulatorMasterKey =
@@ -367,8 +367,8 @@ export class Notebook {
   public static readonly kernelRestartMaxDelayMs = 20000;
   public static readonly autoSaveIntervalMs = 120000;
   public static readonly memoryGuageToGB = 1048576;
-  public static readonly lowMemoryBar = 0.8;
+  public static readonly lowMemoryThreshold = 0.8;
-  public static readonly reminingTimeMin = 10;
+  public static readonly remainingTimeForAlert = 10;
   public static readonly temporarilyDownMsg = "Notebooks is currently not available. We are working on it.";
   public static readonly mongoShellTemporarilyDownMsg =
     "We have identified an issue with the Mongo Shell and it is unavailable right now. We are actively working on the mitigation.";

src/Contracts/DataModels.ts

@@ -430,11 +430,35 @@ export interface NotebookWorkspaceConnectionInfo {
 }
 
 export interface ContainerInfo {
-  durationLeftMin: number;
+  durationLeftInMinutes: number;
   notebookServerInfo: NotebookWorkspaceConnectionInfo;
   status: ContainerStatusType;
 }
 
+export interface IProvisionData {
+  aadToken: string;
+  subscriptionId: string;
+  resourceGroup: string;
+  dbAccountName: string;
+  cosmosEndpoint: string;
+}
+
+export interface IContainerData {
+  dbAccountName: string;
+  forwardingId: string;
+}
+
+export interface IResponse<T> {
+  status: number;
+  data: T;
+}
+
+export interface IPhoenixConnectionInfoResult {
+  readonly notebookAuthToken?: string;
+  readonly notebookServerUrl?: string;
+  readonly forwardingId?: string;
+}
+
 export interface NotebookWorkspaceFeedResponse {
   value: NotebookWorkspace[];
 }

src/Explorer/Explorer.tsx

@@ -14,7 +14,7 @@ import { getErrorMessage, getErrorStack, handleError } from "../Common/ErrorHand
 import * as Logger from "../Common/Logger";
 import { QueriesClient } from "../Common/QueriesClient";
 import * as DataModels from "../Contracts/DataModels";
-import { ContainerConnectionInfo } from "../Contracts/DataModels";
+import { ContainerConnectionInfo, IPhoenixConnectionInfoResult, IResponse } from "../Contracts/DataModels";
 import * as ViewModels from "../Contracts/ViewModels";
 import { GitHubOAuthService } from "../GitHub/GitHubOAuthService";
 import { useSidePanel } from "../hooks/useSidePanel";
@@ -376,7 +376,11 @@ export default class Explorer {
   public async allocateContainer(): Promise<void> {
     const notebookServerInfo = useNotebook.getState().notebookServerInfo;
     const isAllocating = useNotebook.getState().isAllocating;
-    if (isAllocating === false && notebookServerInfo && notebookServerInfo.notebookServerEndpoint === undefined) {
+    if (
+      isAllocating === false &&
+      (notebookServerInfo === undefined ||
+        (notebookServerInfo && notebookServerInfo.notebookServerEndpoint === undefined))
+    ) {
       const provisionData = {
         aadToken: userContext.authorizationToken,
         subscriptionId: userContext.subscriptionId,
@@ -391,11 +395,23 @@ export default class Explorer {
       try {
         useNotebook.getState().setIsAllocating(true);
         const connectionInfo = await this.phoenixClient.containerConnectionInfo(provisionData);
-        if (
+        await this.setNotebookInfo(connectionInfo, connectionStatus);
-          connectionInfo.status === HttpStatusCodes.OK &&
+      } catch (error) {
-          connectionInfo.data &&
+        connectionStatus.status = ConnectionStatusType.Failed;
-          connectionInfo.data.notebookServerUrl
+        useNotebook.getState().resetContainerConnection(connectionStatus);
+        throw error;
+      }
+      this.refreshNotebookList();
+
+      this._isInitializingNotebooks = false;
+    }
+  }
+
+  private async setNotebookInfo(
+    connectionInfo: IResponse<IPhoenixConnectionInfoResult>,
+    connectionStatus: DataModels.ContainerConnectionInfo
   ) {
+    if (connectionInfo.status === HttpStatusCodes.OK && connectionInfo.data && connectionInfo.data.notebookServerUrl) {
       const containerData = {
         forwardingId: connectionInfo.data.forwardingId,
         dbAccountName: userContext.databaseAccount.name,
@@ -412,20 +428,11 @@ export default class Explorer {
       this.notebookManager?.notebookClient
         .getMemoryUsage()
         .then((memoryUsageInfo) => useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo));
-          useNotebook.getState().setIsAllocating(false);
     } else {
       connectionStatus.status = ConnectionStatusType.Failed;
       useNotebook.getState().resetContainerConnection(connectionStatus);
     }
-      } catch (error) {
+    useNotebook.getState().setIsAllocating(false);
-        connectionStatus.status = ConnectionStatusType.Failed;
-        useNotebook.getState().resetContainerConnection(connectionStatus);
-        throw error;
-      }
-      this.refreshNotebookList();
-
-      this._isInitializingNotebooks = false;
-    }
   }
 
   public resetNotebookWorkspace(): void {
@@ -436,11 +443,14 @@ export default class Explorer {
       );
       return;
     }
+    const dialogContent = NotebookUtil.isPhoenixEnabled()
+      ? "Notebooks saved in the temporary workspace will be deleted. Do you want to proceed?"
+      : "This lets you keep your notebook files and the workspace will be restored to default. Proceed anyway?";
 
     const resetConfirmationDialogProps: DialogProps = {
       isModal: true,
       title: "Reset Workspace",
-      subText: "This lets you keep your notebook files and the workspace will be restored to default. Proceed anyway?",
+      subText: dialogContent,
       primaryButtonText: "OK",
       secondaryButtonText: "Cancel",
       onPrimaryButtonClick: this._resetNotebookWorkspace,
@@ -498,16 +508,55 @@ export default class Explorer {
   private _resetNotebookWorkspace = async () => {
     useDialog.getState().closeDialog();
     const clearInProgressMessage = logConsoleProgress("Resetting notebook workspace");
+    let connectionStatus: ContainerConnectionInfo;
     try {
-      await this.notebookManager?.notebookClient.resetWorkspace();
+      const notebookServerInfo = useNotebook.getState().notebookServerInfo;
+      if (!notebookServerInfo || !notebookServerInfo.notebookServerEndpoint) {
+        const error = "No server endpoint detected";
+        Logger.logError(error, "NotebookContainerClient/resetWorkspace");
+        logConsoleError(error);
+        return;
+      }
+
+      if (NotebookUtil.isPhoenixEnabled()) {
+        useTabs.getState().closeAllNotebookTabs(true);
+        connectionStatus = {
+          status: ConnectionStatusType.Connecting,
+        };
+        useNotebook.getState().setConnectionInfo(connectionStatus);
+      }
+      const connectionInfo = await this.notebookManager?.notebookClient.resetWorkspace();
+      if (connectionInfo && connectionInfo.status && connectionInfo.status === HttpStatusCodes.OK) {
+        if (NotebookUtil.isPhoenixEnabled() && connectionInfo.data && connectionInfo.data.notebookServerUrl) {
+          await this.setNotebookInfo(connectionInfo, connectionStatus);
+          useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
+        }
         logConsoleInfo("Successfully reset notebook workspace");
         TelemetryProcessor.traceSuccess(Action.ResetNotebookWorkspace);
+      } else {
+        logConsoleError(`Failed to reset notebook workspace`);
+        TelemetryProcessor.traceFailure(Action.ResetNotebookWorkspace);
+        if (NotebookUtil.isPhoenixEnabled()) {
+          connectionStatus = {
+            status: ConnectionStatusType.Reconnect,
+          };
+          useNotebook.getState().resetContainerConnection(connectionStatus);
+          useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
+        }
+      }
     } catch (error) {
       logConsoleError(`Failed to reset notebook workspace: ${error}`);
       TelemetryProcessor.traceFailure(Action.ResetNotebookWorkspace, {
         error: getErrorMessage(error),
         errorStack: getErrorStack(error),
       });
+      if (NotebookUtil.isPhoenixEnabled()) {
+        connectionStatus = {
+          status: ConnectionStatusType.Failed,
+        };
+        useNotebook.getState().setConnectionInfo(connectionStatus);
+        useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
+      }
       throw error;
     } finally {
       clearInProgressMessage();

src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx

@@ -80,9 +80,7 @@ export function createStaticCommandBarButtons(
     }
 
     notebookButtons.push(createOpenTerminalButton(container));
-    if (userContext.features.phoenix === false) {
     notebookButtons.push(createNotebookWorkspaceResetButton(container));
-    }
     if (
       (userContext.apiType === "Mongo" &&
         useNotebook.getState().isShellEnabled &&

src/Explorer/Menus/CommandBar/ConnectionStatusComponent.tsx

@@ -22,6 +22,7 @@ interface Props {
   container: Explorer;
 }
 export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Element => {
+  const connectionInfo = useNotebook((state) => state.connectionInfo);
   const [second, setSecond] = React.useState("00");
   const [minute, setMinute] = React.useState("00");
   const [isActive, setIsActive] = React.useState(false);
@@ -67,6 +68,12 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
     return () => clearInterval(intervalId);
   }, [isActive, counter]);
 
+  React.useEffect(() => {
+    if (connectionInfo && connectionInfo.status === ConnectionStatusType.Reconnect) {
+      setToolTipContent("Click here to Reconnect to temporary workspace.");
+    }
+  }, [connectionInfo.status]);
+
   const stopTimer = () => {
     setIsActive(false);
     setCounter(0);
@@ -74,15 +81,13 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
     setMinute("00");
   };
 
-  const connectionInfo = useNotebook((state) => state.connectionInfo);
   const memoryUsageInfo = useNotebook((state) => state.memoryUsageInfo);
-
   const totalGB = memoryUsageInfo ? memoryUsageInfo.totalKB / Notebook.memoryGuageToGB : 0;
   const usedGB = totalGB > 0 ? totalGB - memoryUsageInfo.freeKB / Notebook.memoryGuageToGB : 0;
 
   if (
     connectionInfo &&
-    (connectionInfo.status === ConnectionStatusType.Connect || connectionInfo.status === ConnectionStatusType.ReConnect)
+    (connectionInfo.status === ConnectionStatusType.Connect || connectionInfo.status === ConnectionStatusType.Reconnect)
   ) {
     return (
       <ActionButton className="commandReactBtn" onClick={() => container.allocateContainer()}>
@@ -113,11 +118,10 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
     <>
       <TooltipHost
         content={
-          containerInfo.status &&
+          containerInfo?.status === ContainerStatusType.Active &&
-          containerInfo.status === ContainerStatusType.Active &&
+          Math.round(containerInfo.durationLeftInMinutes) <= Notebook.remainingTimeForAlert
-          containerInfo.durationLeftMin <= Notebook.reminingTimeMin
+            ? `Connected to temporary workspace. This temporary workspace will get disconnected in ${Math.round(
-            ? `Connected to temporary workspace. This temporary workspace will get deleted in ${Math.round(
+                containerInfo.durationLeftInMinutes
-                containerInfo.durationLeftMin
               )} minutes.`
             : toolTipContent
         }
@@ -139,16 +143,16 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
             )}
             {connectionInfo.status === ConnectionStatusType.Connected && !isActive && (
               <ProgressIndicator
-                className={usedGB / totalGB > Notebook.lowMemoryBar ? "lowMemory" : ""}
+                className={totalGB !== 0 && usedGB / totalGB > 0.8 ? "lowMemory" : ""}
                 description={usedGB.toFixed(1) + " of " + totalGB.toFixed(1) + " GB"}
-                percentComplete={usedGB / totalGB}
+                percentComplete={totalGB !== 0 ? usedGB / totalGB : 0}
               />
             )}
           </Stack>
           {!isBarDismissed &&
           containerInfo.status &&
           containerInfo.status === ContainerStatusType.Active &&
-          containerInfo.durationLeftMin <= Notebook.reminingTimeMin ? (
+          Math.round(containerInfo.durationLeftInMinutes) <= Notebook.remainingTimeForAlert ? (
             <FocusTrapCallout
               role="alertdialog"
               className={styles.callout}
@@ -158,16 +162,16 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
               setInitialFocus
             >
               <Text block variant="xLarge" className={styles.title}>
-                Remaining time
+                Remaining Time
               </Text>
               <Text block variant="small">
-                This temporary workspace will get deleted in {Math.round(containerInfo.durationLeftMin)} minutes. To
+                This temporary workspace will get disconnected in {Math.round(containerInfo.durationLeftInMinutes)}{" "}
-                save your work permanently, save your notebooks to a GitHub repository or download the notebooks to your
+                minutes. To save your work permanently, save your notebooks to a GitHub repository or download the
-                local machine before the session ends.
+                notebooks to your local machine before the session ends.
               </Text>
               <FocusZone handleTabKey={FocusZoneTabbableElements.all} isCircularNavigation>
                 <Stack className={styles.buttons} gap={8} horizontal>
-                  <DefaultButton onClick={() => setIsBarDismissed(true)}>Dismiss</DefaultButton>
+                  <DefaultButton onClick={() => setIsBarDismissed(true)}>Dimiss</DefaultButton>
                 </Stack>
               </FocusZone>
             </FocusTrapCallout>

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 (
+      if (useNotebook.getState().containerStatus?.status === Constants.ContainerStatusType.Disconnected) {
-        useNotebook.getState().containerStatus?.status &&
-        useNotebook.getState().containerStatus?.status === Constants.ContainerStatusType.InActive
-      ) {
         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,14 +145,41 @@ export class NotebookContainerClient {
 
     const { notebookServerEndpoint, authToken } = this.getNotebookServerConfig();
     try {
-      await fetch(`${notebookServerEndpoint}/api/shutdown`, {
+      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");
+      if (!NotebookUtil.isPhoenixEnabled()) {
         await this.recreateNotebookWorkspaceAsync();
       }
+      throw error;
+    }
   }
 
   private getNotebookServerConfig(): { notebookServerEndpoint: string; authToken: string } {
@@ -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({
+    useNotebook.getState().setNotebookServerInfo(undefined);
-      notebookServerEndpoint: undefined,
-      authToken: undefined,
-      forwardingId: undefined,
-    });
     useNotebook.getState().setIsAllocating(false);
     useNotebook.getState().setContainerStatus({
       status: undefined,
-      durationLeftMin: undefined,
+      durationLeftInMinutes: undefined,
       notebookServerInfo: undefined,
     });
   },

src/Explorer/Tabs/NotebookTabBase.ts

@@ -17,6 +17,7 @@ export interface NotebookTabBaseOptions extends ViewModels.TabOptions {
 
 /**
  * Every notebook-based tab inherits from this class. It holds the static reference to a notebook client (singleton)
+ * Re-initiating the constructor when ever a new container got allocated.
  */
 export default class NotebookTabBase extends TabsBase {
   protected static clientManager: NotebookClientV2;
@@ -27,6 +28,15 @@ export default class NotebookTabBase extends TabsBase {
 
     this.container = options.container;
 
+    useNotebook.subscribe(
+      () => {
+        const notebookServerInfo = useNotebook.getState().notebookServerInfo;
+        if (notebookServerInfo && notebookServerInfo.notebookServerEndpoint) {
+          NotebookTabBase.clientManager = undefined;
+        }
+      },
+      (state) => state.notebookServerInfo
+    );
     if (!NotebookTabBase.clientManager) {
       NotebookTabBase.clientManager = new NotebookClientV2({
         connectionInfo: useNotebook.getState().notebookServerInfo,

src/Explorer/Tabs/NotebookV2Tab.ts

@@ -53,14 +53,16 @@ export default class NotebookTabV2 extends NotebookTabBase {
       onUpdateKernelInfo: this.onKernelUpdate,
     });
   }
-
+  /*
-  public onCloseTabButtonClick(): Q.Promise<any> {
+   * Hard cleaning the workspace(Closing tabs connected with old container connection) when new container got allocated.
+   */
+  public onCloseTabButtonClick(hardClose = false): Q.Promise<any> {
     const cleanup = () => {
       this.notebookComponentAdapter.notebookShutdown();
       super.onCloseTabButtonClick();
     };
 
-    if (this.notebookComponentAdapter.isContentDirty()) {
+    if (this.notebookComponentAdapter.isContentDirty() && hardClose === false) {
       useDialog
         .getState()
         .showOkCancelModalDialog(

src/Explorer/Tree/Collection.ts

@@ -1,4 +1,5 @@
 import { Resource, StoredProcedureDefinition, TriggerDefinition, UserDefinedFunctionDefinition } from "@azure/cosmos";
+import { NotebookUtil } from "Explorer/Notebook/NotebookUtil";
 import * as ko from "knockout";
 import * as _ from "underscore";
 import * as Constants from "../../Common/Constants";
@@ -528,6 +529,9 @@ export default class Collection implements ViewModels.Collection {
   };
 
   public onSchemaAnalyzerClick = async () => {
+    if (NotebookUtil.isPhoenixEnabled()) {
+      await this.container.allocateContainer();
+    }
     useSelectedNode.getState().setSelectedNode(this);
     this.selectedSubnodeKind(ViewModels.CollectionTabKind.SchemaAnalyzer);
     const SchemaAnalyzerTab = await (await import("../Tabs/SchemaAnalyzerTab")).default;

src/Phoenix/PhoenixClient.ts

@@ -2,37 +2,21 @@ import { ContainerStatusType, HttpHeaders, HttpStatusCodes, Notebook } from "../
 import { getErrorMessage } from "../Common/ErrorHandlingUtils";
 import * as Logger from "../Common/Logger";
 import { configContext } from "../ConfigContext";
-import { ContainerInfo } from "../Contracts/DataModels";
+import {
+  ContainerInfo,
+  IContainerData,
+  IPhoenixConnectionInfoResult,
+  IProvisionData,
+  IResponse,
+} from "../Contracts/DataModels";
 import { useNotebook } from "../Explorer/Notebook/useNotebook";
 import { userContext } from "../UserContext";
 import { getAuthorizationHeader } from "../Utils/AuthorizationUtils";
 
-export interface IPhoenixResponse<T> {
-  status: number;
-  data: T;
-}
-export interface IPhoenixConnectionInfoResult {
-  readonly notebookAuthToken?: string;
-  readonly notebookServerUrl?: string;
-  readonly forwardingId?: string;
-}
-export interface IProvisionData {
-  cosmosEndpoint: string;
-  dbAccountName: string;
-  aadToken: string;
-  resourceGroup: string;
-  subscriptionId: string;
-}
-
-export interface IContainerData {
-  dbAccountName: string;
-  forwardingId: string;
-}
-
 export class PhoenixClient {
   public async containerConnectionInfo(
     provisionData: IProvisionData
-  ): Promise<IPhoenixResponse<IPhoenixConnectionInfoResult>> {
+  ): Promise<IResponse<IPhoenixConnectionInfoResult>> {
     try {
       const response = await window.fetch(`${this.getPhoenixContainerPoolingEndPoint()}/allocate`, {
         method: "POST",
@@ -52,6 +36,7 @@ export class PhoenixClient {
       throw error;
     }
   }
+
   public async initiateContainerHeartBeat(containerData: { forwardingId: string; dbAccountName: string }) {
     this.getContainerHealth(Notebook.containerStatusHeartbeatDelayMs, containerData);
   }
@@ -74,22 +59,22 @@ export class PhoenixClient {
       if (response.status === HttpStatusCodes.OK) {
         const containerStatus = await response.json();
         return {
-          durationLeftMin: containerStatus.durationLeftInMinutes,
+          durationLeftInMinutes: containerStatus?.durationLeftInMinutes,
-          notebookServerInfo: containerStatus.notebookServerInfo,
+          notebookServerInfo: containerStatus?.notebookServerInfo,
           status: ContainerStatusType.Active,
         };
       }
       return {
-        durationLeftMin: undefined,
+        durationLeftInMinutes: undefined,
         notebookServerInfo: undefined,
-        status: ContainerStatusType.InActive,
+        status: ContainerStatusType.Disconnected,
       };
     } catch (error) {
       Logger.logError(getErrorMessage(error), "PhoenixClient/getContainerStatus");
       return {
-        durationLeftMin: undefined,
+        durationLeftInMinutes: undefined,
         notebookServerInfo: undefined,
-        status: ContainerStatusType.InActive,
+        status: ContainerStatusType.Disconnected,
       };
     }
   }
@@ -98,10 +83,7 @@ export class PhoenixClient {
     this.getContainerStatusAsync(containerData)
       .then((ContainerInfo) => useNotebook.getState().setContainerStatus(ContainerInfo))
       .finally(() => {
-        if (
+        if (useNotebook.getState().containerStatus?.status === ContainerStatusType.Active) {
-          useNotebook.getState().containerStatus.status &&
-          useNotebook.getState().containerStatus.status === ContainerStatusType.Active
-        ) {
           this.scheduleContainerHeartbeat(delayMs, containerData);
         }
       });

src/hooks/useTabs.ts

@@ -1,5 +1,7 @@
 import create, { UseStore } from "zustand";
 import * as ViewModels from "../Contracts/ViewModels";
+import { CollectionTabKind } from "../Contracts/ViewModels";
+import NotebookTabV2 from "../Explorer/Tabs/NotebookV2Tab";
 import TabsBase from "../Explorer/Tabs/TabsBase";
 
 interface TabsState {
@@ -12,6 +14,7 @@ interface TabsState {
   refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void;
   closeTabsByComparator: (comparator: (tab: TabsBase) => boolean) => void;
   closeTab: (tab: TabsBase) => void;
+  closeAllNotebookTabs: (hardClose: boolean) => void;
 }
 
 export const useTabs: UseStore<TabsState> = create((set, get) => ({
@@ -78,4 +81,31 @@ export const useTabs: UseStore<TabsState> = create((set, get) => ({
 
     set({ openedTabs: updatedTabs });
   },
+  closeAllNotebookTabs: (hardClose): void => {
+    const isNotebook = (tabKind: CollectionTabKind): boolean => {
+      if (
+        tabKind === CollectionTabKind.Notebook ||
+        tabKind === CollectionTabKind.NotebookV2 ||
+        tabKind === CollectionTabKind.SchemaAnalyzer ||
+        tabKind === CollectionTabKind.Terminal
+      ) {
+        return true;
+      }
+      return false;
+    };
+
+    const tabList = get().openedTabs;
+    if (tabList && tabList.length > 0) {
+      tabList.forEach((tab: NotebookTabV2) => {
+        const tabKind: CollectionTabKind = tab.tabKind;
+        if (tabKind && isNotebook(tabKind)) {
+          tab.onCloseTabButtonClick(hardClose);
+        }
+      });
+
+      if (get().openedTabs.length === 0) {
+        set({ activeTab: undefined });
+      }
+    }
+  },
 }));