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) => {