diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index 75639ee..27937ca 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -130,10 +130,21 @@ function App() {
     setIsImportedFromCode(false);
   };
 
-  const stop = () => {
+  const cancelCodeGeneration = () => {
     wsRef.current?.close?.(USER_CLOSE_WEB_SOCKET_CODE);
     // make sure stop can correct the state even if the websocket is already closed
-    setAppState(AppState.CODE_READY);
+    cancelCodeGenerationAndReset();
+  };
+
+  const cancelCodeGenerationAndReset = () => {
+    // When this is the first version, reset the entire app state
+    if (currentVersion === null) {
+      reset();
+    } else {
+      // Otherwise, revert to the last version
+      setGeneratedCode(appHistory[currentVersion].code);
+      setAppState(AppState.CODE_READY);
+    }
   };
 
   function doGenerateCode(
@@ -189,6 +200,11 @@ function App() {
         }
       },
       (line) => setExecutionConsole((prev) => [...prev, line]),
+      // On cancel
+      () => {
+        cancelCodeGenerationAndReset();
+      },
+      // On complete
       () => {
         setAppState(AppState.CODE_READY);
       }
@@ -343,10 +359,10 @@ function App() {
                   </div>
                   <div className="flex mt-4 w-full">
                     <Button
-                      onClick={stop}
+                      onClick={cancelCodeGeneration}
                       className="w-full dark:text-white dark:bg-gray-700"
                     >
-                      Stop
+                      Cancel
                     </Button>
                   </div>
                   <CodePreview code={generatedCode} />
diff --git a/frontend/src/generateCode.ts b/frontend/src/generateCode.ts
index 96fbdc1..55523c4 100644
--- a/frontend/src/generateCode.ts
+++ b/frontend/src/generateCode.ts
@@ -6,7 +6,7 @@ import { FullGenerationSettings } from "./types";
 const ERROR_MESSAGE =
   "Error generating code. Check the Developer Console AND the backend logs for details. Feel free to open a Github issue.";
 
-const STOP_MESSAGE = "Code generation stopped";
+const CANCEL_MESSAGE = "Code generation cancelled";
 
 export function generateCode(
   wsRef: React.MutableRefObject<WebSocket | null>,
@@ -14,6 +14,7 @@ export function generateCode(
   onChange: (chunk: string) => void,
   onSetCode: (code: string) => void,
   onStatusUpdate: (status: string) => void,
+  onCancel: () => void,
   onComplete: () => void
 ) {
   const wsUrl = `${WS_BACKEND_URL}/generate-code`;
@@ -39,15 +40,18 @@ export function generateCode(
       toast.error(response.value);
     }
   });
+
   ws.addEventListener("close", (event) => {
     console.log("Connection closed", event.code, event.reason);
     if (event.code === USER_CLOSE_WEB_SOCKET_CODE) {
-      toast.success(STOP_MESSAGE);
+      toast.success(CANCEL_MESSAGE);
+      onCancel();
     } else if (event.code !== 1000) {
       console.error("WebSocket error code", event);
       toast.error(ERROR_MESSAGE);
+    } else {
+      onComplete();
     }
-    onComplete();
   });
 
   ws.addEventListener("error", (error) => {