coder · BrunoQuaresma · May 15, 2023 · May 15, 2023 · May 15, 2023 · May 15, 2023
diff --git a/site/src/components/WorkspaceActions/Buttons.tsx b/site/src/components/WorkspaceActions/Buttons.tsx
@@ -59,7 +59,12 @@ export const RestartButton: FC<PropsWithChildren<WorkspaceAction>> = ({
   const { t } = useTranslation("workspacePage")
 
   return (
-    <Button size="small" startIcon={<ReplayIcon />} onClick={handleAction}>
+    <Button
+      size="small"
+      startIcon={<ReplayIcon />}
+      onClick={handleAction}
+      data-testid="workspace-restart-button"
+    >
       {t("actionButton.restart")}
     </Button>
   )

diff --git a/site/src/pages/WorkspacePage/WorkspacePage.test.tsx b/site/src/pages/WorkspacePage/WorkspacePage.test.tsx
@@ -20,6 +20,7 @@ import {
   MockDeletedWorkspace,
   MockBuilds,
   MockTemplateVersion3,
+  MockUser,
 } from "testHelpers/entities"
 import * as api from "../../api/api"
 import { Workspace } from "../../api/typesGenerated"
@@ -161,10 +162,41 @@ describe("WorkspacePage", () => {
       .spyOn(api, "stopWorkspace")
       .mockResolvedValueOnce(MockWorkspaceBuild)
 
-    await testButton("Restart", stopWorkspaceMock)
+    // Render
+    await renderWorkspacePage()
+
+    // Actions
+    const user = userEvent.setup()
+    await user.click(screen.getByTestId("workspace-restart-button"))
+    const confirmButton = await screen.findByTestId("confirm-button")
+    await user.click(confirmButton)
+
+    // Assertions
+    await waitFor(() => {
+      expect(stopWorkspaceMock).toBeCalled()
+    })
+  })
+
+  it("requests a stop without confirmation when the user presses Restart", async () => {
+    const stopWorkspaceMock = jest
+      .spyOn(api, "stopWorkspace")
+      .mockResolvedValueOnce(MockWorkspaceBuild)
+    window.localStorage.setItem(
+      `${MockUser.id}_ignoredWarnings`,
+      JSON.stringify({ restart: new Date().toISOString() }),
+    )
+
+    // Render
+    await renderWorkspacePage()
 
-    const button = await screen.findByText("Restarting")
-    expect(button).toBeInTheDocument()
+    // Actions
+    const user = userEvent.setup()
+    await user.click(screen.getByTestId("workspace-restart-button"))
+
+    // Assertions
+    await waitFor(() => {
+      expect(stopWorkspaceMock).toBeCalled()
+    })
   })
 
   it("requests cancellation when the user presses Cancel", async () => {

diff --git a/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx b/site/src/pages/WorkspacePage/WorkspaceReadyPage.tsx
@@ -3,7 +3,7 @@ import { ProvisionerJobLog } from "api/typesGenerated"
 import { useDashboard } from "components/Dashboard/DashboardProvider"
 import dayjs from "dayjs"
 import { useFeatureVisibility } from "hooks/useFeatureVisibility"
-import { useEffect, useState } from "react"
+import { FC, useEffect, useState } from "react"
 import { Helmet } from "react-helmet-async"
 import { useTranslation } from "react-i18next"
 import { useNavigate } from "react-router-dom"
@@ -31,7 +31,13 @@ import { ChangeVersionDialog } from "./ChangeVersionDialog"
 import { useQuery } from "@tanstack/react-query"
 import { getTemplateVersions } from "api/api"
 import { useRestartWorkspace } from "./hooks"
-import { ConfirmDialog } from "components/Dialogs/ConfirmDialog/ConfirmDialog"
+import {
+  ConfirmDialog,
+  ConfirmDialogProps,
+} from "components/Dialogs/ConfirmDialog/ConfirmDialog"
+import { useMe } from "hooks/useMe"
+import Checkbox from "@mui/material/Checkbox"
+import FormControlLabel from "@mui/material/FormControlLabel"
 
 interface WorkspaceReadyPageProps {
   workspaceState: StateFrom<typeof workspaceMachine>
@@ -79,6 +85,9 @@ export const WorkspaceReadyPage = ({
     enabled: changeVersionDialogOpen,
   })
   const [isConfirmingUpdate, setIsConfirmingUpdate] = useState(false)
+  const [isConfirmingRestart, setIsConfirmingRestart] = useState(false)
+  const user = useMe()
+  const { isWarningIgnored, ignoreWarning } = useIgnoreWarnings(user.id)
 
   const {
     mutate: restartWorkspace,
@@ -133,9 +142,21 @@ export const WorkspaceReadyPage = ({
         workspace={workspace}
         handleStart={() => workspaceSend({ type: "START" })}
         handleStop={() => workspaceSend({ type: "STOP" })}
-        handleRestart={() => restartWorkspace(workspace)}
         handleDelete={() => workspaceSend({ type: "ASK_DELETE" })}
-        handleUpdate={() => setIsConfirmingUpdate(true)}
+        handleRestart={() => {
+          if (isWarningIgnored("restart")) {
+            restartWorkspace(workspace)
+          } else {
+            setIsConfirmingRestart(true)
+          }
+        }}
+        handleUpdate={() => {
+          if (isWarningIgnored("update")) {
+            workspaceSend({ type: "UPDATE" })
+          } else {
+            setIsConfirmingUpdate(true)
+          }
+        }}
         handleCancel={() => workspaceSend({ type: "CANCEL" })}
         handleSettings={() => navigate("settings")}
         handleBuildRetry={() => workspaceSend({ type: "RETRY_BUILD" })}
@@ -202,11 +223,12 @@ export const WorkspaceReadyPage = ({
           })
         }}
       />
-      <ConfirmDialog
-        type="info"
-        hideCancel={false}
+      <WarningDialog
         open={isConfirmingUpdate}
-        onConfirm={() => {
+        onConfirm={(shouldIgnore) => {
+          if (shouldIgnore) {
+            ignoreWarning("update")
+          }
           workspaceSend({ type: "UPDATE" })
           setIsConfirmingUpdate(false)
         }}
@@ -215,6 +237,93 @@ export const WorkspaceReadyPage = ({
         confirmText="Update"
         description="Are you sure you want to update your workspace? Updating your workspace will stop all running processes and delete non-persistent data."
       />
+
+      <WarningDialog
+        open={isConfirmingRestart}
+        onConfirm={(shouldIgnore) => {
+          if (shouldIgnore) {
+            ignoreWarning("restart")
+          }
+          restartWorkspace(workspace)
+          setIsConfirmingRestart(false)
+        }}
+        onClose={() => setIsConfirmingRestart(false)}
+        title="Confirm restart"
+        confirmText="Restart"
+        description="Are you sure you want to restart your workspace? Updating your workspace will stop all running processes and delete non-persistent data."
+      />
     </>
   )
 }
+
+type IgnoredWarnings = Record<string, string>
+
+const useIgnoreWarnings = (prefix: string) => {
+  const ignoredWarningsJSON = localStorage.getItem(`${prefix}_ignoredWarnings`)
+  let ignoredWarnings: IgnoredWarnings | undefined
+  if (ignoredWarningsJSON) {
+    ignoredWarnings = JSON.parse(ignoredWarningsJSON)
+  }
+
+  const isWarningIgnored = (warningId: string) => {
+    return Boolean(ignoredWarnings?.[warningId])
+  }
+
+  const ignoreWarning = (warningId: string) => {
+    if (!ignoredWarnings) {
+      ignoredWarnings = {}
+    }
+    ignoredWarnings[warningId] = new Date().toISOString()
+    localStorage.setItem(
+      `${prefix}_ignoredWarnings`,
+      JSON.stringify(ignoredWarnings),
+    )
+  }
+
+  return {
+    isWarningIgnored,
+    ignoreWarning,
+  }
+}
+
+const WarningDialog: FC<
+  Pick<
+    ConfirmDialogProps,
+    "open" | "onClose" | "title" | "confirmText" | "description"
+  > & { onConfirm: (shouldIgnore: boolean) => void }
+> = ({ open, onConfirm, onClose, title, confirmText, description }) => {
+  const [shouldIgnore, setShouldIgnore] = useState(false)
+
+  return (
+    <ConfirmDialog
+      type="info"
+      hideCancel={false}
+      open={open}
+      onConfirm={() => {
+        onConfirm(shouldIgnore)
+      }}
+      onClose={onClose}
+      title={title}
+      confirmText={confirmText}
+      description={
+        <>
+          <div>{description}</div>
+          <FormControlLabel
+            sx={{
+              marginTop: 2,
+            }}
+            control={
+              <Checkbox
+                size="small"
+                onChange={(e) => {
+                  setShouldIgnore(e.target.checked)
+                }}
+              />
+            }
+            label="Don't show me this message again"
+          />
+        </>
+      }
+    />
+  )
+}