From a939132dba75188043e755704d34207152526459 Mon Sep 17 00:00:00 2001
From: xswl <yanghao@ohdat.io>
Date: Fri, 7 Jul 2023 12:57:04 +0800
Subject: [PATCH]  waitmessage add skipMessageId

---
 example/customzoom.ts | 12 ++++++++----
 example/imagine.ts    |  1 +
 src/discord.ws.ts     | 12 +++++++++++-
 src/midjourney.ts     |  1 +
 4 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/example/customzoom.ts b/example/customzoom.ts
index 05429dc..4d85174 100644
--- a/example/customzoom.ts
+++ b/example/customzoom.ts
@@ -63,7 +63,11 @@ async function main() {
   console.log("Custom Zoom", zoomout2x);
   client.Close();
 }
-main().catch((err) => {
-  console.error(err);
-  process.exit(1);
-});
+main()
+  .then(() => {
+    console.log("done");
+  })
+  .catch((err) => {
+    console.error(err);
+    process.exit(1);
+  });
diff --git a/example/imagine.ts b/example/imagine.ts
index 0a31c04..d56972e 100644
--- a/example/imagine.ts
+++ b/example/imagine.ts
@@ -13,6 +13,7 @@ async function main() {
     ChannelId: <string>process.env.CHANNEL_ID,
     SalaiToken: <string>process.env.SALAI_TOKEN,
     Debug: true,
+    Ws: false,
   });
 
   const msg = await client.Imagine(
diff --git a/src/discord.ws.ts b/src/discord.ws.ts
index 6af5e0b..b836777 100644
--- a/src/discord.ws.ts
+++ b/src/discord.ws.ts
@@ -27,6 +27,7 @@ export class WsMessage {
   private event: Array<{ event: string; callback: (message: any) => void }> =
     [];
   private waitMjEvents: Map<string, WaitMjEvent> = new Map();
+  private skipMessageId: string[] = [];
   private reconnectTime: boolean[] = [];
   private heartbeatInterval = 0;
   public UserId = "";
@@ -470,7 +471,12 @@ export class WsMessage {
     this.waitMjEvents.set(nonce, { nonce });
     this.event.push({ event: nonce, callback: once });
   }
-
+  private removeSkipMessageId(messageId: string) {
+    const index = this.skipMessageId.findIndex((id) => id !== messageId);
+    if (index !== -1) {
+      this.skipMessageId.splice(index, 1);
+    }
+  }
   private removeWaitMjEvent(nonce: string) {
     this.waitMjEvents.delete(nonce);
   }
@@ -489,13 +495,16 @@ export class WsMessage {
     nonce,
     prompt,
     onmodal,
+    messageId,
     loading,
   }: {
     nonce: string;
     prompt?: string;
+    messageId?: string;
     onmodal?: OnModal;
     loading?: LoadingHandler;
   }) {
+    if (messageId) this.skipMessageId.push(messageId);
     return new Promise<MJMessage | null>((resolve, reject) => {
       const handleImageMessage = ({ message, error }: MJEmit) => {
         if (error) {
@@ -505,6 +514,7 @@ export class WsMessage {
         }
         if (message && message.progress === "done") {
           this.removeWaitMjEvent(nonce);
+          messageId && this.removeSkipMessageId(messageId);
           resolve(message);
           return;
         }
diff --git a/src/midjourney.ts b/src/midjourney.ts
index 9abf1f1..665ceac 100644
--- a/src/midjourney.ts
+++ b/src/midjourney.ts
@@ -254,6 +254,7 @@ export class Midjourney extends MidjourneyMessage {
       return await wsClient.waitImageMessage({
         nonce,
         loading,
+        messageId: msgId,
         prompt: content,
         onmodal: async (nonde, id) => {
           if (content === undefined || content === "") {