From 4079fc5728475dd77c6c36fe97f805834c198850 Mon Sep 17 00:00:00 2001
From: Kira Pilot <kira@coder.com>
Date: Fri, 14 Apr 2023 08:48:05 -0700
Subject: [PATCH 1/7] 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>
---
 site/src/components/Workspace/Workspace.tsx   |  3 ++
 .../components/WorkspaceActions/Buttons.tsx   | 33 +++++++++++++++----
 .../WorkspaceActions.stories.tsx              |  1 +
 .../WorkspaceActions/WorkspaceActions.tsx     |  4 +++
 .../components/WorkspaceActions/constants.ts  |  3 +-
 site/src/i18n/en/workspacePage.json           |  1 +
 .../WorkspacePage/WorkspaceReadyPage.tsx      |  1 +
 7 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/site/src/components/Workspace/Workspace.tsx b/site/src/components/Workspace/Workspace.tsx
index b0377eeb29847..f7090b45b4eec 100644
--- a/site/src/components/Workspace/Workspace.tsx
+++ b/site/src/components/Workspace/Workspace.tsx
@@ -41,6 +41,7 @@ export interface WorkspaceProps {
   }
   handleStart: () => void
   handleStop: () => void
+  handleRestart: () => void
   handleDelete: () => void
   handleUpdate: () => void
   handleCancel: () => void
@@ -72,6 +73,7 @@ export const Workspace: FC<React.PropsWithChildren<WorkspaceProps>> = ({
   scheduleProps,
   handleStart,
   handleStop,
+  handleRestart,
   handleDelete,
   handleUpdate,
   handleCancel,
@@ -132,6 +134,7 @@ export const Workspace: FC<React.PropsWithChildren<WorkspaceProps>> = ({
               isOutdated={workspace.outdated}
               handleStart={handleStart}
               handleStop={handleStop}
+              handleRestart={handleRestart}
               handleDelete={handleDelete}
               handleUpdate={handleUpdate}
               handleCancel={handleCancel}
diff --git a/site/src/components/WorkspaceActions/Buttons.tsx b/site/src/components/WorkspaceActions/Buttons.tsx
index b8c38469df68f..d6207952a4ac8 100644
--- a/site/src/components/WorkspaceActions/Buttons.tsx
+++ b/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()
diff --git a/site/src/components/WorkspaceActions/WorkspaceActions.stories.tsx b/site/src/components/WorkspaceActions/WorkspaceActions.stories.tsx
index 090e81cb2bb03..b2b2526811d0d 100644
--- a/site/src/components/WorkspaceActions/WorkspaceActions.stories.tsx
+++ b/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"),
diff --git a/site/src/components/WorkspaceActions/WorkspaceActions.tsx b/site/src/components/WorkspaceActions/WorkspaceActions.tsx
index d7508ed64405b..2cc4bbfcac203 100644
--- a/site/src/components/WorkspaceActions/WorkspaceActions.tsx
+++ b/site/src/components/WorkspaceActions/WorkspaceActions.tsx
@@ -12,6 +12,7 @@ import {
   DisabledButton,
   StartButton,
   StopButton,
+  RestartButton,
   UpdateButton,
 } from "./Buttons"
 import {
@@ -28,6 +29,7 @@ export interface WorkspaceActionsProps {
   isOutdated: boolean
   handleStart: () => void
   handleStop: () => void
+  handleRestart: () => void
   handleDelete: () => void
   handleUpdate: () => void
   handleCancel: () => void
@@ -43,6 +45,7 @@ export const WorkspaceActions: FC<WorkspaceActionsProps> = ({
   isOutdated,
   handleStart,
   handleStop,
+  handleRestart,
   handleDelete,
   handleUpdate,
   handleCancel,
@@ -91,6 +94,7 @@ export const WorkspaceActions: FC<WorkspaceActionsProps> = ({
         key={ButtonTypesEnum.stopping}
       />
     ),
+    [ButtonTypesEnum.restart]: <RestartButton handleAction={handleRestart} />,
     [ButtonTypesEnum.deleting]: (
       <ActionLoadingButton
         label={t("actionButton.deleting")}
diff --git a/site/src/components/WorkspaceActions/constants.ts b/site/src/components/WorkspaceActions/constants.ts
index 18923c3f38295..373625ffebad5 100644
--- a/site/src/components/WorkspaceActions/constants.ts
+++ b/site/src/components/WorkspaceActions/constants.ts
@@ -7,6 +7,7 @@ export enum ButtonTypesEnum {
   starting = "starting",
   stop = "stop",
   stopping = "stopping",
+  restart = "restart",
   deleting = "deleting",
   update = "update",
   updating = "updating",
@@ -39,7 +40,7 @@ const statusToActions: Record<WorkspaceStatus, WorkspaceAbilities> = {
     canAcceptJobs: false,
   },
   running: {
-    actions: [ButtonTypesEnum.stop],
+    actions: [ButtonTypesEnum.stop, ButtonTypesEnum.restart],
     canCancel: false,
     canAcceptJobs: true,
   },
diff --git a/site/src/i18n/en/workspacePage.json b/site/src/i18n/en/workspacePage.json
index fa5e9f8335021..fdcfc6ea31a65 100644
--- a/site/src/i18n/en/workspacePage.json
+++ b/site/src/i18n/en/workspacePage.json
@@ -21,6 +21,7 @@
   "actionButton": {
     "start": "Start",
     "stop": "Stop",
+    "restart": "Restart",
     "delete": "Delete",
     "cancel": "Cancel",
     "update": "Update",
diff --git a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
index 6066b75e64b8e..1a910a8a0c6ff 100644
--- a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
+++ b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
@@ -123,6 +123,7 @@ export const WorkspaceReadyPage = ({
         workspace={workspace}
         handleStart={() => workspaceSend({ type: "START" })}
         handleStop={() => workspaceSend({ type: "STOP" })}
+        handleRestart={() => workspaceSend({ type: "START" })}
         handleDelete={() => workspaceSend({ type: "ASK_DELETE" })}
         handleUpdate={() => workspaceSend({ type: "UPDATE" })}
         handleCancel={() => workspaceSend({ type: "CANCEL" })}

From a6cbf3486f571b7e64f8e3b4cc50b34903cbc048 Mon Sep 17 00:00:00 2001
From: Kira Pilot <kira@coder.com>
Date: Tue, 25 Apr 2023 18:27:25 +0000
Subject: [PATCH 2/7] added restart hook

---
 site/src/components/Workspace/Workspace.tsx   |  2 +-
 .../components/WorkspaceActions/Buttons.tsx   |  6 +-
 .../WorkspaceActions/WorkspaceActions.tsx     |  8 ++-
 .../WorkspacePage/WorkspaceReadyPage.tsx      | 11 +++-
 site/src/pages/WorkspacePage/hooks.ts         | 64 +++++++++++++++++++
 5 files changed, 82 insertions(+), 9 deletions(-)
 create mode 100644 site/src/pages/WorkspacePage/hooks.ts

diff --git a/site/src/components/Workspace/Workspace.tsx b/site/src/components/Workspace/Workspace.tsx
index f7090b45b4eec..b79f0dc455d68 100644
--- a/site/src/components/Workspace/Workspace.tsx
+++ b/site/src/components/Workspace/Workspace.tsx
@@ -41,7 +41,7 @@ export interface WorkspaceProps {
   }
   handleStart: () => void
   handleStop: () => void
-  handleRestart: () => void
+  handleRestart: any
   handleDelete: () => void
   handleUpdate: () => void
   handleCancel: () => void
diff --git a/site/src/components/WorkspaceActions/Buttons.tsx b/site/src/components/WorkspaceActions/Buttons.tsx
index d6207952a4ac8..f4bc83a81afb0 100644
--- a/site/src/components/WorkspaceActions/Buttons.tsx
+++ b/site/src/components/WorkspaceActions/Buttons.tsx
@@ -8,6 +8,8 @@ import { LoadingButton } from "components/LoadingButton/LoadingButton"
 import { FC, PropsWithChildren } from "react"
 import { useTranslation } from "react-i18next"
 import { makeStyles } from "@material-ui/core/styles"
+// import { UseMutateFunction } from "@tanstack/react-query"
+// import { WorkspaceBuild } from "api/typesGenerated"
 
 interface WorkspaceAction {
   handleAction: () => void
@@ -67,9 +69,7 @@ export const StopButton: FC<PropsWithChildren<WorkspaceAction>> = ({
   )
 }
 
-export const RestartButton: FC<PropsWithChildren<WorkspaceAction>> = ({
-  handleAction,
-}) => {
+export const RestartButton: FC<PropsWithChildren<any>> = ({ handleAction }) => {
   const { t } = useTranslation("workspacePage")
   const styles = useStyles()
 
diff --git a/site/src/components/WorkspaceActions/WorkspaceActions.tsx b/site/src/components/WorkspaceActions/WorkspaceActions.tsx
index 2cc4bbfcac203..8c0e552b3856e 100644
--- a/site/src/components/WorkspaceActions/WorkspaceActions.tsx
+++ b/site/src/components/WorkspaceActions/WorkspaceActions.tsx
@@ -5,7 +5,7 @@ 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,
@@ -29,7 +29,7 @@ export interface WorkspaceActionsProps {
   isOutdated: boolean
   handleStart: () => void
   handleStop: () => void
-  handleRestart: () => void
+  handleRestart: any
   handleDelete: () => void
   handleUpdate: () => void
   handleCancel: () => void
@@ -133,7 +133,9 @@ export const WorkspaceActions: FC<WorkspaceActionsProps> = ({
         (isUpdating
           ? buttonMapping[ButtonTypesEnum.updating]
           : buttonMapping[ButtonTypesEnum.update])}
-      {actionsByStatus.map((action) => buttonMapping[action])}
+      {actionsByStatus.map((action) => (
+        <span key={action}>{buttonMapping[action]}</span>
+      ))}
       {canCancel && <CancelButton handleAction={handleCancel} />}
       <div>
         <Button
diff --git a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
index 1a910a8a0c6ff..50e144b279233 100644
--- a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
+++ b/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 [restartBuildError, setRestartBuildError] = useState<
+    Error | unknown | undefined
+  >(undefined)
+
+  const { mutate: restartWorkspace } = useRestartWorkspace(setRestartBuildError)
+
   // keep banner machine in sync with workspace
   useEffect(() => {
     bannerSend({ type: "REFRESH_WORKSPACE", workspace })
@@ -123,7 +130,7 @@ export const WorkspaceReadyPage = ({
         workspace={workspace}
         handleStart={() => workspaceSend({ type: "START" })}
         handleStop={() => workspaceSend({ type: "STOP" })}
-        handleRestart={() => workspaceSend({ type: "START" })}
+        handleRestart={() => restartWorkspace(workspace.id)}
         handleDelete={() => workspaceSend({ type: "ASK_DELETE" })}
         handleUpdate={() => workspaceSend({ type: "UPDATE" })}
         handleCancel={() => workspaceSend({ type: "CANCEL" })}
@@ -141,7 +148,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}
diff --git a/site/src/pages/WorkspacePage/hooks.ts b/site/src/pages/WorkspacePage/hooks.ts
new file mode 100644
index 0000000000000..229c7e1620b57
--- /dev/null
+++ b/site/src/pages/WorkspacePage/hooks.ts
@@ -0,0 +1,64 @@
+import { useMutation } from "@tanstack/react-query"
+import {
+  stopWorkspace,
+  getWorkspaceBuildByNumber,
+  startWorkspace,
+} from "api/api"
+import { delay } from "utils/delay"
+import { ProvisionerJob, WorkspaceBuild } from "api/typesGenerated"
+
+function waitForStop(
+  username: string,
+  workspaceName: string,
+  buildNumber: string,
+) {
+  return new Promise((res, reject) => {
+    void (async () => {
+      let latestJobInfo: ProvisionerJob | undefined = undefined
+
+      while (latestJobInfo?.status !== "succeeded") {
+        const { job } = await getWorkspaceBuildByNumber(
+          username,
+          workspaceName,
+          buildNumber,
+        )
+        latestJobInfo = job
+
+        if (
+          ["failed", "canceled"].some((status) =>
+            latestJobInfo?.status.includes(status),
+          )
+        ) {
+          return reject(latestJobInfo)
+        }
+
+        await delay(1000)
+      }
+
+      return res(latestJobInfo)
+    })()
+  })
+}
+
+export const useRestartWorkspace = (
+  setRestartBuildError: (arg: Error | unknown | undefined) => void,
+) => {
+  return useMutation({
+    mutationFn: stopWorkspace,
+    onSuccess: async (data) => {
+      try {
+        await waitForStop(
+          data.workspace_owner_name,
+          data.workspace_name,
+          String(data.build_number),
+        )
+        await startWorkspace(data.workspace_id, data.template_version_id)
+      } catch (error) {
+        if ((error as WorkspaceBuild).status === "canceled") {
+          return
+        }
+        setRestartBuildError(error)
+      }
+    },
+  })
+}

From 1f1d388f0a0f2e3fdbc4efad414ccca65dddce05 Mon Sep 17 00:00:00 2001
From: Kira Pilot <kira@coder.com>
Date: Wed, 26 Apr 2023 13:16:46 +0000
Subject: [PATCH 3/7] added error handling

---
 site/src/api/api.ts                           | 26 ++++++--
 .../WorkspacePage/WorkspaceReadyPage.tsx      | 38 +++++++++--
 site/src/pages/WorkspacePage/hooks.ts         | 64 -------------------
 site/src/pages/WorkspacePage/utils.ts         | 32 ++++++++++
 .../xServices/workspace/workspaceXService.ts  | 11 ++--
 5 files changed, 88 insertions(+), 83 deletions(-)
 delete mode 100644 site/src/pages/WorkspacePage/hooks.ts
 create mode 100644 site/src/pages/WorkspacePage/utils.ts

diff --git a/site/src/api/api.ts b/site/src/api/api.ts
index e778b9fc0a4fb..253b0abcefa29 100644
--- a/site/src/api/api.ts
+++ b/site/src/api/api.ts
@@ -487,14 +487,22 @@ export const postWorkspaceBuild = async (
   return response.data
 }
 
-export const startWorkspace = (
-  workspaceId: string,
-  templateVersionID: string,
-  logLevel?: TypesGen.CreateWorkspaceBuildRequest["log_level"],
-) =>
+// it is necessary to create an interface here for react query
+// as startWorkspace requires multiple parameters
+interface StartWorkspaceParams {
+  workspaceId: string
+  templateVersionId: string
+  logLevel?: TypesGen.CreateWorkspaceBuildRequest["log_level"]
+}
+
+export const startWorkspace = ({
+  workspaceId,
+  templateVersionId,
+  logLevel,
+}: StartWorkspaceParams) =>
   postWorkspaceBuild(workspaceId, {
     transition: "start",
-    template_version_id: templateVersionID,
+    template_version_id: templateVersionId,
     log_level: logLevel,
   })
 export const stopWorkspace = (
@@ -505,6 +513,7 @@ export const stopWorkspace = (
     transition: "stop",
     log_level: logLevel,
   })
+
 export const deleteWorkspace = (
   workspaceId: string,
   logLevel?: TypesGen.CreateWorkspaceBuildRequest["log_level"],
@@ -954,7 +963,10 @@ export const updateWorkspaceVersion = async (
   workspace: TypesGen.Workspace,
 ): Promise<TypesGen.WorkspaceBuild> => {
   const template = await getTemplate(workspace.template_id)
-  return startWorkspace(workspace.id, template.active_version_id)
+  return startWorkspace({
+    workspaceId: workspace.id,
+    templateVersionId: template.active_version_id,
+  })
 }
 
 export const getWorkspaceBuildParameters = async (
diff --git a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
index 50e144b279233..91ddc5b288e4d 100644
--- a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
+++ b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
@@ -1,5 +1,5 @@
 import { useActor } from "@xstate/react"
-import { ProvisionerJobLog } from "api/typesGenerated"
+import { ProvisionerJobLog, WorkspaceBuild } from "api/typesGenerated"
 import { useDashboard } from "components/Dashboard/DashboardProvider"
 import dayjs from "dayjs"
 import { useFeatureVisibility } from "hooks/useFeatureVisibility"
@@ -28,9 +28,9 @@ import {
 } from "../../xServices/workspace/workspaceXService"
 import { UpdateBuildParametersDialog } from "./UpdateBuildParametersDialog"
 import { ChangeVersionDialog } from "./ChangeVersionDialog"
-import { useQuery } from "@tanstack/react-query"
-import { getTemplateVersions } from "api/api"
-import { useRestartWorkspace } from "./hooks"
+import { useQuery, useMutation } from "@tanstack/react-query"
+import { getTemplateVersions, stopWorkspace, startWorkspace } from "api/api"
+import { waitForStop } from "./utils"
 
 interface WorkspaceReadyPageProps {
   workspaceState: StateFrom<typeof workspaceMachine>
@@ -82,7 +82,31 @@ export const WorkspaceReadyPage = ({
     Error | unknown | undefined
   >(undefined)
 
-  const { mutate: restartWorkspace } = useRestartWorkspace(setRestartBuildError)
+  const requestStopWorkspace = useMutation({ mutationFn: stopWorkspace })
+  const requestStartWorkspace = useMutation({ mutationFn: startWorkspace })
+
+  const restartWorkspace = async () => {
+    try {
+      await requestStopWorkspace.mutateAsync(workspace.id)
+
+      if (!requestStopWorkspace.data) {
+        throw new Error("No stop build returned!")
+      }
+
+      await waitForStop(requestStopWorkspace.data)
+      await requestStartWorkspace.mutateAsync({
+        workspaceId: workspace.id,
+        templateVersionId: workspace.latest_build.template_version_id,
+      })
+    } catch (error) {
+      // No error handling needed if the build has been canceled
+      if ((error as WorkspaceBuild).status === "canceled") {
+        return
+      }
+
+      setRestartBuildError(error)
+    }
+  }
 
   // keep banner machine in sync with workspace
   useEffect(() => {
@@ -130,7 +154,7 @@ export const WorkspaceReadyPage = ({
         workspace={workspace}
         handleStart={() => workspaceSend({ type: "START" })}
         handleStop={() => workspaceSend({ type: "STOP" })}
-        handleRestart={() => restartWorkspace(workspace.id)}
+        handleRestart={restartWorkspace}
         handleDelete={() => workspaceSend({ type: "ASK_DELETE" })}
         handleUpdate={() => workspaceSend({ type: "UPDATE" })}
         handleCancel={() => workspaceSend({ type: "CANCEL" })}
@@ -148,7 +172,7 @@ export const WorkspaceReadyPage = ({
         hideVSCodeDesktopButton={featureVisibility["browser_only"]}
         workspaceErrors={{
           [WorkspaceErrors.GET_BUILDS_ERROR]: getBuildsError,
-          [WorkspaceErrors.BUILD_ERROR]: buildError ?? restartBuildError,
+          [WorkspaceErrors.BUILD_ERROR]: buildError || restartBuildError,
           [WorkspaceErrors.CANCELLATION_ERROR]: cancellationError,
         }}
         buildInfo={buildInfo}
diff --git a/site/src/pages/WorkspacePage/hooks.ts b/site/src/pages/WorkspacePage/hooks.ts
deleted file mode 100644
index 229c7e1620b57..0000000000000
--- a/site/src/pages/WorkspacePage/hooks.ts
+++ /dev/null
@@ -1,64 +0,0 @@
-import { useMutation } from "@tanstack/react-query"
-import {
-  stopWorkspace,
-  getWorkspaceBuildByNumber,
-  startWorkspace,
-} from "api/api"
-import { delay } from "utils/delay"
-import { ProvisionerJob, WorkspaceBuild } from "api/typesGenerated"
-
-function waitForStop(
-  username: string,
-  workspaceName: string,
-  buildNumber: string,
-) {
-  return new Promise((res, reject) => {
-    void (async () => {
-      let latestJobInfo: ProvisionerJob | undefined = undefined
-
-      while (latestJobInfo?.status !== "succeeded") {
-        const { job } = await getWorkspaceBuildByNumber(
-          username,
-          workspaceName,
-          buildNumber,
-        )
-        latestJobInfo = job
-
-        if (
-          ["failed", "canceled"].some((status) =>
-            latestJobInfo?.status.includes(status),
-          )
-        ) {
-          return reject(latestJobInfo)
-        }
-
-        await delay(1000)
-      }
-
-      return res(latestJobInfo)
-    })()
-  })
-}
-
-export const useRestartWorkspace = (
-  setRestartBuildError: (arg: Error | unknown | undefined) => void,
-) => {
-  return useMutation({
-    mutationFn: stopWorkspace,
-    onSuccess: async (data) => {
-      try {
-        await waitForStop(
-          data.workspace_owner_name,
-          data.workspace_name,
-          String(data.build_number),
-        )
-        await startWorkspace(data.workspace_id, data.template_version_id)
-      } catch (error) {
-        if ((error as WorkspaceBuild).status === "canceled") {
-          return
-        }
-        setRestartBuildError(error)
-      }
-    },
-  })
-}
diff --git a/site/src/pages/WorkspacePage/utils.ts b/site/src/pages/WorkspacePage/utils.ts
new file mode 100644
index 0000000000000..5ed605c264168
--- /dev/null
+++ b/site/src/pages/WorkspacePage/utils.ts
@@ -0,0 +1,32 @@
+import { getWorkspaceBuildByNumber } from "api/api"
+import { delay } from "utils/delay"
+import { ProvisionerJob, WorkspaceBuild } from "api/typesGenerated"
+
+export function waitForStop(build: WorkspaceBuild) {
+  return new Promise((res, reject) => {
+    void (async () => {
+      let latestJobInfo: ProvisionerJob | undefined = undefined
+
+      while (latestJobInfo?.status !== "succeeded") {
+        const { job } = await getWorkspaceBuildByNumber(
+          build.workspace_owner_name,
+          build.workspace_name,
+          String(build.build_number),
+        )
+        latestJobInfo = job
+
+        if (
+          ["failed", "canceled"].some((status) =>
+            latestJobInfo?.status.includes(status),
+          )
+        ) {
+          return reject(latestJobInfo)
+        }
+
+        await delay(1000)
+      }
+
+      return res(latestJobInfo)
+    })()
+  })
+}
diff --git a/site/src/xServices/workspace/workspaceXService.ts b/site/src/xServices/workspace/workspaceXService.ts
index 818b853960761..623ba82003342 100644
--- a/site/src/xServices/workspace/workspaceXService.ts
+++ b/site/src/xServices/workspace/workspaceXService.ts
@@ -779,11 +779,12 @@ export const workspaceMachine = createMachine(
         },
       startWorkspace: (context) => async (send) => {
         if (context.workspace) {
-          const startWorkspacePromise = await API.startWorkspace(
-            context.workspace.id,
-            context.workspace.latest_build.template_version_id,
-            context.createBuildLogLevel,
-          )
+          const startWorkspacePromise = await API.startWorkspace({
+            workspaceId: context.workspace.id,
+            templateVersionId:
+              context.workspace.latest_build.template_version_id,
+            logLevel: context.createBuildLogLevel,
+          })
           send({ type: "REFRESH_TIMELINE" })
           return startWorkspacePromise
         } else {

From e8e88b45455e7d4debad8f1dbca8d5e64e7c150b Mon Sep 17 00:00:00 2001
From: Kira Pilot <kira@coder.com>
Date: Wed, 26 Apr 2023 14:34:48 +0000
Subject: [PATCH 4/7] going back to chaining in success callback

---
 .../WorkspacePage/WorkspaceReadyPage.tsx      | 36 ++---------
 site/src/pages/WorkspacePage/hooks.ts         | 62 +++++++++++++++++++
 site/src/pages/WorkspacePage/utils.ts         | 32 ----------
 3 files changed, 68 insertions(+), 62 deletions(-)
 create mode 100644 site/src/pages/WorkspacePage/hooks.ts
 delete mode 100644 site/src/pages/WorkspacePage/utils.ts

diff --git a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
index 91ddc5b288e4d..de561d1b3e746 100644
--- a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
+++ b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
@@ -1,5 +1,5 @@
 import { useActor } from "@xstate/react"
-import { ProvisionerJobLog, WorkspaceBuild } from "api/typesGenerated"
+import { ProvisionerJobLog } from "api/typesGenerated"
 import { useDashboard } from "components/Dashboard/DashboardProvider"
 import dayjs from "dayjs"
 import { useFeatureVisibility } from "hooks/useFeatureVisibility"
@@ -28,9 +28,9 @@ import {
 } from "../../xServices/workspace/workspaceXService"
 import { UpdateBuildParametersDialog } from "./UpdateBuildParametersDialog"
 import { ChangeVersionDialog } from "./ChangeVersionDialog"
-import { useQuery, useMutation } from "@tanstack/react-query"
-import { getTemplateVersions, stopWorkspace, startWorkspace } from "api/api"
-import { waitForStop } from "./utils"
+import { useQuery } from "@tanstack/react-query"
+import { getTemplateVersions } from "api/api"
+import { useRestartWorkspace } from "./hooks"
 
 interface WorkspaceReadyPageProps {
   workspaceState: StateFrom<typeof workspaceMachine>
@@ -82,31 +82,7 @@ export const WorkspaceReadyPage = ({
     Error | unknown | undefined
   >(undefined)
 
-  const requestStopWorkspace = useMutation({ mutationFn: stopWorkspace })
-  const requestStartWorkspace = useMutation({ mutationFn: startWorkspace })
-
-  const restartWorkspace = async () => {
-    try {
-      await requestStopWorkspace.mutateAsync(workspace.id)
-
-      if (!requestStopWorkspace.data) {
-        throw new Error("No stop build returned!")
-      }
-
-      await waitForStop(requestStopWorkspace.data)
-      await requestStartWorkspace.mutateAsync({
-        workspaceId: workspace.id,
-        templateVersionId: workspace.latest_build.template_version_id,
-      })
-    } catch (error) {
-      // No error handling needed if the build has been canceled
-      if ((error as WorkspaceBuild).status === "canceled") {
-        return
-      }
-
-      setRestartBuildError(error)
-    }
-  }
+  const { mutate: restartWorkspace } = useRestartWorkspace(setRestartBuildError)
 
   // keep banner machine in sync with workspace
   useEffect(() => {
@@ -154,7 +130,7 @@ export const WorkspaceReadyPage = ({
         workspace={workspace}
         handleStart={() => workspaceSend({ type: "START" })}
         handleStop={() => workspaceSend({ type: "STOP" })}
-        handleRestart={restartWorkspace}
+        handleRestart={() => restartWorkspace(workspace.id)}
         handleDelete={() => workspaceSend({ type: "ASK_DELETE" })}
         handleUpdate={() => workspaceSend({ type: "UPDATE" })}
         handleCancel={() => workspaceSend({ type: "CANCEL" })}
diff --git a/site/src/pages/WorkspacePage/hooks.ts b/site/src/pages/WorkspacePage/hooks.ts
new file mode 100644
index 0000000000000..4c5bd1a49cc7c
--- /dev/null
+++ b/site/src/pages/WorkspacePage/hooks.ts
@@ -0,0 +1,62 @@
+import {
+  getWorkspaceBuildByNumber,
+  stopWorkspace,
+  startWorkspace,
+} from "api/api"
+import { delay } from "utils/delay"
+import { ProvisionerJob, WorkspaceBuild } from "api/typesGenerated"
+import { useMutation } from "@tanstack/react-query"
+
+export function waitForStop(build: WorkspaceBuild) {
+  return new Promise((res, reject) => {
+    void (async () => {
+      let latestJobInfo: ProvisionerJob | undefined = undefined
+
+      while (latestJobInfo?.status !== "succeeded") {
+        const { job } = await getWorkspaceBuildByNumber(
+          build.workspace_owner_name,
+          build.workspace_name,
+          String(build.build_number),
+        )
+        latestJobInfo = job
+        console.log("latest job status", latestJobInfo.status)
+
+        if (
+          ["failed", "canceled"].some((status) =>
+            latestJobInfo?.status.includes(status),
+          )
+        ) {
+          return reject(latestJobInfo)
+        }
+
+        await delay(1000)
+      }
+
+      console.log("resolving status status", latestJobInfo.status)
+
+      return res(latestJobInfo)
+    })()
+  })
+}
+
+export const useRestartWorkspace = (
+  setRestartBuildError: (arg: Error | unknown | undefined) => void,
+) => {
+  return useMutation({
+    mutationFn: stopWorkspace,
+    onSuccess: async (data) => {
+      try {
+        await waitForStop(data)
+        await startWorkspace({
+          workspaceId: data.workspace_id,
+          templateVersionId: data.template_version_id,
+        })
+      } catch (error) {
+        if ((error as WorkspaceBuild).status === "canceled") {
+          return
+        }
+        setRestartBuildError(error)
+      }
+    },
+  })
+}
diff --git a/site/src/pages/WorkspacePage/utils.ts b/site/src/pages/WorkspacePage/utils.ts
deleted file mode 100644
index 5ed605c264168..0000000000000
--- a/site/src/pages/WorkspacePage/utils.ts
+++ /dev/null
@@ -1,32 +0,0 @@
-import { getWorkspaceBuildByNumber } from "api/api"
-import { delay } from "utils/delay"
-import { ProvisionerJob, WorkspaceBuild } from "api/typesGenerated"
-
-export function waitForStop(build: WorkspaceBuild) {
-  return new Promise((res, reject) => {
-    void (async () => {
-      let latestJobInfo: ProvisionerJob | undefined = undefined
-
-      while (latestJobInfo?.status !== "succeeded") {
-        const { job } = await getWorkspaceBuildByNumber(
-          build.workspace_owner_name,
-          build.workspace_name,
-          String(build.build_number),
-        )
-        latestJobInfo = job
-
-        if (
-          ["failed", "canceled"].some((status) =>
-            latestJobInfo?.status.includes(status),
-          )
-        ) {
-          return reject(latestJobInfo)
-        }
-
-        await delay(1000)
-      }
-
-      return res(latestJobInfo)
-    })()
-  })
-}

From 2e2b89da760287733648dcd4a10f4c07ff23504c Mon Sep 17 00:00:00 2001
From: Kira Pilot <kira@coder.com>
Date: Wed, 26 Apr 2023 16:18:47 +0000
Subject: [PATCH 5/7] add restarting btn

---
 site/src/components/Workspace/Workspace.tsx    |  5 ++++-
 .../components/WorkspaceActions/Buttons.tsx    |  6 +++---
 .../WorkspaceActions/WorkspaceActions.tsx      | 18 ++++++++++++++----
 .../components/WorkspaceActions/constants.ts   |  1 +
 .../pages/WorkspacePage/WorkspaceReadyPage.tsx |  8 +++++++-
 site/src/pages/WorkspacePage/hooks.ts          | 18 ++++++++++++------
 6 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/site/src/components/Workspace/Workspace.tsx b/site/src/components/Workspace/Workspace.tsx
index b79f0dc455d68..5be87ec6a95e1 100644
--- a/site/src/components/Workspace/Workspace.tsx
+++ b/site/src/components/Workspace/Workspace.tsx
@@ -41,13 +41,14 @@ export interface WorkspaceProps {
   }
   handleStart: () => void
   handleStop: () => void
-  handleRestart: any
+  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[]
@@ -81,6 +82,7 @@ export const Workspace: FC<React.PropsWithChildren<WorkspaceProps>> = ({
   handleChangeVersion,
   workspace,
   isUpdating,
+  isRestarting,
   resources,
   builds,
   canUpdateWorkspace,
@@ -142,6 +144,7 @@ export const Workspace: FC<React.PropsWithChildren<WorkspaceProps>> = ({
               handleChangeVersion={handleChangeVersion}
               canChangeVersions={canChangeVersions}
               isUpdating={isUpdating}
+              isRestarting={isRestarting}
             />
           </Stack>
         }
diff --git a/site/src/components/WorkspaceActions/Buttons.tsx b/site/src/components/WorkspaceActions/Buttons.tsx
index f4bc83a81afb0..d6207952a4ac8 100644
--- a/site/src/components/WorkspaceActions/Buttons.tsx
+++ b/site/src/components/WorkspaceActions/Buttons.tsx
@@ -8,8 +8,6 @@ import { LoadingButton } from "components/LoadingButton/LoadingButton"
 import { FC, PropsWithChildren } from "react"
 import { useTranslation } from "react-i18next"
 import { makeStyles } from "@material-ui/core/styles"
-// import { UseMutateFunction } from "@tanstack/react-query"
-// import { WorkspaceBuild } from "api/typesGenerated"
 
 interface WorkspaceAction {
   handleAction: () => void
@@ -69,7 +67,9 @@ export const StopButton: FC<PropsWithChildren<WorkspaceAction>> = ({
   )
 }
 
-export const RestartButton: FC<PropsWithChildren<any>> = ({ handleAction }) => {
+export const RestartButton: FC<PropsWithChildren<WorkspaceAction>> = ({
+  handleAction,
+}) => {
   const { t } = useTranslation("workspacePage")
   const styles = useStyles()
 
diff --git a/site/src/components/WorkspaceActions/WorkspaceActions.tsx b/site/src/components/WorkspaceActions/WorkspaceActions.tsx
index 8c0e552b3856e..aacf1f10f085e 100644
--- a/site/src/components/WorkspaceActions/WorkspaceActions.tsx
+++ b/site/src/components/WorkspaceActions/WorkspaceActions.tsx
@@ -29,13 +29,14 @@ export interface WorkspaceActionsProps {
   isOutdated: boolean
   handleStart: () => void
   handleStop: () => void
-  handleRestart: any
+  handleRestart: () => void
   handleDelete: () => void
   handleUpdate: () => void
   handleCancel: () => void
   handleSettings: () => void
   handleChangeVersion: () => void
   isUpdating: boolean
+  isRestarting: boolean
   children?: ReactNode
   canChangeVersions: boolean
 }
@@ -52,6 +53,7 @@ export const WorkspaceActions: FC<WorkspaceActionsProps> = ({
   handleSettings,
   handleChangeVersion,
   isUpdating,
+  isRestarting,
   canChangeVersions,
 }) => {
   const styles = useStyles()
@@ -95,6 +97,12 @@ export const WorkspaceActions: FC<WorkspaceActionsProps> = ({
       />
     ),
     [ButtonTypesEnum.restart]: <RestartButton handleAction={handleRestart} />,
+    [ButtonTypesEnum.restarting]: (
+      <ActionLoadingButton
+        label="Restarting"
+        key={ButtonTypesEnum.restarting}
+      />
+    ),
     [ButtonTypesEnum.deleting]: (
       <ActionLoadingButton
         label={t("actionButton.deleting")}
@@ -133,9 +141,11 @@ export const WorkspaceActions: FC<WorkspaceActionsProps> = ({
         (isUpdating
           ? buttonMapping[ButtonTypesEnum.updating]
           : buttonMapping[ButtonTypesEnum.update])}
-      {actionsByStatus.map((action) => (
-        <span key={action}>{buttonMapping[action]}</span>
-      ))}
+      {isRestarting && buttonMapping[ButtonTypesEnum.restarting]}
+      {!isRestarting &&
+        actionsByStatus.map((action) => (
+          <span key={action}>{buttonMapping[action]}</span>
+        ))}
       {canCancel && <CancelButton handleAction={handleCancel} />}
       <div>
         <Button
diff --git a/site/src/components/WorkspaceActions/constants.ts b/site/src/components/WorkspaceActions/constants.ts
index 373625ffebad5..c4e81ca976ca5 100644
--- a/site/src/components/WorkspaceActions/constants.ts
+++ b/site/src/components/WorkspaceActions/constants.ts
@@ -8,6 +8,7 @@ export enum ButtonTypesEnum {
   stop = "stop",
   stopping = "stopping",
   restart = "restart",
+  restarting = "restarting",
   deleting = "deleting",
   update = "update",
   updating = "updating",
diff --git a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
index de561d1b3e746..8f08b9d8a5fff 100644
--- a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
+++ b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
@@ -82,7 +82,12 @@ export const WorkspaceReadyPage = ({
     Error | unknown | undefined
   >(undefined)
 
-  const { mutate: restartWorkspace } = useRestartWorkspace(setRestartBuildError)
+  const [isRestarting, setIsRestarting] = useState<boolean>(false)
+
+  const { mutate: restartWorkspace } = useRestartWorkspace(
+    setRestartBuildError,
+    setIsRestarting,
+  )
 
   // keep banner machine in sync with workspace
   useEffect(() => {
@@ -127,6 +132,7 @@ export const WorkspaceReadyPage = ({
           ),
         }}
         isUpdating={workspaceState.matches("ready.build.requestingUpdate")}
+        isRestarting={isRestarting}
         workspace={workspace}
         handleStart={() => workspaceSend({ type: "START" })}
         handleStop={() => workspaceSend({ type: "STOP" })}
diff --git a/site/src/pages/WorkspacePage/hooks.ts b/site/src/pages/WorkspacePage/hooks.ts
index 4c5bd1a49cc7c..47788ab9e779d 100644
--- a/site/src/pages/WorkspacePage/hooks.ts
+++ b/site/src/pages/WorkspacePage/hooks.ts
@@ -7,7 +7,7 @@ import { delay } from "utils/delay"
 import { ProvisionerJob, WorkspaceBuild } from "api/typesGenerated"
 import { useMutation } from "@tanstack/react-query"
 
-export function waitForStop(build: WorkspaceBuild) {
+export function waitForBuild(build: WorkspaceBuild) {
   return new Promise((res, reject) => {
     void (async () => {
       let latestJobInfo: ProvisionerJob | undefined = undefined
@@ -40,22 +40,28 @@ export function waitForStop(build: WorkspaceBuild) {
 }
 
 export const useRestartWorkspace = (
-  setRestartBuildError: (arg: Error | unknown | undefined) => void,
+  setBuildError: (arg: Error | unknown | undefined) => void,
+  setLoading: (arg: boolean) => void,
 ) => {
   return useMutation({
     mutationFn: stopWorkspace,
-    onSuccess: async (data) => {
+    onMutate: () => setLoading(true),
+    onSuccess: async (data: WorkspaceBuild) => {
       try {
-        await waitForStop(data)
-        await startWorkspace({
+        await waitForBuild(data)
+        const startBuild = await startWorkspace({
           workspaceId: data.workspace_id,
           templateVersionId: data.template_version_id,
         })
+        await waitForBuild(startBuild)
+
+        setLoading(false)
       } catch (error) {
         if ((error as WorkspaceBuild).status === "canceled") {
           return
         }
-        setRestartBuildError(error)
+        setBuildError(error)
+        setLoading(false)
       }
     },
   })

From 585ece4d4e8c9f44419bd5910234e7f70bac546b Mon Sep 17 00:00:00 2001
From: Kira Pilot <kira@coder.com>
Date: Wed, 26 Apr 2023 19:30:11 +0000
Subject: [PATCH 6/7] added test

---
 site/src/pages/WorkspacePage/WorkspacePage.test.tsx | 11 +++++++++++
 site/src/pages/WorkspacePage/hooks.ts               |  4 +---
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/site/src/pages/WorkspacePage/WorkspacePage.test.tsx b/site/src/pages/WorkspacePage/WorkspacePage.test.tsx
index ac0383ce34b4d..8739d74dfb0fb 100644
--- a/site/src/pages/WorkspacePage/WorkspacePage.test.tsx
+++ b/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(
diff --git a/site/src/pages/WorkspacePage/hooks.ts b/site/src/pages/WorkspacePage/hooks.ts
index 47788ab9e779d..15f223a6a32a5 100644
--- a/site/src/pages/WorkspacePage/hooks.ts
+++ b/site/src/pages/WorkspacePage/hooks.ts
@@ -19,7 +19,6 @@ export function waitForBuild(build: WorkspaceBuild) {
           String(build.build_number),
         )
         latestJobInfo = job
-        console.log("latest job status", latestJobInfo.status)
 
         if (
           ["failed", "canceled"].some((status) =>
@@ -32,8 +31,6 @@ export function waitForBuild(build: WorkspaceBuild) {
         await delay(1000)
       }
 
-      console.log("resolving status status", latestJobInfo.status)
-
       return res(latestJobInfo)
     })()
   })
@@ -58,6 +55,7 @@ export const useRestartWorkspace = (
         setLoading(false)
       } catch (error) {
         if ((error as WorkspaceBuild).status === "canceled") {
+          setLoading(false)
           return
         }
         setBuildError(error)

From 5041d5e22765894564fda82c37c96ffad91828d3 Mon Sep 17 00:00:00 2001
From: Kira Pilot <kira@coder.com>
Date: Thu, 27 Apr 2023 20:00:55 +0000
Subject: [PATCH 7/7] PR feedback

---
 site/src/api/api.ts                           | 69 ++++++++++++++-----
 .../WorkspacePage/WorkspaceReadyPage.tsx      | 17 ++---
 site/src/pages/WorkspacePage/hooks.ts         | 64 +----------------
 .../xServices/workspace/workspaceXService.ts  | 11 ++-
 4 files changed, 66 insertions(+), 95 deletions(-)

diff --git a/site/src/api/api.ts b/site/src/api/api.ts
index 253b0abcefa29..1abd7643fb29c 100644
--- a/site/src/api/api.ts
+++ b/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,
@@ -487,19 +517,11 @@ export const postWorkspaceBuild = async (
   return response.data
 }
 
-// it is necessary to create an interface here for react query
-// as startWorkspace requires multiple parameters
-interface StartWorkspaceParams {
-  workspaceId: string
-  templateVersionId: string
-  logLevel?: TypesGen.CreateWorkspaceBuildRequest["log_level"]
-}
-
-export const startWorkspace = ({
-  workspaceId,
-  templateVersionId,
-  logLevel,
-}: StartWorkspaceParams) =>
+export const startWorkspace = (
+  workspaceId: string,
+  templateVersionId: string,
+  logLevel?: TypesGen.CreateWorkspaceBuildRequest["log_level"],
+) =>
   postWorkspaceBuild(workspaceId, {
     transition: "start",
     template_version_id: templateVersionId,
@@ -532,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> => {
@@ -963,10 +1001,7 @@ export const updateWorkspaceVersion = async (
   workspace: TypesGen.Workspace,
 ): Promise<TypesGen.WorkspaceBuild> => {
   const template = await getTemplate(workspace.template_id)
-  return startWorkspace({
-    workspaceId: workspace.id,
-    templateVersionId: template.active_version_id,
-  })
+  return startWorkspace(workspace.id, template.active_version_id)
 }
 
 export const getWorkspaceBuildParameters = async (
diff --git a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
index 8f08b9d8a5fff..79c8df57716e3 100644
--- a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
+++ b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
@@ -78,16 +78,11 @@ export const WorkspaceReadyPage = ({
     enabled: changeVersionDialogOpen,
   })
 
-  const [restartBuildError, setRestartBuildError] = useState<
-    Error | unknown | undefined
-  >(undefined)
-
-  const [isRestarting, setIsRestarting] = useState<boolean>(false)
-
-  const { mutate: restartWorkspace } = useRestartWorkspace(
-    setRestartBuildError,
-    setIsRestarting,
-  )
+  const {
+    mutate: restartWorkspace,
+    error: restartBuildError,
+    isLoading: isRestarting,
+  } = useRestartWorkspace()
 
   // keep banner machine in sync with workspace
   useEffect(() => {
@@ -136,7 +131,7 @@ export const WorkspaceReadyPage = ({
         workspace={workspace}
         handleStart={() => workspaceSend({ type: "START" })}
         handleStop={() => workspaceSend({ type: "STOP" })}
-        handleRestart={() => restartWorkspace(workspace.id)}
+        handleRestart={() => restartWorkspace(workspace)}
         handleDelete={() => workspaceSend({ type: "ASK_DELETE" })}
         handleUpdate={() => workspaceSend({ type: "UPDATE" })}
         handleCancel={() => workspaceSend({ type: "CANCEL" })}
diff --git a/site/src/pages/WorkspacePage/hooks.ts b/site/src/pages/WorkspacePage/hooks.ts
index 15f223a6a32a5..15da2079767bc 100644
--- a/site/src/pages/WorkspacePage/hooks.ts
+++ b/site/src/pages/WorkspacePage/hooks.ts
@@ -1,66 +1,8 @@
-import {
-  getWorkspaceBuildByNumber,
-  stopWorkspace,
-  startWorkspace,
-} from "api/api"
-import { delay } from "utils/delay"
-import { ProvisionerJob, WorkspaceBuild } from "api/typesGenerated"
+import { restartWorkspace } from "api/api"
 import { useMutation } from "@tanstack/react-query"
 
-export function waitForBuild(build: WorkspaceBuild) {
-  return new Promise((res, reject) => {
-    void (async () => {
-      let latestJobInfo: ProvisionerJob | undefined = undefined
-
-      while (latestJobInfo?.status !== "succeeded") {
-        const { job } = await getWorkspaceBuildByNumber(
-          build.workspace_owner_name,
-          build.workspace_name,
-          String(build.build_number),
-        )
-        latestJobInfo = job
-
-        if (
-          ["failed", "canceled"].some((status) =>
-            latestJobInfo?.status.includes(status),
-          )
-        ) {
-          return reject(latestJobInfo)
-        }
-
-        await delay(1000)
-      }
-
-      return res(latestJobInfo)
-    })()
-  })
-}
-
-export const useRestartWorkspace = (
-  setBuildError: (arg: Error | unknown | undefined) => void,
-  setLoading: (arg: boolean) => void,
-) => {
+export const useRestartWorkspace = () => {
   return useMutation({
-    mutationFn: stopWorkspace,
-    onMutate: () => setLoading(true),
-    onSuccess: async (data: WorkspaceBuild) => {
-      try {
-        await waitForBuild(data)
-        const startBuild = await startWorkspace({
-          workspaceId: data.workspace_id,
-          templateVersionId: data.template_version_id,
-        })
-        await waitForBuild(startBuild)
-
-        setLoading(false)
-      } catch (error) {
-        if ((error as WorkspaceBuild).status === "canceled") {
-          setLoading(false)
-          return
-        }
-        setBuildError(error)
-        setLoading(false)
-      }
-    },
+    mutationFn: restartWorkspace,
   })
 }
diff --git a/site/src/xServices/workspace/workspaceXService.ts b/site/src/xServices/workspace/workspaceXService.ts
index 623ba82003342..818b853960761 100644
--- a/site/src/xServices/workspace/workspaceXService.ts
+++ b/site/src/xServices/workspace/workspaceXService.ts
@@ -779,12 +779,11 @@ export const workspaceMachine = createMachine(
         },
       startWorkspace: (context) => async (send) => {
         if (context.workspace) {
-          const startWorkspacePromise = await API.startWorkspace({
-            workspaceId: context.workspace.id,
-            templateVersionId:
-              context.workspace.latest_build.template_version_id,
-            logLevel: context.createBuildLogLevel,
-          })
+          const startWorkspacePromise = await API.startWorkspace(
+            context.workspace.id,
+            context.workspace.latest_build.template_version_id,
+            context.createBuildLogLevel,
+          )
           send({ type: "REFRESH_TIMELINE" })
           return startWorkspacePromise
         } else {