Merge pull request #63 from clean99/feat/add-reset-button

Feat: add stop button when generating code

frontend/src/App.tsx

@@ -1,4 +1,4 @@
-import { useState, useCallback } from "react";
+import { useRef, useState, useCallback } from "react";
 import ImageUpload from "./components/ImageUpload";
 import CodePreview from "./components/CodePreview";
 import Preview from "./components/Preview";
@@ -14,24 +14,23 @@ import {
   FaUndo,
 } from "react-icons/fa";
 import copy from "copy-to-clipboard";
+import toast from "react-hot-toast";
 import { Button } from "@/components/ui/button";
 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 { Settings } from "./types";
+import { AppState, Settings } from "./types";
 import { IS_RUNNING_ON_CLOUD } from "./config";
 import { PicoBadge } from "./components/PicoBadge";
 import { OnboardingNote } from "./components/OnboardingNote";
 import { usePersistedState } from "./hooks/usePersistedState";
 import { UrlInputSection } from "./components/UrlInputSection";
 import TermsOfServiceDialog from "./components/TermsOfServiceDialog";
-import toast from "react-hot-toast";
+import { USER_CLOSE_WEB_SOCKET_CODE } from "./constants";
 
 function App() {
-  const [appState, setAppState] = useState<"INITIAL" | "CODING" | "CODE_READY">(
-    "INITIAL"
-  );
+  const [appState, setAppState] = useState<AppState>(AppState.INITIAL);
   const [generatedCode, setGeneratedCode] = useState<string>("");
   const [referenceImages, setReferenceImages] = useState<string[]>([]);
   const [executionConsole, setExecutionConsole] = useState<string[]>([]);
@@ -46,6 +45,7 @@ function App() {
     },
     "setting"
   );
+  const wsRef = useRef<WebSocket>(null);
 
   const downloadCode = () => {
     // Create a blob from the generated code
@@ -65,26 +65,31 @@ function App() {
   };
 
   const reset = () => {
-    setAppState("INITIAL");
+    setAppState(AppState.INITIAL);
     setGeneratedCode("");
     setReferenceImages([]);
     setExecutionConsole([]);
     setHistory([]);
   };
 
+  const stop = () => {
+    wsRef.current?.close?.(USER_CLOSE_WEB_SOCKET_CODE);
+  };
+
   function doGenerateCode(params: CodeGenerationParams) {
     setExecutionConsole([]);
-    setAppState("CODING");
+    setAppState(AppState.CODING);
 
     // Merge settings with params
     const updatedParams = { ...params, ...settings };
 
     generateCode(
+      wsRef,
       updatedParams,
       (token) => setGeneratedCode((prev) => prev + token),
       (code) => setGeneratedCode(code),
       (line) => setExecutionConsole((prev) => [...prev, line]),
-      () => setAppState("CODE_READY")
+      () => setAppState(AppState.CODE_READY)
     );
   }
 
@@ -130,7 +135,7 @@ function App() {
             <h1 className="text-2xl ">Screenshot to Code</h1>
             <SettingsDialog settings={settings} setSettings={setSettings} />
           </div>
-          {appState === "INITIAL" && (
+          {appState === AppState.INITIAL && (
             <h2 className="text-sm text-gray-500 mb-2">
               Drag & drop a screenshot to get started.
             </h2>
@@ -138,20 +143,26 @@ function App() {
 
           {IS_RUNNING_ON_CLOUD && !settings.openAiApiKey && <OnboardingNote />}
 
-          {(appState === "CODING" || appState === "CODE_READY") && (
+          {(appState === AppState.CODING ||
+            appState === AppState.CODE_READY) && (
             <>
               {/* Show code preview only when coding */}
-              {appState === "CODING" && (
+              {appState === AppState.CODING && (
                 <div className="flex flex-col">
                   <div className="flex items-center gap-x-1">
                     <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>
               )}
 
-              {appState === "CODE_READY" && (
+              {appState === AppState.CODE_READY && (
                 <div>
                   <div className="grid w-full gap-2">
                     <Textarea
@@ -184,7 +195,7 @@ function App() {
                 <div className="flex flex-col">
                   <div
                     className={classNames({
-                      "scanning relative": appState === "CODING",
+                      "scanning relative": appState === AppState.CODING,
                     })}
                   >
                     <img
@@ -217,7 +228,7 @@ function App() {
       </div>
 
       <main className="py-2 lg:pl-96">
-        {appState === "INITIAL" && (
+        {appState === AppState.INITIAL && (
           <div className="flex flex-col justify-center items-center gap-y-10">
             <ImageUpload setReferenceImages={doCreate} />
             <UrlInputSection
@@ -227,7 +238,7 @@ function App() {
           </div>
         )}
 
-        {(appState === "CODING" || appState === "CODE_READY") && (
+        {(appState === AppState.CODING || appState === AppState.CODE_READY) && (
           <div className="ml-4">
             <Tabs defaultValue="desktop">
               <div className="flex justify-end mr-8 mb-4">

frontend/src/constants.ts

@@ -0,0 +1 @@
+export const USER_CLOSE_WEB_SOCKET_CODE = 4333;

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 "./constants";
 
 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,13 +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();
     }
   });

frontend/src/types.ts

@@ -4,3 +4,9 @@ export interface Settings {
   isImageGenerationEnabled: boolean;
   editorTheme: string;
 }
+
+export enum AppState {
+  INITIAL = "INITIAL",
+  CODING = "CODING",
+  CODE_READY = "CODE_READY",
+}