Phoenix Reconnect Integration (#1123)

* Reconnect integration

* git connection issue

* format issue

* Typo issue

* added constants

* Removed math.round for remainingTime

* code refctor for container status check

* disconnect text change

src/Common/Constants.ts

@@ -347,6 +347,11 @@ export enum ConnectionStatusType {
   ReConnect = "Reconnect",
 }
 
+export enum ContainerStatusType {
+  Active = "Active",
+  InActive = "InActive",
+}
+
 export const EmulatorMasterKey =
   //[SuppressMessage("Microsoft.Security", "CS002:SecretInNextLine", Justification="Well known public masterKey for emulator")]
   "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==";
@@ -357,20 +362,23 @@ export const StyleConstants = require("less-vars-loader!../../less/Common/Consta
 export class Notebook {
   public static readonly defaultBasePath = "./notebooks";
   public static readonly heartbeatDelayMs = 60000;
+  public static readonly containerStatusHeartbeatDelayMs = 30000;
   public static readonly kernelRestartInitialDelayMs = 1000;
   public static readonly kernelRestartMaxDelayMs = 20000;
   public static readonly autoSaveIntervalMs = 120000;
   public static readonly memoryGuageToGB = 1048576;
+  public static readonly lowMemoryBar = 0.8;
+  public static readonly reminingTimeMin = 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.";
   public static readonly cassandraShellTemporarilyDownMsg =
     "We have identified an issue with the Cassandra Shell and it is unavailable right now. We are actively working on the mitigation.";
-  public static saveNotebookModalTitle = "Save Notebook in temporary workspace";
+  public static saveNotebookModalTitle = "Save notebook in temporary workspace";
   public static saveNotebookModalContent =
     "This notebook will be saved in the temporary workspace and will be removed when the session expires. To save your work permanently, save your notebooks to a GitHub repository or download the notebooks to your local machine before the session ends.";
-  public static newNotebookModalTitle = "Create Notebook in temporary workspace";
+  public static newNotebookModalTitle = "Create notebook in temporary workspace";
-  public static newNotebookUploadModalTitle = "Upload Notebook in temporary workspace";
+  public static newNotebookUploadModalTitle = "Upload notebook to temporary workspace";
   public static newNotebookModalContent1 =
     "A temporary workspace will be created to enable you to work with notebooks. When the session expires, any notebooks in the workspace will be removed.";
   public static newNotebookModalContent2 =

src/Contracts/DataModels.ts

@@ -1,4 +1,4 @@
-import { ConnectionStatusType } from "../Common/Constants";
+import { ConnectionStatusType, ContainerStatusType } from "../Common/Constants";
 
 export interface DatabaseAccount {
   id: string;
@@ -426,6 +426,13 @@ export interface OperationStatus {
 export interface NotebookWorkspaceConnectionInfo {
   authToken: string;
   notebookServerEndpoint: string;
+  forwardingId: string;
+}
+
+export interface ContainerInfo {
+  durationLeftMin: number;
+  notebookServerInfo: NotebookWorkspaceConnectionInfo;
+  status: ContainerStatusType;
 }
 
 export interface NotebookWorkspaceFeedResponse {

src/Explorer/Controls/Dialog.tsx

@@ -13,7 +13,6 @@ import {
   Link,
   PrimaryButton,
   ProgressIndicator,
-  Text,
   TextField,
 } from "@fluentui/react";
 import React, { FC } from "react";
@@ -197,7 +196,7 @@ export const Dialog: FC = () => {
           {linkProps.linkText} <FontIcon iconName="NavigateExternalInline" />
         </Link>
       )}
-      {contentHtml && <Text>{contentHtml}</Text>}
+      {contentHtml}
       {progressIndicatorProps && <ProgressIndicator {...progressIndicatorProps} />}
       <DialogFooter>
         <PrimaryButton {...primaryButtonProps} />

src/Explorer/Controls/Notebook/NotebookTerminalComponent.test.tsx

@@ -35,16 +35,19 @@ const testCassandraAccount: DataModels.DatabaseAccount = {
 const testNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo = {
   authToken: "authToken",
   notebookServerEndpoint: "https://testNotebookServerEndpoint.azure.com",
+  forwardingId: "Id",
 };
 
 const testMongoNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo = {
   authToken: "authToken",
   notebookServerEndpoint: "https://testNotebookServerEndpoint.azure.com/mongo",
+  forwardingId: "Id",
 };
 
 const testCassandraNotebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo = {
   authToken: "authToken",
   notebookServerEndpoint: "https://testNotebookServerEndpoint.azure.com/cassandra",
+  forwardingId: "Id",
 };
 
 describe("NotebookTerminalComponent", () => {

src/Explorer/Controls/NotebookViewer/NotebookViewerComponent.tsx

@@ -53,7 +53,7 @@ export class NotebookViewerComponent
     super(props);
 
     this.clientManager = new NotebookClientV2({
-      connectionInfo: { authToken: undefined, notebookServerEndpoint: undefined },
+      connectionInfo: { authToken: undefined, notebookServerEndpoint: undefined, forwardingId: undefined },
       databaseAccountName: undefined,
       defaultExperience: "NotebookViewer",
       isReadOnly: true,

src/Explorer/Explorer.tsx

@@ -3,6 +3,7 @@ import { isPublicInternetAccessAllowed } from "Common/DatabaseAccountUtility";
 import * as ko from "knockout";
 import React from "react";
 import _ from "underscore";
+import shallow from "zustand/shallow";
 import { AuthType } from "../AuthType";
 import { BindingHandlersRegisterer } from "../Bindings/BindingHandlersRegisterer";
 import * as Constants from "../Common/Constants";
@@ -165,11 +166,9 @@ export default class Explorer {
     );
 
     useNotebook.subscribe(
-      async () => {
+      async () => this.initiateAndRefreshNotebookList(),
-        this.initiateAndRefreshNotebookList();
+      (state) => [state.isNotebookEnabled, state.isRefreshed],
-        useNotebook.getState().setIsRefreshed(false);
+      shallow
-      },
-      (state) => state.isNotebookEnabled || state.isRefreshed
     );
 
     this.resourceTree = new ResourceTreeAdapter(this);
@@ -179,6 +178,7 @@ export default class Explorer {
       useNotebook.getState().setNotebookServerInfo({
         notebookServerEndpoint: userContext.features.notebookServerUrl,
         authToken: userContext.features.notebookServerToken,
+        forwardingId: undefined,
       });
     }
 
@@ -364,6 +364,7 @@ export default class Explorer {
       useNotebook.getState().setNotebookServerInfo({
         notebookServerEndpoint: userContext.features.notebookServerUrl || connectionInfo.notebookServerEndpoint,
         authToken: userContext.features.notebookServerToken || connectionInfo.authToken,
+        forwardingId: undefined,
       });
     }
 
@@ -395,11 +396,18 @@ export default class Explorer {
           connectionInfo.data &&
           connectionInfo.data.notebookServerUrl
         ) {
+          const containerData = {
+            forwardingId: connectionInfo.data.forwardingId,
+            dbAccountName: userContext.databaseAccount.name,
+          };
+          await this.phoenixClient.initiateContainerHeartBeat(containerData);
+
           connectionStatus.status = ConnectionStatusType.Connected;
           useNotebook.getState().setConnectionInfo(connectionStatus);
           useNotebook.getState().setNotebookServerInfo({
             notebookServerEndpoint: userContext.features.notebookServerUrl || connectionInfo.data.notebookServerUrl,
             authToken: userContext.features.notebookServerToken || connectionInfo.data.notebookAuthToken,
+            forwardingId: connectionInfo.data.forwardingId,
           });
           this.notebookManager?.notebookClient
             .getMemoryUsage()
@@ -407,11 +415,11 @@ export default class Explorer {
           useNotebook.getState().setIsAllocating(false);
         } else {
           connectionStatus.status = ConnectionStatusType.Failed;
-          useNotebook.getState().resetConatinerConnection(connectionStatus);
+          useNotebook.getState().resetContainerConnection(connectionStatus);
         }
       } catch (error) {
         connectionStatus.status = ConnectionStatusType.Failed;
-        useNotebook.getState().resetConatinerConnection(connectionStatus);
+        useNotebook.getState().resetContainerConnection(connectionStatus);
         throw error;
       }
       this.refreshNotebookList();
@@ -692,7 +700,7 @@ export default class Explorer {
       throw new Error(`Invalid notebookContentItem: ${notebookContentItem}`);
     }
     if (notebookContentItem.type === NotebookContentItemType.Notebook && NotebookUtil.isPhoenixEnabled()) {
-      this.allocateContainer();
+      await this.allocateContainer();
     }
 
     const notebookTabs = useTabs
@@ -1016,8 +1024,8 @@ export default class Explorer {
         useDialog
           .getState()
           .showOkModalDialog(
-            "Failed to Connect",
+            "Failed to connect",
-            "Failed to connect temporary workspace, this could happen because of network issue please refresh and try again."
+            "Failed to connect to temporary workspace. This could happen because of network issues. Please refresh the page and try again."
           );
       }
     } else {

src/Explorer/Menus/CommandBar/ConnectionStatusComponent.tsx

@@ -1,8 +1,20 @@
-import { Icon, ProgressIndicator, Stack, TooltipHost } from "@fluentui/react";
+import {
-import { ActionButton } from "@fluentui/react/lib/Button";
+  FocusTrapCallout,
+  FocusZone,
+  FocusZoneTabbableElements,
+  FontWeights,
+  Icon,
+  mergeStyleSets,
+  ProgressIndicator,
+  Stack,
+  Text,
+  TooltipHost,
+} from "@fluentui/react";
+import { useId } from "@fluentui/react-hooks";
+import { ActionButton, DefaultButton } from "@fluentui/react/lib/Button";
 import * as React from "react";
 import "../../../../less/hostedexplorer.less";
-import { ConnectionStatusType, Notebook } from "../../../Common/Constants";
+import { ConnectionStatusType, ContainerStatusType, Notebook } from "../../../Common/Constants";
 import Explorer from "../../Explorer";
 import { useNotebook } from "../../Notebook/useNotebook";
 import "../CommandBar/ConnectionStatusComponent.less";
@@ -16,6 +28,26 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
   const [counter, setCounter] = React.useState(0);
   const [statusColor, setStatusColor] = React.useState("");
   const [toolTipContent, setToolTipContent] = React.useState("Connect to temporary workspace.");
+  const [isBarDismissed, setIsBarDismissed] = React.useState<boolean>(false);
+  const buttonId = useId("callout-button");
+  const containerInfo = useNotebook((state) => state.containerStatus);
+
+  const styles = mergeStyleSets({
+    callout: {
+      width: 320,
+      padding: "20px 24px",
+    },
+    title: {
+      marginBottom: 12,
+      fontWeight: FontWeights.semilight,
+    },
+    buttons: {
+      display: "flex",
+      justifyContent: "flex-end",
+      marginTop: 20,
+    },
+  });
+
   React.useEffect(() => {
     let intervalId: NodeJS.Timeout;
 
@@ -78,13 +110,25 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
     setToolTipContent("Click here to Reconnect to temporary workspace.");
   }
   return (
+    <>
+      <TooltipHost
+        content={
+          containerInfo.status &&
+          containerInfo.status === ContainerStatusType.Active &&
+          containerInfo.durationLeftMin <= Notebook.reminingTimeMin
+            ? `Connected to temporary workspace. This temporary workspace will get deleted in ${Math.round(
+                containerInfo.durationLeftMin
+              )} minutes.`
+            : toolTipContent
+        }
+      >
         <ActionButton
+          id={buttonId}
           className={connectionInfo.status === ConnectionStatusType.Failed ? "commandReactBtn" : "connectedReactBtn"}
           onClick={(e: React.MouseEvent<HTMLSpanElement>) =>
             connectionInfo.status === ConnectionStatusType.Failed ? container.allocateContainer() : e.preventDefault()
           }
         >
-      <TooltipHost content={toolTipContent}>
           <Stack className="connectionStatusContainer" horizontal>
             <i className={statusColor}></i>
             <span className={connectionInfo.status === ConnectionStatusType.Failed ? "connectionStatusFailed" : ""}>
@@ -95,13 +139,41 @@ export const ConnectionStatus: React.FC<Props> = ({ container }: Props): JSX.Ele
             )}
             {connectionInfo.status === ConnectionStatusType.Connected && !isActive && (
               <ProgressIndicator
-              className={usedGB / totalGB > 0.8 ? "lowMemory" : ""}
+                className={usedGB / totalGB > Notebook.lowMemoryBar ? "lowMemory" : ""}
                 description={usedGB.toFixed(1) + " of " + totalGB.toFixed(1) + " GB"}
                 percentComplete={usedGB / totalGB}
               />
             )}
           </Stack>
-      </TooltipHost>
+          {!isBarDismissed &&
+          containerInfo.status &&
+          containerInfo.status === ContainerStatusType.Active &&
+          containerInfo.durationLeftMin <= Notebook.reminingTimeMin ? (
+            <FocusTrapCallout
+              role="alertdialog"
+              className={styles.callout}
+              gapSpace={0}
+              target={`#${buttonId}`}
+              onDismiss={() => setIsBarDismissed(true)}
+              setInitialFocus
+            >
+              <Text block variant="xLarge" className={styles.title}>
+                Remaining time
+              </Text>
+              <Text block variant="small">
+                This temporary workspace will get deleted in {Math.round(containerInfo.durationLeftMin)} minutes. To
+                save your work permanently, save your notebooks to a GitHub repository or download the 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>
+                </Stack>
+              </FocusZone>
+            </FocusTrapCallout>
+          ) : undefined}
         </ActionButton>
+      </TooltipHost>
+    </>
   );
 };

src/Explorer/Notebook/NotebookContainerClient.ts

@@ -59,6 +59,7 @@ export class NotebookContainerClient {
 
     const { notebookServerEndpoint, authToken } = this.getNotebookServerConfig();
     try {
+      if (this.checkStatus()) {
         const response = await fetch(`${notebookServerEndpoint}api/metrics/memory`, {
           method: "GET",
           headers: {
@@ -78,12 +79,7 @@ export class NotebookContainerClient {
               freeKB: memoryUsageInfo.free,
             };
           }
-      } else if (NotebookUtil.isPhoenixEnabled()) {
+        }
-        const connectionStatus: ContainerConnectionInfo = {
-          status: ConnectionStatusType.ReConnect,
-        };
-        useNotebook.getState().resetConatinerConnection(connectionStatus);
-        useNotebook.getState().setIsRefreshed(true);
       }
       return undefined;
     } catch (error) {
@@ -97,14 +93,30 @@ export class NotebookContainerClient {
         const connectionStatus: ContainerConnectionInfo = {
           status: ConnectionStatusType.Failed,
         };
-        useNotebook.getState().resetConatinerConnection(connectionStatus);
+        useNotebook.getState().resetContainerConnection(connectionStatus);
-        useNotebook.getState().setIsRefreshed(true);
+        useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
       }
       this.onConnectionLost();
       return undefined;
     }
   }
 
+  private checkStatus(): boolean {
+    if (NotebookUtil.isPhoenixEnabled()) {
+      if (
+        useNotebook.getState().containerStatus?.status &&
+        useNotebook.getState().containerStatus?.status === Constants.ContainerStatusType.InActive
+      ) {
+        const connectionStatus: ContainerConnectionInfo = {
+          status: ConnectionStatusType.ReConnect,
+        };
+        useNotebook.getState().resetContainerConnection(connectionStatus);
+        useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
+        return false;
+      }
+    }
+    return true;
+  }
   public async resetWorkspace(): Promise<void> {
     this.isResettingWorkspace = true;
     try {

src/Explorer/Notebook/notebookClientV2.test.ts

@@ -35,6 +35,7 @@ describe("auto start kernel", () => {
         connectionInfo: {
           authToken: "autToken",
           notebookServerEndpoint: "notebookServerEndpoint",
+          forwardingId: "Id",
         },
         databaseAccountName: undefined,
         defaultExperience: undefined,

src/Explorer/Notebook/useNotebook.ts

@@ -7,7 +7,7 @@ import { getErrorMessage } from "../../Common/ErrorHandlingUtils";
 import * as Logger from "../../Common/Logger";
 import { configContext } from "../../ConfigContext";
 import * as DataModels from "../../Contracts/DataModels";
-import { ContainerConnectionInfo } from "../../Contracts/DataModels";
+import { ContainerConnectionInfo, ContainerInfo } from "../../Contracts/DataModels";
 import { IPinnedRepo } from "../../Juno/JunoClient";
 import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
@@ -35,6 +35,7 @@ interface NotebookState {
   notebookFolderName: string;
   isAllocating: boolean;
   isRefreshed: boolean;
+  containerStatus: ContainerInfo;
   setIsNotebookEnabled: (isNotebookEnabled: boolean) => void;
   setIsNotebooksEnabledForAccount: (isNotebooksEnabledForAccount: boolean) => void;
   setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => void;
@@ -53,8 +54,9 @@ interface NotebookState {
   initializeGitHubRepos: (pinnedRepos: IPinnedRepo[]) => void;
   setConnectionInfo: (connectionInfo: ContainerConnectionInfo) => void;
   setIsAllocating: (isAllocating: boolean) => void;
-  resetConatinerConnection: (connectionStatus: ContainerConnectionInfo) => void;
+  resetContainerConnection: (connectionStatus: ContainerConnectionInfo) => void;
   setIsRefreshed: (isAllocating: boolean) => void;
+  setContainerStatus: (containerStatus: ContainerInfo) => void;
 }
 
 export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
@@ -63,6 +65,7 @@ export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
   notebookServerInfo: {
     notebookServerEndpoint: undefined,
     authToken: undefined,
+    forwardingId: undefined,
   },
   sparkClusterConnectionInfo: {
     userName: undefined,
@@ -83,6 +86,11 @@ export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
   notebookFolderName: undefined,
   isAllocating: false,
   isRefreshed: false,
+  containerStatus: {
+    status: undefined,
+    durationLeftMin: undefined,
+    notebookServerInfo: undefined,
+  },
   setIsNotebookEnabled: (isNotebookEnabled: boolean) => set({ isNotebookEnabled }),
   setIsNotebooksEnabledForAccount: (isNotebooksEnabledForAccount: boolean) => set({ isNotebooksEnabledForAccount }),
   setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
@@ -270,13 +278,20 @@ export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
   },
   setConnectionInfo: (connectionInfo: ContainerConnectionInfo) => set({ connectionInfo }),
   setIsAllocating: (isAllocating: boolean) => set({ isAllocating }),
-  resetConatinerConnection: (connectionStatus: ContainerConnectionInfo): void => {
+  resetContainerConnection: (connectionStatus: ContainerConnectionInfo): void => {
     useNotebook.getState().setConnectionInfo(connectionStatus);
     useNotebook.getState().setNotebookServerInfo({
       notebookServerEndpoint: undefined,
       authToken: undefined,
+      forwardingId: undefined,
     });
     useNotebook.getState().setIsAllocating(false);
+    useNotebook.getState().setContainerStatus({
+      status: undefined,
+      durationLeftMin: undefined,
+      notebookServerInfo: undefined,
+    });
   },
   setIsRefreshed: (isRefreshed: boolean) => set({ isRefreshed }),
+  setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }),
 }));

src/Explorer/Tabs/TerminalTab.tsx

@@ -100,6 +100,7 @@ export default class TerminalTab extends TabsBase {
     return {
       authToken: info.authToken,
       notebookServerEndpoint: `${info.notebookServerEndpoint.replace(/\/+$/, "")}/${endpointSuffix}`,
+      forwardingId: info.forwardingId,
     };
   }
 }

src/Phoenix/PhoenixClient.ts

@@ -1,5 +1,9 @@
-import { HttpHeaders, HttpStatusCodes } from "../Common/Constants";
+import { ContainerStatusType, HttpHeaders, HttpStatusCodes, Notebook } from "../Common/Constants";
+import { getErrorMessage } from "../Common/ErrorHandlingUtils";
+import * as Logger from "../Common/Logger";
 import { configContext } from "../ConfigContext";
+import { ContainerInfo } from "../Contracts/DataModels";
+import { useNotebook } from "../Explorer/Notebook/useNotebook";
 import { userContext } from "../UserContext";
 import { getAuthorizationHeader } from "../Utils/AuthorizationUtils";
 
@@ -10,17 +14,24 @@ export interface IPhoenixResponse<T> {
 export interface IPhoenixConnectionInfoResult {
   readonly notebookAuthToken?: string;
   readonly notebookServerUrl?: string;
+  readonly forwardingId?: string;
 }
-export interface IProvosionData {
+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: IProvosionData
+    provisionData: IProvisionData
   ): Promise<IPhoenixResponse<IPhoenixConnectionInfoResult>> {
     try {
       const response = await window.fetch(`${this.getPhoenixContainerPoolingEndPoint()}/allocate`, {
@@ -41,6 +52,60 @@ export class PhoenixClient {
       throw error;
     }
   }
+  public async initiateContainerHeartBeat(containerData: { forwardingId: string; dbAccountName: string }) {
+    this.getContainerHealth(Notebook.containerStatusHeartbeatDelayMs, containerData);
+  }
+
+  private scheduleContainerHeartbeat(delayMs: number, containerData: IContainerData): void {
+    setTimeout(() => {
+      this.getContainerHealth(delayMs, containerData);
+    }, delayMs);
+  }
+
+  private async getContainerStatusAsync(containerData: IContainerData): Promise<ContainerInfo> {
+    try {
+      const response = await window.fetch(
+        `${this.getPhoenixContainerPoolingEndPoint()}/${containerData.dbAccountName}/${containerData.forwardingId}`,
+        {
+          method: "GET",
+          headers: PhoenixClient.getHeaders(),
+        }
+      );
+      if (response.status === HttpStatusCodes.OK) {
+        const containerStatus = await response.json();
+        return {
+          durationLeftMin: containerStatus.durationLeftInMinutes,
+          notebookServerInfo: containerStatus.notebookServerInfo,
+          status: ContainerStatusType.Active,
+        };
+      }
+      return {
+        durationLeftMin: undefined,
+        notebookServerInfo: undefined,
+        status: ContainerStatusType.InActive,
+      };
+    } catch (error) {
+      Logger.logError(getErrorMessage(error), "PhoenixClient/getContainerStatus");
+      return {
+        durationLeftMin: undefined,
+        notebookServerInfo: undefined,
+        status: ContainerStatusType.InActive,
+      };
+    }
+  }
+
+  private getContainerHealth(delayMs: number, containerData: { forwardingId: string; dbAccountName: string }) {
+    this.getContainerStatusAsync(containerData)
+      .then((ContainerInfo) => useNotebook.getState().setContainerStatus(ContainerInfo))
+      .finally(() => {
+        if (
+          useNotebook.getState().containerStatus.status &&
+          useNotebook.getState().containerStatus.status === ContainerStatusType.Active
+        ) {
+          this.scheduleContainerHeartbeat(delayMs, containerData);
+        }
+      });
+  }
 
   public static getPhoenixEndpoint(): string {
     const phoenixEndpoint =

src/Utils/GalleryUtils.ts

@@ -239,7 +239,7 @@ export function downloadItem(
         useDialog
           .getState()
           .showOkModalDialog(
-            "Failed to Connect",
+            "Failed to connect",
             "Failed to connect to temporary workspace. Please refresh the page and try again."
           );
       }