fix(UI): workspace restart button stops build before starting a new one (#7301)

* feat(UI): add workspace restart button (#7137)

* Refactor primary buttons

* refactor(site): Always show the main actions

* Remove tests that are testes on Storybook

* Fix tests

* Fix keys

* added restart btn

---------

Co-authored-by: BrunoQuaresma <bruno_nonato_quaresma@hotmail.com>

* added restart hook

* added error handling

* going back to chaining in success callback

* add restarting btn

* added test

* PR feedback

---------

Co-authored-by: BrunoQuaresma <bruno_nonato_quaresma@hotmail.com>

site/src/api/api.ts

@@ -3,6 +3,7 @@ import dayjs from "dayjs"
 import * as Types from "./types"
 import { DeploymentConfig } from "./types"
 import * as TypesGen from "./typesGenerated"
+import { delay } from "utils/delay"
 
 // Adds 304 for the default axios validateStatus function
 // https://github.com/axios/axios#handling-errors Check status here
@@ -476,6 +477,35 @@ export const getWorkspaceByOwnerAndName = async (
   return response.data
 }
 
+export function waitForBuild(build: TypesGen.WorkspaceBuild) {
+  return new Promise<TypesGen.ProvisionerJob | undefined>((res, reject) => {
+    void (async () => {
+      let latestJobInfo: TypesGen.ProvisionerJob | undefined = undefined
+
+      while (
+        !["succeeded", "canceled"].some((status) =>
+          latestJobInfo?.status.includes(status),
+        )
+      ) {
+        const { job } = await getWorkspaceBuildByNumber(
+          build.workspace_owner_name,
+          build.workspace_name,
+          String(build.build_number),
+        )
+        latestJobInfo = job
+
+        if (latestJobInfo.status === "failed") {
+          return reject(latestJobInfo)
+        }
+
+        await delay(1000)
+      }
+
+      return res(latestJobInfo)
+    })()
+  })
+}
+
 export const postWorkspaceBuild = async (
   workspaceId: string,
   data: TypesGen.CreateWorkspaceBuildRequest,
@@ -489,12 +519,12 @@ export const postWorkspaceBuild = async (
 
 export const startWorkspace = (
   workspaceId: string,
-  templateVersionID: string,
+  templateVersionId: string,
   logLevel?: TypesGen.CreateWorkspaceBuildRequest["log_level"],
 ) =>
   postWorkspaceBuild(workspaceId, {
     transition: "start",
-    template_version_id: templateVersionID,
+    template_version_id: templateVersionId,
     log_level: logLevel,
   })
 export const stopWorkspace = (
@@ -505,6 +535,7 @@ export const stopWorkspace = (
     transition: "stop",
     log_level: logLevel,
   })
+
 export const deleteWorkspace = (
   workspaceId: string,
   logLevel?: TypesGen.CreateWorkspaceBuildRequest["log_level"],
@@ -523,6 +554,22 @@ export const cancelWorkspaceBuild = async (
   return response.data
 }
 
+export const restartWorkspace = async (workspace: TypesGen.Workspace) => {
+  const stopBuild = await stopWorkspace(workspace.id)
+  const awaitedStopBuild = await waitForBuild(stopBuild)
+
+  // If the restart is canceled halfway through, make sure we bail
+  if (awaitedStopBuild?.status === "canceled") {
+    return
+  }
+
+  const startBuild = await startWorkspace(
+    workspace.id,
+    workspace.latest_build.template_version_id,
+  )
+  await waitForBuild(startBuild)
+}
+
 export const cancelTemplateVersionBuild = async (
   templateVersionId: TypesGen.TemplateVersion["id"],
 ): Promise<Types.Message> => {

site/src/components/Workspace/Workspace.tsx

@@ -41,12 +41,14 @@ export interface WorkspaceProps {
   }
   handleStart: () => void
   handleStop: () => void
+  handleRestart: () => void
   handleDelete: () => void
   handleUpdate: () => void
   handleCancel: () => void
   handleSettings: () => void
   handleChangeVersion: () => void
   isUpdating: boolean
+  isRestarting: boolean
   workspace: TypesGen.Workspace
   resources?: TypesGen.WorkspaceResource[]
   builds?: TypesGen.WorkspaceBuild[]
@@ -72,6 +74,7 @@ export const Workspace: FC<React.PropsWithChildren<WorkspaceProps>> = ({
   scheduleProps,
   handleStart,
   handleStop,
+  handleRestart,
   handleDelete,
   handleUpdate,
   handleCancel,
@@ -79,6 +82,7 @@ export const Workspace: FC<React.PropsWithChildren<WorkspaceProps>> = ({
   handleChangeVersion,
   workspace,
   isUpdating,
+  isRestarting,
   resources,
   builds,
   canUpdateWorkspace,
@@ -132,6 +136,7 @@ export const Workspace: FC<React.PropsWithChildren<WorkspaceProps>> = ({
               isOutdated={workspace.outdated}
               handleStart={handleStart}
               handleStop={handleStop}
+              handleRestart={handleRestart}
               handleDelete={handleDelete}
               handleUpdate={handleUpdate}
               handleCancel={handleCancel}
@@ -139,6 +144,7 @@ export const Workspace: FC<React.PropsWithChildren<WorkspaceProps>> = ({
               handleChangeVersion={handleChangeVersion}
               canChangeVersions={canChangeVersions}
               isUpdating={isUpdating}
+              isRestarting={isRestarting}
             />
           </Stack>
         }

site/src/components/WorkspaceActions/Buttons.tsx

@@ -3,8 +3,9 @@ import BlockIcon from "@material-ui/icons/Block"
 import CloudQueueIcon from "@material-ui/icons/CloudQueue"
 import CropSquareIcon from "@material-ui/icons/CropSquare"
 import PlayCircleOutlineIcon from "@material-ui/icons/PlayCircleOutline"
+import ReplayIcon from "@material-ui/icons/Replay"
 import { LoadingButton } from "components/LoadingButton/LoadingButton"
-import { FC } from "react"
+import { FC, PropsWithChildren } from "react"
 import { useTranslation } from "react-i18next"
 import { makeStyles } from "@material-ui/core/styles"
 
@@ -12,7 +13,7 @@ interface WorkspaceAction {
   handleAction: () => void
 }
 
-export const UpdateButton: FC<React.PropsWithChildren<WorkspaceAction>> = ({
+export const UpdateButton: FC<PropsWithChildren<WorkspaceAction>> = ({
   handleAction,
 }) => {
   const { t } = useTranslation("workspacePage")
@@ -30,7 +31,7 @@ export const UpdateButton: FC<React.PropsWithChildren<WorkspaceAction>> = ({
   )
 }
 
-export const StartButton: FC<React.PropsWithChildren<WorkspaceAction>> = ({
+export const StartButton: FC<PropsWithChildren<WorkspaceAction>> = ({
   handleAction,
 }) => {
   const { t } = useTranslation("workspacePage")
@@ -48,7 +49,7 @@ export const StartButton: FC<React.PropsWithChildren<WorkspaceAction>> = ({
   )
 }
 
-export const StopButton: FC<React.PropsWithChildren<WorkspaceAction>> = ({
+export const StopButton: FC<PropsWithChildren<WorkspaceAction>> = ({
   handleAction,
 }) => {
   const { t } = useTranslation("workspacePage")
@@ -66,7 +67,25 @@ export const StopButton: FC<React.PropsWithChildren<WorkspaceAction>> = ({
   )
 }
 
-export const CancelButton: FC<React.PropsWithChildren<WorkspaceAction>> = ({
+export const RestartButton: FC<PropsWithChildren<WorkspaceAction>> = ({
+  handleAction,
+}) => {
+  const { t } = useTranslation("workspacePage")
+  const styles = useStyles()
+
+  return (
+    <Button
+      variant="outlined"
+      startIcon={<ReplayIcon />}
+      onClick={handleAction}
+      className={styles.fixedWidth}
+    >
+      {t("actionButton.restart")}
+    </Button>
+  )
+}
+
+export const CancelButton: FC<PropsWithChildren<WorkspaceAction>> = ({
   handleAction,
 }) => {
   return (
@@ -80,7 +99,7 @@ interface DisabledProps {
   label: string
 }
 
-export const DisabledButton: FC<React.PropsWithChildren<DisabledProps>> = ({
+export const DisabledButton: FC<PropsWithChildren<DisabledProps>> = ({
   label,
 }) => {
   return (
@@ -94,7 +113,7 @@ interface LoadingProps {
   label: string
 }
 
-export const ActionLoadingButton: FC<React.PropsWithChildren<LoadingProps>> = ({
+export const ActionLoadingButton: FC<PropsWithChildren<LoadingProps>> = ({
   label,
 }) => {
   const styles = useStyles()

site/src/components/WorkspaceActions/WorkspaceActions.stories.tsx

@@ -15,6 +15,7 @@ const Template: Story<WorkspaceActionsProps> = (args) => (
 const defaultArgs = {
   handleStart: action("start"),
   handleStop: action("stop"),
+  handleRestart: action("restart"),
   handleDelete: action("delete"),
   handleUpdate: action("update"),
   handleCancel: action("cancel"),

site/src/components/WorkspaceActions/WorkspaceActions.tsx

@@ -5,13 +5,14 @@ import { makeStyles } from "@material-ui/core/styles"
 import MoreVertOutlined from "@material-ui/icons/MoreVertOutlined"
 import { FC, ReactNode, useRef, useState } from "react"
 import { useTranslation } from "react-i18next"
-import { WorkspaceStatus } from "../../api/typesGenerated"
+import { WorkspaceStatus } from "api/typesGenerated"
 import {
   ActionLoadingButton,
   CancelButton,
   DisabledButton,
   StartButton,
   StopButton,
+  RestartButton,
   UpdateButton,
 } from "./Buttons"
 import {
@@ -28,12 +29,14 @@ export interface WorkspaceActionsProps {
   isOutdated: boolean
   handleStart: () => void
   handleStop: () => void
+  handleRestart: () => void
   handleDelete: () => void
   handleUpdate: () => void
   handleCancel: () => void
   handleSettings: () => void
   handleChangeVersion: () => void
   isUpdating: boolean
+  isRestarting: boolean
   children?: ReactNode
   canChangeVersions: boolean
 }
@@ -43,12 +46,14 @@ export const WorkspaceActions: FC<WorkspaceActionsProps> = ({
   isOutdated,
   handleStart,
   handleStop,
+  handleRestart,
   handleDelete,
   handleUpdate,
   handleCancel,
   handleSettings,
   handleChangeVersion,
   isUpdating,
+  isRestarting,
   canChangeVersions,
 }) => {
   const styles = useStyles()
@@ -91,6 +96,13 @@ export const WorkspaceActions: FC<WorkspaceActionsProps> = ({
         key={ButtonTypesEnum.stopping}
       />
     ),
+    [ButtonTypesEnum.restart]: <RestartButton handleAction={handleRestart} />,
+    [ButtonTypesEnum.restarting]: (
+      <ActionLoadingButton
+        label="Restarting"
+        key={ButtonTypesEnum.restarting}
+      />
+    ),
     [ButtonTypesEnum.deleting]: (
       <ActionLoadingButton
         label={t("actionButton.deleting")}
@@ -129,7 +141,11 @@ export const WorkspaceActions: FC<WorkspaceActionsProps> = ({
         (isUpdating
           ? buttonMapping[ButtonTypesEnum.updating]
           : buttonMapping[ButtonTypesEnum.update])}
-      {actionsByStatus.map((action) => buttonMapping[action])}
+      {isRestarting && buttonMapping[ButtonTypesEnum.restarting]}
+      {!isRestarting &&
+        actionsByStatus.map((action) => (
+          <span key={action}>{buttonMapping[action]}</span>
+        ))}
       {canCancel && <CancelButton handleAction={handleCancel} />}
       <div>
         <Button

site/src/components/WorkspaceActions/constants.ts

@@ -7,6 +7,8 @@ export enum ButtonTypesEnum {
   starting = "starting",
   stop = "stop",
   stopping = "stopping",
+  restart = "restart",
+  restarting = "restarting",
   deleting = "deleting",
   update = "update",
   updating = "updating",
@@ -39,7 +41,7 @@ const statusToActions: Record<WorkspaceStatus, WorkspaceAbilities> = {
     canAcceptJobs: false,
   },
   running: {
-    actions: [ButtonTypesEnum.stop],
+    actions: [ButtonTypesEnum.stop, ButtonTypesEnum.restart],
     canCancel: false,
     canAcceptJobs: true,
   },

site/src/i18n/en/workspacePage.json

@@ -21,6 +21,7 @@
   "actionButton": {
     "start": "Start",
     "stop": "Stop",
+    "restart": "Restart",
     "delete": "Delete",
     "cancel": "Cancel",
     "update": "Update",

site/src/pages/WorkspacePage/WorkspacePage.test.tsx

@@ -151,6 +151,17 @@ describe("WorkspacePage", () => {
     )
   })
 
+  it("requests a stop when the user presses Restart", async () => {
+    const stopWorkspaceMock = jest
+      .spyOn(api, "stopWorkspace")
+      .mockResolvedValueOnce(MockWorkspaceBuild)
+
+    await testButton("Restart", stopWorkspaceMock)
+
+    const button = await screen.findByText("Restarting")
+    expect(button).toBeInTheDocument()
+  })
+
   it("requests cancellation when the user presses Cancel", async () => {
     server.use(
       rest.get(

site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx

@@ -30,6 +30,7 @@ import { UpdateBuildParametersDialog } from "./UpdateBuildParametersDialog"
 import { ChangeVersionDialog } from "./ChangeVersionDialog"
 import { useQuery } from "@tanstack/react-query"
 import { getTemplateVersions } from "api/api"
+import { useRestartWorkspace } from "./hooks"
 
 interface WorkspaceReadyPageProps {
   workspaceState: StateFrom<typeof workspaceMachine>
@@ -77,6 +78,12 @@ export const WorkspaceReadyPage = ({
     enabled: changeVersionDialogOpen,
   })
 
+  const {
+    mutate: restartWorkspace,
+    error: restartBuildError,
+    isLoading: isRestarting,
+  } = useRestartWorkspace()
+
   // keep banner machine in sync with workspace
   useEffect(() => {
     bannerSend({ type: "REFRESH_WORKSPACE", workspace })
@@ -120,9 +127,11 @@ export const WorkspaceReadyPage = ({
           ),
         }}
         isUpdating={workspaceState.matches("ready.build.requestingUpdate")}
+        isRestarting={isRestarting}
         workspace={workspace}
         handleStart={() => workspaceSend({ type: "START" })}
         handleStop={() => workspaceSend({ type: "STOP" })}
+        handleRestart={() => restartWorkspace(workspace)}
         handleDelete={() => workspaceSend({ type: "ASK_DELETE" })}
         handleUpdate={() => workspaceSend({ type: "UPDATE" })}
         handleCancel={() => workspaceSend({ type: "CANCEL" })}
@@ -140,7 +149,7 @@ export const WorkspaceReadyPage = ({
         hideVSCodeDesktopButton={featureVisibility["browser_only"]}
         workspaceErrors={{
           [WorkspaceErrors.GET_BUILDS_ERROR]: getBuildsError,
-          [WorkspaceErrors.BUILD_ERROR]: buildError,
+          [WorkspaceErrors.BUILD_ERROR]: buildError || restartBuildError,
           [WorkspaceErrors.CANCELLATION_ERROR]: cancellationError,
         }}
         buildInfo={buildInfo}

site/src/pages/WorkspacePage/hooks.ts

@@ -0,0 +1,8 @@
+import { restartWorkspace } from "api/api"
+import { useMutation } from "@tanstack/react-query"
+
+export const useRestartWorkspace = () => {
+  return useMutation({
+    mutationFn: restartWorkspace,
+  })
+}