From bc73613fb160954fafd4e3c6e13b1afb20ef162c Mon Sep 17 00:00:00 2001
From: clean99 <xff9924@gmail.com>
Date: Tue, 21 Nov 2023 21:15:16 +0800
Subject: [PATCH] feat: update interface and fix close bug

---
 frontend/src/generateCode.ts | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/frontend/src/generateCode.ts b/frontend/src/generateCode.ts
index 2083958..b848b78 100644
--- a/frontend/src/generateCode.ts
+++ b/frontend/src/generateCode.ts
@@ -10,6 +10,7 @@ const STOP_MESSAGE = "Code generation stopped";
 export interface CodeGenerationParams {
   generationType: "create" | "update";
   image: string;
+  resultImg?: string;
   history?: string[];
   // isImageGenerationEnabled: boolean; // TODO: Merge with Settings type in types.ts
 }
@@ -49,13 +50,11 @@ export function generateCode(
     console.log("Connection closed", event.code, event.reason);
     if (event.code === USER_CLOSE_WEB_SOCKET_CODE) {
       toast.success(STOP_MESSAGE);
-      onComplete();
-    } else if (event.code === 1000) {
-      onComplete();
-    } else {
+    } else if (event.code !== 1000) {
       console.error("WebSocket error code", event);
       toast.error(ERROR_MESSAGE);
     }
+    onComplete();
   });
 
   ws.addEventListener("error", (error) => {