From 2a52095ce9e0acccab6acd4fd0b5d57f92609591 Mon Sep 17 00:00:00 2001
From: clean99 <xff9924@gmail.com>
Date: Tue, 21 Nov 2023 16:18:32 +0800
Subject: [PATCH] feat: add client close button

---
 frontend/src/App.tsx         | 18 ++++++++++++++++--
 frontend/src/generateCode.ts | 15 +++++++++++----
 frontend/src/types.ts        |  4 +++-
 3 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index bb71064..6bb0c80 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -1,4 +1,4 @@
-import { useState } from "react";
+import { useRef, useState } from "react";
 import ImageUpload from "./components/ImageUpload";
 import CodePreview from "./components/CodePreview";
 import Preview from "./components/Preview";
@@ -17,7 +17,7 @@ import { Textarea } from "@/components/ui/textarea";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "./components/ui/tabs";
 import CodeMirror from "./components/CodeMirror";
 import SettingsDialog from "./components/SettingsDialog";
-import { AppStatus, Settings } from "./types";
+import { AppStatus, Settings, USER_CLOSE_WEB_SOCKET_CODE } from "./types";
 import { IS_RUNNING_ON_CLOUD } from "./config";
 import { PicoBadge } from "./components/PicoBadge";
 import { OnboardingNote } from "./components/OnboardingNote";
@@ -43,6 +43,7 @@ function App() {
     },
     "setting"
   );
+  const wsRef = useRef<WebSocket>(null);
 
   const downloadCode = () => {
     // Create a blob from the generated code
@@ -69,6 +70,10 @@ function App() {
     setHistory([]);
   };
 
+  const stop = () => {
+    wsRef.current?.close?.(USER_CLOSE_WEB_SOCKET_CODE);
+  }
+
   function doGenerateCode(params: CodeGenerationParams) {
     setExecutionConsole([]);
     setAppState(AppStatus.CODING);
@@ -77,6 +82,7 @@ function App() {
     const updatedParams = { ...params, ...settings };
 
     generateCode(
+      wsRef,
       updatedParams,
       (token) => setGeneratedCode((prev) => prev + token),
       (code) => setGeneratedCode(code),
@@ -139,6 +145,14 @@ function App() {
                     <Spinner />
                     {executionConsole.slice(-1)[0]}
                   </div>
+                  <div className="flex mt-4 w-full">
+                    <Button
+                      onClick={stop}
+                      className="w-full"
+                    >
+                      Stop
+                    </Button>
+                  </div>
                   <CodePreview code={generatedCode} />
                 </div>
               )}
diff --git a/frontend/src/generateCode.ts b/frontend/src/generateCode.ts
index 6fea92d..2083958 100644
--- a/frontend/src/generateCode.ts
+++ b/frontend/src/generateCode.ts
@@ -1,9 +1,12 @@
 import toast from "react-hot-toast";
 import { WS_BACKEND_URL } from "./config";
+import { USER_CLOSE_WEB_SOCKET_CODE } from "./types";
 
 const ERROR_MESSAGE =
   "Error generating code. Check the Developer Console for details. Feel free to open a Github issue";
 
+const STOP_MESSAGE = "Code generation stopped";
+
 export interface CodeGenerationParams {
   generationType: "create" | "update";
   image: string;
@@ -12,6 +15,7 @@ export interface CodeGenerationParams {
 }
 
 export function generateCode(
+  wsRef: React.MutableRefObject<WebSocket | null>,
   params: CodeGenerationParams,
   onChange: (chunk: string) => void,
   onSetCode: (code: string) => void,
@@ -22,6 +26,7 @@ export function generateCode(
   console.log("Connecting to backend @ ", wsUrl);
 
   const ws = new WebSocket(wsUrl);
+  wsRef.current = ws;
 
   ws.addEventListener("open", () => {
     ws.send(JSON.stringify(params));
@@ -40,14 +45,16 @@ export function generateCode(
       toast.error(response.value);
     }
   });
-
   ws.addEventListener("close", (event) => {
     console.log("Connection closed", event.code, event.reason);
-    if (event.code != 1000) {
+    if (event.code === USER_CLOSE_WEB_SOCKET_CODE) {
+      toast.success(STOP_MESSAGE);
+      onComplete();
+    } else if (event.code === 1000) {
+      onComplete();
+    } else {
       console.error("WebSocket error code", event);
       toast.error(ERROR_MESSAGE);
-    } else {
-      onComplete();
     }
   });
 
diff --git a/frontend/src/types.ts b/frontend/src/types.ts
index 10131c0..8797195 100644
--- a/frontend/src/types.ts
+++ b/frontend/src/types.ts
@@ -9,4 +9,6 @@ export enum AppStatus {
   INITIAL = "INITIAL",
   CODING = "CODING",
   CODE_READY = "CODE_READY",
-}
\ No newline at end of file
+}
+
+export const USER_CLOSE_WEB_SOCKET_CODE = 4333;
\ No newline at end of file